#!/usr/bin/env python
"""Fit the differential-rotation profile the archive actually shows.

    fit_rotation.py                      # all six bands, built-in date spread
    fit_rotation.py --bands 195 304      # a subset, e.g. corona vs chromosphere
    fit_rotation.py --out fit.json       # also record the numbers

The warp in :mod:`suvi.fillers` / :mod:`suvi.align` uses the Snodgrass & Ulrich
(1990) *photospheric plasma* profile.  SUVI observes the EUV corona, which is known
to rotate more rigidly than the plasma below it, so this script measures what the
archive itself prefers: for real frame pairs 12 h and 24 h apart it scans

    omega_sidereal(lat) = A + s * (B0 sin^2 lat + C0 sin^4 lat)   [deg/day]

over (A, s) -- with B0, C0 the Snodgrass latitude terms -- warps the earlier frame
forward with each candidate profile, and scores it against the later frame on the
inner disc.  s = 1 recovers Snodgrass & Ulrich exactly; s = 0 is rigid rotation.

Frames are vetted with the production judgment (`detectors.judge_file`) so a bad
frame cannot masquerade as solar evolution.  Dates alternate into a held-out set
that never touches the fit; the fitted profile has to beat both s = 0 and s = 1 on
those dates before the numbers are worth baking into `suvi.align`.
"""

import argparse
import datetime
import glob
import json
import math
import os
import sys

import numpy as np

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from suvi import detectors, fillers, fitsio, paths

#: Snodgrass & Ulrich (1990) latitude terms, deg/day; the shape being rescaled.
B0_TERM = -2.396
C0_TERM = -1.787

#: Candidate equatorial sidereal rates, deg/day.  Snodgrass & Ulrich's 14.713 sits
#: mid-range; the ends are implausible enough to flag a broken measurement.
A_GRID = np.arange(13.75, 16.0, 0.25)
#: Candidate scales on the latitude terms.  1 = full plasma differential, 0 = rigid.
S_GRID = np.arange(0.0, 1.5, 0.25)

#: Gap lengths to measure across, seconds.
GAPS = (12 * 3600, 24 * 3600)

#: Inner-disc mask: far enough from the limb that foreshortening and the corona's
#: non-co-rotating off-disc glow do not dominate the score.
INNER_RHO2 = 0.81

#: How far past the nominal start time to keep looking for a usable frame.
SEARCH_SLOTS = 30


def build_date_list(first="2019-12-15", last="2025-03-15", every_months=3):
    """Mid-month dates every `every_months` months, hitting the B0 extremes.

    March and September (B0 = -/+7.25 deg) are in the cycle, so a B0-sign or
    tilt-handling defect cannot hide behind a mid-latitude-free sample.
    """
    start = datetime.date.fromisoformat(first)
    stop = datetime.date.fromisoformat(last)
    dates, month = [], start
    while month <= stop:
        dates.append(month.isoformat())
        year, next_month = month.year, month.month + every_months
        while next_month > 12:
            year, next_month = year + 1, next_month - 12
        month = datetime.date(year, next_month, 15)
    return dates


def split_holdout(dates, every=3):
    """(fit_dates, holdout_dates): every `every`-th date is held out."""
    fit = [d for i, d in enumerate(dates) if i % every != every - 1]
    held = [d for i, d in enumerate(dates) if i % every == every - 1]
    return fit, held


def profile_rate(a, s, latitude_rad, synodic=True):
    """The candidate profile, same convention as :func:`suvi.fillers.rotation_rate`."""
    sin2 = np.sin(latitude_rad) ** 2
    rate = a + s * (B0_TERM * sin2 + C0_TERM * sin2**2)
    return rate - fillers.EARTH_ORBIT_DEG_PER_DAY if synodic else rate


def refine_minimum(grid, errors):
    """Parabolic refinement of a 1-D sampled minimum; clamped to the grid ends."""
    k = int(np.argmin(errors))
    if k == 0 or k == len(grid) - 1:
        return float(grid[k])
    a, b, c = errors[k - 1], errors[k], errors[k + 1]
    denominator = a - 2 * b + c
    if denominator <= 0:
        return float(grid[k])
    step = grid[1] - grid[0]
    return float(grid[k] + 0.5 * (a - c) / denominator * step)


def find_frame(satellite, wavelength, t_start):
    """First usable frame at or after `t_start`, or (None, ...) if none nearby.

    Walks forward slot by slot so a routine per-slot outage does not discard the
    whole date, and vets with the production judgment so a dark or corrupt frame
    cannot enter the fit as fake solar evolution.
    """
    root = paths.data_root()
    for step in range(SEARCH_SLOTS):
        t = t_start + step * paths.CADENCE
        day = datetime.datetime.fromtimestamp(t, datetime.timezone.utc)
        stamp = day.strftime("%Y%m%dT%H%M%S")
        pattern = os.path.join(
            root, f"goes{satellite:d}", "l2", "data", f"suvi-l2-ci{wavelength:03d}",
            f"{day.year:04d}", f"{day.month:02d}", f"{day.day:02d}",
            f"dr_suvi-l2-ci{wavelength:03d}_g{satellite:d}_s{stamp}Z_*.fits",
        )
        for path in sorted(glob.glob(pattern)):
            slot = (satellite, wavelength, t)
            if detectors.judge_file(slot, path).verdict == detectors.BAD:
                continue
            image, error = fitsio.read_image(path)
            if image is None or image.shape != (1280, 1280):
                continue
            values, _ = fitsio.scan_header(path)
            return fillers._finite(image), values, t
    return None, None, None


def pair_surface(source, target, header, delta_seconds):
    """Relative error surface over (A_GRID, S_GRID) for one frame pair.

    Scores are inner-disc MSE in asinh-stretched units, normalised by the pair's
    own best so dates with different activity levels weigh equally, and the pair's
    time gap is taken from the frames' actual observation times.
    """
    height, width = target.shape
    radius = fillers._disc_radius_pixels(header, target.shape)
    crpix1 = float(header.get("crpix1", (width + 1) / 2.0)) - 1.0
    crpix2 = float(header.get("crpix2", (height + 1) / 2.0)) - 1.0
    grid_x, grid_y = np.meshgrid(np.arange(width), np.arange(height))
    rho2 = ((grid_x - crpix1) ** 2 + (grid_y - crpix2) ** 2) / radius**2
    inner = rho2 < INNER_RHO2

    scale = 0.1 * max(float(np.median(target[inner])), 1e-6)
    stretched_target = np.arcsinh(target / scale)

    surface = np.zeros((len(A_GRID), len(S_GRID)))
    original = fillers.rotation_rate
    try:
        for i, a in enumerate(A_GRID):
            for j, s in enumerate(S_GRID):
                fillers.rotation_rate = (
                    lambda lat, synodic=True, a=a, s=s: profile_rate(a, s, lat, synodic)
                )
                warped, visible = fillers._warp(source, header, delta_seconds, True)
                mask = inner & visible
                difference = np.arcsinh(warped / scale) - stretched_target
                surface[i, j] = float(np.mean(difference[mask] ** 2))
    finally:
        fillers.rotation_rate = original
    minimum = surface.min()
    return surface / minimum if minimum > 0 else surface


def measure_band(satellite, wavelength, dates, verbose=True):
    """Per-pair relative surfaces for one band: {date: {gap_seconds: surface}}."""
    surfaces = {}
    for date in dates:
        t0 = int(datetime.datetime.fromisoformat(date + "T00:00:00+00:00").timestamp())
        source, _, t_source = find_frame(satellite, wavelength, t0)
        if source is None:
            continue
        for gap in GAPS:
            target, header, t_target = find_frame(satellite, wavelength, t_source + gap)
            if target is None:
                continue
            surface = pair_surface(source, target, header, float(t_target - t_source))
            surfaces.setdefault(date, {})[gap] = surface
            if verbose:
                k = np.unravel_index(np.argmin(surface), surface.shape)
                print(f"  ci{wavelength:03d} {date} +{gap // 3600:2d}h: "
                      f"pair argmin A={A_GRID[k[0]]:.2f} s={S_GRID[k[1]]:.2f}",
                      flush=True)
    return surfaces


def fit_from_surfaces(surfaces):
    """Aggregate relative surfaces and refine the (A, s) minimum."""
    total = np.mean([s for by_gap in surfaces.values() for s in by_gap.values()], axis=0)
    i, j = np.unravel_index(np.argmin(total), total.shape)
    a = refine_minimum(A_GRID, total[:, j])
    s = max(0.0, refine_minimum(S_GRID, total[i, :]))
    return a, s, total


def holdout_score(surfaces, a, s):
    """Mean relative error of a profile on held-out surfaces, via interpolation."""
    scores = []
    for by_gap in surfaces.values():
        for surface in by_gap.values():
            ai = np.clip(np.interp(a, A_GRID, np.arange(len(A_GRID))), 0, len(A_GRID) - 1)
            sj = np.clip(np.interp(s, S_GRID, np.arange(len(S_GRID))), 0, len(S_GRID) - 1)
            i0, j0 = int(ai), int(sj)
            i1, j1 = min(i0 + 1, len(A_GRID) - 1), min(j0 + 1, len(S_GRID) - 1)
            fa, fs = ai - i0, sj - j0
            value = (surface[i0, j0] * (1 - fa) * (1 - fs)
                     + surface[i1, j0] * fa * (1 - fs)
                     + surface[i0, j1] * (1 - fa) * fs
                     + surface[i1, j1] * fa * fs)
            scores.append(value)
    return float(np.mean(scores)) if scores else math.nan


def main(argv=None):
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--satellite", type=int, default=16)
    parser.add_argument("--bands", type=int, nargs="*", default=list(paths.WAVELENGTHS))
    parser.add_argument("--first", default="2019-12-15")
    parser.add_argument("--last", default="2025-03-15")
    parser.add_argument("--holdout-every", type=int, default=3)
    parser.add_argument("--out", help="Write the fitted numbers to this JSON file")
    args = parser.parse_args(argv)

    dates = build_date_list(args.first, args.last)
    fit_dates, held_dates = split_holdout(dates, args.holdout_every)
    print(f"{len(fit_dates)} fit dates, {len(held_dates)} held out, "
          f"gaps {[g // 3600 for g in GAPS]} h, satellite G{args.satellite}")

    report = {"satellite": args.satellite, "fit_dates": fit_dates,
              "holdout_dates": held_dates, "bands": {}}
    pooled_fit, pooled_held = [], []
    for wavelength in args.bands:
        fit_surfaces = measure_band(args.satellite, wavelength, fit_dates)
        held_surfaces = measure_band(args.satellite, wavelength, held_dates)
        if not fit_surfaces:
            print(f"ci{wavelength:03d}: no usable pairs, skipped")
            continue
        a, s, _ = fit_from_surfaces(fit_surfaces)
        entry = {
            "A_sidereal": round(a, 3), "s": round(s, 3),
            "pairs_fit": sum(len(v) for v in fit_surfaces.values()),
            "pairs_holdout": sum(len(v) for v in held_surfaces.values()),
            "holdout_fitted": holdout_score(held_surfaces, a, s),
            "holdout_snodgrass": holdout_score(held_surfaces, 14.713, 1.0),
            "holdout_rigid": holdout_score(held_surfaces, 14.713, 0.0),
        }
        report["bands"][wavelength] = entry
        pooled_fit.append(fit_surfaces)
        pooled_held.append(held_surfaces)
        print(f"ci{wavelength:03d}: A={a:.3f} s={s:.3f}  "
              f"holdout fitted/snodgrass/rigid = "
              f"{entry['holdout_fitted']:.4f}/{entry['holdout_snodgrass']:.4f}"
              f"/{entry['holdout_rigid']:.4f}  ({entry['pairs_fit']} pairs)")

    if pooled_fit:
        merged_fit = {f"{k}#{i}": v for i, s in enumerate(pooled_fit) for k, v in s.items()}
        merged_held = {f"{k}#{i}": v for i, s in enumerate(pooled_held) for k, v in s.items()}
        a, s, _ = fit_from_surfaces(merged_fit)
        report["shared"] = {
            "A_sidereal": round(a, 3), "s": round(s, 3),
            "holdout_fitted": holdout_score(merged_held, a, s),
            "holdout_snodgrass": holdout_score(merged_held, 14.713, 1.0),
            "holdout_rigid": holdout_score(merged_held, 14.713, 0.0),
        }
        print(f"\nshared profile: A={a:.3f} s={s:.3f}  "
              f"holdout fitted/snodgrass/rigid = "
              f"{report['shared']['holdout_fitted']:.4f}"
              f"/{report['shared']['holdout_snodgrass']:.4f}"
              f"/{report['shared']['holdout_rigid']:.4f}")
        b, c = round(s * B0_TERM, 3), round(s * C0_TERM, 3)
        print(f"constants for suvi.align:  A={round(a, 3)}  B={b}  C={c}")

    if args.out:
        with open(args.out, "w") as handle:
            json.dump(report, handle, indent=2)
        print(f"wrote {args.out}")


if __name__ == "__main__":
    main()