import pandas as pd
import os
import torch
from PIL import Image
from torch.utils.data import Dataset
import clip
from torch.utils.data import DataLoader
import torchvision.transforms as tf
import torchvision.transforms.functional as TF


try:
    from torchvision.transforms import InterpolationMode
    BICUBIC = InterpolationMode.BICUBIC
except ImportError:
    BICUBIC = Image.BICUBIC
    
    
class ExtractFeaturesDataset(Dataset):
    def __init__(self,
                 annotations,
                 img_path,
                 image_transforms=None,
                 question_transforms=None,
                 tta=False):
        
        
        self.img_path = img_path
        self.image_transforms = image_transforms
        self.question_transforms = question_transforms
        
        self.img_ids = annotations["image_id"].values
        self.split = annotations["split"].values
        self.questions = annotations["question"].values
        
        self.tta = tta
       


    def __getitem__(self, index):

        image_id = self.img_ids[index]
        split = self.split[index]
        
        # image input
        with open(os.path.join(self.img_path, split, image_id), "rb") as f:
            img = Image.open(f)
            
            if self.tta:   
                image_augmentations = []
                
                for transform in self.image_transforms:
                    
                    image_augmentations.append(transform(img))
               
        
                img = torch.stack(image_augmentations, dim=0)
                
            else:
                img = self.image_transforms(img)
            
        question = self.questions[index]
        
        if self.question_transforms: 
            question = self.question_transforms(question)

        # question input
        question = clip.tokenize(question, truncate=True) 
        question = question.squeeze()
            
        return img, question, image_id
    
    def __len__(self):
        return len(self.img_ids)
    
        
def _convert_image_to_rgb(image):
    return image.convert("RGB")


def Sharpen(sharpness_factor=1.0):
    
    def wrapper(x):
    
        return TF.adjust_sharpness(x, sharpness_factor)
    
    return wrapper


def Rotate(angle=0.0):
    
    def wrapper(x):
        return TF.rotate(x, angle)
    
    return wrapper

def transform_crop(n_px):
    return tf.Compose([
        tf.Resize(n_px, interpolation=BICUBIC),
        tf.CenterCrop(n_px),
        _convert_image_to_rgb,
        tf.ToTensor(),
        tf.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
    ])

def transform_crop_rotate(n_px, rotation_angle=0.0):
    return tf.Compose([
        Rotate(angle=rotation_angle),
        tf.Resize(n_px, interpolation=BICUBIC),
        tf.CenterCrop(n_px),
        _convert_image_to_rgb,
        tf.ToTensor(),
        tf.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
    ])


def transform_resize(n_px):
    return tf.Compose([
        tf.Resize((n_px, n_px), interpolation=BICUBIC),
        _convert_image_to_rgb,
        tf.ToTensor(),
        tf.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
    ])


def transform_resize_rotate(n_px, rotation_angle=0.0):
    return tf.Compose([
        Rotate(angle=rotation_angle),
        tf.Resize((n_px, n_px), interpolation=BICUBIC),
        _convert_image_to_rgb,
        tf.ToTensor(),
        tf.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
    ])

def get_tta_preprocess(img_size):
    
    img_preprocess = [
                      transform_crop(img_size),
                      transform_crop_rotate(img_size, rotation_angle=90.0),
                      transform_crop_rotate(img_size, rotation_angle=270.0),
                      transform_resize(img_size),
                      transform_resize_rotate(img_size, rotation_angle=90.0),
                      transform_resize_rotate(img_size, rotation_angle=270.0),
                      ]
    
    return img_preprocess

def question_preprocess(question, debug=False):
    
    question = question.replace("?", ".")

    if question[-1] == " ":
        question = question[:-1]
         

    if question[-1] != ".":
        question = question + "."				
				   
    if debug:
        print("Question:", question)
        
    return question

            
def get_dataloader_extraction(config):


    if config.use_question_preprocess:
        print("Using custom preprocessing: Question")
        question_transforms = question_preprocess
    else:
        question_transforms = None
    
    if config.tta:
        ("Using augmentation transforms:")
        img_preprocess = get_tta_preprocess(config.img_size)
    else:
        ("Using original CLIP transforms:")
        img_preprocess = transform_crop(config.img_size)
           
    

    train_data = pd.read_csv(config.train_annotations_path)
      
    train_dataset = ExtractFeaturesDataset(annotations = train_data,
                                      img_path=config.img_path,
                                      image_transforms=img_preprocess,
                                      question_transforms=question_transforms,
                                      tta=config.tta)
    
    
    
    train_loader = DataLoader(dataset=train_dataset, 
                              batch_size=config.batch_size, 
                              shuffle=False,
                              num_workers=config.num_workers)
    
         
    
    test_data = pd.read_csv(config.test_annotations_path)
      
    test_dataset = ExtractFeaturesDataset(annotations = test_data,
                                      img_path=config.img_path,
                                      image_transforms=img_preprocess,
                                      question_transforms=question_transforms,
                                      tta=config.tta)
    
    
    test_loader = ExtractFeaturesDataset(dataset=test_dataset, 
                              batch_size=config.batch_size, 
                              shuffle=False,
                              num_workers=config.num_workers)
    
    return train_loader, test_loader


def get_dataloader_inference(config):
    
    if config.use_question_preprocess:
        print("Using custom preprocessing: Question")
        question_transforms = question_preprocess
    else:
        question_transforms = None
    
    if config.tta:
        ("Using augmentation transforms:")
        img_preprocess = transform_resize(config.img_size)
    else:
        ("Using original CLIP transforms:")
        img_preprocess = transform_crop(config.img_size)
           
    

    train_data = pd.read_csv(config.train_annotations_path)
      
    train_dataset = ExtractFeaturesDataset(annotations = train_data,
                                      img_path=config.img_path,
                                      image_transforms=img_preprocess,
                                      question_transforms=question_transforms,
                                      tta=config.tta)
    
    
    
    train_loader = DataLoader(dataset=train_dataset, 
                              batch_size=config.batch_size, 
                              shuffle=False,
                              num_workers=config.num_workers)
    
         
    
    test_data = pd.read_csv(config.test_annotations_path)
      
    test_dataset = ExtractFeaturesDataset(annotations = test_data,
                                      img_path=config.img_path,
                                      image_transforms=img_preprocess,
                                      question_transforms=question_transforms,
                                      tta=config.tta)
    
    
    test_loader = ExtractFeaturesDataset(dataset=test_dataset, 
                              batch_size=config.batch_size, 
                              shuffle=False,
                              num_workers=config.num_workers)
    
    return train_loader, test_loader