"""The base architecture.

This is the original model before we switched to the multiple architecture arrangement.
It is a single architecture which implements many different functionalities. Other
architectures can import components from here to reuse them, and over time we should
refactor this architecture to improve reusability.
"""

from __future__ import annotations

from typing import TYPE_CHECKING, Any, Literal

from tabpfn.architectures.base.config import ModelConfig
from tabpfn.architectures.base.encoders import (
    InputNormalizationEncoderStep,
    LinearInputEncoderStep,
    MLPInputEncoderStep,
    MulticlassClassificationTargetEncoder,
    NanHandlingEncoderStep,
    RemoveDuplicateFeaturesEncoderStep,
    RemoveEmptyFeaturesEncoderStep,
    SeqEncStep,
    SequentialEncoder,
    VariableNumFeaturesEncoderStep,
)
from tabpfn.architectures.base.transformer import PerFeatureTransformer

if TYPE_CHECKING:
    from torch import nn

    from tabpfn.architectures.interface import ArchitectureConfig


def parse_config(config: dict[str, Any]) -> tuple[ArchitectureConfig, dict[str, Any]]:
    """Parse the config dict into a ModelConfig, a subclass of ArchitectureConfig.

    This method implements the interface defined in
    tabpfn.architectures.interface.ArchitectureModule.parse_config().

    Unrecognised keys should be ignored during parsing, and returned in the `unused
    config items` dict.

    Args:
        config: Config dict to parse. This function should use Pydantic to
            verify that it matches the expected schema.

    Returns: a tuple (the parsed config, dict containing unused config items)

    Raises:
        pydantic.ValidationError: one or more of the values have the wrong type
    """
    upgraded_dict = ModelConfig.upgrade_config(config)
    parsed_config = ModelConfig(**upgraded_dict)
    return parsed_config, parsed_config.get_unused_config(upgraded_dict)


def get_architecture(
    config: ArchitectureConfig,
    *,
    n_out: int,
    cache_trainset_representation: bool,
) -> PerFeatureTransformer:
    """Construct the base architecture following the given config.

    This factory method implements the interface defined in
    tabpfn.architectures.interface.ArchitectureModule.get_architecture().

    Args:
        config: The config returned by parse_config(). This method should use a
            runtime isinstance() check to downcast the config to this architecture's
            specific config class.
        n_out: The number of output classes that the model should predict.
        cache_trainset_representation: If True, the model should be configured to
            cache the training data during inference to improve speed.

    Returns: the constructed architecture
    """
    assert isinstance(config, ModelConfig)
    return PerFeatureTransformer(
        config=config,
        # Things that were explicitly passed inside `build_model()`
        encoder=get_encoder(
            num_features=config.features_per_group,
            embedding_size=config.emsize,
            remove_empty_features=config.remove_empty_features,
            remove_duplicate_features=config.remove_duplicate_features,
            nan_handling_enabled=config.nan_handling_enabled,
            normalize_on_train_only=config.normalize_on_train_only,
            normalize_to_ranking=config.normalize_to_ranking,
            normalize_x=config.normalize_x,
            remove_outliers=config.remove_outliers,
            normalize_by_used_features=config.normalize_by_used_features,
            encoder_use_bias=config.encoder_use_bias,
            encoder_type=config.encoder_type,
            encoder_mlp_hidden_dim=config.encoder_mlp_hidden_dim,
            encoder_mlp_num_layers=config.encoder_mlp_num_layers,
        ),
        y_encoder=get_y_encoder(
            num_inputs=1,
            embedding_size=config.emsize,
            nan_handling_y_encoder=config.nan_handling_y_encoder,
            max_num_classes=config.max_num_classes,
        ),
        cache_trainset_representation=cache_trainset_representation,
        use_encoder_compression_layer=False,
        n_out=n_out,
        #
        # These are things that had default values from config.get() but were not
        # present in any config.
        layer_norm_with_elementwise_affine=False,
    )


def get_encoder(  # noqa: PLR0913
    *,
    num_features: int,
    embedding_size: int,
    remove_empty_features: bool,
    remove_duplicate_features: bool,
    nan_handling_enabled: bool,
    normalize_on_train_only: bool,
    normalize_to_ranking: bool,
    normalize_x: bool,
    remove_outliers: bool,
    normalize_by_used_features: bool,
    encoder_use_bias: bool,
    encoder_type: Literal["linear", "mlp"] = "linear",
    encoder_mlp_hidden_dim: int | None = None,
    encoder_mlp_num_layers: int = 2,
) -> nn.Module:
    inputs_to_merge = {"main": {"dim": num_features}}

    encoder_steps: list[SeqEncStep] = []
    if remove_empty_features:
        encoder_steps += [RemoveEmptyFeaturesEncoderStep()]

    if remove_duplicate_features:
        encoder_steps += [RemoveDuplicateFeaturesEncoderStep()]

    encoder_steps += [NanHandlingEncoderStep(keep_nans=nan_handling_enabled)]

    if nan_handling_enabled:
        inputs_to_merge["nan_indicators"] = {"dim": num_features}

        encoder_steps += [
            VariableNumFeaturesEncoderStep(
                num_features=num_features,
                normalize_by_used_features=False,
                in_keys=["nan_indicators"],
                out_keys=["nan_indicators"],
            ),
        ]

    encoder_steps += [
        InputNormalizationEncoderStep(
            normalize_on_train_only=normalize_on_train_only,
            normalize_to_ranking=normalize_to_ranking,
            normalize_x=normalize_x,
            remove_outliers=remove_outliers,
        ),
    ]

    encoder_steps += [
        VariableNumFeaturesEncoderStep(
            num_features=num_features,
            normalize_by_used_features=normalize_by_used_features,
        ),
    ]

    num_input_features = sum(i["dim"] for i in inputs_to_merge.values())
    if encoder_type == "mlp":
        encoder_steps += [
            MLPInputEncoderStep(
                num_features=num_input_features,
                emsize=embedding_size,
                hidden_dim=encoder_mlp_hidden_dim,
                activation="gelu",
                num_layers=encoder_mlp_num_layers,
                bias=encoder_use_bias,
                in_keys=tuple(inputs_to_merge),
                out_keys=("output",),
            ),
        ]
    elif encoder_type == "linear":
        encoder_steps += [
            LinearInputEncoderStep(
                num_features=num_input_features,
                emsize=embedding_size,
                bias=encoder_use_bias,
                in_keys=tuple(inputs_to_merge),
                out_keys=("output",),
            ),
        ]
    else:
        raise ValueError(
            f"Invalid encoder type: {encoder_type} (expected 'linear' or 'mlp')"
        )

    return SequentialEncoder(*encoder_steps, output_key="output")


def get_y_encoder(
    *,
    num_inputs: int,
    embedding_size: int,
    nan_handling_y_encoder: bool,
    max_num_classes: int,
) -> nn.Module:
    steps: list[SeqEncStep] = []
    inputs_to_merge = [{"name": "main", "dim": num_inputs}]
    if nan_handling_y_encoder:
        steps += [NanHandlingEncoderStep()]
        inputs_to_merge += [{"name": "nan_indicators", "dim": num_inputs}]

    if max_num_classes >= 2:
        steps += [MulticlassClassificationTargetEncoder()]

    steps += [
        LinearInputEncoderStep(
            num_features=sum([i["dim"] for i in inputs_to_merge]),  # type: ignore
            emsize=embedding_size,
            in_keys=tuple(i["name"] for i in inputs_to_merge),  # type: ignore
            out_keys=("output",),
        ),
    ]
    return SequentialEncoder(*steps, output_key="output")
