"""Inference layer: Flow Matching trajectory generation.

Pipeline:
1. Load model from model_dir (flow_model.pt, click_dist.json,
   duration_dist.json, train_config.json).
2. Compute condition vector: [dist/2000, log(dist/100), log(total_dur/500)].
3. Sample total_duration from duration_dist.json by distance bin (log-normal).
4. 10-step Euler ODE: start from noise, integrate velocity field to get trajectory.
5. Spatial post-processing:
   a. Endpoint snapping: force first=(0,0), last=(1,0), lerp last 6 points.
   b. Smooth start: dampen lateral near start (first 4 points).
   c. Enforce forward monotonicity (prevent x-axis jitter).
   d. 5-point gaussian smooth on lateral (preserve endpoints).
6. Temporal post-processing:
   a. Clip log_dt to [0, 5] to prevent exponential explosion.
   (speed profile and median±1.1 hard clip removed in 2026-05 refactor —
    let the model's learned timing distribution come through naturally.)
7. Decode to pixels via decode_trajectory.
8. Resample to n_points if n_points != model seq_len.
9. Convert log_dt → ms timestamps, scale to total_duration, clip [2, 150].
10. Ensure timestamps monotonically increasing.
11. Append click events sampled from truncated normal.
"""
from __future__ import annotations

import json
import logging
import math
from pathlib import Path

import numpy as np
import torch
from scipy.stats import truncnorm

from ai_mouse.coord import decode_trajectory
from tools.config import GenerateConfig
from tools.models import TrajectoryFlowModel
from tools.utils import resample_arc

logger = logging.getLogger(__name__)

_BUNDLED_MODELS_DIR = Path(__file__).parent.parent / "data" / "models_v2"


# ---------------------------------------------------------------------------
# Duration sampling helper
# ---------------------------------------------------------------------------


def _sample_duration(duration_dist: dict, dist: float) -> float:
    """Sample a total movement duration (ms) for the given pixel distance.

    Uses per-distance-bin log-normal parameters from duration_dist.

    Args:
        duration_dist: dict with "bins" and "params" keys.
        dist: pixel distance between start and end.

    Returns:
        Sampled duration in milliseconds.
    """
    bins = duration_dist["bins"]
    params = duration_dist["params"]
    # Find bin for this distance
    bin_idx = len(bins) - 1
    for i in range(len(bins) - 1):
        if dist < bins[i + 1]:
            bin_idx = i
            break
    # Clamp to valid params index
    bin_idx = min(bin_idx, len(params) - 1)
    mu_log = params[bin_idx]["mu_log"]
    sigma_log = params[bin_idx]["sigma_log"]
    return float(np.exp(np.random.normal(mu_log, sigma_log)))


# ---------------------------------------------------------------------------
# Smoothing helper
# ---------------------------------------------------------------------------


def _gaussian_smooth(x: np.ndarray, sigma: float = 1.0) -> np.ndarray:
    """5-point gaussian smoothing along a 1-D array, preserving endpoints.

    Args:
        x: 1-D input array.
        sigma: Gaussian std (px); larger = more smoothing. Default 1.0 gives
               weights ≈ [0.054, 0.244, 0.403, 0.244, 0.054].

    Returns:
        Smoothed array of the same shape. x[0] and x[-1] are unchanged.
        If len(x) < 5, returns x unchanged (kernel won't fit).
    """
    if len(x) < 5:
        return x.copy()
    kernel = np.exp(-0.5 * (np.arange(-2, 3) / sigma) ** 2)
    kernel /= kernel.sum()
    # Pad with edge values to avoid boundary artifacts, then slice back
    padded = np.pad(x, pad_width=2, mode="edge")
    smoothed = np.convolve(padded, kernel, mode="valid")
    smoothed[0] = x[0]
    smoothed[-1] = x[-1]
    return smoothed


# ---------------------------------------------------------------------------
# Main generate function
# ---------------------------------------------------------------------------


def generate(
    start: tuple[int, int],
    end: tuple[int, int],
    n_points: int = 64,
    speed: float | None = None,
    model_dir: str | None = None,
    config: GenerateConfig | None = None,
) -> list[tuple[int, int, int]]:
    """Generate a human-like mouse trajectory from start to end.

    Uses a Flow Matching model with 4-step Euler ODE integration.

    Args:
        start:      (x, y) starting pixel coordinate.
        end:        (x, y) target pixel coordinate.
        n_points:   number of movement points in the path (default 64).
        speed:      optional speed multiplier; speed=2 halves the duration.
        model_dir:  directory containing flow_model.pt, click_dist.json,
                    duration_dist.json, train_config.json.
                    None → use bundled pre-trained weights.
        config:     GenerateConfig instance; None → use defaults.

    Returns:
        List of (x, y, t_ms) tuples.  All values are ints.
        Last two entries are the mouse-down and mouse-up click events.
    """
    if config is None:
        config = GenerateConfig()

    model_dir_path = Path(model_dir) if model_dir else _BUNDLED_MODELS_DIR

    flow_pt = model_dir_path / "flow_model.pt"
    click_json = model_dir_path / "click_dist.json"
    duration_json = model_dir_path / "duration_dist.json"
    config_json = model_dir_path / "train_config.json"

    if not flow_pt.exists():
        if model_dir is not None:
            raise FileNotFoundError(
                f"Model weights not found in {model_dir_path}. "
                "Run training first or omit model_dir to use bundled weights."
            )
        raise FileNotFoundError(
            f"Bundled model weights missing at {_BUNDLED_MODELS_DIR}. "
            "Run training first."
        )

    # Load train config for model architecture params
    seq_len = config.seq_len
    d_model = 128
    nhead = 4
    num_layers = 4
    dim_feedforward = 256
    cond_dim = 3
    if config_json.exists():
        cfg = json.loads(config_json.read_text())
        seq_len = int(cfg.get("seq_len", seq_len))
        d_model = int(cfg.get("d_model", d_model))
        nhead = int(cfg.get("nhead", nhead))
        num_layers = int(cfg.get("num_layers", num_layers))
        dim_feedforward = int(cfg.get("dim_feedforward", dim_feedforward))
        cond_dim = int(cfg.get("cond_dim", cond_dim))

    # Load model
    model = TrajectoryFlowModel(
        seq_len=seq_len,
        d_model=d_model,
        nhead=nhead,
        num_layers=num_layers,
        dim_feedforward=dim_feedforward,
        cond_dim=cond_dim,
    )
    model.load_state_dict(
        torch.load(flow_pt, map_location="cpu", weights_only=True)
    )
    model.eval()

    # Load auxiliary distributions
    click_params: dict = {"mu": 80.0, "sigma": 30.0, "low": 20.0, "high": 300.0}
    if click_json.exists():
        click_params = json.loads(click_json.read_text())

    duration_dist: dict | None = None
    if duration_json.exists():
        duration_dist = json.loads(duration_json.read_text())

    # Compute pixel distance
    sx, sy = float(start[0]), float(start[1])
    ex, ey = float(end[0]), float(end[1])
    dist = math.hypot(ex - sx, ey - sy)
    dist = max(dist, 1.0)

    # Sample total duration
    if duration_dist is not None:
        total_duration = _sample_duration(duration_dist, dist)
    else:
        # Fallback: simple heuristic ~2px/ms
        total_duration = dist / 2.0
    if speed is not None and speed > 0:
        total_duration /= speed
    total_duration = max(total_duration, 10.0)

    # Build condition vector: [dist_norm, log_dist, log_total_dur]
    cond_arr = np.array(
        [
            dist / 2000.0,
            math.log(dist / 100.0),
            math.log(total_duration / 500.0),
        ],
        dtype=np.float32,
    )
    cond_t = torch.from_numpy(cond_arr).unsqueeze(0)  # (1, 3)

    # -----------------------------------------------------------------------
    # 4-step Euler ODE integration
    # -----------------------------------------------------------------------
    n_steps = config.n_steps
    dt = 1.0 / n_steps

    with torch.no_grad():
        x = torch.randn(1, seq_len, 3)  # start from noise
        for step in range(n_steps):
            t_val = step * dt
            t_tensor = torch.tensor([t_val])
            v = model(x, t_tensor, cond_t)
            x = x + v * dt
        # x is now the generated trajectory in (forward, lateral, log_dt) space

    decoded = x.squeeze(0).numpy()  # (seq_len, 3)

    forward = decoded[:, 0].copy()  # (seq_len,)
    lateral = decoded[:, 1].copy()  # (seq_len,)
    log_dt = decoded[:, 2].copy()   # (seq_len,)

    # ------------------------------------------------------------------
    # Spatial post-processing
    # ------------------------------------------------------------------

    # Endpoint snapping: lerp last 6 points towards (1.0, 0.0)
    n_snap = min(6, seq_len // 4)
    for i in range(n_snap):
        alpha = ((i + 1) / n_snap) ** 2  # quadratic ease-in
        k = seq_len - n_snap + i
        forward[k] = forward[k] * (1.0 - alpha) + 1.0 * alpha
        lateral[k] = lateral[k] * (1.0 - alpha) + 0.0 * alpha

    # Force first and last points to canonical values
    forward[0], lateral[0] = 0.0, 0.0
    forward[-1], lateral[-1] = 1.0, 0.0

    # Smooth start: dampen lateral near start (first 4 points)
    n_start_fix = min(4, seq_len // 4)
    for i in range(1, n_start_fix + 1):
        blend = i / (n_start_fix + 1)  # 0.2, 0.4, 0.6, 0.8
        forward[i] = max(forward[i], forward[i - 1])  # ensure monotonic start
        lateral[i] = lateral[i] * blend  # dampen lateral near start

    # Enforce forward monotonicity with soft correction (prevent x-jitter)
    for i in range(1, seq_len - 1):  # skip last point (already snapped to 1.0)
        if forward[i] < forward[i - 1]:
            forward[i] = forward[i - 1] + 0.001

    # Clamp forward to [0, 1] and re-force endpoints after monotonicity fix
    forward = np.clip(forward, 0.0, 1.0)
    forward[0] = 0.0
    forward[-1] = 1.0

    # Lateral 5-point gaussian smoothing (endpoints preserved)
    lateral = _gaussian_smooth(lateral, sigma=1.0)

    # ------------------------------------------------------------------
    # Temporal post-processing (log_dt)
    # ------------------------------------------------------------------

    # Clip log_dt to prevent extreme values after exp()
    # Training data log_dt = log(Δt_ms + 1), typical range [0, 4.5]
    # (e.g., Δt=1ms → 0.69, Δt=10ms → 2.40, Δt=80ms → 4.39)
    log_dt = np.clip(log_dt, 0.0, 5.0)

    # First point has no interval (padding from training)
    log_dt[0] = 0.0

    # Decode spatial coordinates to pixels
    normalised = np.stack([forward, lateral], axis=1)  # (seq_len, 2)
    pixels = decode_trajectory(normalised, start, end)  # (seq_len, 2)

    # Resample to n_points if needed
    if n_points != seq_len:
        pixels = resample_arc(pixels, n_points)
        # Also resample log_dt via linear interpolation in uniform arc
        log_dt = np.interp(
            np.linspace(0, 1, n_points),
            np.linspace(0, 1, seq_len),
            log_dt,
        )

    xs = pixels[:, 0]
    ys = pixels[:, 1]

    # Convert log_dt → dt (ms), scale to total_duration, clip
    dt_raw = np.exp(log_dt)
    dt_raw = np.clip(dt_raw, 0.0, None)
    dt_sum = dt_raw.sum()
    if dt_sum > 1e-6:
        scale = total_duration / dt_sum
    else:
        scale = total_duration / max(n_points, 1)
    dt_ms = np.clip(
        dt_raw * scale,
        config.dt_clip_min_ms,
        config.dt_clip_max_ms,
    )

    # Cumulative timestamps (start at 0)
    t_abs = np.cumsum(dt_ms)
    t_abs = np.concatenate([[0.0], t_abs[:-1]])  # shift so first point = 0

    # Ensure monotonically increasing
    for i in range(1, len(t_abs)):
        if t_abs[i] <= t_abs[i - 1]:
            t_abs[i] = t_abs[i - 1] + 1.0

    move_points: list[tuple[int, int, int]] = [
        (int(round(xs[i])), int(round(ys[i])), int(round(t_abs[i])))
        for i in range(n_points)
    ]

    # Sample click duration from truncated normal
    mu = float(click_params["mu"])
    sigma = float(click_params["sigma"])
    low = float(click_params["low"])
    high = float(click_params["high"])
    a, b = (low - mu) / sigma, (high - mu) / sigma
    click_duration = int(truncnorm.rvs(a, b, loc=mu, scale=sigma))
    click_duration = max(click_duration, int(low))

    last_t = move_points[-1][2]
    click_x = int(round(xs[-1]))
    click_y = int(round(ys[-1]))
    return move_points + [
        (click_x, click_y, last_t),
        (click_x, click_y, last_t + click_duration),
    ]
