#!/usr/bin/env python
# coding=utf-8
import torch
import torch.nn as nn
import torch.nn.functional as F
from feature_att_model import FeatureAttNet
from time_att_model import TimeAttNet

class HierAttNetnotime(nn.Module):
    def __init__(self, feature_hidden_size, time_hidden_size, batch_size, output_size):
        super(HierAttNetnotime, self).__init__()
        self.batch_size = batch_size
        self.feature_hidden_size = feature_hidden_size
        self.time_hidden_size = time_hidden_size
        # self.feature_att_net = FeatureAttNet(input_size=4, hidden_size=feature_hidden_size, num_layers=1) #crosscheck
        self.feature_att_net = FeatureAttNet(input_size=1, hidden_size=feature_hidden_size, num_layers=1)#parkson
        self.time_att_net = TimeAttNet(input_size=feature_hidden_size, hidden_size=time_hidden_size, num_layers=1)
        self.linear = nn.Linear(time_hidden_size//2, output_size) #park
        # self.linear = nn.Linear(time_hidden_size, output_size) #ppmi


    def forward(self, input):
        output_list = []
        feature_list = []

        input = input['input'].permute(1, 0, 2, 3)

        l1_regularization_list = []
        # print("input",input.size())
        for i in input:

            fan_output,feature_attention,l1_regularization = self.feature_att_net(i.permute(1, 0, 2))

            l1_regularization_list.append(l1_regularization)
            output_list.append(fan_output)
            feature_list.append(feature_attention)
            # print("feature_attention", feature_attention.size())
        output = torch.cat(output_list, 0)

        output = output.permute(1, 0, 2)

        tan_output = output.reshape(output.size(0),-1)


        feature_list = torch.cat(feature_list, 0)
        # print("feature_list", feature_list.size())
        time_attention = 0

        predict = self.linear(tan_output)

        return predict,feature_list,time_attention, l1_regularization_list

    def half(self):
        # for module in self.children():
        #     if type(module) is not nn.BatchNorm2d:
        #         module.half()
        return self

    def get_state(self, mode="full"):
        # return [i for i in self.named_parameters()], []
        return self.state_dict(), []

    def set_state(self, w_server, w_local, mode="full"):
        # sd = self.state_dict()
        # for key, param in w_server:
        #     if key in sd.keys():
        #         sd[key] = param.clone().detach()
        #     else:
        #         print("Server layers mismatch at 'set_state' function.")
        #
        # for key, param in w_local:
        #     if key in sd.keys():
        #         sd[key] = param.clone().detach()
        #     else:
        #         print("Local layers mismatch at 'set_state' function.")

        self.load_state_dict(w_local)


class HierAttNetnofeature(nn.Module):
    def __init__(self, feature_hidden_size, time_hidden_size, batch_size, output_size):
        super(HierAttNetnofeature, self).__init__()
        self.batch_size = batch_size
        self.feature_hidden_size = feature_hidden_size
        self.time_hidden_size = time_hidden_size
        # self.feature_att_net = FeatureAttNet(input_size=4, hidden_size=1028, num_layers=1)
        self.time_att_net = TimeAttNet(input_size=feature_hidden_size, hidden_size=time_hidden_size, num_layers=1)
        self.linear = nn.Linear(time_hidden_size//2, output_size)



    def forward(self, input):


        input = input['input'].permute(1, 0, 2, 3)

        output = input.view(input.size(0),input.size(1), -1)


        l1_regularization_list = []


        feature_list = output

        tan_output,time_attention = self.time_att_net(output)

        #
        tan_output = tan_output[0]
        #
        # print( tan_output.size())
        predict = self.linear(tan_output)



        return predict,feature_list,time_attention, l1_regularization_list

    def half(self):
        # for module in self.children():
        #     if type(module) is not nn.BatchNorm2d:
        #         module.half()
        return self

    def get_state(self, mode="full"):
        # return [i for i in self.named_parameters()], []
        return self.state_dict(), []

    def set_state(self, w_server, w_local, mode="full"):
        # sd = self.state_dict()
        # for key, param in w_server:
        #     if key in sd.keys():
        #         sd[key] = param.clone().detach()
        #     else:
        #         print("Server layers mismatch at 'set_state' function.")
        #
        # for key, param in w_local:
        #     if key in sd.keys():
        #         sd[key] = param.clone().detach()
        #     else:
        #         print("Local layers mismatch at 'set_state' function.")

        self.load_state_dict(w_local)


class HierAttNet(nn.Module):
    def __init__(self, feature_hidden_size, time_hidden_size, batch_size, output_size):
        super(HierAttNet, self).__init__()
        self.batch_size = batch_size
        self.feature_hidden_size = feature_hidden_size
        self.time_hidden_size = time_hidden_size
        # self.feature_att_net = FeatureAttNet(input_size=4, hidden_size=feature_hidden_size, num_layers=1) #crosscheck
        self.feature_att_net = FeatureAttNet(input_size=1, hidden_size=feature_hidden_size, num_layers=1)#parkson
        self.time_att_net = TimeAttNet(input_size=feature_hidden_size, hidden_size=time_hidden_size, num_layers=1)
        self.linear = nn.Linear(time_hidden_size//2, output_size) #正常



    def forward(self, input):
        output_list = []
        feature_list = []
        # print("1", input['input'].size())
        input = input['input'].permute(1, 0, 2, 3)

        l1_regularization_list = []

        for i in input:

            fan_output,feature_attention,l1_regularization = self.feature_att_net(i.permute(1, 0, 2))

            l1_regularization_list.append(l1_regularization)
            output_list.append(fan_output)
            feature_list.append(feature_attention)

        output = torch.cat(output_list, 0)



        feature_list = torch.cat(feature_list, 0)

        tan_output,time_attention = self.time_att_net(output)

        #
        tan_output = tan_output[0]
        #

        predict = self.linear(tan_output)



        return predict,feature_list,time_attention, l1_regularization_list

    def half(self):
        # for module in self.children():
        #     if type(module) is not nn.BatchNorm2d:
        #         module.half()
        return self

    def get_state(self, mode="full"):
        # return [i for i in self.named_parameters()], []
        return self.state_dict(), []

    def set_state(self, w_server, w_local, mode="full"):
        # sd = self.state_dict()
        # for key, param in w_server:
        #     if key in sd.keys():
        #         sd[key] = param.clone().detach()
        #     else:
        #         print("Server layers mismatch at 'set_state' function.")
        #
        # for key, param in w_local:
        #     if key in sd.keys():
        #         sd[key] = param.clone().detach()
        #     else:
        #         print("Local layers mismatch at 'set_state' function.")

        self.load_state_dict(w_local)



