from collections import OrderedDict
from typing import Tuple, Union

import numpy as np
import torch
import torch.nn.functional as F
from torchvision.transforms import functional as FUN
from torch import nn

from PIL import Image
from torchvision.transforms import functional as F1
try:
    from torchvision.transforms import InterpolationMode
    BICUBIC = InterpolationMode.BICUBIC
except ImportError:
    BICUBIC = Image.BICUBIC
from clip import clip
from clip.simple_tokenizer import SimpleTokenizer as _Tokenizer
from modules.Mamba_v3 import PTCOMambaBlock as PTMamba
_tokenizer = _Tokenizer()
from Mutil_prompt_CLIP.LightTrans import LightTransformer, LayerNorm,SPE_prompt_injection,SPA_prompt_injection

import os
def load_clip_to_cpu(basemodel,trainer,deep_prompts_depth,prompt_nctx):
    backbone_name = basemodel
    url = clip._MODELS[backbone_name]
    print(url)
    model_path = clip._download(url)
    print(model_path)
    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(), trainer,deep_prompts_depth,prompt_nctx)

    return model
class AdaIN_trans(nn.Module):
    def __init__(self):
        super().__init__()

    def mu(self, x):
        return torch.sum(x, (1)) / (x.shape[1])

    def sigma(self, x):
        return torch.sqrt(
            (torch.sum((x.permute([1, 0, 2]) - self.mu(x)).permute([1, 0, 2]) ** 2, (1)) + 0.000000023) / (x.shape[1]))

    def forward(self, x, y):
        return (self.sigma(y) * ((x.permute([1, 0, 2]) - self.mu(x)) / self.sigma(x)) + self.mu(y)).permute([1, 0, 2])

# for ViT :
class multi_scale(nn.Module):
    def __init__(self):
        super(multi_scale, self).__init__()
        self.linear = nn.ModuleList(nn.Linear(768, 512) for _ in range(12))
        self.adain = AdaIN_trans()
        self.gap = nn.AdaptiveAvgPool2d((1, 1))
        self.mamba = nn.ModuleList(PTMamba(hidden_dim=768, d_state=16, expand=0.5) for _ in range(12))

    def forward(self, data):
        data_prompt = []
        for i in range(len(data)):
            x_mu = self.adain.mu(data[i])
            x_mu = x_mu.to(torch.float32)
            x_mu = x_mu.unsqueeze(1)
            x_mu = self.mamba[i](x_mu)
            x = self.linear[i](x_mu)
            data_prompt.append(x)
        data_prompt = torch.stack(data_prompt, 1)
        return data_prompt
class MultiScaleConvModule(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_sizes=(1, 3, 5), stride=1, padding=None, dilation=1, groups=1,
                 bias=False):
        super(MultiScaleConvModule, self).__init__()
        if padding is None:
            padding = [k // 2 for k in kernel_sizes]

        self.convs = nn.ModuleList([
            nn.Conv2d(in_channels, out_channels, kernel_size=k, stride=stride, padding=p, dilation=dilation,
                      groups=groups, bias=bias)
            for k, p in zip(kernel_sizes, padding)
        ])

    def forward(self, x):
        return torch.cat([conv(x) for conv in self.convs], dim=1)
class projector(nn.Module):
    def __init__(self):
        super(projector, self).__init__()
        self.adain = AdaIN_trans()

    def forward(self, im_features):
        im_prompt = []
        x_mu = self.adain.mu(im_features)
        im_prompt.append(x_mu)
        im_prompt = torch.stack(im_prompt, 1)
        return im_prompt
class PatchEncoder(nn.Module):
    def __init__(self,in_chans_LIDAR,in_chans):
        super(PatchEncoder, self).__init__()
        self.conv1 = nn.ConvTranspose2d(
            in_channels=in_chans_LIDAR + in_chans,
            out_channels=32,
            kernel_size=6,
            stride=3,
            padding=4,
            # output_padding=1
        )
        self.activation1 = nn.GELU()
        self.bn1 = nn.BatchNorm2d(32)  #
        self.conv3 = nn.ConvTranspose2d(
            in_channels=32,
            out_channels=16,
            kernel_size=2,
            stride=2,
            padding=0
        )
        self.bn2 = nn.BatchNorm2d(16)  #
        self.conv4 = nn.ConvTranspose2d(
            in_channels=16,
            out_channels=8,
            kernel_size=2,
            stride=2,
            padding=0
        )
        self.bn3 = nn.BatchNorm2d(8)  #
        self.conv5 = nn.ConvTranspose2d(
            in_channels=8,
            out_channels=3,
            kernel_size=2,
            stride=2,
            padding=0
        )
        self.bn4 = nn.BatchNorm2d(3)
    def forward(self, x):
        x = self.conv1(x)
        x = self.activation1(x)
        x = self.bn1(x)
        x = self.conv3(x)
        x = self.activation1(x)
        x = self.bn2(x)
        x = self.conv4(x)
        x = self.activation1(x)
        x = self.bn3(x)
        x = self.conv5(x)
        x = self.activation1(x)
        x = self.bn4(x)
        return x
class PatchEncoder7(nn.Module):
    def __init__(self,in_chans_LIDAR,in_chans):
        super(PatchEncoder7, self).__init__()
        self.conv1 = nn.ConvTranspose2d(
            in_channels=in_chans_LIDAR + in_chans,
            out_channels=32,
            kernel_size=4,
            stride=4,
            padding=0,
            # output_padding=1
        )
        self.activation1 = nn.GELU()
        self.bn1 = nn.BatchNorm2d(32)

        self.conv3 = nn.ConvTranspose2d(
            in_channels=32,
            out_channels=16,
            kernel_size=2,
            stride=2,
            padding=0
        )
        self.bn2 = nn.BatchNorm2d(16)
        self.conv4 = nn.ConvTranspose2d(
            in_channels=16,
            out_channels=8,
            kernel_size=2,
            stride=2,
            padding=0
        )
        self.bn3 = nn.BatchNorm2d(8)
        self.conv5 = nn.ConvTranspose2d(
            in_channels=8,
            out_channels=3,
            kernel_size=2,
            stride=2,
            padding=0
        )
        self.bn4 = nn.BatchNorm2d(3)
    def forward(self, x):

        x = self.conv1(x)
        x = self.activation1(x)
        x = self.bn1(x)

        x = self.conv3(x)
        x = self.activation1(x)
        x = self.bn2(x)

        x = self.conv4(x)
        x = self.activation1(x)
        x = self.bn3(x)
        x = self.conv5(x)
        x = self.activation1(x)
        x = self.bn4(x)
        return x
class PatchEncoder9(nn.Module):
    def __init__(self,in_chans_LIDAR,in_chans):
        super(PatchEncoder9, self).__init__()
        self.conv1 = nn.ConvTranspose2d(
            in_channels=in_chans_LIDAR + in_chans,
            out_channels=32,
            kernel_size=4,
            stride=4,
            padding=4,
            # output_padding=1
        )
        self.activation1 = nn.GELU()
        self.bn1 = nn.BatchNorm2d(32)


        self.conv3 = nn.ConvTranspose2d(
            in_channels=32,
            out_channels=16,
            kernel_size=2,
            stride=2,
            padding=0
        )
        self.bn2 = nn.BatchNorm2d(16)

        self.conv4 = nn.ConvTranspose2d(
            in_channels=16,
            out_channels=8,
            kernel_size=2,
            stride=2,
            padding=0
        )
        self.bn3 = nn.BatchNorm2d(8)

        self.conv5 = nn.ConvTranspose2d(
            in_channels=8,
            out_channels=3,
            kernel_size=2,
            stride=2,
            padding=0
        )
        self.bn4 = nn.BatchNorm2d(3)
    def forward(self, x):

        x = self.conv1(x)
        x = self.activation1(x)
        x = self.bn1(x)
        x = self.conv3(x)
        x = self.activation1(x)
        x = self.bn2(x)

        x = self.conv4(x)
        x = self.activation1(x)
        x = self.bn3(x)
        x = self.conv5(x)
        x = self.activation1(x)
        x = self.bn4(x)
        return x
class PatchEncoder13(nn.Module):
    def __init__(self,in_chans_LIDAR,in_chans):
        super(PatchEncoder13, self).__init__()
        self.conv1 = nn.ConvTranspose2d(
            in_channels=in_chans_LIDAR + in_chans,
            out_channels=32,
            kernel_size=4,
            stride=2,
            padding=0,
        )
        self.activation1 = nn.GELU()
        self.bn1 = nn.BatchNorm2d(32)


        self.conv3 = nn.ConvTranspose2d(
            in_channels=32,
            out_channels=16,
            kernel_size=2,
            stride=2,
            padding=0
        )
        self.bn2 = nn.BatchNorm2d(16)

        self.conv4 = nn.ConvTranspose2d(
            in_channels=16,
            out_channels=8,
            kernel_size=2,
            stride=2,
            padding=0
        )
        self.bn3 = nn.BatchNorm2d(8)

        self.conv5 = nn.ConvTranspose2d(
            in_channels=8,
            out_channels=3,
            kernel_size=2,
            stride=2,
            padding=0
        )
        self.bn4 = nn.BatchNorm2d(3)
    def forward(self, x):

        x = self.conv1(x)
        x = self.activation1(x)
        x = self.bn1(x)

        x = self.conv3(x)
        x = self.activation1(x)
        x = self.bn2(x)

        x = self.conv4(x)
        x = self.activation1(x)
        x = self.bn3(x)
        x = self.conv5(x)
        x = self.activation1(x)
        x = self.bn4(x)
        return x
class PatchEncoder23(nn.Module):
    def __init__(self,in_chans_LIDAR,in_chans):
        super(PatchEncoder23, self).__init__()
        self.conv1 = nn.ConvTranspose2d(
            in_channels=in_chans_LIDAR + in_chans,
            out_channels=32,
            kernel_size=6,
            stride=1,
            padding=0,
            # output_padding=1
        )
        self.activation1 = nn.GELU()
        self.bn1 = nn.BatchNorm2d(32)
        self.conv3 = nn.ConvTranspose2d(
            in_channels=32,
            out_channels=16,
            kernel_size=2,
            stride=2,
            padding=0
        )
        self.bn2 = nn.BatchNorm2d(16)
        self.conv4 = nn.ConvTranspose2d(
            in_channels=16,
            out_channels=8,
            kernel_size=2,
            stride=2,
            padding=0
        )
        self.bn3 = nn.BatchNorm2d(8)
        self.conv5 = nn.ConvTranspose2d(
            in_channels=8,
            out_channels=3,
            kernel_size=2,
            stride=2,
            padding=0
        )
        self.bn4 = nn.BatchNorm2d(3)
    def forward(self, x):
        x = self.conv1(x)
        x = self.activation1(x)
        x = self.bn1(x)

        x = self.conv3(x)
        x = self.activation1(x)
        x = self.bn2(x)

        x = self.conv4(x)
        x = self.activation1(x)
        x = self.bn3(x)
        x = self.conv5(x)
        x = self.activation1(x)
        x = self.bn4(x)
        return x
class PatchEncoder17(nn.Module):
    def __init__(self,in_chans_LIDAR,in_chans):
        super(PatchEncoder17, self).__init__()
        self.conv1 = nn.ConvTranspose2d(
            in_channels=in_chans_LIDAR + in_chans,
            out_channels=32,
            kernel_size=2,
            stride=2,
            padding=3,
            # output_padding=1
        )
        self.activation1 = nn.GELU()
        self.bn1 = nn.BatchNorm2d(32)


        self.conv3 = nn.ConvTranspose2d(
            in_channels=32,
            out_channels=16,
            kernel_size=2,
            stride=2,
            padding=0
        )
        self.bn2 = nn.BatchNorm2d(16)

        self.conv4 = nn.ConvTranspose2d(
            in_channels=16,
            out_channels=8,
            kernel_size=2,
            stride=2,
            padding=0
        )
        self.bn3 = nn.BatchNorm2d(8)

        self.conv5 = nn.ConvTranspose2d(
            in_channels=8,
            out_channels=3,
            kernel_size=2,
            stride=2,
            padding=0
        )
        self.bn4 = nn.BatchNorm2d(3)
    def forward(self, x):

        x = self.conv1(x)
        x = self.activation1(x)
        x = self.bn1(x)


        x = self.conv3(x)
        x = self.activation1(x)
        x = self.bn2(x)

        x = self.conv4(x)
        x = self.activation1(x)
        x = self.bn3(x)
        x = self.conv5(x)
        x = self.activation1(x)
        x = self.bn4(x)
        return x
class PatchEncoder32(nn.Module):
    def __init__(self,in_chans_LIDAR,in_chans):
        super(PatchEncoder32, self).__init__()
        self.activation1 = nn.GELU()
        self.bn1 = nn.BatchNorm2d(32)

        self.conv3 = nn.ConvTranspose2d(
            in_channels=in_chans_LIDAR + in_chans,
            out_channels=16,
            kernel_size=2,
            stride=2,
            padding=5
        )
        self.bn2 = nn.BatchNorm2d(16)

        self.conv4 = nn.ConvTranspose2d(
            in_channels=16,
            out_channels=8,
            kernel_size=2,
            stride=2,
            padding=0
        )
        self.bn3 = nn.BatchNorm2d(8)

        self.conv5 = nn.ConvTranspose2d(
            in_channels=8,
            out_channels=3,
            kernel_size=2,
            stride=2,
            padding=0
        )
        self.bn4 = nn.BatchNorm2d(3)
    def forward(self, x):

        x = self.conv3(x)
        x = self.activation1(x)
        x = self.bn2(x)

        x = self.conv4(x)
        x = self.activation1(x)
        x = self.bn3(x)
        x = self.conv5(x)
        x = self.activation1(x)
        x = self.bn4(x)
        return x
class InjectionBlock(nn.Module):
    def __init__(self, vis, ctx):
        super(InjectionBlock, self).__init__()
        self.attention = nn.Sequential(
            nn.Linear(vis, vis // 16, bias=False),
            nn.ReLU(inplace=True),
            nn.Linear(vis // 16, vis, bias=False),
            nn.Sigmoid()
        )

        self.linearlayer1 = nn.Sequential(nn.Linear((vis * 2), vis))
        self.linearlayer2 = nn.Sequential(nn.Linear((vis * 2), ctx))
        self.gap = nn.AdaptiveAvgPool2d((1, 512))

    def forward(self, vis):
        vis_f = self.gap(vis)
        attn1 = self.attention(vis_f.type(torch.float))
        mulattn1 = torch.mul(attn1, vis_f)
        resattn1 = torch.cat((mulattn1, vis_f), 2)
        linear1 = self.linearlayer1(resattn1)

        attn2 = self.attention(linear1.type(torch.float))
        mulattn2 = torch.mul(attn2, vis_f)
        resattn2 = torch.cat((mulattn2, vis_f), 2)
        linear2 = self.linearlayer2(resattn2)

        output = linear2.to(torch.float16)
        return output
class ConvLayer(nn.Module):
    def __init__(self, out_size):
        super(ConvLayer, self).__init__()
        self.out_size = out_size
        self.downsample = nn.Sequential(
            nn.Conv2d(1, 1, kernel_size=3, stride=2, padding=1),  # 1089 → ~545
            nn.BatchNorm2d(1),
            nn.ReLU(),

            nn.Conv2d(1, 1, kernel_size=3, stride=2, padding=1),  # 545 → ~273
            nn.BatchNorm2d(1),
            nn.ReLU(),

            nn.Conv2d(1, 1, kernel_size=3, stride=2, padding=1),  # ... → ~136
            nn.ReLU(),

            nn.AdaptiveAvgPool2d((out_size, out_size))  
        )

    def forward(self, x):
        x = x.view(x.size(0), 1, 1089, 1089)  
        x = self.downsample(x)
        x = x.view(x.size(0), self.out_size, self.out_size)
        return x
class VisionEncoder(nn.Module):
    def __init__(self, clip_model):
        super().__init__()
        self.transformer = clip_model.visual
    def forward(self, x, prompts_F, weight,flag=True):
        x = x.to(torch.float16)
        if flag:
            x, data= self.transformer(x)
        else:
            counter=0
            outputs, data = self.transformer([x,prompts_F,weight,counter])
            x = outputs[0]
        return x, data
class DUTextEncoder(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, deep_prompts):
        x = prompts + self.positional_embedding.type(self.dtype)
        x = x.permute(1, 0, 2)  # NLD -> LND

        combined = [x, deep_prompts, 0]
        outputs = self.transformer(combined)
        x = outputs[0][0]
        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)] @ self.text_projection

        return x
class con_linear(nn.Module):
    def __init__(self, vis, ctx):
        super(con_linear, self).__init__()
        self.linear = nn.Linear(vis, ctx)
        self.gap = nn.AdaptiveAvgPool2d((1, 512))

    def forward(self, vis):
        vis_f = self.gap(vis)
        attn1 = self.linear(vis_f.type(torch.float))
        return attn1
class CrossAttentionModule(nn.Module):
    def __init__(self, channels_HSI, channels_Lidar, reduction_ratio=8):
        super().__init__()


        self.ll_v = nn.Conv2d(channels_HSI + channels_Lidar, channels_HSI + channels_Lidar, 1)


        self.ll_q_HSI = nn.Conv2d(channels_Lidar, channels_Lidar, 1)
        self.ll_k_HSI = nn.Conv2d(channels_HSI, channels_Lidar, 1)

        self.ll_q_Lidar = nn.Conv2d(channels_HSI, channels_HSI, 1)
        self.ll_k_Lidar = nn.Conv2d(channels_Lidar, channels_HSI, 1)
        self.ptconv = ConvLayer(out_size=channels_HSI)
        self.ptpool = nn.MaxPool2d(kernel_size=33, stride=33, padding=0)
    def forward(self, f_HSI, f_Lidar):

        batch_size, _, h, w = f_HSI.shape


        combined = torch.cat([f_HSI, f_Lidar], dim=1)

        value = self.ll_v(combined)
        value = value.view(batch_size, -1, h * w)

        query_HSI = f_Lidar.view(batch_size, -1, h * w) #bx11x121
        key_HSI = self.ll_k_HSI(f_HSI).view(batch_size, -1, h * w) #bx11x121
        attention_HSI_SA = torch.bmm(query_HSI.transpose(1, 2), key_HSI)

        attention_HSI_SA = F.softmax(attention_HSI_SA, dim=-1)
        value_HSI = value[:, f_HSI.size(1):].view(batch_size, -1, h * w)
        updated_Lidar_SA = f_Lidar + torch.bmm(value_HSI, attention_HSI_SA).view_as(f_Lidar)


        query_Lidar = f_HSI.view(batch_size, -1, h * w)
        key_Lidar = self.ll_k_Lidar(f_Lidar).view(batch_size, -1, h * w)
        attention_Lidar_SP = torch.bmm(query_Lidar, key_Lidar.transpose(1, 2))
        attention_Lidar_SP = F.softmax(attention_Lidar_SP, dim=-1)
        value_Lidar = value[:, :f_HSI.size(1)].view(batch_size, -1, h * w)
        updated_HSI_SP = f_HSI + torch.bmm(attention_Lidar_SP, value_Lidar).view_as(f_HSI)

        return updated_Lidar_SA,updated_HSI_SP,attention_HSI_SA,attention_Lidar_SP


class FusionModule(nn.Module):
    def __init__(self, in_channels, in_channels_LIDAR):
        super(FusionModule, self).__init__()
        self.conv1 = nn.Sequential(nn.Conv2d(in_channels=in_channels+in_channels_LIDAR, out_channels=in_channels+in_channels_LIDAR, kernel_size=1),
                                   nn.ReLU(inplace=True),nn.BatchNorm2d(in_channels+in_channels_LIDAR),nn.Dropout(0.1))
        self.conv3 = nn.Sequential(nn.Conv2d(in_channels=in_channels+in_channels_LIDAR, out_channels=in_channels+in_channels_LIDAR, kernel_size=3, padding=1),
                                   nn.ReLU(inplace=True),nn.BatchNorm2d(in_channels+in_channels_LIDAR),nn.Dropout(0.1))
        self.conv5 = nn.Sequential(nn.Conv2d(in_channels=in_channels+in_channels_LIDAR, out_channels=in_channels+in_channels_LIDAR, kernel_size=5, padding=2),
                                   nn.ReLU(inplace=True),nn.BatchNorm2d(in_channels+in_channels_LIDAR),nn.Dropout(0.1))
        self.attention = nn.Sequential(
            nn.Conv2d(in_channels+in_channels_LIDAR, (in_channels+in_channels_LIDAR) // 2, 1),
            nn.ReLU(inplace=True),
            nn.Conv2d((in_channels+in_channels_LIDAR) // 2,in_channels+in_channels_LIDAR, 1),
        )
        self.dropout = nn.Dropout(0.1)
    def forward(self, f1H, f1L):
        f = torch.cat([f1H, f1L], dim=1)
        # Step 1: Compute A_avg and A_max using global average pooling and global max pooling
        A_avg = F.adaptive_avg_pool2d(f, (1, 1))
        A_max = F.adaptive_max_pool2d(f, (1, 1))
        A_max = self.attention(A_max)
        A_avg = self.attention(A_avg)
        A = A_avg + A_max
        A =F.softmax(A,dim=-1)
        # Step 4: Compute f_1fuse
        f_1fuse = self.conv3(f*A) + self.conv5(f*A) + self.conv1(f*A)
        return f_1fuse

class PromptLearner(nn.Module):
    def __init__(self,img_size, input_image_size, init_prompt, num_prompt, classnames, clip_model, deep_prompts_depth,prompt_nctx,in_chans,description,
                     CSC=False):
        super().__init__()
        n_cls = len(classnames)
        n_ctx = prompt_nctx
        t_n_ctx = num_prompt
        ctx_init = init_prompt
        dtype = clip_model.dtype
        ctx_dim = clip_model.ln_final.weight.shape[0]
        vis_dim = clip_model.visual.output_dim
        clip_imsize = clip_model.visual.input_resolution
        cfg_imsize = input_image_size
        visual_ctx_dim = clip_model.visual.positional_embedding.shape[1]

        self.deep_prompts_depth = deep_prompts_depth

        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
            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 + t_n_ctx, :]
            prompt_prefix = " ".join(ctx_init.split()[:n_ctx])
        else:
            # random initialization
            print(f"Initializing A Generic Context")
            ctx_vectors = torch.empty(n_ctx, ctx_dim, dtype=dtype)
            nn.init.normal_(ctx_vectors, std=0.02)
            prompt_prefix = " ".join(['X'] * n_ctx)
        self.ctx = nn.Parameter(ctx_vectors)  # to be optimized

        print(f'Initial context: "{prompt_prefix}"')
        print(f"Number of context words (tokens): {n_ctx}")
        print(f"Depth of deep prompt: {self.deep_prompts_depth}")
        self.deep_prompts = nn.Parameter(torch.empty(self.deep_prompts_depth-1, n_ctx, ctx_dim))
        nn.init.normal_(self.deep_prompts, std=0.02)
        self.self_attn = PTMamba(hidden_dim=ctx_dim, d_state=16, expand=0.25,ctx_dim=visual_ctx_dim,in_chans=in_chans,in_img_size =img_size)
        classnames = [name.replace("_", " ") for name in classnames]
        prompts = [prompt_prefix+" "+ name + "." for name in classnames]#+des+" ",des zip(,description)
        tokenized_prompts = torch.cat([clip.tokenize(p) for p in prompts], dim=0)  # (n_cls, 77)
        with torch.no_grad():
            embedding = clip_model.token_embedding(tokenized_prompts).type(dtype)  # (n_cls, 77, ctx_dim)
        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.dtype = dtype
        self.ctx_dim = ctx_dim
        self.tokenized_prompts = tokenized_prompts  # torch.Tensor

    def construct_prompts(self, ctx, prefix, suffix, label=None):
        if label is not None:
            prefix = prefix[label]
            suffix = suffix[label]

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

        return prompts

    def forward(self,max_spa,max_spe):#bias_t
        ctx = self.ctx  # context parameters (n_ctx, ctx_dim)
        if ctx.dim() == 2:  # generic context (n_ctx, n_dim)
            ctx = ctx.unsqueeze(0).expand(self.n_cls, self.n_ctx, self.ctx_dim)

        prefix = self.token_prefix  # SOS
        suffix = self.token_suffix  # CLS . EOS ~
        ctx = self.construct_prompts(ctx, prefix, suffix)
        # # textual prompts to visual prompts
        visual_prompts = torch.cat([self.ctx.unsqueeze(0), self.deep_prompts], dim=0)
        visual_prompts_s = self.self_attn(visual_prompts,max_spe,max_spa)
        deep_prompts = self.deep_prompts
        return ctx, deep_prompts, visual_prompts_s

class DUCustomCLIP(nn.Module):
    def __init__(self, img_size,input_image_size,init_prompt,num_prompt,classnames, clip_model,deep_prompts_depth,prompt_nctx,in_chans,description):
        super().__init__()
        self. PromptLearner =  PromptLearner(img_size,input_image_size,init_prompt,num_prompt,classnames, clip_model,deep_prompts_depth,prompt_nctx,in_chans,description)
        self.tokenized_prompts = self. PromptLearner.tokenized_prompts  # (n_cls, 77)
        self.text_encoder = DUTextEncoder(clip_model)
        self.image_encoder = clip_model.visual
        self.logit_scale = clip_model.logit_scale
        self.dtype = clip_model.dtype
        ctx_dim = clip_model.ln_final.weight.shape[0]
        width = clip_model.visual.positional_embedding.shape[1]
        visual_ctx_dim = clip_model.visual.positional_embedding.shape[1]
        self.visual_ctx_self_attn = PTMamba(hidden_dim=width, d_state=16, expand=0.25, ctx_dim=ctx_dim,in_chans=in_chans,in_img_size = img_size)
        self.visual_ctx_ln_intra_post = LayerNorm(width)
        self.deep_prompts_depth = deep_prompts_depth
    def forward(self, image,prompt_SA,prompt_SE):#,bias,bias_t
        tokenized_prompts = self.tokenized_prompts  # (n_cls, 77)
        prompts, text_deep_prompts, visual_deep_prompts = self. PromptLearner(prompt_SA,prompt_SE)#,bias_t


        image_features, textual_prompts = self.image_encoder(image.type(self.dtype), visual_deep_prompts[:1, :, :],visual_deep_prompts[1:, :, :])

        textual_prompts_s = self.visual_ctx_self_attn(textual_prompts, prompt_SE, prompt_SA)

        text_deep_prompts = text_deep_prompts + textual_prompts_s

        text_features = self.text_encoder(prompts, tokenized_prompts, text_deep_prompts)  # (n_cls, 1024)

        image_features = image_features / image_features.norm(dim=-1, keepdim=True)
        text_features = text_features / text_features.norm(dim=-1, keepdim=True)

        logit_scale = self.logit_scale.exp()
        logits = logit_scale * image_features @ text_features.t()

        return logits,text_features
class CLIP(nn.Module):
    def __init__(self, img_size, in_chans, in_chans_LIDAR, hid_chans, hid_chans_LIDAR, embed_dim, deep_prompts_depth, num_heads, mlp_ratio,
                 num_classes, global_pool,basemodel,classnames,input_image_size, init_prompt, num_prompt,context_length,
                 vocab_size, transformer_width, transformer_heads, transformer_layers, trainer,prompt_nctx,description
    ):
        super().__init__()
        self.context_length = context_length
        print(f"Loading CLIP (backbone: {basemodel})")
        clip_model = load_clip_to_cpu(basemodel,trainer,deep_prompts_depth,prompt_nctx)
        # if cfg.TRAINER.APPLENET.PREC == "fp32" or cfg.TRAINER.APPLENET.PREC == "amp":
        #     # CLIP's default precision is fp16
        # clip_model.float()
        print("Building custom CLIP")
        # self.clip_model = clip_model
        if img_size == 17:
            self.visual = PatchEncoder17(in_chans,in_chans_LIDAR)
        elif img_size == 11:
            self.visual = PatchEncoder(in_chans,in_chans_LIDAR)
        elif img_size == 9:
            self.visual = PatchEncoder9(in_chans,in_chans_LIDAR)
        elif img_size == 7:
            self.visual = PatchEncoder7(in_chans,in_chans_LIDAR)
        elif img_size == 33:
            self.visual = PatchEncoder32(in_chans,in_chans_LIDAR)
        elif img_size == 13:
            self.visual = PatchEncoder13(in_chans,in_chans_LIDAR)
        elif img_size == 23:
            self.visual = PatchEncoder23(in_chans,in_chans_LIDAR)
        self.in_chans = in_chans
        self.in_chans_LIDAR = in_chans_LIDAR
        self.image_size = img_size
        self.feature_Interaction = CrossAttentionModule(in_chans,in_chans_LIDAR)
        self.feature_Fusion = FusionModule(in_chans,in_chans_LIDAR)
        print("Turning off gradients in both the image and the text encoder")
        self.DUCustomCLIP = DUCustomCLIP(img_size,input_image_size,init_prompt,num_prompt,classnames, clip_model,deep_prompts_depth,prompt_nctx,in_chans,description)
        name_to_update = "PromptLearner" # PromptLearner
        for name, param in self.DUCustomCLIP.named_parameters():
            if name_to_update not in name:
                if "visual_ctx" in name:
                    param.requires_grad_(True)
                else:
                    param.requires_grad_(False)
            # Double check
        enabled = set()
        for name, param in self.DUCustomCLIP.named_parameters():
            if param.requires_grad:
                enabled.add(name)
        print(f"Parameters to be updated: {enabled}")
        device_count = torch.cuda.device_count()
        print(f'gpu detected : {device_count}')
        if device_count > 1:
            print(f"Multiple GPUs detected (n_gpus={device_count}), use all of them!")
            self.DUCustomCLIP = nn.DataParallel(self.DUCustomCLIP)

    def forward(self, image, image_LIDAR):
        image_sp, image_sa ,prompt_SA, prompt_SP = self.feature_Interaction(image, image_LIDAR)
        image_f = self.feature_Fusion(image_sp, image_sa)
        image_f = torch.cat((image_sp, image_sa), 1)+image_f
        image_project = self.visual(image_f)
        logits,text_features = self.DUCustomCLIP(image_project, prompt_SA, prompt_SP)#bias,bias_t
        return logits ,text_features



def build_model(img_size, in_chans, in_chans_LIDAR, hid_chans, hid_chans_LIDAR, embed_dim, deep_prompts_depth, num_heads, mlp_ratio,num_classes,basemodel,classnames,input_image_size, init_prompt, num_prompt,trainer,prompt_nctx,description,global_pool=False):

    context_length = 77
    vocab_size = 49408
    transformer_width = 128
    transformer_heads = 8
    transformer_layers = 3
    model = CLIP(img_size,in_chans, in_chans_LIDAR,hid_chans, hid_chans_LIDAR, embed_dim, deep_prompts_depth, num_heads, mlp_ratio,num_classes,
                 global_pool, basemodel, classnames,input_image_size, init_prompt, num_prompt,context_length, vocab_size,
                 transformer_width, transformer_heads, transformer_layers,trainer,prompt_nctx,description
    )
    return model