import torch
import clip
import os

from tqdm import tqdm
from PIL import Image
from torch.utils.data import DataLoader
from torchvision.datasets import CIFAR100

import numpy as np
from sklearn.linear_model import LogisticRegression



def get_features(dataset):
        all_features = []
        all_labels = []
        
        with torch.no_grad():
            for images, labels in tqdm(DataLoader(dataset, batch_size=100)):
                features = model.encode_image(images.to(device))

                all_features.append(features)
                all_labels.append(labels)

        return torch.cat(all_features).cpu().numpy(), torch.cat(all_labels).cpu().numpy()


if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser()
    
    parser.add_argument("--pt", type=str, help="模型名称")
    
    args = parser.parse_args()
    
    device = "cuda" if torch.cuda.is_available() else "cpu"
    
    if ".pt" in args.pt:
        model, preprocess = clip.load(f"pretrained_models/{args.pt}", device=device)
    else:
        model, preprocess = clip.load(f"{args.pt}", device=device)
        
    # Download the dataset
    cifar100 = CIFAR100(root=os.path.expanduser("~/.cache"), download=True, train=False)
    
    # Load the dataset
    root = os.path.expanduser("~/.cache")
    train = CIFAR100(root, download=True, train=True, transform=preprocess)
    test = CIFAR100(root, download=True, train=False, transform=preprocess)

    # Calculate the image features
    train_features, train_labels = get_features(train)
    test_features, test_labels = get_features(test)

    # Perform logistic regression
    classifier = LogisticRegression(random_state=0, C=0.316, max_iter=1000, verbose=1)
    classifier.fit(train_features, train_labels)

    # Evaluate using the logistic regression classifier
    predictions = classifier.predict(test_features)
    accuracy = np.mean((test_labels == predictions).astype(float)) * 100.
    print(f"Accuracy = {accuracy:.3f}")