"""Stage 4: plate solve the deepest master and copy the WCS to the others.

iTelescope's 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 unnecessary:
the header gives the pointing to arcminutes and usually the plate scale to four
figures, so a catalogue can be fetched for the right patch of sky and matched
directly.

The match is asterism-based, which is invariant to rotation and scale - so the
camera angle never has to be guessed, and a session at an unknown orientation
solves as easily as one at a known angle. It is NOT invariant to a mirror flip,
because a similarity transform cannot include reflection, so both parities are
tried and the one that matches wins.

The seed match returns only a handful of stars. That is enough to establish an
approximate solution but too few to trust, so the solution is then used to pair
every catalogue source with its nearest detection and refit, twice, with a
shrinking tolerance. The residual that comes out of that is the honest measure
of whether the solve is real.

Failure here is not fatal to the pipeline: an unsolved session still produces
images, it just cannot produce positions.
"""
import os

import astroalign as aa
import numpy as np
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
import measure as measure_mod

SEED_STARS = 600


def pointing(header):
    """Approximate field centre from whatever the header offers."""
    for ra_key, dec_key in (("OBJCTRA", "OBJCTDEC"), ("RA", "DEC"),
                            ("CRVAL1", "CRVAL2")):
        ra, dec = header.get(ra_key), header.get(dec_key)
        if ra is None or dec is None:
            continue
        try:
            if isinstance(ra, str) and (":" in ra or " " in ra):
                return SkyCoord(ra, dec, unit=(u.hourangle, u.deg))
            return SkyCoord(float(ra) * u.deg, float(dec) * u.deg)
        except Exception:                                     # noqa: BLE001
            continue
    return None


def plate_scale(header, session=None):
    """Arcsec per pixel: from the header if present, else from the optics."""
    val = header.get("HIERARCH iTelescopePlateScaleH")
    if val:
        try:
            return float(val)
        except (TypeError, ValueError):
            pass
    focal = header.get("FOCALLEN")
    pixel = header.get("XPIXSZ")
    if focal and pixel:
        # 206265 arcsec per radian; XPIXSZ is already binned microns.
        return 206.265 * float(pixel) / float(focal)
    return session.scale if session else None


def fetch_catalogue(centre, radius_deg, cache, limit=SEED_STARS, deep=False):
    """Gaia sources around the pointing, cached so a re-run costs nothing."""
    if os.path.exists(cache):
        z = np.load(cache)
        return z["ra"], z["dec"], z["g"]
    ra = dec = g = None
    mag_cut = 20.5 if deep else 18.0

    # VizieR first, deliberately. It serves the identical DR3 catalogue over a
    # plain HTTP query and answers in under a second. Gaia's own archive is
    # excellent but launch_job_async submits to a job QUEUE and polls for the
    # result, which for a few hundred rows means waiting minutes for work that
    # takes milliseconds - and it hung for ten minutes on a query of 600 stars
    # while the network was demonstrably healthy.
    try:
        from astroquery.vizier import Vizier
        # Ask for far more rows than are needed and pick the brightest
        # locally. VizieR's row_limit truncates the result BEFORE any sort, so
        # asking for 600 rows in a crowded field returns 600 arbitrary stars
        # rather than the 600 brightest - and asterism matching only works if
        # both lists contain the same bright stars. In the LMC field of
        # NGC 2030 that silently produced no match at all.
        v = Vizier(columns=["RA_ICRS", "DE_ICRS", "Gmag"],
                   column_filters={"Gmag": f"<{mag_cut}"}, row_limit=50000)
        v.TIMEOUT = 60
        res = v.query_region(centre, radius=radius_deg * u.deg,
                             catalog="I/355/gaiadr3")
        if res:
            t = res[0]
            ra = np.asarray(t["RA_ICRS"], float)
            dec = np.asarray(t["DE_ICRS"], float)
            g = np.asarray(t["Gmag"], float)
    except Exception as exc:                                  # noqa: BLE001
        print(f"        VizieR unavailable ({type(exc).__name__})")

    if ra is None or len(ra) == 0:
        try:
            from astroquery.gaia import Gaia
            Gaia.ROW_LIMIT = limit
            # Synchronous, not async: no job queue for a query this small.
            job = Gaia.launch_job(f"""
                SELECT TOP {limit} ra, dec, phot_g_mean_mag
                FROM gaiadr3.gaia_source
                WHERE 1 = CONTAINS(POINT('ICRS', ra, dec),
                         CIRCLE('ICRS', {centre.ra.deg}, {centre.dec.deg},
                                {radius_deg}))
                  AND phot_g_mean_mag IS NOT NULL
                  AND phot_g_mean_mag < {mag_cut}
                ORDER BY phot_g_mean_mag ASC""")
            t = job.get_results()
            ra = np.asarray(t["ra"], float)
            dec = np.asarray(t["dec"], float)
            g = np.asarray(t["phot_g_mean_mag"], float)
        except Exception as exc:                              # noqa: BLE001
            print(f"        Gaia archive unavailable ({type(exc).__name__})")
            return None, None, None

    order = np.argsort(g)[:limit]
    ra, dec, g = ra[order], dec[order], g[order]
    np.savez_compressed(cache, ra=ra, dec=dec, g=g)
    return ra, dec, g


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


def run(session, verbose=True):
    """Solve the deepest master; write the WCS into every master."""
    masters_dir = os.path.join(session.root, layout.MASTERS)
    deepest = session.filters[0]
    path = os.path.join(masters_dir, f"master-{deepest}.fit")
    if not os.path.exists(path):
        raise RuntimeError("no masters found - run the register stage first")

    with fits.open(path) as hd:
        image = hd[0].data.astype(np.float32)
        header = hd[0].header
    ny, nx = image.shape

    centre = pointing(header)
    scale = plate_scale(header, session)
    if centre is None or not scale:
        print("        no pointing or plate scale in the header - "
              "cannot solve without a blind solver")
        return None
    radius = 1.15 * 0.5 * np.hypot(nx, ny) * scale / 3600.0
    if verbose:
        print(f"        pointing {centre.to_string('hmsdms')}, "
              f"{scale:.4f} arcsec/px, search {radius:.3f} deg")

    xy, _ = measure_mod.load(session)
    # Detect on the master itself: it is deeper than any single frame.
    from measure import measure_frame
    det = measure_frame(path, max_stars=SEED_STARS)
    src = det["xy"]
    if len(src) < 12:
        print(f"        only {len(src)} stars in the master - too few")
        return None

    cache = os.path.join(session.root, layout.INTERMEDIATES, "_gaia.npz")
    os.makedirs(os.path.dirname(cache), exist_ok=True)
    ra, dec, gmag = fetch_catalogue(centre, radius, cache)
    if ra is None:
        print("        no catalogue available - skipping the solve")
        return None
    if verbose:
        print(f"        {len(ra)} catalogue stars, {len(src)} detected")

    # Matching needs the two lists to contain the SAME stars, and "brightest"
    # does not guarantee that. The detector deliberately rejects saturated
    # cores, so on a well exposed frame the brightest detections are not the
    # brightest stars in the sky - while the catalogue's brightest are. The
    # two top-N lists can therefore barely overlap, which is exactly what
    # happened on the crowded LMC field of NGC 2030.
    #
    # So slide a window down the catalogue's magnitude ranking as well as
    # varying the sample size, and take the best match found.
    best = None
    for offset in (0, 15, 40, 80):
        for n in (120, 200, 60):
            hi = min(offset + n, len(ra))
            if hi - offset < 25 or len(src) < 25:
                continue
            cat_slice = slice(offset, hi)
            for parity in (1.0, -1.0):
                cat_xy = project(ra[cat_slice], dec[cat_slice], centre, scale,
                                 parity)
                try:
                    tform, (cat_m, img_m) = aa.find_transform(
                        cat_xy, src[:max(n, 120)])
                except Exception:                             # noqa: BLE001
                    continue
                if best is None or len(cat_m) > len(best[0]):
                    # Record which catalogue rows matched, in FULL-list terms.
                    best = (cat_m, img_m, parity, offset)
                    if verbose:
                        print(f"        catalogue[{offset}:{hi}] parity "
                              f"{parity:+.0f}: matched {len(cat_m)} stars")
            if best is not None and len(best[0]) >= 8:
                break
        if best is not None and len(best[0]) >= 8:
            break

    if best is None:
        print("        no asterism match in either parity - solve failed")
        return None

    cat_m, img_m, parity, _offset = best
    cat_xy = project(ra, dec, centre, scale, parity)
    idx = [int(np.argmin(np.hypot(cat_xy[:, 0] - px, cat_xy[:, 1] - py)))
           for px, py in cat_m]
    world = SkyCoord(ra[idx] * u.deg, dec[idx] * u.deg)
    wcs = fit_wcs_from_points((img_m[:, 0], img_m[:, 1]), world,
                              proj_point="center", projection="TAN")

    # Refine: the seed match is a handful of stars, so use the approximate
    # solution to pair every catalogue source with its nearest detection.
    all_world = SkyCoord(ra * u.deg, dec * u.deg)
    nmatch, resid = len(cat_m), None
    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(src[:, 0] - cx, src[:, 1] - cy)
            j = int(np.argmin(d))
            if d[j] <= tol:
                pairs.append((i, j))
        if len(pairs) < 20:
            break
        ci = np.array([p[0] for p in pairs])
        ii = np.array([p[1] for p in pairs])
        wcs = fit_wcs_from_points((src[ii, 0], src[ii, 1]), all_world[ci],
                                  proj_point="center", projection="TAN")
        qx, qy = wcs.world_to_pixel(all_world[ci])
        resid = np.hypot(qx - src[ii, 0], qy - src[ii, 1])
        nmatch = len(pairs)

    cd = wcs.pixel_scale_matrix * 3600.0
    solved_scale = float(np.sqrt(abs(np.linalg.det(cd))))
    rot = float(np.degrees(np.arctan2(cd[0, 1], cd[1, 1])))
    med_resid = float(np.median(resid)) if resid is not None else float("nan")
    field = wcs.pixel_to_world(nx / 2.0, ny / 2.0)

    if verbose:
        print(f"        solved on {nmatch} stars, residual "
              f"{med_resid:.2f} px ({med_resid * solved_scale:.2f} arcsec)")
        print(f"        centre {field.to_string('hmsdms')}, "
              f"{solved_scale:.4f} arcsec/px, PA {rot:.2f} deg, "
              f"{nx * solved_scale / 60:.1f}' x {ny * solved_scale / 60:.1f}'")

    whdr = wcs.to_header()
    for filt in session.filters:
        p = os.path.join(masters_dir, f"master-{filt}.fit")
        if not os.path.exists(p):
            continue
        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 / "
                f"{med_resid * solved_scale:.2f} arcsec", "local plate solve")
            hd.flush()

    return dict(nstars=nmatch, residual_px=med_resid,
                residual_arcsec=med_resid * solved_scale,
                scale=solved_scale, position_angle=rot,
                centre=field.to_string("hmsdms"),
                fov_arcmin=(nx * solved_scale / 60, ny * solved_scale / 60))