#!/usr/bin/env python3
"""Export Stage 1 or Stage 3 trajectories to analysis-ready CSV files."""
from __future__ import annotations

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

import numpy as np

from SR_analysis.utils.metrics import legacy_cycle_similarity

SENSOR_NAMES = ("upper_ux", "upper_uy", "center_ux", "center_uy", "lower_ux", "lower_uy")
FORCE_NAMES = ("front_fx", "front_fy", "upper_fx", "upper_fy", "lower_fx", "lower_fy")
ACTION_NAMES = ("front", "upper", "lower")


def load_trajectory(path: Path) -> dict[str, np.ndarray]:
    with np.load(path, allow_pickle=False) as data:
        return {name: np.asarray(data[name]) for name in data.files}


def canonical_arrays(data: Mapping[str, np.ndarray], source: str) -> dict[str, np.ndarray]:
    result: dict[str, np.ndarray] = {}
    sensors_key = "sensors_post" if source == "ppo" and "sensors_post" in data else "sensors"
    forces_key = "forces_post" if source == "ppo" and "forces_post" in data else "forces"
    if sensors_key in data:
        result["sensors"] = np.asarray(data[sensors_key], dtype=float).reshape(-1, 6)
    if forces_key in data:
        result["forces"] = np.asarray(data[forces_key], dtype=float).reshape(-1, 6)
    if source == "ppo":
        if "actions_alpha_commanded" in data:
            result["actions_alpha"] = np.asarray(data["actions_alpha_commanded"], dtype=float).reshape(-1, 3)
        elif "actions_alpha" in data:
            result["actions_alpha"] = np.asarray(data["actions_alpha"], dtype=float).reshape(-1, 3)
        elif "actions" in data:
            result["actions_alpha"] = np.asarray(data["actions"], dtype=float).reshape(-1, 3)
        if "actions_omega_commanded" in data:
            result["actions_omega"] = np.asarray(data["actions_omega_commanded"], dtype=float).reshape(-1, 3)
        elif "actions_omega" in data:
            result["actions_omega"] = np.asarray(data["actions_omega"], dtype=float).reshape(-1, 3)
    else:
        for key in ("actions_alpha", "actions_omega"):
            if key in data:
                result[key] = np.asarray(data[key], dtype=float).reshape(-1, 3)
    for key in ("target_sensors", "target_forces", "target_phase"):
        if key in data and np.asarray(data[key]).size:
            width = 6 if key == "target_sensors" else 2 if key == "target_forces" else 1
            result[key] = np.asarray(data[key]).reshape(-1, width)
    if "target_states" in data:
        target_states = np.asarray(data["target_states"], dtype=float)
        if target_states.shape[1] == 6:
            result["sensors"] = target_states
        elif target_states.shape[1] == 8:
            result["target_forces"] = target_states[:, :2]
            result["sensors"] = target_states[:, 2:8]
        else:
            raise ValueError(f"unsupported target_states width: {target_states.shape[1]}")
    return result


def time_coordinate(n: int, sample_interval: int, u0: float = 0.01, diameter: float = 20.0) -> np.ndarray:
    return np.arange(n, dtype=float) * sample_interval * u0 / diameter


def write_wide_csv(path: Path, arrays: Mapping[str, np.ndarray], scene: str, source: str, sample_interval: int) -> None:
    n = max((len(value) for key, value in arrays.items() if key != "target_sensors"), default=0)
    columns: list[tuple[str, np.ndarray]] = []
    for group, names in (("sensors", SENSOR_NAMES), ("forces", FORCE_NAMES), ("actions_alpha", ACTION_NAMES), ("actions_omega", ACTION_NAMES), ("target_forces", ("target_cd", "target_cl")), ("target_phase", ("target_phase",))):
        if group in arrays:
            values = np.asarray(arrays[group]).reshape(len(arrays[group]), -1)
            columns.extend((f"{group}_{name}", values[:, index]) for index, name in enumerate(names))
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", newline="", encoding="utf-8") as handle:
        writer = csv.writer(handle)
        writer.writerow(("scene", "source", "step", "t_D", *(name for name, _ in columns)))
        times = time_coordinate(n, sample_interval)
        for index in range(n):
            writer.writerow((scene, source, index, times[index], *(values[index] if index < len(values) else "" for _, values in columns)))


def write_long_csv(path: Path, arrays: Mapping[str, np.ndarray], scene: str, source: str, sample_interval: int) -> None:
    groups = (("sensors", SENSOR_NAMES, "velocity"), ("forces", FORCE_NAMES, "force"), ("actions_alpha", ACTION_NAMES, "alpha"), ("actions_omega", ACTION_NAMES, "omega"), ("target_forces", ("target_cd", "target_cl"), "target_force"))
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", newline="", encoding="utf-8") as handle:
        writer = csv.writer(handle)
        writer.writerow(("scene", "source", "step", "t_D", "channel_group", "channel_name", "value"))
        for key, names, group in groups:
            if key not in arrays:
                continue
            values = np.asarray(arrays[key]).reshape(len(arrays[key]), -1)
            times = time_coordinate(len(values), sample_interval)
            for step, row in enumerate(values):
                for channel, value in zip(names, row):
                    writer.writerow((scene, source, step, times[step], group, channel, float(value)))


def legacy_similarity(target: np.ndarray, state: np.ndarray, conv_len: int) -> dict[str, Any]:
    metric = legacy_cycle_similarity(target, state, conv_len)
    return {
        "similarity": metric["similarity"],
        "lag": metric["lag"],
        "per_channel": [metric["per_channel"][str(index)] for index in range(6)],
    }


def write_convergence_csv(path: Path, sensors: np.ndarray, target: np.ndarray, conv_len: int, scene: str, source: str, sample_interval: int) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    endpoints = list(range(max(2 * conv_len, conv_len), len(sensors) + 1, conv_len))
    if not endpoints or endpoints[-1] != len(sensors):
        endpoints.append(len(sensors))
    with path.open("w", newline="", encoding="utf-8") as handle:
        writer = csv.writer(handle)
        writer.writerow(("scene", "source", "end_step", "t_D", "conv_len", "similarity", "lag", *(f"channel_{i}" for i in range(6))))
        for end_step in endpoints:
            if end_step < conv_len or len(target) < 2 * conv_len:
                continue
            metric = legacy_similarity(target, sensors[:end_step], conv_len)
            writer.writerow((scene, source, end_step, end_step * sample_interval * 0.01 / 20.0, conv_len, metric["similarity"], metric["lag"], *metric["per_channel"]))


def export(path: Path, output_dir: Path, scene: str, source: str, sample_interval: int, conv_len: int, target_path: Path | None = None) -> dict[str, Any]:
    arrays = canonical_arrays(load_trajectory(path), source)
    if target_path is not None:
        target_data = canonical_arrays(load_trajectory(target_path), "target")
        if "sensors" in target_data:
            arrays["target_sensors"] = target_data["sensors"]
    write_wide_csv(output_dir / f"{source}_wide.csv", arrays, scene, source, sample_interval)
    write_long_csv(output_dir / f"{source}_long.csv", arrays, scene, source, sample_interval)
    if "target_sensors" in arrays:
        target_arrays = {"sensors": arrays["target_sensors"]}
        write_wide_csv(output_dir / "target_wide.csv", target_arrays, scene, "target", sample_interval)
        write_long_csv(output_dir / "target_long.csv", target_arrays, scene, "target", sample_interval)
        if "sensors" in arrays:
            write_convergence_csv(output_dir / f"{source}_dtw_convergence.csv", arrays["sensors"], arrays["target_sensors"], conv_len, scene, source, sample_interval)
    summary = {"scene": scene, "source": source, "input": str(path), "sample_interval": sample_interval, "conv_len": conv_len, "rows": {key: len(value) for key, value in arrays.items()}}
    (output_dir / f"{source}_manifest.json").write_text(json.dumps(summary, indent=2) + "\n", encoding="utf-8")
    return summary


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--input", type=Path, required=True)
    parser.add_argument("--output-dir", type=Path, required=True)
    parser.add_argument("--scene", required=True)
    parser.add_argument("--source", choices=("ppo", "sr"), required=True)
    parser.add_argument("--sample-interval", type=int, required=True)
    parser.add_argument("--conv-len", type=int, required=True)
    parser.add_argument("--target", type=Path)
    return parser


def main(argv: Sequence[str] | None = None) -> int:
    args = build_parser().parse_args(argv)
    export(args.input, args.output_dir, args.scene, args.source, args.sample_interval, args.conv_len, args.target)
    return 0


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