import threading
import datetime
import torch
import time
import numpy as np
from BResidual import BResidual,CNN
from optimiser import SGD
from torch.optim import Adam
from util import sd_matrixing, PiecewiseLinear, trainable_params, StatsLogger
from hierachical_att_model import HierAttNet,LSTM,MLP,normalCNN,CNNLSTMModel,HierAttNetnofeature,HierAttNetnotime,TransformerModel,spatial_att_LSTM,temporal_att_LSTM
from sklearn.metrics import r2_score,max_error,mean_squared_log_error

def msle_loss(y_pred, y_true):
    epsilon = 1e-10  
    log_pred = torch.where(y_pred >= 0, torch.log(y_pred + 1), torch.log(epsilon))
    log_true = torch.where(y_true >= 0, torch.log(y_true + 1), torch.log(epsilon))
    msle = torch.mean((log_pred - log_true)**2)  
    return msle

def robust_msle(output, target, offset=1.0, min_val=1e-6):
    output = torch.abs(output)
    target = torch.abs(target)
  
    output = torch.clamp(output + offset, min=min_val)
    target = torch.clamp(target + offset, min=min_val)
    return torch.mean((torch.log(output) - torch.log(target)) ** 2)


class Cifar10FedEngine:
    def __init__(self, args, dataloader, global_param, server_param, local_param,
                 outputs, cid, tid, mode, server_state, client_states):
        self.args = args
        self.dataloader = dataloader

        self.global_param = global_param
        self.server_param = server_param
        self.local_param = local_param
        self.server_state = server_state
        self.client_state = client_states

        self.client_id = cid
        self.outputs = outputs
        self.thread = tid

        self.mode = mode

        self.model = self.prepare_model()
        # self.threadLock = threading.Lock()

        self.m1, self.m2, self.m3, self.reg1, self.reg2 = None, None, None, None, None

    def prepare_model(self):

        if self.args.dataset == "crosscheck" or self.args.dataset == "park" or self.args.dataset == "ppmi" or self.args.dataset == "mapss":

            if self.args.dataset == 'park':

                model = HierAttNet(feature_hidden_size=512, time_hidden_size=20, batch_size=8,
                                   output_size=1).to(self.args.device)

            elif self.args.dataset == 'ppmi':
                model = HierAttNet(feature_hidden_size=512, time_hidden_size=20, batch_size=8,
                                   output_size=1).to(self.args.device)

            elif self.args.dataset == 'mapss':
                model = HierAttNet(feature_hidden_size=512, time_hidden_size=20, batch_size=8,
                                   output_size=1).to(self.args.device)

            else:
                model = HierAttNet(feature_hidden_size=512, time_hidden_size=20, batch_size=8,
                                   output_size=1).to(self.args.device)


        else:
            print("Unknown model type ... ")
            model = None

        model.set_state(self.global_param, self.server_param,self.args.agg)
        return model

    def run(self):
        self.model.to(self.args.device)
        if self.args.agg != 'cen':

            output, fea_att, time_att, l1_regularization_list = self.client_run()

        else:
            # output = self.client_run()
            output, fea_att, time_att, l1_regularization_list = self.client_run()
        self.free_memory()
        if self.args.agg != 'central':

            return output, fea_att, time_att
        else:
            # return output
            return output, fea_att, time_att
    def client_run(self):
        # lr_schedule = PiecewiseLinear([0, 5, self.args.client_epochs], [0, 0.4, 0.001])
        lr_schedule = PiecewiseLinear([0, 5, self.args.client_epochs], [0, 0.4, 0.01])
        lr = lambda step: lr_schedule(step / len(self.dataloader)) / self.args.batch_size



        opt = Adam(trainable_params(self.model), lr=0.01, weight_decay=5e-4 * self.args.batch_size
                   )  # our

        mean_loss = []
        mean_rmse = []
        mean_mse = []
        mean_msle = []
        mean_mape = []
        t1 = time.time()
        c_state = None

        if self.mode == "Train":
            # training process
            for epoch in range(self.args.client_epochs):
                if self.args.agg != 'cen':

                    stats,fea_att, time_att, l1_regularization_list = self.batch_run(True, opt.step)
                else:

                    stats, fea_att, time_att, l1_regularization_list = self.batch_run(True, opt.step)

                mean_loss.append(stats['loss'].cpu().detach())
                mean_rmse.append(stats['rmse'].cpu().detach())
                mean_mse.append(stats['mse'].cpu().detach())
                mean_msle.append(stats['msle'].cpu().detach())
                mean_mape.append(stats['mape'].cpu().detach())


        elif self.mode == "Test":
            # validation process
            if self.args.agg != 'cen':

                stats, fea_att, time_att,l1_regularization_list = self.batch_run(True, opt.step)
            else:

                stats, fea_att, time_att, l1_regularization_list = self.batch_run(True, opt.step)

            mean_loss.append(stats['loss'].cpu().detach())
            mean_rmse.append(stats['rmse'].cpu().detach())
            mean_mse.append(stats['mse'].cpu().detach())
            mean_msle.append(stats['msle'].cpu().detach())
            mean_mape.append(stats['mape'].cpu().detach())


        time_cost = time.time() - t1


        log = self.mode + ' - Thread: {:03d}, Client: {:03d}. Average Loss: {:.4f},' \
                          ' Average rmse: {:.4f}, Average mse: {:.4f},Average mape: {:.4f},Average msle: {:.4f},Total Time Cost: {:.4f}'
        self.logger(log.format(self.thread, self.client_id, np.mean(mean_loss), np.mean(mean_rmse), np.mean(mean_mse),np.mean(mean_mape),np.mean(mean_msle),
                               time_cost), False)

        self.model.to("cpu")
        output = {"params": self.model.get_state(),
                  "time": time_cost,
                  "loss": np.mean(mean_loss),
                  "rmse": np.mean(mean_rmse),
                  "mse": np.mean(mean_mse),
                  "msle": np.mean(mean_msle),
                  "mape": np.mean(mean_mape),
                  "client_state": self.client_state,
                  "c_state": c_state}

        # self.outputs[self.thread] = output

        if self.args.agg != 'cen':

            return output, fea_att, time_att,l1_regularization_list
        else:
            # return output
            return output, fea_att, time_att, l1_regularization_list



    def batch_run(self, training, optimizer_step=None, stats=None):
        # stats = stats or StatsLogger(('loss', 'correct'))
        stats = {}
        criterion = torch.nn.L1Loss(reduction='mean')
        ma = torch.nn.MSELoss(reduction='mean')
        self.model.train(training)
        for batch in self.dataloader:
            # print("batch",batch['input'].size())
            if self.args.agg != 'cen':
                output, fea_att, time_att, l1_regularization_list= self.model(batch)

            else:
                output, fea_att, time_att, l1_regularization_list = self.model(batch)

            if training:

                # print("output",output)
                if self.args.agg != 'cen':
                    loss = criterion(output, batch["target"])
                    stats['loss'] = loss
                    stats['mse'] = ma(output, batch["target"])
                    stats['rmse'] = torch.sqrt(stats['mse'])


                    stats['mape'] = torch.mean(torch.where(batch["target"].float() != 0, torch.abs(
                        (batch["target"].float() - output.float()) / batch["target"].float()),
                                                           torch.zeros_like(batch["target"].float()))) * 100
                    if self.args.dataset =='mapss':
                        stats['msle'] = robust_msle(output,batch["target"])
                    else:
                        stats['msle'] = ma(torch.log(output + 16), torch.log(batch["target"] + 16))  #

                else:

                    loss = criterion(output, batch["target"])
                    stats['loss'] = loss
                    stats['mse'] = ma(output, batch["target"])
                    stats['rmse'] = torch.sqrt(stats['mse'])



                    stats['mape'] = torch.mean(torch.where(batch["target"].float() != 0, torch.abs(
                        (batch["target"].float() - output.float()) / batch["target"].float()),
                                                           torch.zeros_like(batch["target"].float()))) * 100
                    stats['msle'] = ma(torch.log(output + 16), torch.log(batch["target"] + 16))  #

                loss.backward()

                optimizer_step()
                self.model.zero_grad()
            else:
                if self.args.agg != 'cen':


                    loss = criterion(output, batch["target"])
                    stats['loss'] = loss
                    stats['mse'] = ma(output, batch["target"])
                    stats['rmse'] = torch.sqrt(stats['mse'])


                    stats['mape'] = torch.mean(torch.where(batch["target"].float() != 0, torch.abs(
                        (batch["target"].float() - output.float()) / batch["target"].float()),
                                                         torch.zeros_like(batch["target"].float()))) * 100
                    stats['msle'] = ma(torch.log(output+ 16), torch.log(batch["target"] + 16))  #

                else:

                    loss = criterion(output, batch["target"])
                    stats['loss'] = loss
                    stats['mse'] = ma(output, batch["target"])
                    stats['rmse'] = torch.sqrt(stats['mse'])



                    stats['mape'] = torch.mean(torch.where(batch["target"].float() != 0, torch.abs(
                        (batch["target"].float() - output.float()) / batch["target"].float()),
                                                           torch.zeros_like(batch["target"].float()))) * 100
                    stats['msle'] = ma(torch.log(output + 16), torch.log(batch["target"] + 16))  #

            batch["input"].to("cpu")
            batch["target"].to("cpu")
        # if self.args.agg == 'attation_gat' or self.args.agg == 'attation_gcn' or self.args.agg == 'fen_attation'or self.args.agg == 'fedap' or self.args.agg == 'avg':
        if self.args.agg != 'cen' :

            return stats,fea_att, time_att,l1_regularization_list
        else:
            # return stats
            return stats, fea_att, time_att, l1_regularization_list

    def criterion(self, loss, mode):
        if self.args.agg == "avg":
            pass
        elif self.args.reg > 0 and mode != "PerTrain" and self.args.clients != 1:
            self.m1 = sd_matrixing(self.model.get_state()[0]).reshape(1, -1).to(self.args.device)
            self.m2 = sd_matrixing(self.server_param).reshape(1, -1).to(self.args.device)
            self.m3 = sd_matrixing(self.global_param).reshape(1, -1).to(self.args.device)
            self.reg1 = torch.nn.functional.pairwise_distance(self.m1, self.m2, p=2)
            self.reg2 = torch.nn.functional.pairwise_distance(self.m1, self.m3, p=2)
            loss = loss + 0.3 * self.reg1 + 0.3 * self.reg2
        return loss

    def free_memory(self):
        if self.m1 is not None:
            self.m1.to("cpu")
        if self.m2 is not None:
            self.m2.to("cpu")
        if self.m3 is not None:
            self.m3.to("cpu")
        if self.reg1 is not None:
            self.reg1.to("cpu")
        if self.reg2 is not None:
            self.reg2.to("cpu")

        torch.cuda.empty_cache()

    def logger(self, buf, p=False):
        if p:
            print(buf)
        # self.threadLock.acquire()
        with open(self.args.logDir, 'a+') as f:
            f.write(str(datetime.datetime.now()) + '\t' + buf + '\n')
        # self.threadLock.release()
