import torch
import torch.nn.functional as F
import yaml
import os
import chess
import chess.engine
import random
import time
from model import ChessTransformer
from data_loader import VOCAB

# Reverse vocab for decoding
INV_VOCAB = {v: k for k, v in VOCAB.items()}

# Load configuration
with open('config.yaml', 'r') as f:
    config = yaml.safe_load(f)

model_cfg = config['model']
data_cfg = config['data']
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

# Initialize Transformer
model = ChessTransformer(
    vocab_size=len(VOCAB),
    d_model=model_cfg['d_model'],
    nhead=model_cfg['nhead'],
    num_layers=model_cfg['num_layers'],
    max_length=data_cfg['max_length']
).to(device)

# Load latest weights
weights_path = 'weights/sf_rl_latest.pth'
if not os.path.exists(weights_path):
    weights_path = 'rl_weights/champion_latest.pth'
    if not os.path.exists(weights_path):
        weights_path = config['training']['weights_path']

if os.path.exists(weights_path):
    checkpoint = torch.load(weights_path, map_location=device, weights_only=True)
    if isinstance(checkpoint, dict) and 'model_state_dict' in checkpoint:
        model.load_state_dict(checkpoint['model_state_dict'])
    else:
        model.load_state_dict(checkpoint)
    print(f"Loaded weights from {weights_path}")
else:
    print("Warning: No pre-trained weights found. Training from scratch!")

model.train()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-5)

# Initialize Stockfish Oracle
stockfish_path = os.path.join("stockfish", "stockfish-windows-x86-64-avx2.exe")
if not os.path.exists(stockfish_path):
    raise FileNotFoundError(f"Stockfish not found at {stockfish_path}")

engine = chess.engine.SimpleEngine.popen_uci(stockfish_path)
# Optional: Set hash size for stockfish to be fast
engine.configure({"Hash": 256, "Threads": 2})
LIMIT = chess.engine.Limit(time=0.1) # 100ms per evaluation

def encode_history(history, max_length=120):
    seq = [VOCAB.get(tok, VOCAB.get("<unk>", 0)) for tok in history]
    if len(seq) > max_length:
        seq = seq[-max_length:]
    else:
        seq = seq + [0] * (max_length - len(seq))
    return torch.tensor(seq, dtype=torch.long, device=device).unsqueeze(0)

def evaluate_position(board):
    info = engine.analyse(board, LIMIT)
    score = info["score"].white()
    # If checkmate is forced, cap at +/- 10000
    if score.is_mate():
        return 10000 if score.mate() > 0 else -10000
    return score.score()

print("\nStarting Dense Policy Gradient RL with Stockfish...")

epoch = 0
os.makedirs("weights", exist_ok=True)

try:
    while True:
        epoch += 1
        board = chess.Board()
        history = ["<bos>"]
        
        states = []
        actions = []
        advantages = []
        
        print(f"\n--- Game {epoch} ---")
        
        current_eval_cp = evaluate_position(board)
        
        while not board.is_game_over() and len(history) < 150:
            # 1. Forward Pass to get Policy
            with torch.no_grad():
                x = encode_history(history)
                policy_logits, _ = model(x)
                p_logits = policy_logits[0, -1, :]
            
            # 2. Sample Move
            temperature = 1.0
            logits = p_logits / temperature
            legal_moves = list(board.legal_moves)
            legal_ucis = [m.uci() for m in legal_moves]
            
            mask = torch.full_like(logits, float('-inf'))
            for idx, token in INV_VOCAB.items():
                if token in legal_ucis:
                    mask[idx] = 0.0
            
            masked_logits = logits + mask
            probs = F.softmax(masked_logits, dim=-1)
            
            if torch.isnan(probs).any() or probs.sum() == 0:
                move = random.choice(legal_moves)
            else:
                m_idx = torch.multinomial(probs, 1).item()
                action_str = INV_VOCAB.get(m_idx)
                if action_str not in legal_ucis:
                    move = random.choice(legal_moves)
                else:
                    move = chess.Move.from_uci(action_str)
                    
            # 3. Calculate Dense Reward using Stockfish
            board.push(move)
            next_eval_cp = evaluate_position(board)
            
            if not board.turn: # It WAS white's turn
                adv = next_eval_cp - current_eval_cp
            else: # It WAS black's turn
                adv = current_eval_cp - next_eval_cp
                
            current_eval_cp = next_eval_cp
            
            # Store for training
            states.append(x.squeeze(0))
            actions.append(VOCAB.get(move.uci(), 0))
            advantages.append(adv)
            
            history.append(move.uci())
            
            if len(history) % 10 == 0:
                print(f"Move {len(history)} | Advantage: {adv} cp | Eval: {current_eval_cp}")
                
        # --- Policy Gradient Update ---
        print(f"Game finished. Training on {len(states)} moves...")
        
        states_t = torch.stack(states).to(device)
        actions_t = torch.tensor(actions, dtype=torch.long, device=device)
        # Normalize advantages to stabilize training (like PPO)
        adv_t = torch.tensor(advantages, dtype=torch.float32, device=device)
        if adv_t.std() > 0:
            adv_t = (adv_t - adv_t.mean()) / (adv_t.std() + 1e-8)
        else:
            adv_t = adv_t - adv_t.mean()
            
        optimizer.zero_grad()
        
        # Forward pass for the whole sequence
        policy_logits, _ = model(states_t)
        
        # We want the logits of the *last* token of each state in the batch
        p_logits = policy_logits[:, -1, :]
        
        # Log probability of the actions we took
        log_probs = F.log_softmax(p_logits, dim=-1)
        action_log_probs = log_probs.gather(1, actions_t.unsqueeze(1)).squeeze(1)
        
        # Loss = -log_prob * advantage
        loss = -(action_log_probs * adv_t).mean()
        
        loss.backward()
        
        # Gradient clipping
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()
        
        print(f"Loss: {loss.item():.4f} | Avg Advantage: {sum(advantages)/len(advantages):.2f} cp")
        
        # Save model periodically
        if epoch % 5 == 0:
            torch.save(model.state_dict(), 'weights/sf_rl_latest.pth')
            print("Saved weights to weights/sf_rl_latest.pth")

except KeyboardInterrupt:
    print("Training stopped manually.")
finally:
    engine.quit()
    print("Stockfish engine closed.")