"""Pydantic models for deprecated and remapped configuration parameters"""

from typing import Any

from pydantic import BaseModel, Field, field_validator

from axolotl.utils.logging import get_logger

LOG = get_logger(__name__)


class DeprecatedParameters(BaseModel):
    """configurations that are deprecated"""

    max_packed_sequence_len: int | None = None
    rope_scaling: Any | None = None
    noisy_embedding_alpha: float | None = None
    dpo_beta: float | None = None
    evaluation_strategy: str | None = None
    eval_table_size: int | None = None
    eval_max_new_tokens: int | None = None
    dpo_use_logits_to_keep: bool | None = None
    dpo_generate_during_eval: bool | None = None
    dpo_norm_loss: bool | None = None
    rpo_alpha: float | None = None

    @field_validator("max_packed_sequence_len")
    @classmethod
    def validate_max_packed_sequence_len(cls, max_packed_sequence_len):
        if max_packed_sequence_len:
            raise DeprecationWarning("`max_packed_sequence_len` is no longer supported")
        return max_packed_sequence_len

    @field_validator("rope_scaling")
    @classmethod
    def validate_rope_scaling(cls, rope_scaling):
        if rope_scaling:
            raise DeprecationWarning(
                "`rope_scaling` is no longer supported, it should now be be a key under `model_config`"
            )
        return rope_scaling

    @field_validator("noisy_embedding_alpha")
    @classmethod
    def validate_noisy_embedding_alpha(cls, noisy_embedding_alpha):
        if noisy_embedding_alpha:
            LOG.warning("noisy_embedding_alpha is deprecated, use neftune_noise_alpha")
        return noisy_embedding_alpha

    @field_validator("dpo_beta")
    @classmethod
    def validate_dpo_beta(cls, dpo_beta):
        if dpo_beta is not None:
            LOG.warning("dpo_beta is deprecated, use rl_beta instead")
        return dpo_beta

    @field_validator("evaluation_strategy")
    @classmethod
    def validate_evaluation_strategy(cls, evaluation_strategy):
        if evaluation_strategy is not None:
            LOG.warning("evaluation_strategy is deprecated, use eval_strategy instead")
        return evaluation_strategy

    @field_validator("eval_table_size")
    @classmethod
    def validate_eval_table_size(cls, eval_table_size):
        if eval_table_size is not None:
            LOG.warning(
                "eval_table_size is deprecated and superseded by generate_samples config. "
                "Please use generate_samples: true and num_generation_samples instead. "
                "The LogPredictionCallback is replaced by the new sample generation feature."
            )
        return eval_table_size

    @field_validator("eval_max_new_tokens")
    @classmethod
    def validate_eval_max_new_tokens(cls, eval_max_new_tokens):
        if eval_max_new_tokens is not None:
            LOG.warning(
                "eval_max_new_tokens is deprecated and superseded by generate_samples config. "
                "Please use generation_max_new_tokens instead."
            )
        return eval_max_new_tokens

    @field_validator("dpo_use_logits_to_keep")
    @classmethod
    def validate_dpo_use_logits_to_keep(cls, dpo_use_logits_to_keep):
        if dpo_use_logits_to_keep is not None:
            raise DeprecationWarning(
                "`dpo_use_logits_to_keep` is no longer supported, "
                "it has been removed in TRL >= 0.29.0"
            )
        return dpo_use_logits_to_keep

    @field_validator("dpo_generate_during_eval")
    @classmethod
    def validate_dpo_generate_during_eval(cls, dpo_generate_during_eval):
        if dpo_generate_during_eval is not None:
            raise DeprecationWarning(
                "`dpo_generate_during_eval` is no longer supported, "
                "it has been removed in TRL >= 0.29.0"
            )
        return dpo_generate_during_eval

    @field_validator("dpo_norm_loss")
    @classmethod
    def validate_dpo_norm_loss(cls, dpo_norm_loss):
        if dpo_norm_loss is not None:
            raise DeprecationWarning(
                "`dpo_norm_loss` is no longer supported, "
                "due to breaking changes in TRL >= 0.29.0"
            )
        return dpo_norm_loss

    @field_validator("rpo_alpha")
    @classmethod
    def validate_rpo_alpha(cls, rpo_alpha):
        if rpo_alpha is not None:
            raise DeprecationWarning(
                "`rpo_alpha` has been deprecated in TRL >= 0.29.0, "
                "and now requires passing multiple loss types, which is not yet supported by Axolotl."
            )  # TODO: change this warning once multiple dpo loss types are supported.
        return rpo_alpha


class RemappedParameters(BaseModel):
    """Parameters that have been remapped to other names"""

    overrides_of_model_config: dict[str, Any] | None = Field(
        default=None,
        alias="model_config",
        json_schema_extra={
            "description": "optional overrides to the base model configuration"
        },
    )
    overrides_of_model_kwargs: dict[str, Any] | None = Field(
        default=None,
        alias="model_kwargs",
        json_schema_extra={
            "description": "optional overrides the base model loading from_pretrained"
        },
    )
    type_of_model: str | None = Field(
        default=None,
        alias="model_type",
        json_schema_extra={
            "description": "If you want to specify the type of model to load, AutoModelForCausalLM is a good choice too"
        },
    )
    revision_of_model: str | None = Field(
        default=None,
        alias="model_revision",
        json_schema_extra={
            "description": "You can specify to choose a specific model revision from huggingface hub"
        },
    )
