"""Detect period and generate phase-aligned field extraction plan.

For a periodic case, this script:
1. Loads sensor data from controlled.npz (or sensors.npz for open-loop)
2. Detects dominant frequency and cycle stability (CV_T)
3. Selects the best 4-cycle stable window
4. Computes which step indices correspond to N_PTS uniform phase points
5. Generates a phase_plan.json for replay_fields.py

Usage:
    python3 src/CCD_analysis/scripts/detect_period.py --scene pinball
    python3 src/CCD_analysis/scripts/detect_period.py --scene illusion_1.0L
    python3 src/CCD_analysis/scripts/detect_period.py --scene target_cylinder_1.0L
    python3 src/CCD_analysis/scripts/detect_period.py --scene karman_re100

Output: data/resampled/{scene_name}/phase_plan.json
"""
from __future__ import annotations

import argparse
import json
import os
import sys

import numpy as np

_SRC = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))
if _SRC not in sys.path:
    sys.path.insert(0, _SRC)

from CCD_analysis.configs import SCENES, DATA_DIR
from CCD_analysis.utils.resampling import (
    detect_dominant_frequency, detect_cycle_stability,
)

N_CYCLES = 4
N_PTS = 24
CV_T_STRICT = 0.10
CV_T_RELAXED = 0.12


def run_single(scene_name: str, n_cycles: int = N_CYCLES, n_pts: int = N_PTS) -> dict:
    if scene_name not in SCENES:
        raise KeyError(f"Unknown scene: {scene_name}")

    cfg = SCENES[scene_name]
    scene_id = cfg["scene_id"]
    data_dir = os.path.join(DATA_DIR, scene_id, scene_name)
    si = cfg["sample_interval"]

    # Load sensor data
    controlled_path = os.path.join(data_dir, "controlled.npz")
    sensors_path = os.path.join(data_dir, "sensors.npz")

    if os.path.isfile(controlled_path):
        d = np.load(controlled_path)
        sensors = d.get("sensors")
        d.close()
    elif os.path.isfile(sensors_path):
        d = np.load(sensors_path)
        sensors = d.get("sensors")
        d.close()
    else:
        raise FileNotFoundError(f"No sensor data found for {scene_name}")

    if sensors is None or len(sensors) < 30:
        raise ValueError(f"Insufficient sensor data ({len(sensors) if sensors is not None else 0})")

    # Use centre sensor v-component for period detection
    signal = sensors[:, 3]

    # Frequency and stability
    f_case, T_case, _ = detect_dominant_frequency(signal, float(si))
    cv_T, mean_T, cy_lengths = detect_cycle_stability(signal, float(si))
    N_raw = mean_T / si if mean_T > 0 else 0
    rho = float(n_pts) / N_raw if N_raw > 0 else 99

    # Gate check (note: delta_f is always <= 0.1 since we self-compare)
    if cv_T <= CV_T_STRICT:
        gate = "strict"
    elif cv_T <= CV_T_RELAXED:
        gate = "relaxed"
    else:
        gate = "auxiliary"

    print(f"  f={f_case:.6f}, T={T_case:.0f}, CV_T={cv_T:.4f}, gate={gate}")
    print(f"  N_raw/cycle={N_raw:.1f}, rho_interp={rho:.2f}")

    if gate not in ("strict", "relaxed"):
        print(f"  WARNING: gate={gate}, case may be too unstable for clean CCD")
        # Still generate plan but flag it

    # Find cycle boundaries via rising zero-crossings
    y = signal - np.mean(signal)
    crossings = np.where((np.sign(y[:-1]) < 0) & (np.sign(y[1:]) > 0))[0]
    if len(crossings) < n_cycles + 1:
        raise ValueError(f"Only {len(crossings)} crossings found, need {n_cycles + 1}")

    # Select the most uniform block of n_cycles
    cycle_lens_steps = np.diff(crossings)
    T_exp_steps = T_case / si if T_case > 0 else float(np.median(cycle_lens_steps))
    best_score, best_start = float("inf"), 0
    for i in range(len(cycle_lens_steps) - n_cycles + 1):
        score = np.sum((cycle_lens_steps[i:i + n_cycles] - T_exp_steps) ** 2)
        if score < best_score:
            best_score, best_start = score, i
    selected_crossings = [int(crossings[best_start + k]) for k in range(n_cycles + 1)]

    # Map each (cycle, phase) to an exact step index
    # For cycle c (0..n_cycles-1), the range is [start, end) in step indices
    # We place n_pts equally spaced within [start, end)
    step_indices = []
    for c in range(n_cycles):
        i_start = selected_crossings[c]
        i_end = selected_crossings[c + 1]
        seg_len = i_end - i_start
        for p in range(n_pts):
            # Phase fraction: p / n_pts
            frac = p / float(n_pts)
            idx = int(i_start + frac * seg_len)
            step_indices.append(idx)

    # The last sample point (end of last cycle) should be included for completeness
    # But we keep exactly n_cycles * n_pts samples; step_indices[p + c * n_pts]

    phase_plan = {
        "scene": scene_name,
        "scene_id": scene_id,
        "n_cycles": n_cycles,
        "n_pts": n_pts,
        "total_steps": n_cycles * n_pts,
        "selected_crossings": selected_crossings,
        "step_indices": step_indices,
        "f_dom": f_case,
        "T_dom_steps": T_case,
        "CV_T": cv_T,
        "N_raw_per_cycle": float(N_raw),
        "rho_interp": rho,
        "gate": gate,
    }

    return phase_plan


def main():
    ap = argparse.ArgumentParser(description="Detect period and generate phase plan")
    ap.add_argument("--scene", type=str, default=None, help="Scene name")
    ap.add_argument("--all-periodic", action="store_true", help="Run all periodic scenes")
    ap.add_argument("--n-cycles", type=int, default=N_CYCLES)
    ap.add_argument("--n-pts", type=int, default=N_PTS)
    args = ap.parse_args()

    if args.all_periodic:
        scenes = [name for name, cfg in SCENES.items() if cfg["target_type"] == "periodic"]
    elif args.scene is not None:
        scenes = [args.scene]
    else:
        ap.print_help()
        return

    for scene_name in scenes:
        print(f"\n=== {scene_name} ===", flush=True)
        try:
            plan = run_single(scene_name, args.n_cycles, args.n_pts)
        except (FileNotFoundError, ValueError, KeyError) as e:
            print(f"  SKIP: {e}", flush=True)
            continue

        out_dir = os.path.join(DATA_DIR, "resampled", scene_name)
        os.makedirs(out_dir, exist_ok=True)
        with open(os.path.join(out_dir, "phase_plan.json"), "w") as f:
            json.dump(plan, f, indent=2)
        print(f"  Saved: {out_dir}/phase_plan.json", flush=True)

    print("\nDone.", flush=True)


if __name__ == "__main__":
    main()
