"""Engine tests: mates, legality, evaluation symmetry, limits."""
import random

import chess
import pytest

from engine import (ALL_FEATURES, MATE, N_FEATURES, Evaluator, Searcher,
                    features, implied_square_values, is_mate_score, leaf_board,
                    pv_to_san, standard_weights)


def searcher(weights=None, **kw):
    return Searcher(Evaluator(weights or standard_weights()), **kw)


MATE_IN_1 = [
    # scholar's mate pattern
    ("r1bqkbnr/pppp1ppp/2n5/4p3/2B1P3/5Q2/PPPP1PPP/RNB1K1NR w KQkq - 0 1", "f3f7"),
    # back-rank mate
    ("6k1/5ppp/8/8/8/8/5PPP/3R2K1 w - - 0 1", "d1d8"),
    # black to move: fool's mate
    ("rnbqkbnr/pppp1ppp/8/4p3/6P1/5P2/PPPPP2P/RNBQKBNR b KQkq g3 0 2", "d8h4"),
]


@pytest.mark.parametrize("fen,expected", MATE_IN_1)
@pytest.mark.parametrize("depth", [1, 3])
def test_mate_in_one(fen, expected, depth):
    board = chess.Board(fen)
    r = searcher().search(board, depth)
    board.push(r.move)
    assert board.is_checkmate()
    if expected:
        assert r.move.uci() == expected
    assert r.score == MATE - 1


MATE_IN_2 = [
    # Morphy-style problem: 1.Ra6! bxa6 2.b7#
    ("kbK5/pp6/1P6/8/8/8/8/R7 w - - 0 1", "a1a6"),
    # Legal's-mate-like: 1.Nf6+ gxf6 2.Bxf7#
    ("r2qkb1r/pp2nppp/3p4/2pNN1B1/2BnP3/3P4/PPP2PPP/R2bK2R w KQkq - 1 1", "d5f6"),
]


@pytest.mark.parametrize("fen,expected", MATE_IN_2)
def test_mate_in_two(fen, expected):
    board = chess.Board(fen)
    r = searcher().search(board, 3)
    assert r.move.uci() == expected
    assert r.score == MATE - 3
    # the PV really is a forced mate line
    b = board.copy()
    for m in r.pv:
        assert m in b.legal_moves
        b.push(m)
    assert b.is_checkmate()


def test_mate_in_two_only_found_with_enough_depth():
    board = chess.Board(MATE_IN_2[0][0])
    r1 = searcher().search(board, 1)
    assert not is_mate_score(r1.score)


def test_search_moves_are_legal_on_random_positions():
    rng = random.Random(1)
    s = searcher(qdepth=4)
    for _ in range(25):
        b = chess.Board()
        for _ in range(rng.randint(0, 60)):
            if b.is_game_over():
                break
            b.push(rng.choice(list(b.legal_moves)))
        if b.is_game_over():
            continue
        r = s.search(b, 2, node_limit=3000)
        assert r.move in b.legal_moves
        bb = b.copy()
        for m in r.pv:                       # every PV move is legal in sequence
            assert m in bb.legal_moves
            bb.push(m)


def test_node_limit_is_respected_and_still_returns_a_move():
    b = chess.Board()
    r = searcher().search(b, 6, node_limit=500)
    assert r.move in b.legal_moves
    assert r.nodes <= 520


def test_no_legal_moves():
    stalemate = chess.Board("7k/5Q2/6K1/8/8/8/8/8 b - - 0 1")
    r = searcher().search(stalemate, 2)
    assert r.move is None and r.score == 0.0


def test_features_are_antisymmetric():
    rng = random.Random(3)
    for _ in range(10):
        b = chess.Board()
        for _ in range(rng.randint(5, 40)):
            if b.is_game_over():
                break
            b.push(rng.choice(list(b.legal_moves)))
        f = features(b)
        g = features(b.mirror())           # colours swapped and board flipped
        assert len(f) == N_FEATURES == len(ALL_FEATURES)
        assert f == pytest.approx([-x for x in g])


def test_start_position_is_balanced_and_material_counts():
    assert features(chess.Board()) == pytest.approx([0.0] * N_FEATURES)
    b = chess.Board("4k3/8/8/8/8/8/PPP5/RNBQK3 w - - 0 1")   # white: 3P N B R Q
    assert features(b, positional=False)[:5] == [3, 1, 1, 1, 1]


def test_quiescence_avoids_hanging_queen():
    # Qxd5?? would lose the queen to ...exd5? No: the black pawn on e6 guards d5.
    b = chess.Board("4k3/8/4p3/3p4/8/8/8/3QK3 w - - 0 1")
    r = searcher().search(b, 1)
    assert r.move.uci() != "d1d5"


def test_implied_square_values_and_pv_text():
    w = standard_weights()
    w[ALL_FEATURES.index("knight_center")] = 0.6
    grid = implied_square_values(w, chess.KNIGHT)
    assert grid[3][3] > grid[0][0]             # d4 better than a1
    assert grid[0][0] == pytest.approx(3.0)      # mobility weight is 0 here
    b = chess.Board()
    s = pv_to_san(b, [chess.Move.from_uci("e2e4"), chess.Move.from_uci("e7e5")])
    assert s == "1. e4 e5"
    assert leaf_board(b, [chess.Move.from_uci("e2e4")]).piece_at(chess.E4) is not None