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


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):
        # x: (B, N, input_dim)
        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 ExpertBlock(nn.Module):
    """
    Mixture-of-Experts block: routes each node's representation to a selected expert or a shared expert.
    """
    def __init__(self, hidden_dim, num_experts):
        super().__init__()
        self.num_experts = num_experts
        # gating network projects to num_experts + 1 (extra shared expert)
        self.gate = nn.Linear(hidden_dim, num_experts + 1)
        # per-expert FFNs
        self.experts = nn.ModuleList([
            nn.Sequential(
                nn.Linear(hidden_dim, hidden_dim * 2),
                nn.ReLU(),
                nn.Linear(hidden_dim * 2, hidden_dim)
            ) for _ in range(num_experts)
        ])
        # shared expert
        self.shared_expert = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim * 2),
            nn.ReLU(),
            nn.Linear(hidden_dim * 2, hidden_dim)
        )

    def forward(self, x):
        # x: (B, N, hidden_dim)
        B, N, D = x.shape
        # flatten to (B*N, D)
        flat = x.view(B * N, D)
        # compute gating scores and select expert per node
        scores = F.softmax(self.gate(flat), dim=-1)  # (B*N, num_experts+1)
        idx = scores.argmax(dim=-1)                  # (B*N,)

        out_flat = torch.zeros_like(flat)
        # apply each expert
        for e in range(self.num_experts):
            mask = (idx == e)
            if mask.any():
                out_flat[mask] = self.experts[e](flat[mask])
        # apply shared expert for last index
        shared_mask = (idx == self.num_experts)
        if shared_mask.any():
            out_flat[shared_mask] = self.shared_expert(flat[shared_mask])

        # reshape back to (B, N, D)
        return out_flat.view(B, N, D)


class MLP(nn.Module):
    def __init__(self, in_dim, hidden_dims, out_dim, activation=nn.ReLU):
        super().__init__()
        dims = [in_dim] + hidden_dims + [out_dim]
        layers = []
        for i in range(len(dims) - 2):
            layers += [nn.Linear(dims[i], dims[i+1]), activation()]
        layers += [nn.Linear(dims[-2], dims[-1])]
        self.net = nn.Sequential(*layers)

    def forward(self, x):
        return self.net(x)


class SandwichBlock(nn.Module):
    def __init__(self, num_nodes, embed_dim, hidden_dim, num_experts):
        super().__init__()
        self.manba1 = MANBA_Block(hidden_dim, hidden_dim * 2)
        self.expert_block = ExpertBlock(hidden_dim, num_experts)
        self.manba2 = MANBA_Block(hidden_dim, hidden_dim * 2)

    def forward(self, h):
        h1 = self.manba1(h)
        h2 = self.expert_block(h1)
        h3 = self.manba2(h2)
        return h3


class EXP(nn.Module):
    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.embed_dim   = args.get('embed_dim', 16)
        self.num_experts = args.get('num_experts', 8)  # number of private experts

        # discrete time embeddings
        self.time_slots    = args.get('time_slots', 24 * 60 // args.get('time_slot', 5))
        self.time_embedding = nn.Embedding(self.time_slots, self.hidden_dim)
        self.day_embedding  = nn.Embedding(7, self.hidden_dim)

        # input projection
        self.input_proj = MLP(
            in_dim      = self.seq_len,
            hidden_dims = [self.hidden_dim],
            out_dim     = self.hidden_dim
        )

        # two Sandwich blocks with MoE
        self.sandwich1 = SandwichBlock(self.num_nodes, self.embed_dim, self.hidden_dim, self.num_experts)
        self.sandwich2 = SandwichBlock(self.num_nodes, self.embed_dim, self.hidden_dim, self.num_experts)

        # output projection
        self.out_proj = MLP(
            in_dim      = self.hidden_dim,
            hidden_dims = [2 * self.hidden_dim],
            out_dim     = self.horizon * self.output_dim
        )

    def forward(self, x):
        """
        x: (B, T, N, D_total)
           x[...,0]= flow, x[...,1]=time_in_day, x[...,2]=day_in_week
        """
        x_flow = x[..., 0]
        x_time = x[..., 1]
        x_day  = x[..., 2]

        B, T, N = x_flow.shape
        assert T == self.seq_len

        # project flow history
        x_flat = x_flow.permute(0, 2, 1).reshape(B * N, T)
        h0 = self.input_proj(x_flat).view(B, N, self.hidden_dim)

        # time & day embeddings at last step
        t_idx = (x_time[:, -1, :,] * (self.time_slots - 1)).long()
        d_idx = x_day[:,  -1, :,].long()
        time_emb = self.time_embedding(t_idx)
        day_emb  = self.day_embedding(d_idx)
        h0 = h0 + time_emb + day_emb

        # two MoE Sandwich blocks + residuals
        h1 = self.sandwich1(h0) + h0
        h2 = self.sandwich2(h1) + h1

        # output
        out = self.out_proj(h2)
        out = out.view(B, N, self.horizon, self.output_dim)
        out = out.permute(0, 2, 1, 3)
        return out