import transformers
from transformers import TrainingArguments, Trainer, HfArgumentParser
from transformers import AutoImageProcessor
from datasets import load_dataset

import torch
import torch.nn.functional as F
import torch.distributed as dist
from torch.utils.data import Dataset

import random
import os
from PIL import Image
from pprint import pformat
from dataclasses import dataclass

from flmr import FLMRQueryEncoderTokenizer, FLMRContextEncoderTokenizer, FLMRModelForRetrieval


@dataclass
class MyArguments:
    model_name_or_path :str = "LinWeizheDragon/PreFLMR_ViT-G"
    image_processor_name :str = "laion/CLIP-ViT-bigG-14-laion2B-39B-b160k"
    dataset_name :str = "BByrneLab/multi_task_multi_modal_knowledge_retrieval_benchmark_M2KR"
    dataset_subset_name :str = "Infoseek"
    query_images_dir :str = "./Infoseek/train_images"
    num_negative_examples :int = 4


@dataclass
class PreFLMRTrainingArguments(TrainingArguments):
    # set according to PreFLMR paper 
    remove_unused_columns :bool = False
    per_device_train_batch_size :int = 8
    gradient_accumulation_steps :int = 8
    logging_steps :int = 1
    eval_strategy :str = 'no'
    save_strategy :str = 'steps'
    save_steps :int = 500
    max_steps :int = 1000
    save_only_model :bool = True
    save_total_limit :int = 5
    seed :int = 42
    # manually set optimizer later
    mapping_structure_lr :float = 1e-4
    non_mapping_structure_lr :float = 1e-5


class PreFLMRTrainer(Trainer):

    def compute_loss(self, model, inputs, return_outputs=False):
        outputs = model(**inputs, return_dict=True)
        
        # replace with ib loss
        ib_loss = outputs["in_batch_negative_loss"]
        outputs["deprecated_loss"] = outputs["loss"]
        outputs["loss"] = ib_loss

        return (ib_loss, outputs) if return_outputs else ib_loss


class PreFLMRDataset(Dataset):
    
    def __init__(self,
                 args,
                 data_df, passages_df, 
                 query_tokenizer, context_tokenizer, image_processor):
        self.args = args
        self.data_df = data_df
        self.passages_df = passages_df
        self.query_tokenizer = query_tokenizer
        self.context_tokenizer = context_tokenizer
        self.image_processor = image_processor
        
        self.unique_passage_ids = set(self.passages_df.index)
    
    def __len__(self):
        return len(self.data_df)
    
    def __getitem__(self, idx):
        row = self.data_df.iloc[idx]
        # NOTE: *infoseek* subset happends to have instructions prepended to `question`.
        #   for other subsets, instructions are required.
        # query = instruction + row['question']
        assert ':' in row['question'], 'Only Infoseek has instruction prepended to question in `question` field.'
        query = row['question']
        
        pos_item_ids = row['pos_item_ids']
        pos_item_id = random.choice(pos_item_ids) # NOTE random choose here?
        pos_passage = self.passages_df.loc[pos_item_id]['passage_content']
        
        neg_item_ids = random.sample(list(self.unique_passage_ids - set(pos_item_ids)), 
                                     self.args.num_negative_examples)  # NOTE random choose here?
        neg_passages = [self.passages_df.loc[neg_item_id]['passage_content'] for neg_item_id in neg_item_ids]
        
        image_path = os.path.join(self.args.query_images_dir, row['img_path'])
        image = Image.open(image_path).convert('RGB')
        pixel_values = self.image_processor(image, return_tensors='pt')['pixel_values'] # [1, 3, 224, 224]
        
        return dict(
            query=query,
            pos_passage=pos_passage,
            neg_passages=neg_passages,
            pixel_values=pixel_values
        )
        
    def collate_fn(self, batch):
        queries = [ex['query'] for ex in batch]
        passages = [] # [pos, neg, neg, neg, pos, ...]
        for ex in batch:
            passages.append(ex['pos_passage'])
            passages.extend(ex['neg_passages'])

        Q_encoding = self.query_tokenizer(queries)
        Q_pixel_values = torch.cat([ex['pixel_values'] for ex in batch], dim=0)
        D_encoding = self.context_tokenizer(passages)
        
        # according to `modeling_flmr.py, FLMRModelForRetrieval.forward`
        inputs = dict(
            query_input_ids=Q_encoding['input_ids'],
            query_attention_mask=Q_encoding['attention_mask'],
            query_pixel_values=Q_pixel_values,
            context_input_ids=D_encoding['input_ids'],
            context_attention_mask=D_encoding['attention_mask'],
            use_in_batch_negatives=True,
            in_batch_negatives_from_all_gpus=False, # NOTE should be False here?
            num_negative_examples=self.args.num_negative_examples
        )
        return inputs


def main():
    parser = HfArgumentParser((MyArguments, PreFLMRTrainingArguments))
    my_args, training_args = parser.parse_args_into_dataclasses()
    
    ## setting up
    assert dist.get_world_size() == 4, 'The paper uses 4 gpus.'
    if dist.get_rank() == 0:
        print('## my_args: ', pformat(my_args))
        print('## training_args: ', pformat(training_args))
    
    ## setting up tokenizer
    query_tokenizer = FLMRQueryEncoderTokenizer.from_pretrained(
        my_args.model_name_or_path, subfolder="query_tokenizer")
    context_tokenizer = FLMRContextEncoderTokenizer.from_pretrained(
        my_args.model_name_or_path, subfolder="context_tokenizer")
    image_processor = AutoImageProcessor.from_pretrained(my_args.image_processor_name)

    ## setting up dataset
    data_df = load_dataset(my_args.dataset_name, f'{my_args.dataset_subset_name}_data')['train']\
        .to_pandas().set_index('question_id')
    passages_df = load_dataset(my_args.dataset_name, f'{my_args.dataset_subset_name}_passages')['train_passages']\
        .to_pandas().set_index('passage_id')
        
    dataset = PreFLMRDataset(args=my_args,
                             data_df=data_df, passages_df=passages_df,
                             query_tokenizer=query_tokenizer, 
                             context_tokenizer=context_tokenizer, 
                             image_processor=image_processor)
    
    if dist.get_rank() == 0:
        print(pformat(f'## dataset[0]: {dataset[0]}'))
    
    ## setting up model
    model = FLMRModelForRetrieval.from_pretrained(
        my_args.model_name_or_path,
        query_tokenizer=query_tokenizer,
        context_tokenizer=context_tokenizer)
    
    ## setting up training
    # build trainables    
    # PreFLMR consists of
    #   pretrained_structure_modules: pretrained text and vision encoder, remain frozen
    #   mapping_structure_modules: a 2-layer MLP_F^MLP and a Transformer block F_M^TR, lr = 1e-4
    #   non_mapping_structure_modules: remaining modules, mostly linears. lr = 1e-5
    pretrained_structure_modules = [
        model.query_text_encoder, # FLMRTextModel
        model.query_vision_encoder, # FLMRVisionModel
        model.context_text_encoder, # FLMRTextModel
        model.context_vision_encoder # FLMRVisionModel
    ]
    mapping_structure_modules = [
        model.query_vision_projection, # FLMRMultiLayerPerceptron
        model.context_vision_projection, # FLMRMultiLayerPerceptron
        model.transformer_mapping_network, # BertEncoder
    ]
    non_mapping_structure_modules = [
        model.query_text_encoder_linear, # Linear
        model.context_text_encoder_linear, # Linear
        model.transformer_mapping_input_linear, # Linear
        model.transformer_mapping_output_linear, # Linear
    ]
    # check included all paramters, nothing left
    assert set(id(p) for p in model.parameters()) == set(id(p) \
        for module in pretrained_structure_modules + mapping_structure_modules + non_mapping_structure_modules
        for p in module.parameters())
    
    for module in pretrained_structure_modules:
        for p in module.parameters():
            p.requires_grad = False
    
    if dist.get_rank() == 0:
        trainables = [pn for pn, p in model.named_parameters() if p.requires_grad]
        n_trainables = sum(p.numel() for p in model.parameters() if p.requires_grad)
        print(pformat(f'## trainables: {trainables}'))
        print(pformat(f'## n_trainables: {n_trainables:,}'))
    
    # build optimizer and constant scheduler
    #   need to deduplicate modules (some modules may share parameters)
    optimizer_groups = []
    
    mapping_structure_modules_dedup = list({id(m) : m for m in mapping_structure_modules}.values())
    for module in mapping_structure_modules_dedup:
        optimizer_groups.append(dict(params=module.parameters(), lr=training_args.mapping_structure_lr))
    
    non_mapping_structure_modules_dedup = list({id(m) : m for m in (non_mapping_structure_modules)}.values())
    for module in non_mapping_structure_modules_dedup:
        optimizer_groups.append(dict(params=module.parameters(), lr=training_args.non_mapping_structure_lr))
    
    optimizer = torch.optim.Adam(optimizer_groups)
    scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lambda epoch: 1.0)
    
    ## start training
    trainer = PreFLMRTrainer(
        model=model,
        args=training_args,
        train_dataset=dataset,
        data_collator=dataset.collate_fn,
        optimizers=[optimizer, scheduler]
    )
    trainer.train()
    
    trainer.save_state()
    trainer.save_model()
    query_tokenizer.save_pretrained(os.path.join(training_args.output_dir, 'query_tokenizer'))
    context_tokenizer.save_pretrained(os.path.join(training_args.output_dir, 'context_tokenizer'))


if __name__ == '__main__':
    main()

from glob import glob
import pickle

## load results
test_results = []
for path in glob('./unzipped/generate_test_index_test_InfoseekDatasetForDPR*'):
    test_results.extend(pickle.load(open(path, 'rb'))['output'])
print('len(test_results)', len(test_results))

## load test passages(which is identical to train passages)
test_passages = pd.read_parquet('./data/multi_task_multi_modal_knowledge_retrieval_benchmark_M2KR/Infoseek_passages/test_passages-00000-of-00001.parquet')
test_passages.set_index('passage_id', inplace=True)

## load test set(which contains `pos_item_ids` that are golden labels)
test = pd.read_parquet('./data/multi_task_multi_modal_knowledge_retrieval_benchmark_M2KR/Infoseek_data/test-00000-of-00001.parquet')
test.set_index('question_id', inplace=True)

## calculate result
topk = 5

n_hit = 0
for result in test_results:
    question_id = result['question_id']
    pos_item_ids = test.loc[question_id]['pos_item_ids']
    
    top_ranking_passage_ids = [ obj['passage_id'] for obj in result['top_ranking_passages']]
    
    if any(pos_item_id in top_ranking_passage_ids[:topk] 
           for pos_item_id in pos_item_ids):
        n_hit += 1

r_at_n = n_hit / len(test_results)
print(f'R@{topk}:', r_at_n)