import os
import torch
from PIL import Image
import open_clip
import sys
import os
import argparse
import json
from torch.utils.data import DataLoader
from torch.utils.data.distributed import DistributedSampler
from tqdm import tqdm
from sklearn.metrics.pairwise import cosine_similarity
from data import HarvardDataset, HarvardVQADataset, IUXrayDataset, IUXrayVQADataset, MimicDataset, MimicVQADataset

def compute_cosine_sim(a, b):
    return float(cosine_similarity(a, b)[0][0])

def retrieve_topk_per_image(logits, val_k_list, retrieve_threshold=''):
    
    if not retrieve_threshold:
        pred_per_val_image = []
        for i, k in enumerate(val_k_list):
            if k == 0:
                pred_per_val_image.append(torch.tensor([-1]))
            elif k == 1:
                pred_per_val_image.append(
                    logits["weighted"][i].argmax(dim=0, keepdim=True)
                )
            else:
                _, topk_indices = logits["weighted"][i].topk(k, dim=0)
                pred_per_val_image.append(topk_indices)
        return pred_per_val_image
    else:
        retrieve_threshold=float(retrieve_threshold)
        pred_per_val_image = []
        logits_per_val_image = []
        for i, k in enumerate(val_k_list):
            if k == 0:
                pred_per_val_image.append(torch.tensor([-1]))
            else:
                logit_values = logits["weighted"][i]
                top1_logit = logit_values.max()
                sorted_logits, sorted_indices = torch.sort(logit_values, descending=True)
                
                # If k is more than the chosen topn, just choose topn
                if sorted_indices.size(0) > k:
                    sorted_indices = sorted_indices[:k]
                    sorted_logits = sorted_logits[:k]
                    
                pred_per_val_image.append(sorted_indices)
                logits_per_val_image.append(sorted_logits)
        return pred_per_val_image, logits_per_val_image


def get_logits(img_ft, question_ft, target_img_ft, target_text_ft, logit_scale, args):

    question2text = (question_ft @ target_text_ft.t()).detach().cpu()
    image2text = (img_ft @ target_text_ft.t()).detach().cpu()
    img2img = (img_ft @ target_img_ft.t()).detach().cpu()
    weighted = question2text * args.alpha + image2text * args.beta + img2img * args.gamma

    logits = {
        "question2text": question2text,
        "image2text": image2text,
        "img2img": img2img,
        "weighted": weighted
    }
    return logits


def clean_data_info(data_info):
    cleaned_data_info = {}
    for key, value in data_info.items():
        if isinstance(value, torch.Tensor):
            if value.numel() == 1:
                cleaned_data_info[key] = value.item()
        elif isinstance(value, list) and len(value) == 1:
            cleaned_data_info[key] = value[0]
        else:
            cleaned_data_info[key] = value
    return cleaned_data_info


def split_and_clean_data_infos(batch_data_infos):
    cleaned_data_infos = []
    num_items = len(next(iter(batch_data_infos.values())))
    for i in range(num_items):
        single_data_info = {}
        for key, value in batch_data_infos.items():
            if isinstance(value, torch.Tensor):
                single_data_info[key] = value[i]
            elif isinstance(value, list) and len(value) == num_items:
                single_data_info[key] = value[i]
            else:
                single_data_info[key] = value
        cleaned_data_infos.append(clean_data_info(single_data_info))
    return cleaned_data_infos

def prepare_data(args, preprocess, tokenizer):
    Dataset = None
    VQADataset = None
    if 'iu_xray' in args.img_root.split('/'):
        Dataset = IUXrayDataset
        VQADataset = IUXrayVQADataset
    elif 'Harvard-FairVLMed' in args.img_root.split('/'):
        Dataset = HarvardDataset
        VQADataset = HarvardVQADataset
    elif 'mimic-cxr-jpg' in args.img_root.split('/'):
        Dataset = MimicDataset
        VQADataset = MimicVQADataset

    train_dataset = Dataset(
        args.img_root,
        args.train_json,
        preprocess,
        tokenizer,
        load_include_path=True,
    )
    eval_dataset = VQADataset(
        args.img_root,
        args.eval_jsonl,
        preprocess,
        tokenizer,
        # test=(args.eval_type=='test'),
        fixed_K=args.fixed_k,
    )

    train_dataloader = DataLoader(
        train_dataset,
        batch_size=32,
        shuffle=False,
        num_workers=4,
        pin_memory=True,
        sampler=None,
        drop_last=False,
    )
    train_dataloader.num_samples = len(train_dataset)
    train_dataloader.num_batches = len(train_dataloader)

    eval_dataloader = DataLoader(
        eval_dataset,
        batch_size=32,
        shuffle=False,
        num_workers=4,
        pin_memory=True,
        sampler=None,
        drop_last=False,
    )
    eval_dataloader.num_samples = len(eval_dataset)
    eval_dataloader.num_batches = len(eval_dataloader)

    return train_dataloader, eval_dataloader


def main(args):
    device = "cuda" if torch.cuda.is_available() else "cpu"
    input_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
    model, _, preprocess = open_clip.create_model_and_transforms(
        args.model_name_or_path,
        pretrained=args.checkpoint_path,
    )
    model = model.to(device)
    model.eval()
    tokenizer = open_clip.get_tokenizer(args.model_name_or_path)
    train_dataloader, eval_dataloader = prepare_data(args, preprocess, tokenizer)
    val_all_image_features, val_all_text_features, val_all_question_features, val_k_list = [], [], [], []
    train_all_image_features, train_all_text_features = [], []
    train_all_image_full_paths, val_all_image_full_paths = [], []
    data_infos_list = []

    with torch.no_grad(), torch.cuda.amp.autocast():
        for batch in tqdm(train_dataloader, desc="Extracting traininig features"):
            images, texts, image_full_paths = batch
            images = images.to(device=device, dtype=input_dtype, non_blocking=True)
            texts = texts.to(device=device, non_blocking=True)
            train_all_image_full_paths.extend(image_full_paths)

            model_out = model(images, texts)
            image_features, text_features, logit_scale = model_out
            train_all_image_features.append(image_features.cpu())
            train_all_text_features.append(text_features.cpu())
            logit_scale = logit_scale.mean()

        for batch in tqdm(eval_dataloader, desc="Extracting validation features"):
            images, texts, image_full_paths, retrieval_ks, data_infos, questions = batch
            cleaned_data_infos = split_and_clean_data_infos(data_infos)
            val_all_image_full_paths.extend(image_full_paths)

            images = images.to(device=device, dtype=input_dtype, non_blocking=True)
            merged_texts = torch.cat([texts, questions], dim=0)  # shape: [2B, L]
            merged_texts = merged_texts.to(device, non_blocking=True)

            model_out = model(images, merged_texts)
            image_features, merged_text_features, logit_scale = model_out

            text_features = merged_text_features[:texts.size(0)] 
            question_features = merged_text_features[texts.size(0):]

            val_all_image_features.append(image_features.cpu())
            val_all_text_features.append(text_features.cpu())
            val_all_question_features.append(question_features.cpu())
            logit_scale = logit_scale.mean()
            data_infos_list.extend(cleaned_data_infos)
            retrival_ks_list = [int(t) for t in retrieval_ks]
            
            val_k_list.extend(retrival_ks_list)
    print(f'val image feature length: {len(torch.cat(val_all_image_features))}')
    print(f'train text feature length: {len(torch.cat(train_all_text_features))}')

    logits = get_logits(
        img_ft=torch.cat(val_all_image_features),
        question_ft=torch.cat(val_all_question_features),
        target_img_ft=torch.cat(train_all_image_features),
        target_text_ft=torch.cat(train_all_text_features),
        logit_scale=logit_scale.cpu(),
        args=args
    )
    
    if args.fixed_k:
        val_k_list = [int(args.fixed_k)]*len(val_k_list)
        print(f'fixed k: {args.fixed_k}')

    pred_per_val_image, logits_per_val_image = retrieve_topk_per_image(logits, val_k_list, retrieve_threshold=args.clip_threshold)
    true_val_k_list = [len(indices) for indices in pred_per_val_image]

    print(f'length of pred_per_val_image: {len(pred_per_val_image)}')
    print(f'length of data_infos_list: {len(data_infos_list)}')
    print(f'length of val_k_list: {len(val_k_list)}')

    output_jsonl = args.output_path
    with open(output_jsonl, "w") as f:
        for topk_indices, data_info, k, topk_logits in zip(
            pred_per_val_image, data_infos_list, true_val_k_list, logits_per_val_image
        ):
            reference_reports = []
            
            for idx in topk_indices:
                if idx.item() == -1:
                    
                    break
                reference_reports.append(
                    train_dataloader.dataset.image_report_pairs[idx.item()][1]
                )
            
            data_info["reference_reports"] = reference_reports
            data_info["retrieve_k"] = k
            data_info["logits"] = topk_logits.tolist()
            f.write(json.dumps(data_info))
            f.write("\n")
        print(f"Finished writing referenced reports to {output_jsonl}.")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Retrieve Top-K Predictions for Images")

    parser.add_argument("--img_root", type=str, default="path/to/harvard/image/root")
    parser.add_argument("--train_json", type=str, default="MR-RAG/data/training/retriever/harvard_train.json")
    parser.add_argument("--eval_jsonl", type=str, default="MR-RAG/data/test/harvard_test.jsonl")
    parser.add_argument("--output_path", type=str, default="MR-RAG/output_retrieve/test.jsonl")

    parser.add_argument("--alpha", type=float, default=0)
    parser.add_argument("--beta", type=float, default=1)
    parser.add_argument("--gamma", type=float, default=0)

    parser.add_argument("--fixed_k", type=str, default='30')
    parser.add_argument("--model_name_or_path", type=str, default="RN50")
    parser.add_argument("--checkpoint_path", type=str, default="MR-RAG/clip/open_clip_pytorch_model.bin")
    parser.add_argument("--config_type", type=str, default="tokenProb")
    parser.add_argument("--eval_type", type=str, default='')
    args = parser.parse_args()
    main(args)
