"""A TabPFN regressor that finetunes the underlying model for a single task.

This module provides the FinetunedTabPFNRegressor class, which wraps TabPFN
and allows fine-tuning on a specific dataset using the familiar scikit-learn
.fit() and .predict() API.
"""

from __future__ import annotations

import logging
from pathlib import Path
from typing import TYPE_CHECKING, Any, Literal
from typing_extensions import override

import numpy as np
import torch
from sklearn.base import RegressorMixin
from sklearn.metrics import mean_squared_error
from sklearn.utils.validation import check_is_fitted

from tabpfn import TabPFNRegressor
from tabpfn.finetuning.finetuned_base import EvalResult, FinetunedTabPFNBase
from tabpfn.finetuning.train_util import clone_model_for_evaluation
from tabpfn.model_loading import get_n_out

logger = logging.getLogger(__name__)

if TYPE_CHECKING:
    from tabpfn.constants import XType, YType
    from tabpfn.finetuning.data_util import RegressorBatch
    from tabpfn.regressor import RegressionResultType


def _compute_regression_loss(  # noqa: C901
    *,
    logits_BQL: torch.Tensor,
    targets_BQ: torch.Tensor,
    bardist_loss_fn: Any,
    ce_loss_weight: float = 1.0,
    crps_loss_weight: float = 0.0,
    crls_loss_weight: float = 0.0,
    mse_loss_weight: float = 0.0,
    mse_loss_clip: float | None = None,
    mae_loss_weight: float = 0.0,
    mae_loss_clip: float | None = None,
) -> torch.Tensor:
    """Compute the regression training loss from bar distribution and auxiliary terms.

    Shapes suffixes:
        B=batch * estimators, L=logits, Q=n_queries.

    Args:
        logits_BQL: Bar distribution logits of shape (B*E, Q, L).
        targets_BQ: Regression targets of shape (B*E, Q).
        bardist_loss_fn: Bar distribution loss function (callable) which also
            exposes a `.mean()` method for converting bar logits to mean
            predictions.
        ce_loss_weight: Weight for the bar distribution negative log-likelihood term
            (cross-entropy-like). Set to 0.0 to disable.
        crps_loss_weight: Weight for a continuous ranked probability score (CRPS) term
            computed on ordered bar probabilities (squared CDF error). Set to 0.0 to
            disable.
        crls_loss_weight: Weight for a continuous ranked logarithmic score (CRLS) term
            computed on ordered bar probabilities (logarithmic score on cumulative
            probabilities). Set to 0.0 to disable.
        mse_loss_weight: Weight for an auxiliary MSE term computed on the mean decoded
            prediction. Set to 0.0 to disable.
        mse_loss_clip: Optional upper bound for the auxiliary mean-MSE term.
        mae_loss_weight: Weight for an auxiliary MAE term computed on the mean
            decoded prediction. Set to 0.0 to disable.
        mae_loss_clip: Optional upper bound for the auxiliary mean-MAE term.

    Returns:
        A scalar loss tensor.
    """
    weights_to_validate = {
        "ce_loss_weight": ce_loss_weight,
        "crps_loss_weight": crps_loss_weight,
        "crls_loss_weight": crls_loss_weight,
        "mse_loss_weight": mse_loss_weight,
        "mae_loss_weight": mae_loss_weight,
    }
    for weight_name, weight in weights_to_validate.items():
        if weight < 0.0:
            raise ValueError(f"{weight_name} must be >= 0.0")

    total_loss = torch.tensor(0.0, device=logits_BQL.device)
    valid_mask_BQ = ~torch.isnan(targets_BQ)

    if ce_loss_weight > 0.0:
        ce_losses_BQ = bardist_loss_fn(logits_BQL, targets_BQ)
        total_loss = total_loss + ce_loss_weight * ce_losses_BQ.mean()

    if crps_loss_weight > 0.0:
        crps_loss = _ranked_probability_score_loss_from_bar_logits(
            logits_BQL=logits_BQL,
            targets_BQ=targets_BQ,
            bardist_loss_fn=bardist_loss_fn,
            loss_type="crps",
        )
        total_loss = total_loss + crps_loss_weight * crps_loss

    if crls_loss_weight > 0.0:
        crls_loss = _ranked_probability_score_loss_from_bar_logits(
            logits_BQL=logits_BQL,
            targets_BQ=targets_BQ,
            bardist_loss_fn=bardist_loss_fn,
            loss_type="crls",
        )
        total_loss = total_loss + crls_loss_weight * crls_loss

    if mse_loss_weight > 0.0 or mae_loss_weight > 0.0:
        predictions_mean_BQ = bardist_loss_fn.mean(logits_BQL)
        diffs_BQ = predictions_mean_BQ - targets_BQ

        if mse_loss_weight > 0.0:
            mse_terms_BQ = torch.where(
                valid_mask_BQ, diffs_BQ.square(), torch.zeros_like(diffs_BQ)
            )
            if mse_loss_clip is not None:
                mse_terms_BQ = mse_terms_BQ.clamp(max=mse_loss_clip)
            total_loss = total_loss + mse_loss_weight * mse_terms_BQ.mean()

        if mae_loss_weight > 0.0:
            mae_terms_BQ = torch.where(
                valid_mask_BQ, diffs_BQ.abs(), torch.zeros_like(diffs_BQ)
            )
            if mae_loss_clip is not None:
                mae_terms_BQ = mae_terms_BQ.clamp(max=mae_loss_clip)
            total_loss = total_loss + mae_loss_weight * mae_terms_BQ.mean()

    return total_loss


def _ranked_probability_score_loss_from_bar_logits(
    *,
    logits_BQL: torch.Tensor,
    targets_BQ: torch.Tensor,
    bardist_loss_fn: Any,
    loss_type: Literal["crps", "crls"] = "crps",
) -> torch.Tensor:
    """Compute a ranked-probability loss from bar distribution logits.

    This implements scoring rules for *ordered categorical* outcomes via the
    cumulative distribution function (CDF), using the bar bins as ordered
    categories. Definitions taken from https://scoringrules.readthedocs.io/en/latest/theory.html

        CRPS (squared): sum_k=1^K w_k (CDF(k) - y_k)^2
        CRLS (log):     -sum_k=1^K w_k log(|CDF(k) + y_k - 1|)

    where K is the number of bins, w_k is the width of bin k, CDF(k) is the predicted
    cumulative probability up to bin k, and y_k is the target cumulative
    probability up to bin k.
    The weighting uses bar/bin widths of the bar distribution.

    Note that this is the RPS/RLS loss weighted by the bar/bin widths, so we refer
    to it as 'Continuous' RPS/CRLS loss, i.e. CRPS/CRLS loss.

    Shapes suffixes:
        B=batch * estimators, L=logits, Q=n_queries.

    Args:
        logits_BQL: Bar distribution logits of shape (B, Q, L).
        targets_BQ: Targets of shape (B, Q)
        bardist_loss_fn: The BarDistribution instance used for bin mapping and
            bucket index mapping.
        loss_type: Which variant to compute. "crps" uses squared CDF differences,
            "crls" uses a log score applied to the cumulative probabilities.

    Returns:
        A scalar mean loss.
    """
    bucket_widths_L = bardist_loss_fn.bucket_widths.to(logits_BQL.device)
    assert bucket_widths_L.shape == (logits_BQL.shape[-1],), (
        f"bucket_widths_L.shape: {bucket_widths_L.shape} "
        f"logits_BQL.shape: {logits_BQL.shape}"
    )
    probs_BQL = torch.softmax(logits_BQL, dim=-1)
    pred_cdf_BQL = torch.cumsum(probs_BQL, dim=-1)

    ignore_loss_mask_BQ = torch.isnan(targets_BQ)
    # Filled with zeros if the target is NaN.
    # Will be ignored in final loss below.
    filled_targets_BQ = torch.where(
        ignore_loss_mask_BQ, torch.zeros_like(targets_BQ), targets_BQ
    )

    target_bins_BQ = bardist_loss_fn.map_to_bucket_idx(filled_targets_BQ).clamp(
        0, bardist_loss_fn.num_bars - 1
    )
    # The target CDF is a step function: 0 for bins < target_bin, 1 for bins >=
    # target_bin.
    bin_indices_L = torch.arange(probs_BQL.shape[-1], device=probs_BQL.device)
    target_cdf_BQL = (bin_indices_L.view(1, 1, -1) >= target_bins_BQ.unsqueeze(-1)).to(
        probs_BQL.dtype
    )

    if loss_type == "crps":
        cdf_diff_BQL = pred_cdf_BQL - target_cdf_BQL
        cdf_term_losses_BQL = cdf_diff_BQL.square()
    else:
        eps = torch.finfo(pred_cdf_BQL.dtype).eps
        cdf = pred_cdf_BQL.clamp(eps, 1 - eps)
        # target_cdf_BQL is binary, so we can expand the log(|CDF(k) + y_k - 1|) term
        # into two separate terms and use log1p. This is numerically more stable.
        cdf_term_losses_BQL = target_cdf_BQL * (-torch.log(cdf)) + (
            1 - target_cdf_BQL
        ) * (-torch.log1p(-cdf))

    weighted_term_losses_BQL = cdf_term_losses_BQL * bucket_widths_L.view(1, 1, -1)
    crps_losses_BQ = weighted_term_losses_BQL.sum(dim=-1)

    if ignore_loss_mask_BQ.any():
        crps_losses_BQ[ignore_loss_mask_BQ] = 0.0

    return crps_losses_BQ.mean()


class FinetunedTabPFNRegressor(FinetunedTabPFNBase, RegressorMixin):
    """A scikit-learn compatible wrapper for fine-tuning the TabPFNRegressor.

    This class encapsulates the fine-tuning loop, allowing you to fine-tune
    TabPFN on a specific dataset using the familiar .fit() and .predict() API.

    Args:
        device: The device to run the model on. Defaults to "cuda".
        epochs: The total number of passes through the fine-tuning data.
            Defaults to 30.
        time_limit: Time limit in seconds for fine-tuning.
            If None, no time limit is applied. Defaults to None.
        learning_rate: The learning rate for the AdamW optimizer. A small value
            is crucial for stable fine-tuning. Defaults to 1e-5.
        weight_decay: The weight decay for the AdamW optimizer. Defaults to 0.01.
        validation_split_ratio: Fraction of the original training data reserved
            as a validation set for early stopping and monitoring. Defaults to 0.1.
        n_finetune_ctx_plus_query_samples: The total number of samples per
            meta-dataset during fine-tuning (context plus query) before applying
            the `finetune_ctx_query_split_ratio`. Defaults to 10_000.
        finetune_ctx_query_split_ratio: The proportion of each fine-tuning
            meta-dataset to use as query samples for calculating the loss. The
            remainder is used as context. Defaults to 0.2.
        n_inference_subsample_samples: The total number of subsampled training
            samples per estimator during validation and final inference.
            Defaults to 50_000.
        random_state: Seed for reproducibility of data splitting and model
            initialization. Defaults to 0.
        early_stopping: Whether to use early stopping based on validation
            performance. Defaults to True.
        early_stopping_patience: Number of epochs to wait for improvement before
            early stopping. Defaults to 8.
        min_delta: Minimum change in metric to be considered as an improvement.
            Defaults to 1e-4.
        grad_clip_value: Maximum norm for gradient clipping. If None, gradient
            clipping is disabled. Gradient clipping helps stabilize training by
            preventing exploding gradients. Defaults to 1.0.
        use_lr_scheduler: Whether to use a learning rate scheduler (linear warmup
            with optional cosine decay) during fine-tuning. Defaults to True.
        lr_warmup_only: If True, only performs linear warmup to the base learning
            rate and then keeps it constant. If False, applies cosine decay after
            warmup. Defaults to False.
        n_estimators_finetune: If set, overrides `n_estimators` of the underlying
            estimator only during fine-tuning to control the number of
            estimators (ensemble size) used in the training loop. If None, the
            value from `kwargs` or the estimator default is used.
            Defaults to 2.
        n_estimators_validation: If set, overrides `n_estimators` only for
            validation-time evaluation during fine-tuning (early-stopping /
            monitoring). If None, the value from `kwargs` or the
            estimator default is used. Defaults to 2.
        n_estimators_final_inference: If set, overrides `n_estimators` only for
            the final fitted inference model that is used after fine-tuning. If
            None, the value from `kwargs` or the estimator default is used.
            Defaults to 8.
        use_activation_checkpointing: Whether to use activation checkpointing to
            reduce memory usage. Defaults to True.
        save_checkpoint_interval: Number of epochs between checkpoint saves. This
            only has an effect if `output_dir` is provided during the `fit()` call.
            If None, no intermediate checkpoints are saved. The best model checkpoint
            is always saved regardless of this setting. Defaults to 10.
        use_fixed_preprocessing_seed: Whether to use a fixed preprocessing seed.
            If True, the preprocessing will always use the same random seed throughout
            data batches. This is helpful in most cases because, e.g., the column order
            will stay the same across batches.
            If False, the preprocessing will use a different random seed for each batch.

        FinetunedTabPFNRegressor specific arguments:

        extra_regressor_kwargs: Additional keyword arguments to pass to the
            underlying `TabPFNRegressor`, such as `n_estimators`.
        eval_metric: The evaluation metric to use for early stopping and
            monitoring.
            For regression, the default is "mse".
            Other metrics are not currently supported.
        ce_loss_weight: Weight for the bar distribution negative log-likelihood term
            (cross-entropy-like). Defaults to 0.0.
        crps_loss_weight: Weight for the continuous ranked probability score (CRPS) term
            computed on ordered bar probabilities (squared CDF error). Defaults to 1.0.
        crls_loss_weight: Weight for the continuous ranked logarithmic score (CRLS) term
            computed on ordered bar probabilities (log score on cumulative
            probabilities). Defaults to 0.0.
        mse_loss_weight: Weight for an auxiliary MSE loss term added to the
            bar distribution loss. Set to 0.0 to disable. Defaults to 1.0.
        mse_loss_clip: Optional upper bound for the auxiliary MSE loss term.
            If None, no clipping is applied. Defaults to None.
        mae_loss_weight: Weight for an auxiliary MAE term computed on the mean
            decoded prediction. Defaults to 0.0.
        mae_loss_clip: Optional upper bound for the auxiliary MAE loss term.
            If None, no clipping is applied. Defaults to None.
    """

    def __init__(  # noqa: PLR0913
        self,
        *,
        device: str = "cuda",
        epochs: int = 30,
        time_limit: int | None = None,
        learning_rate: float = 1e-5,
        weight_decay: float = 0.01,
        validation_split_ratio: float = 0.1,
        n_finetune_ctx_plus_query_samples: int = 10_000,
        finetune_ctx_query_split_ratio: float = 0.2,
        n_inference_subsample_samples: int = 50_000,
        random_state: int = 0,
        early_stopping: bool = True,
        early_stopping_patience: int = 8,
        min_delta: float = 1e-4,
        grad_clip_value: float | None = 1.0,
        use_lr_scheduler: bool = True,
        lr_warmup_only: bool = False,
        n_estimators_finetune: int = 2,
        n_estimators_validation: int = 2,
        n_estimators_final_inference: int = 8,
        use_activation_checkpointing: bool = True,
        save_checkpoint_interval: int | None = 10,
        use_fixed_preprocessing_seed: bool = True,
        extra_regressor_kwargs: dict[str, Any] | None = None,
        ce_loss_weight: float = 0.0,
        crps_loss_weight: float = 1.0,
        crls_loss_weight: float = 0.0,
        mse_loss_weight: float = 1.0,
        mse_loss_clip: float | None = None,
        mae_loss_weight: float = 0.0,
        mae_loss_clip: float | None = None,
        eval_metric: Literal["mse"] | None = None,
    ):
        super().__init__(
            device=device,
            epochs=epochs,
            time_limit=time_limit,
            learning_rate=learning_rate,
            weight_decay=weight_decay,
            validation_split_ratio=validation_split_ratio,
            n_finetune_ctx_plus_query_samples=n_finetune_ctx_plus_query_samples,
            finetune_ctx_query_split_ratio=finetune_ctx_query_split_ratio,
            n_inference_subsample_samples=n_inference_subsample_samples,
            random_state=random_state,
            early_stopping=early_stopping,
            early_stopping_patience=early_stopping_patience,
            min_delta=min_delta,
            grad_clip_value=grad_clip_value,
            use_lr_scheduler=use_lr_scheduler,
            lr_warmup_only=lr_warmup_only,
            n_estimators_finetune=n_estimators_finetune,
            n_estimators_validation=n_estimators_validation,
            n_estimators_final_inference=n_estimators_final_inference,
            use_activation_checkpointing=use_activation_checkpointing,
            save_checkpoint_interval=save_checkpoint_interval,
            use_fixed_preprocessing_seed=use_fixed_preprocessing_seed,
        )
        self.extra_regressor_kwargs = extra_regressor_kwargs
        self.eval_metric = eval_metric
        self.ce_loss_weight = ce_loss_weight
        self.crps_loss_weight = crps_loss_weight
        self.crls_loss_weight = crls_loss_weight
        self.mse_loss_weight = mse_loss_weight
        self.mse_loss_clip = mse_loss_clip
        self.mae_loss_weight = mae_loss_weight
        self.mae_loss_clip = mae_loss_clip

    @property
    @override
    def _estimator_kwargs(self) -> dict[str, Any]:
        """Return the regressor-specific kwargs."""
        return self.extra_regressor_kwargs or {}

    @property
    @override
    def _model_type(self) -> Literal["classifier", "regressor"]:
        """Return the model type string."""
        return "regressor"

    @property
    @override
    def _metric_name(self) -> str:
        """Return the name of the primary metric."""
        return "MSE"

    @override
    def _create_estimator(self, config: dict[str, Any]) -> TabPFNRegressor:
        """Create the TabPFNRegressor with the given config."""
        return TabPFNRegressor(
            **config,
            fit_mode="batched",
            differentiable_input=False,
        )

    @override
    def _setup_estimator(self) -> None:
        """No additional setup needed for regressor at creation time."""

    @override
    def _should_skip_batch(self, batch: RegressorBatch) -> bool:  # type: ignore[override]
        """Never skip a batch for regression."""
        return False

    @override
    def _setup_batch(self, batch: RegressorBatch) -> None:  # type: ignore[override]
        """Set up bar distribution for this batch."""
        self.finetuned_estimator_.raw_space_bardist_ = batch.raw_space_bardist
        self.finetuned_estimator_.bardist_ = batch.znorm_space_bardist
        self._bardist_loss = batch.znorm_space_bardist

    @override
    def _forward_with_loss(self, batch: RegressorBatch) -> torch.Tensor:  # type: ignore[override]
        """Perform forward pass and compute bar distribution loss with optional MSE.

        Args:
            batch: The RegressorBatch containing preprocessed context and query
                data plus bar distribution information.

        Returns:
            The computed loss tensor (bar distribution + optional MSE auxiliary).
        """
        X_query_batch = batch.X_query
        y_query_batch = batch.y_query
        bardist_loss_fn = self._bardist_loss

        _, per_estim_logits, _ = self.finetuned_estimator_.forward(X_query_batch)
        # per_estim_logits is a list (per estimator) of tensors with shape [Q, B(=1), L]

        # shape suffix: Q=n_queries, B=batch(=1), E=estimators, L=logits
        logits_QBEL = torch.stack(per_estim_logits, dim=2)

        Q, B, E, L = logits_QBEL.shape
        num_bars = get_n_out(self.finetuned_estimator_.configs_[0], bardist_loss_fn)
        assert y_query_batch.shape[1] == Q
        assert B == 1
        assert self.n_estimators_finetune == E
        assert num_bars == L

        # Reshape for bar distribution loss: treat estimator dim as batch dim
        # permute to shape (B, E, Q, L) then reshape to (B*E, Q, L)
        logits_BQL = logits_QBEL.permute(1, 2, 0, 3).reshape(B * E, Q, L)

        targets_BQ = y_query_batch.repeat(B * self.n_estimators_finetune, 1).to(
            self.device
        )

        return _compute_regression_loss(
            logits_BQL=logits_BQL,
            targets_BQ=targets_BQ,
            bardist_loss_fn=bardist_loss_fn,
            ce_loss_weight=self.ce_loss_weight,
            crps_loss_weight=self.crps_loss_weight,
            crls_loss_weight=self.crls_loss_weight,
            mse_loss_weight=self.mse_loss_weight,
            mse_loss_clip=self.mse_loss_clip,
            mae_loss_weight=self.mae_loss_weight,
            mae_loss_clip=self.mae_loss_clip,
        )

    @override
    def _evaluate_model(
        self,
        eval_config: dict[str, Any],
        X_train: np.ndarray,
        y_train: np.ndarray,
        X_val: np.ndarray,
        y_val: np.ndarray,
    ) -> EvalResult:
        """Evaluate the regressor using MSE."""
        eval_regressor = clone_model_for_evaluation(
            self.finetuned_estimator_,
            eval_config,
            TabPFNRegressor,
        )
        eval_regressor.fit(X_train, y_train)

        try:
            predictions = eval_regressor.predict(X_val)  # type: ignore
            mse = mean_squared_error(y_val, predictions)
        except (ValueError, RuntimeError, AttributeError) as e:
            logger.warning(f"An error occurred during evaluation: {e}")
            mse = np.nan

        return EvalResult(primary=mse)  # pyright: ignore[reportArgumentType]

    @override
    def _is_improvement(self, current: float, best: float) -> bool:
        """Check if current MSE is better (lower) than best."""
        return current < best - self.min_delta

    @override
    def _get_initial_best_metric(self) -> float:
        """Return inf for minimization."""
        return np.inf

    @override
    def _get_checkpoint_metrics(self, eval_result: EvalResult) -> dict[str, float]:
        """Return metrics for checkpoint saving."""
        return {"mse": eval_result.primary}

    @override
    def _get_valid_finetuning_query_size(
        self,
        *,
        query_size: int,
        **kwargs: Any,
    ) -> int:
        """Calculate a valid finetuning query size."""
        return query_size

    @override
    def _log_epoch_evaluation(
        self, epoch: int, eval_result: EvalResult, mean_train_loss: float | None
    ) -> None:
        """Log evaluation results for regression."""
        mean_train_loss = "N/A" if mean_train_loss is None else f"{mean_train_loss:.4f}"
        logger.info(
            f"📊 Epoch {epoch + 1} Evaluation | Val MSE: {eval_result.primary:.4f}, "
            f"Train Loss: {mean_train_loss}"
        )

    @override
    def _setup_inference_model(
        self, final_inference_eval_config: dict[str, Any]
    ) -> None:
        """Set up the final inference regressor."""
        finetuned_inference_regressor = clone_model_for_evaluation(
            self.finetuned_estimator_,
            final_inference_eval_config,
            TabPFNRegressor,
        )
        self.finetuned_inference_regressor_ = finetuned_inference_regressor
        self.finetuned_inference_regressor_.fit_mode = "fit_preprocessors"  # type: ignore
        self.finetuned_inference_regressor_.fit(self.X_, self.y_)  # type: ignore

    @override
    def fit(
        self,
        X: XType,
        y: YType,
        X_val: XType | None = None,
        y_val: YType | None = None,
        output_dir: Path | None = None,
    ) -> FinetunedTabPFNRegressor:
        """Fine-tune the TabPFN model on the provided training data.

        Args:
            X: The training input samples of shape (n_samples, n_features).
            y: The target values of shape (n_samples,).
            X_val: Optional validation input samples.
            y_val: Optional validation target values.
            output_dir: Directory path for saving checkpoints. If None, no
                checkpointing is performed and progress will be lost if
                training is interrupted.

        Returns:
            The fitted instance itself.
        """
        if self.eval_metric is None:
            self.eval_metric = "mse"

        super().fit(X, y, X_val=X_val, y_val=y_val, output_dir=output_dir)
        return self

    @override
    def predict(self, X: XType, **kwargs) -> RegressionResultType:
        """Predict target values for X.

        Args:
            X: The input samples of shape (n_samples, n_features).
            **kwargs: Additional keyword arguments to pass to the underlying
                inference regressor (e.g., output_type, quantiles).

        Returns:
            The predicted target values with shape (n_samples,).
        """
        check_is_fitted(self)

        return self.finetuned_inference_regressor_.predict(X, **kwargs)
