import datetime
import time

import torch
import numpy as np
from sklearn import metrics
from torch.cuda.amp import GradScaler, autocast
from torch.nn.modules.loss import _Loss
from torch.utils.data import TensorDataset, DataLoader
import torch.nn as nn
from record.generate_pic import feature_map
from tqdm import tqdm




ce_loss = nn.CrossEntropyLoss()
MSE=nn.MSELoss()
def tr_acc(model, image, image_LIDAR, label):


    train_dataset = TensorDataset(torch.tensor(image), torch.tensor(image_LIDAR), torch.tensor(label))

    train_loader = DataLoader(dataset=train_dataset, batch_size=32, shuffle=False)
    train_loss = 0
    corr_num = 0
    T = 1.0
    ce_loss = 0
    # criterion = AppleLoss(T)
    for idx, (image_batch, image_LIDAR_batch, label_batch) in tqdm(enumerate(train_loader),total=len(train_loader),desc="Trainval", unit="epoch"):

        trans_image_batch = image_batch.cuda()
        trans_image_LIDAR_batch = image_LIDAR_batch.cuda()
        label_batch = label_batch.cuda()
        criterion = AppleLoss(T)

        logits,ctx_shifted = model(trans_image_batch, trans_image_LIDAR_batch)

        if isinstance(logits, tuple):
            logits = logits[-1]

        pred = torch.max(logits, dim=1)[1]

        loss,ce_loss = criterion(logits,ctx_shifted, label_batch)

        train_loss = train_loss + loss.cpu().data.numpy()

        corr_num = torch.eq(pred, label_batch).float().sum().cpu().numpy() + corr_num

    return corr_num / image.shape[0], train_loss,ce_loss
class AppleLoss(_Loss):
    def __init__(self, T):
        super(AppleLoss, self).__init__()
        self.T = T

    def forward(self, logits, ctx_shifted, label):
        '''''
        Cross Entropy (CE) Loss
        '''

        ce_loss1 = ce_loss(logits, label)

        '''''
        Context Redundancy Penalizing (CRP) Loss
        '''

        transprompts = ctx_shifted.permute(1, 0)
        pmul = torch.matmul(ctx_shifted / self.T, transprompts / self.T)
        identity = torch.eye(ctx_shifted.size(0)).to(torch.float16).cuda()
        # mask = identity.repeat(pmul.shape[0], 1, 1).cuda()
        # crp_loss = ((pmul/ self.T - identity/ self.T) ** 2).mean()
        # crp_loss = torch.linalg.det(torch.sub(pmul / self.T, mask / self.T)).mean()
        MSE_loss = MSE(pmul / self.T, identity / self.T)
        total_loss = ce_loss1 + 0.02 * MSE_loss
        mse_loss = 0.02 * MSE_loss
        # lambda = 0.05 : weight
        # total_loss = ce_loss1
        # print(0.05 * torch.log10(crp_loss))
        return total_loss,mse_loss

from clip.simple_tokenizer import SimpleTokenizer as _Tokenizer
_tokenizer = _Tokenizer()

def pre_train(
    model,
    train_image,
    train_image_LIDAR,
    train_label,
    validation_image,
    validation_image_LIDAR,
    validation_label,
    epoch,
    optimizer,
    scheduler,
    bs,
    classnames,
    val=False,
):

    global corr_num
    train_dataset = TensorDataset(
        torch.tensor(train_image),
        torch.tensor(train_image_LIDAR),
        torch.tensor(train_label),
    )
    train_loader = DataLoader(
        dataset=train_dataset, batch_size=bs, shuffle=True
    )

    Train_loss = []
    Train_acc = []
    Val_loss = []
    Val_acc = []
    train_losses = []

    BestAcc = 0.0
    Best_tr_loss = 1000
    scaler = GradScaler()
    T = 1.0
    criterion = AppleLoss(T)

    for i in range(epoch):
        model.train()
        train_loss = 0
        corr_num = 0
        crp_loss = 0
        grad_clip = 1.0
        MSE_loss = 0

        for idx, (image_batch, image_LIDAR_batch, label_batch) in tqdm(enumerate(
            train_loader), total=len(train_loader),desc="Training", unit="epoch"):


            trans_image_batch = image_batch.cuda()
            trans_image_LIDAR_batch = image_LIDAR_batch.cuda()
            label_batch = label_batch.cuda()


            # text = text_target(classnames, label_batch).cuda()


            logits_per_image_x, ctx_shifted = model(trans_image_batch, trans_image_LIDAR_batch)

            if isinstance(logits_per_image_x, tuple):
                logits_per_image_x = logits_per_image_x[-1]

            total_loss,mse_loss = criterion(logits_per_image_x, ctx_shifted, label_batch)


            optimizer.zero_grad()
            total_loss.sum().backward()
            # torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)
            optimizer.step()
            MSE_loss = MSE_loss + mse_loss
            train_loss = train_loss + total_loss
            pred = torch.max(logits_per_image_x, dim=1)[1]
            corr_num = torch.eq(pred, label_batch).float().sum().cpu().numpy() + corr_num

        scheduler.step()


        train_loss = train_loss / (idx + 1)
        train_losses.append(train_loss)
        MSE_loss = MSE_loss / (idx + 1)
        train_acc = corr_num / train_image.shape[0]

        toc2 = time.time()


        model.eval()
        val_acc, val_loss,vce_loss= tr_acc(
            model,
            validation_image,
            validation_image_LIDAR,
            validation_label,
        )


        if val_acc > BestAcc:
            if val_acc > 0.5:
                torch.save(model.state_dict(), "Best_val_model/" + "net_params.pkl")
            BestAcc = val_acc
        print(
            "epoch {}, training loss: {:.4f},MSE_loss: {:.4f} train acc:{:.4f}, valid acc:{:.4f},best acc:{:.4f}".format(
                i, train_loss.item(),MSE_loss.item(), train_acc * 100, val_acc * 100,BestAcc*100
            )
        )

        if val:
            Train_loss.append(train_loss.cpu().detach().numpy())
            Val_loss.append(val_loss)
            Train_acc.append(train_acc)
            Val_acc.append(val_acc)


    if val:
        return model, [Train_loss, Train_acc, Val_loss, Val_acc]
    else:
        return model

def test_batch(model, image, image_LIDAR, index, BATCH_SIZE,  nTrain_perClass, nvalid_perClass, halfsize,classnames,is_feature_map):

    ind = index[0][nTrain_perClass[0]:,:]
    nclass = len(index)
    true_label = np.zeros(ind.shape[0], dtype = np.int32)

    for i in range(1, nclass):
        ddd = index[i][nTrain_perClass[i]:,:]
        ind = np.concatenate((ind, ddd), axis = 0)
        tr_label = np.ones(ddd.shape[0], dtype = np.int32) * i
        true_label = np.concatenate((true_label, tr_label), axis = 0)
    test_index = np.copy(ind)
    length = ind.shape[0]

    if length % BATCH_SIZE != 0:
        add_num = BATCH_SIZE - length % BATCH_SIZE
        ff = range(length)
        add_ind = np.random.choice(ff, add_num, replace = False)
        add_ind = ind[add_ind]
        ind = np.concatenate((ind,add_ind), axis =0)

    pred_array = np.zeros([ind.shape[0],nclass], dtype = np.float32)
    n = ind.shape[0] // BATCH_SIZE
    windowsize = 2 * halfsize + 1
    image_batch = np.zeros([BATCH_SIZE, windowsize, windowsize, image.shape[2]], dtype=np.float32)
    image_LIDAR_batch = np.zeros([BATCH_SIZE, windowsize, windowsize, image_LIDAR.shape[2]], dtype=np.float32)

    for i in range(n):
        for j in range(BATCH_SIZE):
            m = ind[BATCH_SIZE*i+j, :]
            image_batch[j,:,:,:] = image[(m[0] - halfsize):(m[0] + halfsize + 1),
                                                   (m[1] - halfsize):(m[1] + halfsize + 1),:]
            image_b = np.transpose(image_batch,(0,3,1,2))
            image_LIDAR_batch[j,:,:,:] = image_LIDAR[(m[0] - halfsize):(m[0] + halfsize + 1),
                                                   (m[1] - halfsize):(m[1] + halfsize + 1),:]
            image_LIDAR_b = np.transpose(image_LIDAR_batch,(0,3,1,2))
        logits,ctx_shifted = model(torch.tensor(image_b).cuda(), torch.tensor(image_LIDAR_b).cuda())
        if isinstance(logits,tuple):
            logits = logits[-1]
        pred_array[i*BATCH_SIZE:(i+1)*BATCH_SIZE] = torch.softmax(logits, dim = 1).cpu().data.numpy()
    pred_array = pred_array[range(length)]
    predict_label  = np.argmax(pred_array, axis=1)
    if is_feature_map:
        feature_map(pred_array,predict_label,nclass,classnames)

    confusion_matrix = metrics.confusion_matrix(true_label, predict_label)
    overall_accuracy = metrics.accuracy_score(true_label, predict_label)
    
    true_cla = np.zeros(nclass,  dtype=np.int64)
    for i in range(nclass):
        true_cla[i] = confusion_matrix[i,i]
    test_num_class = np.sum(confusion_matrix,1)
    test_num = np.sum(test_num_class)
    num1 = np.sum(confusion_matrix,0)
    po = overall_accuracy
    pe = np.sum(test_num_class*num1)/(test_num*test_num)
    kappa = (po-pe)/(1-pe)*100
    true_cla = np.true_divide(true_cla,test_num_class)*100 
    average_accuracy = np.average(true_cla)
    print('overall_accuracy: {0:f}'.format(overall_accuracy*100)) 
    print('average_accuracy: {0:f}'.format(average_accuracy))  
    print('kappa:{0:f}'.format(kappa))
    for i in range(nclass):
        print("Class " + str(i) + ": " + "{:.2f}".format( true_cla[i]))
    return true_cla, overall_accuracy*100, average_accuracy, kappa, true_label, predict_label, test_index, confusion_matrix, pred_array
