import logging
from typing import Optional

import torch

from womac.config import (
    FitMethod,
    MFSMethod,
    ModelConfig,
    ReferencePoolConfig,
    Target,
    TensorLike,
    WomacConfig,
    WomacData,
    WomacResult,
)
from womac.model import AverageModel, LinearModel
from womac.preprocess import (
    _featurize_and_center,
    _missingness_features,
)
from womac.reference_pool import _build_base_reference_pool, _marginal_feature_screening
from womac.utils import (
    _apply_logit_transform,
    _cast_to_tensor,
    _get_indices_without_min_responses,
    impute_data,
    impute_with,
    jackknife_expansion,
    score_mse,
)

logger = logging.getLogger(__name__)

if torch.backends.mps.is_available():
    DEVICE = "mps"
elif torch.cuda.is_available():
    DEVICE = "cuda"
else:
    DEVICE = "cpu"


class Womac:
    def __init__(
        self,
        min_responses_per_task: int = 2,
        min_responses_per_expert: int = 2,
        device: str = DEVICE,
        seed: Optional[int] = None,
    ):
        """
        Initializes the WOMAC class.

        Args:
            min_responses_per_task (int, optional): The minimum number of responses required per task. Defaults to 2.
            min_responses_per_expert (int, optional): The minimum number of responses required per expert. Defaults to 2.
            device (str, optional): The device to be used (e.g., 'cpu' or 'cuda' or 'mps'). Defaults to DEVICE.
            seed (int | None, optional): Random seed for reproducibility. If provided, sets torch seed.
        """

        self.min_responses_per_task = min_responses_per_task
        self.min_responses_per_expert = min_responses_per_expert
        self.device = torch.device(device)
        if seed is not None:
            torch.manual_seed(seed)

        logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")

    def select_reference_pool(
        self,
        X: torch.Tensor,
        Y: torch.Tensor,
        reference_pool_config: ReferencePoolConfig,
    ) -> torch.Tensor:
        """
        Selects a pool of reference experts for each task-expert pair.

        Builds a base reference mask by applying basic configured filters in
        reference_pool_config, then uses marginal feature selection (MFS) and
        refines the mask based on reference_pool_config.mfs_config.

        Mask output dim is (dim1, dim2, n), where:
            - dim1: 1 if reference_pool_config.jackknife_row is False, m otherwise.
            - dim2: 1 if reference_pool_config.jackknife_col is False, n otherwise.
            - n: number of experts.

        Args:
            X (torch.Tensor): Prediction matrix of shape (m, n), where
                m is the number of tasks and n is the number of experts.
                May contain NaNs.
            Y (torch.Tensor): Outcome tensor of shape (m,) or (m, 1). Must have
                the same number of rows (m) as X and contain no NaNs.
            reference_pool_config (ReferencePoolConfig): Configuration object
                specifying the reference pool selection parameters.

        Returns:
            torch.BoolTensor: A boolean mask of shape (dim1, dim2, n), where
                mask[i, j, k] == True indicates that expert k is selected as
                a reference for [task i, expert j] (may be broadcast).

        Raises:
            ValueError: If X or Y fail validation in `_validate_input_data`, or if
                after applying MFS no references remain for any (i, j) pair.
        """
        # Validate inputs
        X_copy, Y_copy = _cast_to_tensor(X, device=self.device), _cast_to_tensor(Y, device=self.device)
        self._validate_input_data(X_copy, Y_copy)

        # build base reference pool
        rows_without_min_resp, cols_without_min_resp = _get_indices_without_min_responses(
            X_copy,
            self.min_responses_per_task,
            self.min_responses_per_expert,
        )

        reference_pool = _build_base_reference_pool(
            X_copy,
            reference_pool_config,
            cols_without_min_resp,
            device=self.device,
        )  # (dim1, dim2, n)

        # if no mfs_config is specified, return default mask
        if reference_pool_config.mfs_config is None:
            return reference_pool

        # score MSE for each (dim1, dim2)
        Y_copy[rows_without_min_resp] = torch.nan  # stops selected Y_i from being used in MSE calculation
        errors = score_mse(
            X_copy,
            Y_copy,
            missing_data_strategy=reference_pool_config.missing_data_strategy,
            jackknife_row=reference_pool_config.jackknife_row,
            apply_logit_transform=False,  # scoring against outcome Y, so no logit transform is applied
        )  # (dim1, n)

        # apply marginal feature screening
        reference_pool = _marginal_feature_screening(reference_pool, errors, reference_pool_config.mfs_config)
        return reference_pool  # (dim1, dim2, n)

    def get_reference_matrix(
        self,
        X: torch.Tensor,
        Y: torch.Tensor,
        womac_config: WomacConfig,
    ) -> torch.Tensor:
        """Compute the WOMAC reference for each [task, expert] pair.

        This method performs the following steps:
        1. Selects a reference-pool mask for each task and expert based on
           minimal-response filters and optional marginal feature selection.
        2. Preprocesses the data (logit transform, missing-data imputation,
           centering, and feature construction) as configured.
        3. Fits a predictive model (uniform average, linear regression, or
           logistic regression) for each hold-out [task, expert] split using
           the reference pool.
        4. Generates predictions for the held-out entries, producing the
           reference matrix. NOTE: If target is binary and feature_config.logit_feats
           is True: then if `womac_config.feature_config.score_as_probs` is False
           the matrix is in logit space; if True, matrix is in probabilities.

        Args:
            X (torch.Tensor): Input prediction matrix of shape (m, n),
                where m is the number of tasks and n is the number of experts.
            Y (torch.Tensor): Outcome tensor of shape (m,) or (m, 1).
            womac_config (WomacConfig, optional): Configuration object
                specifying reference-pool selection, preprocessing, and
                modeling parameters. Defaults to a new WomacConfig.

        Returns:
            torch.Tensor: Reference matrix of shape (m, dim2), where
                dim2 equals the second dimension of the reference-pool mask:
                - If `womac_config.reference_pool_config.jackknife_col` is True,
                  dim2 == n (one prediction per expert).
                - If False, dim2 == 1 (a single aggregated reference per task),
                  which can be broadcast against an (m, n) matrix.

        Raises:
            ValueError: If X or Y fail validation, or if any (task, expert)
                pair has no valid references.
        """
        X, Y = _cast_to_tensor(X, device=self.device), _cast_to_tensor(Y, device=self.device)
        self._validate_input_data(X, Y)

        mask = self.select_reference_pool(
            X,
            reference_pool_config=womac_config.reference_pool_config,
            Y=Y,
        )  # (dim1, dim2, n)
        womac_data = self._preprocess_data(
            X,
            Y,
            mask=mask,
            womac_config=womac_config,
        )  # (dim1, dim2, k, l)

        self._fit(
            womac_data,
            mask=mask,
            target=womac_config.target,
            logit_feats=womac_config.feature_config.logit_feats,
            model_config=womac_config.model_config,
        )
        return self._predict(
            womac_data,
            sigmoid_transform=(womac_config.feature_config.logit_feats and womac_config.feature_config.score_as_probs),
        ).squeeze(-1)  # (m, dim2, 1, l) -> (m, dim2)

    def tune_mfs(self, X: TensorLike, Y: TensorLike, womac_config: WomacConfig) -> float:
        """
        Tunes Marginal Feature Screening (MFS) percentile cutoff to minimize the Mean Squared Error (MSE)
        between the reference matrix and the target tensor `Y`. This method iteratively sweeps through
        a range of cutoff values, computes the outcome MSE for each cutoff, and selects the best cutoff
        that results in the lowest MSE.

        Args:
            X (TensorLike): Input tensor representing the features or data points. It is cast to a tensor
            and validated before processing.
            Y (TensorLike): Target tensor representing the expected outcomes or labels. It is cast to a
            tensor and validated before processing.
            womac_config (WomacConfig): Configuration object containing parameters for the Womac model,
            including the reference pool configuration and MFS tuner settings.

        Returns:
            float: The optimal cutoff value that minimizes the outcome MSE.

        Raises:
            ValueError: If the input tensors `X` and `Y` fail validation checks.
        """
        X, Y = _cast_to_tensor(X, device=self.device), _cast_to_tensor(Y, device=self.device)
        self._validate_input_data(X, Y)

        percentile_scores = []
        start = womac_config.reference_pool_config.mfs_config.start_sweep
        end = womac_config.reference_pool_config.mfs_config.end_sweep
        steps = womac_config.reference_pool_config.mfs_config.step_size

        for c in torch.arange(start, end, steps):
            womac_config.reference_pool_config.mfs_config.mfs_cutoff = c
            ref_mat = self.get_reference_matrix(X, Y, womac_config=womac_config)  # (m, dim2)
            outcome_score = torch.mean((ref_mat - Y) ** 2).cpu().float()
            percentile_scores.append((c, outcome_score))
            logging.info(f"Cutoff {c:.2f} - Outcome MSE: {outcome_score:.4f}")

        # pick the best cutoff based on outcome MSE
        best_cutoff, _ = min(percentile_scores, key=lambda x: x[1])
        logging.info(
            f"Best MFS cutoff: {best_cutoff:.2f} with outcome MSE {min(percentile_scores, key=lambda x: x[1])[1]:.4f}"
        )
        return best_cutoff

    def score_competition(
        self,
        X: TensorLike,
        Y: TensorLike,
        womac_config: WomacConfig,
    ) -> WomacResult:
        """
        Run the full WOMAC algorithm: compute processed predictions and reference values,
        and rank experts by the MSE of their predictions against the reference values and outcomes.

        This method performs the following steps:
          1. Optionally tunes the MFS cutoff if AUTO is selected, updating womac_config in-place.
          2. Imputes missing data from predictions (and applies logit transform if configured), producing
             processed predictions for each task-expert pair (m, n).
          3. Builds a reference matrix for each task/expert via get_reference_matrix (m, dim2),
             where dim2 is 1 or n depending on womac_config.reference_pool_config.jackknife_col.
          4. Computes mean squared errors (MSE) between processed X and reference values (WOMAC scores).
          5. Computes MSE between processed X and Y (outcome scores).
          6. Ranks experts by their WOMAC and outcome MSE scores.

        Args:
            X (TensorLike): Prediction matrix of shape (m, n), where
                m is the number of tasks and n is the number of experts.
            Y (TensorLike): Outcome vector of shape (m,) or matrix of shape (m, 1).
            womac_config (WomacConfig): Configuration object controlling
                reference pool selection, preprocessing options (e.g., logit transform,
                missing-data strategy), and model fitting.

        Returns:
            WomacResult: A result object with fields:
                womac_ranked_indices (torch.Tensor): Expert indices ordered by ascending WOMAC MSE, shape (n,).
                womac_ranked_scores (torch.Tensor): Sorted WOMAC MSE values, shape (n,).
                womac_scores (torch.Tensor): Raw WOMAC MSE for each expert before sorting, shape (n,).
                outcome_ranked_indices (torch.Tensor): Expert indices ordered by ascending outcome MSE, shape (n,).
                outcome_ranked_scores (torch.Tensor): Sorted outcome MSE values, shape (n,).
                outcome_scores (torch.Tensor): Raw outcome MSE for each expert before sorting, shape (n,).
                reference_matrix (torch.Tensor): Reference predictions used for scoring,
                    shape (m, dim2), where dim2 is 1 or n depending on
                    womac_config.reference_pool_config.jackknife_col.
                fitted_model: The fitted model result object.

        Raises:
            ValueError: If X or Y fail validation (wrong shape, insufficient non-NaNs, etc.).
        """
        X, Y = _cast_to_tensor(X, device=self.device), _cast_to_tensor(Y, device=self.device)
        self._validate_input_data(X, Y)

        # 1. Run AUTO MFS if configured. NOTE: this SETS the mfs_cutoff in womac_config.reference_pool_config.mfs_config
        if womac_config.reference_pool_config.mfs_config.mfs_method == MFSMethod.AUTO:
            womac_config.reference_pool_config.mfs_config.mfs_cutoff = self.tune_mfs(X, Y, womac_config=womac_config)

        # 2. Get processed predictions and reference matrix
        reference_matrix = self.get_reference_matrix(X, Y, womac_config=womac_config)  # (m, dim2)

        # 3. Score
        womac_scores = score_mse(
            X,
            reference_matrix,
            womac_config.feature_config.missing_data_strategy,
            jackknife_row=False,
            apply_logit_transform=(
                womac_config.feature_config.logit_feats and not womac_config.feature_config.score_as_probs
            ),
        ).squeeze(0)  # (n,)
        outcome_scores = score_mse(
            X,
            Y,
            womac_config.feature_config.missing_data_strategy,
            jackknife_row=False,
            apply_logit_transform=False,
        ).squeeze(0)  # (n,)

        # 4. Rank
        womac_ranked_indices = torch.argsort(womac_scores)
        womac_ranked_scores = womac_scores[womac_ranked_indices]
        outcome_ranked_indices = torch.argsort(outcome_scores)
        outcome_ranked_scores = outcome_scores[outcome_ranked_indices]

        return WomacResult(
            womac_ranked_indices=womac_ranked_indices.cpu(),
            womac_ranked_scores=womac_ranked_scores.cpu(),
            womac_scores=womac_scores.cpu(),
            outcome_ranked_indices=outcome_ranked_indices.cpu(),
            outcome_ranked_scores=outcome_ranked_scores.cpu(),
            outcome_scores=outcome_scores.cpu(),
            reference_matrix=reference_matrix.cpu(),
            fitted_model=self.model.model_result,
        )

    def _validate_input_data(self, X: torch.Tensor, Y: Optional[torch.Tensor]) -> None:
        """Validates the input data tensors X and optional Y.

        This method ensures that:
          * X is a 2D, non-empty tensor with at least two rows and two columns.
          * Every row in X has at least two non-NaN values, and at least two rows
            have ≥ self.min_responses_per_task non-NaN entries.
          * Every column in X has at least two non-NaN values, and at least two
            columns have ≥ self.min_responses_per_expert non-NaN entries.
          * If provided, Y is a 1D or 2D, non-empty tensor with at least two rows,
            contains no NaNs, and has the same number of rows as X.

        Args:
            X (torch.Tensor): A tensor of shape (tasks, experts) containing
                model predictions, possibly with NaNs.
            Y (Optional[torch.Tensor]): A tensor of shape (tasks,) or
                (tasks, 1) containing target values, or None.

        Raises:
            ValueError: If X is not 2D, is empty, has fewer than two rows or
                columns, has rows or columns with fewer than two non-NaN entries,
                or if fewer than two rows/columns meet the respective
                min_responses_per_task/expert thresholds.
            ValueError: If Y is provided but is empty, has invalid dimensions,
                contains NaNs, has fewer than two rows, or its row count
                does not match X.
        """

        # Checks on X
        if X.ndim != 2:
            raise ValueError(f"X must be a 2D array, but got {X.ndim}D.")
        if X.numel() == 0:
            raise ValueError("X must not be empty.")
        if X.shape[0] < 2:
            raise ValueError("Prediction matrix X must have at least two rows.")
        if X.shape[1] < 2:
            raise ValueError("Prediction matrix X must have at least two columns.")

        row_not_nan_counts = torch.sum(~torch.isnan(X), dim=1)
        if torch.any(row_not_nan_counts < 2):
            problematic_rows = torch.where(row_not_nan_counts < 2)[0]
            raise ValueError(f"Every row must have at least 2 non-NaN values. Problematic rows: {problematic_rows}")
        num_rows_with_min_responses = torch.sum(row_not_nan_counts >= self.min_responses_per_task)
        if num_rows_with_min_responses < 2:
            raise ValueError(
                f"Not enough rows with at least {self.min_responses_per_task} non-NaN values. "
                f"Only {num_rows_with_min_responses} rows have enough data."
            )

        col_not_nan_counts = torch.sum(~torch.isnan(X), dim=0)
        if torch.any(col_not_nan_counts < 2):
            problematic_cols = torch.where(col_not_nan_counts < 2)[0]
            raise ValueError(
                f"Every column must have at least 2 non-NaN values. Problematic columns: {problematic_cols}"
            )
        num_cols_with_min_responses = torch.sum(col_not_nan_counts >= self.min_responses_per_expert)
        if num_cols_with_min_responses < 2:
            raise ValueError(
                f"Not enough columns with at least {self.min_responses_per_expert} non-NaN values. "
                f"Only {num_cols_with_min_responses} columns have enough data."
            )

        # Checks on Y
        if Y is not None:
            if Y.ndim > 2:
                raise ValueError(f"Y must be a 1D or 2D array, but got {Y.ndim}D.")
            if Y.numel() == 0:
                raise ValueError("Y must not be empty.")
            if Y.shape[0] < 2:
                raise ValueError("Target vector Y must have at least two rows.")
            if torch.any(torch.isnan(Y)):
                raise ValueError("Y must not contain NaN values.")

            # Checks on both X and Y
            if X.shape[0] != Y.shape[0]:
                raise ValueError(
                    f"Y must have the same number of rows as X. Y has {Y.shape[0]} rows, but X has {X.shape[0]} rows."
                )

    def _preprocess_data(
        self,
        X: torch.Tensor,
        Y: torch.Tensor,
        mask: torch.Tensor,
        womac_config: WomacConfig,
    ) -> WomacData:
        """
        Preprocess input features and labels for the WOMAC model.

        This method applies, if requested, imputation of missing data according to the
        configured strategy (with logit transform if requested), jackknife expansion of
        the reference pool, centering of selected features, and extraction of missingness
        indicators.

        Args:
            X (torch.Tensor):
                Raw input feature tensor of shape (m, n), where m is the number
                of samples and n is the number of features.
            Y (torch.Tensor):
                Targets tensor of shape (m,) or (m, 1).
            mask (torch.Tensor):
                Boolean mask tensor of shape (dim1, dim2, n) indicating which
                experts can be selected as references for each [task, expert] pair.
            womac_config (WomacConfig): Configuration object containing settings for
                preprocessing,

        Returns:
            WomacData:
                A data container with the following attributes:
                - train_preds (torch.Tensor or None):
                    Preprocessed training predictors.
                    • If fit_method == UNIFORM: None
                    • If non-uniform & jackknife_row == False: shape (1, dim2, m, l)
                    • If non-uniform & jackknife_row == True: shape (m, dim2, m-1, l)
                    l is 1 or n depending on whether feature_config.pred_features is
                    FeaturizeAs.AGGREGATE or FeaturizeAs.SEPARATE. dim2 depends on
                    womac_config.reference_pool_config.jackknife_col.
                - train_missingness (torch.Tensor or None):
                    Missingness indicators for training data, same shape as train_preds.
                    l is 1 or n depending on whether feature_config.missingness_features is
                    FeaturizeAs.AGGREGATE or FeaturizeAs.SEPARATE.
                - train_labels (torch.Tensor or None):
                    Training labels aligned with train_preds.
                    • If fit_method == UNIFORM: None
                    • If non-uniform & jackknife_row == False: shape (1, 1, m)
                    • If non-uniform & jackknife_row == True: shape (m, 1, m-1)
                - test_preds (torch.Tensor):
                    Preprocessed test predictors.
                    • shape (m, 1, n)
                - test_missingness (torch.Tensor or None):
                    Missingness indicators for test data, same shape as test_preds.
                - test_labels (torch.Tensor):
                    Test labels aligned with test_preds.
                    • If reference_pool_config.jackknife_col == True: shape (m, n)
                    • If reference_pool_config.jackknife_col == False: shape (m, 1, 1)
        """

        # Initialize
        X_train, X_train_nans, X_test, X_test_nans, Y_train = None, None, None, None, None

        if womac_config.model_config.fit_method == FitMethod.UNIFORM:
            # 1A. Impute missing data
            X_test = impute_data(
                X,
                missing_data_strategy=womac_config.feature_config.missing_data_strategy,
                Y=Y,
            )[0].unsqueeze(1)  # (m, 1, n)

            # 2A. Apply logit transform if configured
            X_test = _apply_logit_transform(X_test, womac_config.feature_config.logit_feats)  # (m, 1, n)
        else:
            # 1B. Impute missing data
            if womac_config.reference_pool_config.jackknife_row:
                X_expanded = jackknife_expansion(X)  # (m, m-1, n)
                X_train, col_mean = impute_data(
                    X_expanded,
                    missing_data_strategy=womac_config.feature_config.missing_data_strategy,
                    Y=Y,
                )  # (m, m-1, n), (m, 1, n)
                X_test = impute_with(X=X.unsqueeze(1), impute_using=col_mean)  # (m, 1, n)
                Y_train = jackknife_expansion(Y).transpose(1, 2)  # (m, 1, m-1)
            else:
                X_train, col_mean = impute_data(
                    X, missing_data_strategy=womac_config.feature_config.missing_data_strategy, Y=Y
                )  # (m, n), (1, n)
                X_test = impute_with(X, col_mean)  # (m, n)
                X_train = X_train.unsqueeze(0)  # (1, m, n)
                X_test = X_test.unsqueeze(-2)  # (m, 1, n)
                Y_train = Y.view(1, -1).unsqueeze(0)  # (1, 1, m)

            # 2B. Apply logit transform if configured
            X_train = _apply_logit_transform(X_train, womac_config.feature_config.logit_feats)
            X_test = _apply_logit_transform(X_test, womac_config.feature_config.logit_feats)

            # 3. Standardize data by centering (dim1, dim2, k, l), (m, dim2, 1, l)
            X_train, X_test = _featurize_and_center(X_train, X_test, mask, womac_config.feature_config.pred_features)

            # 4. Missingness features (dim1, dim2, k, l), (m, dim2, 1, l)
            X_train_nans, X_test_nans = _missingness_features(
                X,
                mask,
                womac_config.feature_config.missingness_features,
                womac_config.reference_pool_config.jackknife_row,
            )

        if womac_config.reference_pool_config.jackknife_col:
            Y_test = Y.expand(-1, X.shape[1]).unsqueeze(-1)  # (m, n)
        else:
            Y_test = Y.unsqueeze(-1)  # (m, 1, 1)

        return WomacData(
            train_preds=X_train,
            train_missingness=X_train_nans,
            train_labels=Y_train,
            test_preds=X_test,
            test_missingness=X_test_nans,
            test_labels=Y_test,
        )

    def _fit(
        self,
        womac_data: WomacData,
        mask: torch.Tensor,
        target: Target,
        logit_feats: bool,
        model_config: ModelConfig,
    ) -> None:
        """
        Fit a WOMAC reference model according to the specified method.

        Depending on `fit_method`, this will either create an average-based model
        (UNIFORM) or a parametric model (linear or logistic regression) and train it
        on the preprocessed data.

        Args:
            womac_data (WomacData):
                Container holding preprocessed training data.
            mask (torch.Tensor):
                Boolean reference-pool mask of shape (dim1, dim2, n)
                used for uniform weighting when `fit_method` is UNIFORM.
                dim1 and dim2 are determined by jackknife_row and jackknife_col
                in the reference pool configuration.
            logit_feats (bool):
                Whether the data inputs are logit transformed (set by feature_config.logit_feats).
            model_config (ModelConfig):
                Configuration for the regression models. Passed as keyword arguments
                to the model's `.train()` method when using non-uniform fitting.
        """
        if model_config.fit_method == FitMethod.UNIFORM:
            self.model = AverageModel(target=target, logit_feats=logit_feats)
            self.model.fit(mask=mask, val_data=womac_data.test_data, val_labels=womac_data.test_labels)
        else:
            self.model = LinearModel(
                dim1=womac_data.batch_dim1,  # dim1
                dim2=womac_data.batch_dim2,  # dim2
                input_dim=womac_data.feature_dim,  # l
                target=target,
                logit_feats=logit_feats,
                device=self.device,
            )
            self.model.train(
                train_data=womac_data.train_data,
                train_labels=womac_data.train_labels,
                **model_config.kwargs_as_dict(),  # pass other model
                val_data=womac_data.test_data,
                val_labels=womac_data.test_labels,
            )

    def _predict(self, womac_data: WomacData, sigmoid_transform: bool) -> torch.Tensor:
        """Generate reference predictions for the test split.

        Uses the trained model to predict on the preprocessed test data stored
        in the WomacData container.

        Args:
            womac_data (WomacData): Container holding preprocessed tensors.
            sigmoid_transform (bool): Whether to apply a sigmoid transformation
                to the model's output. If True, the model's `predict_proba` method
                is used; otherwise, `predict` is used.

        Returns:
            torch.Tensor: Predicted reference values with shape (m, dim2, 1).
        """
        if sigmoid_transform:
            return self.model.predict_proba(womac_data.test_data)  # (m, dim2, 1)
        else:
            return self.model.predict(womac_data.test_data)  # (m, dim2, 1)
