import torch

import torch.nn as nn

import numpy as np

import torchvision

import torch.nn.functional as F

from torch.utils.data import Dataset, DataLoader

import math

from torch.optim import Optimizer

from torch.optim.optimizer import required

from torch.nn.utils import clip_grad_norm_

import logging

import abc

import sys

from tqdm import tqdm_notebook

import torch.utils.data as D

import torch.nn.functional as F

from apex import amp 

#from .mymodels import out_to_predict,out_to_predict_in,out_to_predict_test,out_to_predict_test_simple

import copy

from scipy.spatial import distance_matrix

import matplotlib.pyplot as plt

from .mymodels import mean_model

def get_model_device(model):

    p = next(model.parameters())

    if p.is_cuda:

        device = torch.device("cuda:{}".format(p.get_device()))

    else:

        device = torch.device('cpu')

    return device

def model_train(model,optimizer,train_dataset,batch_size,num_epochs,loss_func,

                weights=None,accumulation_steps=1,

                weights_func=None,do_apex=True,validate_dataset=None,

                validate_loss=None,metric=None,param_schedualer=None,

                weights_data=False,history=None,return_model=False,model_apexed=False,

                num_workers=7,sampler=None,graph=None,k_lossf=0.01,pre_process=None,

                call_progress=None,use_batchs=True,best_average=1):

    if history is None:

        history = []

    num_average_models=min(num_epochs,best_average)

    best_models=np.empty(num_average_models+1,dtype=object)

    best_val_loss=1e6*np.ones(num_average_models+1,dtype=np.float)

    device = get_model_device(model)

    if do_apex and not model_apexed and (device.type=='cuda'):

        model_apexed=True

        model, optimizer = amp.initialize(model, optimizer, opt_level="O1",verbosity=0)

    model.zero_grad()

    tq_epoch=tqdm_notebook(range(num_epochs))

    lossf=None

    for epoch in tq_epoch:

        best_models[1:]=best_models[:num_average_models]

        best_val_loss[1:]=best_val_loss[:num_average_models]

        torch.cuda.empty_cache()

        model.do_grad()

        if param_schedualer:

            param_schedualer(epoch)

        _=model.train()

        batch_size_= batch_size if use_batchs else None

        if sampler:

            data_loader=D.DataLoader(D.Subset(train_dataset,sampler()),num_workers=num_workers,

                                     batch_size=batch_size if use_batchs else None,

                                     shuffle=use_batchs)

        else:

            data_loader=D.DataLoader(train_dataset,batch_size=batch_size,shuffle=True,num_workers=num_workers)

        sum_loss = 0.

        if metric:

            metric.zero()

        tq_batch = tqdm_notebook(data_loader,leave=True)

        model.zero_grad()

        for i,(batchs) in enumerate(tq_batch):

            x_batch=batchs[0].to(device) if len(batchs)==2 else [x.to(device) for x in batchs[:-1]]  

            y_batch=batchs[-1].to(device)

            if pre_process is not None:

                x_batch,y_batch = pre_process(x_batch,y_batch)

            if weights_data:

                weights=x_batch[-1]

                x_batch=x_batch[:-1]

            if weights_func:

                weights=weights_func(weights,epoch,i)

            y_preds = model(x_batch) if not isinstance(x_batch,list) else model(*x_batch)  

            if weights is not None:

                loss = loss_func(y_preds,y_batch,weights=weights)/accumulation_steps

            else:

                loss = loss_func(y_preds,y_batch)/accumulation_steps

            if model_apexed:

                with amp.scale_loss(loss, optimizer) as scaled_loss:

                    scaled_loss.backward()

            else:

                loss.backward()

            if (i+1) % accumulation_steps == 0:             # Wait for several backward steps

                optimizer.step()                            # Now we can do an optimizer step

                model.zero_grad()

            if lossf:

                lossf = (1-k_lossf)*lossf+k_lossf*loss.detach().item()*accumulation_steps

            else:

                lossf = loss.detach().item()*accumulation_steps

            if graph is not None:

                graph(lossf)

            batch_postfix={'loss':lossf}

            if metric:

                if isinstance(y_preds,tuple):

                    yp = tuple(y_preds[0].detach().cpu())

                else:

                    yp =y_preds.detach().cpu()

                batch_postfix.update(metric.calc(yp,y_batch.cpu().detach()))

            tq_batch.set_postfix(**batch_postfix)

            sum_loss=sum_loss+loss.detach().item()*accumulation_steps

        epoch_postfix={'loss':sum_loss/len(data_loader)}

        if metric:

            epoch_postfix.update(metric.calc_sums())

        tq_epoch.set_postfix(**batch_postfix)

        history.append(batch_postfix)

        if validate_dataset:

            if validate_loss is None:

                vloss = loss_func

                val_weights = weights

            else:

                vloss = validate_loss

                val_weights =None

            res=model_evaluate(model,

                               validate_dataset,

                               batch_size = batch_size if use_batchs else None ,

                               loss_func=vloss,

                               weights=val_weights,

                               metric=metric,

                               do_apex=False,

                               num_workers=num_workers)

            history[-1].update(res[1])

            best_val_loss[0] = res[0]

            best_models[0] = copy.deepcopy(model).to('cpu')

            best_models[0].no_grad()

            best_models=best_models[np.argsort(best_val_loss)]

            best_val_loss=best_val_loss[np.argsort(best_val_loss)]

#            if res[0]<best_val_loss:

#                best_val_loss=res[0]

#                best_model=copy.deepcopy(model)

#                best_model.no_grad()

            tq_epoch.set_postfix(res[1])

        print(history[-1])

        if call_progress is not None:

            call_progress(history)

    if num_average_models>1:

        best_model=mean_model(best_models[:num_average_models])

        model=model.to('cpu')

        best_model=best_model.to(device)

        res=model_evaluate(best_model,

                               validate_dataset,

                               batch_size = batch_size if use_batchs else None ,

                               loss_func=vloss,

                               weights=val_weights,

                               metric=metric,

                               do_apex=False,

                               num_workers=num_workers)

        best_model=best_model.to('cpu')

        model=model.to(device)

        if res[0]>best_val_loss[0]:

            best_model=best_models[0]

            print (best_val_loss[0])

        else:

            print (res)

    else:

        best_model=best_models[0]

        print (best_val_loss)

    return (history,best_model) if return_model else history

def model_run(model,dataset,do_apex=True,batch_size=32,num_workers=6):

    _=model.eval()

    model.no_grad()

    device = get_model_device(model)

    if do_apex and (device.type=='cuda'):

        model = amp.initialize(model, opt_level="O1",verbosity=0)

    res_list=[]

    data_loader=D.DataLoader(dataset,batch_size=batch_size,shuffle=False,num_workers=num_workers)

    for batchs in tqdm_notebook(data_loader):

        y_preds=model(*[x.to(device) for x in batchs]) if isinstance(batchs,tuple) else model(batchs.to(device))

        res_list.append(tuple([y.cpu() for y in y_preds]) if isinstance(y_preds,tuple) else y_preds.cpu())

    return tuple([torch.cat(tens) for tens in map(list, zip(*res_list))]) if isinstance(res_list[0],tuple) else torch.cat(res_list)

def models_run(models,dataset,do_apex=True,batch_size=32,num_workers=6):

    islist = isinstance(models,list)    

    if islist:

        models_=models

    else:

        models_=[models]

    for model in models_:

        _=model.eval()

        model.no_grad()

    device = get_model_device(models_[0])

    if do_apex and (device.type=='cuda'):

        for model in models_:

            model = amp.initialize(model, opt_level="O1",verbosity=0)

    res_list=[]

    for model in models_:        

        res_list.append([])

    data_loader=D.DataLoader(dataset,batch_size=batch_size,shuffle=False,num_workers=num_workers)

    for batchs in tqdm_notebook(data_loader):

        for i,model in enumerate(models_):

            y_preds=model(*[x.to(device) for x in batchs]) if isinstance(batchs,tuple) else model(batchs.to(device))

            res_list[i].append(tuple([y.cpu().detach() for y in y_preds]) if isinstance(y_preds,tuple) else y_preds.cpu().detach())

    res=[]

    for i in range(len(models_)):

        res.append(tuple([torch.cat(tens) for tens in map(list, 

                                                   zip(*res_list[i]))]) if isinstance(res_list[i][0],tuple) else torch.cat(res_list[i]))

    return tuple(zip(*res)) 

def model_evaluate(model,

                   validate_dataset,

                   batch_size,

                   loss_func,

                   weights=None,

                   metric=None,

                   do_apex=False,

                   num_workers=6):

    _=model.eval()

    model.no_grad()

    device = get_model_device(model)

    if do_apex and (device.type=='cuda'):

        model = amp.initialize(model, opt_level="O1",verbosity=0)

    data_loader=D.DataLoader(validate_dataset,batch_size=batch_size,shuffle=False,num_workers=num_workers)

    sum_loss = 0.

    lossf=None

    if metric:

        metric.zero()

    tq_batch = tqdm_notebook(data_loader,leave=True)

    for i,(batchs) in enumerate(tq_batch):

        x_batch=batchs[0].to(device) if len(batchs)==2 else [x.to(device) for x in batchs[:-1]]  

        y_batch=batchs[-1]

        y_preds = model(x_batch) if not isinstance(x_batch,list) else model(*x_batch)

        if weights is None:

            loss = loss_func(y_preds,y_batch.to(device))

        else:

            loss = loss_func(y_preds,y_batch.to(device),weights)

        sum_loss=sum_loss+loss.detach().item()

        if lossf:

            lossf = 0.98*lossf+0.02*loss.detach().item()

        else:

            lossf = loss.item()

        batch_postfix={'val_loss':lossf}

        if metric:

            if isinstance(y_preds,tuple):

                yp = tuple(y_preds[0].detach().cpu())

            else:

                yp =y_preds.detach().cpu()

            batch_postfix.update(metric.calc(yp,y_batch.cpu().detach(),prefix='val_'))

        tq_batch.set_postfix(**batch_postfix)

    epoch_postfix={'val_loss':sum_loss/len(data_loader)}

    if metric:

        epoch_postfix.update(metric.calc_sums('val_'))

    return sum_loss/len(data_loader), epoch_postfix

class loss_graph():

    def __init__(self,fig,ax,num_epoch=1,batch2epoch=100,limits=None):

        self.num_epoch=num_epoch

        self.batch2epoch=batch2epoch

        self.loss_arr=np.zeros(num_epoch*batch2epoch,dtype=np.float)

        self.arr_size=num_epoch*batch2epoch

        self.num_points=0

        self.fig=fig

        self.ax = ax

        self.limits=limits if limits is not None else (-1000,1000)

        self.ticks = (np.arange(0, num_epoch*batch2epoch+1, step=batch2epoch),np.arange(0, num_epoch+1, step=1))

    def __call__(self,loss):

        if self.num_points==self.arr_size:

            new_arr=np.zeros(self.arr_size+self.batch2epoch,dtype=np.float)

            new_arr[:self.arr_size]=self.loss_arr

            self.loss_arr=new_arr

        self.loss_arr[self.num_points]=max(self.limits[0],min(self.limits[1],loss))

        self.num_points=self.num_points+1

        _=self.ax.clear()

        _=self.ax.plot(self.loss_arr[0:self.num_points])

        _=self.ax.set_xlabel('batch')

        _=self.ax.set_ylabel('loss')

        _=plt.xticks(self.ticks[0],self.ticks[1])

        _=self.fig.canvas.draw()