#!/usr/bin/env python3
"""Fit blackbody temperature and apparent brightness to a 2MASS PSC table.

The output is deliberately a small renderer-facing CSV: RA, Dec, blackbody
temperature, and blackbody normalization.  The latter is the fitted apparent
solid angle Omega in F_nu = Omega B_nu(T), expressed in steradians.
"""

from __future__ import annotations

import csv
import argparse
import math
from pathlib import Path

import numpy as np

ROOT = Path(__file__).resolve().parent.parent
INPUT = ROOT / "assets/2mass/raw/2mass_psc_m31_0p5deg.tbl"
OUTPUT = ROOT / "assets/2mass/processed/2mass_psc_m31_0p5deg_stars.csv"

# 2MASS effective wavelengths and Vega zero-magnitude flux densities.
WAVELENGTH_M = np.array([1.235, 1.662, 2.159]) * 1e-6
ZERO_POINT_JY = np.array([1594.0, 1024.0, 666.7])
GOOD_RD_FLAGS = frozenset("123")
MISSING_NUMERIC_TOKENS = frozenset(("", "null"))
MIN_TEMPERATURE_K = 300.0
MAX_TEMPERATURE_K = 100000.0

PLANCK_H = 6.62607015e-34
LIGHT_C = 299792458.0
BOLTZMANN_K = 1.380649e-23


def column_slices(path: Path) -> tuple[list[str], list[tuple[int, int]], int]:
    """Return IPAC fixed-width column metadata and first data-line index."""
    lines = path.read_text(encoding="ascii").splitlines()
    for index, line in enumerate(lines):
        if line.startswith("|") and "designation" in line:
            boundaries = [offset for offset, char in enumerate(line) if char == "|"]
            names = [line[left + 1:right].strip()
                     for left, right in zip(boundaries, boundaries[1:])]
            slices = [(left + 1, right)
                      for left, right in zip(boundaries, boundaries[1:])]
            return names, slices, index + 4
    raise ValueError(f"no IPAC table header in {path}")


def load_required_columns(path: Path) -> dict[str, np.ndarray]:
    names, slices, first_data_line = column_slices(path)
    required = ("ra", "dec", "j_m", "h_m", "k_m", "rd_flg")
    indices = {name: names.index(name) for name in required}
    values: dict[str, list[str]] = {name: [] for name in required}

    with path.open(encoding="ascii") as table:
        for line_number, line in enumerate(table):
            if line_number < first_data_line or not line.strip():
                continue
            for name, index in indices.items():
                left, right = slices[index]
                values[name].append(line[left:right].strip())

    result: dict[str, np.ndarray] = {}
    for name in ("ra", "dec", "j_m", "h_m", "k_m"):
        result[name] = np.array(
            [float(value) if value.lower() not in MISSING_NUMERIC_TOKENS else math.nan
             for value in values[name]], dtype=float
        )
    result["rd_flg"] = np.array(values["rd_flg"], dtype="U3")
    return result


def log_planck_nu_jy_per_sr(log_temperature: np.ndarray) -> np.ndarray:
    """Evaluate log B_nu in Jy/sr at the three 2MASS effective wavelengths."""
    temperature = np.exp(log_temperature)[:, None]
    frequency = LIGHT_C / WAVELENGTH_M
    exponent = PLANCK_H * frequency / (BOLTZMANN_K * temperature)
    radiance_si = 2.0 * PLANCK_H * frequency**3 / LIGHT_C**2 / np.expm1(exponent)
    return np.log(radiance_si / 1e-26)


def fit_blackbody(log_flux_jy: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    """Unweighted least-squares fit in log F_nu for T and apparent solid angle."""
    count = len(log_flux_jy)
    lower = np.full(count, math.log(MIN_TEMPERATURE_K))
    upper = np.full(count, math.log(MAX_TEMPERATURE_K))
    golden = (math.sqrt(5.0) - 1.0) / 2.0
    first = upper - golden * (upper - lower)
    second = lower + golden * (upper - lower)

    def objective(log_temperature: np.ndarray) -> np.ndarray:
        model = log_planck_nu_jy_per_sr(log_temperature)
        residual = (log_flux_jy - log_flux_jy.mean(axis=1, keepdims=True)
                    - (model - model.mean(axis=1, keepdims=True)))
        return np.sum(residual * residual, axis=1)

    first_value = objective(first)
    second_value = objective(second)
    for _ in range(64):
        keep_left = first_value <= second_value
        upper = np.where(keep_left, second, upper)
        second = np.where(keep_left, first, second)
        second_value = np.where(keep_left, first_value, second_value)
        lower = np.where(keep_left, lower, first)
        first = np.where(keep_left, upper - golden * (upper - lower), second)
        first_value = np.where(keep_left, objective(first), second_value)
        second = np.where(keep_left, second, lower + golden * (upper - lower))
        second_value = np.where(keep_left, second_value, objective(second))

    log_temperature = 0.5 * (lower + upper)
    log_radiance = log_planck_nu_jy_per_sr(log_temperature)
    log_solid_angle = np.mean(log_flux_jy - log_radiance, axis=1)
    return np.exp(log_temperature), np.exp(log_solid_angle)


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Fit a renderer-facing blackbody catalog from a 2MASS PSC IPAC table."
    )
    parser.add_argument("--input", type=Path, default=INPUT,
                        help="raw 2MASS PSC IPAC table")
    parser.add_argument("--output", type=Path, default=OUTPUT,
                        help="renderer-facing CSV to write")
    args = parser.parse_args()
    columns = load_required_columns(args.input)
    magnitudes = np.column_stack((columns["j_m"], columns["h_m"], columns["k_m"]))
    valid_photometry = np.isfinite(magnitudes).all(axis=1)
    valid_position = (np.isfinite(columns["ra"]) & np.isfinite(columns["dec"]) &
                      (columns["ra"] >= 0.0) & (columns["ra"] < 360.0) &
                      (columns["dec"] >= -90.0) & (columns["dec"] <= 90.0))
    valid_rd_flag = np.array(
        [len(flag) == 3 and all(value in GOOD_RD_FLAGS for value in flag)
         for flag in columns["rd_flg"]]
    )
    keep = valid_position & valid_photometry & valid_rd_flag
    if not np.any(keep):
        raise ValueError("no sources retain valid J/H/Ks photometry")

    flux_jy = ZERO_POINT_JY * np.power(10.0, -0.4 * magnitudes[keep])
    temperature_k, amplitude_sr = fit_blackbody(np.log(flux_jy))

    args.output.parent.mkdir(parents=True, exist_ok=True)
    with args.output.open("w", newline="", encoding="ascii") as output:
        writer = csv.writer(output, lineterminator="\n")
        writer.writerow(("ra_deg", "dec_deg", "temperature_K", "amplitude_sr"))
        for ra, dec, temperature, amplitude in zip(
            columns["ra"][keep], columns["dec"][keep], temperature_k, amplitude_sr
        ):
            writer.writerow((f"{ra:.6f}", f"{dec:.6f}", f"{temperature:.8g}",
                             f"{amplitude:.9e}"))

    print(f"input_rows={len(keep)}")
    print(f"discarded_invalid_position={np.count_nonzero(~valid_position)}")
    print(f"discarded_invalid_photometry={np.count_nonzero(~valid_photometry)}")
    print(f"discarded_rd_flg_not_123={np.count_nonzero(~valid_rd_flag)}")
    print(f"retained_rows={np.count_nonzero(keep)}")
    print(f"wrote={args.output}")


if __name__ == "__main__":
    main()