import os
import os.path as osp
import time
import math
import datetime
import numpy as np
import torch
import torch.nn as nn
from torch.nn import functional as F
from torch.cuda.amp import GradScaler, autocast
import torch.optim as optim
import json
from dassl.engine import TRAINER_REGISTRY, TrainerX
from dassl.metrics import compute_accuracy
from dassl.utils import load_pretrained_weights, load_checkpoint, MetricMeter, AverageMeter
from dassl.optim import build_optimizer, build_lr_scheduler

from clip import clip
from clip.simple_tokenizer import SimpleTokenizer as _Tokenizer

_tokenizer = _Tokenizer()

idx_l = [0, 1000, (1000 + 196), (1000 + 196 + 397), (1000 + 196 + 397 + 100), (1000 + 196 + 397 + 100 + 37),
         (1000 + 196 + 397 + 100 + 37 + 102), (1000 + 196 + 397 + 100 + 37 + 102 + 47),
         (1000 + 196 + 397 + 100 + 37 + 102 + 47 + 100),
         (1000 + 196 + 397 + 100 + 37 + 102 + 47 + 100 + 10), (1000 + 196 + 397 + 100 + 37 + 102 + 47 + 100 + 10 + 101),
         (1000 + 196 + 397 + 100 + 37 + 102 + 47 + 100 + 10 + 101 + 101)]
DG_IDX = {"ImageNet": [idx_l[0], idx_l[1]], "StanfordCars": [idx_l[1], idx_l[2]], "SUN397": [idx_l[2], idx_l[3]],
          "Caltech101": [idx_l[3], idx_l[4]], 'OxfordPets': [idx_l[4], idx_l[5]], "OxfordFlowers": [idx_l[5], idx_l[6]],
          "DescribableTextures": [idx_l[6], idx_l[7]], "FGVCAircraft": [idx_l[7], idx_l[8]], "EuroSAT": [idx_l[8], idx_l[9]],
          "UCF101": [idx_l[9], idx_l[10]], "Food101": [idx_l[10], idx_l[11]]}

CUSTOM_TEMPLATES = {
    "OxfordPets": "a photo of a {}, a type of pet.",
    "OxfordFlowers": "a photo of a {}, a type of flower.",
    "FGVCAircraft": "a photo of a {}, a type of aircraft.",
    "DescribableTextures": "{} texture.",
    "EuroSAT": "a centered satellite photo of {}.",
    "StanfordCars": "a photo of a {}.",
    "Food101": "a photo of {}, a type of food.",
    "SUN397": "a photo of a {}.",
    "Caltech101": "a photo of a {}.",
    "UCF101": "a photo of a person doing {}.",
    "ImageNet": "a photo of a {}.",
    "ImageNetSketch": "a photo of a {}.",
    "ImageNetV2": "a photo of a {}.",
    "ImageNetA": "a photo of a {}.",
    "ImageNetR": "a photo of a {}.",
}

CLASS_N = {"StanfordCars": 196, "SUN397": 397, "Caltech101": 100, "OxfordPets": 37, 'OxfordFlowers': 102,
           "DescribableTextures": 47, "FGVCAircraft": 100, "EuroSAT": 10, "UCF101": 101, "Food101": 101,
           "ImageNet": 1000, } # "ImageNetV2": 1000, "ImageNetA": 1000, "ImageNetR": 1000 , "ImageNetSketch": 1000

def load_clip_to_cpu(cfg):
    backbone_name = cfg.MODEL.BACKBONE.NAME
    url = clip._MODELS[backbone_name]
    model_path = clip._download(url)

    try:
        # loading JIT archive
        model = torch.jit.load(model_path, map_location="cpu").eval()
        state_dict = None

    except RuntimeError:
        state_dict = torch.load(model_path, map_location="cpu")

    model = clip.build_model(state_dict or model.state_dict())

    return model


class TextEncoder(nn.Module):
    def __init__(self, clip_model):
        super().__init__()
        self.transformer = clip_model.transformer
        self.positional_embedding = clip_model.positional_embedding
        self.ln_final = clip_model.ln_final
        self.text_projection = clip_model.text_projection
        self.dtype = clip_model.dtype

    def forward(self, prompts, tokenized_prompts, no_projection=False):
        x = prompts + self.positional_embedding.type(self.dtype)
        x = x.permute(1, 0, 2)  # NLD -> LND
        x = self.transformer(x)
        x = x.permute(1, 0, 2)  # LND -> NLD
        x = self.ln_final(x).type(self.dtype)

        x = x[torch.arange(x.shape[0]), tokenized_prompts.argmax(dim=-1)]
        if no_projection:
            return x

        x = x @ self.text_projection
        return x


class CustomCLIP(nn.Module):
    def __init__(self, cfg, classnames, clip_model, all_class_names, class_taken_idxs=None):
        super().__init__()
        self.cfg = cfg
        self.class_taken_idxs = class_taken_idxs
        classnames = all_class_names
        print(classnames)
        self.n_cls = len(classnames)
        self.n_base = math.ceil(self.n_cls / 2)
        self.n_new = self.n_cls - math.ceil(self.n_cls / 2)
        print("n_cls, n_base, n_new", self.n_cls, self.n_base, self.n_new)
        clip_imsize = clip_model.visual.input_resolution
        cfg_imsize = cfg.INPUT.SIZE[0]
        assert cfg_imsize == clip_imsize, f"cfg_imsize ({cfg_imsize}) must equal to clip_imsize ({clip_imsize})"

        temp = CUSTOM_TEMPLATES[cfg.DATASET.NAME]
        with open("gpt_file/"+cfg.DATASET.NAME+".json") as f:
            gpt3_prompt = json.load(f)
        if cfg.DG:
            for root, dirs, files in os.walk("/gpt_file/", topdown=False):
                for file in files:
                    print(file)
                    if file == cfg.DATASET.NAME+".json":
                        continue
                    with open(root+file) as f:
                        gpt3_prompt.update(json.load(f))

        prompts = [[temp.format(c.replace("_", " "))] + gpt3_prompt[c.replace("_", " ")][:cfg.TRAINER.APPLE.n_text-1]
                                                                    for c in classnames]
        prompts = torch.cat([clip.tokenize(p) for p in prompts])
        clip_model.cuda()

        with torch.no_grad():
            all_class_text_f = torch.zeros(prompts.shape[0], 512).half()
            if prompts.shape[0] > 10000:
                for i in range(int(prompts.shape[0]/(self.n_cls*5))):
                    all_class_text_f[i*(self.n_cls*5):(i+1)*(self.n_cls*5)] = clip_model.encode_text(prompts[i*(self.n_cls*5):(i+1)*(self.n_cls*5)].cuda())
            else:
                all_class_text_f = clip_model.encode_text(prompts.cuda())
            text_features = all_class_text_f / all_class_text_f.norm(dim=-1, keepdim=True)
            self.register_buffer("text_features", all_class_text_f.cpu())
        self.register_parameter("text_features_learnable", nn.Parameter(text_features))  # SOS

        ones = torch.ones(self.n_cls, cfg.TRAINER.APPLE.n_text).half()
        ones = nn.Parameter(ones, requires_grad=True)
        self.register_parameter("attention", ones)

        self.image_encoder = clip_model.visual
        self.text_encoder = TextEncoder(clip_model)
        self.logit_scale = clip_model.logit_scale
        self.dtype = clip_model.dtype

        dset = cfg.ori_data_name
        self.idx_cls_1 = DG_IDX[dset][0]
        self.idx_cls_2 = DG_IDX[dset][1]

    def forward(self, image, label=None):
        if self.cfg.DATASET.NAME in ["ImageNetA", "ImageNetR"] and self.attention.shape[0] == 1000:
            self.text_features_learnable = torch.nn.Parameter(
                self.text_features_learnable.view(1000, 50, 512)[self.class_taken_idxs].view(-1, 512))
            self.attention = torch.nn.Parameter(self.attention[self.class_taken_idxs])

        image_features = self.image_encoder(image.type(self.dtype))
        image_features = image_features / image_features.norm(dim=-1, keepdim=True)

        logit_scale = self.logit_scale.exp()

        if self.cfg.eval_only:
            if self.cfg.DATASET.SUBSAMPLE_CLASSES == "base":
                text_features = self.text_features_learnable[:self.n_base*self.cfg.TRAINER.APPLE.n_text] / \
                                self.text_features_learnable[:self.n_base*self.cfg.TRAINER.APPLE.n_text].norm(dim=-1, keepdim=True)
            elif self.cfg.DATASET.SUBSAMPLE_CLASSES == "new":
                text_features = self.text_features_learnable[self.n_base * self.cfg.TRAINER.APPLE.n_text:] / \
                                self.text_features_learnable[self.n_base * self.cfg.TRAINER.APPLE.n_text:].norm(dim=-1, keepdim=True)
            elif self.cfg.DATASET.SUBSAMPLE_CLASSES == "all":
                text_features = self.text_features_learnable / \
                                self.text_features_learnable.norm(dim=-1, keepdim=True)
        else:
            text_features = self.text_features_learnable / \
                            self.text_features_learnable.norm(dim=-1, keepdim=True)

        logits = image_features @ text_features.t()
        logits = logits.reshape(logits.shape[0], -1, self.cfg.TRAINER.APPLE.n_text)

        if self.cfg.eval_only:
            if self.cfg.DATASET.SUBSAMPLE_CLASSES == "base":
                logits = logit_scale * torch.mul(logits, self.attention[:self.n_base:])
            if self.cfg.DATASET.SUBSAMPLE_CLASSES == "new":
                logits = logit_scale * torch.mul(logits, self.attention[self.n_base:])
        else:
            logits = torch.mul(logits, self.attention)
            proto_l2_norm = torch.norm(logits, p=2, dim=2)
            l2_loss = self.cfg.TRAINER.APPLE.l2_w * proto_l2_norm.mean()
            logits = logit_scale * logits

        logits_max = F.max_pool1d(logits, kernel_size=self.cfg.TRAINER.APPLE.n_text).squeeze()
        if self.training:
            logits = logits.mean(dim=2)
            return F.cross_entropy(logits, label) + self.cfg.TRAINER.APPLE.max_w * F.cross_entropy(logits_max, label) + l2_loss
        if self.cfg.DG and self.cfg.eval_only:
            logits = logits[:, self.idx_cls_1:self.idx_cls_2].clone()
            logits_max = logits_max[:, self.idx_cls_1:self.idx_cls_2]
        logits = 0.8*logits.mean(dim=2) + 0.2*logits_max #logits.max(dim=2) logits.mean(dim=2) #

        if self.cfg.eval_only and self.cfg.DATASET.SUBSAMPLE_CLASSES == "new":
            logits_ = torch.zeros(logits.shape[0], self.n_cls).cuda()
            logits_[:, self.n_base:] = logits
            logits = logits_
        return logits

@TRAINER_REGISTRY.register()
class APPLe(TrainerX):
    """
    Adaptive Prompt Prototype Learning (APPLe).
    """

    def check_cfg(self, cfg):
        assert cfg.TRAINER.APPLE.PREC in ["fp16", "fp32", "amp"]

    def build_model(self):
        cfg = self.cfg
        classnames = self.dm.dataset.classnames

        all_class_names = self.dm.dataset.all_class_names

        self.class_n = CLASS_N[cfg.DATASET.NAME]

        assert self.class_n == len(all_class_names), f"all class names ({len(all_class_names)}) must " \
                                                     f"equal to predefined classes ({self.class_n})"
        cfg.ori_data_name = self.cfg.DATASET.NAME
        if cfg.DG:
            print("************DG*****************")
            # all_class_names = []
            # all_class_names.extend(classnames)
            print("self.cfg.DATASET.NAME: ", self.cfg.DATASET.NAME)
            from dassl.data import DataManager
            cfg.DATASET.NAME = "ImageNet"
            dm = DataManager(cfg)
            ds_class_names = dm.dataset.all_class_names
            all_class_names = ds_class_names

            for ds in CLASS_N.keys():
                print("DS:", ds, "DATASETNAME", cfg.DATASET.NAME)
                print("")
                if ds.strip() == "ImageNet":
                    continue
                cfg.DATASET.NAME = ds
                dm = DataManager(cfg)
                ds_class_names = dm.dataset.all_class_names
                all_class_names.extend(ds_class_names)

            self.class_n = len(all_class_names)
            print("class_n", self.class_n)
            assert self.class_n == len(all_class_names), f"all class names ({len(all_class_names)}) must " \
                                                     f"equal to predefined classes ({self.class_n})"
            #############

        print(f"Loading CLIP (backbone: {cfg.MODEL.BACKBONE.NAME})")
        clip_model = load_clip_to_cpu(cfg)

        if cfg.TRAINER.APPLE.PREC == "fp32" or cfg.TRAINER.APPLE.PREC == "amp":
            # CLIP's default precision is fp16
            clip_model.float()
        self.clip_model = clip_model
        print("Building custom CLIP")
        if self.cfg.DATASET.NAME in ["ImageNetA", "ImageNetR"]:
            self.model = CustomCLIP(cfg, classnames, clip_model, all_class_names, class_taken_idxs=self.dm.dataset.class_taken_idxs)
        else:
            self.model = CustomCLIP(cfg, classnames, clip_model, all_class_names)

        print("Turning off gradients in both the image and the text encoder")

        update_list = ['text_features_learnable', 'attention']

        for name, param in self.model.named_parameters():
            if name not in update_list and "prompt_learner" not in name:
                param.requires_grad_(False)

        # Double check
        enabled = set()
        for name, param in self.model.named_parameters():
            if param.requires_grad:
                enabled.add(name)
        print(f"Parameters to be updated: {enabled}")

        if cfg.MODEL.INIT_WEIGHTS:
            load_pretrained_weights(self.model, cfg.MODEL.INIT_WEIGHTS)

        self.model.to(self.device)

        self.scaler = GradScaler() if cfg.TRAINER.APPLE.PREC == "amp" else None
        self.model.device = self.device

        self.optim = build_optimizer(self.model, cfg.OPTIM)
        self.sched = build_lr_scheduler(self.optim, cfg.OPTIM)
        self.register_model("text_features", self.model, self.optim, self.sched)

    def forward_backward(self, batch):
        image, label = self.parse_batch_train(batch)
        model = self.model
        optim = self.optim
        scaler = self.scaler

        prec = self.cfg.TRAINER.APPLE.PREC
        if prec == "amp":
            with autocast():
                output = self.model(image)
                loss = F.cross_entropy(output, label)
            self.optim.zero_grad()
            self.scaler.scale(loss).backward()
            self.scaler.step(self.optim)
            self.scaler.update()
        else:
            loss = model(image, label)
            optim.zero_grad()
            loss.backward()
            optim.step()

        loss_summary = {
            "loss": loss.item(),
        }

        if (self.batch_idx + 1) == self.num_batches:
            self.update_lr()

        return loss_summary

    def parse_batch_train(self, batch):
        input = batch["img"]
        label = batch["label"]
        input = input.to(self.device)
        label = label.to(self.device)
        return input, label

    def load_model(self, directory, epoch=None):
        if not directory:
            print("Note that load_model() is skipped as no pretrained model is given")
            return

        names = self.get_model_names()

        # By default, the best model is loaded
        model_file = "model-best.pth.tar"

        if epoch is not None:
            model_file = "model.pth.tar-" + str(epoch)

        for name in names:
            model_path = osp.join(directory, name, model_file)

            if not osp.exists(model_path):
                print(self.cfg.DATASET)
                if self.cfg.DG:
                    model_path = ""
                elif self.cfg.DATASET.NAME in ['ImageNetV2', "ImageNetA", "ImageNetR", "ImageNetSketch"]:
                    model_path = ""
                else:
                    raise FileNotFoundError('Model not found at "{}"'.format(model_path))
                print("No data")
            checkpoint = load_checkpoint(model_path)
            state_dict = checkpoint["state_dict"]
            epoch = checkpoint["epoch"]

            # Ignore fixed token vectors
            if "token_prefix" in state_dict:
                del state_dict["token_prefix"]

            if "token_suffix" in state_dict:
                del state_dict["token_suffix"]

            print("Loading weights to {} " 'from "{}" (epoch = {})'.format(name, model_path, epoch))
            # set strict=False
            self._models[name].load_state_dict(state_dict, strict=False)
