"""Pass 3: plate solve the luminance master and copy the WCS to every master.

iTelescope's calibrated frames arrive with a PinPoint HISTORY line but no WCS
keywords at all, so the astrometry has to be redone locally. A blind solve is
not needed: the header gives the pointing to arcminutes and the plate scale to
four figures, so this fetches a Gaia DR3 catalogue for that patch of sky and
matches it to the detected stars.

The match itself is asterism-based (astroalign), which is invariant to rotation
and scale, so the roll angle never has to be guessed. It is NOT invariant to a
mirror flip, so both parities are tried and the one that matches wins. The
final WCS is a least-squares TAN fit to the matched pairs, and the residual it
reports is the honest measure of whether the solve is real.
"""
import os
import sys

import astroalign as aa
import numpy as np
import sep
from astropy import units as u
from astropy.coordinates import SkyCoord
from astropy.io import fits
from astropy.wcs import WCS
from astropy.wcs.utils import fit_wcs_from_points

import layout

SRC = layout.SESSION
OUT = layout.SESSION
MASTERS = ["Luminance", "Red", "Green", "Blue"]
CATCACHE = layout.path("_gaia.npz")


def detect(image, nmax=600):
    bkg = sep.Background(image, bw=64, bh=64, fw=3, fh=3)
    sub = image - bkg.back()
    objs = sep.extract(sub, 8.0, err=bkg.globalrms, minarea=9,
                       deblend_cont=0.005)
    objs = objs[(objs["flag"] == 0) & (objs["npix"] > 12) &
                (objs["npix"] < 3000)]
    objs = objs[np.argsort(objs["flux"])[::-1][:nmax]]
    return np.column_stack([objs["x"], objs["y"]]), objs["flux"]


def gaia_catalogue(ra_deg, dec_deg, radius_deg, nmax=600):
    """Gaia DR3 sources around the pointing, brightest first, cached to disk."""
    if os.path.exists(CATCACHE):
        z = np.load(CATCACHE)
        print(f"catalogue: {len(z['ra'])} cached Gaia sources")
        return z["ra"], z["dec"], z["g"]
    from astroquery.gaia import Gaia

    Gaia.ROW_LIMIT = nmax
    query = f"""
        SELECT TOP {nmax} ra, dec, phot_g_mean_mag
        FROM gaiadr3.gaia_source
        WHERE 1 = CONTAINS(POINT('ICRS', ra, dec),
                           CIRCLE('ICRS', {ra_deg}, {dec_deg}, {radius_deg}))
          AND phot_g_mean_mag IS NOT NULL
        ORDER BY phot_g_mean_mag ASC
    """
    tbl = Gaia.launch_job_async(query).get_results()
    ra = np.asarray(tbl["ra"], dtype=float)
    dec = np.asarray(tbl["dec"], dtype=float)
    g = np.asarray(tbl["phot_g_mean_mag"], dtype=float)
    np.savez_compressed(CATCACHE, ra=ra, dec=dec, g=g)
    print(f"catalogue: {len(ra)} Gaia DR3 sources, G {g.min():.1f}-{g.max():.1f}")
    return ra, dec, g


def project(ra, dec, ra0, dec0, scale_arcsec, parity):
    """Gnomonic projection to pixel-like coordinates for asterism matching."""
    c = SkyCoord(ra * u.deg, dec * u.deg)
    centre = SkyCoord(ra0 * u.deg, dec0 * u.deg)
    dx, dy = centre.spherical_offsets_to(c)
    x = dx.to_value(u.arcsec) / scale_arcsec * parity
    y = dy.to_value(u.arcsec) / scale_arcsec
    return np.column_stack([x, y])


def main():
    path = layout.path("master-Luminance.fit")
    with fits.open(path) as hd:
        image = hd[0].data.astype(np.float32)
        hdr = hd[0].header
    ny, nx = image.shape

    centre = SkyCoord(hdr["OBJCTRA"], hdr["OBJCTDEC"],
                      unit=(u.hourangle, u.deg))
    scale = float(hdr["HIERARCH iTelescopePlateScaleH"])
    radius = 1.15 * 0.5 * np.hypot(nx, ny) * scale / 3600.0
    print(f"pointing {centre.to_string('hmsdms')}  scale {scale:.4f}\"/px  "
          f"search radius {radius:.3f} deg")

    xy, flux = detect(image)
    print(f"detected {len(xy)} stars in the luminance master")

    ra, dec, gmag = gaia_catalogue(centre.ra.deg, centre.dec.deg, radius)

    best = None
    for parity in (-1.0, 1.0):
        cat_xy = project(ra, dec, centre.ra.deg, centre.dec.deg, scale, parity)
        try:
            tform, (src, dst) = aa.find_transform(cat_xy, xy)
        except Exception as exc:                                  # noqa: BLE001
            print(f"  parity {parity:+.0f}: no match ({exc})")
            continue
        print(f"  parity {parity:+.0f}: matched {len(src)} stars, "
              f"rotation {np.degrees(tform.rotation):.3f} deg, "
              f"scale {tform.scale:.5f}")
        if best is None or len(src) > best[0]:
            best = (len(src), parity, src, dst)
    if best is None:
        sys.exit("plate solve failed: no asterism match in either parity")

    nmatch, parity, cat_pts, img_pts = best
    # Recover which catalogue rows were matched so the fit uses sky coordinates
    # rather than the projected proxy.
    cat_xy = project(ra, dec, centre.ra.deg, centre.dec.deg, scale, parity)
    idx = [int(np.argmin(np.hypot(cat_xy[:, 0] - px, cat_xy[:, 1] - py)))
           for px, py in cat_pts]
    world = SkyCoord(ra[idx] * u.deg, dec[idx] * u.deg)

    wcs = fit_wcs_from_points((img_pts[:, 0], img_pts[:, 1]), world,
                              proj_point="center", projection="TAN")
    pred = wcs.world_to_pixel(world)
    resid = np.hypot(pred[0] - img_pts[:, 0], pred[1] - img_pts[:, 1])
    print(f"seed fit on {nmatch} stars: residual median {np.median(resid):.2f} px "
          f"({np.median(resid) * scale:.2f}\"), max {resid.max():.2f} px")

    # The asterism match only ever returns a handful of stars. Now that an
    # approximate solution exists, every catalogue source can be pushed through
    # it and paired with the nearest detection, which grows the fit from a
    # dozen stars to hundreds and averages down the centroid noise. Two passes
    # with a shrinking tolerance is enough to converge.
    all_world = SkyCoord(ra * u.deg, dec * u.deg)
    for tol in (4.0, 2.0):
        px, py = wcs.world_to_pixel(all_world)
        pairs = []
        for i, (cx, cy) in enumerate(zip(px, py)):
            if not (0 <= cx < nx and 0 <= cy < ny):
                continue
            d = np.hypot(xy[:, 0] - cx, xy[:, 1] - cy)
            j = int(np.argmin(d))
            if d[j] <= tol:
                pairs.append((i, j))
        if len(pairs) < 20:
            print(f"  refine (tol {tol} px): only {len(pairs)} pairs, kept seed")
            break
        ci = np.array([p[0] for p in pairs])
        ii = np.array([p[1] for p in pairs])
        wcs = fit_wcs_from_points((xy[ii, 0], xy[ii, 1]), all_world[ci],
                                  proj_point="center", projection="TAN")
        qx, qy = wcs.world_to_pixel(all_world[ci])
        resid = np.hypot(qx - xy[ii, 0], qy - xy[ii, 1])
        nmatch = len(pairs)
        print(f"  refine (tol {tol} px): {nmatch} stars, residual median "
              f"{np.median(resid):.2f} px ({np.median(resid) * scale:.2f}\"), "
              f"max {resid.max():.2f} px")

    cen = wcs.pixel_to_world(nx / 2.0, ny / 2.0)
    cd = wcs.pixel_scale_matrix * 3600.0
    solved_scale = np.sqrt(abs(np.linalg.det(cd)))
    rot = np.degrees(np.arctan2(cd[0, 1], cd[1, 1]))
    print(f"field centre {cen.to_string('hmsdms')}")
    print(f"solved scale {solved_scale:.4f}\"/px, position angle {rot:.2f} deg")
    print(f"field of view {nx * solved_scale / 60:.1f}' x "
          f"{ny * solved_scale / 60:.1f}'")

    whdr = wcs.to_header()
    for name in MASTERS:
        p = layout.path(f"master-{name}.fit")
        with fits.open(p, mode="update") as hd:
            for card in whdr.cards:
                hd[0].header[card.keyword] = (card.value, card.comment)
            hd[0].header["ASTRSOLV"] = (
                f"Gaia DR3 / {nmatch} stars / {np.median(resid) * scale:.2f} arcsec",
                "local plate solution")
            hd.flush()
        print(f"  WCS written to {os.path.basename(p)}")


if __name__ == "__main__":
    main()