"""Markdown report generation for eval results.

Outputs a self-contained .md file with embedded PNG plots in a sibling
'plots/' directory.
"""
from __future__ import annotations

import logging
from datetime import datetime
from pathlib import Path

import numpy as np

import matplotlib
matplotlib.use("Agg")  # headless
import matplotlib.pyplot as plt  # noqa: E402

from tools.eval.metrics import (
    compute_acceleration,
    compute_jerk,
    compute_speed,
    compute_stats,
    fft_spectrum,
    kl_divergence_histograms,
)

logger = logging.getLogger(__name__)


def _aggregate_kinematics(traces: list[dict]) -> dict[str, np.ndarray]:
    """Concatenate per-trace speed/accel/jerk arrays from a list of traces.

    Args:
        traces: list of {"xs", "ys", "ts"} dicts (1-D ndarrays).

    Returns:
        dict with keys "speed", "accel", "jerk", "dt" — each a flat ndarray.
    """
    speeds, accels, jerks, dts = [], [], [], []
    for tr in traces:
        xs, ys, ts = tr["xs"], tr["ys"], tr["ts"]
        if len(xs) < 4:
            continue
        v = compute_speed(xs, ys, ts)
        a = compute_acceleration(v, ts)
        j = compute_jerk(a, ts)
        speeds.append(v)
        accels.append(a)
        jerks.append(j)
        dts.append(np.diff(ts))
    return {
        "speed": np.concatenate(speeds) if speeds else np.array([]),
        "accel": np.concatenate(accels) if accels else np.array([]),
        "jerk": np.concatenate(jerks) if jerks else np.array([]),
        "dt": np.concatenate(dts) if dts else np.array([]),
    }


def _plot_distribution(
    gen: np.ndarray,
    ref: np.ndarray,
    title: str,
    output: Path,
    xlabel: str,
    bins: int = 50,
) -> None:
    """Side-by-side histogram of gen vs ref."""
    fig, ax = plt.subplots(figsize=(8, 4), dpi=100)
    if len(gen) > 0:
        ax.hist(gen, bins=bins, alpha=0.5, label="生成", density=True)
    if len(ref) > 0:
        ax.hist(ref, bins=bins, alpha=0.5, label="参考 (Balabit)", density=True)
    ax.set_title(title)
    ax.set_xlabel(xlabel)
    ax.set_ylabel("密度")
    ax.legend()
    fig.tight_layout()
    fig.savefig(output)
    plt.close(fig)


def _plot_fft_overlay(
    gen_traces: list[dict],
    ref_traces: list[dict],
    output: Path,
    sample_rate_hz: float = 100.0,
) -> None:
    """Average FFT magnitude over lateral component for gen vs ref."""
    def _avg_spectrum(traces: list[dict]) -> tuple[np.ndarray, np.ndarray]:
        all_freqs = None
        all_mags = []
        for tr in traces:
            xs, ys = tr["xs"], tr["ys"]
            if len(xs) < 8:
                continue
            sig = ys - np.linspace(ys[0], ys[-1], len(ys))
            f, m = fft_spectrum(sig, sample_rate_hz)
            if all_freqs is None:
                all_freqs = f
                all_mags.append(m)
            elif len(m) == len(all_freqs):
                all_mags.append(m)
        if not all_mags:
            return np.array([]), np.array([])
        return all_freqs, np.mean(all_mags, axis=0)

    fig, ax = plt.subplots(figsize=(8, 4), dpi=100)
    f_gen, m_gen = _avg_spectrum(gen_traces)
    f_ref, m_ref = _avg_spectrum(ref_traces)
    if len(f_gen) > 0:
        ax.plot(f_gen, m_gen, label="生成", alpha=0.7)
    if len(f_ref) > 0:
        ax.plot(f_ref, m_ref, label="参考 (Balabit)", alpha=0.7)
    ax.axvspan(4, 12, alpha=0.1, color="green", label="生理震颤区间 4–12 Hz")
    ax.set_title("FFT 频谱（横向偏移信号）")
    ax.set_xlabel("Hz")
    ax.set_ylabel("|FFT|")
    ax.set_xlim(0, sample_rate_hz / 2)
    ax.legend()
    fig.tight_layout()
    fig.savefig(output)
    plt.close(fig)


def _plot_paths_overlay(traces: list[dict], output: Path, max_traces: int = 5) -> None:
    """Plot up to N generated trajectories on the same axes."""
    fig, ax = plt.subplots(figsize=(6, 5), dpi=100)
    for i, tr in enumerate(traces[:max_traces]):
        ax.plot(tr["xs"], tr["ys"], alpha=0.6, label=f"路径 {i+1}")
    ax.invert_yaxis()  # screen coords
    ax.set_title(f"前 {min(max_traces, len(traces))} 条生成轨迹")
    ax.set_xlabel("x (px)")
    ax.set_ylabel("y (px)")
    ax.set_aspect("equal", adjustable="datalim")
    ax.legend()
    fig.tight_layout()
    fig.savefig(output)
    plt.close(fig)


def build_report(
    generated_traces: list[dict],
    reference_traces: list[dict],
    output_md: Path,
    tag: str,
    model_dir: str,
    sample_rate_hz: float = 100.0,
) -> None:
    """Build a Markdown eval report with embedded plots.

    Args:
        generated_traces: list of {"xs","ys","ts"} from the generator under test.
        reference_traces: list of {"xs","ys","ts"} from Balabit (ground truth).
        output_md: destination .md path. plots/ created in same dir.
        tag: short identifier (e.g. "baseline", "post-finetune").
        model_dir: model directory path string (for provenance).
        sample_rate_hz: nominal sample rate for FFT.
    """
    plot_dir = output_md.parent / "plots"
    plot_dir.mkdir(parents=True, exist_ok=True)

    gen_kin = _aggregate_kinematics(generated_traces)
    ref_kin = _aggregate_kinematics(reference_traces)

    # --- KL divergences ---
    kl_speed = kl_divergence_histograms(gen_kin["speed"], ref_kin["speed"])
    kl_accel = kl_divergence_histograms(gen_kin["accel"], ref_kin["accel"])
    kl_jerk = kl_divergence_histograms(gen_kin["jerk"], ref_kin["jerk"])
    kl_dt = kl_divergence_histograms(gen_kin["dt"], ref_kin["dt"])

    # --- Stats ---
    stats_gen = {k: compute_stats(v) for k, v in gen_kin.items()}
    stats_ref = {k: compute_stats(v) for k, v in ref_kin.items()}

    # --- Plots ---
    _plot_distribution(gen_kin["speed"], ref_kin["speed"],
                       "速度分布", plot_dir / f"{tag}-speed.png", "px/ms")
    _plot_distribution(gen_kin["accel"], ref_kin["accel"],
                       "加速度分布", plot_dir / f"{tag}-accel.png", "px/ms²")
    _plot_distribution(gen_kin["jerk"], ref_kin["jerk"],
                       "Jerk 分布", plot_dir / f"{tag}-jerk.png", "px/ms³")
    _plot_distribution(gen_kin["dt"], ref_kin["dt"],
                       "Δt 分布", plot_dir / f"{tag}-dt.png", "ms")
    _plot_fft_overlay(generated_traces, reference_traces,
                      plot_dir / f"{tag}-fft.png", sample_rate_hz)
    _plot_paths_overlay(generated_traces, plot_dir / f"{tag}-paths.png")

    # --- Markdown ---
    now = datetime.now().strftime("%Y-%m-%d %H:%M")
    lines = [
        f"# Eval Report: {tag} ({now})",
        "",
        "## 模型信息",
        f"- Checkpoint dir: `{model_dir}`",
        f"- 生成样本数: {len(generated_traces)}",
        f"- 参考样本数: {len(reference_traces)}",
        "",
        "## KL 散度（生成 vs 参考，越小越好）",
        "| 指标 | KL |",
        "|---|---|",
        f"| 速度分布 | {kl_speed:.4f} |",
        f"| 加速度分布 | {kl_accel:.4f} |",
        f"| Jerk 分布 | {kl_jerk:.4f} |",
        f"| Δt 分布 | {kl_dt:.4f} |",
        "",
        "## 摘要统计",
        "| 指标 | 生成 mean | 参考 mean | 生成 CV | 参考 CV |",
        "|---|---|---|---|---|",
    ]
    for key, label in [("speed", "速度"), ("accel", "加速度"), ("jerk", "jerk"), ("dt", "Δt")]:
        lines.append(
            f"| {label} | {stats_gen[key]['mean']:.4g} | {stats_ref[key]['mean']:.4g} | "
            f"{stats_gen[key]['cv']:.3f} | {stats_ref[key]['cv']:.3f} |"
        )
    lines += [
        "",
        "## 直方图",
        f"![速度](plots/{tag}-speed.png)",
        f"![加速度](plots/{tag}-accel.png)",
        f"![Jerk](plots/{tag}-jerk.png)",
        f"![Δt](plots/{tag}-dt.png)",
        "",
        "## FFT 频谱（横向偏移）",
        f"![FFT](plots/{tag}-fft.png)",
        "",
        "## 生成轨迹示例",
        f"![轨迹](plots/{tag}-paths.png)",
        "",
    ]
    output_md.write_text("\n".join(lines), encoding="utf-8")
    logger.info("Report written to %s", output_md)