from __future__ import annotations

import argparse
import json
from pathlib import Path
from typing import Any

import numpy as np
from scipy.ndimage import gaussian_filter

from .io import load_image, percentile_normalize, read_jsonl, resolve_manifest_path, write_json


def _q(vals: list[float], q: float, default: float) -> float:
    if not vals:
        return float(default)
    return float(np.quantile(np.asarray(vals, dtype=np.float64), float(q)))


def _clipf(v: float, lo: float, hi: float) -> float:
    return float(min(max(float(v), float(lo)), float(hi)))


def _clipi(v: float, lo: int, hi: int) -> int:
    return int(min(max(int(round(float(v))), int(lo)), int(hi)))


def _normalize_image01(img: np.ndarray) -> np.ndarray:
    raw = np.asarray(img)
    arr = raw.astype(np.float32, copy=False)
    if arr.size == 0:
        return arr
    if np.issubdtype(raw.dtype, np.integer):
        maxv = float(np.iinfo(raw.dtype).max)
        if maxv > 0:
            arr = arr / maxv
    arr = np.nan_to_num(arr, copy=False)
    a_min = float(np.min(arr))
    a_max = float(np.max(arr))
    if a_max > 1.0 or a_min < 0.0:
        # Match the train/inference representation used by datasets.
        arr = percentile_normalize(arr, low=1.0, high=99.0)
    return np.clip(arr, 0.0, 1.0).astype(np.float32)


def _parse_track_points(track: dict[str, Any]) -> tuple[list[tuple[int, float]], int | None]:
    pts: list[tuple[int, float]] = []
    raw_path = track.get("path_tx")
    if isinstance(raw_path, list):
        for node in raw_path:
            if not isinstance(node, (list, tuple)) or len(node) < 2:
                continue
            try:
                t = int(round(float(node[0])))
                x = float(node[1])
            except Exception:
                continue
            if np.isfinite(x):
                pts.append((t, x))

    if not pts:
        anchors = track.get("anchors")
        if isinstance(anchors, list):
            for a in anchors:
                if isinstance(a, dict):
                    try:
                        t = int(round(float(a.get("t"))))
                        x = float(a.get("x"))
                    except Exception:
                        continue
                elif isinstance(a, (list, tuple)) and len(a) >= 2:
                    try:
                        t = int(round(float(a[0])))
                        x = float(a[1])
                    except Exception:
                        continue
                else:
                    continue
                if np.isfinite(x):
                    pts.append((t, x))

    if not pts:
        try:
            t0 = int(round(float(track.get("t0"))))
            t1 = int(round(float(track.get("t1"))))
            x0 = float(track.get("x0"))
            slope = float(track.get("slope", 0.0))
            if np.isfinite(x0) and np.isfinite(slope) and t1 >= t0:
                pts = [(int(t), float(x0 + slope * (t - t0))) for t in range(t0, t1 + 1)]
        except Exception:
            pts = []

    max_gap = track.get("max_interp_gap", None)
    if max_gap is None:
        max_gap_i = None
    else:
        try:
            max_gap_i = int(max_gap)
        except Exception:
            max_gap_i = None

    pts.sort(key=lambda z: z[0])
    return pts, max_gap_i


def _dense_track(points: list[tuple[int, float]], max_interp_gap: int | None) -> dict[int, float]:
    if len(points) < 2:
        return {}
    by_t: dict[int, list[float]] = {}
    for (t0, x0), (t1, x1) in zip(points[:-1], points[1:]):
        by_t.setdefault(int(t0), []).append(float(x0))
        if t1 <= t0:
            continue
        dt = int(t1 - t0)
        if max_interp_gap is not None and max_interp_gap > 0 and dt > int(max_interp_gap):
            continue
        for t in range(int(t0) + 1, int(t1)):
            a = float((t - t0) / max(1, dt))
            x = float((1.0 - a) * x0 + a * x1)
            by_t.setdefault(int(t), []).append(float(x))
    by_t.setdefault(int(points[-1][0]), []).append(float(points[-1][1]))
    return {int(t): float(np.mean(xs)) for t, xs in by_t.items()}


def _track_class(track: dict[str, Any]) -> str:
    cls = str(track.get("class", "moving")).strip().lower()
    if cls in {"moving", "static", "diffusive", "crossing"}:
        return cls
    return "moving"


def _count_crossings(dense_tracks: list[dict[int, float]], tol_px: float = 1.0) -> int:
    n = len(dense_tracks)
    if n < 2:
        return 0
    out = 0
    tol = float(max(0.0, tol_px))
    for i in range(n):
        a = dense_tracks[i]
        if not a:
            continue
        for j in range(i + 1, n):
            b = dense_tracks[j]
            if not b:
                continue
            ts = set(a.keys()).intersection(b.keys())
            if len(ts) < 3:
                continue
            diffs = np.asarray([float(a[t] - b[t]) for t in ts], dtype=np.float32)
            has_pos = bool(np.any(diffs > tol))
            has_neg = bool(np.any(diffs < -tol))
            if has_pos and has_neg:
                out += 1
    return int(out)


def _contiguous_run_lengths(mask_1d: np.ndarray) -> list[int]:
    m = np.asarray(mask_1d, dtype=bool)
    out: list[int] = []
    n = int(m.shape[0])
    i = 0
    while i < n:
        if not m[i]:
            i += 1
            continue
        j = i + 1
        while j < n and m[j]:
            j += 1
        out.append(int(j - i))
        i = j
    return out


def _bg_axis_slopes(img01: np.ndarray, bg_mask: np.ndarray) -> tuple[float, float]:
    tdim, xdim = int(img01.shape[0]), int(img01.shape[1])
    row_means = np.empty((tdim,), dtype=np.float32)
    col_means = np.empty((xdim,), dtype=np.float32)
    for t in range(tdim):
        m = bg_mask[t]
        row_means[t] = float(np.mean(img01[t, m])) if np.any(m) else float(np.mean(img01[t]))
    for x in range(xdim):
        m = bg_mask[:, x]
        col_means[x] = float(np.mean(img01[m, x])) if np.any(m) else float(np.mean(img01[:, x]))
    xt = np.linspace(0.0, 1.0, tdim, dtype=np.float32)
    xx = np.linspace(0.0, 1.0, xdim, dtype=np.float32)
    slope_t = float(np.polyfit(xt, row_means, deg=1)[0]) if tdim >= 2 else 0.0
    slope_x = float(np.polyfit(xx, col_means, deg=1)[0]) if xdim >= 2 else 0.0
    return slope_t, slope_x


def parse_args() -> argparse.Namespace:
    ap = argparse.ArgumentParser(description="Calibrate synthetic generator profile from real labeled data.")
    ap.add_argument("--manifest", required=True, help="Real-data manifest JSONL with image/label/meta")
    ap.add_argument("--out", required=True, help="Output profile JSON")
    ap.add_argument("--max_samples", type=int, default=0, help="Use at most this many rows (0=all)")
    ap.add_argument("--log_every", type=int, default=50)
    return ap.parse_args()


def main() -> None:
    args = parse_args()
    manifest_path = Path(args.manifest)
    base_dir = manifest_path.resolve().parent
    rows = read_jsonl(manifest_path)
    if int(args.max_samples) > 0:
        rows = rows[: int(args.max_samples)]
    log_every = max(1, int(args.log_every))

    T_vals: list[float] = []
    X_vals: list[float] = []
    n_moving_per_img: list[float] = []
    n_static_per_img: list[float] = []
    n_diff_per_img: list[float] = []
    crossings_per_img: list[float] = []

    slope_abs_vals: list[float] = []
    track_span_vals: list[float] = []
    pause_events_vals: list[float] = []
    pause_dur_vals: list[float] = []
    residual_std_vals: list[float] = []
    track_missing_frac_vals: list[float] = []
    moving_has_pause_flags: list[float] = []
    diff_has_pause_flags: list[float] = []

    width_vals: list[float] = []
    img_median_vals: list[float] = []
    fg_mean_vals: list[float] = []
    bg_mean_vals: list[float] = []
    bg_std_vals: list[float] = []
    hp_std_vals: list[float] = []
    bg_p10_vals: list[float] = []
    bg_p50_vals: list[float] = []
    slope_t_vals: list[float] = []
    slope_x_vals: list[float] = []
    snr_vals: list[float] = []

    kept = 0
    for i, row in enumerate(rows, start=1):
        if "image" not in row or "label" not in row:
            continue
        img_path = resolve_manifest_path(base_dir, row["image"])
        lbl_path = resolve_manifest_path(base_dir, row["label"])
        img01 = _normalize_image01(load_image(img_path))
        lbl = np.asarray(load_image(lbl_path), dtype=np.int64)
        if img01.ndim != 2 or lbl.ndim != 2:
            continue
        if img01.shape != lbl.shape:
            continue

        kept += 1
        T_vals.append(float(img01.shape[0]))
        X_vals.append(float(img01.shape[1]))
        img_median_vals.append(float(np.median(img01)))

        moving_mask = lbl == 1
        static_mask = lbl == 2
        bg_mask = lbl == 0
        if not np.any(bg_mask):
            bg_mask = np.ones_like(lbl, dtype=bool)
        bg_vals = img01[bg_mask]
        bg_mean = float(np.mean(bg_vals))
        bg_std = float(np.std(bg_vals) + 1e-9)
        bg_mean_vals.append(bg_mean)
        bg_std_vals.append(bg_std)
        bg_p10_vals.append(float(np.quantile(bg_vals, 0.10)))
        bg_p50_vals.append(float(np.quantile(bg_vals, 0.50)))
        slope_t, slope_x = _bg_axis_slopes(img01, bg_mask)
        slope_t_vals.append(float(abs(slope_t)))
        slope_x_vals.append(float(abs(slope_x)))

        smooth = gaussian_filter(img01, sigma=(1.0, 1.0))
        hp = img01 - smooth
        hp_std_vals.append(float(np.std(hp[bg_mask]) + 1e-9))

        if np.any(moving_mask):
            fg_vals = img01[moving_mask]
            fg_mean = float(np.mean(fg_vals))
            fg_mean_vals.append(fg_mean)
            snr_vals.append(float((fg_mean - bg_mean) / (bg_std + 1e-9)))
            for t in range(moving_mask.shape[0]):
                width_vals.extend(_contiguous_run_lengths(moving_mask[t]))

        tracks_obj: list[dict[str, Any]] = []
        meta_rel = row.get("meta")
        if meta_rel:
            try:
                meta_path = resolve_manifest_path(base_dir, meta_rel)
                meta = json.loads(Path(meta_path).read_text(encoding="utf-8"))
                if isinstance(meta, dict):
                    tracks = meta.get("tracks", [])
                    if isinstance(tracks, list):
                        tracks_obj = [t for t in tracks if isinstance(t, dict)]
            except Exception:
                tracks_obj = []

        dense_for_cross: list[dict[int, float]] = []
        n_moving = 0
        n_static = 0
        n_diff = 0
        for tr in tracks_obj:
            cls = _track_class(tr)
            points, max_gap = _parse_track_points(tr)
            if len(points) < 2:
                continue
            dense = _dense_track(points, max_gap)
            if len(dense) < 2:
                continue
            ts = np.asarray(sorted(dense.keys()), dtype=np.float32)
            xs = np.asarray([dense[int(t)] for t in ts], dtype=np.float32)
            if ts.size >= 2:
                m, b = np.polyfit(ts, xs, deg=1)
                slope_abs_vals.append(float(abs(m)))
                fit = (m * ts) + b
                residual_std_vals.append(float(np.std(xs - fit)))
            track_span_vals.append(float(max(2, int(ts[-1] - ts[0] + 1))))
            dense_span = float(max(2, int(ts[-1] - ts[0] + 1)))
            track_missing_frac_vals.append(float(max(0.0, 1.0 - (len(dense) / dense_span))))

            gaps = [max(0, int(points[k + 1][0] - points[k][0] - 1)) for k in range(len(points) - 1)]
            gaps = [g for g in gaps if g > 0]
            pause_events_vals.append(float(len(gaps)))
            if gaps:
                pause_dur_vals.extend(float(g) for g in gaps)

            if cls in {"moving", "crossing"}:
                n_moving += 1
                moving_has_pause_flags.append(float(len(gaps) > 0))
                dense_for_cross.append(dense)
            elif cls == "diffusive":
                n_diff += 1
                diff_has_pause_flags.append(float(len(gaps) > 0))
                dense_for_cross.append(dense)
            elif cls == "static":
                n_static += 1

        if not tracks_obj:
            # Weak fallback from labels when metadata is unavailable.
            n_moving = int(np.any(moving_mask))
            n_static = int(np.any(static_mask))

        n_moving_per_img.append(float(n_moving))
        n_static_per_img.append(float(n_static))
        n_diff_per_img.append(float(n_diff))
        crossings_per_img.append(float(_count_crossings(dense_for_cross, tol_px=1.0)))

        if (i % log_every) == 0 or i == len(rows):
            print(f"[info] processed {i}/{len(rows)} rows (kept={kept})", flush=True)

    if kept == 0:
        raise RuntimeError("No valid samples were found for calibration.")

    width_med = _q(width_vals, 0.50, 2.2)
    sigma_center = _clipf(width_med / 2.35, 0.55, 1.80)
    snr_med = _q(snr_vals, 0.50, 5.0)
    gap_frac_med = _q(track_missing_frac_vals, 0.50, 0.03)
    read_noise_med = _q(hp_std_vals, 0.50, 0.018)
    bg_std_med = _q(bg_std_vals, 0.50, 0.02)
    slope_abs_med = _q(slope_abs_vals, 0.50, 1.10)
    img_median = _q(img_median_vals, 0.50, 0.45)
    gamma_center = _clipf(np.log(max(1e-3, img_median)) / np.log(0.5), 0.60, 1.80)
    amp_scale = _clipf(snr_med / 5.0, 0.65, 1.80)
    high_freq_ratio = _clipf(read_noise_med / max(1e-6, bg_std_med), 0.05, 1.50)

    pause_prob_moving = _clipf(
        float(np.mean(np.asarray(moving_has_pause_flags, dtype=np.float32))) if moving_has_pause_flags else 0.55,
        0.05,
        0.98,
    )
    pause_prob_diffusive = _clipf(
        float(np.mean(np.asarray(diff_has_pause_flags, dtype=np.float32))) if diff_has_pause_flags else 0.45,
        0.05,
        0.98,
    )
    fast_frac = _clipf(
        float(np.mean(np.asarray(slope_abs_vals, dtype=np.float32) >= 1.40)) if slope_abs_vals else 0.35,
        0.10,
        0.95,
    )

    synth_overrides = {
        "T": _clipi(_q(T_vals, 0.50, 220), 32, 2000),
        "X": _clipi(_q(X_vals, 0.50, 180), 32, 2000),
        "T_min": _clipi(_q(T_vals, 0.10, _q(T_vals, 0.50, 220)), 32, 2000),
        "T_max": _clipi(_q(T_vals, 0.90, _q(T_vals, 0.50, 220)), 32, 2000),
        "X_min": _clipi(_q(X_vals, 0.10, _q(X_vals, 0.50, 180)), 32, 2000),
        "X_max": _clipi(_q(X_vals, 0.90, _q(X_vals, 0.50, 180)), 32, 2000),
        "moving_min": _clipi(_q(n_moving_per_img, 0.10, 6), 1, 128),
        "moving_max": _clipi(_q(n_moving_per_img, 0.90, 12), 1, 256),
        "static_min": _clipi(_q(n_static_per_img, 0.10, 1), 0, 64),
        "static_max": _clipi(_q(n_static_per_img, 0.90, 5), 0, 128),
        "diffusive_min": _clipi(_q(n_diff_per_img, 0.10, 2), 0, 64),
        "diffusive_max": _clipi(_q(n_diff_per_img, 0.90, 6), 0, 128),
        "fast_track_fraction": fast_frac,
        "crossing_pairs_max": _clipi(_q(crossings_per_img, 0.90, 2), 0, 8),
        "frame_rate_scale": _clipf(1.4 * (1.1 / max(0.25, slope_abs_med)), 0.8, 2.6),
        "pause_prob_moving": pause_prob_moving,
        "pause_prob_diffusive": pause_prob_diffusive,
        "pause_events_min": _clipi(_q(pause_events_vals, 0.20, 1), 0, 6),
        "pause_events_max": _clipi(_q(pause_events_vals, 0.90, 3), 0, 10),
        "pause_dur_min": _clipi(_q(pause_dur_vals, 0.15, 4), 1, 48),
        "pause_dur_max": _clipi(_q(pause_dur_vals, 0.90, 24), 2, 96),
        "pause_jitter": _clipf(_q(residual_std_vals, 0.50, 0.06), 0.005, 0.20),
        "moving_amp_min": _clipf(0.45 * amp_scale, 0.15, 2.0),
        "moving_amp_max": _clipf(1.15 * amp_scale, 0.35, 3.2),
        "static_amp_min": _clipf(0.30 * amp_scale, 0.10, 1.8),
        "static_amp_max": _clipf(0.90 * amp_scale, 0.20, 2.6),
        "diffusive_amp_min": _clipf(0.28 * amp_scale, 0.10, 1.8),
        "diffusive_amp_max": _clipf(0.95 * amp_scale, 0.20, 2.8),
        "sigma_x_min": _clipf(0.75 * sigma_center, 0.45, 1.8),
        "sigma_x_max": _clipf(1.35 * sigma_center, 0.70, 2.8),
        "thickness_min": _clipi(width_med - 1.0, 1, 6),
        "thickness_max": _clipi(width_med + 1.0, 1, 8),
        "blink_moving_min": 0.0,
        "blink_moving_max": _clipf(0.008 + (0.55 * gap_frac_med), 0.008, 0.12),
        "blink_static_min": 0.0,
        "blink_static_max": _clipf(0.005 + (0.30 * gap_frac_med), 0.005, 0.08),
        "blink_diffusive_min": 0.0,
        "blink_diffusive_max": _clipf(0.010 + (0.70 * gap_frac_med), 0.010, 0.14),
        "blink_crossing_min": 0.0,
        "blink_crossing_max": _clipf(0.008 + (0.45 * gap_frac_med), 0.008, 0.10),
        "bg_offset_min": _clipf(_q(bg_p10_vals, 0.10, 0.02), 0.0, 0.4),
        "bg_offset_max": _clipf(_q(bg_p50_vals, 0.90, 0.15), 0.02, 0.7),
        "bg_grad_t_abs_max": _clipf(_q(slope_t_vals, 0.90, 0.06), 0.01, 0.20),
        "bg_grad_x_abs_max": _clipf(_q(slope_x_vals, 0.90, 0.06), 0.01, 0.20),
        "texture_sigma_t": _clipf(_q(T_vals, 0.50, 220) / 24.0, 4.0, 18.0),
        "texture_sigma_x": _clipf(_q(X_vals, 0.50, 180) / 20.0, 4.0, 18.0),
        # Keep 2D texture weak to avoid unrealistic leopard-like blotches.
        "texture_amp_min": _clipf(_q(hp_std_vals, 0.30, 0.010) * 0.30, 0.001, 0.020),
        "texture_amp_max": _clipf(_q(hp_std_vals, 0.80, 0.040) * 0.75, 0.003, 0.060),
        "texture_row_amp_min": _clipf(bg_std_med * 0.10, 0.001, 0.020),
        "texture_row_amp_max": _clipf(bg_std_med * 0.55, 0.003, 0.080),
        "texture_col_amp_min": _clipf(bg_std_med * 0.06, 0.001, 0.015),
        "texture_col_amp_max": _clipf(bg_std_med * 0.35, 0.002, 0.050),
        "texture_2d_frac_min": _clipf(0.04 + (0.05 * high_freq_ratio), 0.03, 0.20),
        "texture_2d_frac_max": _clipf(0.12 + (0.12 * high_freq_ratio), 0.08, 0.35),
        "psf_sigma_t_min": _clipf(0.60 * sigma_center, 0.20, 1.20),
        "psf_sigma_t_max": _clipf(1.20 * sigma_center, 0.35, 1.80),
        "psf_sigma_x_min": _clipf(0.55 * sigma_center, 0.20, 1.20),
        "psf_sigma_x_max": _clipf(1.10 * sigma_center, 0.35, 1.80),
        "poisson_gain_min": _clipf((25.0 + (12.0 * snr_med)) * 0.70, 12.0, 160.0),
        "poisson_gain_max": _clipf((25.0 + (12.0 * snr_med)) * 1.35, 20.0, 240.0),
        "read_noise_min": _clipf(read_noise_med * 0.50, 0.002, 0.06),
        "read_noise_max": _clipf(read_noise_med * 1.80, 0.006, 0.12),
        "gamma_min": _clipf(gamma_center - 0.25, 0.55, 1.80),
        "gamma_max": _clipf(gamma_center + 0.25, 0.70, 2.10),
    }

    # Keep min<=max relations safe.
    for lo_key, hi_key in (
        ("moving_min", "moving_max"),
        ("static_min", "static_max"),
        ("diffusive_min", "diffusive_max"),
        ("pause_events_min", "pause_events_max"),
        ("pause_dur_min", "pause_dur_max"),
        ("T_min", "T_max"),
        ("X_min", "X_max"),
        ("moving_amp_min", "moving_amp_max"),
        ("static_amp_min", "static_amp_max"),
        ("diffusive_amp_min", "diffusive_amp_max"),
        ("sigma_x_min", "sigma_x_max"),
        ("bg_offset_min", "bg_offset_max"),
        ("texture_amp_min", "texture_amp_max"),
        ("texture_row_amp_min", "texture_row_amp_max"),
        ("texture_col_amp_min", "texture_col_amp_max"),
        ("texture_2d_frac_min", "texture_2d_frac_max"),
        ("psf_sigma_t_min", "psf_sigma_t_max"),
        ("psf_sigma_x_min", "psf_sigma_x_max"),
        ("poisson_gain_min", "poisson_gain_max"),
        ("read_noise_min", "read_noise_max"),
        ("gamma_min", "gamma_max"),
    ):
        if float(synth_overrides[lo_key]) > float(synth_overrides[hi_key]):
            synth_overrides[lo_key], synth_overrides[hi_key] = synth_overrides[hi_key], synth_overrides[lo_key]

    payload = {
        "source_manifest": str(manifest_path),
        "n_rows": int(len(rows)),
        "n_used": int(kept),
        "summary": {
            "T_median": _q(T_vals, 0.50, 220),
            "X_median": _q(X_vals, 0.50, 180),
            "moving_per_image_median": _q(n_moving_per_img, 0.50, 8),
            "static_per_image_median": _q(n_static_per_img, 0.50, 2),
            "diffusive_per_image_median": _q(n_diff_per_img, 0.50, 3),
            "crossings_per_image_median": _q(crossings_per_img, 0.50, 0),
            "track_slope_abs_median": _q(slope_abs_vals, 0.50, 1.1),
            "track_width_median_px": width_med,
            "snr_median": snr_med,
            "bg_std_median": _q(bg_std_vals, 0.50, 0.02),
            "hp_std_median": read_noise_med,
            "gap_fraction_median": gap_frac_med,
        },
        "synth_overrides": synth_overrides,
    }
    write_json(args.out, payload)

    print(f"[done] wrote synth profile: {args.out}", flush=True)
    print("[next] generate calibrated synth with:", flush=True)
    print(
        "       python -m kymo_autopick.synth --out <OUT_DIR> --n 8000 --seed 123 "
        f"--profile_json {args.out}",
        flush=True,
    )


if __name__ == "__main__":
    main()
