"""
UAD-CMPT: Unified Face Attack Detection via Cross-Modal Prompt Tuning (ICLR-26)
"""
import os.path as osp
import torch, copy
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
from dassl.utils import load_pretrained_weights, load_checkpoint
from dassl.optim import build_optimizer, build_lr_scheduler
from dassl.data.datasets import build_dataset
from clip import clip
from clip.simple_tokenizer import SimpleTokenizer as _Tokenizer
from util.utils_FAS import cross_entropy
from util.utils_FAS import calc_mean_std, mean_variance_norm
_tokenizer = _Tokenizer()
from trainers.trainer_fas import TrainerX



class SynonymExplorer(nn.Module):
    """
    Learnable synonym explorer in embedding space (using fixed vocab_head).
    - Input: class word embedding e (B x D)
    - Output: augmented embedding (B x D), plus optional synonym scores
    """

    def __init__(self, emb_dim, vocab_head, hidden_dim=256, topk=10, normalize=False):
        """
        Args:
            emb_dim: dimension of class embedding
            vocab_head: frozen vocab head (nn.Linear with weight = vocab_emb)
            hidden_dim: dimension of intermediate MLP
            topk: number of nearest neighbors to consider
            normalize: whether to normalize embeddings before similarity
        """
        super().__init__()
        self.emb_dim = emb_dim
        self.vocab_head = vocab_head  # Linear
        self.topk = topk
        self.normalize = normalize

        # learnable transformation f_theta
        self.mlp = nn.Sequential(
            nn.Linear(emb_dim, hidden_dim),
            nn.ReLU(inplace=True),
            nn.Linear(hidden_dim, emb_dim)
        )

        if normalize:
            self.register_buffer("vocab_emb_norm", F.normalize(self.vocab_head.weight.data, dim=-1))
        else:
            self.register_buffer("vocab_emb_norm", self.vocab_head.weight.data)

    def forward(self, e):
        """
        e: (B x D) class embedding
        return:
            aug_e: (B x D) augmented embedding with synonyms
            scores: (B x topk) softmax weights for chosen synonyms
            neighbors: (B x topk x D) synonym embeddings
        """
        B, D = e.shape

        # 1. transform class embedding into query
        q = self.mlp(e)  # (B x D)

        if self.normalize:
            q = F.normalize(q, dim=-1)
        vocab = self.vocab_emb_norm

        # 2. similarity to all vocab embeddings
        sim = torch.matmul(q, vocab.t())  # (B x V)

        # 3. get top-k nearest neighbors
        topk_sim, topk_idx = torch.topk(sim, self.topk, dim=-1)  # (B x k), (B x k)

        # 4. fetch embeddings of topk words
        neighbors = vocab[topk_idx]  # (B x k x D)

        # 5. compute soft weights over top-k neighbors
        scores = F.softmax(topk_sim, dim=-1)  # (B x k)

        # 6. weighted average to form synonym-enhanced embedding
        aug_e = e + torch.sum(neighbors * scores.unsqueeze(-1), dim=1)  # (B x D)

        return aug_e, scores, neighbors


class HighFrequencyExtractor(nn.Module):
    def __init__(self):
        super(HighFrequencyExtractor, self).__init__()
        self.conv1 = nn.Sequential(
            nn.Conv2d(768, 768, kernel_size=1, stride=1, bias=False),
            nn.ReLU(inplace=True),
        )
        self.realconv1 = nn.Conv2d(768, 768, kernel_size=1, stride=1, bias=False)
        self.imagconv1 = nn.Conv2d(768, 768, kernel_size=1, stride=1, bias=False)
        self.conv2 = nn.Conv2d(768, 768, kernel_size=1, stride=1, bias=False)

    def highFreqWH(self, x, scale):
        input_type = x.dtype
        x = x.to(torch.float32)
        x = torch.fft.fft2(x, norm="ortho")
        x = torch.fft.fftshift(x, dim=[-2, -1])
        b, c, h, w = x.shape

        x[:, :, h // 2 - int(h // 2 * scale):h // 2 + int(h // 2 * scale),
        w // 2 - int(w // 2 * scale):w // 2 + int(w // 2 * scale)] = 0.0

        x = torch.fft.ifftshift(x, dim=[-2, -1])
        x = torch.fft.ifft2(x, norm="ortho")
        x = torch.real(x)
        x = F.relu(x, inplace=True)
        return x.to(input_type)

    def forward(self, x):
        x = self.highFreqWH(x, 0.25)
        x = self.conv1(x)
        x = x.to(torch.float32)

        x = torch.fft.fft2(x, norm="ortho")
        x = torch.fft.fftshift(x, dim=[-2, -1]).to(torch.float32)
        x = x.to(torch.complex64)

        real_part = self.realconv1(x.real).to(torch.float32)
        phase_part = self.imagconv1(x.imag).to(torch.float32)
        x = torch.complex(real_part, phase_part)

        x = torch.fft.ifftshift(x, dim=[-2, -1])
        x = torch.fft.ifft2(x, norm="ortho")

        x = torch.real(x)
        x = F.relu(x, inplace=True)

        x = self.conv2(x)
        # x = F.relu(x, inplace=True)
        return x


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")
    design_details = {"trainer": 'CMPT',
                      "vision_depth": 0, "language_depth": 0,
                      "vision_ctx": 0, "language_ctx": 0,
                      "cmpt_da_length": cfg.TRAINER.CMPT.N_DACTX,
                      "cmpt_ds_length": cfg.TRAINER.CMPT.N_DSTKN
                      }
    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, compound_prompts_text):
        x = prompts + self.positional_embedding.type(self.dtype)
        x = x.permute(1, 0, 2)  # NLD -> LND
        # Pass as the list, as nn.sequential cannot process multiple arguments in the forward pass
        combined = [x, compound_prompts_text, 0]  # third argument is the counter which denotes depth of prompt
        outputs = self.transformer(combined)
        x = outputs[0]  # extract the x back from here
        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 MultiModalPromptLearner(nn.Module):
    def __init__(self, cfg, device, classnames, clip_model):
        super().__init__()
        self.device = device
        n_cls = len(classnames)
        n_dactx = cfg.TRAINER.CMPT.N_DACTX
        n_dstkn = cfg.TRAINER.CMPT.N_DSTKN
        ctx_init = cfg.TRAINER.CMPT.CTX_INIT
        L = cfg.TRAINER.CMPT.L
        self.W = cfg.TRAINER.CMPT.W

        dtype = clip_model.dtype
        ctx_dim = clip_model.ln_final.weight.shape[0]
        vis_dim = clip_model.visual.ln_post.weight.shape[0]  ## vision_width
        clip_imsize = clip_model.visual.input_resolution
        cfg_imsize = cfg.INPUT.SIZE[0]
        self.vocab_size = clip_model.vocab_size  # 49408

        # Default is 1, which is compound shallow prompting
        assert cfg.TRAINER.CMPT.PROMPT_DEPTH >= 1, "For CMPT, PROMPT_DEPTH should be >= 1"
        self.compound_prompts_depth = cfg.TRAINER.CMPT.PROMPT_DEPTH  # max=12, but will create 11 such shared prompts
        assert cfg_imsize == clip_imsize, f"cfg_imsize ({cfg_imsize}) must equal to clip_imsize ({clip_imsize})"

        if ctx_init and (n_dactx) <= 4:
            # use given words to initialize context vectors
            ctx_init = ctx_init.replace("_", " ")
            prompt = clip.tokenize(ctx_init)
            with torch.no_grad():
                embedding = clip_model.token_embedding(prompt).type(dtype)
            ctx_vectors = embedding[0, 1: 1 + n_dactx, :]
            prompt_prefix = ctx_init
        else:
            # random initialization
            ctx_vectors = torch.empty(n_dactx, ctx_dim, dtype=dtype)
            nn.init.normal_(ctx_vectors, std=0.02)
            prompt_prefix = " ".join(["X"] * n_dactx)
        self.ctx = nn.Parameter(ctx_vectors)
        print('CMPT design: Cross-modal Prompt Tuning')
        print(f'Initial context: "{prompt_prefix}"')
        print(f"Number of CMPT context words (tokens): {n_dactx}")
        # Vocab Head
        self.vocab_head = nn.Linear(ctx_dim, self.vocab_size, bias=False)
        self.vocab_head.weight.data = clip_model.token_embedding.weight.data.clone()
        self.vocab_head.weight.requires_grad_(False)

        # Domain-specific prompts (DSPrompts) in vision branch. Domain-agnostic prompts (DAPrompts) in language branch.
        self.one_livefakes, self.map_freqs, self.scores, self.neighbors, self.synonm ={}, {}, {}, {}, {}
        # Define parameters
        bases_style = torch.empty(1, vis_dim, 14, 14, dtype=dtype)  # [1, vis_dim, 14, 14]
        nn.init.normal_(bases_style, std=0.02)
        self.bases_style = nn.Parameter(bases_style)

        # Use for Getting prompted parameters.
        self.class_explorer = SynonymExplorer(emb_dim=ctx_dim, vocab_head=self.vocab_head, hidden_dim=256, topk=L)
        # self.compound_prompt_class_explorer = _get_clones(self.class_explorer, self.compound_prompts_depth)
        self.freq_extractor = HighFrequencyExtractor()
        # self.compound_prompt_freq_extractor = _get_clones(self.freq_extractor, self.compound_prompts_depth)

        # Also make corresponding projection layers from vision to text (and text to vision) branch for each prompt.
        self.proj_v2t = nn.Linear(vis_dim, ctx_dim)
        single_layer = nn.Linear(vis_dim, ctx_dim)
        self.compound_prompt_v2tProjections = _get_clones(single_layer, self.compound_prompts_depth - 1)
        self.proj_t2v = nn.Linear(ctx_dim, vis_dim)
        single_layer = nn.Linear(ctx_dim, vis_dim)
        self.compound_prompt_t2vProjections = _get_clones(single_layer, self.compound_prompts_depth - 1)

        # These below parameters related to the text prompts. Define the compound prompts for the deeper layers
        # Minimum can be 1, which defaults to shallow CMPT compound test prompts. 11: [2, 512]
        self.compound_prompts_text = nn.ParameterList(
            [nn.Parameter(torch.empty(n_dactx, ctx_dim)) for _ in range(self.compound_prompts_depth - 1)])
        for single_para in self.compound_prompts_text:
            nn.init.normal_(single_para, std=0.02)

        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)
    
        with torch.no_grad():
            embedding = clip_model.token_embedding(tokenized_prompts).type(dtype)  # (n_cls, n_tkn, ctx_dim)
        # 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_dactx:, :])  # CLS, EOS

        self.n_cls = n_cls
        self.vis_dim = vis_dim
        self.ctx_dim = ctx_dim
        self.n_dactx = n_dactx
        self.n_dstkn = n_dstkn
        self.tokenized_prompts = tokenized_prompts  # torch.Tensor
        self.embedding = embedding

    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)
        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, patch_embedding=None, word_embedding=None, map_freqs=None, one_livefakes=None, counter=-1):
        if patch_embedding is not None:
            # DSPrompts
            # [normalized_c, normalized_cs], [mu_s, sig_s, mu_cs, sig_cs], [O_ss, style], \
            #     ds_prompt_vis = self.attn_style(patch_embedding, self.bases_style)
            # self.contents[counter + 1] = [normalized_c, normalized_cs]
            # self.styles[counter + 1] = [mu_s, sig_s, mu_cs, sig_cs]
            # self.identities[counter + 1] = [O_ss, style]

            [b, n, c] = patch_embedding.shape
            grid_size = int(n**(0.5))
            patch_embedding = patch_embedding.reshape(-1, grid_size, grid_size, c)
            patch_embedding = patch_embedding.permute(0, 3, 1, 2)
            # patch_embedding = self.compound_prompt_freq_extractor[counter+1](patch_embedding)
            patch_embedding = self.freq_extractor(patch_embedding)
            ds_prompt_vis = patch_embedding.reshape(-1, c, n).permute(0, 2, 1)
            self.map_freqs[counter + 1] = ds_prompt_vis

        if word_embedding is not None:
            # DAPrompts
            # da_prompt_ctx, scores, neighbors = self.compound_prompt_class_explorer[counter+1](word_embedding)  # [2, ctx_dim]
            da_prompt_ctx, scores, neighbors = self.class_explorer(word_embedding)  # [2, ctx_dim]
            self.one_livefakes[counter + 1] = da_prompt_ctx
            self.scores[counter + 1] = scores
            self.neighbors[counter + 1] = neighbors
            self.synonm[counter + 1] = [da_prompt_ctx, word_embedding]

        if counter == -1:  # the shallow layer
            prefix = self.token_prefix
            suffix = self.token_suffix
            # ctx_shifted = self.proj_v2t(ds_prompt_vis).expand(self.n_cls, -1, -1)
            # prompts_text = self.construct_prompts(ctx_shifted, prefix, suffix)   # (n_cls, n_tkn, ctx_dim)
            # prompts_vision = self.proj_t2v(da_prompt_ctx)
            ctx_text = self.ctx
            bias = torch.sum(self.proj_v2t(ds_prompt_vis), dim=1)                # (1, ctx_dim)
            ctx_shifted = ctx_text + bias
            ctx_shifted = ctx_shifted.unsqueeze(0).expand(self.n_cls, -1, -1)    # (n_cls, n_ctx, ctx_dim)
            prompts_text = self.construct_prompts(ctx_shifted, prefix, suffix)   # (n_cls, n_tkn, ctx_dim)
            prompts_vision = self.proj_t2v(da_prompt_ctx)
            return prompts_text, prompts_vision  # pass here original, as for visual 768 is required
        
        else:  # the deeper layers
            if map_freqs is not None:
                if len(map_freqs) > 1:
                    # ds_prompt_vis = map_musigs[counter].float()
                    # prompts_text = self.compound_prompt_v2tProjections[counter](ds_prompt_vis).expand(self.n_cls, -1,-1)
                    ds_prompt_vis = torch.sum(map_freqs[counter].float(), dim=1)               # [1, 2, 768] - [1, 768]
                    ctx_text = self.compound_prompts_text[counter]
                    bias = self.compound_prompt_v2tProjections[counter](ds_prompt_vis)  # (1, ctx_dim)
                    ctx_shifted = ctx_text + bias                                       # (n_ctx, ctx_dim)
                    # print(ctx_text, ctx_text.shape)
                    # print(bias, bias.shape)
                    # input()

                    prompts_text = ctx_shifted.unsqueeze(0).expand(self.n_cls, -1, -1)  # (n_cls, n_ctx, ctx_dim)
                else:
                    prompts_text = self.compound_prompts_text[counter].unsqueeze(0).expand(self.n_cls, -1, -1)
                return prompts_text

            if one_livefakes is not None:
                da_prompt_ctx = one_livefakes[counter]
                prompts_vision = self.compound_prompt_t2vProjections[counter](da_prompt_ctx)
                return prompts_vision


class CustomCLIP(nn.Module):
    def __init__(self, cfg, device, classnames, clip_model):
        super().__init__()
        self.prompt_learner = MultiModalPromptLearner(cfg, device, classnames, clip_model)
        self.tokenized_prompts = self.prompt_learner.tokenized_prompts
        self.embedding = self.prompt_learner.embedding

        self.n_dactx = self.prompt_learner.n_dactx
        self.n_dstkn = self.prompt_learner.n_dstkn
        self.vis_dim = self.prompt_learner.vis_dim
        self.W = self.prompt_learner.W
        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.device = device
        self.clip_model = clip_model
        self.cos = torch.nn.CosineSimilarity(dim=1, eps=1e-07)

    def total_loss(self, scores, neighbors, synonm):
        loss_su, loss_nd = 0, 0
        for counter in scores:
            [aug_e, e] = synonm[counter]
            # loss_se += (scores[counter] * scores[counter].log()).sum(dim=-1).mean()
            uniform = torch.full_like(scores[counter], 1.0 / scores[counter].size(-1))
            loss_su += F.kl_div(scores[counter].log(), uniform, reduction="batchmean")

            sim_matrix = torch.matmul(neighbors[counter], neighbors[counter].transpose(-1, -2))  # (B, k, k)
            loss_nd += sim_matrix.mean()
        
        loss_su /= (counter + 1)
        loss_nd /= (counter + 1)
        return loss_su*0.01 + loss_nd*0.01

    def forward(self, image, label=None):
        tokenized_prompts = self.tokenized_prompts
        logit_scale = self.logit_scale.exp()

        with torch.no_grad():
            # patch
            patch_embedding = self.image_encoder.conv1(image.type(self.dtype))
            patch_embedding = patch_embedding.reshape(patch_embedding.shape[0], patch_embedding.shape[1], -1)
            patch_embedding = patch_embedding.permute(0, 2, 1)
            patch_embedding = patch_embedding + self.image_encoder.positional_embedding[1:, :].to(patch_embedding.dtype)
        # text
        word_embedding = \
            self.embedding[torch.arange(self.embedding.shape[0]), self.tokenized_prompts.argmax(dim=-1)].to(self.device)

        # shallow
        prompts_text, prompts_vision = self.prompt_learner(patch_embedding, word_embedding)

        with torch.no_grad():
            self.text_encoder(prompts_text, tokenized_prompts, self.prompt_learner)  # fill one_livefakes for vision

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

        text_features = self.text_encoder(prompts_text, tokenized_prompts, self.prompt_learner)
        text_features = text_features / text_features.norm(dim=-1, keepdim=True)
        logits = logit_scale * image_features @ text_features.t()

        if self.prompt_learner.training:
            loss_constraint = self.total_loss(
                self.prompt_learner.scores,
                self.prompt_learner.neighbors,
                self.prompt_learner.synonm)
            return logits, F.cross_entropy(logits, label) + loss_constraint
        return logits

def _get_clones(module, N):
    return nn.ModuleList([copy.deepcopy(module) for i in range(N)])

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

    def build_data_loader(self):
        """Create essential data-related attributes.
        A re-implementation of this method must create the
        same attributes (self.dm is optional).
        """
        dataset = build_dataset(self.cfg)
        self.train_loader_x = dataset.train_loader
        self.val_loader = dataset.dev_loader
        self.test_loader = dataset.test_loader
        self.lab2cname = dataset.lab2cname
        self.classnames = dataset.classnames

    def build_model(self):
        cfg = self.cfg
        self.device = torch.device('cuda:%d' % cfg.TRAINER.GPU[0])
        # self.device = torch.device('cpu')
        self.prec = cfg.TRAINER.CMPT.PREC

        print(f"Loading CLIP (backbone: {cfg.MODEL.BACKBONE.NAME})")
        clip_model = load_clip_to_cpu(cfg)
        if self.prec == "fp32" or self.prec == "amp":
            # CLIP's default precision is fp16
            clip_model.float()
    
        print("Building custom CLIP")
        self.model = CustomCLIP(cfg, self.device, self.classnames, clip_model)

        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)

        # Double check
        enabled = set()
        for name, param in self.model.named_parameters():
            if param.requires_grad:
                enabled.add(name)
                print(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)
        # NOTE: only give prompt_learner to the optimizer
        self.optim = build_optimizer(self.model, cfg.OPTIM)
        self.sched = build_lr_scheduler(self.optim, cfg.OPTIM)
        self.register_model("MultiModalPromptLearner", self.model, self.optim, self.sched)
        self.scaler = GradScaler() if self.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):
        XY_R, XY_L, XY_D = self.parse_batch_train(batch)
        model, optim, scaler = self.model, self.optim, self.scaler

        if self.prec == "amp":
            with autocast():
                logit, loss = model(XY_R, XY_L)
            optim.zero_grad()
            scaler.scale(loss).backward()
            scaler.step(optim)
            scaler.update()
        else:
            logit, loss = model(XY_R, XY_L)
            optim.zero_grad()
            loss.backward()
            optim.step()

        loss_summary = {
            "loss": loss.item(),
            "acc": cross_entropy(logit, XY_L)[0]
        }

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


    def model_inference(self, batch):
        logits = []
        with autocast():
            for idx, XY_R in enumerate(batch):
                logit = self.model(XY_R)
                logits.append(logit)
        logit = sum(logits) / len(logits)
        return logit

    def parse_batch_train(self, batch):
        X_R, X_L, X_D = batch['X_R'].to(self.device), batch['X_L'].to(self.device), batch['X_D'].to(self.device)
        Y_R, Y_L, Y_D = batch['Y_R'].to(self.device), batch['Y_L'].to(self.device), batch['Y_D'].to(self.device)

        XY_R = torch.cat([X_R, Y_R], dim=0)
        XY_L = torch.cat([X_L, Y_L], dim=0)
        XY_D = torch.cat([X_D, Y_D], dim=0)
        i = torch.randint(0, len(XY_L), (1,)).item()
        XY_R, XY_L, XY_D = XY_R[i:i + 1], XY_L[i:i + 1], XY_D[i:i + 1]
        return XY_R, XY_L, XY_D

    def parse_batch_test(self, batch):
        frame1, frame2, label, path = \
            batch['frame1'].to(self.device), batch['frame2'].to(self.device), \
            batch['label'].to(self.device), batch['path']
        return [frame1, frame2], 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 "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)

