import torch.nn as nn

from model.STGCN import layers
from data.get_adj import get_gso


class STGCNChebGraphConv(nn.Module):
    # STGCNChebGraphConv contains 'TGTND TGTND TNFF' structure
    # ChebGraphConv is the graph convolution from ChebyNet.
    # Using the Chebyshev polynomials of the first kind as a graph filter.

    # T: Gated Temporal Convolution Layer (GLU or GTU)
    # G: Graph Convolution Layer (ChebGraphConv)
    # T: Gated Temporal Convolution Layer (GLU or GTU)
    # N: Layer Normolization
    # D: Dropout

    # T: Gated Temporal Convolution Layer (GLU or GTU)
    # G: Graph Convolution Layer (ChebGraphConv)
    # T: Gated Temporal Convolution Layer (GLU or GTU)
    # N: Layer Normolization
    # D: Dropout

    # T: Gated Temporal Convolution Layer (GLU or GTU)
    # N: Layer Normalization
    # F: Fully-Connected Layer
    # F: Fully-Connected Layer

    def __init__(self, args):
        super(STGCNChebGraphConv, self).__init__()
        gso = get_gso(args)
        Ko = args['n_his'] - (args['Kt'] - 1) * 2 * args['stblock_num']
        blocks = [[1]]
        for l in range(args['stblock_num']):
            blocks.append([64, 16, 64])
        if Ko == 0:
            blocks.append([128])
        elif Ko > 0:
            blocks.append([128, 128])
        blocks.append([12])
        modules = []
        for l in range(len(blocks) - 3):
            modules.append(layers.STConvBlock(args['Kt'], args['Ks'], args['num_nodes'], blocks[l][-1], blocks[l + 1],
                                              args['act_func'], args['graph_conv_type'], gso,
                                              args['enable_bias'], args['droprate']))
        self.st_blocks = nn.Sequential(*modules)
        Ko = args['n_his'] - (len(blocks) - 3) * 2 * (args['Kt'] - 1)
        self.Ko = Ko
        if self.Ko > 1:
            self.output = layers.OutputBlock(Ko, blocks[-3][-1], blocks[-2], blocks[-1][0],
                                             args['num_nodes'], args['act_func'], args['enable_bias'], args['droprate'])
        elif self.Ko == 0:
            self.fc1 = nn.Linear(in_features=blocks[-3][-1], out_features=blocks[-2][0], bias=args['enable_bias'])
            self.fc2 = nn.Linear(in_features=blocks[-2][0], out_features=blocks[-1][0], bias=args['enable_bias'])
            self.relu = nn.ReLU()
            self.dropout = nn.Dropout(p=args['droprate'])

    def forward(self, x):
        x = x[..., 0:1]
        x = x.permute(0, 3, 1, 2)
        x = self.st_blocks(x)  # 64,12,307,3
        if self.Ko > 1:
            x = self.output(x)
        elif self.Ko == 0:
            x = self.fc1(x.permute(0, 2, 3, 1))
            x = self.relu(x)
            x = self.fc2(x).permute(0, 3, 1, 2)

        x = x.permute(0, 1, 3, 2)
        return x


class STGCNGraphConv(nn.Module):
    # STGCNGraphConv contains 'TGTND TGTND TNFF' structure
    # GraphConv is the graph convolution from GCN.
    # GraphConv is not the first-order ChebConv, because the renormalization trick is adopted.
    # Be careful about over-smoothing.

    # T: Gated Temporal Convolution Layer (GLU or GTU)
    # G: Graph Convolution Layer (GraphConv)
    # T: Gated Temporal Convolution Layer (GLU or GTU)
    # N: Layer Normolization
    # D: Dropout

    # T: Gated Temporal Convolution Layer (GLU or GTU)
    # G: Graph Convolution Layer (GraphConv)
    # T: Gated Temporal Convolution Layer (GLU or GTU)
    # N: Layer Normolization
    # D: Dropout

    # T: Gated Temporal Convolution Layer (GLU or GTU)
    # N: Layer Normalization
    # F: Fully-Connected Layer
    # F: Fully-Connected Layer

    def __init__(self, args):
        super(STGCNGraphConv, self).__init__()
        gso = get_gso(args)
        Ko = args['n_his'] - (args['Kt'] - 1) * 2 * args['stblock_num']
        blocks = [[1]]
        for l in range(args['stblock_num']):
            blocks.append([64, 16, 64])
        if Ko == 0:
            blocks.append([128])
        elif Ko > 0:
            blocks.append([128, 128])
        blocks.append([1])
        modules = []
        for l in range(len(blocks) - 3):
            modules.append(layers.STConvBlock(args['Kt'], args['Ks'], args['num_nodes'], blocks[l][-1], blocks[l + 1],
                                              args['act_func'], args['graph_conv_type'], gso,
                                              args['enable_bias'], args['droprate']))
        self.st_blocks = nn.Sequential(*modules)
        Ko = args['n_his'] - (len(blocks) - 3) * 2 * (args['Kt'] - 1)
        self.Ko = Ko
        if self.Ko > 1:
            self.output = layers.OutputBlock(Ko, blocks[-3][-1], blocks[-2], blocks[-1][0], args['num_nodes'],
                                             args['act_func'], args['enable_bias'], args['droprate'])
        elif self.Ko == 0:
            self.fc1 = nn.Linear(in_features=blocks[-3][-1], out_features=blocks[-2][0], bias=args['enable_bias'])
            self.fc2 = nn.Linear(in_features=blocks[-2][0], out_features=blocks[-1][0], bias=args['enable_bias'])
            self.relu = nn.ReLU()
            self.do = nn.Dropout(p=args['droprate'])

    def forward(self, x):
        x = self.st_blocks(x)
        if self.Ko > 1:
            x = self.output(x)
        elif self.Ko == 0:
            x = self.fc1(x.permute(0, 2, 3, 1))
            x = self.relu(x)
            x = self.fc2(x).permute(0, 3, 1, 2)
        return x