# src/models/context.py
from __future__ import annotations

import config as cfg
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from typing import Optional, Tuple, Dict, List


class ContextMemory:
    """Stores hidden states, predictions and features across sequences
    to create a long-term memory beyond SEQ_LEN.

    All stored tensors live on CPU to avoid repeated GPU↔CPU sync
    in the hot training loop.
    """

    def __init__(
        self,
        capacity: int = 128,
        hidden_dim: int = 48,
        context_feature_dim: int = 64,
        device: str = "cpu",
    ):
        self.capacity = capacity
        self.hidden_dim = hidden_dim
        self.context_feature_dim = context_feature_dim
        self.device = device

        self.hidden_states: List[torch.Tensor] = []
        self.cell_states: List[torch.Tensor] = []
        self.predictions: List[torch.Tensor] = []
        self.features: List[torch.Tensor] = []
        self.timestep: int = 0

    def push(
        self,
        hidden: torch.Tensor,
        cell: torch.Tensor,
        pred_long: torch.Tensor,
        pred_short: torch.Tensor,
        features: torch.Tensor,
    ) -> None:
        """Append one timestep to the context memory.

        If the inputs are batched (batch > 1) we aggregate by mean
        first, then move to CPU.
        """
        def _to_1d(t: torch.Tensor) -> torch.Tensor:
            """Collapse single-element dimensions → (D,)"""
            if t.ndim == 0:
                return t.unsqueeze(0)
            if t.ndim == 2 and t.shape[-1] == 1:
                return t.squeeze(-1)
            return t

        # When batch > 1, average across the batch dimension
        if hidden.ndim == 2 and hidden.shape[0] > 1:
            h = hidden.mean(dim=0)
            c = cell.mean(dim=0)
            p_long = _to_1d(pred_long.mean(dim=0))
            p_short = _to_1d(pred_short.mean(dim=0))
            p = torch.cat([p_long, p_short])
            f = features.mean(dim=0)
        else:
            h = hidden.squeeze(0)
            c = cell.squeeze(0)
            p_long = _to_1d(pred_long)
            p_short = _to_1d(pred_short)
            p = torch.cat([p_long, p_short])
            f = features.squeeze(0)

        # Detach and move to CPU once — keeps the hot loop GPU-bound
        h = h.detach().cpu()
        c = c.detach().cpu()
        p = p.detach().cpu()
        f = f.detach().cpu()

        self.hidden_states.append(h)
        self.cell_states.append(c)
        self.predictions.append(p)
        self.features.append(f)
        self.timestep += 1

        if len(self.hidden_states) > self.capacity:
            del self.hidden_states[0]
            del self.cell_states[0]
            del self.predictions[0]
            del self.features[0]

    def get_context_sequence(self, seq_len: Optional[int] = None, device: Optional[torch.device] = None) -> Optional[torch.Tensor]:
        """Return a (T, concat_dim) context tensor.

        Each row = [h, c, pred, features].
        """
        if not self.hidden_states:
            return None

        n = min(seq_len or len(self.hidden_states), len(self.hidden_states))

        h_stack = torch.stack(self.hidden_states[-n:], dim=0)
        c_stack = torch.stack(self.cell_states[-n:], dim=0)
        p_stack = torch.stack(self.predictions[-n:], dim=0)
        f_stack = torch.stack(self.features[-n:], dim=0)

        context_vec = torch.cat([h_stack, c_stack, p_stack, f_stack], dim=-1)

        if device is not None:
            context_vec = context_vec.to(device)
        return context_vec

    def get_prediction_history(self, seq_len: Optional[int] = None) -> Optional[torch.Tensor]:
        """Return history of predictions (long, short) with shape (T, 2)."""
        if not self.predictions:
            return None

        n = min(seq_len or len(self.predictions), len(self.predictions))
        p_stack = torch.cat(self.predictions[-n:], dim=0)
        return p_stack.to(self.device)

    def clear(self) -> None:
        """Clear the memory."""
        self.hidden_states.clear()
        self.cell_states.clear()
        self.predictions.clear()
        self.features.clear()
        self.timestep = 0

    def __len__(self) -> int:
        return len(self.hidden_states)


class ContextEncoder(nn.Module):
    """Encodes context memory into a fixed-size vector via a Transformer
    encoder with positional encoding and attention pooling."""

    def __init__(
        self,
        input_dim: int,
        hidden_dim: int = 64,
        num_heads: int = 4,
        num_layers: int = 2,
        dropout: float = 0.15,
        max_seq_len: int = 128,
    ):
        super().__init__()
        self.input_proj = nn.Linear(input_dim, hidden_dim)
        self.positional_encoding = nn.Parameter(
            torch.randn(max_seq_len, hidden_dim) * 0.02
        )
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=hidden_dim,
            nhead=num_heads,
            dim_feedforward=hidden_dim * 4,
            dropout=dropout,
            batch_first=True,
            activation="gelu",
            layer_norm_eps=1e-5,
        )
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
        self.norm = nn.LayerNorm(hidden_dim, eps=1e-5)
        self.dropout = nn.Dropout(dropout)
        self.cross_attention = nn.MultiheadAttention(
            embed_dim=hidden_dim,
            num_heads=num_heads,
            dropout=dropout,
            batch_first=True,
        )

    def forward(self, context_seq: torch.Tensor) -> torch.Tensor:
        """
        Args:
            context_seq: (batch, seq_len, input_dim) or (seq_len, input_dim)
        Returns:
            context_vector: (batch, hidden_dim)
        """
        if context_seq.ndim == 2:
            context_seq = context_seq.unsqueeze(0)

        batch, seq_len, _ = context_seq.shape

        x = self.input_proj(context_seq)
        pos = self.positional_encoding[:seq_len].unsqueeze(0).expand(batch, -1, -1)
        x = x + pos
        x = self.dropout(x)

        encoded = self.transformer(x)

        # Cross-attention with query = first element
        # [CLS]-like token
        query = encoded[:, 0:1, :]
        attended, _ = self.cross_attention(query, encoded, encoded)
        enriched = self.dropout(attended.squeeze(1))

        encoded = encoded + 0.3 * enriched.unsqueeze(1)
        encoded = self.norm(encoded)

        # Attention pooling
        pooling_weights = torch.softmax(encoded[:, :, 0:1], dim=1)
        context_vector = torch.sum(pooling_weights * encoded, dim=1)

        return context_vector


class ContextAwareAttention(nn.Module):
    """Context-aware attention where query = current LSTM hidden state
    and Key/Value = encoded context vector."""

    def __init__(self, hidden_dim: int, context_dim: int, num_heads: int = 4):
        super().__init__()
        self.num_heads = num_heads
        self.scale = (hidden_dim // num_heads) ** -0.5

        self.q_proj = nn.Linear(hidden_dim, hidden_dim)
        self.k_proj = nn.Linear(context_dim, hidden_dim)
        self.v_proj = nn.Linear(context_dim, hidden_dim)
        self.out_proj = nn.Linear(hidden_dim, hidden_dim)

        self.gate = nn.Sequential(
            nn.Linear(hidden_dim * 2, hidden_dim),
            nn.GELU(),
            nn.Dropout(0.1),
            nn.Linear(hidden_dim, hidden_dim),
            nn.Sigmoid(),
        )

        self.residual_gate = nn.Sequential(
            nn.Linear(hidden_dim * 2, hidden_dim),
            nn.Sigmoid(),
        )

        self.norm = nn.LayerNorm(hidden_dim, eps=1e-5)
        self.scale_proj = nn.Parameter(torch.ones(1) * 0.5)

    def forward(
        self,
        query: torch.Tensor,
        context: torch.Tensor,
    ) -> torch.Tensor:
        """
        Args:
            query: (batch, seq_len, hidden_dim)
            context: (batch, context_dim) — or (context_dim,) when broadcast by expand
        Returns:
            enriched: (batch, seq_len, hidden_dim)
        """
        batch, seq_len, _ = query.shape

        q = self.q_proj(query)
        k = self.k_proj(context).unsqueeze(1).expand(-1, seq_len, -1)
        v = self.v_proj(context).unsqueeze(1).expand(-1, seq_len, -1)

        attn_weights = torch.sum(q * k, dim=-1, keepdim=True) * self.scale * self.scale_proj
        attn_weights = F.softmax(attn_weights, dim=1)

        # Emphasise important context elements
        attn_weights = attn_weights ** 1.5
        attn_weights = attn_weights / (attn_weights.sum(dim=1, keepdim=True) + 1e-9) * attn_weights.shape[1]
        attn_weights = torch.clamp(attn_weights, max=3.0)

        attended = attn_weights * v
        attended = self.out_proj(attended)

        gate_input = torch.cat([query, attended], dim=-1)
        gate = self.gate(gate_input)
        residual_gate = self.residual_gate(gate_input)

        enriched = query + 0.7 * gate * attended + 0.3 * residual_gate * query
        enriched = self.norm(enriched)

        return enriched


class ContextEnhancedLSTMModel(nn.Module):
    """Architecture with long-term context memory.

    Pipeline:
    1. Conv1D for local patterns
    2. Self-Attention for intra-sequence context
    3. LSTM for temporal modelling
    4. Context Memory for history
    5. Context Encoder → fixed vector
    6. Context-Aware Attention to enrich LSTM output
    7. Dual Head for LONG/SHORT logits
    """

    def __init__(
        self,
        input_dim: int,
        hidden_dim: int = 48,
        num_layers: int = 2,
        dropout: float = 0.35,
        use_conv: bool = True,
        use_attention: bool = True,
        use_context: bool = True,
        context_capacity: int = 64,
        context_hidden_dim: int = 64,
    ):
        super().__init__()
        self.use_conv = use_conv
        self.use_attention = use_attention
        self.use_context = use_context
        self.hidden_dim = hidden_dim

        self.input_proj = nn.Linear(input_dim, hidden_dim)

        if use_conv:
            from src.models.lstm import ConvBlock
            self.conv_block = ConvBlock(hidden_dim, kernel_size=cfg.DEFAULT_KERNEL_SIZE)

        if use_attention:
            from src.models.lstm import SelfAttention
            self.attention = SelfAttention(hidden_dim, num_heads=cfg.DEFAULT_NUM_HEADS)

        self.lstm = nn.LSTM(
            hidden_dim, hidden_dim, num_layers,
            batch_first=True,
            dropout=dropout if num_layers > 1 else 0.0,
        )
        self.lstm_dropout = nn.Dropout(dropout * 0.5)
        self.bn = nn.BatchNorm1d(hidden_dim)

        self._context_capacity: Optional[int] = None
        self._context_hidden_dim: Optional[int] = None
        self._context_input_dim: Optional[int] = None
        self.context_encoder: Optional[ContextEncoder] = None
        self.context_attention: Optional[ContextAwareAttention] = None
        self._context_memory: Optional[ContextMemory] = None
        self._context_enabled: bool = True

        if use_context:
            self._context_capacity = context_capacity
            self._context_hidden_dim = context_hidden_dim
            self._context_input_dim = hidden_dim * 2 + 2 + input_dim
            self.context_encoder = ContextEncoder(
                input_dim=self._context_input_dim,
                hidden_dim=context_hidden_dim,
                num_heads=4,
                num_layers=2,
                dropout=dropout * 0.5,
                max_seq_len=context_capacity,
            )
            self.context_attention = ContextAwareAttention(
                hidden_dim=hidden_dim,
                context_dim=context_hidden_dim,
                num_heads=4,
            )

        self.fc_long = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim // 2),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim // 2, 1),
        )
        self.fc_short = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim // 2),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim // 2, 1),
        )

    def _get_context_memory(self) -> ContextMemory:
        if self._context_memory is None:
            device_str = str(next(self.parameters()).device)
            self._context_memory = ContextMemory(
                capacity=self._context_capacity or 64,
                hidden_dim=self.hidden_dim,
                context_feature_dim=self._context_hidden_dim or 64,
                device=device_str,
            )
        return self._context_memory

    def disable_context(self) -> None:
        self._context_enabled = False

    def enable_context(self) -> None:
        self._context_enabled = True

    def forward(
        self,
        x: torch.Tensor,
        features_raw: Optional[torch.Tensor] = None,
        store_context: bool = True,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        batch = x.shape[0]
        x = self.input_proj(x)

        if self.use_conv:
            x = self.conv_block(x)

        if self.use_attention:
            x = self.attention(x)

        out, (h_n, c_n) = self.lstm(x)
        out = self.lstm_dropout(out)

        if self.use_context and self._context_enabled:
            _mem = self._get_context_memory()
            min_ctx = min(8, (_mem.capacity // 4))

            if features_raw is not None and len(_mem) >= min_ctx:
                ctx_seq = _mem.get_context_sequence(device=x.device)
                if ctx_seq is not None:
                    ctx_seq = ctx_seq.unsqueeze(0).expand(batch, -1, -1)
                    ctx_vec = self.context_encoder(ctx_seq)
                    enriched = self.context_attention(out, ctx_vec)
                    features = enriched[:, -1, :]
                    if features.size(0) > 1:
                        features = self.bn(features)
                else:
                    features = out[:, -1, :]
                    if features.size(0) > 1:
                        features = self.bn(features)
            else:
                features = out[:, -1, :]
                if features.size(0) > 1:
                    features = self.bn(features)

            logits_long = self.fc_long(features).squeeze(-1)
            logits_short = self.fc_short(features).squeeze(-1)

            if store_context:
                f_raw = features_raw if features_raw is not None else features.detach()
                _mem.push(
                    h_n[-1], c_n[-1],
                    logits_long.detach(), logits_short.detach(), f_raw,
                )
        else:
            features = self.bn(out[:, -1, :])
            logits_long = self.fc_long(features).squeeze(-1)
            logits_short = self.fc_short(features).squeeze(-1)

        return logits_long, logits_short

    def get_context_info(self) -> Dict[str, Any]:
        if not self.use_context:
            return {"enabled": False}
        _mem = self._get_context_memory()
        return {
            "enabled": self._context_enabled,
            "memory_size": len(_mem),
            "capacity": _mem.capacity,
            "timestep": _mem.timestep,
        }

    def clear_context(self) -> None:
        if self.use_context and self._context_memory is not None:
            self._context_memory.clear()