import logging

import os

import re

import pandas as pd

import torch

from torch.utils.data import Dataset

from src.utils import processed_data_path

class InstructionTuningDataset(Dataset):

    def __init__(self, split, source):

        assert split in ["train", "val", "test", "test_subset"]

        assert source in ["event", "note", "joint", "joint_all"]

        self.split = split

        self.source = source

        self.data_path = os.path.join(processed_data_path, f"mimic4")

        self.cohort = pd.read_csv(os.path.join(self.data_path, f"cohort_{split}.csv"))

        qa_note = pd.read_json(os.path.join(self.data_path, "qa_note.jsonl"), lines=True)

        qa_note = qa_note[qa_note.hadm_id.isin(self.cohort.hadm_id.unique())]

        qa_event = pd.read_json(os.path.join(self.data_path, f"qa_event.jsonl"), lines=True)

        qa_event = qa_event[qa_event.hadm_id.isin(self.cohort.hadm_id.unique())]

        if source == "note":

            qa = qa_note

        elif source == "event":

            qa = qa_event

        else:

            if source == "joint":

                logging.warning(f"Subsample event QA to {len(qa_note)}")

                qa_event = qa_event.sample(n=len(qa_note), replace=False, random_state=42)

            else:

                logging.warning(f"Use all event QA")

            qa = pd.concat([qa_note, qa_event], ignore_index=True)

        self.qa = qa

        logging.warning(f"Loaded {len(qa)} {source} QA samples for {split} on {event}")

    def _get_event_list(self, hadm_id):

        df = pd.read_csv(os.path.join(self.data_path, f"event_selected/event_{hadm_id}.csv"))

        event_list = []

        for i, row in df.iterrows():

            event_list.append((row.timestamp, row.event_type, row.event_value))

        return event_list

    def _get_event_emb(self, hadm_id):

        return torch.load(os.path.join(self.data_path, f"pt_event_selected_no_time_type/event_{hadm_id}.pt"))

    def __len__(self):

        return len(self.qa)

    @staticmethod

    def _extract_digits(event_tuple):

        timestamp, event_type, event_value = event_tuple

        try:

            if event_type == "patient demographics" or event_type == "patient_demographics":

                value_match = re.search(r"age:\s*([\d.]+)", event_value)

                if value_match:

                    value = float(value_match.group(1))

                else:

                    value = 0

                duration = 0

            elif event_type == "admission info" or event_type == "admission_info":

                value, duration = 0, 0

            elif event_type == "diagnoses_icd":

                value, duration = 0, 0

            elif event_type == "labevents":

                value_match = re.search(r":\s*([\d.]+)", event_value)

                if value_match:

                    value = float(value_match.group(1))

                else:

                    value = 0

                duration = 0

            elif event_type == "microbiologyevents":

                value, duration = 0, 0

            elif event_type == "prescriptions":

                value_match = re.search(r"prescribed dose:\s*([\d.]+)", event_value)

                if value_match:

                    value = float(value_match.group(1))

                else:

                    value = 0

                duration_match = re.search(r"duration:\s*([\d.]+)", event_value)

                if duration_match:

                    duration = float(duration_match.group(1))

                else:

                    duration = 0

            elif event_type == "transfers":

                value, duration = 0, 0

            elif event_type == "procedureevents":

                value = 0

                duration_match = re.search(r"for\s*([\d.]+)\s*hour", event_value)

                if duration_match:

                    duration = float(duration_match.group(1))

                else:

                    duration = 0

            else:

                raise ValueError(f"Unknown event type: {event_type}")

        except Exception as e:

            value, duration = 0, 0

            logging.warning(f"Error {e} in extracting digits from event tuple: {event_tuple}")

        return value, duration

    def __getitem__(self, index):

        data = self.qa.iloc[index]

        q = data["q"]

        a = data["a"]

        event_emb = self._get_event_emb(data["hadm_id"])

        num_events = event_emb.shape[0]

        event_list = self._get_event_list(data["hadm_id"])

        assert len(event_list) == num_events

        time_tensor = torch.tensor([[e[0]] for e in event_list], dtype=torch.float32)

        value_duration_tensor = torch.tensor([self._extract_digits(e) for e in event_list], dtype=torch.float32)

        event_emb = torch.cat(

            [

                event_emb,

                time_tensor,

                value_duration_tensor,

            ],

            dim=1

        )

        final_q = "\n".join(["<image>" * num_events, q])

        return final_q, a, event_emb