import argparse
import random
from pathlib import Path

import numpy as np
from draughts import RussianBoard
from draughts.models import Color

from network import FixedNetwork
from rules import Adjudicator


def encode(board,perspective):
    masks=(board.white_men,board.white_kings,board.black_men,board.black_kings)
    if perspective==Color.BLACK:
        masks=masks[2:]+masks[:2]
    x=np.zeros(128,np.float32)
    for p,mask in enumerate(masks):
        while mask:
            bit=mask&-mask
            square=bit.bit_length()-1
            x[p*32+(square if perspective==Color.WHITE else 31-square)]=1
            mask^=bit
    return x


class Agent:
    def __init__(self,model_dir='.',seed=1):
        self.network=FixedNetwork.load(model_dir)
        self.rng=random.Random(seed)

    def choose(self,board):
        legal=board.legal_moves
        if not legal:
            raise ValueError('No legal moves')
        who=board.turn
        batch=[]
        for move in legal:
            board.push(move)
            batch.append(encode(board,who))
            board.pop()
        scores=self.network.forward(np.stack(batch))
        best=np.flatnonzero(np.isclose(scores,scores.max(),atol=1e-5,rtol=0))
        return legal[self.rng.choice(best.tolist())]


def notation(move):
    return (':' if move.captured_list else '-').join(RussianBoard.SQUARE_NAMES[s] for s in move.square_list)


def main():
    ap=argparse.ArgumentParser(description='Fixed-weight neural Russian draughts')
    ap.add_argument('--model-dir',default=str(Path(__file__).resolve().parent))
    ap.add_argument('--fen')
    ap.add_argument('--seed',type=int,default=1)
    ap.add_argument('--human',choices=['white','black'])
    ap.add_argument('--one-move',action='store_true')
    args=ap.parse_args()
    board=RussianBoard.from_fen(args.fen) if args.fen else RussianBoard()
    agent=Agent(args.model_dir,args.seed)
    if args.one_move:
        print(notation(agent.choose(board)))
        return
    human=Color.WHITE if args.human=='white' else Color.BLACK if args.human=='black' else None
    ai_side=Color.BLACK if human==Color.WHITE else Color.WHITE
    random_opponent=random.Random(args.seed+1)
    referee=Adjudicator(board,max_plies=300)
    while True:
        print(board)
        legal=board.legal_moves
        outcome=referee.outcome(board,legal)
        if outcome is not None:
            winner,reason=outcome
            print('Draw' if winner==0 else 'White wins' if winner==-1 else 'Black wins',reason)
            return
        side=board.turn
        if side==human:
            for i,move in enumerate(legal):
                print(i+1,notation(move))
            raw=input('Move number (q to quit): ').strip()
            if raw.lower()=='q':
                return
            try:
                index=int(raw)-1
                if index<0:
                    raise ValueError()
                move=legal[index]
            except (ValueError,IndexError):
                print('Choose a move number from the list')
                continue
        else:
            move=agent.choose(board) if side==ai_side else random_opponent.choice(legal)
            print('Move:',notation(move))
        piece=board._get(move.square_list[0])
        board.push(move)
        referee.update(board,move,piece,side)


if __name__=='__main__':
    main()