import torch
from PIL import Image
import pandas as pd
from transformers import PreTrainedTokenizer, BaseImageProcessor
from typing import Dict, Sequence, List
from torch.utils.data import Dataset
from random import randrange
from src.custom_args import CustomTrainingArguments
from src.custom_utils import (
    get_logger,
    sample_mbeir_instruction,
    is_siglip_tokenizer,
    is_siglip2_tokenizer
)
from pathlib import Path

logger = get_logger(__name__)


class DataCollatorForCLIPFinetuning(object):
    def __init__(self, tokenizer: PreTrainedTokenizer, image_processor: BaseImageProcessor):
        self.tokenizer = tokenizer
        self.image_processor = image_processor

    def __call__(self, instances: Sequence[Dict]) -> Dict[str, torch.Tensor]:
        batch = {}
        for k in instances[0]:
            values = [(x[k] if k in x else '') for x in instances]
            if 'pixel_values' in k:
                batch[k] = torch.stack(values)

            elif k == 'query_input_text':
                if is_siglip_tokenizer(self.tokenizer):
                    padding = 'max_length'
                else:
                    padding = True
                if self.tokenizer is not None:
                    tik = self.tokenizer(
                        values, truncation=True, padding=padding, return_tensors='pt')
                    batch['query_input_ids'] = tik['input_ids']
                    batch['query_attention_mask'] = tik.get('attention_mask')

            elif k == 'passage_input_text':
                if is_siglip_tokenizer(self.tokenizer):
                    padding = 'max_length'
                else:
                    padding = True
                if self.tokenizer is not None:
                    tok = self.tokenizer(
                        values, truncation=True, padding=padding, return_tensors='pt')
                    batch['passage_input_ids'] = tok['input_ids']
                    batch['passage_attention_mask'] = tok.get('attention_mask')

            elif 'pixel_values' in k:
                batch[k] = torch.stack(values)

            elif 'vision_mask' in k or 'text_mask' in k:
                batch[k] = torch.tensor(
                    values, dtype=torch.int64)

            else:
                batch[k] = values
        return batch


class CustomDataset(Dataset):
    def __init__(self,
                 args: CustomTrainingArguments,
                 data_path: List[str],
                 tokenizer: PreTrainedTokenizer,
                 image_processor: BaseImageProcessor,
                 training: bool = True,
                 dataset_ratio: List[float] = None,
                 shuffle: bool = None,
                 ):
        super(CustomDataset, self).__init__()
        self.args = args
        self.tokenizer = tokenizer
        self.image_processor = image_processor
        self.tokenizer = tokenizer
        self.training = training
        self.shuffle = training if shuffle is None else shuffle
        self.dataset_ratio = dataset_ratio

        prefix_str = '[TRAIN]' if training else '[EVAL]'
        final_json = []
        if self.dataset_ratio is None:
            dataset_ratio = [1.0] * len(data_path)

        for el, ratio in zip(data_path, dataset_ratio):
            data = pd.read_json(el, lines=True, nrows=args.dataset_n_samples)
            if not self.training and args.dataset_eval_n_samples:
                data = data.sample(frac=1.0, ignore_index=True)
                data = data.head(args.dataset_eval_n_samples)

            # to avoid reshuffling when ratio is 1.0
            if ratio != 1.0:
                if ratio < 1.0:
                    data = data.sample(frac=ratio, ignore_index=True)
                else:
                    data = data.sample(
                        frac=ratio, replace=True, ignore_index=True)

            final_json.append(data)
            logger.info(
                f"{prefix_str} Loaded {len(data)} samples from {el}")

        self.data = pd.concat(final_json, ignore_index=True, axis=0)
        logger.info(
            f"{prefix_str} Loaded {len(self.data)} samples from {data_path}")

        if self.shuffle:
            self.data = self.data.sample(frac=1.0, ignore_index=True)
            logger.info(f"{prefix_str} Dataset shuffled")
        else:
            logger.info(f"{prefix_str} Dataset not shuffled")

        self.instructions = None
        if args.dataset_sample_mbeir_instructions:
            self.instructions = pd.read_csv(args.dataset_mbeir_instructions_path, sep='\t')
            logger.info(
                f"Loaded {len(self.instructions)} instructions from {args.dataset_mbeir_instructions_path}")

    @property
    def entity_labels(self):
        return self.data.passage_id.apply(lambda x: x[0] if type(x) is list else x)

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        sample = self.data.loc[idx]
        ret = {}

        if isinstance(sample.passage_id, list):
            passage_idx = 0
            if self.training and len(sample.passage_id) > 1:
                passage_idx = randrange(len(sample.passage_id))
            passage_id = sample.passage_id[passage_idx]
            passage_text = sample.passage_text[passage_idx]
            if isinstance(sample.passage_image_path, list):
                passage_image_path = sample.passage_image_path[passage_idx]
            else:
                passage_image_path = sample.passage_image_path
        else:
            passage_id = sample.passage_id
            passage_text = sample.passage_text
            if isinstance(passage_text, list):
                passage_text = passage_text[0]
            passage_image_path = sample.passage_image_path
            if isinstance(passage_image_path, list):
                passage_image_path = passage_image_path[0]

        instruction = sample.instruction
        if pd.isna(instruction):
            # TODO: assuming that we are training on M-BEIR
            assert self.args.dataset_sample_mbeir_instructions
            instruction = sample_mbeir_instruction(sample, self.instructions)
        if self.tokenizer is not None:
            if sample.question and not pd.isna(sample.question):
                query_input_text = instruction + \
                    ' ' + sample.question
            else:
                query_input_text = instruction
            if is_siglip2_tokenizer(self.tokenizer):
                query_input_text = query_input_text.lower()
        else:
            query_input_text = ''
        ret['query_input_text'] = query_input_text

        if self.tokenizer is not None:
            if pd.isna(passage_text):
                passage_input_text = ''
                passage_text_mask = 0
            else:
                passage_input_text = passage_text
                passage_text_mask = 1
            if is_siglip2_tokenizer(self.tokenizer):
                passage_input_text = passage_input_text.lower()
            ret['passage_input_text'] = passage_input_text
        ret['label_id'] = passage_id

        img_size = 336
        query_vision_mask = 0
        if self.image_processor is not None:
            if pd.isna(sample.image_path):
                query_pil = Image.new(
                    'RGB', (img_size, img_size), color='black').convert('RGB')
            else:
                query_pil = Image.open(sample.image_path).convert('RGB')
                query_vision_mask = 1
            ret['query_pixel_values'] = self.image_processor(
                query_pil, return_tensors='pt')['pixel_values'][0]

        passage_vision_mask = 0
        if self.image_processor is not None:
            if self.args.no_passage_images:
                passage_pil = Image.new(
                    'RGB', (img_size, img_size), color='black').convert('RGB')
            else:
                try:
                    passage_pil = Image.open(passage_image_path).convert('RGB')
                    passage_vision_mask = 1
                except:
                    # TODO: handle multiple models
                    # e.g. DiNO image_processor has no shortest_edge
                    # img_size = self.image_processor.size['shortest_edge']
                    passage_pil = Image.new(
                        'RGB', (img_size, img_size), color='black').convert('RGB')
            ret['passage_pixel_values'] = self.image_processor(
                passage_pil, return_tensors='pt')['pixel_values'][0]

        assert query_input_text or self.tokenizer is None
        query_text_mask = 1 if query_input_text else 0
        ret['query_text_mask'] = query_text_mask
        ret['query_vision_mask'] = query_vision_mask
        ret['passage_text_mask'] = passage_text_mask
        ret['passage_vision_mask'] = passage_vision_mask

        return ret


def build_data_module(args, data_paths, **kwargs):
    dataset_common_args = dict(
        args=args,
        **kwargs,
    )

    if not args.do_only_eval:
        dataset_train_args = dict(
            data_path=args.dataset_train, training=True, dataset_ratio=args.dataset_ratio, **dataset_common_args)

        ret_dict = dict(
            train_dataset=CustomDataset(**dataset_train_args),
            data_collator=DataCollatorForCLIPFinetuning(
                tokenizer=kwargs['tokenizer'],
                image_processor=kwargs['image_processor'])
        )
    else:
        ret_dict = dict(
            train_dataset=None,
            data_collator=DataCollatorForCLIPFinetuning(
                tokenizer=kwargs['tokenizer'],
                image_processor=kwargs['image_processor'])
        )

    if args.dataset_eval is not None:
        ret_dict['eval_dataset'] = {}
        for dataset_path in args.dataset_eval:
            dataset_name = Path(dataset_path).stem
            dataset_eval_args = dict(
                data_path=[dataset_path], training=False, **dataset_common_args)

            ret_dict['eval_dataset'][dataset_name] = CustomDataset(
                **dataset_eval_args)
    else:
        ret_dict['eval_dataset'] = None
    return ret_dict
