"""Training pipeline for Conditional Flow Matching mouse trajectory model.

Pipeline:
1. Load traces from JSONL, convert to rotated coordinate frame
2. Apply 6× data augmentation
3. Train TrajectoryFlowModel with OT-Conditional Flow Matching:
   - x1 = real data, x0 = randn_like(x1), t = rand(B)
   - x_t = (1-t)*x0 + t*x1
   - v_target = x1 - x0
   - v_pred = model(x_t, t, cond)
   - loss = MSE(v_pred, v_target)
4. Save: flow_model.pt, click_dist.json, duration_dist.json, train_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
from torch.utils.data import DataLoader

from ai_mouse._coord import encode_trajectory
from tools.config import TrainConfig
from tools.models import TrajectoryFlowModel
from tools.utils import resample_arc

logger = logging.getLogger(__name__)

# Distance bins for duration distribution (in pixels)
_DIST_BINS: list[float] = [0, 50, 100, 200, 400, 600, 800, 1200, float("inf")]


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


def load_and_prepare_data(
    data_path: Path,
    seq_len: int = 64,
) -> tuple[np.ndarray, np.ndarray, list[float]]:
    """Load JSONL traces and convert to rotated-frame tensors.

    Args:
        data_path: path to traces.jsonl
        seq_len: number of time steps to resample each trajectory to

    Returns:
        seq:        (N, seq_len, 3) float32 — (forward, lateral, log_Δt)
        cond:       (N, 3) float32        — [dist_norm, log_dist, log_dur]
        click_durs: list of float click durations in ms
    """
    data_path = Path(data_path)
    seq_list: list[np.ndarray] = []
    cond_list: list[np.ndarray] = []
    click_durs: list[float] = []

    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:
            continue

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

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

        sx, sy = meta["start"]
        ex, ey = meta["end"]

        # Extract move events
        moves = [(e["x"], e["y"], e["t"]) for e in events if e["type"] == "move"]
        if len(moves) < 2:
            continue

        xs = np.array([m[0] for m in moves], dtype=float)
        ys = np.array([m[1] for m in moves], dtype=float)
        ts = np.array([m[2] for m in moves], dtype=float)

        xy_raw = np.stack([xs, ys], axis=1)

        # Reject degenerate (zero-length) trajectories
        total_arc = float(np.linalg.norm(np.diff(xy_raw, axis=0), axis=1).sum())
        if total_arc < 1.0:
            continue

        # Resample spatial positions to seq_len via arc-length
        xy_resampled = resample_arc(xy_raw, seq_len)  # (seq_len, 2)

        # Resample timestamps along the same arc-length grid
        arc_dist = np.concatenate(
            [[0.0], np.cumsum(np.linalg.norm(np.diff(xy_raw, axis=0), axis=1))]
        )
        s_uniform = np.linspace(0.0, arc_dist[-1], seq_len)
        ts_resampled = np.interp(s_uniform, arc_dist, ts)

        # Convert spatial coords to rotated frame (forward, lateral)
        fl = encode_trajectory(xy_resampled, (sx, sy), (ex, ey))  # (seq_len, 2)

        # Compute Δt intervals (length seq_len-1) → log(Δt+1), pad 0 at front
        dt_raw = np.diff(ts_resampled).clip(0.0)
        log_dt = np.log(dt_raw + 1.0)  # (seq_len-1,)
        log_dt_padded = np.concatenate([[0.0], log_dt])  # (seq_len,) — first step has no interval

        # Stack into (seq_len, 3)
        seq_arr = np.stack([fl[:, 0], fl[:, 1], log_dt_padded], axis=1).astype(np.float32)

        # Condition vector
        dist = float(meta["dist"]) if meta["dist"] > 0 else float(
            math.hypot(ex - sx, ey - sy)
        )
        dist = max(dist, 1.0)
        total_dur = float(ts_resampled[-1] - ts_resampled[0])
        total_dur = max(total_dur, 1.0)

        cond_arr = np.array(
            [
                dist / 2000.0,                   # dist_norm
                math.log(dist / 100.0),           # log_dist
                math.log(total_dur / 500.0),      # log_dur
            ],
            dtype=np.float32,
        )

        seq_list.append(seq_arr)
        cond_list.append(cond_arr)

        # Click duration (down→up)
        downs = [e for e in events if e["type"] == "down"]
        ups = [e for e in events if e["type"] == "up"]
        if downs and ups:
            click_durs.append(float(ups[-1]["t"] - downs[-1]["t"]))

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

    return (
        np.stack(seq_list, axis=0),   # (N, seq_len, 3)
        np.stack(cond_list, axis=0),  # (N, 3)
        click_durs,
    )


# ---------------------------------------------------------------------------
# Data augmentation  (6×)
# ---------------------------------------------------------------------------


def _augment(
    seq: np.ndarray,
    cond: np.ndarray,
) -> tuple[np.ndarray, np.ndarray]:
    """6× augmentation operating in the rotated (forward, lateral, log_dt) frame.

    Variants:
      0 — original
      1 — lateral flip:       lateral → −lateral
      2 — speed ×0.8:         log_Δt[1:] += log(1.25)
      3 — speed ×1.2:         log_Δt[1:] += log(1/1.2)
      4 — temporal noise:     log_Δt[1:] += N(0, 0.05)
      5 — combined:           lateral flip + speed ×0.9

    Args:
        seq:  (N, T, 3) — (forward, lateral, log_dt)
        cond: (N, 3)    — [dist_norm, log_dist, log_dur]

    Returns:
        seq_aug:  (6N, T, 3)
        cond_aug: (6N, 3)
    """
    log_1_25 = math.log(1.25)
    log_inv_1_2 = math.log(1.0 / 1.2)
    log_1_1 = math.log(1.0 / 0.9)

    seqs = [seq]
    conds = [cond]

    # 1. Lateral flip
    s1 = seq.copy()
    s1[:, :, 1] = -s1[:, :, 1]
    seqs.append(s1)
    conds.append(cond.copy())

    # 2. Speed ×0.8 (longer duration: log_dt += log(1.25))
    s2 = seq.copy()
    s2[:, 1:, 2] += log_1_25
    c2 = cond.copy()
    c2[:, 2] += log_1_25   # log_dur updated
    seqs.append(s2)
    conds.append(c2)

    # 3. Speed ×1.2 (shorter duration: log_dt += log(1/1.2))
    s3 = seq.copy()
    s3[:, 1:, 2] += log_inv_1_2
    c3 = cond.copy()
    c3[:, 2] += log_inv_1_2
    seqs.append(s3)
    conds.append(c3)

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

    # 5. Lateral flip + speed ×0.9 (i.e. log_dt += log(1/0.9))
    s5 = seq.copy()
    s5[:, :, 1] = -s5[:, :, 1]
    s5[:, 1:, 2] += log_1_1
    c5 = cond.copy()
    c5[:, 2] += log_1_1
    seqs.append(s5)
    conds.append(c5)

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


class TrajectoryDataset(torch.utils.data.Dataset):
    """Trajectory dataset with on-the-fly 6× augmentation.

    Replaces the old eager `_augment(seq, cond)` which expanded the dataset
    6× in memory before training. With this class, the original (N, T, 3)
    arrays stay as-is and each `__getitem__` call computes one of the 6
    augmentation variants on demand.

    Augmentation variants (matching legacy `_augment` semantics):
      0 — original
      1 — lateral flip (lateral → −lateral)
      2 — speed ×0.8  (log_dt[1:] += log(1.25),  cond[2] += log(1.25))
      3 — speed ×1.2  (log_dt[1:] += log(1/1.2), cond[2] += log(1/1.2))
      4 — temporal noise (log_dt[1:] += N(0, 0.05))
      5 — flip + speed ×0.9 (lateral flip, log_dt[1:] += log(1/0.9), cond[2] += log(1/0.9))
    """

    _LOG_1_25 = math.log(1.25)
    _LOG_INV_1_2 = math.log(1.0 / 1.2)
    _LOG_INV_0_9 = math.log(1.0 / 0.9)

    def __init__(self, seq: np.ndarray, cond: np.ndarray, augment: bool = True):
        self.seq = seq
        self.cond = cond
        self.augment = augment
        self._n_aug = 6 if augment else 1

    def __len__(self) -> int:
        return len(self.seq) * self._n_aug

    def __getitem__(self, idx: int) -> tuple[torch.Tensor, torch.Tensor]:
        base = idx // self._n_aug
        aug_id = idx % self._n_aug

        s = self.seq[base].copy()
        c = self.cond[base].copy()

        if aug_id == 1:
            s[:, 1] = -s[:, 1]
        elif aug_id == 2:
            s[1:, 2] += self._LOG_1_25
            c[2] += self._LOG_1_25
        elif aug_id == 3:
            s[1:, 2] += self._LOG_INV_1_2
            c[2] += self._LOG_INV_1_2
        elif aug_id == 4:
            noise = np.random.normal(0.0, 0.05, size=s[1:, 2].shape).astype(np.float32)
            s[1:, 2] += noise
        elif aug_id == 5:
            s[:, 1] = -s[:, 1]
            s[1:, 2] += self._LOG_INV_0_9
            c[2] += self._LOG_INV_0_9

        return torch.from_numpy(s), torch.from_numpy(c)


# ---------------------------------------------------------------------------
# Duration distribution (per distance bin)
# ---------------------------------------------------------------------------


def _compute_duration_dist(data_path: Path) -> dict:
    """Compute per-distance-bin log-normal parameters for trace duration.

    Args:
        data_path: path to traces.jsonl

    Returns:
        dict with keys "bins" (list of floats) and "params" (list of dicts
        with "mu_log" and "sigma_log" per bin).
    """
    bin_durations: list[list[float]] = [[] for _ in range(len(_DIST_BINS) - 1)]

    for raw_line in Path(data_path).read_text(encoding="utf-8").splitlines():
        line = raw_line.strip()
        if not line:
            continue
        try:
            trace = json.loads(line)
        except json.JSONDecodeError:
            continue

        meta = trace["meta"]
        events = trace["events"]
        dist = float(meta.get("dist", 0))

        moves = [e for e in events if e["type"] == "move"]
        if len(moves) < 2:
            continue
        dur = float(moves[-1]["t"] - moves[0]["t"])
        if dur <= 0:
            continue

        # Find bin
        for i in range(len(_DIST_BINS) - 1):
            if _DIST_BINS[i] <= dist < _DIST_BINS[i + 1]:
                bin_durations[i].append(dur)
                break

    params = []
    for durs in bin_durations:
        if len(durs) >= 2:
            log_durs = np.log(np.array(durs, dtype=float))
            mu_log = float(np.mean(log_durs))
            sigma_log = float(np.std(log_durs, ddof=1))
        else:
            mu_log = float(np.log(500.0))
            sigma_log = 0.5
        params.append({"mu_log": mu_log, "sigma_log": max(sigma_log, 0.05)})

    return {"bins": _DIST_BINS, "params": params}


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


def train(
    data_path: Path,
    output_dir: Path,
    epochs: int = 300,
    batch_size: int = 64,
    lr: float = 3e-4,
    seq_len: int = 64,
    progress_callback: Callable[[dict], None] | None = None,
    config: TrainConfig | None = None,
    resume_from: Path | None = None,
) -> None:
    """Train TrajectoryFlowModel with OT-Conditional Flow Matching.

    Args:
        data_path:          path to traces.jsonl
        output_dir:         directory where artefacts are written
        epochs:             training epochs
        batch_size:         mini-batch size
        lr:                 AdamW learning rate
        seq_len:            number of time steps per trajectory
        progress_callback:  optional callable invoked each epoch with
                            {"epoch": n, "total": N, "loss": f}.
                            Called once at the end with {"done": True, "mu", "sigma"}.
        config:             optional TrainConfig for model hyperparameters
        resume_from:        if given, load model weights from this checkpoint
                            directory (must contain flow_model.pt). Used for
                            two-stage training (pretrain → fine-tune).
    """
    data_path = Path(data_path)
    output_dir = Path(output_dir)

    if config is None:
        config = TrainConfig(
            epochs=epochs, batch_size=batch_size, lr=lr, seq_len=seq_len
        )

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

    # ---- Load & prepare ----
    logger.info("Loading data from %s", data_path)
    seq_np, cond_np, click_durs = load_and_prepare_data(data_path, seq_len=seq_len)
    logger.info("Loaded %d traces", len(seq_np))

    # ---- Build streaming dataset (on-the-fly 6× augmentation) ----
    if config.augment:
        logger.info("Using on-the-fly 6× augmentation, base samples: %d", len(seq_np))
    ds = TrajectoryDataset(seq_np, cond_np, augment=config.augment)
    logger.info("Effective dataset size: %d", len(ds))

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

    # ---- Model & optimiser ----
    model = TrajectoryFlowModel(
        seq_len=seq_len,
        d_model=config.d_model,
        nhead=config.nhead,
        num_layers=config.num_layers,
        dim_feedforward=config.dim_feedforward,
        dropout=config.dropout,
        cond_dim=3,
    )

    # ---- Resume from checkpoint if requested ----
    if resume_from is not None:
        resume_path = Path(resume_from) / "flow_model.pt"
        if not resume_path.exists():
            raise FileNotFoundError(
                f"resume_from checkpoint not found: {resume_path}"
            )
        logger.info("Resuming from checkpoint: %s", resume_path)
        state_dict = torch.load(resume_path, map_location="cpu", weights_only=True)
        model.load_state_dict(state_dict)

    optimiser = torch.optim.AdamW(model.parameters(), lr=lr)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimiser, T_max=epochs)

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

    # ---- Training loop: OT-Conditional Flow Matching ----
    model.train()
    for epoch in range(epochs):
        epoch_loss = 0.0
        n_batches = 0

        for x1_batch, cond_batch in loader:
            B = x1_batch.shape[0]

            # Sample noise x0 ~ N(0, I)
            x0 = torch.randn_like(x1_batch)

            # Sample random timestep t ~ U[0, 1]
            t = torch.rand(B)

            # Interpolate: x_t = (1-t)*x0 + t*x1
            t_expand = t[:, None, None]  # (B, 1, 1) for broadcasting
            x_t = (1.0 - t_expand) * x0 + t_expand * x1_batch

            # Target velocity: v = x1 - x0
            v_target = x1_batch - x0

            # Predict velocity
            v_pred = model(x_t, t, cond_batch)

            # MSE loss
            loss = torch.nn.functional.mse_loss(v_pred, v_target)

            optimiser.zero_grad()
            loss.backward()
            torch.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)
        logger.debug("Epoch %d/%d loss=%.6f", epoch + 1, epochs, epoch_loss)

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

    # ---- Save model ----
    torch.save(model.state_dict(), output_dir / "flow_model.pt")
    logger.info("Saved flow_model.pt to %s", output_dir)

    # ---- Click duration distribution ----
    if click_durs:
        arr = np.array(click_durs)
        arr = arr[(arr >= 20) & (arr <= 500)]
        if len(arr) >= 2:
            mu_c = float(arr.mean())
            sigma_c = max(float(arr.std()), 1.0)
        else:
            mu_c, sigma_c = 80.0, 30.0
    else:
        mu_c, sigma_c = 80.0, 30.0

    click_dist = {"mu": mu_c, "sigma": sigma_c, "low": 20.0, "high": 500.0}
    (output_dir / "click_dist.json").write_text(json.dumps(click_dist, indent=2))

    # ---- Duration distribution (per distance bin) ----
    dur_dist = _compute_duration_dist(data_path)
    (output_dir / "duration_dist.json").write_text(json.dumps(dur_dist, indent=2))

    # ---- Train config ----
    train_cfg = {
        "seq_len": seq_len,
        "epochs": epochs,
        "batch_size": batch_size,
        "lr": lr,
        "d_model": config.d_model,
        "nhead": config.nhead,
        "num_layers": config.num_layers,
        "dim_feedforward": config.dim_feedforward,
        "dropout": config.dropout,
        "cond_dim": 3,
    }
    (output_dir / "train_config.json").write_text(json.dumps(train_cfg, indent=2))

    if progress_callback is not None:
        progress_callback({"done": True, "mu": mu_c, "sigma": sigma_c})
