import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn import init
from utils.get_adj import get_adj
import numbers


# --- 基础算子 ---
class NConv(nn.Module):
    def forward(self, x, adj):
        return torch.einsum("ncwl,vw->ncvl", (x, adj)).contiguous()


class DyNconv(nn.Module):
    def forward(self, x, adj):
        return torch.einsum("ncvl,nvwl->ncwl", (x, adj)).contiguous()


class Linear(nn.Module):
    def __init__(self, c_in, c_out, bias=True):
        super().__init__()
        self.mlp = nn.Conv2d(c_in, c_out, kernel_size=1, bias=bias)

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


class Prop(nn.Module):
    def __init__(self, c_in, c_out, gdep, dropout, alpha):
        super().__init__()
        self.nconv = NConv()
        self.mlp = Linear(c_in, c_out)
        self.gdep, self.dropout, self.alpha = gdep, dropout, alpha

    def forward(self, x, adj):
        adj = adj + torch.eye(adj.size(0), device=x.device)
        d = adj.sum(1)
        a = adj / d.view(-1, 1)
        h = x
        for _ in range(self.gdep):
            h = self.alpha * x + (1 - self.alpha) * self.nconv(h, a)
        return self.mlp(h)


class MixProp(nn.Module):
    def __init__(self, c_in, c_out, gdep, dropout, alpha):
        super().__init__()
        self.nconv = NConv()
        self.mlp = Linear((gdep + 1) * c_in, c_out)
        self.gdep, self.dropout, self.alpha = gdep, dropout, alpha

    def forward(self, x, adj):
        adj = adj + torch.eye(adj.size(0), device=x.device)
        d = adj.sum(1)
        a = adj / d.view(-1, 1)
        out = [x]
        h = x
        for _ in range(self.gdep):
            h = self.alpha * x + (1 - self.alpha) * self.nconv(h, a)
            out.append(h)
        return self.mlp(torch.cat(out, dim=1))


class DyMixprop(nn.Module):
    def __init__(self, c_in, c_out, gdep, dropout, alpha):
        super().__init__()
        self.nconv = DyNconv()
        self.mlp1 = Linear((gdep + 1) * c_in, c_out)
        self.mlp2 = Linear((gdep + 1) * c_in, c_out)
        self.gdep, self.dropout, self.alpha = gdep, dropout, alpha
        self.lin1, self.lin2 = Linear(c_in, c_in), Linear(c_in, c_in)

    def forward(self, x):
        x1 = torch.tanh(self.lin1(x))
        x2 = torch.tanh(self.lin2(x))
        adj = self.nconv(x1.transpose(2, 1), x2)
        adj0 = torch.softmax(adj, dim=2)
        adj1 = torch.softmax(adj.transpose(2, 1), dim=2)
        # 两条分支
        out1, out2 = [x], [x]
        h = x
        for _ in range(self.gdep):
            h = self.alpha * x + (1 - self.alpha) * self.nconv(h, adj0)
            out1.append(h)
        h = x
        for _ in range(self.gdep):
            h = self.alpha * x + (1 - self.alpha) * self.nconv(h, adj1)
            out2.append(h)
        return self.mlp1(torch.cat(out1, dim=1)) + self.mlp2(torch.cat(out2, dim=1))


class DilatedInception(nn.Module):
    def __init__(self, cin, cout, dilation_factor=2):
        super().__init__()
        self.kernels = [2, 3, 6, 7]
        cout_each = int(cout / len(self.kernels))
        self.convs = nn.ModuleList(
            [
                nn.Conv2d(
                    cin, cout_each, kernel_size=(1, k), dilation=(1, dilation_factor)
                )
                for k in self.kernels
            ]
        )

    def forward(self, x):
        outs = [conv(x)[..., -self.convs[-1](x).size(3) :] for conv in self.convs]
        return torch.cat(outs, dim=1)


class GraphConstructor(nn.Module):
    def __init__(self, nnodes, k, dim, device, alpha=3, static_feat=None):
        super().__init__()
        self.nnodes, self.k, self.dim, self.alpha, self.device = (
            nnodes,
            k,
            dim,
            alpha,
            device,
        )
        self.static_feat = static_feat
        if static_feat is not None:
            xd = static_feat.shape[1]
            self.lin1, self.lin2 = nn.Linear(xd, dim), nn.Linear(xd, dim)
        else:
            self.emb1 = nn.Embedding(nnodes, dim)
            self.emb2 = nn.Embedding(nnodes, dim)
            self.lin1, self.lin2 = nn.Linear(dim, dim), nn.Linear(dim, dim)

    def forward(self, idx):
        if self.static_feat is None:
            vec1, vec2 = self.emb1(idx), self.emb2(idx)
        else:
            vec1 = vec2 = self.static_feat[idx, :]
        vec1 = torch.tanh(self.alpha * self.lin1(vec1))
        vec2 = torch.tanh(self.alpha * self.lin2(vec2))
        a = torch.mm(vec1, vec2.transpose(1, 0)) - torch.mm(vec2, vec1.transpose(1, 0))
        adj = F.relu(torch.tanh(self.alpha * a))
        mask = torch.zeros(idx.size(0), idx.size(0), device=self.device)
        s1, t1 = adj.topk(self.k, 1)
        mask.scatter_(1, t1, s1.new_ones(s1.size()))
        return adj * mask


class LayerNorm(nn.Module):
    __constants__ = ["normalized_shape", "eps", "elementwise_affine"]

    def __init__(self, normalized_shape, eps=1e-5, elementwise_affine=True):
        super().__init__()
        if isinstance(normalized_shape, numbers.Integral):
            normalized_shape = (normalized_shape,)
        self.normalized_shape, self.eps, self.elementwise_affine = (
            tuple(normalized_shape),
            eps,
            elementwise_affine,
        )
        if elementwise_affine:
            self.weight = nn.Parameter(torch.Tensor(*normalized_shape))
            self.bias = nn.Parameter(torch.Tensor(*normalized_shape))
            init.ones_(self.weight)
            init.zeros_(self.bias)
        else:
            self.register_parameter("weight", None)
            self.register_parameter("bias", None)

    def forward(self, x, idx):
        if self.elementwise_affine:
            return F.layer_norm(
                x,
                tuple(x.shape[1:]),
                self.weight[:, idx, :],
                self.bias[:, idx, :],
                self.eps,
            )
        else:
            return F.layer_norm(x, tuple(x.shape[1:]), self.weight, self.bias, self.eps)

    def extra_repr(self):
        return f"{self.normalized_shape}, eps={self.eps}, elementwise_affine={self.elementwise_affine}"


# --- 合并后的模型类，支持 teacher 与 stmlp 两种分支 ---
class STMLP(nn.Module):
    def __init__(self, args):
        super().__init__()
        # 参数从字典中读取
        self.adj_mx = get_adj(args)
        self.num_nodes = args["num_nodes"]
        self.feature_dim = args["input_dim"]

        self.input_window = args["input_window"]
        self.output_window = args["output_window"]
        self.output_dim = args["output_dim"]
        self.device = args["device"]

        self.gcn_true = args["gcn_true"]
        self.buildA_true = args["buildA_true"]
        self.gcn_depth = args["gcn_depth"]
        self.dropout = args["dropout"]
        self.subgraph_size = args["subgraph_size"]
        self.node_dim = args["node_dim"]
        self.dilation_exponential = args["dilation_exponential"]

        self.conv_channels = args["conv_channels"]
        self.residual_channels = args["residual_channels"]
        self.skip_channels = args["skip_channels"]
        self.end_channels = args["end_channels"]

        self.layers = args["layers"]
        self.propalpha = args["propalpha"]
        self.tanhalpha = args["tanhalpha"]
        self.layer_norm_affline = args["layer_norm_affline"]

        self.model_type = args["model_type"]  # 'teacher' 或 'stmlp'
        self.idx = torch.arange(self.num_nodes).to(self.device)
        self.predefined_A = (
            None
            if self.adj_mx is None
            else (torch.tensor(self.adj_mx) - torch.eye(self.num_nodes)).to(self.device)
        )
        self.static_feat = None

        # transformer（保留原有结构）
        self.encoder_layer = nn.TransformerEncoderLayer(
            d_model=12, nhead=4, batch_first=True
        )
        self.transformer_encoder = nn.TransformerEncoder(
            self.encoder_layer, num_layers=3
        )

        # 构建各层
        self.start_conv = nn.Conv2d(
            self.feature_dim, self.residual_channels, kernel_size=1
        )
        self.gc = GraphConstructor(
            self.num_nodes,
            self.subgraph_size,
            self.node_dim,
            self.device,
            alpha=self.tanhalpha,
            static_feat=self.static_feat,
        )
        # 计算 receptive_field
        kernel_size = 7
        if self.dilation_exponential > 1:
            self.receptive_field = int(
                self.output_dim
                + (kernel_size - 1)
                * (self.dilation_exponential**self.layers - 1)
                / (self.dilation_exponential - 1)
            )
        else:
            self.receptive_field = self.layers * (kernel_size - 1) + self.output_dim

        self.filter_convs = nn.ModuleList()
        self.gate_convs = nn.ModuleList()
        self.residual_convs = nn.ModuleList()
        self.skip_convs = nn.ModuleList()
        self.norm = nn.ModuleList()
        self.stu_mlp = nn.ModuleList(
            [
                nn.Sequential(nn.Linear(c, c), nn.Linear(c, c), nn.Linear(c, c))
                for c in [13, 7, 1]
            ]
        )
        if self.gcn_true:
            self.gconv1 = nn.ModuleList()
            self.gconv2 = nn.ModuleList()

        new_dilation = 1
        for i in range(1):
            rf_size_i = (
                int(
                    1
                    + i
                    * (kernel_size - 1)
                    * (self.dilation_exponential**self.layers - 1)
                    / (self.dilation_exponential - 1)
                )
                if self.dilation_exponential > 1
                else i * self.layers * (kernel_size - 1) + 1
            )
            for j in range(1, self.layers + 1):
                rf_size_j = (
                    int(
                        rf_size_i
                        + (kernel_size - 1)
                        * (self.dilation_exponential**j - 1)
                        / (self.dilation_exponential - 1)
                    )
                    if self.dilation_exponential > 1
                    else rf_size_i + j * (kernel_size - 1)
                )
                self.filter_convs.append(
                    DilatedInception(
                        self.residual_channels,
                        self.conv_channels,
                        dilation_factor=new_dilation,
                    )
                )
                self.gate_convs.append(
                    DilatedInception(
                        self.residual_channels,
                        self.conv_channels,
                        dilation_factor=new_dilation,
                    )
                )
                self.residual_convs.append(
                    nn.Conv2d(self.conv_channels, self.residual_channels, kernel_size=1)
                )
                k_size = (
                    (1, self.input_window - rf_size_j + 1)
                    if self.input_window > self.receptive_field
                    else (1, self.receptive_field - rf_size_j + 1)
                )
                self.skip_convs.append(
                    nn.Conv2d(
                        self.conv_channels, self.skip_channels, kernel_size=k_size
                    )
                )
                if self.gcn_true:
                    self.gconv1.append(
                        MixProp(
                            self.conv_channels,
                            self.residual_channels,
                            self.gcn_depth,
                            self.dropout,
                            self.propalpha,
                        )
                    )
                    self.gconv2.append(
                        MixProp(
                            self.conv_channels,
                            self.residual_channels,
                            self.gcn_depth,
                            self.dropout,
                            self.propalpha,
                        )
                    )
                norm_size = (
                    (
                        self.residual_channels,
                        self.num_nodes,
                        self.input_window - rf_size_j + 1,
                    )
                    if self.input_window > self.receptive_field
                    else (
                        self.residual_channels,
                        self.num_nodes,
                        self.receptive_field - rf_size_j + 1,
                    )
                )
                self.norm.append(
                    LayerNorm(norm_size, elementwise_affine=self.layer_norm_affline)
                )
                new_dilation *= self.dilation_exponential

        self.end_conv_1 = nn.Conv2d(
            self.skip_channels, self.end_channels, kernel_size=1, bias=True
        )
        self.end_conv_2 = nn.Conv2d(
            self.end_channels, self.output_window, kernel_size=1, bias=True
        )
        k0 = (
            (1, self.input_window)
            if self.input_window > self.receptive_field
            else (1, self.receptive_field)
        )
        self.skip0 = nn.Conv2d(
            self.feature_dim, self.skip_channels, kernel_size=k0, bias=True
        )
        kE = (
            (1, self.input_window - self.receptive_field + 1)
            if self.input_window > self.receptive_field
            else (1, 1)
        )
        self.skipE = nn.Conv2d(
            self.residual_channels, self.skip_channels, kernel_size=kE, bias=True
        )
        # 最后输出分支，根据模型类型选择不同的头
        if self.model_type == "teacher":
            self.tt_linear1 = nn.Linear(self.residual_channels, self.input_window)
            self.tt_linear2 = nn.Linear(1, 32)
            self.ss_linear1 = nn.Linear(self.residual_channels, self.input_window)
            self.ss_linear2 = nn.Linear(1, 32)
        else:  # stmlp
            self.out_linear1 = nn.Linear(self.residual_channels, self.input_window)
            self.out_linear2 = nn.Linear(1, 32)

    def forward(self, source, idx=None):
        source = source[..., 0:1]
        sout, tout = [], []
        inputs = source.transpose(1, 3)
        assert inputs.size(3) == self.input_window, "input sequence length mismatch"
        if self.input_window < self.receptive_field:
            inputs = F.pad(inputs, (self.receptive_field - self.input_window, 0, 0, 0))
        if self.gcn_true:
            adp = (
                self.gc(self.idx if idx is None else idx)
                if self.buildA_true
                else self.predefined_A
            )
        x = self.start_conv(inputs)
        skip = self.skip0(F.dropout(inputs, self.dropout, training=self.training))
        for i in range(self.layers):
            residual = x
            filters = torch.tanh(self.filter_convs[i](x))
            gate = torch.sigmoid(self.gate_convs[i](x))
            x = F.dropout(filters * gate, self.dropout, training=self.training)
            tout.append(x)
            s = self.skip_convs[i](x)
            skip = s + skip
            if self.gcn_true:
                x = self.gconv1[i](x, adp) + self.gconv2[i](x, adp.transpose(1, 0))
            else:
                x = self.stu_mlp[i](x)
            x = x + residual[:, :, :, -x.size(3) :]
            x = self.norm[i](x, self.idx if idx is None else idx)
            sout.append(x)
        skip = self.skipE(x) + skip
        x = F.relu(skip)
        x = F.relu(self.end_conv_1(x))
        x = self.end_conv_2(x)
        if self.model_type == "teacher":
            ttout = self.tt_linear2(
                self.tt_linear1(tout[-1].transpose(1, 3)).transpose(1, 3)
            )
            ssout = self.ss_linear2(
                self.ss_linear1(sout[-1].transpose(1, 3)).transpose(1, 3)
            )
            return x, ttout, ssout
        else:
            x_ = self.out_linear2(
                self.out_linear1(tout[-1].transpose(1, 3)).transpose(1, 3)
            )
            return x, x_, x