import random

import numpy as np

import pandas as pd

import pickle

import sys

import torch

import time

import re

from torch.utils.data import Dataset

from collections import defaultdict

from project_utils import read_patients

import project_config

class PatientDataset(Dataset):

    def __init__(self, filepath, gp_spl=None, raw_data=False, mondo_map_file=str(project_config.PROJECT_DIR / 'mondo_references.csv'), needs_disease_mapping=False, time=False): 

        self.filepath = filepath

        self.patients = read_patients(filepath)

        print('Dataset filepath: ', filepath)

        print('Number of patients: ', len(self.patients))

        # add placeholder for true genes/diseases if they don't exist

        for patient in self.patients:

            if 'true_genes' not in patient: patient['true_genes'] = []

            if 'true_diseases' not in patient: patient['true_diseases'] = []

        self.raw_data = raw_data

        self.needs_disease_mapping = needs_disease_mapping

        self.time = time

        # create HPO to node_idx map

        with open(project_config.KG_DIR / f'hpo_to_idx_dict_{project_config.CURR_KG}.pkl', 'rb') as handle:

            self.hpo_to_idx_dict = pickle.load(handle)

        with open(project_config.KG_DIR / f'hpo_to_name_dict_{project_config.CURR_KG}.pkl', 'rb') as handle:

            self.hpo_to_name_dict = pickle.load(handle)

        self.idx_to_hpo_dict = {v:self.hpo_to_name_dict[k] if k in self.hpo_to_name_dict else k for k, v in self.hpo_to_idx_dict.items()}

        # create ensembl to node_idx map

        # NOTE: assumes conversion from gene symbols to ensembl IDs has already occurred

        with open(str(project_config.KG_DIR  / f'ensembl_to_idx_dict_{project_config.CURR_KG}.pkl'), 'rb') as handle:

            self.ensembl_to_idx_dict = pickle.load(handle)

        self.idx_to_ensembl_dict = {v:k for k, v in self.ensembl_to_idx_dict.items()}

        # orphanet to mondo disease map

        with open(str(project_config.PROJECT_DIR / 'preprocess' / 'orphanet' / 'orphanet_to_mondo_dict.pkl'), 'rb') as handle:

            self.orpha_mondo_map = pickle.load(handle)

        with open(project_config.KG_DIR / f'mondo_to_idx_dict_{project_config.CURR_KG}.pkl', 'rb') as handle:

            self.disease_to_idx_dict = pickle.load(handle)

        with open(project_config.KG_DIR / f'mondo_to_name_dict_{project_config.CURR_KG}.pkl', 'rb') as handle:

            self.disease_to_name_dict = pickle.load(handle)

        self.idx_to_disease_dict = {v:self.disease_to_name_dict[k] if k in self.disease_to_name_dict else k for k, v in self.disease_to_idx_dict.items()}

        # degree dict from idx to degree - used for debugging

        #NOTE: may need to subtract 1 to index into this dict

        with open(str(project_config.KG_DIR / f'degree_dict_{project_config.CURR_KG}.pkl'), 'rb') as handle:

            self.degree_dict = pickle.load(handle)

        # get patients with similar genes

        if all(['true_genes' in patient for patient in self.patients]): # first check to make sure all patients have true genes

            genes_to_patients = defaultdict(list)

            for patient in self.patients:

                for g in patient['true_genes']:

                    genes_to_patients[g].append(patient['id'])

            self.patients_with_same_gene = defaultdict(list)

            for patients in genes_to_patients.values():

                for p in patients:

                    self.patients_with_same_gene[p].extend([pat for pat in patients if pat != p])

        # get patients with similar diseases

        if all(['true_diseases' in patient for patient in self.patients]): # first check to make sure all patients have true diseases

            dis_to_patients = defaultdict(list)

            for patient in self.patients:

                patient_diseases = patient['true_diseases']

                for d in patient_diseases: 

                    dis_to_patients[d].append(patient['id'])

            self.patients_with_same_disease = defaultdict(list)

            for patients in dis_to_patients.values():

                for p in patients:

                    self.patients_with_same_disease[p].extend([pat for pat in patients if pat != p])

        # map from patient id to index in dataset

        self.patient_id_to_index = {p['id']:i for i, p in enumerate(self.patients)}

        print('Finished initalizing dataset')

    def __len__(self):

        '''

        Returns the length of the dataset

        '''

        return len(self.patients)

    def __getitem__(self, idx):

        '''

        Returns a single example from the dataset

        '''

        t0 = time.time()

        patient = self.patients[idx]

        additional_labels_dict = {}

        if 'additional_labels' in patient:

            for label, values in patient['additional_labels'].items():

                if label == "n_hops_cand_g_p": continue

                if values == None: values = [[-1]]

                if type(values) != list: values = [[values]] #  wrap in list if needed

                if type(values) == list and (len(values)==0 or type(values[0]) != list): values = [values]

                additional_labels_dict[label] = values 

            if 'max_percent_phen_overlap_train' not in patient['additional_labels']: additional_labels_dict['max_percent_phen_overlap_train'] = [[-1]] 

            if 'max_phen_overlap_train' not in patient['additional_labels']: additional_labels_dict['max_phen_overlap_train'] = [[-1]]

        phenotype_node_idx = [self.hpo_to_idx_dict[p] for p in patient['positive_phenotypes'] if p in self.hpo_to_idx_dict ]

        correct_genes_node_idx = [self.ensembl_to_idx_dict[g] for g in patient['true_genes'] if g in self.ensembl_to_idx_dict ]

        if 'all_candidate_genes' in patient:

            candidate_gene_node_idx = [self.ensembl_to_idx_dict[g] for g in patient['all_candidate_genes'] if g in self.ensembl_to_idx_dict ]

        else: candidate_gene_node_idx = []

        if 'true_diseases' in patient:

            if self.needs_disease_mapping:

                orpha_diseases = [ int(d) if len(re.match("^[0-9]*", d)[0]) > 0 else d for d in patient['true_diseases']]

                mondo_diseases = [mondo_d for orpha_d in set(orpha_diseases).intersection(set(self.orpha_mondo_map.keys())) for mondo_d in self.orpha_mondo_map[orpha_d]]

            else:

                mondo_diseases = [str(d) for d in patient['true_diseases']]

            disease_node_idx = [self.disease_to_idx_dict[d] for d in mondo_diseases if d in self.disease_to_idx_dict]

        else: disease_node_idx = None

        if not self.raw_data:

            phenotype_node_idx = torch.LongTensor(phenotype_node_idx)

            correct_genes_node_idx = torch.LongTensor(correct_genes_node_idx)

            candidate_gene_node_idx = torch.LongTensor(candidate_gene_node_idx)

            if 'true_diseases' in patient: disease_node_idx = torch.LongTensor(disease_node_idx)

        assert len(phenotype_node_idx) >= 1, f'There are no phenotypes for patient: {patient}'

        #NOTE: assumes that patient has a single causal/correct gene (the model still outputs a score for each candidate gene)

        if not self.raw_data:

            if len(correct_genes_node_idx) > 1:

                #print('NOTE: The patient has multiple correct genes, but we\'re only selecting the first.')

                correct_genes_node_idx = correct_genes_node_idx[0].unsqueeze(-1)

        # get index of correct gene

        if len(correct_genes_node_idx) == 0: # no correct genes available for the patient

            label_idx = None

        else:

            if self.raw_data:

                label_idx = [candidate_gene_node_idx.index(g) for g in correct_genes_node_idx]

            else:

                label_idx = (candidate_gene_node_idx == correct_genes_node_idx[0]).nonzero(as_tuple=True)[0]

        if self.time:

            t1 = time.time()

            print(f'It takes {t1-t0:0.4f}s to get an item from the dataset')

        return (phenotype_node_idx, candidate_gene_node_idx, correct_genes_node_idx, disease_node_idx, label_idx, additional_labels_dict, patient['id'])

    def node_idx_to_degree(self, idx):

        if idx in self.degree_dict:

            return self.degree_dict[idx]

        else:

            return -1

    def node_idx_to_name(self, idx):

        if idx in self.idx_to_hpo_dict:

            return self.idx_to_hpo_dict[idx]

        elif idx in self.idx_to_ensembl_dict:

            return self.idx_to_ensembl_dict[idx]

        elif idx in self.idx_to_disease_dict:

            return self.idx_to_disease_dict[idx]

        elif idx == 0:

            return 'padding'

        else:

            print("Exception on:", idx)

            raise Exception

    def get_similar_patients(self, patient_id, similarity_type='gene'):

        if similarity_type == 'gene':

            sim_pats = np.array(self.patients_with_same_gene[patient_id])

            np.random.shuffle(sim_pats)

            return sim_pats

        elif similarity_type == 'disease':

            sim_pats = np.array(self.patients_with_same_disease[patient_id])

            np.random.shuffle(sim_pats)

            return sim_pats

        else:

            raise NotImplementedError

    def get_candidate_diseases(self, cand_type='all_kg_nodes'):

        if cand_type == 'all_kg_nodes':

            all_kg_diseases_idx = np.unique(list(self.disease_to_idx_dict.values())) # get idx of all diseases in KG

            return torch.LongTensor(all_kg_diseases_idx)

        elif cand_type == 'orphanet':

            orpha_mondo_diseases = [d[0] for d in list(self.orpha_mondo_map.values())]

            orpha_mondo_idx = np.unique([self.disease_to_idx_dict[d] for d in orpha_mondo_diseases if d in self.disease_to_idx_dict])

            return torch.LongTensor(orpha_mondo_idx)

        else:

            raise NotImplementedError