"""シミュレーションの物差しの試験（sim/schelling.py）。

陽性対照: 完全に分かれた街は分離度 1、でたらめな街は理論の期待値に一致する。
"""
import os
import sys

import numpy as np
import pytest

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, os.path.join(ROOT, "sim"))
import schelling as S  # noqa: E402


def _random_grid(rng, size=100, empty=0.10, p=0.5):
    g = np.where(rng.random((size, size)) < p, 1, 2).astype(np.int8)
    g[rng.random((size, size)) < empty] = 0
    return g


def test_fully_segregated_scores_one_two_colours():
    g = np.ones((100, 100), dtype=np.int8)
    g[:, 50:] = 2
    rng = np.random.default_rng(1)
    g[rng.random((100, 100)) < 0.1] = 0          # 空き家がどこにあっても地区は1色
    assert S.dissimilarity(g, 2, 10) == pytest.approx(1.0, abs=1e-12)


def test_fully_segregated_scores_one_three_colours():
    g = np.zeros((100, 100), dtype=np.int8)
    for bi in range(10):
        for bj in range(10):
            g[bi * 10:(bi + 1) * 10, bj * 10:(bj + 1) * 10] = 1 + (bi * 10 + bj) % 3
    assert S.dissimilarity(g, 3, 10) == pytest.approx(1.0, abs=1e-12)


def test_perfectly_even_scores_zero():
    g = np.ones((100, 100), dtype=np.int8)
    g[:, 1::2] = 2                                 # どの地区も半々
    assert S.dissimilarity(g, 2, 10) == pytest.approx(0.0, abs=1e-12)


def test_two_colour_D_equals_duncan_formula():
    rng = np.random.default_rng(7)
    g = _random_grid(rng)
    blocks = g.reshape(10, 10, 10, 10)
    b = (blocks == 1).sum(axis=(1, 3)).ravel().astype(float)
    w = (blocks == 2).sum(axis=(1, 3)).ravel().astype(float)
    duncan = 0.5 * np.abs(b / b.sum() - w / w.sum()).sum()
    assert S.dissimilarity(g, 2, 10) == pytest.approx(duncan, abs=1e-12)


def test_random_grid_matches_expected_D_and_same_share():
    rng = np.random.default_rng(2026)
    Ds, ms = [], []
    sizes = None
    for _ in range(40):
        g = _random_grid(rng)
        Ds.append(S.dissimilarity(g, 2, 10))
        ms.append(S.measure(g, 2, False)["mean_same"])
        sz = (g.reshape(10, 10, 10, 10) > 0).sum(axis=(1, 3)).ravel()
        sizes = sz if sizes is None else sizes + sz
    expected = S.expected_dissimilarity_random((sizes / 40).round().astype(int).tolist(), 0.5)
    assert np.mean(Ds) == pytest.approx(expected, abs=0.005)
    assert np.mean(ms) == pytest.approx(0.5, abs=0.005)


def test_neighbour_counts_match_brute_force():
    rng = np.random.default_rng(3)
    g = _random_grid(rng, size=20)
    for torus in (False, True):
        c = S.neighbour_counts(g, 2, torus)
        for y in range(20):
            for x in range(20):
                for col in (1, 2):
                    n = 0
                    for dy, dx in S.NBR_OFFSETS:
                        yy, xx = y + dy, x + dx
                        if torus:
                            yy %= 20
                            xx %= 20
                        elif not (0 <= yy < 20 and 0 <= xx < 20):
                            continue
                        n += g[yy, xx] == col
                    assert c[col - 1, y, x] == n


def test_run_conserves_population_and_converges_to_zero_unhappy():
    cond = S.Condition(key="t", tol=0.30, size=40, max_steps=200)
    r = S.run(cond, 11, keep_display=True)
    init = np.array(list(r["display"]["init"]), dtype=int)
    fin = np.array(list(r["final_grid"]), dtype=int)
    for v in (0, 1, 2):
        assert (init == v).sum() == (fin == v).sum()
    assert r["converged"]
    assert r["final"]["unhappy"] == 0.0
    # 引越しの記録を再生すると最終の街になる（画面の再生がそのまま正しい）
    g = init.copy()
    for st in r["display"]["step_moves"]:
        for a, b in st:
            g[b], g[a] = g[a], 0
    assert (g == fin).all()


def test_same_seed_same_result_and_030_equals_033():
    a = S.run(S.Condition(key="a", tol=0.30, size=40), 5)
    b = S.run(S.Condition(key="b", tol=0.30, size=40), 5)
    c = S.run(S.Condition(key="c", tol=0.33, size=40), 5)
    assert a["final_grid_sha256"] == b["final_grid_sha256"]
    assert a["final_grid_sha256"] == c["final_grid_sha256"]


def test_indifferent_agents_never_move():
    cond = S.Condition(key="i", tol=0.30, indifferent=0.3, size=40)
    r = S.run(cond, 9, keep_display=True)
    ind = set(r["indifferent_cells_init"])
    moved_from = {a for st in r["display"]["step_moves"] for a, _ in st}
    assert ind and not (ind & moved_from)
    assert set(r["indifferent_cells_final"]) == ind