import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import PreTrainedModel

# リモート読み込みとローカル直接インポートの両方に対応する互換インポート
try:
    from .configuration_dumbc import DumbcConfig
except ImportError:
    from configuration_dumbc import DumbcConfig

class DumbChessRetina(nn.Module):
    def __init__(self, dim=192):
        super().__init__()
        self.piece_embed = nn.Embedding(14, 32)
        self.conv_rank_file = nn.Conv1d(32, 32, kernel_size=3, padding=1, groups=32)
        self.conv_diag = nn.Conv1d(32, 32, kernel_size=3, padding=1, groups=32)
        self.conv_antidiag = nn.Conv1d(32, 32, kernel_size=3, padding=1, groups=32)
        self.conv_knight = nn.Conv1d(32, 32, kernel_size=3, padding=1, groups=32)
        self.tension_mlp = nn.Sequential(
            nn.Linear(32 * 4 + 1, 64),
            nn.GELU(),
            nn.Linear(64, dim)
        )

    def forward(self, board_state, material_weights):
        x = self.piece_embed(board_state)
        x_t = x.transpose(1, 2)
        f1 = self.conv_rank_file(x_t)
        f2 = self.conv_diag(x_t)
        f3 = self.conv_antidiag(x_t)
        f4 = self.conv_knight(x_t)
        f_all = torch.cat([f1, f2, f3, f4], dim=1).transpose(1, 2)
        tension_in = torch.cat([f_all, material_weights.unsqueeze(-1)], dim=-1)
        return self.tension_mlp(tension_in)

class DumbAttention(nn.Module):
    def __init__(self, dim=192, heads=8, bottleneck=32):
        super().__init__()
        self.dim = dim
        self.heads = heads
        self.head_dim = dim // heads
        self.qkv_proj = nn.Linear(dim, dim * 3, bias=False)
        self.out_proj = nn.Linear(dim, dim, bias=False)
        self.hadamard_mlp = nn.Sequential(
            nn.Linear(dim, bottleneck),
            nn.SiLU(),
            nn.Linear(bottleneck, 1)
        )
        self.trinity_g = nn.Linear(dim, 16, bias=False)
        self.temp_mlp = nn.Sequential(
            nn.Linear(dim, 16),
            nn.SiLU(),
            nn.Linear(16, 1)
        )
        self.threat_weight = nn.Parameter(torch.ones(1) * 0.5)

    def forward(self, x, mat_diff_matrix):
        B, N, C = x.shape
        q, k, v = self.qkv_proj(x).chunk(3, dim=-1)
        q_h = q.view(B, N, self.heads, self.head_dim).transpose(1, 2)
        k_h = k.view(B, N, self.heads, self.head_dim).transpose(1, 2)
        v_h = v.view(B, N, self.heads, self.head_dim).transpose(1, 2)
        S_base = (q_h @ k_h.transpose(-2, -1)) / math.sqrt(self.head_dim)
        
        q_k_hadamard = q.unsqueeze(2) * k.unsqueeze(1)
        S_tensor = self.hadamard_mlp(q_k_hadamard).squeeze(-1).unsqueeze(1)
        
        g_q = torch.sigmoid(self.trinity_g(q))
        g_k = torch.sigmoid(self.trinity_g(k))
        S_trinity = (g_q @ g_k.transpose(-2, -1)).unsqueeze(1)
        
        B_material = F.relu(mat_diff_matrix).unsqueeze(1) * self.threat_weight
        
        tau = torch.sigmoid(self.temp_mlp(x.mean(dim=1))) * 0.5 + 0.75
        tau = tau.unsqueeze(-1).unsqueeze(-1)
        
        S_total = (S_base + S_tensor + S_trinity + B_material) / tau
        A = F.softmax(S_total, dim=-1)
        out = (A @ v_h).transpose(1, 2).reshape(B, N, C)
        return self.out_proj(out)

class DumbFractalFFN(nn.Module):
    def __init__(self, dim=192, hidden_dim=288):
        super().__init__()
        self.w1 = nn.Linear(dim, hidden_dim, bias=False)
        self.w2 = nn.Linear(dim, hidden_dim, bias=False)
        self.w3 = nn.Linear(hidden_dim, dim, bias=False)

    def forward(self, x):
        h1 = F.silu(self.w1(x))
        h2 = self.w2(x)
        chunk_size = h2.shape[-1] // 2
        h2_a, h2_b = torch.split(h2, chunk_size, dim=-1)
        fractal_interaction = torch.cat([h2_a * h2_b, h2_b**2], dim=-1)
        return self.w3(h1 * fractal_interaction)

class DumbBlock(nn.Module):
    def __init__(self, dim=192):
        super().__init__()
        self.norm1 = nn.LayerNorm(dim)
        self.attn = DumbAttention(dim=dim)
        self.norm2 = nn.LayerNorm(dim)
        self.ffn = DumbFractalFFN(dim=dim)
        self.gate = nn.Parameter(torch.ones(1) * 0.1)

    def forward(self, x, mat_diff_matrix):
        x = x + self.gate * self.attn(self.norm1(x), mat_diff_matrix)
        x = x + self.gate * self.ffn(self.norm2(x))
        return x

class DumbcPreTrainedModel(PreTrainedModel):
    config_class = DumbcConfig
    base_model_prefix = "dumbc"

    def _init_weights(self, module):
        if isinstance(module, (nn.Linear, nn.Conv1d)):
            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
            if module.bias is not None:
                module.bias.data.zero_()

class DumbcModel(DumbcPreTrainedModel):
    def __init__(self, config):
        super().__init__(config)
        self.config = config
        self.retina = DumbChessRetina(dim=config.dim)
        self.layers = nn.ModuleList([DumbBlock(dim=config.dim) for _ in range(config.num_unique_layers)])
        self.step_embed = nn.Parameter(torch.randn(config.num_loops, 1, 1, config.dim) * 0.02)
        
        self.from_head = nn.Linear(config.dim, 64)
        self.to_head = nn.Linear(config.dim, 64)
        self.value_head = nn.Sequential(
            nn.Linear(config.dim, 64),
            nn.GELU(),
            nn.Linear(64, 3)
        )
        self.post_init()

    def forward(self, board_state, mat_diff_matrix, material_weights, **kwargs):
        x = self.retina(board_state, material_weights)
        for loop_idx in range(self.config.num_loops):
            x = x + self.step_embed[loop_idx]
            for layer in self.layers:
                x = layer(x, mat_diff_matrix)
                
        from_logits = self.from_head(x)
        to_logits = self.to_head(x)
        policy_matrix = torch.bmm(from_logits, to_logits.transpose(1, 2))
        
        global_pool = x.mean(dim=1)
        value_logits = self.value_head(global_pool)
        
        return {"policy_matrix": policy_matrix, "value_logits": value_logits}