import torch
import torch.nn as nn
import torch.nn.functional as F

"""
加入混合专家
"""


class DynamicGraphConstructor(nn.Module):
    def __init__(self, node_num, embed_dim):
        super().__init__()
        self.nodevec1 = nn.Parameter(torch.randn(node_num, embed_dim))
        self.nodevec2 = nn.Parameter(torch.randn(node_num, embed_dim))

    def forward(self):
        adj = torch.matmul(self.nodevec1, self.nodevec2.T)
        adj = F.relu(adj)
        adj = F.softmax(adj, dim=-1)
        return adj


class GraphConvBlock(nn.Module):
    def __init__(self, input_dim, output_dim):
        super().__init__()
        self.theta = nn.Linear(input_dim, output_dim)
        self.residual = input_dim == output_dim
        if not self.residual:
            self.res_proj = nn.Linear(input_dim, output_dim)

    def forward(self, x, adj):
        res = x
        x = torch.matmul(adj, x)
        x = self.theta(x)
        if self.residual:
            x = x + res
        else:
            x = x + self.res_proj(res)
        return F.relu(x)


class MANBA_Block(nn.Module):
    def __init__(self, input_dim, hidden_dim):
        super().__init__()
        self.attn = nn.MultiheadAttention(
            embed_dim=input_dim, num_heads=4, batch_first=True
        )
        self.ffn = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, input_dim),
        )
        self.norm1 = nn.LayerNorm(input_dim)
        self.norm2 = nn.LayerNorm(input_dim)

    def forward(self, x):
        res = x
        x_attn, _ = self.attn(x, x, x)
        x = self.norm1(res + x_attn)

        res2 = x
        x_ffn = self.ffn(x)
        x = self.norm2(res2 + x_ffn)
        return x


class EXPExpert(nn.Module):  # 原 EXP 改名
    def __init__(self, args):
        super().__init__()
        self.horizon = args["horizon"]
        self.output_dim = args["output_dim"]
        self.seq_len = args.get("in_len", 12)
        self.hidden_dim = args.get("hidden_dim", 64)
        self.num_nodes = args["num_nodes"]

        self.graph = DynamicGraphConstructor(self.num_nodes, embed_dim=16)
        self.input_proj = nn.Linear(self.seq_len, self.hidden_dim)
        self.gc = GraphConvBlock(self.hidden_dim, self.hidden_dim)
        self.manba = MANBA_Block(self.hidden_dim, self.hidden_dim * 2)
        self.out_proj = nn.Linear(self.hidden_dim, self.horizon * self.output_dim)

    def forward(self, x):
        x = x[..., 0]  # (B, T, N)
        B, T, N = x.shape
        x = x.permute(0, 2, 1).reshape(B * N, T)
        h = self.input_proj(x).view(B, N, -1)
        adj = self.graph()
        h = self.gc(h, adj)
        h = self.manba(h)
        out = self.out_proj(h)
        return out.view(B, N, self.horizon, self.output_dim).permute(0, 2, 1, 3)


class EXP(nn.Module):
    def __init__(self, args, num_experts=4, top_k=2):
        super().__init__()
        self.num_experts = num_experts
        self.top_k = top_k
        self.experts = nn.ModuleList([EXPExpert(args) for _ in range(num_experts)])
        self.gate = nn.Sequential(
            nn.Linear(args["in_len"] * args["num_nodes"], 128),
            nn.ReLU(),
            nn.Linear(128, num_experts),
        )

    def forward(self, x):
        B = x.size(0)
        # Flatten input for gating
        gate_input = x[..., 0].reshape(B, -1)  # (B, T*N)
        gate_logits = self.gate(gate_input)  # (B, num_experts)
        gate_scores = F.softmax(gate_logits, dim=-1)  # soft selection

        # Get top-k experts
        topk_val, topk_idx = torch.topk(gate_scores, self.top_k, dim=-1)  # (B, k)

        outputs = torch.zeros_like(self.experts[0](x))  # (B, H, N, D_out)
        for i in range(self.top_k):
            idx = topk_idx[:, i]
            for expert_id in torch.unique(idx):
                mask = idx == expert_id
                if mask.sum() == 0:
                    continue
                selected_x = x[mask]
                expert_output = self.experts[expert_id](selected_x)
                outputs[mask] += (
                    topk_val[mask, i].unsqueeze(1).unsqueeze(1).unsqueeze(1)
                    * expert_output
                )

        return outputs  # (B, H, N, D_out)