import torch
import torch.nn as nn
from typing import Dict, Union, Literal, Optional

class BaseTimeSeriesModel(nn.Module):
    """
    Base class for models using an embedding client.
    
    This class handles normalization, embedding extraction, and device logic
    for time series models that use external embedding services.
    
    Parameters:
    -----------
    embedding_client : object
        Client for accessing embedding services.
    dropout : float, default=0.1
        Dropout rate for regularization.
    """

    def __init__(self, embedding_client, dropout: float = 0.1):
        super().__init__()
        self.client = embedding_client
        self.dropout = dropout

        # Buffers for normalization
        self.register_buffer("mean", torch.zeros(1))
        self.register_buffer("std", torch.ones(1))

        # Dummy data to infer embedding params
        B, T, C = 1, 128, 5  # assumes patch_size <= 128 and power of 2
        dummy_data = {"ts_array": torch.randn(B, T, C), "descriptions": B*[[f"dummy_{i}" for i in range(C)]]}
        self._infer_embedding_params(dummy_data["ts_array"], dummy_data["descriptions"])

    @property
    def device(self):
        """
        Get the device of the model.
        
        Returns:
        --------
        torch.device
            Device on which the model parameters are stored.
        """
        return next(self.parameters()).device

    def _get_stats(self, x: torch.Tensor):
        """
        Calculate and store mean and standard deviation of input data.
        
        Parameters:
        -----------
        x : torch.Tensor
            Input tensor of shape [B, T, C].
        """
        assert x.ndim == 3, f"Input must be [B, T, C], got {x.shape}"
        self.mean = x.mean(dim=1, keepdim=True)
        self.std = x.std(dim=1, keepdim=True).clamp_min(1e-8)

    @torch.no_grad()
    def get_embeddings(self, data: Dict[str, torch.Tensor]) -> torch.Tensor:
        """
         Extract embeddings via the external client (CharmClient).
        
        Parameters:
        -----------
        data : Dict[str, torch.Tensor]
            Dictionary containing 'ts_array' and 'descriptions'.
            
        Returns:
        --------
        torch.Tensor
            Embeddings with shape [B, T_, C, D] where:
            - B is batch size
            - T_ is number of patches (T / patch_size)
            - C is number of channels
            - D is embedding dimension
        """
        ts_array = data["ts_array"]
        descriptions = data["descriptions"]
        B, T, C = ts_array.shape
        T_ = T // self.patch_size

        response = self.client.embeddings.create(
            descriptions=descriptions,
            ts_array=ts_array.cpu().numpy().tolist(),
            return_tensors="torch",
            progress=False
        )

        embeds = response.embeds.to(device=self.device, non_blocking=True)
        embeds = embeds.reshape(B, T_, C, self.D)
        return embeds
    
    @torch.no_grad()
    def _infer_embedding_params(self, ts_array: torch.Tensor, descriptions: list):
        """
        Infer patch size and embedding dimension from the embedding client.
        
        Parameters:
        -----------
        ts_array : torch.Tensor
            Sample time series data.
        descriptions : list
            Descriptions for each channel.
        """
        B, T, C = ts_array.shape
        response = self.client.embeddings.create(
            descriptions=descriptions,
            ts_array=ts_array.cpu().numpy().tolist(),
            return_tensors="torch",
            batch_size=64
        )
        embeds = response.embeds
        _, T_, _, D = embeds.shape
        self.patch_size = T // T_
        self.D = D
        
    def forward(self, data: Dict[str, torch.Tensor]) -> torch.Tensor:
        """
        Obtains embeddings from the input, using cached embeddings when
        available or computing them from raw time series otherwise.
    
        Parameters:
        -----------
        data : Dict[str, torch.Tensor]
            Must contain 'descriptions' and either:
            - 'embeds': precomputed embeddings [B, T_, C, D], or
            - 'ts_array': raw time series [B, T, C].
    
        Returns:
        --------
        torch.Tensor
            Embedding tensor to be consumed by child models.
        """
    
        if "embeds" in data:
            embeds = data["embeds"]
            embeds = embeds.to(self.device, non_blocking=True)
        else:
            embeds = self.get_embeddings(data)
    
        return embeds


        
class ReconstructionModel(BaseTimeSeriesModel):
    """
    Model for time series reconstruction tasks.
    
    This model uses embeddings to reconstruct the original time series data.
    
    Parameters:
    -----------
    embedding_client : object
        Client for accessing embedding services.
    reconstructor : str or nn.Module, default="linear"
        Type of reconstructor network to use ("linear", "mlp") or a custom module.
    hidden_dim : int, default=128
        Hidden dimension size for MLP reconstructor.
    dropout : float, default=0.1
        Dropout rate for regularization.
    """
    def __init__(
        self,
        embedding_client,
        reconstructor: Union[str, nn.Module] = "linear",
        hidden_dim: int = 128,
        dropout: float = 0.1,
    ):
        super().__init__(embedding_client, dropout)
        self.hidden_dim = hidden_dim
        self.reconstructor_type = reconstructor

        # Build reconstructor
        self.reconstructor = self._build_reconstructor(reconstructor)

        # Validate reconstructor output
        self._validate_reconstructor()

    def _build_reconstructor(self, reconstructor: Union[str, nn.Module]) -> nn.Module:
        """
        Builds the reconstructor network based on provided type or module.
        
        Parameters:
        -----------
        reconstructor : str or nn.Module
            Type of reconstructor ("linear", "mlp") or a custom module.
            
        Returns:
        --------
        nn.Module
            Reconstructor network.
            
        Raises:
        -------
        ValueError
            If an invalid reconstructor type is provided.
        """
        if reconstructor is None or reconstructor == "linear":
            return nn.Sequential(
                nn.Dropout(p=self.dropout),
                nn.Linear(self.D, self.patch_size),
            )
        elif reconstructor == "mlp":
            return nn.Sequential(
                nn.Linear(self.D, self.hidden_dim),
                nn.GELU(),
                nn.Dropout(p=self.dropout),
                nn.Linear(self.hidden_dim, self.patch_size),
            )
        elif isinstance(reconstructor, nn.Module):
            return reconstructor
        else:
            raise ValueError(f"Invalid reconstructor type: {reconstructor}")

    def _validate_reconstructor(self):
        """
        Ensures reconstructor output matches expected patch size.
        
        Raises:
        -------
        ValueError
            If the reconstructor output dimension doesn't match the patch size.
        """
        dummy = torch.zeros(1, self.D)
        out = self.reconstructor(dummy)
        if out.shape[-1] != self.patch_size:
            raise ValueError(
                f"Invalid reconstructor output shape {out.shape[-1]}, expected {self.patch_size}"
            )

    def forward(self, data: Dict[str, torch.Tensor]) -> torch.Tensor:
        """
        Reconstructs original time series from latent embeddings.
        
        Parameters:
        -----------
        data : Dict[str, torch.Tensor]
            Dictionary containing 'ts_array' and 'descriptions'.
            
        Returns:
        --------
        torch.Tensor
            Reconstructed time series with shape [B, T, C].
        """
        x = data["ts_array"].to(self.device, non_blocking=True)
        self._get_stats(x)  # caches mean/std for inverse scaling
        B, T, C = x.shape
        T_ = T // self.patch_size

        # --- Encode ---
        embeds = super().forward(data)     # [B, T_, C, D]
        embeds = embeds.reshape(B * T_ * C, self.D)

        # --- Reconstruct ---
        preds = self.reconstructor(embeds)  # [B*T_*C, patch_size]
        preds = preds.reshape(B, T_, C, self.patch_size)
        preds = preds.permute(0, 1, 3, 2).reshape(B, T, C)

        # --- Denormalize ---
        return preds * self.std + self.mean



class ForecastingModel(BaseTimeSeriesModel):
    """
    Model for time series forecasting tasks.
    
    This model uses embeddings to predict future values of time series data.
    
    Parameters:
    -----------
    embedding_client : object
        Client for accessing embedding services.
    head : str or nn.Module, default="linear"
        Type of forecasting head to use ("linear", "mlp") or a custom module.
    hidden_dim : int, default=128
        Hidden dimension size for MLP head.
    dropout : float, default=0.0
        Dropout rate for regularization.
    horizon : int
        Number of future time steps to predict.
    mode : str, default="last"
        Pooling mode for embeddings ("last", "avg", "none").
    per_channel : bool, default=True
        Whether to use separate forecasting heads for each channel.
    num_channels : int, optional
        Number of channels in the time series (required if per_channel=True).
    input_size: int
    Number of time steps to use as input
    """
    def __init__(
        self,
        embedding_client,
        head: Union[str, nn.Module] = "linear",  # "linear", "mlp", or custom nn.Module
        hidden_dim: int = 128,
        dropout: float = 0.0,
        horizon: int = None,
        mode: Literal["last", "avg", "none"] = "last",
        per_channel: bool = True,
        num_channels: int = None,
        input_size: int = None,
    ):
        super().__init__(embedding_client, dropout)

        # --- Core attributes ---
        self.hidden_dim = hidden_dim
        self.horizon = horizon
        self.mode = mode
        self.per_channel = per_channel
        self.head_type = head
        self.num_channels = num_channels
        self.input_size =input_size

        # --- Sanity checks ---
        if self.horizon is None:
            raise ValueError("`horizon` must be provided.")
        if self.input_size is None:
            raise ValueError("`input_size` must be provided.")
        if self.mode not in ("last", "avg", "none"):
            raise ValueError("`mode` must be one of: 'last', 'avg', or 'none'.")
        if self.per_channel and self.num_channels is None:
            raise ValueError("`num_channels` must be specified for per-channel heads.")

        # --- Build forecasting heads ---
        channels = 1 if num_channels is None else num_channels
        B = 2
        dummy_embeds = torch.zeros(B, input_size, channels, self.D)  # [B, T_, C, D]
        dummy_pooled = self._pool_embeddings(dummy_embeds)
        in_dim = dummy_pooled.shape[-1] // self.patch_size if self.mode == "none" else self.D

        self.forecaster = self._build_forecaster(in_dim)

        # --- Validate head output dimension ---
        self._validate_head(in_dim)

    def _make_head(self, in_dim: int) -> nn.Module:
        """
        Factory method for constructing a forecasting head.
        
        Parameters:
        -----------
        in_dim : int
            Input dimension for the head.
            
        Returns:
        --------
        nn.Module
            Forecasting head module.
            
        Raises:
        -------
        ValueError
            If an unsupported head type is provided.
        """
        if isinstance(self.head_type, nn.Module):
            return self.head_type
        elif self.head_type == "linear":
            return nn.Sequential(
                nn.Dropout(p=self.dropout),
                nn.Linear(in_dim, self.horizon),
            )
        elif self.head_type == "mlp":
            return nn.Sequential(
                nn.Linear(in_dim, self.hidden_dim),
                nn.GELU(),
                nn.Dropout(p=self.dropout),
                nn.Linear(self.hidden_dim, self.horizon),
            )
        else:
            raise ValueError(f"Unsupported head type: {self.head_type}")

    def _build_forecaster(self, in_dim: int) -> nn.Module:
        """
        Builds shared or per-channel forecasting heads.
        
        Parameters:
        -----------
        in_dim : int
            Input dimension for the head.
            
        Returns:
        --------
        nn.Module or nn.ModuleList
            Forecasting head(s).
        """
        if self.per_channel:
            return nn.ModuleList([self._make_head(in_dim) for _ in range(self.num_channels)])
        else:
            return self._make_head(in_dim)

    def _validate_head(self, in_dim: int):
        """
        Runs a dummy forward pass to ensure the forecasting head is dimensionally consistent.
        
        Parameters:
        -----------
        in_dim : int
            Input dimension for the head.
            
        Raises:
        -------
        ValueError
            If the head output dimension doesn't match the horizon.
        """
        dummy_input = torch.zeros(1, in_dim)
        dummy_out = (
            self.forecaster[0](dummy_input)
            if self.per_channel
            else self.forecaster(dummy_input)
        )
        if dummy_out.shape[-1] != self.horizon:
            raise ValueError(
                f"Forecast head output {dummy_out.shape[-1]} != expected horizon {self.horizon}"
            )

    def _pool_embeddings(self, embeds: torch.Tensor) -> torch.Tensor:
        """
        Applies temporal pooling to embeddings.
        
        Parameters:
        -----------
        embeds : torch.Tensor
            Embeddings with shape [B, T_, C, D].
            
        Returns:
        --------
        torch.Tensor
            Pooled embeddings with shape depending on pooling mode.
            
        Raises:
        -------
        ValueError
            If an invalid pooling mode is provided.
        """
        if self.mode == "last":
            return embeds[:, -1, :, :]  # [B, C, D]
        elif self.mode == "avg":
            return embeds.mean(dim=1)   # [B, C, D]
        elif self.mode == "none":
            # Flatten temporal patches per channel
            embeds = embeds.permute(0, 2, 1, 3)  # [B, C, T_, D]
            return embeds.reshape(embeds.shape[0], embeds.shape[1], -1)
        else:
            raise ValueError(f"Invalid pooling mode: {self.mode}")

    def forward(self, data: Dict[str, torch.Tensor]) -> torch.Tensor:
        """
        Generate forecasts for future time steps.
        
        Parameters:
        -----------
        data : Dict[str, torch.Tensor]
            Dictionary containing 'ts_array' and 'descriptions'.
            
        Returns:
        --------
        torch.Tensor
            Forecasted values with shape [B, horizon, C].
        """
        x = data["ts_array"].to(self.device, non_blocking=True)
        self._get_stats(x)
        B, T, C = x.shape
        if T != self.input_size:
            raise ValueError(f"Input size of {self.input_size} was specified, but {T} is now being passed.")
        
        embeds = super().forward(data)          # [B, T_, C, D]
        pooled = self._pool_embeddings(embeds)  # [B, C, D_in]

        # --- Forecast ---
        if self.per_channel:
            # Efficient vectorized per-channel head application
            preds = torch.stack(
                [head(pooled[:, c, :]) for c, head in enumerate(self.forecaster)],
                dim=2,
            )  # [B, horizon, C]
        else:
            preds = self.forecaster(pooled.reshape(B * C, -1))
            preds = preds.reshape(B, C, self.horizon).permute(0, 2, 1)  # [B, horizon, C]

        # --- Denormalize ---
        return preds * self.std + self.mean


class ClassificationModel(BaseTimeSeriesModel):
    """
    Flexible classification model for time series embeddings.
    
    This model uses embeddings to classify time series data.
    
    Parameters:
    -----------
    embedding_client : object
        Client for accessing embedding services.
    num_classes : int
        Number of classes to predict.
    hidden_dim : int, default=128
        Hidden dimension size for MLP classifier.
    dropout : float, default=0.1
        Dropout rate for regularization.
    pooling_over_t : str, default="mean"
        Pooling mode for temporal dimension ("flatten", "mean", "last").
    pooling_over_channels : str, default="mean"
        Pooling mode for channel dimension ("flatten", "mean").
    classifier_type : str or nn.Module, default="mlp"
        Type of classifier to use ("linear", "mlp", "svm") or a custom module.
    kernel : str, optional
        Kernel type for SVM classifier.
    gamma : float, default=0.1
        Gamma parameter for RBF kernel in SVM.
    T : int, optional
        Expected temporal dimension (required if pooling_over_t="flatten").
    num_channels : int, optional
        Number of channels (required if pooling_over_channels="flatten").
    """

    def __init__(
        self,
        embedding_client,
        num_classes: int,
        hidden_dim: int = 128,
        dropout: float = 0.1,
        pooling_over_t: str = "mean",        # "flatten", "mean", "last"
        pooling_over_channels: str = "mean", # "flatten", "mean"
        classifier_type: Union[str, nn.Module] = "mlp",  # "linear", "mlp", "svm" or custom nn.Module
        kernel: Optional[str] = None,
        gamma: float = 0.1,
        T: Optional[int] = None,
        num_channels: Optional[int] = None,
    ):
        super().__init__(embedding_client, dropout)

        # --- Config ---
        self.num_classes = num_classes
        self.hidden_dim = hidden_dim
        self.dropout = dropout
        self.pooling_over_t = pooling_over_t
        self.pooling_over_channels = pooling_over_channels
        self.classifier_type = classifier_type
        self.kernel = kernel
        self.gamma = gamma
        self.expected_T = T
        self.expected_C = num_channels

        # --- Validate pooling requirements ---
        if pooling_over_channels == "flatten" and num_channels is None:
            raise ValueError("num_channels must be specified for 'flatten' channel pooling.")
        if pooling_over_t == "flatten" and T is None:
            raise ValueError("T must be specified for 'flatten' temporal pooling.")

        # --- Compute expected input dimensionality ---
        self.input_dim = self._compute_input_dim()
        self.classifier = self._build_classifier(self.input_dim)

        # --- Sanity check ---
        dummy = torch.zeros(1, self.input_dim)
        dummy_out = self.classifier(dummy)
        if dummy_out.shape[-1] != num_classes:
            raise ValueError("Classifier output dimensionality mismatch with num_classes.")

    def _get_pooled(self, embeds: torch.Tensor) -> torch.Tensor:
        """
        Apply pooling over time and channels.
        
        Parameters:
        -----------
        embeds : torch.Tensor
            Embeddings with shape [B, T_, C, D].
            
        Returns:
        --------
        torch.Tensor
            Pooled embeddings with shape [B, input_dim].
            
        Raises:
        -------
        ValueError
            If an invalid pooling mode is provided.
        """
        B, T, C, D = embeds.shape

        # --- Time pooling ---
        if self.pooling_over_t == "mean":
            embeds = embeds.mean(dim=1)  # (B, C, D)
        elif self.pooling_over_t == "last":
            embeds = embeds[:, -1]       # (B, C, D)
        elif self.pooling_over_t == "flatten":
            embeds = embeds
        else:
            raise ValueError(f"Invalid pooling_over_t: {self.pooling_over_t}")
        embeds = embeds.reshape(B,-1, C, D)
        
        # --- Channel pooling ---
        if self.pooling_over_channels == "mean":
            embeds = embeds.mean(dim=2)  # (B, D)
        elif self.pooling_over_channels == "flatten":
            embeds = embeds
        else:
            raise ValueError(f"Invalid pooling_over_channels: {self.pooling_over_channels}")

        return embeds.reshape(B, -1)  # final shape (B, input_dim)

    def _build_classifier(self, input_dim: int) -> nn.Module:
        """
        Construct classifier layer or module.
        
        Parameters:
        -----------
        input_dim : int
            Input dimension for the classifier.
            
        Returns:
        --------
        nn.Module
            Classifier module.
            
        Raises:
        -------
        ValueError
            If an unknown classifier type is provided.
        """
        if isinstance(self.classifier_type, nn.Module):
            return self.classifier_type  # custom module directly

        if self.classifier_type == "linear":
            return nn.Linear(input_dim, self.num_classes)
        elif self.classifier_type == "mlp":
            return nn.Sequential(
                nn.Linear(input_dim, self.hidden_dim),
                nn.GELU(),
                nn.Dropout(self.dropout),
                nn.Linear(self.hidden_dim, self.num_classes),
            )
        elif self.classifier_type == "svm":
            return SVMHead(input_dim, self.num_classes, kernel=self.kernel or "linear", gamma=self.gamma)
        else:
            raise ValueError(f"Unknown classifier_type: {self.classifier_type}")

    def _compute_input_dim(self) -> int:
        """
        Derive input dimensionality based on pooling choices.
        
        Returns:
        --------
        int
            Input dimension for the classifier.
        """
        T = self.expected_T or 1
        C = self.expected_C or 1

        if self.pooling_over_t == "flatten":
            t_factor = T // self.patch_size
        else:
            t_factor = 1

        if self.pooling_over_channels == "flatten":
            c_factor = C
        else:
            c_factor = 1

        return t_factor * c_factor * self.D

    def forward(self, data: Dict[str, torch.Tensor]) -> torch.Tensor:
        """
        Classify time series data.
        
        Parameters:
        -----------
        data : Dict[str, torch.Tensor]
            Dictionary containing 'ts_array' and 'descriptions'.
            
        Returns:
        --------
        torch.Tensor
            Classification logits with shape [B, num_classes].
            
        Raises:
        -------
        ValueError
            If the input dimensions don't match the expected dimensions.
        """
        x = data["ts_array"]
        _, T, C = x.shape

        if self.expected_T is not None and T != self.expected_T:
            raise ValueError(f"Temporal dimension changed from expected {self.expected_T} to {T}")
        if self.expected_C is not None and C != self.expected_C:
            raise ValueError(f"Channel dimension changed from expected {self.expected_C} to {C}")

        embeds = super().forward(data)     # [B, T, C, D]
        pooled = self._get_pooled(embeds)
        return self.classifier(pooled)


class SVMHead(nn.Module):
    """
    Support Vector Machine head for classification tasks.
    
    This module implements a differentiable SVM-like classifier that can be used
    as a head in neural networks.
    
    Parameters:
    -----------
    input_dim : int
        Input dimension.
    num_classes : int
        Number of classes to predict.
    kernel : str, default="linear"
        Kernel type ("linear" or "rbf").
    gamma : float, default=0.1
        Gamma parameter for RBF kernel.
    """
    def __init__(self, input_dim, num_classes, kernel="linear", gamma=0.1):
        super().__init__()
        self.kernel = kernel
        self.gamma = gamma
        self.num_classes = num_classes

        self.support_vectors = nn.Parameter(torch.randn(num_classes, input_dim))
        self.weights = nn.Parameter(torch.randn(num_classes))
        self.bias = nn.Parameter(torch.zeros(num_classes))

        nn.init.xavier_uniform_(self.support_vectors)
        nn.init.zeros_(self.bias)

    def forward(self, x):
        """
        Compute classification logits.
        
        Parameters:
        -----------
        x : torch.Tensor
            Input tensor with shape [B, input_dim].
            
        Returns:
        --------
        torch.Tensor
            Classification logits with shape [B, num_classes].
            
        Raises:
        -------
        ValueError
            If an unsupported kernel is provided.
        """
        # x: [B, D]
        if self.kernel == "linear":
            sims = x @ self.support_vectors.T  # [B, num_classes]
        elif self.kernel == "rbf":
            x_norm = (x ** 2).sum(dim=1, keepdim=True)
            sv_norm = (self.support_vectors ** 2).sum(dim=1)
            dists = x_norm + sv_norm - 2 * x @ self.support_vectors.T
            sims = torch.exp(-self.gamma * dists)
        else:
            raise ValueError(f"Unsupported kernel: {self.kernel}")

        logits = sims * self.weights + self.bias
        return logits
