"""
Export selected output/display.txt fields into browser-friendly JSON files.
"""

from __future__ import annotations

import json
import math
import os
import sys
from pathlib import Path

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import compute


SCRIPT_DIR = Path(__file__).resolve().parent
OUTPUT_DIR = Path(compute.get_output_dir(SCRIPT_DIR))
DISPLAY_TXT = OUTPUT_DIR / "display.txt"
DISPLAY_JSON = OUTPUT_DIR / "display.json"
DISPLAY_JS = OUTPUT_DIR / "display.js"


def rounded_list(values, digits=6):
    rounded = []
    for value in values:
        if isinstance(value, float):
            rounded.append(round(value, digits))
        else:
            rounded.append(value)
    return rounded


def build_payload(arrays):
    disp_t = rounded_list(arrays["disp_t"])
    disp_x = rounded_list(arrays["disp_x"])
    disp_y = rounded_list(arrays["disp_y"])
    disp_z = rounded_list(arrays["disp_z"])
    disp_vx = rounded_list(arrays["disp_vx"])
    disp_vy = rounded_list(arrays["disp_vy"])
    disp_vz = rounded_list(arrays["disp_vz"])

    z_values = [float(value) for value in arrays["disp_z"]]
    speed_values = [
        math.sqrt(
            float(vx) ** 2 + float(vy) ** 2 + float(vz) ** 2
        )
        for vx, vy, vz in zip(arrays["disp_vx"], arrays["disp_vy"], arrays["disp_vz"])
    ]

    return {
        "metadata": {
            "method": arrays["method"],
            "coord_file": arrays["coord_file"],
            "plot_atom_id": int(arrays["plot_atom_id"]),
            "plot_atom_row": int(arrays["plot_atom_row"]),
            "n_frames": int(arrays["n_frames"]),
            "NT": int(arrays["NT"]),
            "DT": float(arrays["DT"]),
            "NSTEP": int(arrays["NSTEP"]),
            "warmup_steps": int(arrays["warmup_steps"]),
            "sample_start": int(arrays["sample_start"]),
            "sample_end": int(arrays["sample_end"]),
            "bounds": {
                "x": [float(arrays["X_MIN"]), float(arrays["X_MAX"])],
                "y": [float(arrays["Y_MIN"]), float(arrays["Y_MAX"])],
                "z": [float(arrays["Z_MIN"]), float(arrays["Z_MAX"])],
            },
            "ball_radius": float(arrays["ball_radius"]),
            "ball_color": [
                float(arrays["ball_color_r"]),
                float(arrays["ball_color_g"]),
                float(arrays["ball_color_b"]),
            ],
            "box_color": [
                float(arrays["box_color_r"]),
                float(arrays["box_color_g"]),
                float(arrays["box_color_b"]),
            ],
            "alpha": float(arrays["alpha"]),
        },
        "series": {
            "t": disp_t,
            "x": disp_x,
            "y": disp_y,
            "z": disp_z,
            "vx": disp_vx,
            "vy": disp_vy,
            "vz": disp_vz,
            "step": [int(value) for value in arrays["disp_step"]],
        },
        "summary": {
            "z_min": round(min(z_values), 6),
            "z_max": round(max(z_values), 6),
            "max_speed": round(max(speed_values), 6),
            "final_position": [
                disp_x[-1],
                disp_y[-1],
                disp_z[-1],
            ],
            "final_velocity": [
                disp_vx[-1],
                disp_vy[-1],
                disp_vz[-1],
            ],
        },
    }


def main():
    OUTPUT_DIR.mkdir(exist_ok=True)
    arrays = compute.load_text_data(DISPLAY_TXT)
    payload = build_payload(arrays)
    json_text = json.dumps(payload, ensure_ascii=False, separators=(",", ":"))
    DISPLAY_JSON.write_text(json_text, encoding="utf-8")
    DISPLAY_JS.write_text(
        f"window.__DYNAMICS_DISPLAY__ = {json_text};\n",
        encoding="utf-8",
    )
    print(f"[export] Wrote {DISPLAY_JSON}")
    print(f"[export] Wrote {DISPLAY_JS}")


if __name__ == "__main__":
    main()