#!venv/bin/python3
from chess import Board
from chess.pgn import read_game
from time import time
from re import findall


def timer(func, *args, **kwargs):
	def new_func(*args, **kwargs):
		clock = time()
		res = func(*args, **kwargs)
		clock = round(time() - clock, 3)
		print(f'{clock} s')
		return res

	return new_func



INF = 9999
good_bots = {
	'Ethereal 12.75 64-bit 4CPU',
	'Komodo 14 64-bit 4CPU',
	'Stockfish 12 64-bit 4CPU',
	'Stockfish 2019-10-09 64-bit 4CPU',
	'ShashChess 15.0 64-bit',
	'asmFish 051117 64-bit 4CPU',
	'Stockfish 11 64-bit 4CPU',
	'Houdini 6 64-bit 4CPU',
	'SugaR NN 1.1 64-bit 4CPU',
	'Stockfish+NNUE 150720 64-bit',
	'Nemorino 6.00 64-bit 4CPU',
	'Stockfish 12 64-bit',
	'ShashChess 15.0 64-bit 4CPU'
}


def to_cp(s):
	s = s.strip()
	if '+M' in s:
		return (-INF if s[s.find('+M') + 2] == '-' else INF)
	reg = r'[\+\-]\d+\.\d+'
	s = findall(reg, s)
	if not s:
		return 0
	return round(float(s[0]) * 100)


@timer
def parse(filename):
	with open(filename, encoding='utf-8') as book, open('fens.txt', 'w', encoding='utf-8') as fens:
		cnt = 0
		while game := read_game(book):
			cnt += 1
			white, black = game.headers['White'], game.headers['Black']
			good_white = white in good_bots
			good_black = black in good_bots
			cur_white = True
			board = game.board()
			text = str(game.mainline_moves())
			moves = list(game.mainline_moves())
			comms = list(map(to_cp, findall(r'\{[^}]+\}', text)))
			comms = [0] * (len(moves) - len(comms)) + comms
			for move, com in zip(moves, comms):
				if good_white and cur_white:
					print(board.fen(), com, sep=';', file=fens)
				elif good_black and not cur_white:
					print(board.fen(), com, sep=';', file=fens)
				cur_white = not cur_white
				board.push(move)
			if not cnt % 1000:
				print(f'[+] reached {cnt}')


if __name__ == '__main__':
	try:
		parse('game.pgn')
	except KeyboardInterrupt:
		exit()
	except Exception as exp:
		print(exp)
#!venv/bin/python3
from chess import Board
from chess.pgn import read_game
from time import time
from re import findall


def timer(func, *args, **kwargs):
	def new_func(*args, **kwargs):
		clock = time()
		res = func(*args, **kwargs)
		clock = round(time() - clock, 3)
		print(f'{clock} s')
		return res

	return new_func



INF = 9999
good_bots = {
	'Ethereal 12.75 64-bit 4CPU',
	'Komodo 14 64-bit 4CPU',
	'Stockfish 12 64-bit 4CPU',
	'Stockfish 2019-10-09 64-bit 4CPU',
	'ShashChess 15.0 64-bit',
	'asmFish 051117 64-bit 4CPU',
	'Stockfish 11 64-bit 4CPU',
	'Houdini 6 64-bit 4CPU',
	'SugaR NN 1.1 64-bit 4CPU',
	'Stockfish+NNUE 150720 64-bit',
	'Nemorino 6.00 64-bit 4CPU',
	'Stockfish 12 64-bit',
	'ShashChess 15.0 64-bit 4CPU'
}


def to_cp(s):
	s = s.strip()
	if '+M' in s:
		return (-INF if s[s.find('+M') + 2] == '-' else INF)
	reg = r'[\+\-]\d+\.\d+'
	s = findall(reg, s)
	if not s:
		return 0
	return round(float(s[0]) * 100)


@timer
def parse(filename):
	with open(filename, encoding='utf-8') as book, open('fens.txt', 'w', encoding='utf-8') as fens:
		cnt = 0
		while game := read_game(book):
			cnt += 1
			white, black = game.headers['White'], game.headers['Black']
			good_white = white in good_bots
			good_black = black in good_bots
			cur_white = True
			board = game.board()
			text = str(game.mainline_moves())
			moves = list(game.mainline_moves())
			comms = list(map(to_cp, findall(r'\{[^}]+\}', text)))
			comms = [0] * (len(moves) - len(comms)) + comms
			for move, com in zip(moves, comms):
				if good_white and cur_white:
					print(board.fen(), com, sep=';', file=fens)
				elif good_black and not cur_white:
					print(board.fen(), com, sep=';', file=fens)
				cur_white = not cur_white
				board.push(move)
			if not cnt % 1000:
				print(f'[+] reached {cnt}')


if __name__ == '__main__':
	try:
		parse('game.pgn')
	except KeyboardInterrupt:
		exit()
	except Exception as exp:
		print(exp)
#!venv/bin/python3