#!/usr/bin/env python3
"""Parallax-Chess v3: Predict move as single index (from_sq * 64 + to_sq).
Fixes the v2 bug where only from-square was supervised.
"""
import sys, math, time, json, random
from pathlib import Path
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
import chess

ROOT = Path(__file__).parent.parent
sys.path.insert(0, str(ROOT / "scripts"))
sys.path.insert(0, str(ROOT / "src"))
sys.path.insert(0, str(ROOT.parent))
from centauri.models.base import SmallLMConfig
from centauri.models.torch.small_lm import SmallLM
from chess_encoding import encode_board, VOCAB_SIZE

random.seed(42)

# Move index: from_sq * 64 + to_sq for non-promo, then promo moves
# Total: 64*64 = 4096 normal + 4*8*8 = 256 promo = 4352 total
MOVE_INDEX_SIZE = 4352  # 64*64 + 4*64 (promotions)


def move_to_index(move):
    """Convert chess.Move to index 0..4351."""
    if move.promotion:
        promo_map = {chess.QUEEN: 0, chess.ROOK: 1, chess.BISHOP: 2, chess.KNIGHT: 3}
        return 4096 + promo_map[move.promotion] * 64 + move.to_square
    return move.from_square * 64 + move.to_square


def index_to_move(idx, board):
    """Convert index back to chess.Move, only if legal."""
    if idx >= 4096:
        promo_offset = idx - 4096
        promo_type = promo_offset // 64
        to_sq = promo_offset % 64
        promo_map = {0: chess.QUEEN, 1: chess.ROOK, 2: chess.BISHOP, 3: chess.KNIGHT}
        promotion = promo_map[promo_type]
        # Find from_square (must be a pawn on correct file)
        for from_sq in chess.SQUARES:
            piece = board.piece_at(from_sq)
            if piece and piece.piece_type == chess.PAWN and piece.color == board.turn:
                move = chess.Move(from_sq, to_sq, promotion=promotion)
                if move in board.legal_moves:
                    return move
        return None
    else:
        from_sq = idx // 64
        to_sq = idx % 64
        move = chess.Move(from_sq, to_sq)
        if move in board.legal_moves:
            return move
        return None


class ChessDataset(Dataset):
    """Load SF-evaluated positions."""

    def __init__(self, data_paths):
        self.samples = []
        for path in data_paths:
            path = Path(path)
            if not path.exists():
                continue
            with open(path, "r", encoding="utf-8") as f:
                for line in f:
                    try:
                        item = json.loads(line)
                        board = chess.Board(item["fen"])
                        move_key = "best_move" if "best_move" in item else "move"
                        move = chess.Move.from_uci(item[move_key])
                        if move not in board.legal_moves:
                            continue

                        board_tokens = encode_board(board)
                        move_idx = move_to_index(move)

                        if "eval_cp" in item:
                            value = max(-1.0, min(1.0, item["eval_cp"] / 1000.0))
                        elif "result" in item:
                            r = item["result"]
                            value = 1.0 if r == "1-0" else -1.0 if r == "0-1" else 0.0
                        else:
                            value = 0.0

                        self.samples.append((board_tokens, move_idx, value))

                        if random.random() < 0.5:
                            flipped = flip_board(board_tokens)
                            flipped_move = flip_move_index(move_idx)
                            self.samples.append((flipped, flipped_move, value))
                    except:
                        pass
        print("Loaded %d samples" % len(self.samples))

    def __len__(self):
        return len(self.samples)

    def __getitem__(self, i):
        board_tokens, move_idx, value = self.samples[i]
        return (
            torch.tensor(board_tokens, dtype=torch.long),
            torch.tensor(move_idx, dtype=torch.long),
            torch.tensor(value, dtype=torch.float),
        )


def flip_board(board_tokens):
    flipped = list(board_tokens)
    for r in range(8):
        for f in range(8):
            flipped[r * 8 + (7 - f)] = board_tokens[r * 8 + f]
    return flipped


def flip_sq(sq):
    """Mirror a square horizontally: file f -> 7-f."""
    rank = sq // 8
    file = sq % 8
    return rank * 8 + (7 - file)


def flip_move_index(idx):
    """Flip move index when board is mirrored horizontally."""
    if idx >= 4096:
        promo_offset = idx - 4096
        promo_type = promo_offset // 64
        to_sq = promo_offset % 64
        return 4096 + promo_type * 64 + flip_sq(to_sq)
    from_sq = idx // 64
    to_sq = idx % 64
    return flip_sq(from_sq) * 64 + flip_sq(to_sq)


class ParallaxChessV3(nn.Module):
    """Board -> move index (4352 classes) + value."""

    def __init__(self, cfg):
        super().__init__()
        self.cfg = cfg
        self.board_encoder = nn.Embedding(VOCAB_SIZE, cfg.d_model)
        from centauri.models.torch.small_lm import SmallLM
        self.backbone = SmallLM(cfg)
        self.policy_head = nn.Linear(cfg.d_model, MOVE_INDEX_SIZE)
        self.value_head = nn.Sequential(
            nn.Linear(cfg.d_model, 256), nn.ReLU(), nn.Dropout(0.1),
            nn.Linear(256, 1), nn.Tanh())

    def forward(self, board_tokens):
        x = self.board_encoder(board_tokens)
        freq, _ = __import__('centauri.models.torch.small_lm', fromlist=['precompute_rope']).precompute_rope(
            self.cfg, x.device, x.dtype)
        for layer in self.backbone.layers:
            x, _ = layer(x, freq, None)
        pooled = x.mean(dim=1)
        return {"logits": self.policy_head(pooled), "value": self.value_head(pooled).squeeze(-1)}

    def predict_move(self, board, device=None):
        if device is None:
            device = next(self.parameters()).device
        board_tokens = encode_board(board)
        inp = torch.tensor([board_tokens], dtype=torch.long).to(device)
        with torch.no_grad():
            out = self(inp)
            probs = torch.softmax(out["logits"], dim=-1)

        # Score only legal moves
        best_move, best_score = None, -float("inf")
        for move in board.legal_moves:
            idx = move_to_index(move)
            score = probs[0, idx].item()
            if score > best_score:
                best_score = score
                best_move = move
        return best_move

    def evaluate(self, board, device=None):
        if device is None:
            device = next(self.parameters()).device
        board_tokens = encode_board(board)
        inp = torch.tensor([board_tokens], dtype=torch.long).to(device)
        with torch.no_grad():
            return self(inp)["value"].item() * 1000


def main():
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("--batch_size", type=int, default=64)
    parser.add_argument("--lr", type=float, default=1e-3)
    parser.add_argument("--max_steps", type=int, default=80000)
    parser.add_argument("--save_dir", default=str(ROOT / "checkpoints" / "parallax_chess_v3"))
    args = parser.parse_args()

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print("Device:", device)

    cfg_dict = {
        "vocab_size": VOCAB_SIZE, "d_model": 512, "n_heads": 8, "n_kv_heads": 4,
        "n_layers": 8, "intermediate_size": 2048, "max_seq_len": 128,
        "norm_type": "rms", "rope_type": "neox", "n_experts": 0,
        "n_loops": 1, "loop_mode": "per_layer",
    }
    cfg = SmallLMConfig(**cfg_dict)
    model = ParallaxChessV3(cfg).to(device)
    n_params = sum(p.numel() for p in model.parameters())
    print("Params: %d (%.1fM)" % (n_params, n_params / 1e6))

    data_paths = [
        ROOT / "data" / "chess_train_sf.jsonl",
        ROOT / "data" / "chess_train_large.jsonl",
    ]
    ds = ChessDataset(data_paths)
    loader = DataLoader(ds, batch_size=args.batch_size, shuffle=True, num_workers=0, pin_memory=True)

    save_dir = Path(args.save_dir)
    save_dir.mkdir(parents=True, exist_ok=True)

    step = 0
    for f in sorted(save_dir.glob("step_*.pt")):
        try:
            ckpt = torch.load(str(f), map_location=device)
            model.load_state_dict(ckpt["model"])
            step = ckpt.get("step", 0)
            print("Resumed from %s at step %d" % (f.name, step))
            break
        except:
            pass

    total_steps = min(args.max_steps, len(ds) // args.batch_size * 5)
    opt = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=0.01, betas=(0.9, 0.98))

    def lr_schedule(s):
        warmup = 2000
        if s < warmup:
            return s / warmup
        progress = (s - warmup) / max(1, total_steps - warmup)
        return max(0.05, 0.5 * (1.0 + math.cos(math.pi * progress)))

    sched = torch.optim.lr_scheduler.LambdaLR(opt, lr_schedule)

    model.train()
    t0 = time.time()

    for epoch in range(999):
        for board_tokens, move_indices, values in loader:
            board_tokens = board_tokens.to(device)
            move_indices = move_indices.to(device)
            values = values.to(device)

            out = model(board_tokens)

            policy_loss = F.cross_entropy(out["logits"], move_indices)
            value_loss = F.mse_loss(out["value"], values)
            loss = policy_loss + 1.0 * value_loss

            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            opt.step()
            opt.zero_grad()
            sched.step()

            if step % 100 == 0:
                lr = sched.get_last_lr()[0]
                print("Step %6d | P:%.4f V:%.4f | Loss:%.4f | LR %.1e" % (
                    step, policy_loss.item(), value_loss.item(), loss.item(), lr))

            if step > 0 and step % 10000 == 0:
                ckpt = save_dir / ("step_%d.pt" % step)
                torch.save({"model": model.state_dict(), "step": step, "config": cfg_dict, "n_params": n_params}, str(ckpt))
                print("Saved %s" % ckpt.name)

            step += 1
            if step >= total_steps:
                break

    final = save_dir / "final.pt"
    torch.save({"model": model.state_dict(), "step": step, "config": cfg_dict, "n_params": n_params}, str(final))
    print("Done! %s (%d steps)" % (final, step))


if __name__ == "__main__":
    main()