import os

os.environ['CUDA_VISIBLE_DEVICES'] = '0'
from scipy.io import savemat
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset

import time
import datetime
import numpy as np
from config_test import load_args
from datapre.data_read import readdata
from record.generate_pic import generate
from record.generate_pic import feature_map

from record.record import record_output
from datapre.hyper_dataset import HyperData
from datapre.augment import CenterResizeCrop
from util_CNN_clip import test_batch
from Mutil_prompt_CLIP.MutilCLIP import build_model
from collections import OrderedDict
# get class name
def read_classnames(text_file):
    """Return a dictionary containing
    key-value pairs of <folder name>: <class name>.
    """
    classnames = OrderedDict()
    with open(text_file, "r") as f:
        lines = f.readlines()
        for line in lines:
            line = line.strip().split(" ")
            folder = line[0]
            classname = " ".join(line[1:])
            classnames[folder] = classname
    return classnames
print(torch)
args = load_args()
mask_ratio = args.mask_ratio
windowsize = args.windowsize
dataset = args.dataset
type = args.type
num_epoch = args.epochs
num_fine_tuned = args.fine_tuned_epochs
lr = args.lr
train_num_per = args.train_num_perclass
num_of_ex = 16
batch_size = args.batch_size
basemodel = args.basemodel
proportion = args.proportion
net_name = 'MutilCLIP'
day = datetime.datetime.now()
day_str = day.strftime('%m_%d_%H_%M')
halfsize = int((windowsize - 1) / 2)
val_num = args.valid_num
Seed = 0
_, _, _, validation_image, validation_image_LIDAR, validation_label, _, _, _, _, _, _, image_LiDAR, gt, s = readdata(type, dataset, windowsize, train_num_per, val_num, 0,proportion)
nclass = np.max(gt).astype(np.int64)
print(nclass)
classnames_dict = read_classnames("classname/classnames_" + dataset + ".txt")
description_dict = read_classnames("classname/des_" + dataset + ".txt")
classnames = list(classnames_dict.values())
description = list(description_dict.values())
is_feature_map = args.is_feature_map
if is_feature_map:
    feature_map(validation_image_LIDAR,validation_label,nclass,classnames)
in_chans_LIDAR = image_LiDAR.shape[2]
KAPPA = []
OA = []
AA = []
TRAINING_TIME = []
TESTING_TIME = []
ELEMENT_ACC = np.zeros((num_of_ex, nclass))
af_result = np.zeros([nclass + 3, num_of_ex])
criterion = nn.CrossEntropyLoss()

for num in range(4, num_of_ex):
    print('num:', num)
    train_image, train_image_LIDAR, train_label, validation_image1, validation_image_LIDAR1, validation_label1, nTrain_perClass, nvalid_perClass, \
        train_index, val_index, index, image, image_LiDAR, gt, s = readdata(type, dataset, windowsize, train_num_per,val_num, num,proportion)
    nvalid_perClass = np.zeros_like(nvalid_perClass)
    nband = train_image.shape[3]

    train_num = train_image.shape[0]
    train_image = np.transpose(train_image, (0, 3, 1, 2))
    train_image_LIDAR = np.transpose(train_image_LIDAR, (0, 3, 1, 2))

    validation_image = np.transpose(validation_image, (0, 3, 1, 2))
    validation_image1 = np.transpose(validation_image1, (0, 3, 1, 2))
    validation_image_LIDAR = np.transpose(validation_image_LIDAR, (0, 3, 1, 2))
    validation_image_LIDAR1 = np.transpose(validation_image_LIDAR1, (0, 3, 1, 2))

    if args.augment:
        transform_train = [CenterResizeCrop(scale_begin=args.scale, windowsize=windowsize)]
        untrain_dataset = HyperData((train_image, train_image_LIDAR, train_label), transform_train)
    else:
        untrain_dataset = TensorDataset(torch.tensor(train_image), torch.tensor(train_image_LIDAR),
                                        torch.tensor(train_label))
    untrain_loader = DataLoader(dataset=untrain_dataset, batch_size=batch_size, shuffle=True)

    print("=> creating model '{}'".format(net_name))
    basemodel = args.basemodel
    # ########################   finetune    # ########################

    model = build_model(img_size=windowsize, in_chans=nband, in_chans_LIDAR=in_chans_LIDAR,
                        hid_chans=args.hid_chans, hid_chans_LIDAR=args.hid_chans_LIDAR, embed_dim=args.encoder_dim,num_classes=nclass,
                        deep_prompts_depth=args.deep_prompts_depth, num_heads=args.encoder_num_heads, mlp_ratio=args.mlp_ratio, classnames=classnames,
                        global_pool=False, basemodel=basemodel, input_image_size=224, init_prompt="a multimodal fusion patch of", num_prompt=6,trainer=args.trainer,prompt_nctx=args.prompt_nctx,description=description).cuda()

    tic2 = time.time()
    optimizer = optim.Adam(model.parameters(), lr=5e-4, weight_decay=1e-4, eps=1e-4)
    scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=[70,10], gamma=0.2, last_epoch=-1)
    toc2 = time.time()
    model.load_state_dict(torch.load('../Best_val_model/net_params.pkl'))
    true_cla, overall_accuracy, average_accuracy, kappa, true_label, test_pred, test_index, cm, pred_array = test_batch(
        model.eval(), image, image_LiDAR, index, 32, nTrain_perClass, nvalid_perClass, halfsize, classnames,
        is_feature_map)
    toc3 = time.time()

    af_result[:nclass, num] = true_cla
    af_result[nclass, num] = overall_accuracy
    af_result[nclass + 1, num] = average_accuracy
    af_result[nclass + 2, num] = kappa
    OA.append(overall_accuracy)
    AA.append(average_accuracy)
    KAPPA.append(kappa)
    TRAINING_TIME.append(toc2 - tic2)
    print("TRAINING_TIME"+str(toc2 - tic2))
    TESTING_TIME.append(toc3 - toc2)
    print("TESTING_TIME"+str(toc3 - toc2))
    ELEMENT_ACC[num, :] = true_cla
    classification_map, gt_map = generate(image, gt, index, nTrain_perClass, nvalid_perClass, test_pred,
                                          overall_accuracy, halfsize, dataset, day_str, num, net_name)
    file_name = 'result/' + dataset + '/'  + str(overall_accuracy) + '.mat'
    savemat(file_name, {'map': classification_map})
result = np.mean(af_result, axis=1)
print(result)

print("--------" + net_name + " Training Finished-----------")
record_output(OA, AA, KAPPA, ELEMENT_ACC, TRAINING_TIME, TESTING_TIME,
                     'records/' + dataset + '/' + net_name + '_' + day_str + '_' + str(args.epochs) + '_' + str(
                         args.fine_tuned_epochs) + '_train_num:' + str(train_image.shape[0]) +'batch_size' + str(batch_size)+
                     '_augment_' + str(args.augment) + '_aug_scale_' + str(args.scale) + '_decoder_depth_' + str(
                         args.deep_prompts_depth) + '.txt')