import sys
sys.path.insert(0, '')
sys.path.append('..')
import math
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

from utils import import_class, count_params
from model.ms_gcn import MultiScale_GraphConv as MS_GCN
from model.ms_tcn import MultiScale_TemporalConv as MS_TCN
from model.ms_gtcn import SpatialTemporal_MS_GCN, UnfoldTemporalWindows
from model.mlp import MLP
from model.activation import activation_factory


class MS_G3D(nn.Module):
    def __init__(self,
                 in_channels,
                 out_channels,
                 A_binary,
                 num_scales,
                 window_size,
                 window_stride,
                 window_dilation,
                 embed_factor=1,
                 activation='relu'):
        super().__init__()
        # print(in_channels, out_channels, A_binary.shape, num_scales,
        # window_size, window_stride, window_dilation)
        self.window_size = window_size
        self.out_channels = out_channels
        self.embed_channels_in = self.embed_channels_out = out_channels // embed_factor
        if embed_factor == 1:
            self.in1x1 = nn.Identity()
            self.embed_channels_in = self.embed_channels_out = in_channels
            # The first STGC block changes channels right away; others change at collapse
            if in_channels == 3:
                self.embed_channels_out = out_channels
        else:
            self.in1x1 = MLP(in_channels, [self.embed_channels_in])

        self.gcn3d = nn.Sequential(
            UnfoldTemporalWindows(window_size, window_stride, window_dilation),
            SpatialTemporal_MS_GCN(
                in_channels=self.embed_channels_in,
                out_channels=self.embed_channels_out,
                A_binary=A_binary,
                num_scales=num_scales,
                window_size=window_size,
                use_Ares=True
            )
        )

        self.out_conv = nn.Conv3d(
            self.embed_channels_out, out_channels, kernel_size=(1, self.window_size, 1))
        self.out_bn = nn.BatchNorm2d(out_channels)

    def forward(self, x):
        N, _, T, V = x.shape
        # print(x.shape)
        x = self.in1x1(x)
        # Construct temporal windows and apply MS-GCN
        x = self.gcn3d(x)

        # Collapse the window dimension
        x = x.view(N, self.embed_channels_out, -1, self.window_size, V)
        x = self.out_conv(x).squeeze(dim=3)
        x = self.out_bn(x)

        # no activation
        return x


class MultiWindow_MS_G3D(nn.Module):
    def __init__(self,
                 in_channels,
                 out_channels,
                 A_binary,
                 num_scales,
                 window_sizes=[3, 5],
                 window_stride=1,
                 window_dilations=[1, 1]):

        super().__init__()
        # print(in_channels, out_channels, num_scales,
        # window_sizes, window_stride, window_dilations)
        self.gcn3d = nn.ModuleList([
            MS_G3D(
                in_channels,
                out_channels,
                A_binary,
                num_scales,
                window_size,
                window_stride,
                window_dilation
            )
            for window_size, window_dilation in zip(window_sizes, window_dilations)
        ])

    def forward(self, x):
        # Input shape: (N, C, T, V)
        # print(x.shape)
        out_sum = 0
        for gcn3d in self.gcn3d:
            out_sum += gcn3d(x)
        # no activation
        return out_sum

class Pass(nn.Module):
    """docstring for myBatchNorm"""
    def __init__(self,shape):
        super(Pass, self).__init__()
        self.placeholder=nn.Parameter(torch.zeros(shape))



    def forward(self,x):
        return x

class Model(nn.Module):
    def __init__(self,
                 num_class,
                 num_point,
                 num_person,
                 num_gcn_scales,
                 num_g3d_scales,
                 graph,
                 in_channels=3):
        super(Model, self).__init__()

        

        Graph = import_class(graph)
        A_binary = Graph().A_binary
        T = 50
        self.passx=Pass((num_person * in_channels * num_point,T))
        self.data_bn = nn.BatchNorm1d(num_person * in_channels * num_point)


        # channels
        c1 = 96
        c2 = c1 * 2     # 192
        c3 = c2 * 2     # 384

        # r=3 STGC blocks
        self.gcn3d1 = MultiWindow_MS_G3D(
            3, c1, A_binary, num_g3d_scales, window_stride=1)
        self.sgcn1 = nn.Sequential(
            MS_GCN(num_gcn_scales, 3, c1, A_binary, disentangled_agg=True),
            MS_TCN(c1, c1),
            MS_TCN(c1, c1))
        self.sgcn1[-1].act = nn.Identity()
        self.tcn1 = MS_TCN(c1, c1)

        self.gcn3d2 = MultiWindow_MS_G3D(
            c1, c2, A_binary, num_g3d_scales, window_stride=2)
        self.sgcn2 = nn.Sequential(
            MS_GCN(num_gcn_scales, c1, c1, A_binary, disentangled_agg=True),
            MS_TCN(c1, c2, stride=2),
            MS_TCN(c2, c2))
        self.sgcn2[-1].act = nn.Identity()
        self.tcn2 = MS_TCN(c2, c2)

        self.gcn3d3 = MultiWindow_MS_G3D(
            c2, c3, A_binary, num_g3d_scales, window_stride=2)
        self.sgcn3 = nn.Sequential(
            MS_GCN(num_gcn_scales, c2, c2, A_binary, disentangled_agg=True),
            MS_TCN(c2, c3, stride=2),
            MS_TCN(c3, c3))
        self.sgcn3[-1].act = nn.Identity()
        self.tcn3 = MS_TCN(c3, c3)

        self.fc = nn.Linear(c3, num_class)

    def forward(self, x):
        
        C, T, V, M = x.shape

        x = x.permute(3,2,0,1).contiguous().view(1, M * V * C, T)
        x= self.passx(x)

        x = self.data_bn(x)
        # x = self.data_bn(x.unsqueeze（-1).squeeze（-1)
        x = x.view(M, V, C, T).permute(0, 2, 3, 1).contiguous()

        # Apply activation to the sum of the pathways
        x = F.relu(self.sgcn1(x) + self.gcn3d1(x), inplace=True)
        x = self.tcn1(x)

        x = F.relu(self.sgcn2(x) + self.gcn3d2(x), inplace=True)
        x = self.tcn2(x)

        x = F.relu(self.sgcn3(x) + self.gcn3d3(x), inplace=True)
        x = self.tcn3(x)

        out = x
        # print(x.shape)
        out_channels = out.size(1)
        # out = out.view(M, out_channels, -1)
        # out = out.mean(2)   # Global Average Pooling (Spatial+Temporal)
        # out = out.mean(0)   # Average pool number of bodies in the sequence
        # 不能少那一维！
        out = out.view(1, M, out_channels, -1)
        out = out.mean(3) 
        out = out.mean(1)   

        out = self.fc(out)
        return out

    # def init_reshape(self, x):
    #     # x = x.unsqueeze(0)
    #     C, T, V, M = 3, 50, 18, 2
    #     return x.permute(3,2,0,1).contiguous().view(1, M * V * C, T)


if __name__ == "__main__":
    # from graph.kinetics import AdjMatrixGraph
    # graph = AdjMatrixGraph()
    # A_binary = graph.A_binary
    N, C, T, V, M = 1, 3, 50, 18, 2
    # m1 = MultiWindow_MS_G3D(192, 384, A_binary, 8, [3, 5], 2, [1, 1])
    # x1 = torch.ones(2, 192, 50, 18)

    m = Model(num_class=400, num_point=18, num_person=2,
              num_gcn_scales=8, num_g3d_scales=8, graph='graph.kinetics.AdjMatrixGraph').eval()
    state=m.state_dict()

    x = torch.ones(C, T, V, M)
    pt_model_path = './pretrained-models/kinetics-joint.pt'
    
    pretrained_state_dict=torch.load('../pretrained-models/kinetics-joint.pt')
    
    state.update(pretrained_state_dict)
    m.load_state_dict(state)
    # print(m.state_dict().get('passx.placeholder'))
    # print(m.state_dict())
    
    print(m(x).argmax())
    # from torchsummaryX import summary
    # summary(m, torch.zeros(108))
    # torch.onnx.export(m, x, 'msg3d.onnx', opset_version=12)
    # import onnx
    # omx = onnx.load('msg3d.onnx')
    # from onnx_tf.backend import prepare
    # tfx = prepare(omx)
    # tfx.export_graph('msg3d.pb')

    # # # # # print('Model total # params:', count_params(model))
