from typing import Optional

import numpy as np
import torch

from womac.config import MissingDataStrategy, TensorLike


def _cast_to_tensor(X: TensorLike, device: torch.device = torch.device("cpu")) -> torch.Tensor | None:
    """
    Converts the input data to a torch tensor. Reshapes 1D arrays to 2D with shape (-1, 1).

    Args:
        X (TensorLike): The input data to be converted. It can be either a numpy array,
            a torch tensor, or None.
        device (torch.device): The device on which the resulting tensor should be placed.
            Defaults to CPU.

    Returns:
        torch.Tensor | None: A 2D torch tensor representation of the input data, or None

    Raises:
        TypeError: If the input is not a numpy array, torch tensor, or None.
    """
    if isinstance(X, np.ndarray):
        if X.ndim == 1:
            X = X.reshape(-1, 1)
        return torch.tensor(X, device=device, dtype=torch.float32)
    if isinstance(X, torch.Tensor):
        if X.ndim == 1:
            X = X.view(-1, 1)
        return X.clone().to(device)
    if X is None:
        return None
    raise TypeError(f"Expected input to be a numpy array, torch tensor, or None, but got {type(X)}.")


def _apply_logit_transform(X: torch.Tensor, logit_transform: bool) -> torch.Tensor:
    """
    Applies a logit transformation to input tensor values if requested.

    Args:
        X (torch.Tensor): Tensor containing values strictly between 0 and 1.
        logit_transform (bool): Whether to apply the logit transform.

    Returns:
        torch.Tensor: Logit-transformed tensor when logit_transform is True; otherwise original tensor.

    Raises:
        ValueError: If logit_transform is True and any values in X are <= 0 or >= 1.
    """
    if logit_transform:
        if torch.any((X <= 0) | (X >= 1)):  # scalar bool
            raise ValueError(
                "All entries must be strictly between 0 and 1 for logit transform. Check your data or disable logit transformation."
            )
        return torch.logit(X)
    return X


def _get_indices_without_min_responses(
    X: torch.Tensor,
    min_responses_per_row: int,
    min_responses_per_col: int,
) -> tuple[torch.Tensor, torch.Tensor]:
    """
    Compute boolean masks indicating which rows and columns of the input
    tensor have fewer valid (non-NaN) entries than the configured minimum thresholds:
    `self.min_responses_per_task` for rows (tasks) and
    `self.min_responses_per_expert` for columns (experts).

    Args:
        X (torch.Tensor): A 2D tensor of shape (m, n). Missing or invalid responses
            should be encoded as NaN.

    Returns:
        Tuple[torch.Tensor, torch.Tensor]: A tuple containing:
            rows_without_min_resp (torch.Tensor): A 1D boolean tensor of length m tasks.
                True at position i indicates that task i has fewer than
                `self.min_responses_per_task` valid responses.
            cols_without_min_resp (torch.Tensor): A 1D boolean tensor of length n experts.
                True at position j indicates that expert j has fewer than
                `self.min_responses_per_expert` valid responses.
    """
    is_not_nan = ~torch.isnan(X)
    rows_without_min_resp = torch.sum(is_not_nan, dim=1) < min_responses_per_row
    cols_without_min_resp = torch.sum(is_not_nan, dim=0) < min_responses_per_col

    return rows_without_min_resp, cols_without_min_resp


######### Jackknife Expansion Function #########


def jackknife_expansion(X: torch.Tensor) -> torch.Tensor:
    """
    Generate jackknife samples by excluding each row in turn.

    Args:
        X (torch.Tensor): A 2D tensor of shape (m, *).

    Returns:
        torch.Tensor: A tensor of shape (m, m-1, *) where each slice along the first
            dimension is X with the corresponding row removed.
    """
    m = X.size(0)
    X_rep = X.unsqueeze(0).expand(m, -1, *X.shape[1:])  # (m, m, *)
    mask = ~torch.eye(m, dtype=torch.bool, device=X.device)  # (m, m)
    return X_rep[mask].view(m, m - 1, *X.shape[1:])


def score_mse(
    X: torch.Tensor,
    Y: torch.Tensor,
    missing_data_strategy: MissingDataStrategy,
    jackknife_row: bool = False,
    apply_logit_transform: bool = False,
) -> torch.Tensor:
    """
    Compute mean squared error (MSE) between predictions and targets for each expert.

    If `jackknife_row` is True, performs leave-one-row-out (jackknife) MSE for each row.
    Otherwise, computes MSE for each expert over all rows.

    Args:
        X (torch.Tensor): Prediction matrix (m, n), may contain NaNs.
        Y (torch.Tensor): Outcome vector or matrix (m,) or (m, 1), may contain NaNs.
        missing_data_strategy (MissingDataStrategy): How to handle NaNs in X.
        jackknife_row (bool): If True, use jackknife; else, compute over all rows.
        apply_logit_transform (bool): Whether to apply logit transformation to X before scoring.

    Returns:
        torch.Tensor: MSE values, shape (m, n) if jackknife, else (1, n).

    Raises:
        ValueError: If NaNs present in errors (due to too many NaNs in Y).
    """
    X, Y = _cast_to_tensor(X, device=X.device), _cast_to_tensor(Y, device=Y.device)
    # Optionally expand for jackknife
    if jackknife_row:
        X = jackknife_expansion(X)  # (m, m-1, n)
        Y = jackknife_expansion(Y)  # (m, m-1, 1)
    else:
        X = X.unsqueeze(0)  # (1, m, n)
        Y = Y.unsqueeze(0)  # (1, m, 1)

    # Impute or ignore NaNs in X_sc (apply_logit_transform=False since scoring directly against targets Y_sc)
    X_sc, _ = impute_data(X, missing_data_strategy, Y=Y)  # (m, m-1, n) or (1, m, n)

    # Apply logit transformation if configured
    X_sc = _apply_logit_transform(X_sc, apply_logit_transform)  # (m, m-1, n) or (1, m, n)

    # Compute mean squared errors
    errors = torch.nanmean((X_sc - Y) ** 2, dim=-2)  # (m, n) (if jackknife_row=True) or (1, n) (if jackknife_row=False)

    # Ensure no NaNs remain
    if torch.isnan(errors).any():
        raise ValueError("NaN values in errors matrix. Reduce `min_responses_per_task` or provide more data.")

    return errors


######### Imputation Functions #########


def impute_column_mean(X: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    """
    Imputes missing entries in X with the mean of each column (second-to-last dimension).

    Args:
        X (torch.Tensor): Input tensor of any shape with NaN values along the second-to-last
          dimension.

    Returns:
        Tuple[torch.Tensor, torch.Tensor]:
            - imputed_X (torch.Tensor): Tensor with NaNs replaced by column means.
            - column_mean (torch.Tensor): Tensor of column means with shape (*, 1, n).
    """
    column_mean = torch.nanmean(X, dim=-2, keepdims=True)  # (*, 1, n)
    return torch.where(torch.isnan(X), column_mean, X), column_mean  # (*, k, n)


def impute_outcome_mean(X: torch.Tensor, Y: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    """
    Impute missing entries in X using the mean of outcome tensor Y.

    Args:
        X (torch.Tensor): Input tensor with NaN values to be imputed.
        Y (torch.Tensor): Outcome tensor used to compute the fill-in mean.

    Returns:
        Tuple[torch.Tensor, torch.Tensor]:
            - imputed_X (torch.Tensor): Tensor with NaNs replaced by the mean of Y.
            - Y_mean (torch.Tensor): Scalar or tensor containing the mean of Y.
    """
    Y_mean = torch.nanmean(Y, dim=-2, keepdim=True)  # (1, 1)
    imputed = torch.where(torch.isnan(X), Y_mean, X)  # (*, k, n)
    return imputed, Y_mean


def impute_data(
    X: torch.Tensor,
    missing_data_strategy: MissingDataStrategy,
    Y: Optional[torch.Tensor],
) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
    """
    Impute missing data in tensor X according to a specified strategy.

    Args:
        X (torch.Tensor): Input tensor with potential NaN values.
        missing_data_strategy (MissingDataStrategy): Strategy for handling missing data.
        Y (Optional[torch.Tensor]): Optional outcome tensor required by some strategies.

    Returns:
        Tuple[torch.Tensor, Optional[torch.Tensor]]: A tuple containing:
            - imputed_X (torch.Tensor): Tensor with missing values handled.
            - imputation_info (Optional[torch.Tensor]): The tensor used for imputation.

    Raises:
        ValueError: If an invalid missing_data_strategy is provided.
    """
    if missing_data_strategy == MissingDataStrategy.IMPUTE_COLUMN_MEAN:
        return impute_column_mean(X)
    elif missing_data_strategy == MissingDataStrategy.IMPUTE_OUTCOME_MEAN:
        return impute_outcome_mean(X, Y)
    elif missing_data_strategy == MissingDataStrategy.IGNORE_NANS:
        return X, None
    else:
        raise ValueError(f"Invalid missing data strategy: {missing_data_strategy}")


def impute_with(X: torch.Tensor, impute_using: torch.Tensor) -> torch.Tensor:
    """
    Impute missing entries in X with values from another tensor.

    Args:
        X (torch.Tensor): Tensor with NaN values to be replaced.
        impute_using (torch.Tensor): Tensor providing replacement values wherever X is NaN.

    Returns:
        torch.Tensor: Tensor with NaNs in X replaced by corresponding values from impute_using.
    """
    return torch.where(torch.isnan(X), impute_using, X)
