#!/usr/bin/env python3
"""Compare cumulative adaptive DIS vs fixed initial-reference DIS (no integration).

Reproduces drift on static scenes where adaptive flow integrates incremental
displacement: final_flow = queue[-1] + prev_fine_flow (see orisys.tracker).

Usage:
  python tests/DIS_no_integrate.py --source /path/to/static_clip.mp4 --frames 500
  python tests/DIS_no_integrate.py --source 0 --frames 300 --plot drift.png
  python tests/DIS_no_integrate.py --config box --source 0 --no-show

Modes compared on identical grayscale frames:
  cumulative    Tracker(adaptive=True)  — production default
  non_adaptive  Tracker(adaptive=False) — single ref, warm-started DIS
  fixed_ref     DIS calc(ref0, frame, None) — initial frame only, no queue
"""

from __future__ import annotations

import argparse
import os
import sys
from dataclasses import dataclass, field
from pathlib import Path
from typing import Optional

import cv2
import numpy as np

REPO_ROOT = Path(__file__).resolve().parents[1]
SRC_ROOT = REPO_ROOT / "src"
for path in (REPO_ROOT, SRC_ROOT):
    if str(path) not in sys.path:
        sys.path.insert(0, str(path))

from orisys.tracker import Tracker  # noqa: E402


DEFAULT_CONFIG = "box"


def parse_source(value: str):
    text = str(value).strip()
    for prefix in ("camera:", "cam:", "webcam:"):
        if text.lower().startswith(prefix):
            text = text[len(prefix) :].strip()
            break
    try:
        return int(text)
    except ValueError:
        return text


def load_tracker_config(config_ref: str) -> dict:
    """Load tracker section from an Orisys config stem or JSON path."""
    path = Path(config_ref)
    if path.suffix.lower() == ".json" and path.is_file():
        import json

        data = json.loads(path.read_text(encoding="utf-8"))
    else:
        import orisys

        data = orisys.configs.load_config(config_ref)
    tracker = dict(data.get("tracker", {}))
    tracker.setdefault("algorithm", "dis")
    tracker.setdefault("dis_preset", "medium")
    tracker.setdefault("dis_mode", 3)
    tracker.setdefault("lower_threshold", 12.0)
    tracker.setdefault("upper_threshold", 18.0)
    tracker.setdefault("loop_closure_interval", 1)
    return tracker


def dis_preset_from_name(name: str) -> int:
    preset = (name or "medium").lower()
    if preset == "ultrafast":
        return cv2.DISOPTICAL_FLOW_PRESET_ULTRAFAST
    if preset == "fast":
        return cv2.DISOPTICAL_FLOW_PRESET_FAST
    return cv2.DISOPTICAL_FLOW_PRESET_MEDIUM


def flow_magnitude_stats(flow: np.ndarray) -> tuple[float, float]:
    """Return (mean magnitude, max magnitude) for HxWx2 flow."""
    if flow is None:
        return float("nan"), float("nan")
    mag = np.linalg.norm(flow.astype(np.float32), axis=-1)
    return float(mag.mean()), float(mag.max())


@dataclass
class FrameMetrics:
    mean_mag: float
    max_mag: float
    queue_len: int = 0


@dataclass
class ModeSeries:
    name: str
    mean_mags: list[float] = field(default_factory=list)
    max_mags: list[float] = field(default_factory=list)
    queue_lens: list[int] = field(default_factory=list)

    def append(self, metrics: FrameMetrics) -> None:
        self.mean_mags.append(metrics.mean_mag)
        self.max_mags.append(metrics.max_mag)
        self.queue_lens.append(metrics.queue_len)

    def slope(self) -> float:
        if len(self.mean_mags) < 2:
            return float("nan")
        x = np.arange(len(self.mean_mags), dtype=np.float64)
        y = np.asarray(self.mean_mags, dtype=np.float64)
        if not np.all(np.isfinite(y)):
            y = np.nan_to_num(y, nan=0.0)
        return float(np.polyfit(x, y, 1)[0])


class FixedRefDIS:
    """DIS from a fixed initial reference frame; no integration or queue."""

    def __init__(self, config_tracker: dict):
        preset = dis_preset_from_name(config_tracker.get("dis_preset", "medium"))
        self._dis = cv2.DISOpticalFlow_create(preset)
        self._dis.setFinestScale(int(config_tracker.get("dis_mode", 3)))
        self._ref: Optional[np.ndarray] = None

    def reset(self) -> None:
        self._ref = None

    def track(self, gray: np.ndarray) -> np.ndarray:
        if self._ref is None:
            self._ref = gray.copy()
            h, w = gray.shape[:2]
            return np.zeros((h, w, 2), dtype=np.float32)
        return self._dis.calc(self._ref, gray, None)


def make_tracker(config_tracker: dict, *, adaptive: bool) -> Tracker:
    cfg = dict(config_tracker)
    return Tracker(
        adaptive=adaptive,
        cuda=False,
        verbose=False,
        config_tracker=cfg,
    )


def read_gray_frame(cap: cv2.VideoCapture, static_image: Optional[np.ndarray]) -> tuple[bool, Optional[np.ndarray]]:
    if static_image is not None:
        return True, static_image.copy()
    ok, frame = cap.read()
    if not ok or frame is None:
        return False, None
    if frame.ndim == 3:
        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
    return True, frame


def run_comparison(
    source,
    *,
    config_tracker: dict,
    max_frames: int,
    reset_every: Optional[int],
    show: bool,
) -> dict[str, ModeSeries]:
    path = Path(str(source)) if isinstance(source, str) else None
    static_image: Optional[np.ndarray] = None
    if isinstance(source, str) and path is not None and path.is_file():
        img = cv2.imread(str(path), cv2.IMREAD_GRAYSCALE)
        if img is not None:
            static_image = img

    cap = None if static_image is not None else cv2.VideoCapture(source)
    if cap is not None and not cap.isOpened():
        raise RuntimeError(f"Cannot open source: {source}")

    cumulative = make_tracker(config_tracker, adaptive=True)
    non_adaptive = make_tracker(config_tracker, adaptive=False)
    fixed_ref = FixedRefDIS(config_tracker)

    series = {
        "cumulative": ModeSeries("cumulative"),
        "non_adaptive": ModeSeries("non_adaptive"),
        "fixed_ref": ModeSeries("fixed_ref"),
    }

    window = "DIS compare (cumulative | non_adaptive | fixed_ref)"
    if show:
        cv2.namedWindow(window, cv2.WINDOW_NORMAL)

    try:
        for frame_idx in range(max_frames):
            ok, gray = read_gray_frame(cap, static_image)
            if not ok or gray is None:
                break

            if reset_every and reset_every > 0 and frame_idx > 0 and frame_idx % reset_every == 0:
                cumulative.reset()

            flow_c = cumulative.track(gray.copy())
            flow_n = non_adaptive.track(gray.copy())
            flow_f = fixed_ref.track(gray)

            mean_c, max_c = flow_magnitude_stats(flow_c)
            mean_n, max_n = flow_magnitude_stats(flow_n)
            mean_f, max_f = flow_magnitude_stats(flow_f)

            series["cumulative"].append(
                FrameMetrics(mean_c, max_c, queue_len=cumulative.getFrameQueueLength())
            )
            series["non_adaptive"].append(FrameMetrics(mean_n, max_n, queue_len=1))
            series["fixed_ref"].append(FrameMetrics(mean_f, max_f, queue_len=1))

            if show and frame_idx % 2 == 0:
                h, w = gray.shape[:2]
                panels = []
                for flow, label, mean_val in (
                    (flow_c, "cumulative", mean_c),
                    (flow_n, "non_adaptive", mean_n),
                    (flow_f, "fixed_ref", mean_f),
                ):
                    vis = cv2.cvtColor(gray, cv2.COLOR_GRAY2BGR)
                    step = max(8, min(h, w) // 40)
                    ys, xs = np.mgrid[step // 2 : h : step, step // 2 : w : step]
                    for y, x in zip(ys.ravel(), xs.ravel()):
                        dx, dy = flow[y, x]
                        x2 = int(round(x + dx * 2))
                        y2 = int(round(y + dy * 2))
                        cv2.arrowedLine(vis, (x, y), (x2, y2), (0, 255, 0), 1, tipLength=0.3)
                    cv2.putText(
                        vis,
                        f"{label} m={mean_val:.3f}",
                        (8, 20),
                        cv2.FONT_HERSHEY_SIMPLEX,
                        0.5,
                        (0, 255, 255),
                        1,
                        cv2.LINE_AA,
                    )
                    panels.append(vis)
                combined = np.hstack(panels)
                cv2.imshow(window, combined)
                key = cv2.waitKey(1) & 0xFF
                if key == ord("q"):
                    break
    finally:
        if cap is not None:
            cap.release()
        if show:
            cv2.destroyAllWindows()

    return series


def print_summary(series: dict[str, ModeSeries], frames_processed: int) -> None:
    print(f"\nProcessed {frames_processed} frame(s)\n")
    header = f"{'mode':<14} {'initial':>10} {'final':>10} {'mean':>10} {'slope':>12} {'max_queue':>10}"
    print(header)
    print("-" * len(header))

    for key in ("cumulative", "non_adaptive", "fixed_ref"):
        s = series[key]
        if not s.mean_mags:
            print(f"{key:<14} {'(no data)':>10}")
            continue
        initial = s.mean_mags[0]
        final = s.mean_mags[-1]
        mean_all = float(np.mean(s.mean_mags))
        slope = s.slope()
        max_q = max(s.queue_lens) if s.queue_lens else 0
        print(
            f"{key:<14} {initial:10.4f} {final:10.4f} {mean_all:10.4f} {slope:12.6f} {max_q:10d}"
        )

    print("\nInterpretation on a static scene:")
    print("  cumulative   — rising mean/slope indicates integration drift")
    print("  fixed_ref    — should stay near a small bias if drift is removed")


def save_plot(series: dict[str, ModeSeries], plot_path: str) -> None:
    import matplotlib.pyplot as plt

    fig, ax = plt.subplots(figsize=(10, 5))
    for key, color in (
        ("cumulative", "C0"),
        ("non_adaptive", "C1"),
        ("fixed_ref", "C2"),
    ):
        s = series[key]
        if not s.mean_mags:
            continue
        ax.plot(s.mean_mags, label=key, color=color)
    ax.set_xlabel("frame")
    ax.set_ylabel("mean flow magnitude")
    ax.set_title("DIS drift comparison (static scene expected: flat fixed_ref)")
    ax.legend()
    ax.grid(True, alpha=0.3)
    fig.tight_layout()
    path = Path(plot_path)
    path.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(path, dpi=120)
    plt.close(fig)
    print(f"Saved plot: {path.resolve()}")


def parse_args():
    parser = argparse.ArgumentParser(
        description="Compare cumulative vs fixed-reference DIS optical flow drift"
    )
    parser.add_argument(
        "--source",
        default="0",
        help="Camera index, video path, or image path (default: 0)",
    )
    parser.add_argument(
        "--config",
        "-c",
        default=DEFAULT_CONFIG,
        help=f"Orisys config stem or JSON path for DIS settings (default: {DEFAULT_CONFIG})",
    )
    parser.add_argument(
        "--frames",
        type=int,
        default=300,
        help="Maximum frames to process (default: 300)",
    )
    parser.add_argument(
        "--reset-every",
        type=int,
        default=0,
        help="Reset cumulative tracker every N frames (0 = disabled)",
    )
    parser.add_argument(
        "--plot",
        nargs="?",
        const="dis_drift.png",
        default=None,
        help="Save matplotlib plot (optional path, default: dis_drift.png)",
    )
    parser.add_argument(
        "--no-show",
        action="store_true",
        help="Skip OpenCV preview window",
    )
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    source = parse_source(args.source)
    config_tracker = load_tracker_config(args.config)

    if config_tracker.get("algorithm", "dis") != "dis":
        print(f"Warning: config algorithm is {config_tracker.get('algorithm')}; forcing dis for test.")

    print("DIS no-integrate comparison")
    print(f"  source: {source}")
    print(f"  config: {args.config}")
    print(f"  dis_preset: {config_tracker.get('dis_preset')}  dis_mode: {config_tracker.get('dis_mode')}")
    print(f"  frames: {args.frames}  reset_every: {args.reset_every or 'off'}")

    series = run_comparison(
        source,
        config_tracker=config_tracker,
        max_frames=max(args.frames, 1),
        reset_every=args.reset_every or None,
        show=not args.no_show,
    )

    n = len(series["cumulative"].mean_mags)
    print_summary(series, n)

    if args.plot:
        save_plot(series, args.plot)

    return 0


def test_dis_no_integrate_smoke() -> None:
    """Optional pytest entry: synthetic static frames, no camera required."""
    if os.getenv("ORISYS_RUN_DIS_DRIFT") != "1":
        import pytest

        pytest.skip("Set ORISYS_RUN_DIS_DRIFT=1 to run DIS drift smoke test")

    gray = np.full((64, 64), 128, dtype=np.uint8)
    # Slight per-frame noise to mimic sensor variation without true motion.
    cap_frames = []
    rng = np.random.default_rng(0)
    for _ in range(60):
        noisy = gray.astype(np.int16) + rng.integers(-2, 3, gray.shape, dtype=np.int16)
        cap_frames.append(np.clip(noisy, 0, 255).astype(np.uint8))

    config_tracker = load_tracker_config(DEFAULT_CONFIG)
    cumulative = make_tracker(config_tracker, adaptive=True)
    fixed_ref = FixedRefDIS(config_tracker)

    cum_mags = []
    fix_mags = []
    for frame in cap_frames:
        fc = cumulative.track(frame.copy())
        ff = fixed_ref.track(frame)
        cum_mags.append(flow_magnitude_stats(fc)[0])
        fix_mags.append(flow_magnitude_stats(ff)[0])

    assert len(cum_mags) >= 10
    # Fixed ref should not drift more than cumulative on identical static+noise input.
    cum_slope = float(np.polyfit(np.arange(len(cum_mags)), cum_mags, 1)[0])
    fix_slope = float(np.polyfit(np.arange(len(fix_mags)), fix_mags, 1)[0])
    assert abs(fix_slope) <= abs(cum_slope) + 0.05, (
        f"fixed_ref slope {fix_slope:.4f} worse than cumulative {cum_slope:.4f}"
    )


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