import torch

from womac.config import MFSConfig, MFSMethod, ReferencePoolConfig


def _build_base_reference_pool(
    X: torch.Tensor,
    reference_pool_config: ReferencePoolConfig,
    cols_without_min_resp: torch.Tensor,
    device: torch.device = torch.device("cpu"),
) -> torch.Tensor:
    """
    Create initial reference pool mask for tasks and experts.

    Exclusions applied for:
        1. Insufficient responses per task/expert.
        2. Held-out expert column if jackknife_col is True.
        3. NaN entries in held-out rows if drop_row_nans is True.

    Args:
        X (torch.Tensor): Data matrix shape (m, n) with NaNs for missing.
        reference_pool_config (ReferencePoolConfig): Flags controlling jackknife and NaN rules.
        cols_without_min_resp (torch.Tensor): bool mask of shape (n,) marking insufficient experts.
        device (torch.device): Device for mask tensor allocation.

    Returns:
        torch.Tensor: Bool mask of shape (dim1, dim2, n), where
            dim1 = m if jackknife_row else 1,
            dim2 = n if jackknife_col else 1.
    """
    m, n = X.shape
    nan_locations = torch.where(torch.isnan(X))
    dim1 = m if reference_pool_config.jackknife_row else 1
    dim2 = n if reference_pool_config.jackknife_col else 1
    reference_pool = torch.ones((dim1, dim2, n), device=device, dtype=torch.bool)

    # 1. Exclude rows/cols with insufficient responses
    reference_pool[:, :, cols_without_min_resp] = False

    # 2. if jackknife_col, exclude all expert j at held-out col j
    if reference_pool_config.jackknife_col:
        reference_pool[:, torch.arange(n), torch.arange(n)] = False

    # 3. if drop_row_nans, exclude NaNs in the held-out row
    if reference_pool_config.drop_row_nans:
        reference_pool[nan_locations[0], :, nan_locations[1]] = False

    return reference_pool


def _marginal_feature_screening(
    base_pool: torch.Tensor,
    errors: torch.Tensor,
    mfs_config: MFSConfig,
) -> torch.Tensor:
    """
    Apply marginal feature screening to refine reference pool.

    Filters base_pool entries based on error thresholds or percentiles.

    Args:
        base_pool (torch.Tensor): Bool mask shape (dim1, dim2, n) to update with MFS.
        errors (torch.Tensor): Error scores shape (dim1, n).
        mfs_config (MFSConfig): Configuration for screening method and cutoff.

    Returns:
        torch.Tensor: Bool mask shape (dim1, dim2, n) after screening.
        dim1 = m if jackknife_row else 1,
        dim2 = n if jackknife_col else 1.

    Raises:
        ValueError: For invalid mfs_method or if any (i,j) has no valid references.
    """
    if mfs_config.mfs_method == MFSMethod.THRESHOLD:
        threshold_mask = errors <= mfs_config.mfs_cutoff  # (dim1, n)
        reference_pool = base_pool & threshold_mask.unsqueeze(1)  # (dim1, dim2, n)

    elif mfs_config.mfs_method == MFSMethod.PERCENTILE or mfs_config.mfs_method == MFSMethod.AUTO:
        errors_expanded = errors.unsqueeze(1).repeat(1, base_pool.shape[1], 1)  # (dim1, dim2, n)
        errors_expanded[~base_pool] = torch.nan  # (dim1, dim2, n)

        threshold_mask = torch.nanquantile(
            errors_expanded,
            mfs_config.mfs_cutoff / 100.0,
            dim=-1,
            keepdim=True,
        )  # shape (dim1, dim2, 1)
        reference_pool = base_pool & (errors_expanded <= threshold_mask)  # (dim1, dim2, 1)
    elif mfs_config.mfs_method is MFSMethod.NONE:
        reference_pool = base_pool
    else:
        raise ValueError(f"Invalid mfs_config.mfs_method: {mfs_config.mfs_method}")

    if torch.any(torch.sum(reference_pool, dim=-1) == 0):
        raise ValueError("Some indices have no reference pool. Try raising the mfs_cutoff.")

    return reference_pool
