"""Tests for lead-block target reduction, incl. the heatwave block-straddle rule."""
import os, sys
import numpy as np
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src"))
from ehd.targets import block_targets, _run_membership

BLOCKS = {"W1": (0, 10), "W2": (11, 17), "W3": (18, 24),
          "W4": (25, 31), "W5": (32, 38), "W6": (39, 45)}


def _series(true_days):
    a = np.zeros(46, dtype=float)
    a[list(true_days)] = 1.0
    return a


def test_primary_any_ehd():
    s = _series([3])                       # one EHD on day 3 (W1)
    prim, _ = block_targets(s, BLOCKS)
    assert prim["W1"] == 1.0
    assert all(prim[w] == 0.0 for w in ["W2", "W3", "W4", "W5", "W6"])


def test_heatwave_needs_three_consecutive():
    prim, sec = block_targets(_series([5, 6]), BLOCKS)      # run of 2
    assert prim["W1"] == 1.0 and sec["W1"] == 0.0
    prim, sec = block_targets(_series([5, 6, 7]), BLOCKS)   # run of 3
    assert sec["W1"] == 1.0


def test_heatwave_straddling_boundary_flags_both_blocks():
    # run on days 10,11,12: day10 in W1, days 11,12 in W2 -> both flagged
    _, sec = block_targets(_series([10, 11, 12]), BLOCKS)
    assert sec["W1"] == 1.0
    assert sec["W2"] == 1.0
    assert sec["W3"] == 0.0


def test_run_membership_general_min_run():
    e = np.array([0, 1, 1, 1, 0, 1, 1, 0], dtype=bool)
    m = _run_membership(e, 3)
    assert list(m.astype(int)) == [0, 1, 1, 1, 0, 0, 0, 0]


def test_missing_day_makes_block_nan():
    s = _series([3]); s[7] = np.nan          # a missing day inside W1
    prim, sec = block_targets(s, BLOCKS)
    assert np.isnan(prim["W1"]) and np.isnan(sec["W1"])
    assert prim["W2"] == 0.0                  # other blocks unaffected


def test_vectorized_over_grid():
    grid = np.zeros((46, 2, 2), dtype=float)
    grid[5:8, 0, 0] = 1.0                     # heatwave at cell (0,0)
    grid[3, 1, 1] = 1.0                       # single EHD at cell (1,1)
    prim, sec = block_targets(grid, BLOCKS)
    assert prim["W1"][0, 0] == 1.0 and sec["W1"][0, 0] == 1.0
    assert prim["W1"][1, 1] == 1.0 and sec["W1"][1, 1] == 0.0