#!/usr/bin/env python3
"""Rebuild the neutral baryonic-profile corpus and chamber ladders from the three SPARC source files."""

from __future__ import annotations

from pathlib import Path
import numpy as np
import pandas as pd

KPC_M = 3.0856775814913673e19
ACC_FACTOR = 1.0e6 / KPC_M
UPSILON_DISK = 0.5
UPSILON_BULGE = 0.7

CATALOGUE_COLUMNS = [
    "galaxy_id", "hubble_type", "distance_mpc", "distance_error_mpc",
    "distance_method", "inclination_deg", "inclination_error_deg",
    "luminosity_3p6_1e9Lsun", "luminosity_error_3p6_1e9Lsun",
    "effective_radius_kpc", "effective_surface_brightness_Lsun_pc2",
    "disk_scale_length_kpc", "disk_central_surface_brightness_Lsun_pc2",
    "MHI_1e9Msun", "HI_radius_kpc", "vflat_kms", "vflat_error_kms",
    "quality_flag", "references",
]
MASS_COLUMNS = [
    "galaxy_id", "distance_mpc_massmodel", "radius_kpc", "v_obs_kms",
    "v_obs_error_kms", "v_gas_kms", "v_disk_kms", "v_bulge_kms",
    "sb_disk_Lsun_pc2", "sb_bulge_Lsun_pc2",
]

def parse_catalogue(path: Path) -> pd.DataFrame:
    rows = []
    for line in path.read_text(encoding="utf-8").splitlines():
        tokens = line.split()
        if len(tokens) < 19:
            continue
        try:
            int(tokens[1]); float(tokens[2]); int(tokens[17])
        except ValueError:
            continue
        rows.append(tokens[:18] + [" ".join(tokens[18:])])
    frame = pd.DataFrame(rows, columns=CATALOGUE_COLUMNS)
    for column in CATALOGUE_COLUMNS[1:18]:
        if column != "distance_method":
            frame[column] = pd.to_numeric(frame[column], errors="raise")
    return frame

def parse_mass_models(path: Path) -> pd.DataFrame:
    rows = []
    for line in path.read_text(encoding="utf-8").splitlines():
        tokens = line.split()
        if len(tokens) != 10:
            continue
        try:
            [float(value) for value in tokens[1:]]
        except ValueError:
            continue
        rows.append(tokens)
    frame = pd.DataFrame(rows, columns=MASS_COLUMNS)
    for column in MASS_COLUMNS[1:]:
        frame[column] = pd.to_numeric(frame[column], errors="raise")
    return frame

def main() -> int:
    root = Path(__file__).resolve().parents[1]
    raw = root / "data/raw/SPARC"
    prepared = root / "data/prepared"
    ladders_i = root / "inputs/ladders/struc_i"
    ladders_p = root / "inputs/ladders/struc_perc_i"
    prepared.mkdir(parents=True, exist_ok=True)
    ladders_i.mkdir(parents=True, exist_ok=True)
    ladders_p.mkdir(parents=True, exist_ok=True)

    catalogue = parse_catalogue(raw / "SPARC_Lelli2016c.mrt")
    mass = parse_mass_models(raw / "MassModels_Lelli2016c.mrt")
    frame = mass.merge(catalogue, on="galaxy_id", how="left", validate="many_to_one")
    frame["gas_velocity_squared_signed_kms2"] = frame["v_gas_kms"] * frame["v_gas_kms"].abs()
    frame["disk_velocity_squared_scaled_kms2"] = UPSILON_DISK * frame["v_disk_kms"] ** 2
    frame["bulge_velocity_squared_scaled_kms2"] = UPSILON_BULGE * frame["v_bulge_kms"] ** 2
    frame["baryonic_velocity_squared_signed_kms2"] = (
        frame["gas_velocity_squared_signed_kms2"]
        + frame["disk_velocity_squared_scaled_kms2"]
        + frame["bulge_velocity_squared_scaled_kms2"]
    )
    frame["baryonic_acceleration_ms2"] = (
        frame["baryonic_velocity_squared_signed_kms2"] / frame["radius_kpc"]
    ) * ACC_FACTOR
    finite = np.isfinite(frame["baryonic_acceleration_ms2"]) & np.isfinite(frame["radius_kpc"])
    frame["valid_point"] = (
        finite & (frame["radius_kpc"] > 0)
        & (frame["baryonic_velocity_squared_signed_kms2"] > 0)
        & (frame["baryonic_acceleration_ms2"] > 0)
    )
    frame["log10_baryonic_acceleration"] = np.nan
    frame.loc[frame["valid_point"], "log10_baryonic_acceleration"] = np.log10(
        frame.loc[frame["valid_point"], "baryonic_acceleration_ms2"]
    )
    frame["radius_normalized"] = (
        frame["radius_kpc"] / frame.groupby("galaxy_id")["radius_kpc"].transform("max")
    )
    frame.to_csv(prepared / "baryonic_profiles_rebuilt.csv", index=False)
    catalogue.to_csv(prepared / "galaxy_catalog_rebuilt.csv", index=False)

    for galaxy_id, group in frame.loc[frame["valid_point"]].groupby("galaxy_id"):
        values = np.unique(np.sort(group["log10_baryonic_acceleration"].to_numpy(float)))
        pd.DataFrame({"value": values}).to_csv(ladders_i / f"{galaxy_id}.csv", index=False)
        (ladders_p / f"{galaxy_id}.txt").write_text(
            "\n".join(f"{value:.17g}" for value in values) + "\n", encoding="utf-8"
        )
    print(f"Prepared {frame['galaxy_id'].nunique()} galaxies.")
    return 0

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