#!/usr/bin/env python3
"""Compute quantitative metrics for steady cloak.

Measures:
  - Mean wake restoration (downstream ux profile)
  - Fluctuation (RMS) suppression ratio
  - Recirculation zone length (centreline ux < 0)
  - dq_ctl + dq_blk cancellation quality
  - Force / power bookkeeping (if forces available)

Usage:
    conda run -n pycuda_3_10 python correction_analysis/run_steady_metrics.py
"""
from __future__ import annotations

import json
import os
import sys
from typing import Any

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 DATA_DIR, NX, NY, L0, CENTER_Y, U0, SCENES
from CCD_analysis.correction_analysis.compute_correction_fields import (
    compute_correction,
)
from CCD_analysis.correction_analysis.process_legacy_steady import (
    load_legacy_steady,
)

# ---------------------------------------------------------------------------
# Sensor / geometry constants (steady_cloak layout)
# ---------------------------------------------------------------------------
# pinball_front_x = 30.0 * L0 = 600
# pinball_rear_x  = 31.3 * L0 = 626
# sensor_x        = 40.0 * L0 = 800
SENSOR_X_PX = int(SCENES["steady_cloak"]["sensor_x"] * L0)  # ~800
FRONT_X_PX = int(SCENES["steady_cloak"]["pinball_front_x"] * L0)  # ~600
REAR_X_PX = int(SCENES["steady_cloak"]["pinball_rear_x"] * L0)  # ~626
CY = int(round(CENTER_Y))  # centreline row index


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _crop_sensor_zone(ux: np.ndarray) -> np.ndarray:
    """Crop ux to sensor-zone column range around SENSOR_X_PX.

    Handles both 2D (NY, NX) and 3D (N, NY, NX) arrays by always
    cropping the last (x) axis.
    """
    half = int(NX * 0.1)  # ~10 % of total width on each side
    x0 = max(0, SENSOR_X_PX - half)
    x1 = min(NX, SENSOR_X_PX + half)
    # Ellipsis crops the last axis regardless of dimensionality
    return ux[..., x0:x1]


def _crop_streamwise_centreline(ux_mean: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    """Extract centreline ux profile from mean field, trimmed to downstream region.

    Returns
    -------
    x_vals : (NX_trim,) pixel indices
    ux_cl  : (NX_trim,) centreline ux values
    """
    # Downstream region: from body trailing edge to domain end
    x0 = REAR_X_PX - 20  # start a bit before body
    x_end = min(NX, int(NX * 0.95))
    ux_cl = ux_mean[CY, x0:x_end]
    x_vals = np.arange(x0, x_end)
    return x_vals, ux_cl


# ---------------------------------------------------------------------------
# Metrics
# ---------------------------------------------------------------------------

def compute_mean_wake_restoration(
    dq_blk: dict, dq_ctl: dict, q_in: dict, q_ctl_raw: dict
) -> dict[str, Any]:
    """Compare downstream ux profiles for blocked vs controlled flow.

    A perfect steady cloak restores the wake to the uniform-channel profile.
    """
    # Mean fields from correction differences
    ux_blk_mean = np.mean(dq_blk["ux"], axis=0)  # (NY, NX) — blockage perturbation
    ux_ctl_mean = np.mean(dq_ctl["ux"], axis=0)  # (NY, NX) — control perturbation

    # Actual mean fields
    # q_in (target_channel) is the ideal undisturbed profile
    ux_in_mean = np.mean(q_in["ux"], axis=0)  # (NY, NX)

    # q_ctl raw = steady_cloak mean
    ux_sc_mean = np.mean(q_ctl_raw["ux"], axis=0)  # (NY, NX)

    # Centreline profiles
    x_vals, ux_cl_in = _crop_streamwise_centreline(ux_in_mean)
    _, ux_cl_sc = _crop_streamwise_centreline(ux_sc_mean)
    _, ux_cl_blk = _crop_streamwise_centreline(ux_blk_mean)
    _, ux_cl_ctl = _crop_streamwise_centreline(ux_ctl_mean)

    # Sensor-zone averaged ux (over the full field, not just centreline)
    sz_blk = np.mean(_crop_sensor_zone(dq_blk["ux"]))
    sz_ctl = np.mean(_crop_sensor_zone(dq_ctl["ux"]))

    # Wake restoration metric: RMS deviation from target channel in sensor zone
    _, ux_sz_in = _crop_sensor_zone(ux_in_mean), None  # not used for deviation
    ux_sz_sc = _crop_sensor_zone(ux_sc_mean)
    ux_sz_tc = _crop_sensor_zone(ux_in_mean)
    dev_sc = np.std(ux_sz_sc - ux_sz_tc)

    return {
        "sensor_zone_mean_ux_blk": float(np.mean(sz_blk)),
        "sensor_zone_mean_ux_ctl": float(np.mean(sz_ctl)),
        "centreline_ux_blk_mean": float(np.mean(ux_cl_blk)),
        "centreline_ux_ctl_mean": float(np.mean(ux_cl_ctl)),
        "sensor_zone_deviation_from_channel": float(dev_sc),
    }


def compute_rms_suppression(dq_blk: dict, dq_ctl: dict) -> dict[str, Any]:
    """Compute RMS fluctuation suppression ratio.

    suppression_ratio = 1 - RMS(rms_ctl) / RMS(rms_blk)
    where rms is computed per-pixel over the snapshot dimension.

    A value of 1.0 = perfect suppression, 0.0 = no suppression.
    """
    rms_blk = np.std(dq_blk["ux"], axis=0)  # (NY, NX)
    rms_ctl = np.std(dq_ctl["ux"], axis=0)  # (NY, NX)

    global_rms_blk = np.sqrt(np.mean(rms_blk**2))
    global_rms_ctl = np.sqrt(np.mean(rms_ctl**2))

    suppression_ratio = 1.0 - global_rms_ctl / max(global_rms_blk, 1e-15)

    # Sensor-zone specific
    sz_blk = _crop_sensor_zone(rms_blk)
    sz_ctl = _crop_sensor_zone(rms_ctl)
    sz_suppression = 1.0 - np.mean(sz_ctl) / max(np.mean(sz_blk), 1e-15)

    return {
        "global_RMS_blk": float(global_rms_blk),
        "global_RMS_ctl": float(global_rms_ctl),
        "suppression_ratio": float(suppression_ratio),
        "sensor_zone_RMS_blk": float(np.mean(sz_blk)),
        "sensor_zone_RMS_ctl": float(np.mean(sz_ctl)),
        "sensor_zone_suppression": float(sz_suppression),
    }


def compute_recirculation_zone(q_ctl_raw: dict) -> dict[str, Any]:
    """Find recirculation zone length from mean ux of steady_cloak field.

    Recirculation length: streamwise distance from body trailing edge
    to the point where centreline ux recovers to >= 0.
    """
    ux_mean = np.mean(q_ctl_raw["ux"], axis=0)  # (NY, NX)
    x_vals, ux_cl = _crop_streamwise_centreline(ux_mean)

    # Find first point (downstream of body) where ux returns to >= 0
    neg = ux_cl < 0
    if not np.any(neg):
        recirc_len = 0.0
        x_recovery = None
    else:
        # Find the last negative index in this trimmed region
        neg_indices = np.where(neg)[0]
        last_neg = neg_indices[-1]
        x_recovery = int(x_vals[last_neg])
        # Distance from rear cylinder in pixel units, convert to L0
        recirc_len = (x_recovery - REAR_X_PX) / L0

    # Also report min centreline ux
    min_ux = float(np.min(ux_cl))

    return {
        "recirculation_length_L0": float(recirc_len) if recirc_len is not None else 0.0,
        "recirculation_x_recovery_px": x_recovery,
        "centreline_min_ux": min_ux,
    }


def compute_cancellation_quality(dq_blk: dict, dq_ctl: dict) -> dict[str, Any]:
    """Compute residual cancellation quality.

    For perfect steady cloak: dq_ctl ≈ -dq_blk (control cancels blockage).
    measured by: cancellation_ratio = RMS(dq_ctl + dq_blk) / RMS(dq_blk)
    (lower is better, 0.0 = perfect cancellation)
    """
    residual_ux = dq_ctl["ux"] + dq_blk["ux"]  # (N, NY, NX)
    rms_residual = np.std(residual_ux)
    rms_blk = np.std(dq_blk["ux"])

    cancel_ratio = rms_residual / max(rms_blk, 1e-15)

    # Sensor-zone specific
    sz_res = _crop_sensor_zone(residual_ux)
    sz_blk_rms = np.std(_crop_sensor_zone(dq_blk["ux"]))
    sz_cancel = np.std(sz_res) / max(sz_blk_rms, 1e-15)

    return {
        "cancellation_ratio": float(cancel_ratio),
        "sensor_zone_cancellation_ratio": float(sz_cancel),
        "residual_RMS": float(rms_residual),
        "blockage_RMS": float(rms_blk),
    }


def compute_force_bookkeeping(
    dq_blk: dict, dq_ctl: dict
) -> dict[str, Any]:
    """Estimate drag from field data.

    Since steady_cloak sensors.npz does not contain force telemetry,
    we estimate drag proxy from the momentum deficit in the wake.

    drag_proxy = integral of (U0 - ux) across a wake profile
    (qualitative comparison only, not calibrated to actual drag)
    """
    # Use mean ux from blockage (pinball - channel) and control (steady_cloak - pinball)
    ux_blk_m = np.mean(dq_blk["ux"], axis=0)  # blockage perturbation
    ux_ctl_m = np.mean(dq_ctl["ux"], axis=0)  # control perturbation

    # Mean flow = blockage + channel for pinball; control restores toward channel
    # Take a wake profile at sensor_x location
    # Drag proxy: momentum deficit across channel height
    # Positive deficit means flow slower than free-stream
    deficit_blk = float(np.trapz(-ux_blk_m[:, SENSOR_X_PX])) if SENSOR_X_PX < NX else 0.0
    deficit_ctl = float(np.trapz(-ux_ctl_m[:, SENSOR_X_PX])) if SENSOR_X_PX < NX else 0.0
    deficit_ratio = deficit_ctl / max(abs(deficit_blk), 1e-15)

    return {
        "drag_proxy_blockage": deficit_blk,
        "drag_proxy_control": deficit_ctl,
        "drag_proxy_ratio": deficit_ratio,
        "note": "drag proxy from ux deficit at sensor plane; no actual force telemetry available",
    }


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def run() -> dict[str, Any]:
    print("=" * 60)
    print("Steady Cloak Quantitative Metrics")
    print("=" * 60)

    # -- 1. Load correction fields --
    print("\n--- Loading correction fields ---")
    corr = compute_correction("steady_cloak")
    dq_blk = corr["dq_blk"]  # q_blk - q_in (pinball blockage)
    dq_ctl = corr["dq_ctl"]  # q_ctl - q_blk (control correction)
    q_in = corr["q_in"]      # target_channel (undisturbed)
    N = corr["N"]
    print(f"  Aligned frames: {N}")

    # -- 2. Load steady_cloak raw data (for true mean field) --
    print("\n--- Loading steady_cloak raw fields ---")
    sc_raw = load_legacy_steady("steady_cloak")
    print(f"  ux shape: {sc_raw['ux'].shape}")

    # -- 3. Compute metrics --
    print("\n--- Computing wake restoration ---")
    wake = compute_mean_wake_restoration(dq_blk, dq_ctl, q_in, sc_raw)

    print("\n--- Computing RMS suppression ---")
    rms = compute_rms_suppression(dq_blk, dq_ctl)

    print("\n--- Computing recirculation zone ---")
    recirc = compute_recirculation_zone(sc_raw)

    print("\n--- Computing cancellation quality ---")
    cancel = compute_cancellation_quality(dq_blk, dq_ctl)

    print("\n--- Computing drag bookkeeping ---")
    drag = compute_force_bookkeeping(dq_blk, dq_ctl)

    # -- 4. Assemble --
    metrics = {
        "scene": "steady_cloak",
        "N_frames": int(N),
        "N_raw": int(sc_raw["ux"].shape[0]),
        "recirculation_zone": recirc,
        "wake_restoration": wake,
        "rms_suppression": rms,
        "cancellation_quality": cancel,
        "force_bookkeeping": drag,
    }

    # -- 5. Save --
    os.makedirs(os.path.join(DATA_DIR, "ccd"), exist_ok=True)
    out_path = os.path.join(DATA_DIR, "ccd", "steady_metrics.json")
    with open(out_path, "w") as f:
        json.dump(metrics, f, indent=2)
    print(f"\nMetrics saved to {out_path}")

    # -- 6. Print summary --
    _print_summary(metrics)

    return metrics


def _print_summary(m: dict):
    r = m["recirculation_zone"]
    w = m["wake_restoration"]
    rms = m["rms_suppression"]
    c = m["cancellation_quality"]
    d = m["force_bookkeeping"]

    print("\n" + "=" * 60)
    print("=== Steady Cloak Metrics ===")
    print("=" * 60)

    # Drag bookkeeping
    print(f"Drag proxy (deficit area): blockage={d['drag_proxy_blockage']:.4f}, "
          f"control={d['drag_proxy_control']:.4f}")

    # Fluctuation suppression
    print(f"Fluctuation suppression (global): {rms['suppression_ratio']*100:.1f}%")
    print(f"Fluctuation suppression (sensor zone): {rms['sensor_zone_suppression']*100:.1f}%")

    # Recirculation
    print(f"Recirculation length: {r['recirculation_length_L0']:.2f} (L0 units)")
    print(f"Centreline min ux: {r['centreline_min_ux']:.6f}")

    # Cancellation
    print(f"dq_ctl + dq_blk cancellation ratio: {c['cancellation_ratio']:.4f}")
    print(f"Sensor-zone cancellation ratio: {c['sensor_zone_cancellation_ratio']:.4f}")

    # Wake restoration
    print(f"Sensor-zone deviation from channel (RMS): {w['sensor_zone_deviation_from_channel']:.6f}")
    print(f"Note: {d['note']}")


if __name__ == "__main__":
    run()
