"""Scroll training pipeline for ScrollCVAE.

Pipeline:
1. Load scroll traces from JSONL -> (seq, cond) tensors
2. Apply 4x data augmentation
3. Train ScrollCVAE with: MSE(delta_norm) + 1.5*MSE(log_dt) + beta*KL
4. Save: scroll_model.pt, scroll_config.json
"""
from __future__ import annotations

import json
import logging
import math
from collections.abc import Callable
from pathlib import Path

import numpy as np
import torch
import torch.nn as nn
from torch.distributions import Normal, kl_divergence
from torch.utils.data import DataLoader, TensorDataset

from tools.config import ScrollTrainConfig
from tools.scroll.models import ScrollCVAE

logger = logging.getLogger(__name__)

# Mode -> one-hot index
_MODE_INDEX = {"target": 0, "fast": 1, "precise": 2}


# ---------------------------------------------------------------------------
# Data loading
# ---------------------------------------------------------------------------


def load_scroll_data(
    data_path: Path,
    seq_len: int = 32,
) -> tuple[np.ndarray, np.ndarray]:
    """Load scroll JSONL and return (seq, cond) arrays.

    Args:
        data_path: path to scroll_traces.jsonl
        seq_len:   number of wheel-event steps to pad/truncate to

    Returns:
        seq:  (N, seq_len, 2) float32 -- (delta_norm, log_dt)
        cond: (N, 7) float32 -- [dist/5000, log(dist/500), direction, viewport_norm, mode_onehot*3]
    """
    data_path = Path(data_path)
    seq_list: list[np.ndarray] = []
    cond_list: list[np.ndarray] = []

    for i, raw_line in enumerate(data_path.read_text(encoding="utf-8").splitlines(), 1):
        line = raw_line.strip()
        if not line:
            continue
        try:
            trace = json.loads(line)
        except json.JSONDecodeError:
            logger.warning("Skipping line %d: invalid JSON", i)
            continue

        if "meta" not in trace or "events" not in trace:
            logger.warning("Skipping line %d: missing meta or events", i)
            continue

        meta = trace["meta"]
        events = trace["events"]

        if not events:
            continue

        distance = float(meta.get("distance", 0))
        if distance <= 0:
            start_y = float(meta.get("start_scrollY", 0))
            target_y = float(meta.get("target_scrollY", 0))
            distance = abs(target_y - start_y)
        if distance <= 0:
            distance = 1.0

        direction_str = meta.get("direction", "down")
        direction = 1.0 if direction_str == "down" else -1.0

        viewport_height = float(meta.get("viewport_height", 900))
        viewport_norm = viewport_height / 1000.0

        mode_str = meta.get("mode", "target")
        mode_idx = _MODE_INDEX.get(mode_str, 0)
        mode_onehot = np.zeros(3, dtype=np.float32)
        mode_onehot[mode_idx] = 1.0

        deltas = np.array([float(e.get("deltaY", 0)) for e in events], dtype=np.float32)
        times = np.array([float(e.get("t", 0)) for e in events], dtype=np.float32)

        if len(deltas) == 0:
            continue

        delta_norm = deltas / distance

        dt_raw = np.diff(times).clip(0.0)
        log_dt = np.log(dt_raw + 1.0)
        log_dt_padded = np.concatenate([[0.0], log_dt])

        seq_raw = np.stack([delta_norm, log_dt_padded], axis=1).astype(np.float32)

        n = len(seq_raw)
        if n >= seq_len:
            seq_out = seq_raw[:seq_len]
        else:
            pad = np.zeros((seq_len - n, 2), dtype=np.float32)
            seq_out = np.concatenate([seq_raw, pad], axis=0)

        dist_norm = distance / 5000.0
        log_dist = math.log(max(distance, 1.0) / 500.0)
        cond_arr = np.array(
            [dist_norm, log_dist, direction, viewport_norm, *mode_onehot],
            dtype=np.float32,
        )

        seq_list.append(seq_out)
        cond_list.append(cond_arr)

    if not seq_list:
        raise ValueError(f"No valid scroll traces found in {data_path}")

    logger.info("Loaded %d scroll traces from %s", len(seq_list), data_path)
    return (
        np.stack(seq_list, axis=0),
        np.stack(cond_list, axis=0),
    )


# ---------------------------------------------------------------------------
# Data augmentation  (4x)
# ---------------------------------------------------------------------------


def _augment_scroll(
    seq: np.ndarray,
    cond: np.ndarray,
) -> tuple[np.ndarray, np.ndarray]:
    """4x augmentation for scroll sequences.

    Variants:
      0 -- original
      1 -- speed x0.8:       log_dt[1:] += log(1.25)
      2 -- speed x1.2:       log_dt[1:] += log(1/1.2)
      3 -- temporal noise:   log_dt[1:] += N(0, 0.05)
    """
    log_1_25 = math.log(1.25)
    log_inv_1_2 = math.log(1.0 / 1.2)

    seqs = [seq]
    conds = [cond]

    s1 = seq.copy()
    s1[:, 1:, 1] += log_1_25
    seqs.append(s1)
    conds.append(cond.copy())

    s2 = seq.copy()
    s2[:, 1:, 1] += log_inv_1_2
    seqs.append(s2)
    conds.append(cond.copy())

    s3 = seq.copy()
    noise = np.random.normal(0.0, 0.05, size=s3[:, 1:, 1].shape).astype(np.float32)
    s3[:, 1:, 1] += noise
    seqs.append(s3)
    conds.append(cond.copy())

    return np.concatenate(seqs, axis=0), np.concatenate(conds, axis=0)


# ---------------------------------------------------------------------------
# Main training function
# ---------------------------------------------------------------------------


def train_scroll(
    data_path: Path,
    output_dir: Path,
    epochs: int = 100,
    batch_size: int = 32,
    lr: float = 5e-4,
    seq_len: int = 32,
    progress_callback: Callable[[dict], None] | None = None,
    config: ScrollTrainConfig | None = None,
) -> None:
    """Train ScrollCVAE and save artefacts to output_dir."""
    cfg = config or ScrollTrainConfig()
    epochs = epochs or cfg.epochs
    batch_size = batch_size or cfg.batch_size
    lr = lr or cfg.lr
    seq_len = seq_len or cfg.seq_len

    data_path = Path(data_path)
    output_dir = Path(output_dir)

    if not data_path.exists():
        raise FileNotFoundError(f"Data file not found: {data_path}")

    logger.info("Starting scroll training: epochs=%d, data=%s", epochs, data_path)

    seq_np, cond_np = load_scroll_data(data_path, seq_len=seq_len)
    seq_np, cond_np = _augment_scroll(seq_np, cond_np)

    seq_t = torch.from_numpy(seq_np)
    cond_t = torch.from_numpy(cond_np)

    output_dir.mkdir(parents=True, exist_ok=True)

    model = ScrollCVAE(seq_len=seq_len, latent_dim=16, hidden=64, cond_dim=7)
    optimiser = torch.optim.AdamW(model.parameters(), lr=lr)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimiser, T_max=max(epochs, 1))

    ds = TensorDataset(seq_t, cond_t)
    loader = DataLoader(ds, batch_size=batch_size, shuffle=True, drop_last=False)

    model.train()
    for epoch in range(epochs):
        beta = min(cfg.beta_max, cfg.beta_max * epoch / max(cfg.beta_warmup_epochs, 1))

        epoch_loss = 0.0
        n_batches = 0

        for seq_b, cond_b in loader:
            optimiser.zero_grad()

            recon, mu, logvar = model(seq_b, cond_b)

            delta_loss = nn.functional.mse_loss(recon[:, :, 0], seq_b[:, :, 0])
            logdt_loss = nn.functional.mse_loss(recon[:, :, 1], seq_b[:, :, 1])
            recon_loss = cfg.weight_delta * delta_loss + cfg.weight_log_dt * logdt_loss

            std = torch.exp(0.5 * logvar)
            q = Normal(mu, std)
            p = Normal(torch.zeros_like(mu), torch.ones_like(std))
            kl_loss = kl_divergence(q, p).mean()

            loss = recon_loss + beta * kl_loss
            loss.backward()

            nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
            optimiser.step()

            epoch_loss += loss.item()
            n_batches += 1

        scheduler.step()
        epoch_loss /= max(n_batches, 1)

        if progress_callback is not None:
            progress_callback({
                "epoch": epoch + 1,
                "total": epochs,
                "loss": epoch_loss,
                "stage": "scroll",
            })

    torch.save(model.state_dict(), output_dir / "scroll_model.pt")

    scroll_cfg = {
        "seq_len": seq_len,
        "latent_dim": model.latent_dim,
        "hidden": model.hidden,
        "cond_dim": model.cond_dim,
        "epochs": epochs,
        "batch_size": batch_size,
        "lr": lr,
        "beta_max": cfg.beta_max,
        "beta_anneal_epochs": cfg.beta_warmup_epochs,
        "w_delta": cfg.weight_delta,
        "w_logdt": cfg.weight_log_dt,
    }
    (output_dir / "scroll_config.json").write_text(json.dumps(scroll_cfg, indent=2))

    logger.info("Scroll training complete. Model saved to %s", output_dir)

    if progress_callback is not None:
        progress_callback({"done": True})
