"""Catalogue of synthetic frame corruptions used by the test bench.

Each entry models a failure mode actually seen in (or plausible for) this archive.
Corruptions are seeded and reproducible: the same case seed and frame always yield
byte-identical output, so a bench run can be repeated exactly.

Two kinds, because they break a frame at different layers:

* **array** corruptions damage the pixels; the bench re-writes a valid FITS around
  the result.
* **file** corruptions damage the bytes on disk, producing files that are not valid
  FITS at all -- truncated downloads, missing HDUs, bit rot.

The ``recompute_stats`` flag is the subtle part.  A real eclipse frame has a header
whose ``IMG_MEAN`` agrees with its dim pixels, because NOAA computed it from them; a
frame damaged in transit keeps the *original* header over broken pixels.  Setting
this correctly per mode is what makes the bench's verdict on header-only detection
honest -- otherwise header_v1 would appear to catch corruptions it could never see.
"""

from dataclasses import dataclass

import numpy as np

from . import fitsio

#: Corruption strength in [0, 1].  0 is a barely-perceptible defect, 1 is total loss.
DEFAULT_SEVERITY = 1.0


@dataclass(frozen=True)
class Corruption:
    """One failure mode."""

    name: str
    #: Which of the four groups this belongs to, for per-group reporting.
    group: str
    #: 'array' (damages pixels) or 'file' (damages bytes on disk).
    kind: str
    apply: callable
    #: Whether the header's radiance statistics are recomputed from the damaged
    #: pixels.  True models a fault upstream of NOAA's header generation.
    recompute_stats: bool = False
    #: Whether this mode needs a second frame to draw from.
    needs_donor: bool = False


# ------------------------------------------------------------------ dropout / blackout


def _eclipse_dim(image, rng, severity=DEFAULT_SEVERITY, donor=None):
    """The archive's most common real failure: Earth shadow drops radiance ~1000x."""
    factor = 10.0 ** (-4.0 * severity)
    dimmed = image * factor
    # Real eclipse frames keep sensor read noise, so they are not exactly zero.
    noise = rng.normal(0.0, float(np.abs(image).mean()) * 1e-4, image.shape)
    return (dimmed + noise).astype(np.float32), {"DEGRADED": True, "ECLIPSE": 2}


def _all_zero(image, rng, severity=DEFAULT_SEVERITY, donor=None):
    return np.zeros_like(image), {"EMPTY": True}


def _nan_fill(image, rng, severity=DEFAULT_SEVERITY, donor=None):
    """Undefined pixels over part or all of the frame."""
    out = image.copy()
    if severity >= 1.0:
        out[:] = np.nan
    else:
        mask = rng.random(image.shape) < severity
        out[mask] = np.nan
    return out, {}


def _zblank_fill(image, rng, severity=DEFAULT_SEVERITY, donor=None):
    """Fill with the FITS blank sentinel rather than NaN."""
    out = image.copy()
    mask = rng.random(image.shape) < severity if severity < 1.0 else np.ones(image.shape, bool)
    out[mask] = fitsio.ZBLANK
    return out, {}


# -------------------------------------------------------------------------- structural


def _truncate(raw, rng, severity=DEFAULT_SEVERITY):
    """A download cut short.  Keeps at least the primary header."""
    keep = max(fitsio.BLOCK, int(len(raw) * (1.0 - 0.9 * severity)))
    return raw[:keep]


def _drop_image_hdu(raw, rng, severity=DEFAULT_SEVERITY):
    """Only the primary header survives -- the blank-HDU case the old filter hit."""
    return raw[: fitsio.BLOCK]


def _block_corruption(raw, rng, severity=DEFAULT_SEVERITY):
    """Random bytes overwritten inside the data unit, leaving the header intact."""
    out = bytearray(raw)
    start = fitsio.BLOCK * 8  # past the headers
    if len(out) <= start:
        return bytes(out)
    span = max(1, int((len(out) - start) * 0.02 * severity))
    offset = int(rng.integers(start, len(out) - span))
    out[offset : offset + span] = rng.integers(0, 256, span, dtype=np.uint8).tobytes()
    return bytes(out)


def _torn_frame(image, rng, severity=DEFAULT_SEVERITY, donor=None):
    """Part of the frame comes from another observation -- a bad merge."""
    if donor is None:
        return image, {}
    out = image.copy()
    split = int(image.shape[0] * (1.0 - severity * 0.5))
    out[split:, :] = donor[split:, :]
    return out, {}


def _dropped_rows(image, rng, severity=DEFAULT_SEVERITY, donor=None):
    """Whole scan lines lost, as from a dropped packet."""
    out = image.copy()
    count = max(1, int(image.shape[0] * 0.3 * severity))
    rows = rng.choice(image.shape[0], size=count, replace=False)
    out[rows, :] = 0.0
    return out, {}


# ------------------------------------------------------------------------ radiometric


def _gain_shift(image, rng, severity=DEFAULT_SEVERITY, donor=None):
    """Calibration drift: everything scaled by a constant factor."""
    factor = 1.0 + 4.0 * severity * (1 if rng.random() < 0.5 else -0.2)
    return (image * factor).astype(np.float32), {}


def _offset_shift(image, rng, severity=DEFAULT_SEVERITY, donor=None):
    """A constant added to every pixel -- a dark-current or bias fault."""
    return (image + float(np.abs(image).mean()) * 5.0 * severity).astype(np.float32), {}


def _saturate(image, rng, severity=DEFAULT_SEVERITY, donor=None):
    """A blowout that drives a large fraction of the frame to the ceiling."""
    ceiling = float(np.nanmax(image)) or 1.0
    boosted = image * (1.0 + 50.0 * severity)
    return np.minimum(boosted, ceiling * 50.0).astype(np.float32), {}


def _gaussian_noise(image, rng, severity=DEFAULT_SEVERITY, donor=None):
    scale = float(np.nanstd(image)) * severity
    return (image + rng.normal(0.0, scale, image.shape)).astype(np.float32), {}


def _salt_pepper(image, rng, severity=DEFAULT_SEVERITY, donor=None):
    """Cosmic-ray hits and dead pixels."""
    out = image.copy()
    fraction = 0.05 * severity
    mask = rng.random(image.shape) < fraction
    extreme = float(np.nanmax(image)) or 1.0
    out[mask] = np.where(rng.random(int(mask.sum())) < 0.5, 0.0, extreme * 10.0)
    return out, {}


# ------------------------------------------------------------------ geometric/temporal


def _translate(image, rng, severity=DEFAULT_SEVERITY, donor=None):
    """Mispointing: the solar disc sits off centre."""
    shift = int(200 * severity)
    dx = int(rng.integers(-shift, shift + 1)) if shift else 0
    dy = int(rng.integers(-shift, shift + 1)) if shift else 0
    out = np.zeros_like(image)
    h, w = image.shape
    xs, xd = (max(0, -dx), max(0, dx))
    ys, yd = (max(0, -dy), max(0, dy))
    height, width = h - abs(dy), w - abs(dx)
    out[yd : yd + height, xd : xd + width] = image[ys : ys + height, xs : xs + width]
    return out, {"CRPIX1": (w + 1) / 2.0 + dx, "CRPIX2": (h + 1) / 2.0 + dy}


def _rotate(image, rng, severity=DEFAULT_SEVERITY, donor=None):
    """Wrong roll angle -- the disc is round, so only structure reveals this."""
    import cv2 as cv

    angle = 180.0 * severity
    centre = ((image.shape[1] - 1) / 2.0, (image.shape[0] - 1) / 2.0)
    matrix = cv.getRotationMatrix2D(centre, angle, 1.0)
    rotated = cv.warpAffine(
        np.nan_to_num(image), matrix, (image.shape[1], image.shape[0]), flags=cv.INTER_LINEAR
    )
    return rotated.astype(np.float32), {"CROTA": angle}


def _yaw_flip(image, rng, severity=DEFAULT_SEVERITY, donor=None):
    """The spacecraft's twice-yearly yaw flip applied when it should not be."""
    return np.flip(np.flip(image, 0), 1).copy(), {"YAW_FLIP": 1}


def _frozen(image, rng, severity=DEFAULT_SEVERITY, donor=None):
    """The feed stalled: this frame is a byte-for-byte repeat of a neighbour."""
    return (donor.copy() if donor is not None else image), {}


def _wrong_time(image, rng, severity=DEFAULT_SEVERITY, donor=None):
    """A frame from a different observation filed under this timestamp."""
    return (donor.copy() if donor is not None else image), {}


CATALOG = {
    corruption.name: corruption
    for corruption in (
        # Dropout / blackout -- the header follows the pixels, as NOAA computes it
        # from the image it actually produced.
        Corruption("eclipse_dim", "dropout", "array", _eclipse_dim, recompute_stats=True),
        Corruption("all_zero", "dropout", "array", _all_zero, recompute_stats=True),
        Corruption("nan_fill", "dropout", "array", _nan_fill, recompute_stats=False),
        Corruption("zblank_fill", "dropout", "array", _zblank_fill, recompute_stats=False),
        # Structural -- damage after the file was written, so the header is stale.
        Corruption("truncate", "structural", "file", _truncate),
        Corruption("drop_image_hdu", "structural", "file", _drop_image_hdu),
        Corruption("block_corruption", "structural", "file", _block_corruption),
        Corruption("torn_frame", "structural", "array", _torn_frame, needs_donor=True),
        Corruption("dropped_rows", "structural", "array", _dropped_rows),
        # Radiometric -- an instrument fault upstream of header generation.
        Corruption("gain_shift", "radiometric", "array", _gain_shift, recompute_stats=True),
        Corruption("offset_shift", "radiometric", "array", _offset_shift, recompute_stats=True),
        Corruption("saturate", "radiometric", "array", _saturate, recompute_stats=True),
        Corruption("gaussian_noise", "radiometric", "array", _gaussian_noise),
        Corruption("salt_pepper", "radiometric", "array", _salt_pepper),
        # Geometric / temporal -- the modes single-frame detectors are worst at.
        Corruption("translate", "geometric", "array", _translate, recompute_stats=True),
        Corruption("rotate", "geometric", "array", _rotate, recompute_stats=True),
        Corruption("yaw_flip", "geometric", "array", _yaw_flip),
        Corruption("frozen", "geometric", "array", _frozen, needs_donor=True),
        Corruption("wrong_time", "geometric", "array", _wrong_time, needs_donor=True),
    )
}

GROUPS = sorted({corruption.group for corruption in CATALOG.values()})


def by_group(group):
    return [name for name, c in CATALOG.items() if c.group == group]


def apply_array(name, image, seed, severity=DEFAULT_SEVERITY, donor=None):
    """Apply an array corruption.  Returns (image, header_overrides).

    Deterministic in `seed`, so a bench case is exactly reproducible.
    """
    corruption = CATALOG[name]
    if corruption.kind != "array":
        raise ValueError(f"{name} is a {corruption.kind} corruption, not an array one")
    if corruption.needs_donor and donor is None:
        raise ValueError(f"{name} requires a donor frame")
    rng = np.random.default_rng(seed)
    return corruption.apply(np.asarray(image, dtype=np.float32), rng,
                            severity=severity, donor=donor)


def apply_file(name, raw, seed, severity=DEFAULT_SEVERITY):
    """Apply a file-level corruption to raw FITS bytes."""
    corruption = CATALOG[name]
    if corruption.kind != "file":
        raise ValueError(f"{name} is a {corruption.kind} corruption, not a file one")
    rng = np.random.default_rng(seed)
    return corruption.apply(raw, rng, severity=severity)


def recomputed_stats(image):
    """Header statistics consistent with `image`, as NOAA would have written them."""
    finite = image[np.isfinite(image)]
    if finite.size == 0:
        return {"IMG_MIN": 0.0, "IMG_MAX": 0.0, "IMG_MEAN": 0.0, "IMG_SDEV": 0.0}
    return {
        "IMG_MIN": float(finite.min()),
        "IMG_MAX": float(finite.max()),
        "IMG_MEAN": float(finite.mean()),
        "IMG_SDEV": float(finite.std()),
    }