import os
import random
import argparse
import yaml
from tqdm import tqdm

import torch
import torch.nn.functional as F
import torch.nn as nn
import torchvision.transforms as transforms

from datasets import build_dataset
from datasets.utils import build_data_loader
import clip
from utils_eurosat import *
from utils_eurosat import *

def get_arguments():
    parser = argparse.ArgumentParser()
    parser.add_argument('--config', default='/home/moming/code/VITCLIP/configs/eurosat.yaml', dest='config', help='settings of Tip-Adapter in yaml format')
    args = parser.parse_args()

    return args


def run_tip_adapter(cfg, cache_keys, cache_values, val_features, val_labels, test_features, test_labels, clip_weights):
    print("\n-------- Searching hyperparameters on the val set. --------")

    # Zero-shot CLIP
    clip_logits = 100. * val_features @ clip_weights
    acc = cls_acc(clip_logits, val_labels)
    print("\n**** Zero-shot CLIP's val accuracy: {:.2f}. ****\n".format(acc))

    # Tip-Adapter
    beta, alpha = cfg['init_beta'], cfg['init_alpha']

    affinity = val_features @ cache_keys
    cache_logits = ((-1) * (beta - beta * affinity)).exp() @ cache_values

    tip_logits = clip_logits + cache_logits * alpha
    acc = cls_acc(tip_logits, val_labels)
    print("**** Tip-Adapter's val accuracy: {:.2f}. ****\n".format(acc))

    # Search Hyperparameters
    best_beta, best_alpha = search_hp(cfg, cache_keys, cache_values, val_features, val_labels, clip_weights)

    print("\n-------- Evaluating on the test set. --------")

    # Zero-shot CLIP
    clip_logits = 100. * test_features @ clip_weights
    acc = cls_acc(clip_logits, test_labels)
    print("\n**** Zero-shot CLIP's test accuracy: {:.2f}. ****\n".format(acc))

    # Tip-Adapter
    affinity = test_features @ cache_keys
    cache_logits = ((-1) * (best_beta - best_beta * affinity)).exp() @ cache_values

    tip_logits = clip_logits + cache_logits * best_alpha
    acc = cls_acc(tip_logits, test_labels)
    print("**** Tip-Adapter's test accuracy: {:.2f}. ****\n".format(acc))


def run_tip_adapter_F(cfg, clss_head, test_features, test_labels, clip_weights,
                      clip_model, train_loader_F, max_epoch=200):
    device_id = clip_model.positional_embedding.device
    # Enable the cached keys to be learnable
    adapter = nn.Linear(clss_head.shape[1], clss_head.shape[0], bias=False).to(clip_model.dtype).cuda(device_id)
    adapter.weight = nn.Parameter(clss_head)
    # adapter.weight = nn.Parameter(clip_weights.t())

    optimizer = torch.optim.AdamW(
        [{'params': adapter.parameters()},
         # {'params': clip_weights},
         ],
        lr=cfg['lr'], eps=1e-4) # 5e-3 89.19
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, cfg['train_epoch'] * len(train_loader_F))

    # beta, alpha = cfg['init_beta'], cfg['init_alpha']
    best_acc, best_epoch = 0.0, 0

    ### zero shot XtremeCLIP
    adapter.eval()
    test_img_feat = test_features.unsqueeze(1).repeat(1, 10, 1)
    test_txt_feat = clip_weights.permute(1, 0).unsqueeze(0).repeat(test_features.size(0), 1, 1)
    test_final_feat = fusion_feat(test_txt_feat, test_img_feat)
    test_logits = adapter(test_final_feat).squeeze(-1)
    test_logits = torch.stack([torch.diag(l) for l in test_logits], dim=0)
    acc = cls_acc(test_logits, test_labels)
    ###
    for train_idx in range(cfg['train_epoch']):
        # Train
        adapter.train()
        correct_samples, all_samples = 0, 0
        loss_list = []
        print('Train Epoch: {:} / {:}'.format(train_idx, cfg['train_epoch']))

        for i, (images, target) in enumerate(tqdm(train_loader_F)):
            images, target = images.cuda(device_id), target.cuda(device_id)
            with torch.no_grad():
                image_features = clip_model.encode_image(images)
                image_features /= image_features.norm(dim=-1, keepdim=True)

            img_feat = image_features.unsqueeze(1).repeat(1, 10, 1)
            txt_feat = clip_weights.permute(1, 0).unsqueeze(0).repeat(images.size(0), 1, 1)
            final_feat = fusion_feat(txt_feat, img_feat)
            logits = adapter(final_feat).squeeze(-1)
            logits = torch.stack([torch.diag(l) for l in logits], dim=0)
            # cache_logits = ((-1) * (beta - beta * affinity)).exp() @ cache_values
            # clip_logits = image_features @ clip_weights
            # tip_logits = clip_logits + cache_logits * alpha
            log_prob = torch.softmax(logits, dim=-1)  # torch.log_softmax(logits, dim=-1)
            pos_score = log_prob.gather(-1, target.unsqueeze(-1).repeat(1, logits.size(1)))
            ones = torch.ones(pos_score.size()).cuda(device_id)
            loss_func = torch.nn.MarginRankingLoss(0.0)
            contr_loss = loss_func(pos_score, log_prob, ones)
            # logits = logits * alpha + clip_logits
            loss = F.cross_entropy(logits, target) + contr_loss

            acc = cls_acc(logits, target)
            correct_samples += acc / 100 * len(logits)
            all_samples += len(logits)
            loss_list.append(loss.item())

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            scheduler.step()

        current_lr = scheduler.get_last_lr()[0]
        print('LR: {:.6f}, Acc: {:.4f} ({:}/{:}), Loss: {:.4f}'.format(current_lr, correct_samples / all_samples,
                                                                       correct_samples, all_samples,
                                                                       sum(loss_list) / len(loss_list)))

        # Eval
        adapter.eval()
        test_img_feat = test_features.unsqueeze(1).repeat(1, 10, 1)
        test_txt_feat = clip_weights.permute(1, 0).unsqueeze(0).repeat(test_features.size(0), 1, 1)
        test_final_feat = fusion_feat(test_txt_feat, test_img_feat)
        test_logits = adapter(test_final_feat).squeeze(-1)
        test_logits = torch.stack([torch.diag(l) for l in test_logits], dim=0)

        acc = cls_acc(test_logits, test_labels)

        print("**** Tip-Adapter-F's test accuracy: {:.2f}. ****\n".format(acc))
        if acc > best_acc:
            best_acc = acc
            best_epoch = train_idx
            torch.save(adapter.weight, cfg['cache_dir'] + "/best_F_" + str(cfg['shots']) + "shots.pt")

    adapter.weight = torch.load(cfg['cache_dir'] + "/best_F_" + str(cfg['shots']) + "shots.pt")
    print(f"**** After fine-tuning, Tip-Adapter-F's best test accuracy: {best_acc:.2f}, at epoch: {best_epoch}. ****\n")

    print("\n-------- Searching hyperparameters on the val set. --------")


def main():
    # Load config file
    args = get_arguments()
    assert (os.path.exists(args.config))

    cfg = yaml.load(open(args.config, 'r'), Loader=yaml.Loader)

    cache_dir = os.path.join('./caches', cfg['dataset'])
    os.makedirs(cache_dir, exist_ok=True)
    cfg['cache_dir'] = cache_dir

    print("\nRunning configs.")
    print(cfg, "\n")

    # CLIP
    clip_model, preprocess = clip.load(cfg['backbone'], cfg['gpu'])
    clip_model.eval()

    # Prepare dataset
    random.seed(1)
    torch.manual_seed(1)

    print("Preparing dataset.")
    dataset = build_dataset(cfg['dataset'], cfg['root_path'], cfg['shots'])

    val_loader = build_data_loader(data_source=dataset.val, batch_size=64, is_train=False, tfm=preprocess,
                                   shuffle=False)
    test_loader = build_data_loader(data_source=dataset.test, batch_size=64, is_train=False, tfm=preprocess,
                                    shuffle=False)

    train_tranform = transforms.Compose([
        transforms.RandomResizedCrop(size=224, scale=(0.5, 1), interpolation=transforms.InterpolationMode.BICUBIC),
        transforms.RandomHorizontalFlip(p=0.5),
        transforms.ToTensor(),
        transforms.Normalize(mean=(0.48145466, 0.4578275, 0.40821073), std=(0.26862954, 0.26130258, 0.27577711))
    ])

    train_loader_cache = build_data_loader(data_source=dataset.train_x, batch_size=256, tfm=train_tranform,
                                           is_train=True, shuffle=False)
    train_loader_F = build_data_loader(data_source=dataset.train_x, batch_size=16, tfm=train_tranform, is_train=True,
                                       shuffle=True)

    # Textual features
    print("\nGetting textual features as CLIP's classifier.")
    clip_weights = clip_classifier(dataset.classnames, dataset.template, clip_model)

    # Construct the cache model by few-shot training set
    print("\nConstructing cache model by few-shot visual features and labels.")
    cache_keys, cache_values, clss_head = build_cache_model(cfg, clip_model, train_loader_cache, clip_weights)

    # Pre-load val features
    print("\nLoading visual features and labels from val set.")
    val_features, val_labels = pre_load_features(cfg, "val", clip_model, val_loader)

    # Pre-load test features
    print("\nLoading visual features and labels from test set.")
    test_features, test_labels = pre_load_features(cfg, "test", clip_model, test_loader)

    # ------------------------------------------ Tip-Adapter ------------------------------------------
    # run_tip_adapter(cfg, cache_keys, cache_values, val_features, val_labels, test_features, test_labels, clip_weights)

    # ------------------------------------------ Tip-Adapter-F ------------------------------------------
    run_tip_adapter_F(cfg, clss_head, test_features, test_labels, clip_weights,
                      clip_model, train_loader_F)


if __name__ == '__main__':
    main()