import logging
from typing import Optional

import torch
import torch.nn as nn
from tqdm import tqdm

from womac.config import ModelResult, Target

logger = logging.getLogger(__name__)


class AverageModel:
    """Model that computes predictions by averaging selected experts with uniform weights.

    Attributes:
        coef_ (None or torch.Tensor): Coefficients of the model, initialized as None.
        target (Target): The target object passed during initialization.
        model_result (None or ModelResult): The result of the model computation, initialized as None.
        logit_feats (bool): Indicates whether the inputs are logit-transformed.
    """

    def __init__(self, target: Target, logit_feats: bool):
        """
        Initialize the AverageModel class.

        Args:
            target (Target): The target object that defines the model's target variable.
            logit_feats (bool): A flag indicating whether the target variable is logit-transformed.

        Attributes:
            coef_ (None or torch.Tensor): Coefficients of the model, initialized as None.
            target (Target): The target object passed during initialization.
            model_result (None or ModelResult): The result of the model computation, initialized as None.
            logit_feats (bool): Indicates whether the inputs are logit-transformed (ONLY for binary targets).
        """
        self.coef_ = None
        self.target = target
        self.model_result = None
        self.logit_feats = logit_feats
        if self.logit_feats and self.target != Target.BINARY:
            raise ValueError("logit_feats can only be True for binary targets.")

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

    def fit(
        self,
        mask: torch.Tensor,
        val_data: Optional[torch.Tensor] = None,
        val_labels: Optional[torch.Tensor] = None,
    ) -> None:
        """Set coefficients for averaging based on a one-hot mask.
        If the target is BINARY:
            - If logit_feats is True, model will average over logits (geometric pooling).
            Val loss computed as cross-entropy loss (after a sigmoid) against a binary label.
            - If logit_feats is False, model will average over probabilities (arithmetic pooling).
            Val loss computed directly as cross-entropy loss against a binary label.
        If the target is CONTINUOUS, model will average over continuous values (arithmetic pooling).
        Val loss computed as MSE loss against a continuous label.

        Args:
            mask (Tensor): One-hot tensor of shape (dim1, dim2, n) where mask[i, j, k] = 1
                indicates that for task i and expert j, expert k is included in the average.
                dim1 is m if jackknife_row is True, otherwise 1;
                dim2 is n if jackknife_col is True, otherwise 1.
            val_data (Optional[Tensor]): Validation data tensor of shape [m, dim2, 1, l].
            val_targets (Optional[Tensor]): Validation targets tensor of shape [m, dim2, 1].

        Returns:
            None. Sets self.coef_ to a Tensor of shape (m, n, n) containing normalized weights.
        """
        val_losses = []

        self.coef_ = mask / torch.sum(mask, dim=2, keepdims=True)  # (dim1, dim2, n)

        if val_data is not None and val_labels is not None:
            if self.target == Target.BINARY and self.logit_feats:
                criterion = nn.BCEWithLogitsLoss()
            elif self.target == Target.BINARY and not self.logit_feats:
                criterion = nn.BCELoss()
            else:
                criterion = nn.MSELoss()
            val_preds = self.predict(val_data)
            val_loss = criterion(val_preds, val_labels)
            val_losses.append(val_loss.detach().cpu().float())

        self.model_result = ModelResult(
            coef_=self.coef_.cpu(),
            train_losses=[],
            bias_=None,
            val_losses=val_losses,
        )

    def predict(self, X: torch.Tensor) -> torch.Tensor:
        """Predict averaged outputs using mask coefficients.

        Args:
            X (Tensor): Input tensor of shape (m, 1, n), where
                m: number of tasks,
                n: number of experts,

        Returns:
            Tensor of shape (m, dim2, 1) containing averaged predictions.
        """
        # (dim1, dim2, n) @ (m, 1, n).T -> (m, dim2, 1)
        return torch.matmul(self.coef_, X.transpose(1, 2))  # (m, dim2, 1)

    def predict_proba(self, X: torch.Tensor) -> torch.Tensor:
        """Compute probabilities for binary classification.

        Args:
            X (Tensor): Input tensor of shape [m, dim2, 1, l].

        Returns:
            Tensor of shape [m, dim2, 1]: Probabilities for binary classification.

        Raises:
            ValueError: If the target is not binary or if logit transformation was not applied.
        """
        if not self.logit_feats:
            raise ValueError("predict_proba is only available when inputs are logit transformed (logit_feats=True).")
        logits = self.predict(X)
        return torch.sigmoid(logits)


class LinearModel(nn.Module):
    """Parallel linear or logistic regression models across batch dimensions.

    Implements dim1 x dim2 separate regressions (or classifiers) in parallel.
    dim1 is m if jackknife_row is True, otherwise 1;
    dim2 is n if jackknife_col is True, otherwise 1.

    Attributes:
        W (Parameter): Weight tensor of shape (dim1, dim2, input_dim, 1).
        b (Parameter): Bias tensor of shape (m, n, 1).
        target (Target): Whether the target labels/outcomes are BINARY or CONTINUOUS.
        logit_feats (bool): Whether the inputs are logit-transformed (ONLY for binary targets).
        model_result (None or ModelResult): The result of the model computation, initialized as None.
    """

    def __init__(
        self,
        dim1: int,
        dim2: int,
        input_dim: int,
        target: Target,
        logit_feats: bool,
        device: str = "cpu",
    ):
        """Initialize batch-wise linear or logistic models.

        Args:
            dim1 (int): Number of tasks (first batch dimension, m).
            dim2 (int): Number of experts (second batch dimension, n).
            input_dim (int): Number of input features (l).
            target (Target): Whether the target labels/outcomes are BINARY or CONTINUOUS.
            logit_feats (bool): Whether the inputs are logit-transformed (ONLY for binary targets).
            device (str): Device to allocate parameters.
        """
        super().__init__()
        self.W = nn.Parameter(torch.randn(dim1, dim2, input_dim, 1, device=device))  # (dim1, dim2, l, 1)
        self.b = nn.Parameter(torch.zeros(dim1, dim2, 1, device=device))  # (dim1, dim2, 1)
        self.target = target
        self.logit_feats = logit_feats
        if self.logit_feats and self.target != Target.BINARY:
            raise ValueError("logit_feats can only be True for binary targets.")
        self.model_result = None  # Will be set after training
        logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")

    def forward(self, X: torch.Tensor) -> torch.Tensor:
        """Compute forward pass for parallel linear models.

        Args:
            X (Tensor): Input tensor of shape [dim1, dim2, k, l] during training, where
                dim1: batch dim (tasks),
                dim2: batch dim (experts),
                k: number of samples to predict,
                l: features
                During testing, shape is [m, dim2, 1, l] where m is number of tasks.
        Returns:
            Tensor of shape [m, n, k]: Raw outputs (logits).
        """

        Z = torch.matmul(X, self.W).squeeze(-1)  # [dim1, dim2, k] or [m, dim2, 1]
        Z = Z + self.b  # [dim1, dim2, k] or [m, dim2, 1]
        return Z  # [dim1, dim2, k] (train) or [m, dim2, 1] (test)

    def train(
        self,
        train_data: torch.Tensor,
        train_labels: torch.Tensor,
        epochs: int,
        lr: float,
        l2: float,
        log_interval: int,
        val_data: Optional[torch.Tensor] = None,
        val_labels: Optional[torch.Tensor] = None,
    ) -> None:
        """Train all parallel models using SGD with optional L2 regularization.

        Args:
            train_data (Tensor): Training data tensor of shape [dim1, dim2, k, l].
            train_labels (Tensor): Targets tensor of shape [dim1, 1, k].
            epochs (int): Number of training epochs.
            lr (float): Learning rate.
            l2 (float): L2 regularization coefficient on weights. Defaults to 0.0.
            log_interval (int): Interval for logging training progress.
            val_data (Optional[Tensor]): Validation data tensor of shape [m, dim2, 1, l].
            val_targets (Optional[Tensor]): Validation targets tensor of shape [m, dim2, 1].
        """

        y_b = train_labels.expand(-1, train_data.shape[1], -1).float()  # [dim1, dim2, k]
        if self.target == Target.BINARY:
            criterion = nn.BCEWithLogitsLoss()
        else:
            criterion = nn.MSELoss()
        optimizer = torch.optim.SGD(self.parameters(), lr=lr)
        train_losses, val_losses = [], []

        logger.info(
            f"Training {self.__class__.__name__} model with {epochs} epochs and learning rate {lr}, l2 penalty={l2}."
        )
        for epoch in tqdm(range(epochs), desc="Training Progress"):
            optimizer.zero_grad()
            preds = self.forward(train_data)  # [dim1, dim2, k]
            loss = criterion(preds, y_b)
            if l2 > 0:
                loss = loss + l2 * self.W.pow(2).sum()
            loss.backward()
            optimizer.step()
            train_losses.append(loss.detach().cpu().float())

            if val_data is not None and val_labels is not None:
                with torch.no_grad():
                    val_preds = self.forward(val_data)
                    val_loss = criterion(val_preds, val_labels)
                    reg_loss = (l2 * self.W.pow(2).sum()).item() if l2 > 0 else 0
                val_losses.append(val_loss.detach().cpu().float())

            if (epoch + 1) % log_interval == 0 or epoch == 0:
                val_loss_msg = (
                    f"| Val Loss: {val_loss.item():.6f} | Reg Loss: {reg_loss:.6f} | Total Val Loss: {(val_loss + reg_loss).item():.6f}"
                    if val_data is not None
                    else ""
                )
                tqdm.write(f"Epoch {epoch + 1}/{epochs}, Train Loss: {loss.item():.6f}{val_loss_msg}")

        self.model_result = ModelResult(
            coef_=self.W.detach().cpu(),
            train_losses=train_losses,
            bias_=self.b.detach().cpu(),
            val_losses=val_losses,
        )

    def predict(self, X: torch.Tensor) -> torch.Tensor:
        """Compute predictions without gradient tracking.

        Args:
            X (Tensor): Input tensor of shape [m, dim2, 1, l].

        Returns:
            Tensor of shape [m, dim2, 1]: Predictions (values or probabilities).
        """
        with torch.no_grad():
            return self.forward(X)  # [m, dim2, 1]

    def predict_proba(self, X: torch.Tensor) -> torch.Tensor:
        """Compute probabilities for binary classification.

        Args:
            X (Tensor): Input tensor of shape [m, dim2, 1, l].

        Returns:
            Tensor of shape [m, dim2, 1]: Probabilities for binary classification.
        """
        if not self.logit_feats:
            raise ValueError("predict_proba is only available when inputs are logit transformed (logit_feats=True).")
        logits = self.predict(X)
        return torch.sigmoid(logits)
