"""Learning tests: the TD-Leaf update moves weights in the right direction."""
import math

import pytest

from engine import ALL_FEATURES, N_FEATURES, N_MATERIAL, equal_material_weights
from learn import (TrainConfig, average_checkpoints, play_selfplay_game, sigmoid,
                   td_leaf_update, train)

QUEEN = ALL_FEATURES.index("queen")
PAWN = ALL_FEATURES.index("pawn")
MATERIAL_MASK = [i < N_MATERIAL for i in range(N_FEATURES)]


def one_hot(i, value=1.0):
    f = [0.0] * N_FEATURES
    f[i] = value
    return f


def synthetic_game(feat, n, w, result, beta):
    """n plies whose leaves all have feature vector `feat`, then `result`."""
    v = sum(a * b for a, b in zip(w, feat))
    return {"scores": [v] * n, "feats": [feat] * n, "result": result}


def test_extra_queen_and_win_increases_queen_value():
    cfg = TrainConfig(alpha=0.1, lam=0.7, beta=0.5)
    w = equal_material_weights()
    game = synthetic_game(one_hot(QUEEN), 10, w, 1.0, cfg.beta)
    w2 = td_leaf_update(w, game, cfg, MATERIAL_MASK)
    assert w2[QUEEN] > w[QUEEN]
    assert [w2[i] for i in range(N_FEATURES) if i != QUEEN] == \
           [w[i] for i in range(N_FEATURES) if i != QUEEN]


def test_extra_pawn_and_loss_decreases_pawn_value():
    cfg = TrainConfig(alpha=0.1)
    w = equal_material_weights()
    game = synthetic_game(one_hot(PAWN), 10, w, 0.0, cfg.beta)
    assert td_leaf_update(w, game, cfg, MATERIAL_MASK)[PAWN] < w[PAWN]


def test_black_advantage_and_black_win_also_increases_value():
    # features are White - Black, so a Black extra queen is -1 and a Black
    # win (result 0) should push the queen weight *up* as well.
    cfg = TrainConfig(alpha=0.1)
    w = equal_material_weights()
    game = synthetic_game(one_hot(QUEEN, -1.0), 10, w, 0.0, cfg.beta)
    assert td_leaf_update(w, game, cfg, MATERIAL_MASK)[QUEEN] > w[QUEEN]


def test_update_matches_hand_computation_lambda_zero():
    # lambda = 0 -> each step uses only its own one-step TD error.
    cfg = TrainConfig(alpha=1.0, lam=0.0, beta=1.0)
    w = equal_material_weights()
    f = one_hot(QUEEN)
    game = {"scores": [1.0], "feats": [f], "result": 1.0}
    J0 = sigmoid(1.0)
    expected = w[QUEEN] + 1.0 * (1.0 - J0) * J0 * (1.0 - J0)
    assert td_leaf_update(w, game, cfg, MATERIAL_MASK)[QUEEN] == pytest.approx(expected)


def test_masked_weights_are_frozen():
    cfg = TrainConfig(alpha=0.1)
    w = equal_material_weights()
    game = synthetic_game(one_hot(QUEEN), 5, w, 1.0, cfg.beta)
    frozen = [False] * N_FEATURES
    assert td_leaf_update(w, game, cfg, frozen) == w


def test_tiny_selfplay_run_learns_heavy_pieces_win():
    """Start every game with White having Q+2R vs a bare king.  The 1-ply
    learner mates quickly, so the queen and rook values must go up."""
    cfg = TrainConfig(games=3, workers=1, max_plies=80, epsilon=0.0,
                      random_opening_max=0, checkpoint_every=1,
                      start_fen="4k3/8/8/8/8/8/8/RR1QK3 w - - 0 1")
    w0 = equal_material_weights()
    out = train(w0, cfg, MATERIAL_MASK, positional=False, verbose=False)
    w = out["final"]
    assert out["stats"]["results"]["1-0"] == 3
    assert w[QUEEN] > w0[QUEEN]
    assert w[ALL_FEATURES.index("rook")] > w0[ALL_FEATURES.index("rook")]
    assert all(math.isfinite(x) for x in w)
    assert len(out["checkpoints"]) == 4          # games 0,1,2,3
    assert len(average_checkpoints(out["checkpoints"], 0.5)) == N_FEATURES


def test_selfplay_game_from_start_position_is_well_formed():
    cfg = TrainConfig(max_plies=30, qdepth=2)
    g = play_selfplay_game(equal_material_weights(), cfg, seed=5, positional=True)
    assert g["result"] in (0.0, 0.5, 1.0)
    assert len(g["scores"]) == len(g["feats"]) <= 30
    assert all(f is None or len(f) == N_FEATURES for f in g["feats"])