import math
import random
import sys

import chess
import chess.pgn
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer


def clean_san(s):
    return s.replace("+", "").replace("#", "")


class Model:
    def __init__(self, path, device, temp=1.0):
        self.tok = AutoTokenizer.from_pretrained(path)
        self.model = AutoModelForCausalLM.from_pretrained(path).to(device).eval()
        self.device = device
        self.temp = temp
        self.name = path.rstrip("/").split("/")[-1]
        self.prompt_ids = self.tok("1.", return_tensors="pt").input_ids.to(device)

    @torch.no_grad()
    def pick(self, board):
        legal = list(board.legal_moves)
        seqs = [tuple(self.tok(clean_san(board.san(m)) + " ",
                               add_special_tokens=False).input_ids) for m in legal]
        active = list(range(len(seqs)))
        step = 0
        cur, past = self.prompt_ids, None
        while True:
            out = self.model(cur, past_key_values=past, use_cache=True)
            past = out.past_key_values
            logits = out.logits[0, -1]
            allowed = {}
            for i in active:
                if step < len(seqs[i]):
                    allowed.setdefault(seqs[i][step], []).append(i)
            if not allowed:
                return legal[active[0]]
            toks = list(allowed.keys())
            probs = torch.softmax(logits[torch.tensor(toks, device=self.device)] / self.temp, 0)
            choice = toks[torch.multinomial(probs, 1).item()]
            active = allowed[choice]
            step += 1
            done = [i for i in active if len(seqs[i]) == step]
            if done:
                return legal[done[0]]
            cur = torch.tensor([[choice]], device=self.device)


class Rand:
    name = "random"

    def pick(self, board):
        return random.choice(list(board.legal_moves))


def make(p, d):
    return Rand() if p == "random" else Model(p, d)


def main():
    a, b, n = sys.argv[1], sys.argv[2], int(sys.argv[3])
    pgn_out = sys.argv[4] if len(sys.argv) > 4 else None
    d = "cuda" if torch.cuda.is_available() else "cpu"
    A, B = make(a, d), make(b, d)
    print("blind arena: %s vs %s, %d games" % (A.name, B.name, n), flush=True)
    w = l = dr = 0
    pgns = []
    for g in range(n):
        board = chess.Board()
        a_white = g % 2 == 0
        white, black = (A, B) if a_white else (B, A)
        plies = 0
        while not board.is_game_over(claim_draw=True) and plies < 250:
            board.push((white if board.turn == chess.WHITE else black).pick(board))
            plies += 1
        res = board.result(claim_draw=True)
        apt = 0.5 if res == "1/2-1/2" or res == "*" else (
            1.0 if (res == "1-0") == a_white else 0.0)
        w += apt == 1.0
        l += apt == 0.0
        dr += apt == 0.5
        if pgn_out is not None:
            game = chess.pgn.Game.from_board(board)
            game.headers.update(White=A.name if a_white else B.name,
                                 Black=B.name if a_white else A.name, Result=res)
            pgns.append(str(game))
        print("game %d/%d: %s  [%s W%d L%d D%d]" % (g + 1, n, res, A.name, int(w), int(l), int(dr)), flush=True)
    score = w + 0.5 * dr
    p = score / n
    se = math.sqrt(0.25 / n)
    z = (p - 0.5) / se
    print("\nFINAL: %s vs %s = %dW %dL %dD over %d games" % (A.name, B.name, int(w), int(l), int(dr), n))
    print("  score %.1f/%d = %.1f%%  z=%.2f" % (score, n, 100 * p, z))
    if pgn_out is not None:
        with open(pgn_out, "w") as f:
            f.write("\n\n".join(pgns))
        print("  PGNs -> %s" % pgn_out)


if __name__ == "__main__":
    main()