from dataclasses import asdict, dataclass, field
from enum import Enum
from typing import Optional

import numpy as np
import torch

TensorLike = torch.Tensor | np.ndarray

###############################################################################
# Enumerations for configuration options
###############################################################################


# missing data strategies
class MissingDataStrategy(Enum):
    IGNORE_NANS = "ignore_nans"
    IMPUTE_OUTCOME_MEAN = "impute_outcome_mean"
    IMPUTE_COLUMN_MEAN = "impute_column_mean"


# methods for marginal feature screening
class MFSMethod(Enum):
    AUTO = "auto"
    THRESHOLD = "threshold"
    PERCENTILE = "percentile"
    NONE = "none"


# Whether to aggregate or separate features
class FeaturizeAs(Enum):
    AGGREGATE = "aggregate"
    SEPARATE = "separate"


# Fit Methods
class FitMethod(Enum):
    UNIFORM = "uniform"
    WEIGHTED = "weighted"


class Target(Enum):
    BINARY = "binary"
    CONTINUOUS = "continuous"


###############################################################################
# Configuration dataclasses
###############################################################################


@dataclass
class MFSConfig:
    mfs_method: MFSMethod = MFSMethod.PERCENTILE
    mfs_cutoff: float = 15
    start_sweep: Optional[float] = 1
    end_sweep: Optional[float] = 100
    step_size: Optional[float] = 1

    def __post_init__(self):
        if self.mfs_cutoff < 0:
            raise ValueError("mfs_cutoff must be positive.")
        if self.mfs_method == MFSMethod.PERCENTILE and (self.mfs_cutoff > 100 or self.mfs_cutoff < 0):
            raise ValueError("mfs_cutoff for PERCENTILE method must be <= 100.")
        if self.start_sweep <= 0:
            raise ValueError("start_sweep must be greater than 0.")
        if self.end_sweep <= self.start_sweep or self.end_sweep > 100:
            raise ValueError("end_sweep must be greater than start_sweep and less than or equal to 100.")
        if self.step_size <= 0:
            raise ValueError("step_size must be at greater than 0.")


@dataclass
class ReferencePoolConfig:
    jackknife_row: bool = True  # (m, 1, n)
    jackknife_col: bool = True  # (1, n, n)
    drop_row_nans: bool = True
    missing_data_strategy: MissingDataStrategy = MissingDataStrategy.IMPUTE_OUTCOME_MEAN
    mfs_config: Optional[MFSConfig] = field(default_factory=MFSConfig)

    def __post_init__(self):
        if self.drop_row_nans and not self.jackknife_row:
            raise ValueError("Cannot drop row NaNs without jackknife_row enabled.")
        if self.jackknife_row and self.mfs_config is None:
            raise ValueError(
                "jackknife_row cannot be enabled without a marginal feature screening method; "
                "otherwise all jackknifed rows will have the same reference pool.",
            )


@dataclass
class FeatureConfig:
    missing_data_strategy: MissingDataStrategy = MissingDataStrategy.IMPUTE_OUTCOME_MEAN
    missingness_features: Optional[FeaturizeAs] = None
    pred_features: FeaturizeAs = FeaturizeAs.AGGREGATE
    logit_feats: bool = False  # Only used for binary outcomes
    score_as_probs: bool = False  # Only relevant for binary outcomes and when logit_feats is True (ignored otherwise)

    def __post_init__(self):
        if self.missing_data_strategy == MissingDataStrategy.IGNORE_NANS:
            raise ValueError("missing_data_strategy for FeatureConfig cannot be IGNORE_NANS!")


@dataclass
class ModelConfig:
    fit_method: FitMethod = FitMethod.UNIFORM
    epochs: int = 500
    lr: float = 0.1
    l2: float = 0.1
    log_interval: int = 50

    def kwargs_as_dict(self) -> dict:
        """Return model config as dict (without fit_method)."""
        d = asdict(self)
        d.pop("fit_method", None)
        return d


@dataclass
class WomacConfig:
    target: Target
    reference_pool_config: ReferencePoolConfig = field(default_factory=ReferencePoolConfig)
    feature_config: FeatureConfig = field(default_factory=FeatureConfig)
    model_config: ModelConfig = field(default_factory=ModelConfig)

    def __post_init__(self):
        if self.feature_config.score_as_probs and self.target == Target.CONTINUOUS:
            raise ValueError("score_as_probs requires target to be binary.")
        if self.feature_config.logit_feats and self.target == Target.CONTINUOUS:
            raise ValueError("logit_feats cannot be enabled for continuous outcomes.")

    @classmethod
    def binary_outcome(cls) -> "WomacConfig":
        """Configure for binary outcomes with probabilistic predictions.

        Returns:
            WomacConfig: Configuration for logistic regression on binary outcomes with logit transformation enabled.
        """
        return cls(
            target=Target.BINARY,
            feature_config=FeatureConfig(
                logit_feats=True,
                score_as_probs=True,
            ),
        )

    @classmethod
    def continuous_outcome(cls) -> "WomacConfig":
        """Configure for continuous outcomes with linear predictions.

        Returns:
            WomacConfig: Configuration for linear regression on continuous outcomes.
        """
        return cls(
            target=Target.CONTINUOUS,
        )


###############################################################################
# Data container for WOMAC algorithm
###############################################################################


@dataclass(frozen=True, slots=True)
class WomacData:
    """Container for WOMAC algorithm input data.

    Attributes:
        train_preds (Optional[torch.Tensor]): Prediction tensor for training of shape (dim1, dim2, k, l),
            where dim1 = m if jackknife_row else 1, dim2 = n if jackknife_col else 1,
            and k ∈ {m-1, m} (depending on jackknife_row) and l ∈ {1, n} (depending on pred_features).
        train_missingness (Optional[torch.Tensor]): Missingness indicator tensor for training of same shape as train_preds.
        train_labels (torch.Tensor): Labels tensor for training of shape (dim1, 1, k).
        test_preds (torch.Tensor): Prediction tensor for testing of shape (m, dim2, 1, l).
        test_missingness (Optional[torch.Tensor]): Missingness indicator tensor for testing of same shape as test_preds.
        test_labels (Optional[torch.Tensor]): Labels tensor for testing of shape (m, dim2, 1) if provided.
    """

    train_preds: Optional[torch.Tensor]
    train_missingness: Optional[torch.Tensor]
    train_labels: torch.Tensor
    test_preds: torch.Tensor
    test_missingness: Optional[torch.Tensor]
    test_labels: Optional[torch.Tensor] = None

    @property
    def batch_dim1(self) -> int:
        """Size of the first batch dimension (m). This is m if jackknife_row
        is True, otherwise 1.
        """
        if self.train_preds is not None:
            return self.train_preds.shape[0]
        raise ValueError("Invalid training predictions provided.")

    @property
    def batch_dim2(self) -> int:
        """Size of the second batch dimension (n). This is n if jackknife_col
        is True, otherwise 1.
        """
        if self.train_preds is not None:
            return self.train_preds.shape[1]
        raise ValueError("Invalid training predictions provided.")

    @property
    def data_dim(self) -> int:
        """Size of the data dimension (k1).

        This is the number of tasks to predict over at a time.
        During training, this is m-1 if jackknife_row is True, otherwise m.
        During testing, this is 1.
        """
        if self.train_preds is not None:
            return self.train_preds.shape[2]
        raise ValueError("Invalid training predictions provided.")

    @property
    def feature_dim(self) -> int:
        """Size of the feature dimension (pred plus missingness features).

        This is:
        1 if pred_features is AGGREGATE and missingness_features is None,
        2 if pred_features is AGGREGATE and missingness_features is AGGREGATE,
        n if pred_features is SEPARATE and missingness_features is None,
        n + 1 if pred_features is SEPARATE and missingness_features is AGGREGATE OR
              if pred_features is AGGREGATE and missingness_features is SEPARATE.
        2 * n if pred_features is SEPARATE and missingness_features is SEPARATE.
        """
        train_data = self.train_data
        if train_data is not None:
            return train_data.shape[3]
        raise ValueError("Invalid training predictions provided.")

    @property
    def train_data(self) -> torch.Tensor:
        """Construct combined training data tensor by concatenating predictions and missingness
        features.

        If train_missingness is provided, returns a tensor of shape (dim1, dim2, k, l1 + l2) by
        concatenating along dim=3; otherwise returns train_preds unchanged.
        """
        if self.train_missingness is None:
            return self.train_preds
        else:
            return torch.cat([self.train_preds, self.train_missingness], dim=3)

    @property
    def test_data(self) -> torch.Tensor:
        """Construct combined test data tensor by concatenating predictions and missingness
        features.

        If test_missingness is provided, returns a tensor of shape (m, n, 1, l1 + l2) by
        concatenating along dim=3; otherwise returns test_preds unchanged.
        """
        if self.test_missingness is None:
            return self.test_preds
        else:
            return torch.cat([self.test_preds, self.test_missingness], dim=3)

    def __post_init__(self):
        if self.train_preds is None and self.train_missingness is not None:
            raise ValueError("train_missingness cannot be provided without train_preds.")
        if self.test_preds is None and self.test_missingness is not None:
            raise ValueError("test_missingness cannot be provided without test_preds.")


###############################################################################
# Results dataclasses
###############################################################################


@dataclass
class ModelResult:
    coef_: torch.Tensor
    train_losses: list[TensorLike]
    bias_: Optional[torch.Tensor] = None
    val_losses: Optional[list[TensorLike]] = None  # This is ONLY val loss; ignores regularization loss.


@dataclass
class WomacResult:
    """Container for WOMAC algorithm output.

    Attributes:
        womac_ranked_indices (torch.Tensor): Expert indices sorted by WOMAC performance, shape (n experts,).
        womac_ranked_scores (torch.Tensor): Corresponding scores for womac_ranked_indices, shape (n experts,).
        womac_scores (torch.Tensor): Raw WOMAC scores per expert, shape (n experts,).
        outcome_ranked_indices (torch.Tensor): Expert indices sorted by outcome scores, shape (n experts,).
        outcome_ranked_scores (torch.Tensor): Corresponding scores for outcome_ranked_indices, shape (n experts,).
        outcome_scores (torch.Tensor): Raw outcome scores per expert, shape (n experts,).
        reference_matrix (torch.Tensor): Reference solutions per expert, shape (m tasks, dim2).
        dim2 is n if jackknife_col is True, otherwise 1.
    """

    womac_ranked_indices: torch.Tensor
    womac_ranked_scores: torch.Tensor
    womac_scores: torch.Tensor
    outcome_ranked_indices: torch.Tensor
    outcome_ranked_scores: torch.Tensor
    outcome_scores: torch.Tensor
    reference_matrix: torch.Tensor
    fitted_model: ModelResult

    @property
    def winner(self, by_womac: bool = True) -> int:
        """Get the index of the best-performing expert.

        Args:
            by_womac (bool): Whether to rank by WOMAC scores or outcome scores.

        Returns:
            int: Expert index with highest performance.
        """
        ranked_indices = self.womac_ranked_indices if by_womac else self.outcome_ranked_indices
        return int(ranked_indices[0].item())

    @property
    def reference_solution(self) -> torch.Tensor:
        """Compute the averaged reference solution across experts."""
        return torch.mean(self.reference_matrix, dim=1)  # (m,)

    def reference_for_expert(self, index: int) -> torch.Tensor:
        """Retrieve the reference solution for a specific expert based
        on their original index.

        Args:
            index (int): Expert index (0 <= index < num_experts).

        Returns:
            torch.Tensor: Reference solution vector for the specified expert, shape (m,).

        Raises:
            IndexError: If index is out of bounds.
        """
        if index < 0 or index >= self.womac_scores.shape[0]:
            raise IndexError("Index out of bounds.")
        if self.reference_matrix.ndim == 1:
            return self.reference_matrix
        return self.reference_matrix[:, index]

    def top_k(self, k: int, by_womac: bool = True) -> torch.Tensor:
        """Return the indices of the top k experts.

        Args:
            k (int): Number of top experts to retrieve.
            by_womac (bool): Whether to rank by WOMAC scores or outcome scores.

        Returns:
            torch.Tensor: Tensor of top k expert indices, shape (k,).
        """
        ranked_indices = self.womac_ranked_indices if by_womac else self.outcome_ranked_indices
        return ranked_indices[:k]

    def score_of(self, index: int, by_womac: bool = True) -> float:
        """Get the score for a specific expert based on their
        original index.

        Args:
            index (int): Rank index in ranked_scores (0-based).
            by_womac (bool): Whether to use WOMAC scores or outcome scores.

        Returns:
            float: Score for the expert at the specified position.

        Raises:
            IndexError: If index is out of bounds.
        """
        scores = self.womac_scores if by_womac else self.outcome_scores
        if index < 0 or index >= scores.shape[0]:
            raise IndexError(f"Index out of bounds; should be between 0 and {scores.shape[0] - 1}.")
        return float(scores[index].item())

    def rank_of(self, index: int, by_womac: bool = True) -> int:
        """Get the performance rank for a specific expert based
        on their original index.

        Args:
            index (int): Expert index to query.
            by_womac (bool): Whether to rank by WOMAC scores or outcome scores.

        Returns:
            int: 1-based rank of the expert.

        Raises:
            IndexError: If index is out of bounds.
        """
        ranked_indices = self.womac_ranked_indices if by_womac else self.outcome_ranked_indices
        if index < 0 or index >= len(ranked_indices):
            raise IndexError(f"Index out of bounds; should be between 0 and {len(ranked_indices) - 1}.")
        return int((ranked_indices == index).nonzero(as_tuple=True)[0].item() + 1)

    def leaderboard(self, limit: Optional[int] = None, by_womac: bool = True) -> str:
        """Generate a human-readable leaderboard of experts showing both WOMAC and outcome scores,
        sorted by the selected metric.

        Args:
            limit (Optional[int]): Maximum number of top experts to include. If None, includes all.
            by_womac (bool): Whether to sort by WOMAC scores or outcome scores.

        Returns:
            str: Multiline string table with columns Rank, Expert Index, WOMAC MSE, and Outcome MSE.
        """
        # Choose the ranking order
        order = self.womac_ranked_indices if by_womac else self.outcome_ranked_indices

        # Determine how many entries to show
        if limit is None:
            limit = len(order)

        # Header with both columns
        lines = ["Rank | Expert Index | WOMAC MSE  | Outcome MSE"]
        for rank, idx in enumerate(order[:limit], start=1):
            idx_int = int(idx.item())
            womac_score = self.womac_scores[idx_int].item()
            outcome_score = self.outcome_scores[idx_int].item()
            lines.append(f"{rank:4} | {idx_int:12} | {womac_score:10.7f} | {outcome_score:11.7f}")

        return "\n".join(lines)
