import os

import pandas as pd
import torch

from .clip import clip
from PIL import Image
import torch.nn as nn


CHANNELS = {
    "RN50" : 1024,
    "ViT-L/14" : 768
}

class CLIPModel(nn.Module):
    def __init__(self, name,hyperLoRA, num_classes=1):
        super(CLIPModel, self).__init__()
        self.config = {
            'example_key': 'example_value'
        }
        self.hyperLoRA=hyperLoRA
        self.model, self.preprocess = clip.load(name,self.hyperLoRA, device="cuda") # self.preprecess will not be used during training, which is handled in Dataset class
        # self.fc=nn.Linear(CHANNELS[name],3)
        self.fc = nn.Linear( CHANNELS[name], num_classes )

    def forward(self, x, task,return_feature=False):
        features = self.model.encode_image(x,task)
        if return_feature:
            return features
        featuresOrg=features
        lora_x = self.hyperLoRA(features, task, 8, 2)
        return self.fc(featuresOrg) +  0.01*lora_x


