# Copyright 2024 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# pylint: disable=protected-access
"""Classes for model tuning based on distillation."""

from typing import Optional

from google.cloud.aiplatform.utils import gcs_utils
from google.cloud.aiplatform_v1beta1.types import (
    tuning_job as gca_tuning_job_types,
)

from vertexai import generative_models
from vertexai.tuning import _tuning


def distill_model(
    *,
    student_model: str,
    teacher_model: str,
    training_dataset: str,
    validation_dataset: Optional[str] = None,
    epoch_count: Optional[int] = None,
    learning_rate_multiplier: Optional[float] = None,
    tuned_model_display_name: Optional[str] = None,
) -> "DistillationJob":
    """Tunes a model using distillation.

    Args:
        student_model:
            Student model name for distillation, e.g., "gemma-1.1-2b-it".
        teacher_model:
            Teacher model name for distillation, e.g., "gemini-1.5-flash-001".
        training_dataset: Cloud Storage path to file containing training dataset for distillation.
            The dataset should be in JSONL format.
        validation_dataset: Cloud Storage path to file containing validation dataset for distillation.
            The dataset should be in JSONL format.
        epoch_count: Number of training epoches for this tuning job.
        learning_rate_multiplier: Learning rate multiplier for tuning.
        tuned_model_display_name: The display name of the
            [TunedModel][google.cloud.aiplatform.v1.Model]. The name can
            be up to 128 characters long and can consist of any UTF-8 characters.

    Returns:
        A `TuningJob` object.
    """

    if isinstance(student_model, generative_models.GenerativeModel):
        student_model = student_model._prediction_resource_name

    student_model = student_model.rpartition("/")[-1]
    teacher_model = teacher_model.rpartition("/")[-1]

    pipeline_root = (
        gcs_utils.create_gcs_bucket_for_pipeline_artifacts_if_it_does_not_exist()
    )

    distillation_spec = gca_tuning_job_types.DistillationSpec(
        student_model=student_model,
        base_teacher_model=teacher_model,
        training_dataset_uri=training_dataset,
        validation_dataset_uri=validation_dataset,
        hyper_parameters=gca_tuning_job_types.DistillationHyperParameters(
            epoch_count=epoch_count,
            learning_rate_multiplier=learning_rate_multiplier,
        ),
        pipeline_root_directory=pipeline_root,
    )

    return DistillationJob._create(  # pylint: disable=protected-access
        base_model=None,
        tuning_spec=distillation_spec,
        tuned_model_display_name=tuned_model_display_name,
    )


class DistillationJob(_tuning.TuningJob):
    pass