import numpy as np
import pytest

from conftest import solar_disc, write_fits
from suvi import corruptions, fitsio

ARRAY_MODES = [n for n, c in corruptions.CATALOG.items() if c.kind == "array"]
FILE_MODES = [n for n, c in corruptions.CATALOG.items() if c.kind == "file"]


@pytest.fixture
def image():
    return solar_disc(peak=3.0)


@pytest.fixture
def donor():
    return solar_disc(peak=3.0, centre=(700, 640))


def apply(name, image, donor=None, seed=7, severity=1.0):
    return corruptions.apply_array(name, image, seed, severity, donor)


# ------------------------------------------------------------------------ catalogue


def test_catalogue_covers_all_four_groups():
    assert set(corruptions.GROUPS) == {"dropout", "structural", "radiometric", "geometric"}
    for group in corruptions.GROUPS:
        assert corruptions.by_group(group), f"{group} has no modes"


def test_every_entry_is_self_consistent():
    for name, corruption in corruptions.CATALOG.items():
        assert corruption.name == name
        assert corruption.kind in ("array", "file")
        assert callable(corruption.apply)


# -------------------------------------------------------------------- determinism


@pytest.mark.parametrize("name", ARRAY_MODES)
def test_array_corruptions_are_deterministic(name, image, donor):
    first, header_a = apply(name, image, donor)
    second, header_b = apply(name, image, donor)
    np.testing.assert_array_equal(np.nan_to_num(first), np.nan_to_num(second))
    assert header_a == header_b


@pytest.mark.parametrize("name", ARRAY_MODES)
def test_array_corruptions_change_the_data(name, image, donor):
    damaged, _ = apply(name, image, donor)
    assert damaged.shape == image.shape
    assert not np.array_equal(np.nan_to_num(damaged), np.nan_to_num(image))


@pytest.mark.parametrize("name", ARRAY_MODES)
def test_different_seeds_stay_reproducible(name, image, donor):
    first, _ = apply(name, image, donor, seed=1)
    again, _ = apply(name, image, donor, seed=1)
    np.testing.assert_array_equal(np.nan_to_num(first), np.nan_to_num(again))


@pytest.mark.parametrize("name", FILE_MODES)
def test_file_corruptions_are_deterministic(name, tmp_path, image):
    path = write_fits(str(tmp_path / "src.fits"), image)
    raw = open(path, "rb").read()
    first = corruptions.apply_file(name, raw, 7, 1.0)
    second = corruptions.apply_file(name, raw, 7, 1.0)
    assert first == second
    assert first != raw


# ------------------------------------------------------------------ specific modes


def test_eclipse_dim_reproduces_the_real_failure(image):
    """The archive's commonest fault: radiance collapses, flags are set."""
    damaged, header = apply("eclipse_dim", image)
    assert header["DEGRADED"] is True and header["ECLIPSE"] == 2
    assert abs(float(damaged.mean())) < float(image.mean()) / 1000


def test_all_zero_sets_the_empty_flag(image):
    damaged, header = apply("all_zero", image)
    assert not damaged.any() and header["EMPTY"] is True


def test_nan_fill_at_full_severity_covers_everything(image):
    damaged, _ = apply("nan_fill", image, severity=1.0)
    assert np.isnan(damaged).all()


def test_nan_fill_at_partial_severity_is_partial(image):
    damaged, _ = apply("nan_fill", image, severity=0.3)
    fraction = np.isnan(damaged).mean()
    assert 0.1 < fraction < 0.5


def test_zblank_fill_uses_the_fits_sentinel(image):
    damaged, _ = apply("zblank_fill", image, severity=1.0)
    assert (damaged == fitsio.ZBLANK).all()


def test_frozen_returns_the_donor_exactly(image, donor):
    damaged, _ = apply("frozen", image, donor)
    np.testing.assert_array_equal(damaged, donor)


def test_torn_frame_mixes_two_observations(image, donor):
    damaged, _ = apply("torn_frame", image, donor, severity=1.0)
    assert np.array_equal(damaged[-1], donor[-1])
    assert np.array_equal(damaged[0], image[0])


def test_dropped_rows_blanks_whole_scan_lines(image):
    damaged, _ = apply("dropped_rows", image, severity=1.0)
    blank = [row for row in range(damaged.shape[0]) if not damaged[row].any()]
    assert len(blank) >= 1


def test_yaw_flip_is_a_180_degree_rotation(image):
    damaged, header = apply("yaw_flip", image)
    assert header["YAW_FLIP"] == 1
    np.testing.assert_array_equal(damaged, np.flip(np.flip(image, 0), 1))


def test_translate_moves_the_disc_and_updates_the_wcs(image):
    damaged, header = apply("translate", image, severity=1.0)
    assert "CRPIX1" in header and "CRPIX2" in header
    assert not np.array_equal(damaged, image)


def test_saturate_raises_the_ceiling(image):
    damaged, _ = apply("saturate", image, severity=1.0)
    assert float(damaged.max()) > float(image.max())


def test_modes_requiring_a_donor_say_so(image):
    for name, corruption in corruptions.CATALOG.items():
        if corruption.needs_donor:
            with pytest.raises(ValueError, match="donor"):
                corruptions.apply_array(name, image, 1, 1.0, None)


def test_truncate_keeps_at_least_one_block(tmp_path, image):
    path = write_fits(str(tmp_path / "src.fits"), image)
    raw = open(path, "rb").read()
    damaged = corruptions.apply_file("truncate", raw, 1, 1.0)
    assert fitsio.BLOCK <= len(damaged) < len(raw)


def test_drop_image_hdu_leaves_only_the_primary_header(tmp_path, image):
    path = write_fits(str(tmp_path / "src.fits"), image)
    raw = open(path, "rb").read()
    damaged = corruptions.apply_file("drop_image_hdu", raw, 1, 1.0)
    assert len(damaged) == fitsio.BLOCK
    out = str(tmp_path / "damaged.fits")
    open(out, "wb").write(damaged)
    data, error = fitsio.read_image(out)
    assert data is None and error


def test_block_corruption_leaves_the_header_readable(tmp_path, image):
    """Simulates bit rot: the header still parses, the data no longer matches it."""
    path = write_fits(str(tmp_path / "src.fits"), image)
    raw = open(path, "rb").read()
    damaged = corruptions.apply_file("block_corruption", raw, 3, 1.0)
    out = str(tmp_path / "damaged.fits")
    open(out, "wb").write(damaged)
    cards, error = fitsio.read_header(out)
    assert cards.get("WAVELNTH") == 171
    assert "datasum mismatch" in fitsio.verify_datasums(out)


# ---------------------------------------------------------------------- interface


def test_wrong_kind_is_rejected(image, tmp_path):
    with pytest.raises(ValueError, match="not an array one"):
        corruptions.apply_array("truncate", image, 1)
    with pytest.raises(ValueError, match="not a file one"):
        corruptions.apply_file("eclipse_dim", b"x" * 3000, 1)


def test_unknown_mode_raises(image):
    with pytest.raises(KeyError):
        corruptions.apply_array("does_not_exist", image, 1)


def test_recomputed_stats_match_the_damaged_pixels(image):
    stats = corruptions.recomputed_stats(image)
    assert stats["IMG_MEAN"] == pytest.approx(float(image.mean()), rel=1e-5)
    assert stats["IMG_MAX"] == pytest.approx(float(image.max()), rel=1e-5)


def test_recomputed_stats_survive_an_all_nan_frame():
    stats = corruptions.recomputed_stats(np.full((4, 4), np.nan, np.float32))
    assert stats == {"IMG_MIN": 0.0, "IMG_MAX": 0.0, "IMG_MEAN": 0.0, "IMG_SDEV": 0.0}


def test_stats_are_recomputed_only_where_the_header_would_follow():
    """A real eclipse has a matching header; bit rot does not.  See the module docs."""
    assert corruptions.CATALOG["eclipse_dim"].recompute_stats is True
    assert corruptions.CATALOG["gain_shift"].recompute_stats is True
    assert corruptions.CATALOG["block_corruption"].recompute_stats is False
    assert corruptions.CATALOG["salt_pepper"].recompute_stats is False