#!/usr/bin/env python3
"""Static Illusion proxy/phase diagnostic without introducing lag features."""
from __future__ import annotations

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

import numpy as np

from SR_analysis.configs import get_scene
from SR_analysis.utils.data_contracts import build_batch_dataset
from SR_analysis.utils.provenance import atomic_write_json, hash_file

SCHEMA_VERSION = "sr-illusion-static-proxy-diagnostic-v1"
FEATURE_GROUPS = {
    "lift_only": ("Cl_F", "Cl_T", "Cl_B", "target_Cl"),
    "drag_only": ("Cd_F", "Cd_T", "Cd_B", "target_Cd"),
    "actual_only": ("Cd_F", "Cl_F", "Cd_T", "Cl_T", "Cd_B", "Cl_B"),
    "target_only": ("target_Cd", "target_Cl"),
    "error_only": ("Cd_err", "Cl_err"),
    "actual_target": (
        "Cd_F", "Cl_F", "Cd_T", "Cl_T", "Cd_B", "Cl_B", "target_Cd", "target_Cl",
    ),
}


def _fit_block(X: np.ndarray, y: np.ndarray, split: int) -> dict[str, Any]:
    train_X = np.column_stack((np.ones(split), X[:split]))
    test_X = np.column_stack((np.ones(len(X) - split), X[split:]))
    coefficients, *_ = np.linalg.lstsq(train_X, y[:split], rcond=None)
    prediction = test_X @ coefficients
    truth = y[split:]
    residual = truth - prediction
    denominator = np.sum((truth - np.mean(truth, axis=0)) ** 2, axis=0)
    r2 = 1.0 - np.sum(residual**2, axis=0) / np.where(denominator > 1e-12, denominator, 1.0)
    return {
        "coefficients": coefficients.tolist(),
        "test_r2": r2.tolist(),
        "test_mae": np.mean(np.abs(residual), axis=0).tolist(),
        "train_indices": [0, split],
        "test_indices": [split, len(X)],
    }


def diagnose(path: Path, *, split_fraction: float = 0.6) -> dict[str, Any]:
    cfg = get_scene("illusion_1L")
    with np.load(path, allow_pickle=False) as data:
        sensors = np.asarray(data["sensors"])
        forces = np.asarray(data["forces"])
        actions = np.asarray(data["actions"])
        target = np.asarray(data["target_forces"])
    all_features = tuple(dict.fromkeys(key for group in FEATURE_GROUPS.values() for key in group))
    dataset = build_batch_dataset(
        sensors, forces, actions, all_features,
        mu=float(cfg["mu"]), dt_c=float(cfg["control_dt"]),
        alignment="historical_post_state_same_action", action_mode="normalized",
        action_scale=float(cfg["action_scale"]), action_bias=cfg["action_bias"],
        target_forces=target, illusion=True, u0=float(cfg["u0"]),
        include_cos_sin=False, warmup=2,
    )
    split = max(2, min(len(dataset.X) - 2, int(len(dataset.X) * split_fraction)))
    index = {name: i for i, name in enumerate(all_features)}
    diagnostics = {
        name: {
            "feature_names": list(features),
            **_fit_block(dataset.X[:, [index[key] for key in features]], dataset.Y, split),
        }
        for name, features in FEATURE_GROUPS.items()
    }
    return {
        "schema_version": SCHEMA_VERSION,
        "scene": "illusion_1L",
        "status": "diagnostic_only",
        "triggered_by": [
            "mapped-shared rear static fit collapsed to a constant",
            "independent rear fit selected drag terms while front selected lift/target-lift terms",
            "PPO empirical G defect is large",
        ],
        "trajectory": {"path": str(path), "sha256": hash_file(path)},
        "action_layout": ["front", "upper", "lower"],
        "ablations": diagnostics,
        "evidence_boundary": (
            "All regressions use current-time static signals only. Differences identify statistical proxy value; "
            "they do not establish transport delay, and no cross-correlation-selected lag is used."
        ),
    }


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--trajectory", type=Path, required=True)
    parser.add_argument("--split-fraction", type=float, default=0.6)
    parser.add_argument("--output", type=Path, required=True)
    return parser


def main(argv: Sequence[str] | None = None) -> int:
    args = build_parser().parse_args(argv)
    if not 0.1 <= args.split_fraction <= 0.9:
        raise ValueError("split fraction must lie in [0.1, 0.9]")
    document = diagnose(args.trajectory.resolve(), split_fraction=args.split_fraction)
    atomic_write_json(args.output.resolve(), document)
    print(json.dumps(document, indent=2, sort_keys=True))
    return 0


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