import os.path as osp
import json

import torch
import torch.nn as nn
from torch.nn import functional as F
from torch.cuda.amp import GradScaler, autocast
from collections import OrderedDict

from dassl.engine import TRAINER_REGISTRY, TrainerX
from dassl.metrics import compute_accuracy
from dassl.utils import load_pretrained_weights, load_checkpoint
from dassl.optim import build_optimizer, build_lr_scheduler

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

_tokenizer = _Tokenizer()

DESC_LLM = "gpt-4.1"
DESC_TOPK = 4


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

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": "a photo of a {}, a type of texture.",
    "EuroSAT": "a centered satellite photo of {}.",
    #"EuroSAT": "a photo of a {}.",
    "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 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):
        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.shape = [batch_size, n_ctx, transformer.width]
        # take features from the eot embedding (eot_token is the highest number in each sequence)
        x = x[torch.arange(x.shape[0]), tokenized_prompts.argmax(dim=-1)] @ self.text_projection

        return x


class PromptLearner(nn.Module):
    def __init__(self, cfg, classnames, clip_model):
        super().__init__()
        n_cls = len(classnames)
        n_ctx = cfg.TRAINER.COOP.N_CTX
        ctx_init = cfg.TRAINER.COOP.CTX_INIT
        dtype = clip_model.dtype
        ctx_dim = clip_model.ln_final.weight.shape[0]
        clip_imsize = clip_model.visual.input_resolution
        cfg_imsize = cfg.INPUT.SIZE[0]
        n_prompts = cfg.TRAINER.COOP.N_PROMPTS
        assert cfg_imsize == clip_imsize, f"cfg_imsize ({cfg_imsize}) must equal to clip_imsize ({clip_imsize})"

        if ctx_init:
            # use given words to initialize context vectors
            temp = 'a photo of a'
            ctx_init = temp.replace("_", " ")
            n_ctx = len(ctx_init.split(" "))
            prompt = clip.tokenize(ctx_init)
            with torch.no_grad():
                embedding = clip_model.token_embedding(prompt).type(dtype)
            prompt_embeddings = embedding[0, 1:1 + n_ctx, :]
            ctx_vectors = torch.zeros(n_prompts, n_ctx, ctx_dim, dtype=dtype)
            ctx_vectors[:, n_ctx - n_ctx:, :] = prompt_embeddings.expand(n_prompts, -1, -1)
            # ctx_vectors = embedding[0, 1 : 1 + n_ctx, :]
            prompt_prefix = ctx_init

        else:
            # random initialization
            if cfg.TRAINER.COOP.CSC:
                print("Initializing class-specific contexts")
                ctx_vectors = torch.empty(n_prompts, n_cls, n_ctx, ctx_dim, dtype=dtype)
            else:
                print("Initializing a generic context")
                ctx_vectors = torch.empty(n_prompts, n_ctx, ctx_dim, dtype=dtype)
            nn.init.normal_(ctx_vectors, std=0.02)
            prompt_prefix = " ".join(["X"] * n_ctx)


        print(f'Initial context: "{prompt_prefix}"')
        print(f"Number of context words (tokens): {n_ctx}")

        self.ctx = nn.Parameter(ctx_vectors)  # to be optimized
        self.n_prompts = n_prompts

        bias_vectors = torch.empty(1, 512, dtype=dtype)
        nn.init.normal_(bias_vectors, std=0.02)
        self.bias_vectors = nn.Parameter(bias_vectors)

        classnames = [name.replace("_", " ") for name in classnames]
        name_lens = [len(_tokenizer.encode(name)) for name in classnames]
        prompts = [prompt_prefix + " " + name + "." for name in classnames]

        #print(f"Loading CLIP (backbone: {cfg.MODEL.BACKBONE.NAME})")
        clip_model_ = load_clip_to_cpu(cfg)
        clip_model_.cuda()
        
        #prompts_ = [prompt_prefix + " " + name + "." for name in classnames]        
        temp = CUSTOM_TEMPLATES[cfg.DATASET.NAME]
        prompts_ = [temp.format(c.replace("_", " ")) for c in classnames]
        print(f"Prompts: {prompts_}")
        prompts_ = torch.cat([clip.tokenize(p) for p in prompts_])
        prompts_ = prompts_.cuda()

        with torch.no_grad():
            text_features = clip_model_.encode_text(prompts_)
            text_features = text_features / text_features.norm(dim=-1, keepdim=True)

        text_features = []
        desc_file = f"./desc/{DESC_LLM}/descriptions_top{DESC_TOPK}/{cfg.DATASET.NAME}.json"
        with open(desc_file, "r") as f:
            all_desc = json.load(f)
        for cls in classnames:
            cls_descs = [f"a photo of {cls}, {desc}" for desc in all_desc[cls]]
            cls_token = torch.cat([clip.tokenize(cls_desc) for cls_desc in cls_descs]).cuda()
            with torch.no_grad():
                cls_feature = clip_model_.encode_text(cls_token)
                cls_feature = cls_feature / cls_feature.norm(dim=-1, keepdim=True)
                text_features.append(torch.mean(cls_feature, dim=0))

        self.text_features = torch.stack(text_features)

        self.meta_net = nn.Sequential(OrderedDict([
            ("linear1", nn.Linear(512, 512)),
            ("relu", nn.ReLU(inplace=True))
            #("linear2", nn.Linear(128, 512))
        ]))


        if cfg.TRAINER.COCOOP.PREC == "fp16":
            self.meta_net.half()


        tokenized_prompts = torch.cat([clip.tokenize(p) for p in prompts])
        with torch.no_grad():
            embedding = clip_model.token_embedding(tokenized_prompts).type(dtype)
        # These token vectors will be saved when in save_model(),
        # but they should be ignored in load_model() as we want to use
        # those computed using the current class names
        self.register_buffer("token_prefix", embedding[:, :1, :])  # SOS
        self.register_buffer("token_suffix", embedding[:, 1 + n_ctx :, :])  # CLS, EOS


        self.n_cls = n_cls
        self.n_ctx = n_ctx
        self.tokenized_prompts = tokenized_prompts  # torch.Tensor
        self.name_lens = name_lens
        self.class_token_position = cfg.TRAINER.COOP.CLASS_TOKEN_POSITION

    def forward(self):
        ctx = self.ctx

        if ctx.dim() == 3:
            ctx = ctx.unsqueeze(1).expand(-1, self.n_cls, -1, -1)
        
        prefix = self.token_prefix
        suffix = self.token_suffix

        prompts = torch.stack([
            torch.cat([prefix, ctx[i], suffix], dim=1) 
            for i in range(self.n_prompts)
        ])

        return prompts


class Adapter(nn.Module):
    def __init__(self, c_in, reduction=4):
        super(Adapter, self).__init__()
        self.fc = nn.Sequential(
            nn.Linear(c_in, c_in // reduction, bias=False),
            nn.ReLU(inplace=True),
            nn.Linear(c_in // reduction, c_in, bias=False),
            nn.ReLU(inplace=True)
        )

    def forward(self, x):
        x = self.fc(x)
        return x

class UncertaintyPromptIntegrator(nn.Module):
    def __init__(self, temperature=1.0):
        """
        基于预测熵的不确定性加权集成器
        
        Args:
            temperature: 控制权重分布的平滑度，值越大权重分布越平均
        """
        super().__init__()
        self.temperature = temperature
    
    def forward(self, all_logits):
        """
        Args:
            all_logits: [n_prompts, batch_size, n_classes]
        
        Returns:
            integrated_logits: [batch_size, n_classes]
            prompt_weights: [n_prompts, batch_size]
            entropy: [n_prompts, batch_size]
        """
        n_prompts, batch_size, n_classes = all_logits.shape
        
        log_probs = F.log_softmax(all_logits, dim=-1)
        probs = log_probs.exp()
        
        entropy = -(probs * log_probs).sum(dim=-1)
        
        temperature = max(self.temperature, 1e-8)
        weights = F.softmax(-entropy / temperature, dim=0)
        
        integrated_logits = torch.einsum('pb,pbc->bc', weights, all_logits)
        
        return integrated_logits, weights, entropy

class CustomCLIP(nn.Module):
    def __init__(self, cfg, classnames, clip_model):
        super().__init__()
        self.prompt_learner = PromptLearner(cfg, classnames, clip_model)
        self.tokenized_prompts = self.prompt_learner.tokenized_prompts
        self.n_prompts = self.prompt_learner.n_prompts
        self.ori_embedding = self.prompt_learner.text_features
        self.image_encoder = clip_model.visual
        self.text_encoder = TextEncoder(clip_model)
        self.logit_scale = clip_model.logit_scale
        self.dtype = clip_model.dtype
        self.meta_net = self.prompt_learner.meta_net
        self.adapter = Adapter(512, 4).to(clip_model.dtype)
        
        self.use_uncertainty_integration = cfg.TRAINER.COOP.get('UNC_ENABLED', False)
        self.unc_temperature = cfg.TRAINER.COOP.get('UNC_TEMPERATURE', 1.0)
        
        if self.use_uncertainty_integration:
            self.unc_integrator = UncertaintyPromptIntegrator(
                temperature=self.unc_temperature
            )

    def compute_diversity_loss(self, text_features):
        if self.n_prompts == 1:
            return torch.tensor(0.0, device=text_features.device)
        
        n_prompts = text_features.shape[0]
        n_cls = text_features.shape[1]
        
        total_orthogonal_loss = 0.0
        
        for c in range(n_cls):
            class_features = text_features[:, c, :]  # [n_prompts, feature_dim]
            class_features = F.normalize(class_features, p=2, dim=-1)
            sim_matrix = class_features @ class_features.T
            triu_indices = torch.triu_indices(n_prompts, n_prompts, offset=1)
            similarities = sim_matrix[triu_indices[0], triu_indices[1]]
            orthogonal_loss = (similarities ** 2).mean()
            total_orthogonal_loss += orthogonal_loss
        
        return total_orthogonal_loss / n_cls

    def forward(self, image):
        prompts = self.prompt_learner()
        image_features = self.image_encoder(image.type(self.dtype))

        tokenized_prompts = self.tokenized_prompts
        all_text_features = torch.stack([
            self.text_encoder(p, tokenized_prompts) 
            for p in prompts
        ])
        diversity_loss = self.compute_diversity_loss(all_text_features)

        text_features_old = self.ori_embedding
        image_features = image_features / image_features.norm(dim=-1, keepdim=True)

        mean_text_features = all_text_features.mean(dim=0)
        mean_text_features_norm = mean_text_features / mean_text_features.norm(dim=-1, keepdim=True)
        text_features_old = self.ori_embedding / self.ori_embedding.norm(dim=-1, keepdim=True)
        cos_sim = torch.nn.CosineSimilarity(dim=1, eps=1e-07)
        score = 1.0 - torch.mean(cos_sim(mean_text_features_norm, text_features_old))
        logit_scale = self.logit_scale.exp()
        all_logits = []
        for i in range(self.n_prompts):
            text_features_i = all_text_features[i, :, :]
            text_features_i = text_features_i / text_features_i.norm(dim=-1, keepdim=True)
            logits_i = logit_scale * image_features @ text_features_i.t()
            all_logits.append(logits_i)
        
        all_logits = torch.stack(all_logits)
        
        if self.use_uncertainty_integration:
            logits, prompt_weights, entropy = self.unc_integrator(all_logits)
            self.last_prompt_weights = prompt_weights.detach()
            self.last_entropy = entropy.detach()
        else:
            logits = all_logits.mean(dim=0)

        return logits, score, diversity_loss


@TRAINER_REGISTRY.register()
class MSGCoOp(TrainerX):

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

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

        print(f"Loading CLIP (backbone: {cfg.MODEL.BACKBONE.NAME})")
        clip_model = load_clip_to_cpu(cfg)
        
        if cfg.TRAINER.COOP.PREC == "fp32" or cfg.TRAINER.COOP.PREC == "amp":
            # CLIP's default precision is fp16
            clip_model.float()

        print("Building custom CLIP")
        self.model = CustomCLIP(cfg, classnames, clip_model)
        self.w = cfg.TRAINER.COOP.W
        self.diversity_weight = cfg.TRAINER.COOP.DIV_WEIGHT

        print("Turning off gradients in both the image and the text encoder")
        for name, param in self.model.named_parameters():
            if "ctx" not in name: 
                param.requires_grad_(False)
            else:
                print(name)

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

        self.model.to(self.device)
        # NOTE: only give prompt_learner to the optimizer
        self.optim = build_optimizer(self.model.prompt_learner, cfg.OPTIM)
        self.sched = build_lr_scheduler(self.optim, cfg.OPTIM)
        self.register_model("prompt_learner", self.model.prompt_learner, self.optim, self.sched)
        
        #self.optim_ = build_optimizer(self.model.adapter, cfg.OPTIM)
        #self.sched_ = build_lr_scheduler(self.optim, cfg.OPTIM)
        #self.register_model('clip_adapter', self.model.adapter, self.optim_, self.sched_)

        self.scaler = GradScaler() if cfg.TRAINER.COOP.PREC == "amp" else None

        # Note that multi-gpu training could be slow because CLIP's size is
        # big, which slows down the copy operation in DataParallel
        device_count = torch.cuda.device_count()
        if device_count > 1:
            print(f"Multiple GPUs detected (n_gpus={device_count}), use all of them!")
            self.model = nn.DataParallel(self.model)

    def forward_backward(self, batch):
        image, label = self.parse_batch_train(batch)
        prec = self.cfg.TRAINER.COOP.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:
            output, score, diversity_loss = self.model(image)
            loss = F.cross_entropy(output, label)+self.w*score + diversity_loss * self.diversity_weight
            self.model_backward_and_update(loss)

        loss_summary = {
            "loss": loss.item(),
            "acc": compute_accuracy(output, label)[0].item(),
        }

        if (self.batch_idx + 1) == self.num_batches:
            #self.update_lr()
            self.sched.step()
            #self.sched_.step()
        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 model_inference(self, input):
        return self.model(input)[0]


    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()
        print(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):
                raise FileNotFoundError('Model not found at "{}"'.format(model_path))

            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"]

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

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