"""Seq2Seq LSTM forecaster for NDVI time-series."""

from __future__ import annotations

import random

import torch
import torch.nn as nn

from ndvi_forecast._utils import gather_last, last_time_and_step_raw, maybe_jitter_step


class Seq2SeqNDVI(nn.Module):
    """Seq2Seq LSTM with residual prediction, season descriptor, and stitch-point continuity.

    Encoder: 2-layer LSTM on [NDVI, RAW_DAP].
    Season head: maps encoder hidden state to a latent season descriptor.
    Decoder: multi-layer LSTM on [last_ndvi, group_stats, dap_t, season_desc].
    Output: delta prediction via tanh-bounded MLP + optional affine correction.
    Stitch: hard C0 (first pred = last obs) and soft C1 (slope blending) at the transition.
    """

    def __init__(
        self,
        input_dim: int = 2,
        hidden_dim: int = 64,
        fallback_dap_step: float = 1.0,
        min_step: float = 0.5,
        max_step: float = 5.0,
        step_jitter_std: float = 0.05,
        use_output_affine: bool = True,
        clamp01: bool = True,
        max_delta_mag: float = 0.25,
        decoder_layers: int = 2,
        decoder_dropout: float = 0.1,
        season_dim: int = 32,
        default_target_len: int = 30,
        stitch_slope_blend_init: float = 0.0,
        stitch_slope_blend_scale: float = 0.75,
    ):
        super().__init__()
        self.fallback_dap_step = float(fallback_dap_step)
        self.min_step = float(min_step)
        self.max_step = float(max_step)
        self.step_jitter_std = float(step_jitter_std)
        self.use_output_affine = bool(use_output_affine)
        self.clamp01 = bool(clamp01)
        self.max_delta_mag = float(max_delta_mag)
        self.default_target_len = int(default_target_len)
        self.season_dim = int(season_dim)
        self.stitch_slope_blend_scale = float(stitch_slope_blend_scale)

        # Encoder: [ndvi, raw_dap] -> hidden
        self.encoder = nn.LSTM(input_dim, hidden_dim, num_layers=2, batch_first=True)

        # Season descriptor from top encoder hidden state
        self.season_head = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, self.season_dim),
            nn.Tanh(),
        )

        # Decoder input: [last_ndvi, group_mean, group_std, dap_t, season_desc]
        dec_input_dim = 4 + self.season_dim
        self.decoder = nn.LSTM(
            input_size=dec_input_dim,
            hidden_size=hidden_dim,
            num_layers=decoder_layers,
            batch_first=True,
            dropout=decoder_dropout if decoder_layers > 1 else 0.0,
        )

        # Output head context: [decoder_out, group_stats, dap_t, last_ndvi, season_desc]
        head_in_dim = hidden_dim + 4 + self.season_dim
        mid = max(hidden_dim // 2, 32)
        self.delta_head = nn.Sequential(
            nn.Linear(head_in_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, mid),
            nn.ReLU(),
            nn.Linear(mid, 1),
        )

        # Stitch gate: blends predicted delta toward last observed slope at t=1
        self.stitch_gate = nn.Linear(head_in_dim, 1)
        nn.init.zeros_(self.stitch_gate.weight)
        nn.init.constant_(self.stitch_gate.bias, float(stitch_slope_blend_init))

        if self.use_output_affine:
            self.affine_head_gain = nn.Linear(head_in_dim, 1)
            self.affine_head_bias = nn.Linear(head_in_dim, 1)
            nn.init.zeros_(self.affine_head_gain.weight)
            nn.init.zeros_(self.affine_head_gain.bias)
            nn.init.zeros_(self.affine_head_bias.weight)
            nn.init.zeros_(self.affine_head_bias.bias)

    # --- helpers ---

    def _last_ndvi(self, input_seq, input_lengths):
        return gather_last(input_seq, input_lengths, feat_idx=0).unsqueeze(1)

    def _last_slope(self, input_seq, input_lengths):
        """Last observed NDVI slope (y_last - y_prev). Returns (B, 1), 0 if length < 2."""
        device = input_seq.device
        B = input_seq.size(0)
        idx_last = (input_lengths.to(device) - 1).clamp(min=0)
        idx_prev = (idx_last - 1).clamp(min=0)
        bidx = torch.arange(B, device=device)
        y_last = input_seq[bidx, idx_last, 0]
        y_prev = input_seq[bidx, idx_prev, 0]
        has_prev = (input_lengths.to(device) > 1).float()
        return ((y_last - y_prev) * has_prev).unsqueeze(1)

    def _decode_step(self, last_val, group_stats, time_t, season_desc, h, c,
                     *, stitch=False, last_slope=None):
        """Single autoregressive decode step. Returns (pred_ndvi, h, c)."""
        dec_in = torch.cat([last_val, group_stats, time_t, season_desc], dim=1).unsqueeze(1)
        dec_out, (h, c) = self.decoder(dec_in, (h, c))

        out_ctx = torch.cat([dec_out.squeeze(1), group_stats, time_t, last_val, season_desc], dim=1)

        delta_raw = self.delta_head(out_ctx)
        delta = torch.tanh(delta_raw) * self.max_delta_mag

        if stitch:
            g = torch.sigmoid(self.stitch_gate(out_ctx)) * self.stitch_slope_blend_scale
            if last_slope is None:
                last_slope = torch.zeros_like(delta)
            delta = (1.0 - g) * delta + g * last_slope

        pred = last_val + delta

        if self.use_output_affine:
            gain = 1.0 + torch.tanh(self.affine_head_gain(out_ctx)) * 0.05
            bias = torch.tanh(self.affine_head_bias(out_ctx)) * 0.02
            pred = gain * pred + bias

        if self.clamp01:
            pred = pred.clamp(0.0, 1.0)

        return pred.squeeze(1), h, c

    def forward(
        self,
        input_seq: torch.Tensor,
        input_lengths: torch.Tensor,
        group_stats: torch.Tensor,
        target_seq: torch.Tensor | None = None,
        teacher_forcing_ratio: float = 0.5,
    ) -> torch.Tensor:
        """Forward pass. Returns (B, L_out) predicted NDVI."""
        packed = nn.utils.rnn.pack_padded_sequence(
            input_seq, input_lengths.cpu(), batch_first=True, enforce_sorted=False,
        )
        _, (h, c) = self.encoder(packed)

        # Season descriptor from top encoder layer
        season_desc = self.season_head(h[-1])

        lv = self._last_ndvi(input_seq, input_lengths)
        last_slope = self._last_slope(input_seq, input_lengths)
        t_last, step = last_time_and_step_raw(
            input_seq, input_lengths, self.fallback_dap_step, self.min_step, self.max_step,
        )
        step = maybe_jitter_step(step, self.training, self.step_jitter_std, self.min_step, self.max_step)

        target_len = target_seq.size(1) if target_seq is not None else self.default_target_len

        # Hard C0: first prediction equals last observed NDVI
        preds = [lv.squeeze(1)]

        if target_len <= 1:
            return torch.stack(preds, dim=1)

        cur_time = t_last
        for t in range(1, target_len):
            cur_time = cur_time + step
            pred_ndvi, h, c = self._decode_step(
                lv, group_stats, cur_time, season_desc, h, c,
                stitch=(t == 1), last_slope=last_slope,
            )
            preds.append(pred_ndvi)

            if target_seq is not None and random.random() < teacher_forcing_ratio:
                lv = target_seq[:, t].unsqueeze(1)
            else:
                lv = pred_ndvi.unsqueeze(1)

        return torch.stack(preds, dim=1)
