import argparse

import dataclasses

import json

import os

from enum import auto, Enum

from pathlib import Path

from typing import List, Any

import random

import numpy as np

import pandas as pd

import torch

from omegaconf import OmegaConf

from torch.utils.data import Dataset, DataLoader

from tqdm import tqdm

from transformers import AutoTokenizer

from torch.utils.data.sampler import Sampler

from data.instruct_tasks import create_direct_task_data, create_cp_task_data, create_correction_task_data, create_nle_task_data

from local_config import VIS_ROOT, PATH_TO_MIMIC_CXR

from model.lavis.models.blip2_models.modeling_llama_imgemb import LlamaForCausalLM

class MyReportProcessor():

    def __init__(self, prompt="", max_words=50, prompt_neg=""):

        self.prompt = prompt

        self.max_words = max_words

        self.prompt_neg = prompt_neg

    def __call__(self, findings, no_labels=False):

        prompt = self.prompt

        if no_labels:

            findings = "no common findings"  # cannot write which findings as we don't no them

        prompt = prompt.format(findings=findings)

        return prompt

    @classmethod

    def from_config(cls, cfg=None):

        if cfg is None:

            cfg = OmegaConf.create()

        prompt = cfg.get("prompt", "")

        max_words = cfg.get("max_words", 50)

        return cls(prompt=prompt, max_words=max_words)

class SeparatorStyle(Enum):

    """Different separator style."""

    SINGLE = auto()

    TWO = auto()

@dataclasses.dataclass

class Conversation:

    """A class that keeps all conversation history."""

    system: str

    roles: List[str]

    messages: List[List[str]]

    offset: int

    sep_style: SeparatorStyle = SeparatorStyle.SINGLE

    sep: str = "###"

    sep2: str = None

    # Used for gradio server

    skip_next: bool = False

    conv_id: Any = None

    def get_prompt(self):

        if self.sep_style == SeparatorStyle.SINGLE:

            ret = self.system

            for role, message in self.messages:

                if message:

                    ret += self.sep + " " + role + ": " + message

                else:

                    ret += self.sep + " " + role + ":"

            return ret

        elif self.sep_style == SeparatorStyle.TWO:

            seps = [self.sep, self.sep2]

            ret = self.system + seps[0]

            for i, (role, message) in enumerate(self.messages):

                if message:

                    ret += role + ": " + message + seps[i % 2]

                else:

                    ret += role + ":"

            return ret

        else:

            raise ValueError(f"Invalid style: {self.sep_style}")

    def append_message(self, role, message):

        self.messages.append([role, message])

    def dict(self):

        return {

            "system": self.system,

            "roles": self.roles,

            "messages": self.messages,

            "offset": self.offset,

            "sep": self.sep,

            "sep2": self.sep2,

            "conv_id": self.conv_id,

        }

def create_conv():

    conv = Conversation(

        system="A chat between a curious user and an artificial intelligence assistant acting as an experienced radiologist. "

               "The assistant gives professional, detailed, and polite answers to the user's questions.",

        roles=["USER", "ASSISTANT"],

        messages=[],

        offset=0,

        sep_style=SeparatorStyle.TWO,

        sep=" ",

        sep2="</s>",

    )

    return conv

class MIMIC_Text_Dataset(Dataset):

    def __init__(self, split, truncate=None, prompt_type="basic", use_indication=False):

        super().__init__()

        # load csv file

        self.split = pd.read_csv(

            f'{PATH_TO_MIMIC_CXR}/mimic-cxr-jpg/2.0.0/mimic-cxr-2.0.0-split.csv')

        self.reports = pd.read_csv('mimic-cxr/reports_processed/mimic_cxr_sectioned.csv')

        # drop reports where findings are nan

        self.reports = self.reports.dropna(subset=['findings'])

        self.img_ids = {img_id: i for i, img_id in enumerate(self.reports['dicom_id'])}

        self.chexpert = pd.read_csv(f'data/data_files/finding_chexbert_labels.csv')

        self.chexpert_cols = ["No Finding", "Enlarged Cardiomediastinum",

                              "Cardiomegaly", "Lung Opacity",

                              "Lung Lesion", "Edema",

                              "Consolidation", "Pneumonia",

                              "Atelectasis", "Pneumothorax",

                              "Pleural Effusion", "Pleural Other",

                              "Fracture", "Support Devices"]

        self.use_indication = use_indication

        self.vis_root = VIS_ROOT

        self.prompt_type = prompt_type

        self.split_ids = set(self.split.loc[self.split['split'] == split]['dicom_id'])

        self.train_ids = set(self.split.loc[self.split['split'] == 'train']['dicom_id'])

        # get all dicom_ids where "split" is split

        self.annotation = self.reports.loc[self.reports['dicom_id'].isin(self.split_ids)]

        if truncate is not None:

            self.annotation = self.annotation[:truncate]

        self.annotation['findings'] = self.annotation['findings'].apply(lambda x: x.replace('\n', ''))

        # Extract patient_id from Img_Folder (3rd part) and study_id is the name of the notefile without the pre-pending 's'

        self.annotation['subject_id'] = self.annotation['Img_Folder'].apply(lambda x: int(x.split('/')[2].lstrip('p')))

        self.annotation['study_id'] = self.annotation['Note_file'].apply(lambda x: int(x.lstrip('s').rstrip('.txt')))

        # Merge chexpert labels with annotation dataframe

        self.annotation = pd.merge(self.annotation, self.chexpert, how='left', left_on=['dicom_id'],

                                   right_on=['dicom_id'])

        # for every row add a string of comma-separated positive labels

        self.annotation['positive_labels'] = self.annotation.apply(lambda x: self.convert_to_finding_labels(x[self.chexpert_cols].values,

                                                                                                            self.chexpert_cols), axis=1)

        # maybe use transforms from here: ResNet50_Weights.IMAGENET1K_V2.transforms

        # read prompt from json

        prompts = json.loads(Path(f"vicuna_prompts.json").read_text(encoding="UTF-8"))

        self.text_processor = MyReportProcessor(

            prompt=prompts[prompt_type], max_words=1000,

            prompt_neg=prompts[prompt_type.replace("matching_examples", "neg_matching_examples")])

    def convert_to_finding_labels(self, chexpert_labels, columns, label=1):

        # Get indices where value is 1

        indices = np.where(chexpert_labels == label)

        # Get the corresponding column names and join them into a string

        labels = ", ".join([columns[i] for i in indices[0]])

        return labels

    def __getitem__(self, index):

        ann = self.annotation.iloc[index]

        # if self.use_indication:

        #     indication = self.indications[study_id]

        #     if indication == "":

        #         indication = "Indication not given."

        caption = ann["findings"].strip()

        chexpert_labels = ann[self.chexpert_cols].astype(float).values

        chexpert_label_str = ann["positive_labels"]

        dicom_id = ann["dicom_id"]

        # check if all columns are in (nan, 0) -> no labels

        no_labels = np.all((np.isnan(chexpert_labels)) | (chexpert_labels == 0) | (chexpert_labels == -1.))

        finding_string = chexpert_label_str.lower().strip()

        input_text = self.text_processor(findings=finding_string, no_labels=no_labels)

        # if self.use_indication:

        #     input_text = "Indication: " + indication + " " + input_text

        # template for vicuna v1.3

        conv = Conversation(

            system="A chat between a curious user and an artificial intelligence assistant acting as an experienced radiologist. "

                   "The assistant gives professional, detailed, and polite answers to the user's questions.",

            roles=["USER", "ASSISTANT"],

            messages=[],

            offset=0,

            sep_style=SeparatorStyle.TWO,

            sep=" ",

            sep2="</s>",

        )

        conv.append_message(conv.roles[0], input_text)

        conv.append_message(conv.roles[1], None)

        prompt = conv.get_prompt()

        return {

            "text_input": prompt,

            "text_target": caption,

            "ig_label_string": finding_string,

            "chexpert_labels": chexpert_labels,

            "chexpert_cols": self.chexpert_cols,

            "dicom": dicom_id,

            "img_path": ann["Img_Folder"] + "/" + ann["Img_Filename"],

        }

    def __len__(self):

        return len(self.annotation)

class SubsetSampler(Sampler):

    def __init__(self, indices):

        self.indices = indices

    def __iter__(self):

        return (self.indices[i] for i in range(len(self.indices)))

    def __len__(self):

        return len(self.indices)

def stratified_sample(df, simulated_epochs=1):

    # We want to reduce the number of examples with no finding to 1/14th of the dataset. We achieve this easily by first seperating the dataset into 2 groups: no finding and finding.

    # either no finding, or nothing is considered a no finding

    no_findings_indices = df.annotation[((df.annotation['No Finding'] == 1) | ((df.annotation[df.chexpert_cols] == 1).sum(1) == 0) == 1)].index

    finding_indices = df.annotation.index.difference(no_findings_indices)

    no_findings_indices = no_findings_indices.tolist()

    finding_indices = finding_indices.tolist()

    # we are striving to lose as little no_finding data as possible. So instead of just reducing the number of no_finding examples, we will increase the number of finding examples. Just clone and extend dataset

    finding_indices = finding_indices * simulated_epochs

    # subsample the no finding examples to be 1/14th of the new dataset

    new_dataset_size = len(finding_indices) * 14 / 13

    new_no_finding_count = int(new_dataset_size / 14)

    # merge considering the new dataset size

    all_indices = finding_indices + random.sample(no_findings_indices, new_no_finding_count)

    return all_indices

def create_report_data_vicuna_specific_stratified(prompt_type):

    val_dataset = MIMIC_Text_Dataset(split="train", truncate=None, prompt_type=prompt_type)

    stratified_indices = stratified_sample(val_dataset, simulated_epochs=2)

    sampler = SubsetSampler(stratified_indices)

    data_loader = DataLoader(val_dataset, batch_size=200, num_workers=200, sampler=sampler)

    report_jsons = []

    for _, batch in tqdm(enumerate(data_loader)):

        # iterate over batch elements

        for i in range(len(batch["text_input"])):

            text_input = batch["text_input"][i]

            text_target = batch["text_target"][i]

            dicom = batch["dicom"][i]

            # sample random prompt for every report

            reports_json = {

                "instruction": text_input,

                "input": "",

                "output": text_target,

                "dicom": dicom,

            }

            report_jsons.append(reports_json)

    # Save the JSON data to a file

    with open("data/data_files/mimic_cxr_reports_stratified.json", "w") as f:

        json.dump(report_jsons, f, ensure_ascii=False, indent=4)

'''

this method saves instruct data jsons for all the different tasks we defined:

- easy language: EL DONE

- correction: CO DONE

- summerization: SU DONE

- reasoning: RE (based on MIMIC-NLE) DONE

- region QA: RQA DONE

- CP binary QA: CPbQA DONE

- CP all QA: CPaQA DONE

for every report we sample one task and one prompt and save the report, the question (task) and the answer generated by vicuna (or from dataset groundtruth)

'''

def create_report_data_vicuna_instruct_large():

    lang_model = LlamaForCausalLM.from_pretrained("lmsys/vicuna-13b-v1.3", torch_dtype=torch.float16, device_map='auto', load_in_8bit=False)

    tokenizer = AutoTokenizer.from_pretrained("lmsys/vicuna-13b-v1.3", use_fast=False, truncation_side="left", padding_side="left")

    tokenizer.pad_token = tokenizer.unk_token

    val_dataset = MIMIC_Text_Dataset(split="train", truncate=None, prompt_type="img_matching_examples_ig2_noexamples")

    # split in 6 portions of 1/6th each, randomly

    split_size = len(val_dataset) // 6

    remainder = len(val_dataset) % 6

    val_dataset_EL, _, val_dataset_SU, val_dataset_EX, val_dataset_RQA, val_dataset_CPQA = torch.utils.data.random_split(val_dataset,

                                                                                                                         [split_size + (i < remainder)

                                                                                                                          for i in range(

                                                                                                                             6)])  # correction is samples somewhere else

    # split val_dataset_CPQA in 2

    split_size = len(val_dataset_CPQA) // 2

    remainder = len(val_dataset_CPQA) % 2

    val_dataset_CPbQA, val_dataset_CPaQA = torch.utils.data.random_split(val_dataset_CPQA, [split_size + (i < remainder) for i in range(2)])

    # create directory

    if not os.path.exists("data/large_instruct_data"):

        os.makedirs("data/large_instruct_data")

    # create data

    create_direct_task_data(lang_model, tokenizer, val_dataset_EL, task_name="EL")

    create_direct_task_data(lang_model, tokenizer, val_dataset_SU, task_name="SU")

    create_direct_task_data(lang_model, tokenizer, val_dataset_RQA, task_name="RQA")

    create_cp_task_data(val_dataset_CPbQA, task_name="CPbQA")

    create_cp_task_data(val_dataset_CPaQA, task_name="CPaQA")

    create_correction_task_data(lang_model, tokenizer)

    create_nle_task_data()

'''

fuse instruct data with report generation task into one dataset json

'''

def fuse_instruct_dataset(prompt_type="img_matching_examples_ig2_noexamples_IMG_findings"):

    # get report generation data

    val_dataset = MIMIC_Text_Dataset(split="train", truncate=None, prompt_type=prompt_type)

    stratified_indices = stratified_sample(val_dataset, simulated_epochs=2)

    sampler = SubsetSampler(stratified_indices)

    data_loader = DataLoader(val_dataset, batch_size=200, sampler=sampler, num_workers=200)

    report_jsons = []

    for _, batch in tqdm(enumerate(data_loader)):

        # iterate over batch elements

        for i in range(len(batch["text_input"])):

            text_input = batch["text_input"][i]

            text_target = batch["text_target"][i]

            dicom = batch["dicom"][i]

            # sample random prompt for every report

            reports_json = {

                "instruction": text_input,

                "input": "",

                "output": text_target,

                "dicom": dicom,

            }

            report_jsons.append(reports_json)

    task_jsons = []

    with open(f"vicuna_prompts.json", "r") as f:

        prompts = json.load(f)

    report_prompt = prompts[prompt_type]

    # get instruct data

    for task in ["EL", "RE", "CO", "SU", "RQA", "CPbQA", "CPaQA"]:

        print("Creating data for " + task)

        with open(f"data/large_instruct_data/instruct_large_{task}.json", "r") as f:

            task_data = json.load(f)

        for elem in tqdm(task_data):

            report = elem["gt_report"] if task != "CO" else elem["incorrect_report"]

            conv = create_conv()

            conv.append_message(conv.roles[0], report_prompt)

            conv.append_message(conv.roles[1], report)

            conv.append_message(conv.roles[0], elem["task"])

            conv.append_message(conv.roles[1], None)

            instruction = conv.get_prompt()

            # get elem directly from val_dataset.train_annotation with same dicom

            orig_elem = val_dataset.annotation[val_dataset.annotation["dicom_id"] == elem["dicom"]].iloc[0]

            if type(orig_elem['positive_labels']) == float and np.isnan(orig_elem['positive_labels']):

                finding_str = "no common findings"

            else:

                finding_str = orig_elem['positive_labels'].lower().strip()

            instruction = instruction.format(findings=finding_str)

            task_json = {

                "instruction": instruction,

                "input": "",

                "output": elem["output"].lower().strip() if task == "CPaQA" else elem["output"].strip(),

                "dicom": elem["dicom"],

            }

            task_jsons.append(task_json)

    # combine and shuffle report and task jsons

    combined_jsons = report_jsons + task_jsons

    random.shuffle(combined_jsons)

    # save to json

    with open(f"data/data_files/mimic_cxr_instruct_stratified.json", "w") as f:

        json.dump(combined_jsons, f, indent=4)

if __name__ == '__main__':

    # args parser

    parser = argparse.ArgumentParser()

    parser.add_argument('--mode', type=str, default='RG', help='RG or INS')

    args = parser.parse_args()

    ''' Create data to train RaDialog-RG model'''

    if args.mode == 'RG':

        create_report_data_vicuna_specific_stratified(prompt_type="img_matching_examples_ig2_noexamples_IMG_findings")

    ''' Create data to train RaDialog-INS model'''

    if args.mode == 'INS':

        create_report_data_vicuna_instruct_large()

        fuse_instruct_dataset()

    # This code is meant for understanding how our instruct dataset is created.

    # Due to randomness in the sampling and model predictions, a newly generated dataset could be slightly different.

    # To exactly reproduce our results, please use the instruct dataset we published and use 'fuse_instruct_dataset' to merge with your MIMIC data.