#!/usr/bin/env python3
"""Plot range-vs-2D-feature correlation figures for one tracked object."""

from __future__ import annotations

import argparse
import csv
import json
from pathlib import Path

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np


DEFAULT_FEATURES = ("h", "area", "y2")
FEATURE_LABELS = {
    "x_ego": "X Ego (m)",
    "w": "BBox Width (px)",
    "h": "BBox Height (px)",
    "area": "BBox Area (px^2)",
    "x1": "BBox x1 (px)",
    "x2": "BBox x2 (px)",
    "y1": "BBox y1 (px)",
    "y2": "BBox y2 (px)",
    "cx": "BBox cx (px)",
    "cy": "BBox cy (px)",
    "aspect": "BBox Aspect",
}
FEATURE_COLORS = {
    "x_ego": "#1d4ed8",
    "h": "#059669",
    "area": "#b45309",
    "y2": "#b91c1c",
    "y1": "#7c3aed",
    "cy": "#db2777",
    "w": "#0f766e",
    "x1": "#475569",
    "x2": "#9333ea",
    "cx": "#ea580c",
    "aspect": "#0891b2",
}


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--tracking", required=True, type=Path, help="Tracking JSON, e.g. merge.json")
    parser.add_argument("--series-csv", required=True, type=Path, help="Per-track temporal observation CSV")
    parser.add_argument("--track-id", required=True, type=int)
    parser.add_argument("--output-dir", required=True, type=Path)
    parser.add_argument(
        "--feature",
        action="append",
        dest="features",
        help="2D feature to include. Can repeat. Defaults to h, area, y2.",
    )
    parser.add_argument("--smooth-window", type=int, default=21)
    parser.add_argument("--max-lag", type=int, default=20)
    return parser.parse_args()


def load_series_csv(path: Path) -> list[dict]:
    rows = list(csv.DictReader(path.open()))
    for row in rows:
        row["frame_id"] = int(row["frame_id"])
        row["frame_idx"] = int(row["frame_idx"])
        row["x_ego"] = float(row["x_ego"])
    return rows


def load_track_feature_series(tracking_path: Path, frame_ids: set[int], track_id: int) -> list[dict]:
    frames = json.loads(tracking_path.read_text())
    out = []
    for frame in frames:
        det = next((item for item in frame.get("detections", []) if int(item.get("track_id", -1)) == track_id), None)
        if det is None:
            continue
        frame_id = int(det.get("frameId", det.get("frame_id")))
        if frame_id not in frame_ids:
            continue
        x1, y1, x2, y2 = map(float, det["bbox"])
        w = x2 - x1
        h = y2 - y1
        out.append(
            {
                "frame_id": frame_id,
                "x1": x1,
                "y1": y1,
                "x2": x2,
                "y2": y2,
                "cx": (x1 + x2) * 0.5,
                "cy": (y1 + y2) * 0.5,
                "w": w,
                "h": h,
                "area": w * h,
                "aspect": w / h if h else 0.0,
            }
        )
    out.sort(key=lambda item: item["frame_id"])
    return out


def moving_average(values: np.ndarray, window: int) -> np.ndarray:
    if window <= 1:
        return values.copy()
    pad = window // 2
    padded = np.pad(values, (pad, pad), mode="edge")
    kernel = np.ones(window, dtype=np.float64) / window
    return np.convolve(padded, kernel, mode="valid")


def detrend(values: np.ndarray) -> np.ndarray:
    x = np.arange(len(values), dtype=np.float64)
    coef = np.polyfit(x, values, 1)
    return values - np.polyval(coef, x)


def zscore(values: np.ndarray) -> np.ndarray:
    std = float(np.std(values))
    if std == 0.0:
        return np.zeros_like(values, dtype=np.float64)
    return (values - float(np.mean(values))) / std


def corr(a: np.ndarray, b: np.ndarray) -> float:
    a = a.astype(np.float64)
    b = b.astype(np.float64)
    a = a - a.mean()
    b = b - b.mean()
    denom = float(np.sqrt(np.dot(a, a) * np.dot(b, b)))
    if denom == 0.0:
        return 0.0
    return float(np.dot(a, b) / denom)


def build_arrays(series_rows: list[dict], feature_rows: list[dict], features: list[str]) -> dict[str, np.ndarray]:
    feature_map = {row["frame_id"]: row for row in feature_rows}
    frame_ids = [row["frame_id"] for row in series_rows if row["frame_id"] in feature_map]
    result = {
        "frame_id": np.asarray(frame_ids, dtype=np.int64),
        "x_ego": np.asarray([row["x_ego"] for row in series_rows if row["frame_id"] in feature_map], dtype=np.float64),
    }
    for feat in features:
        result[feat] = np.asarray([feature_map[fid][feat] for fid in frame_ids], dtype=np.float64)
    return result


def save_summary(path: Path, payload: dict) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(payload, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")


def plot_raw_curves(arrays: dict[str, np.ndarray], features: list[str], output_path: Path) -> None:
    frame_id = arrays["frame_id"]
    fig, axes = plt.subplots(len(features) + 1, 1, figsize=(14, 3.2 * (len(features) + 1)), sharex=True)
    if len(features) == 0:
        axes = [axes]

    axes[0].plot(frame_id, arrays["x_ego"], color=FEATURE_COLORS["x_ego"], linewidth=1.4)
    axes[0].set_ylabel(FEATURE_LABELS["x_ego"])
    axes[0].set_title("Raw Curves")
    axes[0].grid(True, alpha=0.3)

    for idx, feat in enumerate(features, start=1):
        axes[idx].plot(frame_id, arrays[feat], color=FEATURE_COLORS.get(feat, "#374151"), linewidth=1.2)
        axes[idx].set_ylabel(FEATURE_LABELS.get(feat, feat))
        axes[idx].grid(True, alpha=0.3)

    axes[-1].set_xlabel("Frame ID")
    fig.tight_layout()
    fig.savefig(output_path, dpi=150)
    plt.close(fig)


def plot_smoothed_detrended_overlay(
    arrays: dict[str, np.ndarray],
    features: list[str],
    smooth_window: int,
    output_path: Path,
) -> dict[str, np.ndarray]:
    frame_id = arrays["frame_id"]
    processed = {}
    fig, ax = plt.subplots(figsize=(14, 6))

    x_vals = zscore(detrend(moving_average(arrays["x_ego"], smooth_window)))
    processed["x_ego"] = x_vals
    ax.plot(frame_id, x_vals, color=FEATURE_COLORS["x_ego"], linewidth=2.0, label="x_ego")

    for feat in features:
        vals = zscore(detrend(moving_average(arrays[feat], smooth_window)))
        processed[feat] = vals
        ax.plot(frame_id, vals, linewidth=1.6, label=feat, color=FEATURE_COLORS.get(feat, "#374151"))

    ax.set_title(f"Smoothed + Detrended + Z-scored Curves (window={smooth_window})")
    ax.set_xlabel("Frame ID")
    ax.set_ylabel("Z-score")
    ax.grid(True, alpha=0.3)
    ax.legend()
    fig.tight_layout()
    fig.savefig(output_path, dpi=150)
    plt.close(fig)
    return processed


def plot_delta_overlay(arrays: dict[str, np.ndarray], features: list[str], output_path: Path) -> dict[str, np.ndarray]:
    frame_id = arrays["frame_id"][1:]
    processed = {}
    fig, ax = plt.subplots(figsize=(14, 6))

    dx = np.diff(arrays["x_ego"])
    processed["x_ego"] = dx
    ax.plot(frame_id, zscore(dx), color=FEATURE_COLORS["x_ego"], linewidth=2.0, label="Δx_ego")

    for feat in features:
        vals = np.diff(arrays[feat])
        processed[feat] = vals
        ax.plot(frame_id, zscore(vals), linewidth=1.4, label=f"Δ{feat}", color=FEATURE_COLORS.get(feat, "#374151"))

    ax.set_title("Frame-to-Frame Delta Curves (Z-scored)")
    ax.set_xlabel("Frame ID")
    ax.set_ylabel("Z-score")
    ax.grid(True, alpha=0.3)
    ax.legend()
    fig.tight_layout()
    fig.savefig(output_path, dpi=150)
    plt.close(fig)
    return processed


def plot_delta_scatter(delta_arrays: dict[str, np.ndarray], features: list[str], output_path: Path) -> dict[str, dict[str, float]]:
    dx = delta_arrays["x_ego"]
    fig, axes = plt.subplots(1, len(features), figsize=(5.3 * len(features), 4.6))
    if len(features) == 1:
        axes = [axes]

    summary = {}
    for ax, feat in zip(axes, features):
        vals = delta_arrays[feat]
        signed = corr(dx, vals)
        magnitude = corr(np.abs(dx), np.abs(vals))
        ax.scatter(vals, dx, s=12, alpha=0.55, color=FEATURE_COLORS.get(feat, "#374151"))
        ax.set_xlabel(f"Δ{feat}")
        ax.set_ylabel("Δx_ego")
        ax.set_title(f"{feat}\nr={signed:.3f}, |r|={magnitude:.3f}")
        ax.grid(True, alpha=0.25)
        summary[feat] = {
            "corr_signed": signed,
            "corr_magnitude": magnitude,
        }

    fig.tight_layout()
    fig.savefig(output_path, dpi=150)
    plt.close(fig)
    return summary


def plot_cross_correlation(
    processed: dict[str, np.ndarray],
    features: list[str],
    max_lag: int,
    output_path: Path,
) -> dict[str, list[dict[str, float]]]:
    ref = processed["x_ego"]
    fig, ax = plt.subplots(figsize=(14, 6))
    summary = {}

    for feat in features:
        vals = processed[feat]
        points = []
        for lag in range(-max_lag, max_lag + 1):
            if lag < 0:
                corr_value = corr(ref[-lag:], vals[:lag])
            elif lag > 0:
                corr_value = corr(ref[:-lag], vals[lag:])
            else:
                corr_value = corr(ref, vals)
            points.append({"lag": lag, "corr": corr_value})
        summary[feat] = points
        ax.plot([item["lag"] for item in points], [item["corr"] for item in points], label=feat, linewidth=1.8)

    ax.axhline(0.0, color="#6b7280", linestyle="--", linewidth=1)
    ax.set_title("Cross-correlation of Low-frequency Components")
    ax.set_xlabel("Lag (frames)")
    ax.set_ylabel("Correlation")
    ax.grid(True, alpha=0.3)
    ax.legend()
    fig.tight_layout()
    fig.savefig(output_path, dpi=150)
    plt.close(fig)
    return summary


def main() -> int:
    args = parse_args()
    features = args.features or list(DEFAULT_FEATURES)
    tracking_path = args.tracking.resolve()
    series_csv = args.series_csv.resolve()
    output_dir = args.output_dir.resolve()
    output_dir.mkdir(parents=True, exist_ok=True)

    series_rows = load_series_csv(series_csv)
    frame_ids = {row["frame_id"] for row in series_rows}
    feature_rows = load_track_feature_series(tracking_path, frame_ids, args.track_id)
    arrays = build_arrays(series_rows, feature_rows, features)

    plot_raw_curves(arrays, features, output_dir / "raw_curves.png")
    smooth_processed = plot_smoothed_detrended_overlay(
        arrays,
        features,
        smooth_window=args.smooth_window,
        output_path=output_dir / "smoothed_detrended_zscore_overlay.png",
    )
    delta_arrays = plot_delta_overlay(arrays, features, output_dir / "delta_zscore_overlay.png")
    scatter_summary = plot_delta_scatter(delta_arrays, features, output_dir / "delta_scatter.png")
    cross_corr_summary = plot_cross_correlation(
        smooth_processed,
        features,
        max_lag=args.max_lag,
        output_path=output_dir / "cross_correlation.png",
    )

    summary = {
        "track_id": args.track_id,
        "features": features,
        "tracking_json": str(tracking_path),
        "series_csv": str(series_csv),
        "smooth_window": args.smooth_window,
        "max_lag": args.max_lag,
        "scatter_summary": scatter_summary,
        "cross_correlation": cross_corr_summary,
    }
    save_summary(output_dir / "summary.json", summary)
    print(output_dir)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
