import gradio as gr
import numpy as np
import torch
import json
from huggingface_hub import hf_hub_download
import torch.nn as nn



class BlokusTransformer(nn.Module):
    def __init__(self, 
                 d_max=20,      # maximum board dimension
                 embed_dim=128, # dimension of the transformer embeddings
                 num_heads=4,
                 mlp_dim=256,   # dimension of feedforward layer in TransformerEncoderLayer
                 num_layers=4,  # number of transformer layers
                 dropout=0.1):
        super().__init__()
        
        self.input_proj = nn.Linear(5, embed_dim)
        self.pos_embed = nn.Parameter(torch.zeros(d_max * d_max, embed_dim))
        nn.init.trunc_normal_(self.pos_embed, std=0.02)

        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
        
        encoder_layer = nn.TransformerEncoderLayer(d_model=embed_dim,
                                                   nhead=num_heads,
                                                   dim_feedforward=mlp_dim,
                                                   dropout=dropout,
                                                   batch_first=True)
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
        self.policy_head = nn.Linear(embed_dim, 1)
        self.value_head = nn.Linear(embed_dim, 4)
        
    def forward(self, x):
        """
        board_tensor: shape [batch_size, d, input_dim]
                      each cell has input_dim=5 features (4 occupancy bits + 1 legal bit)
        
        Returns:
            policy_logits: [batch_size, d, d]
            value:         [batch_size, 4]
        """
        batch_size, d, _ = x.shape 
        
        # Project to embedding dim: [batch_size, d*d, embed_dim]
        x = self.input_proj(x)
        
        # Add positional embeddings
        pos_embed_slice = self.pos_embed[:d, :].unsqueeze(0)  # [1, d, embed_dim]
        x = x + pos_embed_slice

        init_cls_tokens = self.cls_token.expand(batch_size, -1, -1)
        x = torch.cat([init_cls_tokens, x], dim=1)
        
        x = self.transformer(x)  # [batch_size, 1 + d, embed_dim]
        board_tokens = x[:, 1:, :]
        cls_tokens = x[:, 0, :]

        policy_logits = self.policy_head(board_tokens).view(batch_size, d) # [batch_size, d]
        value_logits = self.value_head(cls_tokens)  # [batch_size, 4]
        
        return policy_logits, value_logits

        
# Download your model from the Hub
model_path = hf_hub_download(
    repo_id="aracape/Blokus-Transformer", 
    filename="20x20.pt"
)

# Load your model
config = {
    "d_max": 20,
    "embed_dim": 128,
    "num_heads": 4,
    "mlp_dim": 256,
    "num_layers": 4,
    "dropout": 0.1
}    
model = BlokusTransformer(**config)
model.load_state_dict(torch.load(model_path, weights_only=True, map_location='cpu'))
model.eval()

def predict(input_data):
    """
    Process input and return prediction
    """
    try:
        # Parse input if it's JSON string
        if isinstance(input_data, str):
            input_data = json.loads(input_data)
        
        # Convert input
        input_np = np.array(input_data, dtype=np.float32)
        if input_np.size == 20 * 20 * 5:
            # Reshape and add batch dimension
            input_tensor = torch.from_numpy(input_np.reshape(1, 20 * 20, 5))
        else:
            return "Error: Input must be 20x20x5 = 2000 elements", "Error: Invalid input size"
        
        # Run inference
        with torch.no_grad():
            policies, values = model(input_tensor)

        policy = policies.squeeze().tolist()
        value = values.squeeze().tolist()
        
        return policy, value
        
    except Exception as e:
        return f"Error: {str(e)}"

# Create Gradio interface
iface = gr.Interface(
    fn=predict,
    inputs=gr.Textbox(label="Input (JSON format)", placeholder='{"data": [1, 2, 3, 4]}'),
    outputs=gr.Textbox(label="Prediction"),
    title="My PyTorch Model API"
)

# Launch with API access enabled
iface.launch(share=True)