import json
import os

import numpy as np
import pytest

from hunt import images as im
from hunt import levels

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))


@pytest.fixture(scope="module")
def all_levels():
    return levels.make_all()


def test_generation_is_reproducible(all_levels):
    again = levels.make_all()
    assert json.dumps(all_levels, sort_keys=True) == json.dumps(again, sort_keys=True)


def test_shipped_files_match_generator(all_levels):
    for lv in all_levels:
        path = os.path.join(ROOT, "web", "levels", lv["id"] + ".json")
        assert os.path.exists(path), "run python data/make_levels.py"
        with open(path, encoding="utf-8") as f:
            assert json.load(f) == json.loads(json.dumps(lv))
        assert os.path.getsize(path) < 300 * 1024


def test_every_level_collapses_without_its_shortcut(all_levels):
    for lv in all_levels:
        r = levels.run_level(lv)
        if lv["fix"]["type"] == "metric":
            assert r["before"]["acc"] > 0.9
            assert r["after"]["balanced_acc"] < 0.6 and r["after"]["recall"] < 0.1
        else:
            assert r["before"]["acc"] >= 0.85, lv["id"]
            assert r["before"]["acc"] - r["after"]["acc"] >= 0.15, lv["id"]


def test_id_level_stump_picks_the_id(all_levels):
    r = levels.run_level(all_levels[0])
    assert r["before"]["feature"] == "id"
    assert r["after"]["feature"] != "id"


def test_dedupe_keeps_only_fresh_rows(all_levels):
    lv = next(l for l in all_levels if l["id"] == "dejavu")
    r = levels.run_level(lv)
    assert r["dropped"] == 240 and r["kept"] == 80


def test_time_level_month_is_the_shortcut(all_levels):
    lv = next(l for l in all_levels if l["id"] == "timestamp")
    r = levels.run_level(lv)
    assert r["before"]["features"][:12] == ["month=%d" % m for m in range(1, 13)]
    assert r["after_no_time"]["acc"] == pytest.approx(r["after"]["acc"], abs=0.05)


def test_image_roundtrip_and_counterfactuals(all_levels):
    m = im.shape_mask("triangle", 14, 14, 6)
    assert im.unpack_mask(im.pack_mask(m)).tolist() == m.tolist()
    img = im.compose(m, 12, 2, mark=True)
    assert img[2, 23] == im.MARK_VALUE and img.max() == 12 and img.min() == 2
    wm = next(l for l in all_levels if l["id"] == "watermark")
    r = levels.run_level(wm)
    cf = {c["y"]: c for c in r["counterfactual"]}
    assert cf[0]["p_before"] < 0.5 < cf[0]["p_after"]        # adding the mark flips a circle to "square"
    assert cf[1]["p_after"] < 0.5 < cf[1]["p_before"]        # removing it flips a square back
    assert sum(1 for rec in wm["train"] if rec["mk"]) == sum(1 for rec in wm["train"] if rec["y"] == 1)


def test_background_fix_breaks_the_correlation(all_levels):
    lv = next(l for l in all_levels if l["id"] == "background")
    y = np.array([r["y"] for r in lv["train"]])
    bg = np.array([r["bg"] for r in lv["train"]])
    fixed = np.array(lv["bg_fixed"]["train"])
    assert abs(np.corrcoef(y, bg)[0, 1]) > 0.7
    assert abs(np.corrcoef(y, fixed)[0, 1]) < 0.25
    assert sorted(bg.tolist()) == sorted(fixed.tolist())