import berserk
import torch
from torch import nn
from torch import optim
import chess
import os
import pandas as pd

device = torch.device("cpu")

def board_to_tensor(board):
    piece_encoding = {
        'P': 1, 'N': 2, 'B': 3, 'R': 4, 'Q': 5, 'K': 6,
        'p': 7, 'n': 8, 'b': 9, 'r': 10, 'q': 11, 'k': 12
    }

    tensor = torch.zeros(64, dtype=torch.long)
    for square in chess.SQUARES:
        piece = board.piece_at(square)
        if piece:
            tensor[square] = piece_encoding[piece.symbol()]
        else:
            tensor[square] = 0

    return tensor.unsqueeze(0)

class BOT(nn.Module):
    def __init__(self):
        super().__init__()
        self.embedding = nn.Embedding(13,64)
        self.attention = nn.MultiheadAttention(embed_dim=64,num_heads=16)
        self.neurons = nn.Sequential(
            nn.Linear(4096,128),
            nn.ReLU(),
            nn.Linear(128,128),
            nn.ReLU(),
            nn.Linear(128,128),
            nn.ReLU(),
            nn.Linear(128,64),
            nn.ReLU(),
            nn.Linear(64,1)
        )
    def forward(self,x):
        x = self.embedding(x)
        x = x.permute(1, 0, 2)
        attn_output, _ = self.attention(x, x, x)
        x = attn_output.permute(1, 0, 2).contiguous()
        x = x.view(x.size(0), -1)
        x = self.neurons(x)
        return x
model = BOT().to(device)
model = torch.compile(model,mode="max-autotune",dynamic=False)
if os.path.exists("booty.pth"):
    file = torch.load("booty.pth",map_location=device,weights_only=True)
    model.load_state_dict(file)
model.train()
optimizer = optim.Adam(model.parameters(),lr=1e-4)
criterion = nn.MSELoss()
num_epochs = 1
df = pd.read_csv("lichess_db_puzzle.csv",nrows=1000)
df = df.sort_values(by="Rating", ascending=True)
t1 = torch.tensor([10.0], dtype=torch.float32, device=device)
t2 = torch.tensor([-10.0], dtype=torch.float32, device=device)
t3 = torch.tensor([0.0], dtype=torch.float32, device=device)
for i in range(num_epochs):
    total_loss = 0.0
    for puzzle in df.iloc:
     board = chess.Board(puzzle["FEN"])
     print(f"Rating: {puzzle['Rating']} elo.")
     n = 0
     for move in puzzle["Moves"].split():
         if n % 2 == 0:
             for movey in list(board.legal_moves):
                 if str(movey.uci()) == move:
                     b = True
                 else:
                     b = False
                 board.push(movey)
                 tensor = board_to_tensor(board).to(device)
                 evaling = model(tensor)
                 board.pop()
                 optimizer.zero_grad()
                 if b and board.turn == chess.WHITE:
                     loss = criterion(t1, evaling)
                 elif b and board.turn == chess.BLACK:
                     loss = criterion(t2, evaling)
                 else:
                     loss = criterion(t3, evaling)
                 total_loss += loss.item()
                 loss.backward()
                 optimizer.step()
         n += 1
         board.push_uci(move)
    print(f"Epoch [{i+1}/{num_epochs}], Loss: {total_loss}.")
    torch.save(model.state_dict(),"booty.pth")