import torch

from womac.config import FeaturizeAs
from womac.utils import jackknife_expansion


def _featurize_and_center(
    X_train: torch.Tensor,
    X_test: torch.Tensor,
    mask: torch.Tensor,
    pred_features: FeaturizeAs,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Featurize (aggregate predictions or keep separate) and center training and test data.

    Args:
        X_train: Tensor of shape (dim1, k, n). If jackknife_row is True,
            k is m-1 and dim1 is m; otherwise k is m and dim1 is 1.
        X_test: Tensor of shape (m, 1, n), for m tasks and n experts.
        mask: Tensor of shape (dim1, dim2, n). Mask indicating which experts can be used for
          each task. If jackknife_row is True, dim1 is m; otherwise dim1 is 1.
        If jackknife_col is True, dim2 is n; otherwise dim2 is 1.
        pred_features: FeaturizeAs. Specifies AGGREGATE or SEPARATE featurization.

    Returns:
        A tuple of (X_train, X_test):
            X_train: Centered training tensor. Shapes:
                - AGGREGATE: (dim1, dim2, k, 1)
                - SEPARATE:  (dim1, dim2, k, n)
            X_test: Centered test tensor. Shapes:
                - AGGREGATE: (m, dim2, 1, 1)
                - SEPARATE:  (m, dim2, 1, n)
    """

    if pred_features == FeaturizeAs.AGGREGATE:
        average_mask = mask / torch.sum(mask, dim=-1, keepdims=True)  # (dim1, dim2, n)

        # (dim1, dim2, n) @ (dim1, n, k) -> (dim1, dim2, k) -> (dim1, dim2, k, 1)
        X_train = torch.matmul(average_mask, X_train.transpose(1, 2)).unsqueeze(3)
        # (dim1, dim2, 1, n) @ (m, 1, n, 1) -> (m, dim2, 1, 1)
        X_test = torch.matmul(average_mask.unsqueeze(-2), X_test.unsqueeze(-1))

    else:
        # (dim1, 1, k, n) * (dim1, dim2, 1, n) -> (dim1, dim2, k, n)
        X_train = X_train.unsqueeze(1) * mask.unsqueeze(2)
        # (m, 1, 1, n) * (dim1, dim2, 1, n) -> (m, dim2, 1, n)
        X_test = X_test.unsqueeze(1) * mask.unsqueeze(2)

    X_mean = torch.mean(X_train, dim=[2, 3], keepdim=True)  # (dim1, dim2, 1, 1)
    X_std = torch.std(X_train, dim=[2, 3], keepdim=True)  # (dim2, dim2, 1, 1)
    X_train = (X_train - X_mean) / X_std  # (dim1, dim2, k, l)
    X_test = (X_test - X_mean) / X_std  # (m, dim2, 1, l)

    return X_train, X_test


def _missingness_features(
    X: torch.Tensor,
    mask: torch.Tensor,
    missingness_features: FeaturizeAs,
    jackknife_row: bool,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compute missingness indicator features for training and test data.

    Args:
        X: Tensor of shape (m, n) with input observations that may contain NaNs.
        mask: Tensor of shape (dim1, dim2, n). Mask indicating which experts can be used for
            each task. If jackknife_row is True, dim1 is m; otherwise dim1 is 1. If
            jackknife_col is True, dim2 is n; otherwise dim2 is 1.
        missingness_features: FeaturizeAs enum specifying featurization method
            None or (AGGREGATE or SEPARATE).
        jackknife_row: Boolean flag indicating whether to apply jackknife expansion
            to rows of X, producing k = m-1 training rows if True.

    Returns:
        Tuple of (X_train_nans, X_test_nans):
            X_train_nans: Tensor or None containing training missingness features:
                - AGGREGATE: shape (dim1, dim2, k, 1)
                - SEPARATE: shape (dim1, dim2, k, n).
                - NONE: None.
                If jackknife_row is True, k = m-1; otherwise k = m.
            X_test_nans: Tensor or None containing test missingness features:
                - AGGREGATE: shape (m, dim2, 1, 1).
                - SEPARATE: shape (m, dim2, 1, n).
                - NONE: None.
    """

    if jackknife_row:
        X_transf = jackknife_expansion(X)  # (m, m-1, n)
        train_nan_mask = torch.isnan(X_transf).float()  # (m, m-1, n)
    else:
        X_transf = X.unsqueeze(0)  # (1, m, n)
        train_nan_mask = torch.isnan(X_transf).float()  # (1, m, n)
    test_nan_mask = torch.isnan(X).float().unsqueeze(1)  #  (m, 1, n)

    if missingness_features == FeaturizeAs.AGGREGATE:
        average_mask = mask / torch.sum(mask, dim=-1, keepdim=True)  # (dim1, dim2, n)

        # (dim1, dim2, n) @ (dim1, n, k) -> (dim1, dim2, k) -> (dim1, dim2, k, 1)
        X_train_nans = torch.matmul(average_mask, train_nan_mask.transpose(1, 2)).unsqueeze(3)
        # (dim1, dim2, 1, n) @ (m, 1, n, 1) -> (m, dim2, 1, 1)
        X_test_nans = torch.matmul(average_mask.unsqueeze(-2), test_nan_mask.unsqueeze(-1))

    elif missingness_features == FeaturizeAs.SEPARATE:
        # (dim1, 1, k, n) * (dim1, dim2, 1, n) -> (dim1, dim2, k, n)
        X_train_nans = train_nan_mask.unsqueeze(1) * mask.unsqueeze(2)
        # (m, 1, 1, n) * (dim1, dim2, 1, n) -> (m, dim2, 1, n)
        X_test_nans = test_nan_mask.unsqueeze(1) * mask.unsqueeze(2)

    else:
        X_train_nans = None
        X_test_nans = None

    return X_train_nans, X_test_nans  # (dim1, dim2, k, l), (m, dim2, 1, l) or None, None
