"""Methods for reconstructing a missing or rejected frame.

Every filler takes the same :class:`FillContext` and returns a replacement array, so
the bench can swap them without knowing which one it is holding.  All are pure: they
read the context and return an array, nothing else.

The methods span a deliberate range of physical sophistication, from "repeat the last
good frame" (what the pipeline does today) to a differential-rotation warp that models
how the Sun actually moves.  The bench exists to say which of them is worth the cost
at which gap length.
"""

from dataclasses import dataclass, field

import cv2 as cv
import numpy as np

#: Nominal solar radius in metres (IAU 2015).
R_SUN = 6.957e8
#: Snodgrass (1983) sidereal differential rotation, degrees per day, by latitude.
SNODGRASS_A = 14.713
SNODGRASS_B = -2.396
SNODGRASS_C = -1.787
#: Earth's mean orbital motion, subtracted to get the rotation an Earth-orbiting
#: observer actually sees.
EARTH_ORBIT_DEG_PER_DAY = 0.9856
SECONDS_PER_DAY = 86400.0


@dataclass
class FillContext:
    """Everything a filler may draw on to reconstruct one frame."""

    #: Nearest good frame before the gap, and how many seconds back it sits.
    before: np.ndarray | None = None
    dt_before: float = 0.0
    #: Nearest good frame after the gap, and how many seconds forward.
    after: np.ndarray | None = None
    dt_after: float = 0.0
    #: The other satellite's view of this same instant, if it has one.
    counterpart: np.ndarray | None = None
    #: Header of the frame being reconstructed, for the WCS a rotation warp needs.
    header: dict = field(default_factory=dict)

    @property
    def alpha(self):
        """Position within the gap: 0 at `before`, 1 at `after`."""
        span = self.dt_before + self.dt_after
        if span <= 0:
            return 0.0
        return self.dt_before / span

    @property
    def gap_frames(self):
        """Gap width in 4-minute slots, for reporting quality against gap length."""
        return int(round((self.dt_before + self.dt_after) / 240.0))


def _finite(image):
    return np.nan_to_num(np.asarray(image, dtype=np.float32), nan=0.0, posinf=0.0, neginf=0.0)


# ---------------------------------------------------------------- simple baselines


def hold_last(context):
    """Repeat the last good frame.

    What ``merger_FITS.py`` does today (up to ``max_time_gap`` slots).  Cheap, never
    invents structure, but freezes the Sun and then jumps -- the visible stutter in
    the current videos.  The baseline every other method must beat.
    """
    if context.before is not None:
        return _finite(context.before)
    if context.after is not None:
        return _finite(context.after)
    return None


def linear_blend(context):
    """Cross-fade between the frames bracketing the gap.

    Removes the jump that ``hold_last`` leaves, at the cost of ghosting: moving
    features appear twice, faintly, rather than moving.
    """
    if context.before is None:
        return hold_last(context)
    if context.after is None or context.before.shape != context.after.shape:
        # Frames of differing size cannot be mixed; the nearer one is the best
        # available answer.  This is also the fallback the other fillers unwind to.
        return _finite(context.before)
    alpha = context.alpha
    return ((1.0 - alpha) * _finite(context.before) + alpha * _finite(context.after)).astype(
        np.float32
    )


# ------------------------------------------------------------------- optical flow


def _for_flow(image):
    """Compress radiance into a range optical flow can work with.

    Radiance is heavy-tailed -- a flare can be 1000x the quiet corona -- so raw
    values make flow chase the brightest pixels only.  log1p plus a percentile
    stretch keeps faint structure in play.
    """
    scaled = np.log1p(np.clip(_finite(image), 0.0, None))
    high = np.percentile(scaled, 99.5)
    if high <= 0:
        return np.zeros(scaled.shape, dtype=np.uint8)
    return np.clip(scaled / high * 255.0, 0, 255).astype(np.uint8)


def optical_flow(context, use_dis=True):
    """Motion-compensated interpolation between the bracketing frames.

    Estimates dense flow both ways and warps each bracket forward to the target
    instant, then blends.  Unlike ``linear_blend`` this moves features instead of
    dissolving between them, which is what the eye reads as smooth motion.
    """
    if context.before is None or context.after is None:
        return linear_blend(context)

    before, after = _finite(context.before), _finite(context.after)
    if before.shape != after.shape:
        return linear_blend(context)

    first, second = _for_flow(before), _for_flow(after)
    if use_dis:
        engine = cv.DISOpticalFlow_create(cv.DISOPTICAL_FLOW_PRESET_MEDIUM)
        forward = engine.calc(first, second, None)
        backward = engine.calc(second, first, None)
    else:
        forward = cv.calcOpticalFlowFarneback(
            first, second, None, 0.5, 3, 15, 3, 5, 1.2, 0
        )
        backward = cv.calcOpticalFlowFarneback(
            second, first, None, 0.5, 3, 15, 3, 5, 1.2, 0
        )

    alpha = context.alpha
    height, width = before.shape
    grid_x, grid_y = np.meshgrid(
        np.arange(width, dtype=np.float32), np.arange(height, dtype=np.float32)
    )
    # cv.remap samples the source *at* the map coordinates, so to place a feature
    # where it should be at time alpha we read from where it was: a feature at x in
    # `before` sits at x + forward(x) in `after`, hence at p - alpha*forward(p) when
    # looking back from the interpolated frame.  Adding the flow instead of
    # subtracting it moves every feature the wrong way, which is worse than not
    # compensating at all.
    warped_before = cv.remap(
        before,
        grid_x - forward[..., 0] * alpha,
        grid_y - forward[..., 1] * alpha,
        cv.INTER_LINEAR,
        borderMode=cv.BORDER_REPLICATE,
    )
    warped_after = cv.remap(
        after,
        grid_x - backward[..., 0] * (1.0 - alpha),
        grid_y - backward[..., 1] * (1.0 - alpha),
        cv.INTER_LINEAR,
        borderMode=cv.BORDER_REPLICATE,
    )
    return ((1.0 - alpha) * warped_before + alpha * warped_after).astype(np.float32)


# ------------------------------------------------------------ cross-satellite fill


def gain_match(source, reference):
    """Put `source` on `reference`'s radiance scale by least squares.

    GOES-16 and GOES-18 carry different SUVI flight models, so their radiances
    differ by a roughly affine factor even when both are healthy.
    """
    x = _finite(source).ravel().astype(np.float64)
    y = _finite(reference).ravel().astype(np.float64)
    variance = float(((x - x.mean()) ** 2).sum())
    if variance <= 0:
        return np.asarray(source, dtype=np.float32)
    gain = float(((x - x.mean()) * (y - y.mean())).sum() / variance)
    offset = float(y.mean() - gain * x.mean())
    return (np.asarray(source, dtype=np.float32) * gain + offset).astype(np.float32)


def crosssat(context, align=True):
    """Substitute the other satellite's view of the same instant.

    The two spacecraft see the same Sun from 1 AU, so the substitute is a real
    observation of the real Sun at the right time -- not an interpolation.  It
    should dominate every temporal method whenever it is available, which is the
    thing worth quantifying: it is unavailable in the 31% of slots where both
    satellites are out simultaneously.

    Residual differences are instrument calibration (removed by gain matching) and
    a few pixels of geostationary parallax (removed by alignment).
    """
    if context.counterpart is None:
        return None
    counterpart = _finite(context.counterpart)
    reference = context.before if context.before is not None else context.after
    if reference is None:
        return counterpart
    reference = _finite(reference)
    if counterpart.shape != reference.shape:
        return counterpart

    if align:
        window = cv.createHanningWindow(
            (counterpart.shape[1], counterpart.shape[0]), cv.CV_64F
        )
        (dx, dy), _ = cv.phaseCorrelate(
            counterpart.astype(np.float64), reference.astype(np.float64), window
        )
        matrix = np.array([[1.0, 0.0, dx], [0.0, 1.0, dy]], dtype=np.float32)
        counterpart = cv.warpAffine(
            counterpart,
            matrix,
            (counterpart.shape[1], counterpart.shape[0]),
            flags=cv.INTER_LINEAR,
            borderMode=cv.BORDER_REPLICATE,
        )
    return gain_match(counterpart, reference)


# --------------------------------------------------------- solar rotation warping


def rotation_rate(latitude_rad, synodic=True):
    """Snodgrass differential rotation in degrees per day at a given latitude."""
    sin2 = np.sin(latitude_rad) ** 2
    rate = SNODGRASS_A + SNODGRASS_B * sin2 + SNODGRASS_C * sin2**2
    return rate - EARTH_ORBIT_DEG_PER_DAY if synodic else rate


def _disc_radius_pixels(header, shape):
    """Solar radius in pixels, from the header if possible."""
    diameter = header.get("diam_sun")
    if diameter:
        return float(diameter) / 2.0
    distance, scale = header.get("dsun_obs"), header.get("cdelt1")
    if distance and scale:
        return float(np.degrees(np.arcsin(R_SUN / distance)) * 3600.0 / scale)
    return shape[0] * 0.3  # falls back to the archive's typical disc fraction


def _rotation_map(shape, header, delta_seconds, synodic=True):
    """Inverse map: for each output pixel, where in the input it came from.

    Works in heliographic coordinates -- de-project each pixel onto the sphere, undo
    the rotation that happened over `delta_seconds`, re-project.  Returns
    (map_x, map_y, on_disc) with NaN where the source point is not visible.
    """
    height, width = shape
    radius = _disc_radius_pixels(header, shape)
    crpix1 = float(header.get("crpix1", (width + 1) / 2.0)) - 1.0
    crpix2 = float(header.get("crpix2", (height + 1) / 2.0)) - 1.0
    b0 = np.radians(float(header.get("solar_b0", 0.0)))

    grid_x, grid_y = np.meshgrid(np.arange(width), np.arange(height))
    x = (grid_x - crpix1) / radius
    y = (grid_y - crpix2) / radius

    rho2 = x**2 + y**2
    on_disc = rho2 < 1.0
    z = np.sqrt(np.clip(1.0 - rho2, 0.0, None))

    # Plane-of-sky -> heliographic, undoing the observer's B0 tilt.
    sin_lat = y * np.cos(b0) + z * np.sin(b0)
    sin_lat = np.clip(sin_lat, -1.0, 1.0)
    latitude = np.arcsin(sin_lat)
    longitude = np.arctan2(x, z * np.cos(b0) - y * np.sin(b0))

    # Step the longitude back to where this material was `delta_seconds` ago.
    days = delta_seconds / SECONDS_PER_DAY
    source_longitude = longitude - np.radians(rotation_rate(latitude, synodic)) * days

    # Heliographic -> plane-of-sky.
    cos_lat = np.cos(latitude)
    sx = cos_lat * np.sin(source_longitude)
    sy = sin_lat * np.cos(b0) - cos_lat * np.cos(source_longitude) * np.sin(b0)
    sz = sin_lat * np.sin(b0) + cos_lat * np.cos(source_longitude) * np.cos(b0)

    visible = on_disc & (sz > 0)
    map_x = (sx * radius + crpix1).astype(np.float32)
    map_y = (sy * radius + crpix2).astype(np.float32)
    return map_x, map_y, visible


def solar_rotation(context, synodic=True):
    """Warp the bracketing frames by differential solar rotation, then blend.

    The Sun is not a rigid body: the equator turns in about 25 days, the poles in
    about 35.  Over a short gap that is a sub-pixel effect, but across a multi-hour
    outage it is the difference between features landing where they belong and
    smearing.  This is the only method here that uses a physical model of the scene.

    Applies on-disc only.  The corona above the limb does not co-rotate with the
    photosphere, so off-disc pixels fall back to a plain cross-fade.
    """
    if context.before is None and context.after is None:
        return None
    if context.before is None or context.after is None:
        source = context.before if context.before is not None else context.after
        delta = context.dt_before if context.before is not None else -context.dt_after
        warped, visible = _warp(source, context.header, delta, synodic)
        blended = np.where(visible, warped, _finite(source))
        return blended.astype(np.float32)

    before, after = _finite(context.before), _finite(context.after)
    if before.shape != after.shape:
        return linear_blend(context)

    # Roll `before` forward to the target instant and `after` backward to it.
    warped_before, visible_before = _warp(before, context.header, context.dt_before, synodic)
    warped_after, visible_after = _warp(after, context.header, -context.dt_after, synodic)

    alpha = context.alpha
    rotated = (1.0 - alpha) * warped_before + alpha * warped_after
    faded = (1.0 - alpha) * before + alpha * after
    visible = visible_before & visible_after
    return np.where(visible, rotated, faded).astype(np.float32)


def _warp(image, header, delta_seconds, synodic):
    image = _finite(image)
    map_x, map_y, visible = _rotation_map(image.shape, header, delta_seconds, synodic)
    warped = cv.remap(
        image, map_x, map_y, cv.INTER_LINEAR, borderMode=cv.BORDER_CONSTANT, borderValue=0.0
    )
    return warped, visible


FILLERS = {
    "hold_last": hold_last,
    "linear_blend": linear_blend,
    "optical_flow": optical_flow,
    "crosssat": crosssat,
    "solar_rotation": solar_rotation,
}

# TODO: learned filler.  Train a model to predict a frame from its preceding frames,
# following frames, and the other satellite's view, then evaluate it here across
# severities of missing data and prediction horizons (single-frame gaps through
# multi-hour outages, one satellite out versus both).  It plugs in as another entry
# in FILLERS and reuses the bench's existing cases and metrics unchanged.