import copy
import os.path as osp
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

from dassl.engine import TRAINER_REGISTRY, TrainerX
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
from .imagenet_templates import IMAGENET_TEMPLATES
from tqdm import tqdm
import math

from clip.model import VisionTransformer, convert_weights

_tokenizer = _Tokenizer()

class Feature_Trans_Module_two_layer(nn.Module):
    def __init__(self, input_dim=100, out_dim=256):
        super(Feature_Trans_Module_two_layer, self).__init__()

        self.conv1 = nn.Sequential(
            nn.Conv2d(input_dim, out_dim, 1),
            nn.BatchNorm2d(out_dim),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_dim, out_dim, 1)
        )
    def forward(self, input_feat):
        
        final_feat = self.conv1(input_feat.unsqueeze(-1).unsqueeze(-1))
        
        return final_feat.squeeze(-1).squeeze(-1)
        
def load_clip_to_cpu_teacher(cfg, zero_shot_model=False):
    backbone_name = cfg.TRAINER.PROMPTKD.TEACHER_NAME
    # url = clip._MODELS[backbone_name]
    
    if backbone_name == "ViT-B/16":
        model_path = './clip/ViT-B-16.pt'
    elif backbone_name == "ViT-L/14":
        model_path = './clip/ViT-L-14.pt'
    elif backbone_name == "ViT-B/32":
        model_path = './clip/ViT-B-32.pt'
    else:
        print('enter the wrong teacher name.')
    
    print(f"CLIP Teacher name is {backbone_name}")
    
    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")

    # We default use PromptSRC to pretrain our teacher model
    design_details = {"trainer": 'IVLP',
                        "vision_depth": 9,
                        "language_depth": 9,
                        "vision_ctx": 4,
                        "language_ctx": 4}
    
    model = clip.build_model(state_dict or model.state_dict(), design_details)

    return model


def load_clip_to_cpu(cfg, zero_shot_model=False):
    backbone_name = cfg.MODEL.BACKBONE.NAME
    # url = clip._MODELS[backbone_name]
    model_path = './clip/ViT-B-16.pt'
    
    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")

    design_details = {"trainer": 'IVLP',
                      "vision_depth": cfg.TRAINER.PROMPTKD.PROMPT_DEPTH_VISION,
                      "language_depth": cfg.TRAINER.PROMPTKD.PROMPT_DEPTH_TEXT,
                      "vision_ctx": cfg.TRAINER.PROMPTKD.N_CTX_VISION,
                      "language_ctx": cfg.TRAINER.PROMPTKD.N_CTX_TEXT}
    model = clip.build_model(state_dict or model.state_dict(), design_details)

    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):
        
        # print(f'------prompts size is {prompts.size()}------')
        # print(f'------tokenized prompts size is {tokenized_prompts.size()}------')

        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 VLPromptLearner(nn.Module):
    def __init__(self, cfg, classnames, clip_model, is_teacher):
        super().__init__()
        n_cls = len(classnames)
        # Make sure Language depth >= 1
        assert cfg.TRAINER.PROMPTKD.PROMPT_DEPTH_TEXT >= 1, "In Independent VL prompting, Language prompt depth should be >=1" \
                                                        "\nPlease use VPT trainer if you want to learn only vision " \
                                                        "branch"
        n_ctx = cfg.TRAINER.PROMPTKD.N_CTX_TEXT
        ctx_init = cfg.TRAINER.PROMPTKD.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]
        assert cfg_imsize == clip_imsize, f"cfg_imsize ({cfg_imsize}) must equal to clip_imsize ({clip_imsize})"
        
        self.trainer_name = cfg.TRAINER.NAME
        self.train_modal = cfg.TRAINER.MODAL
        
        if ctx_init and n_ctx <= 4:
            # use given words to initialize context vectors
            ctx_init = ctx_init.replace("_", " ")
            n_ctx = n_ctx
            prompt = clip.tokenize(ctx_init)
            with torch.no_grad():
                embedding = clip_model.token_embedding(prompt).type(dtype)
            ctx_vectors = embedding[0, 1: 1 + n_ctx, :]
            prompt_prefix = ctx_init
        else:
            # random initialization
            ctx_vectors = torch.empty(n_ctx, ctx_dim, dtype=dtype)
            nn.init.normal_(ctx_vectors, std=0.02)
            prompt_prefix = " ".join(["X"] * n_ctx)
        print(f"Independent V-L design")
        print(f'Initial text context: "{prompt_prefix}"')
        print(f"Number of context words (tokens) for Language prompting: {n_ctx}")
        print(f"Number of context words (tokens) for Vision prompting: {cfg.TRAINER.PROMPTKD.N_CTX_VISION}")
        self.ctx = nn.Parameter(ctx_vectors)

        classnames = [name.replace("_", " ") for name in classnames]
        prompts = [prompt_prefix + " " + name + "." for name in classnames]
        tokenized_prompts = torch.cat([clip.tokenize(p) for p in prompts])  # (n_cls, n_tkn)
        
        print(f'classnames size is {len(classnames)}')

        with torch.no_grad():
            embedding = clip_model.token_embedding(tokenized_prompts).type(dtype)
        
        self.n_cls = n_cls
        self.n_ctx = n_ctx
        self.tokenized_prompts = tokenized_prompts  # torch.Tensor
        # self.name_lens = name_lens

        if self.train_modal == "base2novel":
            self.register_buffer("token_prefix", embedding[:math.ceil(self.n_cls / 2), :1, :])  # SOS
            self.register_buffer("token_suffix", embedding[:math.ceil(self.n_cls / 2), 1 + n_ctx:, :])  # CLS, EOS

            self.register_buffer("token_prefix2", embedding[math.ceil(self.n_cls / 2):, :1, :])  # SOS
            self.register_buffer("token_suffix2", embedding[math.ceil(self.n_cls / 2):, 1 + n_ctx:, :])  # CLS, EOS
            
        elif self.train_modal == "cross":
            self.register_buffer("token_prefix", embedding[:, :1, :])  # SOS
            self.register_buffer("token_suffix", embedding[:, 1 + n_ctx:, :])  # CLS, EOS
            
            self.register_buffer("token_prefix2", embedding[:, :1, :])  # SOS
            self.register_buffer("token_suffix2", embedding[:, 1 + n_ctx:, :])  # CLS, EOS

    def construct_prompts(self, ctx, prefix, suffix, label=None):
        # dim0 is either batch_size (during training) or n_cls (during testing)
        # ctx: context tokens, with shape of (dim0, n_ctx, ctx_dim)
        # prefix: the sos token, with shape of (n_cls, 1, ctx_dim)
        # suffix: remaining tokens, with shape of (n_cls, *, ctx_dim)

        # print(f'label is {label}')
        # if label is not None:
        #     prefix = prefix[label]
        #     suffix = suffix[label]

        prompts = torch.cat(
            [
                prefix,  # (dim0, 1, dim)
                ctx,  # (dim0, n_ctx, dim)
                suffix,  # (dim0, *, dim)
            ],
            dim=1,
        )

        return prompts

    def forward(self):
        ctx = self.ctx
        if ctx.dim() == 2:
            ctx = ctx.unsqueeze(0).expand(self.n_cls, -1, -1)
        # print(f'ctx size is {ctx.size()}')

        prefix = self.token_prefix
        # print(f'prefix size is {prefix.size()}')
        
        suffix = self.token_suffix
        # print(f'suffix size is {suffix.size()}')

        if self.trainer_name in ["PromptKD", "PromptKD_LMC"] and self.train_modal == "base2novel":
            # print(f'n_cls is {self.n_cls}')
            prefix = torch.cat([prefix, self.token_prefix2], dim=0)
            suffix = torch.cat([suffix, self.token_suffix2], dim=0)

        prompts = self.construct_prompts(ctx, prefix, suffix)

        return prompts

class CustomCLIP(nn.Module):
    def __init__(self, cfg, classnames, clip_model):
        super().__init__()
        self.image_encoder = clip_model.visual
        
        self.logit_scale = clip_model.logit_scale
        self.dtype = clip_model.dtype
        self.total_epochs = cfg.OPTIM.MAX_EPOCH
        self.n_cls = len(classnames)
        
        self.VPT_image_trans = Feature_Trans_Module_two_layer(512, 768)
       
        self.cfg = cfg
        
        self.VPT_image_trans = self.VPT_image_trans.cuda()
        convert_weights(self.VPT_image_trans)

    def forward(self, image, label=None):
        logit_scale = self.logit_scale.exp()
        
        image_features = self.image_encoder(image.type(self.dtype))
        image_features = self.VPT_image_trans(image_features)
        image_features = image_features / image_features.norm(dim=-1, keepdim=True)

        return image_features, logit_scale


class CustomCLIP_teacher(nn.Module):
    def __init__(self, cfg, classnames, clip_model):
        super().__init__()
        self.prompt_learner = VLPromptLearner(cfg, classnames, clip_model, True)
        self.tokenized_prompts = self.prompt_learner.tokenized_prompts
        self.image_encoder = clip_model.visual
        self.text_encoder = TextEncoder(clip_model).cuda()
        self.logit_scale = clip_model.logit_scale
        self.dtype = clip_model.dtype
  
    def forward(self, image=None, label=None):
        
        prompts = self.prompt_learner()
        # Compute the prompted image and text features
        tokenized_prompts = self.tokenized_prompts
        text_features = self.text_encoder(prompts.cuda(), tokenized_prompts.cuda())
        text_features = text_features / text_features.norm(dim=-1, keepdim=True)
        
        logit_scale = self.logit_scale.exp()
        
        image_features = self.image_encoder(image.type(self.dtype))
        image_features = image_features / image_features.norm(dim=-1, keepdim=True)
        # Compute the prompted logits
        
        logits = logit_scale * image_features @ text_features.t()
        
        return image_features, text_features, logits


@TRAINER_REGISTRY.register()
class PromptKD_LMC(TrainerX):
    def check_cfg(self, cfg):
        assert cfg.TRAINER.PROMPTKD.PREC in ["fp16", "fp32", "amp"]

    def build_model(self):
        cfg = self.cfg
        
        classnames = self.dm.dataset.classnames
        self.n_cls = len(classnames)
        
        print(f"Loading CLIP (backbone: {cfg.MODEL.BACKBONE.NAME})")
        clip_model = load_clip_to_cpu(cfg)

        clip_model_teacher = load_clip_to_cpu_teacher(cfg)

        if cfg.TRAINER.PROMPTKD.PREC == "fp32" or cfg.TRAINER.PROMPTKD.PREC == "amp":
            # CLIP's default precision is fp16
            clip_model.float()

        print("Building custom CLIP")
        self.model = CustomCLIP(cfg, classnames, clip_model)

        self.model_teacher = CustomCLIP_teacher(cfg, classnames, clip_model_teacher)
        
        if cfg.TRAINER.MODAL == "base2novel":
            model_path = './teacher_model/'+str(cfg.DATASET.NAME)+'/VLPromptLearner/model-best.pth.tar'
        elif cfg.TRAINER.MODAL == "cross":
            model_path = './teacher_model/ImageNet-xd/VLPromptLearner_large/model.pth.tar-20'
            
        self.train_modal = cfg.TRAINER.MODAL
        
        checkpoint = load_checkpoint(model_path)
        state_dict = checkpoint["state_dict"]
        
        if "prompt_learner.token_prefix" in state_dict:
            del state_dict["prompt_learner.token_prefix"]
        if "prompt_learner.token_prefix2" in state_dict:
            del state_dict["prompt_learner.token_prefix2"]

        if "prompt_learner.token_suffix" in state_dict:
            del state_dict["prompt_learner.token_suffix"]
        if "prompt_learner.token_suffix2" in state_dict:
            del state_dict["prompt_learner.token_suffix2"]
        
        self.model_teacher.load_state_dict(state_dict, strict=False)
        self.model_teacher.to(self.device)
        self.model_teacher.eval()
        
        print("Turning off gradients in both the image and the text encoder")
        name_to_update = "prompt_learner"

        for name, param in self.model.named_parameters():
            if name_to_update not in name:
                # Make sure that VPT prompts are updated
                if "VPT" in name:
                    param.requires_grad_(True)
                else:
                    param.requires_grad_(False)
            else:
                if "ZS_image_encoder" 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}")
        print(f"Parameters count: {len(enabled)}")
        if cfg.MODEL.INIT_WEIGHTS:
            load_pretrained_weights(self.model, cfg.MODEL.INIT_WEIGHTS)

        self.model.to(self.device)
        # NOTE: only give prompt_learner to the optimizer

        self.trainable_list = nn.ModuleList([])
        self.trainable_list.append(self.model)

        self.optim = build_optimizer(self.trainable_list, cfg.OPTIM)
        self.sched = build_lr_scheduler(self.optim, cfg.OPTIM)
        self.register_model("VLPromptLearner", self.model, self.optim, self.sched)
        
        # Cosine scheduler
        self.total_epochs = cfg.OPTIM.MAX_EPOCH
        self.step_counter = 1
        N = cfg.OPTIM.MAX_EPOCH
        
        self.scaler = GradScaler() if cfg.TRAINER.PROMPTKD.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)

        self.temperature = cfg.TRAINER.PROMPTKD.TEMPERATURE
        
        # Initialize pretrained student model for LMC regularization
        self.model_pretrained_student = None
        if hasattr(cfg.TRAINER.PROMPTKD, 'USE_LMC') and cfg.TRAINER.PROMPTKD.USE_LMC:
            # Try to load pretrained student model if path is provided
            pretrained_path = getattr(cfg.TRAINER.PROMPTKD, 'PRETRAINED_STUDENT_PATH', None)
            if pretrained_path and self.load_pretrained_student_model(pretrained_path):
                print("Loaded pretrained student model for LMC regularization")
                
                # IMPORTANT: Also initialize the current training model with the same pretrained weights
                print("Initializing current training model with pretrained weights...")
                pretrained_state_dict = {k: v.clone() for k, v in self.model_pretrained_student.state_dict().items()}
                self.model.load_state_dict(pretrained_state_dict, strict=False)
                print("Current training model initialized from pretrained student model")
            else:
                # Use current model as pretrained baseline
                self.model_pretrained_student = copy.deepcopy(self.model)
                # Freeze pretrained student parameters
                for param in self.model_pretrained_student.parameters():
                    param.requires_grad = False
                print("Initialized pretrained student model (copy of current) for LMC regularization")
    
    def calculate_l2_regularization(self):
        """Calculate L2 regularization between current student and pretrained student weights"""
        if self.model_pretrained_student is None:
            return torch.tensor(0.0, device=self.device)
        
        l2_loss = 0.0
        total_elements = 0
        
        for (name1, param1), (name2, param2) in zip(
            self.model.named_parameters(), 
            self.model_pretrained_student.named_parameters()
        ):
            if param1.requires_grad and name1 == name2:
                # Calculate mean squared difference for this parameter
                param_diff = param1 - param2
                l2_loss += torch.sum(param_diff ** 2)
                total_elements += param_diff.numel()
        
        # Return mean squared difference (normalized by total number of parameters)
        if total_elements > 0:
            return l2_loss / total_elements
        else:
            return torch.tensor(0.0, device=self.device)
    
    def calculate_lmc_connectivity_loss(self, image, tea_text_features):
        """Calculate LMC connectivity loss by interpolating between student, pretrained student, and teacher"""
        if self.model_pretrained_student is None:
            return torch.tensor(0.0, device=self.device)
        
        # Get configurations
        num_samples = getattr(self.cfg.TRAINER.PROMPTKD, 'NUM_SAMPLES', 10)
        line_samples = np.arange(0.1, 1.01, 1.0 / float(num_samples))
        
        # Get features from different models
        with torch.no_grad():
            pretrained_student_features, pretrained_logit_scale = self.model_pretrained_student(image)
            pretrained_student_logits = pretrained_logit_scale * pretrained_student_features @ tea_text_features.t()
        
        current_student_features, current_logit_scale = self.model(image)
        current_student_logits = current_logit_scale * current_student_features @ tea_text_features.t()
        
        total_lmc_loss = 0.0
        
        # Interpolate between pretrained student and current student
        for alpha in line_samples:
            # Linear interpolation in feature space
            interpolated_features = (1 - alpha) * pretrained_student_features + alpha * current_student_features
            interpolated_logit_scale = (1 - alpha) * pretrained_logit_scale + alpha * current_logit_scale
            interpolated_logits = interpolated_logit_scale * interpolated_features @ tea_text_features.t()
            
            # Use teacher's soft targets for connectivity loss
            with torch.no_grad():
                _, _, tea_logits = self.model_teacher(image)
            
            # Calculate connectivity loss at this interpolation point
            lmc_loss_point = F.kl_div(
                F.log_softmax(interpolated_logits / self.temperature, dim=1),
                F.softmax(tea_logits / self.temperature, dim=1),
                reduction='sum'
            ) * (self.temperature * self.temperature) / interpolated_logits.numel()
            
            total_lmc_loss += lmc_loss_point / len(line_samples)
        
        return total_lmc_loss

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

        with torch.no_grad():
            tea_image_features, tea_text_features, tea_logits = self.model_teacher(image)

        model = self.model
        optim = self.optim
        scaler = self.scaler
        
        prec = self.cfg.TRAINER.PROMPTKD.PREC
        if prec == "amp":
            with autocast():
                image_ft, logit_scale = model(image, label)
                stu_logits = logit_scale * image_ft @ tea_text_features.t().detach()
                
                # Knowledge Distillation Loss
                L_ukd = F.kl_div(
                    F.log_softmax(stu_logits / self.temperature, dim=1),
                    F.softmax(tea_logits / self.temperature, dim=1),
                    reduction='sum',
                ) * (self.temperature * self.temperature) / stu_logits.numel()
                
                # Build final loss with LMC regularization
                loss = self.cfg.TRAINER.PROMPTKD.KD_WEIGHT * L_ukd
                
                # Add L2 regularization if enabled
                l2_reg_loss = torch.tensor(0.0, device=self.device)
                lmc_connectivity_loss = torch.tensor(0.0, device=self.device)
                
                if hasattr(self.cfg.TRAINER.PROMPTKD, 'USE_LMC') and self.cfg.TRAINER.PROMPTKD.USE_LMC:
                    l2_reg_loss = self.calculate_l2_regularization()
                    lmc_connectivity_loss = self.calculate_lmc_connectivity_loss(image, tea_text_features)
                    
                    lambda_reg = getattr(self.cfg.TRAINER.PROMPTKD, 'LAMBDA_REG', 0.1)
                    beta_lmc = getattr(self.cfg.TRAINER.PROMPTKD, 'BETA_LMC', 0.5)
                    
                    loss = loss + lambda_reg * l2_reg_loss + beta_lmc * lmc_connectivity_loss
                
            optim.zero_grad()
            scaler.scale(loss).backward()
            scaler.step(optim)
            scaler.update()
        else:
            image_ft, logit_scale = model(image, label)
            
            stu_logits = logit_scale * image_ft @ tea_text_features.t().detach()
            
            # Knowledge Distillation Loss
            L_ukd = F.kl_div(
                F.log_softmax(stu_logits / self.temperature, dim=1),
                F.softmax(tea_logits / self.temperature, dim=1),
                reduction='sum',
            ) * (self.temperature * self.temperature) / stu_logits.numel()
            
            # Build final loss with LMC regularization
            loss = self.cfg.TRAINER.PROMPTKD.KD_WEIGHT * L_ukd
            
            # Initialize components for loss tracking
            l2_reg_loss = torch.tensor(0.0, device=self.device)
            lmc_connectivity_loss = torch.tensor(0.0, device=self.device)
            
            # Add LMC regularization if enabled
            if hasattr(self.cfg.TRAINER.PROMPTKD, 'USE_LMC') and self.cfg.TRAINER.PROMPTKD.USE_LMC:
                l2_reg_loss = self.calculate_l2_regularization()
                lmc_connectivity_loss = self.calculate_lmc_connectivity_loss(image, tea_text_features)
                
                lambda_reg = getattr(self.cfg.TRAINER.PROMPTKD, 'LAMBDA_REG', 0.1)
                beta_lmc = getattr(self.cfg.TRAINER.PROMPTKD, 'BETA_LMC', 0.5)
                
                loss = loss + lambda_reg * l2_reg_loss + beta_lmc * lmc_connectivity_loss
            
            optim.zero_grad()
            loss.backward()
            optim.step()

        # Get the weight parameters for logging
        kd_weight = self.cfg.TRAINER.PROMPTKD.KD_WEIGHT
        lambda_reg = getattr(self.cfg.TRAINER.PROMPTKD, 'LAMBDA_REG', 0.1)
        beta_lmc = getattr(self.cfg.TRAINER.PROMPTKD, 'BETA_LMC', 0.5)
        
        loss_summary = {
            "loss": loss.item(),
            "loss_kd": L_ukd.item(),
            "loss_l2_reg": l2_reg_loss.item(), 
            "loss_lmc_connectivity": lmc_connectivity_loss.item(),
            "weighted_kd": (kd_weight * L_ukd).item(),
            "weighted_l2_reg": (lambda_reg * l2_reg_loss).item(),
            "weighted_lmc": (beta_lmc * lmc_connectivity_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):
                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 "prompt_learner.token_prefix" in state_dict:
                del state_dict["prompt_learner.token_prefix"]
            if "prompt_learner.token_prefix2" in state_dict:
                del state_dict["prompt_learner.token_prefix2"]
                
            if "prompt_learner.token_suffix" in state_dict:
                del state_dict["prompt_learner.token_suffix"]
            if "prompt_learner.token_suffix2" in state_dict:
                del state_dict["prompt_learner.token_suffix2"]
                
            print("Loading weights to {} " 'from "{}" (epoch = {})'.format(name, model_path, epoch))
            # set strict=False
            self._models[name].load_state_dict(state_dict, strict=False)

    @torch.no_grad()
    def test(self, split=None):
        """A generic testing pipeline."""
        self.set_model_mode("eval")
        self.evaluator.reset()

        if split is None:
            split = self.cfg.TEST.SPLIT

        if split == "val" and self.val_loader is not None:
            data_loader = self.val_loader
        elif split == "train":
            data_loader = self.train_loader
        else:
            split = "test"  # in case val_loader is None
            data_loader = self.test_loader

        print(f"Evaluate on the *{split}* set")
        
        for batch_idx, batch in enumerate(tqdm(data_loader)):
            image, label = self.parse_batch_test(batch)
            
            with torch.no_grad():
                tea_image_features, tea_text_features, tea_logits = self.model_teacher(image, label)
                
            image_ft, logit_scale = self.model(image, label)
            
            if self.train_modal == "base2novel":
                if split == "val":
                    output = logit_scale * image_ft @ tea_text_features[:math.ceil(self.n_cls / 2),:].t()
                elif split == "test":
                    output = logit_scale * image_ft @ tea_text_features[math.ceil(self.n_cls / 2):,:].t()
            elif self.train_modal == "cross" :
                output = logit_scale * image_ft @ tea_text_features.t()
            
            self.evaluator.process(output, label) 

        results = self.evaluator.evaluate()

        for k, v in results.items():
            tag = f"{split}/{k}"
            self.write_scalar(tag, v, self.epoch)

        return list(results.values())[0]
    
    def save_pretrained_student_model(self, directory):
        """Save the current student model as pretrained model for future LMC regularization"""
        if not directory:
            return
        
        import os
        os.makedirs(directory, exist_ok=True)
        model_path = os.path.join(directory, "pretrained_student_model.pth")
        
        # Save only the state dict
        torch.save(self.model.state_dict(), model_path)
        print(f"Pretrained student model saved to {model_path}")
    
    def load_pretrained_student_model(self, directory):
        """Load pretrained student model for LMC regularization"""
        if not directory:
            return False
        
        import os
        # Look for the standard PromptKD model format
        model_path = os.path.join(directory, "VLPromptLearner", "model-best.pth.tar")
        
        if not os.path.exists(model_path):
            print(f"Pretrained student model not found at {model_path}")
            return False
        
        try:
            # Create a copy of current model structure
            self.model_pretrained_student = copy.deepcopy(self.model)
            
            # Load the pretrained checkpoint (same format as load_model)
            checkpoint = load_checkpoint(model_path)
            state_dict = checkpoint["state_dict"]
            epoch = checkpoint["epoch"]
            
            # Remove fixed token vectors (same as load_model)
            if "prompt_learner.token_prefix" in state_dict:
                del state_dict["prompt_learner.token_prefix"]
            if "prompt_learner.token_prefix2" in state_dict:
                del state_dict["prompt_learner.token_prefix2"]
            if "prompt_learner.token_suffix" in state_dict:
                del state_dict["prompt_learner.token_suffix"]
            if "prompt_learner.token_suffix2" in state_dict:
                del state_dict["prompt_learner.token_suffix2"]
            
            # Load the state dict
            self.model_pretrained_student.load_state_dict(state_dict, strict=False)
            
            # Freeze parameters
            for param in self.model_pretrained_student.parameters():
                param.requires_grad = False
            
            print(f"Pretrained student model loaded from {model_path} (epoch = {epoch})")
            return True
        except Exception as e:
            print(f"Failed to load pretrained student model: {e}")
            return False

