#!/usr/bin/env python
"""Parse DOLPHOT's native ASCII .phot output into an Astropy Table.

Per the `.columns` sidecar, the fixed layout is:
  cols  1-12: geometric/combined measurements (Ext, Chip, X, Y, Chi, SN, ...)
  cols 13-25: per-filter combined measurements (one block per filter)
  cols  26+: per-exposure measurements (13 columns each)

For the Sickle benchmark each DOLPHOT run is a single filter (all 24 NRCB
exposures of one filter) so there is exactly ONE combined-filter block and
N per-exposure blocks. This loader extracts the combined-filter block as
the per-filter catalog, converts pixel coords to RA/Dec via the reference
image's WCS, and writes a FITS table with a normalised schema.

Output column schema (matches bandmerge.load_tool_catalog expectations):
  ra, dec       float64  degrees
  x_ref, y_ref  float64  pixel coords on reference image (1-indexed DOLPHOT)
  flux          float64  combined count rate (col 15, "Normalized count rate")
  flux_err      float64  combined count rate uncertainty (col 16)
  mag           float64  VEGAMAG (col 17)
  magerr        float64  (col 19)
  snr           float64  combined S/N (col 6)
  chi           float64  combined Chi (col 5)
  sharpness     float64  col 7
  roundness     float64  col 8
  crowding      float64  col 10
  object_type   int      col 11 (1=bright, 2=faint, 3=elongated, 4=hot, 5=ext)
  qfit          float64  = col 5 (Chi — closest analogue to qfit)
"""
from __future__ import annotations

import argparse
import re
from pathlib import Path

import numpy as np
from astropy.io import fits
from astropy.table import Table
from astropy.wcs import WCS


def parse_columns_file(columns_path: Path) -> list[str]:
    descs = []
    with open(columns_path) as fh:
        for line in fh:
            m = re.match(r"^\s*(\d+)[.)]?\s+(.*)$", line.strip())
            if m:
                descs.append(m.group(2).strip())
    return descs


def load_dolphot_phot(phot_path: Path, ref_fits: Path) -> Table:
    phot_path = Path(phot_path)
    cols_path = Path(str(phot_path) + ".columns")
    descs = parse_columns_file(cols_path)
    if len(descs) < 25:
        raise ValueError(f"unexpected columns file: only {len(descs)} rows")

    # Infer filter from col 13's description (e.g. "Total counts, NIRCAM_F480M")
    m = re.search(r"NIRCAM_(F\d+[A-Z])$", descs[12])
    filt = m.group(1) if m else None

    arr = np.loadtxt(phot_path)
    if arr.ndim == 1:
        arr = arr[None, :]
    n = arr.shape[0]

    # DOLPHOT column indices are 1-based in the columns file.
    # Numpy slicing below is 0-based.
    col = lambda k: arr[:, k - 1]

    tab = Table()
    tab["extension"] = col(1).astype(int)
    tab["chip"] = col(2).astype(int)
    tab["x_ref"] = col(3)
    tab["y_ref"] = col(4)
    tab["chi"] = col(5)
    tab["snr"] = col(6)
    tab["sharpness"] = col(7)
    tab["roundness"] = col(8)
    tab["crowding"] = col(10)
    tab["object_type"] = col(11).astype(int)
    tab["pass_detected"] = col(12).astype(int)
    # Combined-filter block (cols 13-25)
    tab["total_counts"] = col(13)
    tab["sky"] = col(14)
    tab["flux"] = col(15)       # normalised count rate
    tab["flux_err"] = col(16)   # normalised count rate uncertainty
    tab["mag"] = col(17)        # instrumental VEGAMAG
    tab["magerr"] = col(19)
    tab["quality_flag"] = col(25).astype(int)
    tab["qfit"] = tab["chi"]    # alias Chi -> qfit for the unified schema

    # WCS from the reference image's SCI extension
    with fits.open(ref_fits) as h:
        ext = "SCI" if "SCI" in h else 1
        wcs = WCS(h[ext].header)
    ra, dec = wcs.all_pix2world(tab["x_ref"] - 1.0, tab["y_ref"] - 1.0, 0)
    tab["ra"] = ra
    tab["dec"] = dec

    tab.meta["TOOL"] = "dolphot"
    tab.meta["FILTER"] = filt or ""
    tab.meta["NROWS"] = n
    tab.meta["REFFITS"] = str(ref_fits)
    return tab


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("phot_file", help="e.g. benchmark/dolphot/work/F480M/f480m_nrcb")
    ap.add_argument("--ref-fits", required=True, help="reference image (img0) from which WCS is read")
    ap.add_argument("-o", "--output", default=None)
    args = ap.parse_args()
    phot = Path(args.phot_file)
    tab = load_dolphot_phot(phot, Path(args.ref_fits))
    print(f"Loaded {len(tab)} rows from {phot}  (filter={tab.meta['FILTER']})")
    out = Path(args.output) if args.output else phot.parent / f"{tab.meta['FILTER'].lower()}_dolphot_nrcb.fits"
    tab.write(out, overwrite=True)
    print(f"Wrote {out}")


if __name__ == "__main__":
    main()