#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
from math import radians
from pathlib import Path

import healpy as hp
import pandas as pd

"""
Minimal bias query tool.

Inputs:
- ra: right ascension in degree.
- dec: declination in degree.
- catalog: catalog key (e.g. ucac4, usno_a2, gsc1_2).
- time: observation time in Julian year (t in the correction formula).

Formula (mas):
  dRA(t)  = dRA0 + (time - 2000.0) * pmRA
  dDEC(t) = dDEC0 + (time - 2000.0) * pmDEC
"""

# ---------------------------------------------------------------------------
# Default demo values (used when no arguments are supplied on the command line)
# ---------------------------------------------------------------------------
_DEFAULT_RA = 83.82  # deg  – random pick near Orion
_DEFAULT_DEC = -5.39  # deg  – random pick near Orion
_DEFAULT_CATALOG = "ucac4"
_DEFAULT_TIME = 2015.5  # Julian year


def ang_to_ipix(ra_deg: float, dec_deg: float, nside: int) -> int:
    theta = radians(90.0 - dec_deg)
    phi = radians(ra_deg % 360.0)
    return int(hp.ang2pix(nside, theta, phi, nest=True))


def query_subtable(
    base: Path, catalog: str, nside: int, ipix: int, time_year: float
) -> dict:
    path = base / "tables" / f"nside{nside}" / f"{catalog}_nside{nside}.csv"
    if not path.exists():
        zip_hint = [
            str(base / "tables_nside64.zip"),
            str(base / "tables_nside256.zip"),
        ]
        raise FileNotFoundError(
            "table directory not found. Please unzip tables_nside64.zip and "
            f"tables_nside256.zip under {base} before running query_bias.py. "
            f"Expected files live under {base / 'tables'}. Zip files: {', '.join(zip_hint)}"
        )

    df = pd.read_csv(path)
    row = df[df["ipix"] == ipix]
    if row.empty:
        raise ValueError(f"ipix={ipix} not found in {path.name}")

    r = row.iloc[0]
    dt = float(time_year) - 2000.0
    dra = float(r["dRA(mas)"]) + dt * float(r["pmRA[mas/yr]"])
    ddec = float(r["dDEC(mas)"]) + dt * float(r["pmDEC[mas/yr]"])

    return {
        "nside": nside,
        "ipix": ipix,
        "dRA_mas": dra,
        "dDEC_mas": ddec,
        "dRA0_mas": float(r["dRA(mas)"]),
        "dDEC0_mas": float(r["dDEC(mas)"]),
        "pmRA_masyr": float(r["pmRA[mas/yr]"]),
        "pmDEC_masyr": float(r["pmDEC[mas/yr]"]),
    }


def main() -> int:
    parser = argparse.ArgumentParser(
        description="Query bias values from release tables.",
        epilog=(
            "time means observation time (Julian year), consistent with t in the bias formula.\n\n"
            "Example (default demo run – no arguments needed):\n"
            f"  python query_bias.py\n"
            f"    => catalog={_DEFAULT_CATALOG}  ra={_DEFAULT_RA}  dec={_DEFAULT_DEC}  time={_DEFAULT_TIME}\n\n"
            "Example (explicit arguments):\n"
            f"  python query_bias.py --catalog {_DEFAULT_CATALOG} --ra {_DEFAULT_RA} --dec {_DEFAULT_DEC} --time {_DEFAULT_TIME}\n"
            "    => same result as above, printed as JSON to stdout."
        ),
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument(
        "--ra",
        type=float,
        default=None,
        help=f"right ascension (deg)  [default: {_DEFAULT_RA}]",
    )
    parser.add_argument(
        "--dec",
        type=float,
        default=None,
        help=f"declination (deg)       [default: {_DEFAULT_DEC}]",
    )
    parser.add_argument(
        "--catalog",
        type=str,
        default=None,
        help=f"catalog key, e.g. ucac4 [default: {_DEFAULT_CATALOG}]",
    )
    parser.add_argument(
        "--time",
        type=float,
        default=None,
        help=f"observation time (Julian year) [default: {_DEFAULT_TIME}]",
    )
    parser.add_argument(
        "--base",
        type=Path,
        default=Path(__file__).resolve().parent,
        help="release root",
    )
    args = parser.parse_args()

    # Fill in defaults and notify the user when falling back to demo values
    using_defaults = False
    if args.ra is None:
        args.ra = _DEFAULT_RA
        using_defaults = True
    if args.dec is None:
        args.dec = _DEFAULT_DEC
        using_defaults = True
    if args.catalog is None:
        args.catalog = _DEFAULT_CATALOG
        using_defaults = True
    if args.time is None:
        args.time = _DEFAULT_TIME
        using_defaults = True

    if using_defaults:
        print(
            f"[INFO] No arguments supplied – running with built-in demo values:\n"
            f"       catalog={args.catalog}  ra={args.ra}  dec={args.dec}  time={args.time}\n"
            f"       (use --help for usage information)\n"
        )

    catalog = args.catalog.strip().lower().replace(".", "_")
    if not (-90.0 <= args.dec <= 90.0):
        raise ValueError("--dec must be in [-90, 90]")

    ipix64 = ang_to_ipix(args.ra, args.dec, 64)
    ipix256 = ang_to_ipix(args.ra, args.dec, 256)

    out = {
        "input": {
            "ra_deg": args.ra,
            "dec_deg": args.dec,
            "catalog": catalog,
            "time": args.time,
        },
        "result": {
            "nside64": query_subtable(args.base, catalog, 64, ipix64, args.time),
            "nside256": query_subtable(args.base, catalog, 256, ipix256, args.time),
        },
    }
    print(json.dumps(out, ensure_ascii=False, indent=2))
    return 0


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