"""Pretraining loop for Seq2SeqNDVI with blended TF/AR loss."""

import numpy as np
import torch
from torch.utils.data import DataLoader
from sklearn.metrics import mean_absolute_error, r2_score
import matplotlib.pyplot as plt

from ndvi_forecast.model import Seq2SeqNDVI
from ndvi_forecast.datasets import BatchCollator
from ndvi_forecast.losses import (
    horizon_weights,
    masked_mse,
    masked_slope_mse,
    masked_curvature_mse,
    simple_ema_smooth_targets,
    cosine_decay,
)
from ndvi_forecast._utils import gather_last


class Trainer:
    """Trains Seq2SeqNDVI with blended TF/AR loss, slope/curvature regularization, and early stopping."""

    def __init__(
        self,
        train_dataset,
        val_dataset,
        batch_size: int = 32,
        hidden_dim: int = 64,
        lr: float = 1e-3,
        device: str = "cpu",
        best_model_path: str = "./best_model.pt",
        alpha_ar_init: float = 0.30,
        alpha_ar_final: float = 0.50,
        lambda_slope_init: float = 0.25,
        lambda_slope_final: float = 0.05,
        lambda_curv_init: float = 0.15,
        lambda_curv_final: float = 0.03,
        lambda_anchor: float = 5.0,
        horizon_gamma: float = 0.6,
        ema_alpha: float = 0.1,
        use_ema_targets: bool = False,
        fallback_dap_step: float = 1.0,
    ):
        self.device = torch.device(device)
        self.model = Seq2SeqNDVI(hidden_dim=hidden_dim, fallback_dap_step=fallback_dap_step).to(self.device)
        self.optimizer = torch.optim.Adam(self.model.parameters(), lr=lr)

        collator = BatchCollator()
        self.train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, collate_fn=collator)
        self.val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, collate_fn=collator)

        self.alpha_ar_init = float(alpha_ar_init)
        self.alpha_ar_final = float(alpha_ar_final)
        self.lambda_slope_init = float(lambda_slope_init)
        self.lambda_slope_final = float(lambda_slope_final)
        self.lambda_curv_init = float(lambda_curv_init)
        self.lambda_curv_final = float(lambda_curv_final)
        self.lambda_anchor = float(lambda_anchor)
        self.horizon_gamma = float(horizon_gamma)
        self.ema_alpha = float(ema_alpha)
        self.use_ema_targets = bool(use_ema_targets)

        self.train_losses, self.val_losses = [], []
        self.val_maes, self.val_r2s = [], []
        self.best_val_loss = float("inf")
        self.best_model_path = best_model_path

        self.alpha_ar = self.alpha_ar_init
        self.lambda_slope = self.lambda_slope_init
        self.lambda_curv = self.lambda_curv_init

    def train(self, num_epochs: int = 5, patience: int = 5):
        """Main training loop with cosine-decayed schedules and early stopping."""
        patience_counter = 0
        for epoch in range(num_epochs):
            p = 0.0 if num_epochs <= 1 else epoch / (num_epochs - 1)
            self.lambda_slope = cosine_decay(self.lambda_slope_init, self.lambda_slope_final, p)
            self.lambda_curv = cosine_decay(self.lambda_curv_init, self.lambda_curv_final, p)
            self.alpha_ar = cosine_decay(self.alpha_ar_init, self.alpha_ar_final, p)
            tf = 0.95 * (1.0 - p) + 0.05

            train_loss, _, _ = self._run_epoch(self.train_loader, training=True, tf_ratio=tf)
            val_loss, mae, r2 = self._run_epoch(self.val_loader, training=False, tf_ratio=0.0)

            self.train_losses.append(train_loss)
            self.val_losses.append(val_loss)
            self.val_maes.append(mae)
            self.val_r2s.append(r2)

            pct = int(100 * (epoch + 1) / num_epochs)
            bar_len = 20
            filled = int(bar_len * (epoch + 1) / num_epochs)
            bar = "\u2588" * filled + "\u2591" * (bar_len - filled)
            print(
                f"Epoch {epoch+1}/{num_epochs} [{bar}] {pct}% \u2014 "
                f"Train: {train_loss:.4f} Val: {val_loss:.4f} "
                f"(\u03b1_AR={self.alpha_ar:.2f} \u03bb_s={self.lambda_slope:.3f} \u03bb_c={self.lambda_curv:.3f})"
            )

            if val_loss < self.best_val_loss:
                self.best_val_loss = val_loss
                patience_counter = 0
                torch.save(self.model.state_dict(), self.best_model_path)
            else:
                patience_counter += 1
                if patience_counter >= patience:
                    print(f"Early stopping at epoch {epoch+1}")
                    break

        return self.plot_losses()

    def _run_epoch(self, dataloader, training: bool, tf_ratio: float):
        self.model.train() if training else self.model.eval()
        total_loss = 0.0
        preds_all, targets_all = [], []

        ctx = torch.enable_grad() if training else torch.no_grad()
        with ctx:
            for xb, yb, in_lens, tgt_lens, gs in dataloader:
                xb = xb.to(self.device)
                yb = yb.to(self.device)
                in_lens = in_lens.to(self.device)
                tgt_lens = tgt_lens.to(self.device)
                gs = gs.to(self.device)

                if training:
                    self.optimizer.zero_grad()

                y_train = simple_ema_smooth_targets(yb, tgt_lens, self.ema_alpha) if (training and self.use_ema_targets) else yb

                preds_tf = self.model(xb, in_lens, gs, target_seq=y_train, teacher_forcing_ratio=tf_ratio)
                if preds_tf.dim() == 3:
                    preds_tf = preds_tf.squeeze(-1)

                preds_ar = self.model(xb, in_lens, gs, target_seq=y_train, teacher_forcing_ratio=0.0)
                if preds_ar.dim() == 3:
                    preds_ar = preds_ar.squeeze(-1)

                T = yb.size(1)
                w_val = horizon_weights(T, self.horizon_gamma, device=yb.device)
                w_slope = horizon_weights(max(T - 1, 1), self.horizon_gamma, device=yb.device)
                w_curv = horizon_weights(max(T - 2, 1), self.horizon_gamma, device=yb.device)

                loss_tf = masked_mse(preds_tf, y_train, tgt_lens, w=w_val)
                loss_ar = masked_mse(preds_ar, y_train, tgt_lens, w=w_val)
                loss_slp = masked_slope_mse(preds_tf, y_train, tgt_lens, w=w_slope)
                loss_crv = masked_curvature_mse(preds_tf, y_train, tgt_lens, w=w_curv)

                loss = (1.0 - self.alpha_ar) * loss_tf + self.alpha_ar * loss_ar
                loss += self.lambda_slope * loss_slp + self.lambda_curv * loss_crv

                last_val = gather_last(xb, in_lens, feat_idx=0)
                has_step0 = (tgt_lens > 0).float()
                anchor_err = (preds_tf[:, 0] - last_val) ** 2
                loss += self.lambda_anchor * (anchor_err * has_step0).sum() / has_step0.sum().clamp_min(1.0)

                if training:
                    loss.backward()
                    torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
                    self.optimizer.step()
                else:
                    for pi, yi, tl in zip(preds_ar, yb, tgt_lens):
                        tl = int(tl.item())
                        preds_all.append(pi[:tl].detach().cpu().numpy())
                        targets_all.append(yi[:tl].detach().cpu().numpy())

                total_loss += float(loss.item())

        if not training:
            pf = np.concatenate(preds_all) if preds_all else np.array([])
            tf_ = np.concatenate(targets_all) if preds_all else np.array([])
            mae = mean_absolute_error(tf_, pf) if pf.size else float("nan")
            r2 = r2_score(tf_, pf) if pf.size else float("nan")
            print(f"    Val MAE: {mae:.4f} | R\u00b2: {r2:.4f}")
            return total_loss / max(1, len(dataloader)), mae, r2
        return total_loss / max(1, len(dataloader)), None, None

    def plot_losses(self) -> plt.Figure:
        """Plot training curves. Returns the figure."""
        fig, axs = plt.subplots(3, 1, figsize=(8, 10))
        axs[0].plot(self.train_losses, label="Train")
        axs[0].plot(self.val_losses, label="Val")
        axs[0].set_ylabel("Loss")
        axs[0].set_title("Training & Validation Loss")
        axs[0].legend()
        axs[1].plot(self.val_maes, color="orange")
        axs[1].set_ylabel("MAE")
        axs[1].set_title("Validation MAE")
        axs[2].plot(self.val_r2s, color="green")
        axs[2].set_ylabel("R\u00b2")
        axs[2].set_title("Validation R\u00b2")
        axs[2].set_xlabel("Epoch")
        plt.tight_layout()
        return fig

    def load_best_model(self):
        """Load the best checkpoint and return the model in eval mode."""
        self.model.load_state_dict(torch.load(self.best_model_path, map_location=self.device))
        self.model.eval()
        return self.model
