#!/usr/bin/env python3
"""Prepare the real SPARC corpus without executing either UNNS chamber.

Inputs are the three frozen local SPARC files in data/raw_galaxies/SPARC/.
Outputs are the prepared radial table, galaxy/sample table, exclusions,
source and cross-check manifests, and a preparation summary.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import re
import zipfile
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 sha256_file(path: Path) -> str:
    h = hashlib.sha256()
    with path.open("rb") as stream:
        for block in iter(lambda: stream.read(1024 * 1024), b""):
            h.update(block)
    return h.hexdigest()


def parse_catalogue(path: Path) -> pd.DataFrame:
    rows: list[list[str]] = []
    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)
    if len(frame) != 175 or frame["galaxy_id"].nunique() != 175:
        raise ValueError(f"Expected 175 catalogue galaxies, found {len(frame)} rows / {frame['galaxy_id'].nunique()} IDs")
    for column in CATALOGUE_COLUMNS[1:-1]:
        frame[column] = pd.to_numeric(frame[column], errors="raise")
    if not frame["quality_flag"].isin([1, 2, 3]).all():
        raise ValueError("Unexpected SPARC quality flag")
    return frame


def parse_mass_models(path: Path) -> pd.DataFrame:
    rows: list[list[str]] = []
    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")
    if len(frame) != 3391 or frame["galaxy_id"].nunique() != 175:
        raise ValueError(f"Expected 3391 radial rows and 175 galaxies, found {len(frame)} / {frame['galaxy_id'].nunique()}")
    if frame.duplicated(["galaxy_id", "radius_kpc"]).any():
        duplicates = frame.loc[frame.duplicated(["galaxy_id", "radius_kpc"], keep=False), ["galaxy_id", "radius_kpc"]]
        raise ValueError(f"Duplicate galaxy-radius rows detected:\n{duplicates.to_string(index=False)}")
    return frame


def parse_rotmod_member(data: bytes) -> tuple[float, np.ndarray]:
    text = data.decode("utf-8")
    match = re.search(r"Distance\s*=\s*([0-9.]+)", text)
    if not match:
        raise ValueError("Distance header missing")
    distance = float(match.group(1))
    rows = []
    for line in text.splitlines():
        if not line.strip() or line.lstrip().startswith("#"):
            continue
        values = [float(token) for token in line.split()]
        if len(values) != 8:
            raise ValueError(f"Expected 8 data columns, got {len(values)}")
        rows.append(values)
    return distance, np.asarray(rows, dtype=float)


def crosscheck_zip(zip_path: Path, mass: pd.DataFrame) -> pd.DataFrame:
    expected = {f"{galaxy_id}_rotmod.dat" for galaxy_id in mass["galaxy_id"].unique()}
    records = []
    with zipfile.ZipFile(zip_path) as archive:
        actual = {name for name in archive.namelist() if name.endswith("_rotmod.dat")}
        if actual != expected:
            raise ValueError(f"ZIP membership mismatch: missing={sorted(expected-actual)}, extra={sorted(actual-expected)}")
        for galaxy_id, group in mass.groupby("galaxy_id", sort=True):
            member = f"{galaxy_id}_rotmod.dat"
            distance, array = parse_rotmod_member(archive.read(member))
            combined = group[
                ["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"]
            ].to_numpy(dtype=float)
            if array.shape != combined.shape:
                raise ValueError(f"{galaxy_id}: ZIP/combined shape mismatch {array.shape} != {combined.shape}")
            max_abs = float(np.max(np.abs(array - combined))) if array.size else 0.0
            combined_distance = float(group["distance_mpc_massmodel"].iloc[0])
            records.append({
                "galaxy_id": galaxy_id,
                "zip_member": member,
                "rows_zip": int(len(array)),
                "rows_combined": int(len(combined)),
                "distance_zip_mpc": distance,
                "distance_combined_mpc": combined_distance,
                "max_abs_numeric_difference": max_abs,
                "status": "EXACT_MATCH" if max_abs == 0.0 and distance == combined_distance else "MISMATCH",
            })
    result = pd.DataFrame(records)
    if not (result["status"] == "EXACT_MATCH").all():
        raise ValueError("Rotmod ZIP cross-check failed")
    return result


def build_prepared(catalogue: pd.DataFrame, mass: pd.DataFrame) -> tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]:
    frame = mass.merge(catalogue, on="galaxy_id", how="left", validate="many_to_one")
    if frame["quality_flag"].isna().any():
        raise ValueError("Mass-model galaxy missing from catalogue")

    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["g_obs"] = (frame["v_obs_kms"] ** 2 / frame["radius_kpc"]) * ACC_FACTOR
    frame["g_bar"] = (frame["baryonic_velocity_squared_signed_kms2"] / frame["radius_kpc"]) * ACC_FACTOR

    finite = np.ones(len(frame), dtype=bool)
    for column in ["radius_kpc", "v_obs_kms", "v_obs_error_kms", "v_gas_kms",
                   "v_disk_kms", "v_bulge_kms", "g_obs", "g_bar"]:
        finite &= np.isfinite(frame[column].to_numpy(dtype=float))
    frame["valid_point"] = (
        finite & (frame["radius_kpc"] > 0) & (frame["v_obs_kms"] > 0)
        & (frame["g_obs"] > 0)
        & (frame["baryonic_velocity_squared_signed_kms2"] > 0)
        & (frame["g_bar"] > 0)
    )

    reason = np.full(len(frame), "", dtype=object)
    reason[~finite] = "NONFINITE_SOURCE_OR_DERIVED_VALUE"
    reason[finite & (frame["radius_kpc"] <= 0)] = "NONPOSITIVE_RADIUS"
    reason[finite & (frame["v_obs_kms"] <= 0)] = "NONPOSITIVE_VOBS"
    reason[
        finite & (frame["radius_kpc"] > 0) & (frame["v_obs_kms"] > 0)
        & (frame["baryonic_velocity_squared_signed_kms2"] <= 0)
    ] = "NONPOSITIVE_SIGNED_BARYONIC_VELOCITY_SQUARED"
    frame["exclusion_reason"] = reason

    frame["g_obs_ms2"] = frame["g_obs"]
    frame["g_bar_ms2"] = frame["g_bar"]
    frame["log10_g_obs"] = np.nan
    frame["log10_g_bar"] = np.nan
    valid_mask = frame["valid_point"].to_numpy(dtype=bool)
    frame.loc[valid_mask, "log10_g_obs"] = np.log10(frame.loc[valid_mask, "g_obs"])
    frame.loc[valid_mask, "log10_g_bar"] = np.log10(frame.loc[valid_mask, "g_bar"])
    frame["log10_discrepancy"] = frame["log10_g_obs"] - frame["log10_g_bar"]
    frame["radius_normalized"] = frame["radius_kpc"] / frame.groupby("galaxy_id")["radius_kpc"].transform("max")

    raw_counts = frame.groupby("galaxy_id").size().rename("raw_radial_points")
    valid_counts = frame.loc[frame["valid_point"]].groupby("galaxy_id").size().rename("valid_radial_points")
    negative_counts = (frame["v_gas_kms"] < 0).groupby(frame["galaxy_id"]).sum().rename("negative_vgas_points")
    invalid_counts = (~frame["valid_point"]).groupby(frame["galaxy_id"]).sum().rename("invalid_radial_points")

    galaxies = catalogue.merge(raw_counts, on="galaxy_id").merge(valid_counts, on="galaxy_id", how="left")
    galaxies = galaxies.merge(negative_counts, on="galaxy_id").merge(invalid_counts, on="galaxy_id")
    galaxies["valid_radial_points"] = galaxies["valid_radial_points"].fillna(0).astype(int)
    galaxies["is_primary"] = (galaxies["quality_flag"] <= 2) & (galaxies["valid_radial_points"] >= 10)
    galaxies["is_inclination_sensitivity"] = galaxies["is_primary"] & (galaxies["inclination_deg"] >= 30)
    galaxies["is_resolution_sensitivity"] = (galaxies["quality_flag"] <= 2) & (galaxies["valid_radial_points"] >= 16)
    galaxies["is_strict_sensitivity"] = galaxies["is_resolution_sensitivity"] & (galaxies["inclination_deg"] >= 30)
    galaxies["is_broad_q12"] = (galaxies["quality_flag"] <= 2) & (galaxies["valid_radial_points"] >= 3)
    galaxies["is_q3_sensitivity"] = (galaxies["quality_flag"] == 3) & (galaxies["valid_radial_points"] >= 3)

    membership_cols = [
        "galaxy_id", "quality_flag", "inclination_deg", "raw_radial_points",
        "valid_radial_points", "invalid_radial_points", "negative_vgas_points",
        "is_primary", "is_inclination_sensitivity", "is_resolution_sensitivity",
        "is_strict_sensitivity", "is_broad_q12", "is_q3_sensitivity",
    ]
    frame = frame.merge(
        galaxies[membership_cols],
        on=["galaxy_id", "quality_flag", "inclination_deg"],
        how="left",
        validate="many_to_one",
    )

    exclusions = frame.loc[~frame["valid_point"]].copy()
    valid = frame.loc[frame["valid_point"]].copy()
    valid = valid.sort_values(["galaxy_id", "radius_kpc"], kind="mergesort").reset_index(drop=True)
    exclusions = exclusions.sort_values(["galaxy_id", "radius_kpc"], kind="mergesort").reset_index(drop=True)
    galaxies = galaxies.sort_values("galaxy_id").reset_index(drop=True)
    return valid, galaxies, exclusions


def source_manifest(root: Path) -> pd.DataFrame:
    rows = [
        {
            "source_name": "SPARC galaxy catalogue",
            "source_url": "https://astroweb.case.edu/SPARC/SPARC_Lelli2016c.mrt",
            "source_version": "Lelli et al. 2016",
            "source_file": "data/raw_galaxies/SPARC/SPARC_Lelli2016c.mrt",
            "role": "galaxy metadata, quality, inclination, provenance",
        },
        {
            "source_name": "SPARC combined Newtonian mass models",
            "source_url": "https://astroweb.case.edu/SPARC/MassModels_Lelli2016c.mrt",
            "source_version": "Lelli et al. 2016",
            "source_file": "data/raw_galaxies/SPARC/MassModels_Lelli2016c.mrt",
            "role": "radial observed and baryonic velocity components",
        },
        {
            "source_name": "SPARC individual Newtonian mass models",
            "source_url": "https://astroweb.case.edu/SPARC/Rotmod_LTG.zip",
            "source_version": "Lelli et al. 2016",
            "source_file": "data/raw_galaxies/SPARC/Rotmod_LTG.zip",
            "role": "independent exact cross-check of combined table",
        },
    ]
    for row in rows:
        path = root / row["source_file"]
        row["bytes"] = path.stat().st_size
        row["sha256"] = sha256_file(path)
        row["license"] = "Not stated on source page; scientific citation requested"
        row["accessed_date"] = "2026-07-26"
    return pd.DataFrame(rows)


def write_report(path: Path, summary: dict) -> None:
    path.write_text(f"""# SPARC corpus preparation report

**Status:** DATA PREPARED — CHAMBERS NOT RUN

## Source integrity

- Catalogue galaxies: {summary['catalogue_galaxies']}
- Combined mass-model rows: {summary['raw_radial_rows']}
- Individual ZIP members: {summary['rotmod_members']}
- ZIP-to-combined exact matches: {summary['rotmod_exact_matches']}

## Derived corpus

- Valid radial rows: {summary['valid_radial_rows']}
- Excluded radial rows: {summary['excluded_radial_rows']}
- Galaxies with at least one negative `Vgas` value: {summary['galaxies_with_negative_vgas']}
- Negative-`Vgas` radial rows retained with signed-square treatment: {summary['negative_vgas_rows']}
- Rows excluded because signed baryonic velocity squared was non-positive: {summary['nonpositive_signed_baryon_rows']}

## Frozen baryonic calculation

```text
g_obs = (Vobs^2 / R) * 1e6 / kpc_m

g_bar = [Vgas*abs(Vgas) + 0.5*Vdisk^2 + 0.7*Vbul^2]
        / R * 1e6 / kpc_m
```

`Vgas*abs(Vgas)` preserves the direction encoded by negative SPARC gas
contributions. The stellar mass-to-light ratios are fixed, not fitted.

## Sample branches

- PRIMARY (`Q <= 2`, valid points >= 10): {summary['primary_galaxies']}
- INCLINATION_SENSITIVITY (PRIMARY and inclination >= 30 deg): {summary['inclination_sensitivity_galaxies']}
- RESOLUTION_SENSITIVITY (`Q <= 2`, valid points >= 16): {summary['resolution_sensitivity_galaxies']}
- STRICT_SENSITIVITY (resolution branch and inclination >= 30 deg): {summary['strict_sensitivity_galaxies']}
- BROAD_Q12 (`Q <= 2`, valid points >= 3): {summary['broad_q12_galaxies']}
- Q3_SENSITIVITY (`Q = 3`, valid points >= 3): {summary['q3_sensitivity_galaxies']}

## Chamber boundary

This preparation generated chamber input ladders from retained
`log10(g_bar)` values at actual measured radii. It performed no interpolation,
no extrapolation, and no chamber execution.
""", encoding="utf-8")


def parse_args() -> argparse.Namespace:
    root = Path(__file__).resolve().parents[1]
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--project-root", type=Path, default=root)
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    root = args.project_root.resolve()
    raw = root / "data/raw_galaxies/SPARC"
    prepared = root / "data/prepared_profiles"
    prepared.mkdir(parents=True, exist_ok=True)

    catalogue_path = raw / "SPARC_Lelli2016c.mrt"
    mass_path = raw / "MassModels_Lelli2016c.mrt"
    zip_path = raw / "Rotmod_LTG.zip"
    for path in (catalogue_path, mass_path, zip_path):
        if not path.exists():
            raise FileNotFoundError(path)

    catalogue = parse_catalogue(catalogue_path)
    mass = parse_mass_models(mass_path)
    crosscheck = crosscheck_zip(zip_path, mass)
    profiles, galaxies, exclusions = build_prepared(catalogue, mass)

    profiles.to_csv(prepared / "sparc_profiles.csv", index=False, float_format="%.17g")
    profiles.to_csv(prepared / "galaxy_profiles.csv", index=False, float_format="%.17g")
    galaxies.to_csv(prepared / "sparc_galaxies.csv", index=False, float_format="%.17g")
    exclusions.to_csv(prepared / "sparc_point_exclusions.csv", index=False, float_format="%.17g")
    crosscheck.to_csv(prepared / "rotmod_crosscheck.csv", index=False, float_format="%.17g")
    source_manifest(root).to_csv(root / "data/source_manifest.csv", index=False)

    summary = {
        "status": "DATA_PREPARED_CHAMBERS_NOT_RUN",
        "catalogue_galaxies": int(len(catalogue)),
        "raw_radial_rows": int(len(mass)),
        "valid_radial_rows": int(len(profiles)),
        "excluded_radial_rows": int(len(exclusions)),
        "rotmod_members": int(len(crosscheck)),
        "rotmod_exact_matches": int((crosscheck["status"] == "EXACT_MATCH").sum()),
        "negative_vgas_rows": int((mass["v_gas_kms"] < 0).sum()),
        "galaxies_with_negative_vgas": int(mass.loc[mass["v_gas_kms"] < 0, "galaxy_id"].nunique()),
        "nonpositive_signed_baryon_rows": int((exclusions["exclusion_reason"] == "NONPOSITIVE_SIGNED_BARYONIC_VELOCITY_SQUARED").sum()),
        "primary_galaxies": int(galaxies["is_primary"].sum()),
        "inclination_sensitivity_galaxies": int(galaxies["is_inclination_sensitivity"].sum()),
        "resolution_sensitivity_galaxies": int(galaxies["is_resolution_sensitivity"].sum()),
        "strict_sensitivity_galaxies": int(galaxies["is_strict_sensitivity"].sum()),
        "broad_q12_galaxies": int(galaxies["is_broad_q12"].sum()),
        "q3_sensitivity_galaxies": int(galaxies["is_q3_sensitivity"].sum()),
        "disk_mass_to_light": UPSILON_DISK,
        "bulge_mass_to_light": UPSILON_BULGE,
        "acceleration_factor_mps2_per_kms2_per_kpc": ACC_FACTOR,
        "chambers_run": False,
    }
    (prepared / "sparc_preparation_summary.json").write_text(json.dumps(summary, indent=2) + "\n", encoding="utf-8")
    write_report(prepared / "PREPARATION_REPORT.md", summary)
    print(json.dumps(summary, indent=2))
    return 0


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