import os
import time
import ujson as json
import torch
import sys
import pickle
import numpy as np
from torch.utils.data import Dataset
import torch.distributed as dist
import torch.nn.functional as F

from bootleg.symbols.alias_entity_table import AliasEntityTable
from bootleg.symbols.constants import *
from bootleg.prep import prep_data
from bootleg.utils import logging_utils, data_utils, train_utils
from bootleg.utils.utils import import_class
from bootleg.utils import utils

# https://github.com/pytorch/pytorch/issues/37581#issuecomment-624516586
import warnings
warnings.filterwarnings("ignore", message=".*The given NumPy array is not writeable.*")

class WikiDataset(Dataset):
    """
    Main dataset class that handles preparing a batch of input.

    Things to note
    **Input is a sentence with mentions that are both true and false golds. A true gold is one that was directly
    mined with Wikipedia. A false gold is one that was generated by weak labelling.
    **We determine entities that are in a slice by if the true entity index is -1 or not. During train, if use_weak_label is true,
    we allow the model to leverage true and false golds. During eval, we only score true enchors.
    **Some embeddings require more expensive processing. E.g., extracting the pairs of candidate entities that are connected
    in a KG. When this processing is done in the dataloader where is can benefit from multiple dataloader threads,
    the embedding is stored in batch_on_the_fly. This embedding must have a batch_prep method
    When this processing is done during data prep, the embedding is stored in batch_prep.
    **If training a NIL model, we support randomly removing the true entity from the candidate list and setting the true
    entity index to be the NIL entity.
    **We support data slices (subsets of data) for both training (if slice model) and eval. If using slices for training model,
    we supports probabilistic slice indices.

    Attributes:
        batch_prepped_emb_file_names: embedding that are batch prepped in advance
        batch_on_the_fly_embs: embedding where the batch_prep method is called in the __get_item__ method
        random_nil: whether to do NIL candidate random generation

    Batch Inputs:
        start_idx_in_sent: first token index of a mention,
        end_idx_in_sent: last token index of a mention,
        alias_idx: the alias (mention) index in our alias dictionary,
        word_indices: word indexes into the word emeddings (e.g., BERT token indices),
        sent_idx: unique sentence index,
        subsent_idx: unique subsentence index in the case of sentence windowing,
        entity_indices: the entity indices in our entity dictionary,
        alias_list_pos: keeps track of the original alias position in the list of all aliases in case the sentence
                        is split via windowing
        true_entity_idx_for_train: entity indices for true and false golds, as seen during train
        slice_indices (optional): if slice dataset, we pass in matrix where each row is alias and each column
                                  is 0/1 if that mention is in the slice or not
        <ind_task_name> (option): probabilistic labels of if an mention is in a slice or not (used in slicing model)
        <pred_task_name>: NED prediction labels; for slice model, predictions of aliases not in the slice are masked
        <embs>: all batch prep or batch on the fly emeddings
    """
    def __init__(self, args, use_weak_label, input_src, dataset_name,
                 is_writer, distributed, word_symbols, entity_symbols,
                 slice_dataset=None, dataset_is_eval=False):
        # Need to save args to reinstantiate logger
        self.args = args
        self.logger = logging_utils.get_logger(args)
        # Number of candidates, including NIL if a NIL model (train_in_candidates is False)
        self.K = entity_symbols.max_candidates + (not args.data_config.train_in_candidates)
        self.num_entities_with_pad_and_nocand = entity_symbols.num_entities_with_pad_and_nocand
        self.dataset_name = dataset_name
        self.slice_dataset = slice_dataset
        self.dataset_is_eval = dataset_is_eval
        # Slice names used for eval slices and a slicing model
        self.slice_names = train_utils.get_data_slices(args, dataset_is_eval)
        self.storage_type_file = data_utils.get_storage_file(self.dataset_name)
        # Mappings from sent_idx to row_id in dataset
        self.sent_idx_file = os.path.splitext(dataset_name)[0] + "_sent_idx.json"
        self.type_pred = False
        if args.data_config.type_prediction.use_type_pred:
            self.type_pred = True
            self.eid2typeid, self.num_types_with_pad = self.load_coarse_type_table(args, entity_symbols)
        # Load memory mapped file
        self.logger.info("Loading dataset...")
        self.logger.debug("Seeing if " + dataset_name + " exists")
        if (args.data_config.overwrite_preprocessed_data or
            (not os.path.exists(self.dataset_name)) or
            (not os.path.exists(self.sent_idx_file)) or
            (not os.path.exists(self.storage_type_file)) or
            (not os.path.exists(data_utils.get_batch_prep_config(self.dataset_name)))):
            start = time.time()
            self.logger.debug(f"Building dataset with {input_src}")
            # Only prep data once per node
            if is_writer:
                prep_data(args, use_weak_label=use_weak_label, dataset_is_eval=self.dataset_is_eval,
                input_src=input_src, dataset_name=dataset_name,
                prep_dir=data_utils.get_data_prep_dir(args))
            if distributed:
                # Make sure all processes wait for data to be created
                dist.barrier()
            self.logger.debug(f"Finished building and saving dataset in {round(time.time() - start, 2)}s.")

        start = time.time()

        # Storage type for loading memory mapped file of dataset
        self.storage_type = pickle.load(open(self.storage_type_file, 'rb'))

        self.data = np.memmap(self.dataset_name, dtype=self.storage_type, mode='r')
        self.data_len = len(self.data)

        # Mapping from sentence idx to rows in the dataset (indices).
        # Needed when sampling sentence indices from slices for evaluation.
        sent_idx_to_idx_str = utils.load_json_file(self.sent_idx_file)
        self.sent_idx_to_idx = {int(i):val for i,val in sent_idx_to_idx_str.items()}
        self.logger.info(f"Finished loading dataset.")

        # Stores info about the batch prepped embedding memory mapped files and their shapes and datatypes
        # so we can load them
        self.batch_prep_config = utils.load_json_file(data_utils.get_batch_prep_config(self.dataset_name))
        self.batch_prepped_emb_files = {}
        self.batch_prepped_emb_file_names = {}
        for emb in args.data_config.ent_embeddings:
            if 'batch_prep' in emb and emb['batch_prep']:
                assert emb.key in self.batch_prep_config, f'Need to prep {emb.key}. Please call prep instead of run with batch_prep_embeddings set to true.'
                self.batch_prepped_emb_file_names[emb.key] = os.path.join(os.path.dirname(self.dataset_name),
                    os.path.basename(self.batch_prep_config[emb.key]['file_name']))
                self.batch_prepped_emb_files[emb.key] = np.memmap(
                    self.batch_prepped_emb_file_names[emb.key],
                    dtype=self.batch_prep_config[emb.key]['dtype'],
                    shape=tuple(self.batch_prep_config[emb.key]['shape']),
                    mode='r')
                assert len(self.batch_prepped_emb_files[emb.key]) == self.data_len,\
                    f'Preprocessed emb data file {self.batch_prep_config[emb.key]["file_name"]} does not match length of main data file.'

        # Stores embeddings that we compute on the fly; these are embeddings where batch_on_the_fly is set to true.
        self.batch_on_the_fly_embs = {}
        for emb in args.data_config.ent_embeddings:
            if 'batch_on_the_fly' in emb and emb['batch_on_the_fly'] is True:
                mod, load_class = import_class("bootleg.embeddings", emb.load_class)
                try:
                    self.batch_on_the_fly_embs[emb.key] = getattr(mod, load_class)(main_args=args,
                        emb_args=emb['args'], entity_symbols=entity_symbols,
                        model_device=None, word_symbols=None, key=emb.key)
                except AttributeError as e:
                    self.logger.warning(f'No prep method found for {emb.load_class} with error {e}')
                except Exception as e:
                    print("ERROR", e)
        # The data in this table shouldn't be pickled since we delete it in the class __getstate__
        self.alias2entity_table = AliasEntityTable(args=args, entity_symbols=entity_symbols)
        # Random NIL percent
        self.mask_perc = args.train_config.random_nil_perc
        self.random_nil = False
        # Don't want to random mask for eval
        if not dataset_is_eval:
            # Whether to use a random NIL training regime
            self.random_nil = args.train_config.random_nil
            if self.random_nil:
                self.logger.info(f'Using random nils during training with {self.mask_perc} percent')

    def __len__(self):
        return self.data_len

    def __getitem__(self, key):
        # start = time.time()
        example = self.data[key]
        entity_indices = self.alias2entity_table(example['alias_idx'])
        # True entities will be true and false golds for train (if use_weak_label in config is true) and just true golds for eval
        true_entities = torch.from_numpy(example['true_entity_idx'])
        M = true_entities.shape
        if self.random_nil:
            # example['true_entity_idx'] is  M -> we want to sample some % of these and set them to not in candidate list
            # randomly mask each entity embedding
            bern_prob = (torch.ones(M) * self.mask_perc)
            keep_mask = torch.bernoulli(bern_prob) < 1
            # whichever we sample, we want to set corresponding true candidate to -1 and mask it out
            # to simulate not being in the candidate list
            # can't have negatives for one hot so we temporarily cast padded values to 0
            padded_entities = true_entities == -1
            true_entities = true_entities.masked_fill(padded_entities, 0)
            one_hot_true_entities = F.one_hot(true_entities, num_classes=self.K)
            one_hot_true_entities[keep_mask.unsqueeze(-1).expand_as(one_hot_true_entities)] = 0
            one_hot_true_entities[padded_entities.unsqueeze(-1).expand_as(one_hot_true_entities)] = 0
            entity_indices = entity_indices.masked_fill(one_hot_true_entities, -1)
            # set new true label to 0 ('not in candidate')
            true_entities = true_entities.masked_fill(~keep_mask, 0)
            # make sure original padded entities are padded
            true_entities = true_entities.masked_fill(padded_entities, -1)

        start_idx_in_sent = example['start_idx_in_sent']
        end_idx_in_sent = example['end_idx_in_sent']
        example_dict = {'start_idx_in_sent': start_idx_in_sent,
                        'end_idx_in_sent': end_idx_in_sent,
                        'alias_idx': example['alias_idx'],
                        'word_indices': example['word_indices'],
                        'sent_idx': example['sent_idx'],
                        'subsent_idx': example['subsent_idx'],
                        'entity_indices': entity_indices,
                        # due to subsentence split, we need to keep track of the original alias position in the list
                        # to do eval over slices when distributed
                        # (examples from a sentence may be distributed across different GPUs)
                        'alias_list_pos': example['alias_list_pos'],
                        # true entities of the mentions seen during train (true and false golds); in eval, we only keep
                        # true entities of true golds
                        'true_entity_idx_for_train': example['true_entity_idx_for_train']}


        # If this dataset is associated with slices, slice_indices is a incidence matrix indicating
        # for each alias in the batch, which ones participate in which slice (slices keep track of sentence indexes and aliases to predict)
        # Slices are not windowed like that are for training data.
        if self.slice_dataset is not None:
            # -1 is pad and should not be in the mapping from sentence index to row in array.
            assert -1 != self.slice_dataset.sent_idx_arr[example["sent_idx"]]
            # One row per mention and one column per slice
            slice_indices = np.hstack([self.slice_dataset.data[slice_name][self.slice_dataset.sent_idx_arr[example["sent_idx"]]].alias_to_predict.T
                              for slice_name in self.slice_names])
            prob_labels_arr = np.hstack([self.slice_dataset.data[slice_name][self.slice_dataset.sent_idx_arr[example["sent_idx"]]].prob_labels.T
                              for slice_name in self.slice_names])
            # alias_list_pos will have -1 for no alias; we want these to become zero in slice_indices.
            # Therefore we add a pad row to the bottom of slice_indices
            slice_indices = np.vstack([slice_indices, np.zeros(slice_indices.shape[1])]).astype(int)
            slice_indices = slice_indices[example['alias_list_pos']]
            # Probabilistic slice labels for slice indicator head training
            prob_labels_arr = np.vstack([prob_labels_arr, np.zeros(prob_labels_arr.shape[1])]).astype(float)
            prob_labels_arr = prob_labels_arr[example['alias_list_pos']]

            # If this is an eval dataset, keep slice indices intact for eval_wrapper
            example_dict['slice_indices'] = slice_indices
            # Assign true entity idx to -1 if example alias doesn't participate in slice
            for i, slice_name in enumerate(self.slice_names):
                prob_labels = prob_labels_arr[:,i]
                bin_in_slice_labels = slice_indices[:,i]

                # NED prediction labels; set predictions to be -1 for masking for mentions not in a slice
                pred_labels = np.copy(true_entities)
                pred_labels[~(bin_in_slice_labels).astype(bool)] = -1

                # Mask out slice alias labels for which we don't want to make a prediction
                # We need to use true_entity_idx to account for subsentences which indicate
                # which alias to predict
                prob_labels[true_entities == -1] = -1

                ind_task_name = train_utils.get_slice_head_ind_name(slice_name)
                pred_task_name = train_utils.get_slice_head_pred_name(slice_name)

                # Add indicator head and prediction head labels
                example_dict[ind_task_name] = prob_labels
                example_dict[pred_task_name] = pred_labels

        else:
            example_dict[train_utils.get_slice_head_pred_name(FINAL_LOSS)] = example['true_entity_idx']
        # Add type preds
        if self.type_pred:
            example_dict["type_labels"] = self.eid2typeid[true_entities]
        # Add embeddings to example forward
        for emb_name in self.batch_prepped_emb_files:
            example_dict[emb_name] = np.asarray(self.batch_prepped_emb_files[emb_name][key])
        # Prep the embeddings (this will call the batch_prep method for the embedding)
        for emb_name, emb in self.batch_on_the_fly_embs.items():
            example_dict[emb_name] = emb.batch_prep(example['alias_idx'], entity_indices)
        return example_dict

    def __getstate__(self):
        state = self.__dict__.copy()
        # Not picklable
        del state['data']
        del state['logger']
        # the sent_idx mapping is expensive to pickle so remove
        # also not needed in dataloader workers so we don't need to setstate for it
        del state['sent_idx_to_idx']
        del state['batch_prepped_emb_files']
        return state

    def __setstate__(self, state):
        self.__dict__.update(state)
        self.data = np.memmap(self.dataset_name, dtype=self.storage_type, mode='r')
        self.batch_prepped_emb_files = {}
        for emb_name, file_name in self.batch_prepped_emb_file_names.items():
            self.batch_prepped_emb_files[emb_name] = np.memmap(self.batch_prepped_emb_file_names[emb_name],
                    dtype=self.batch_prep_config[emb_name]['dtype'],
                    shape=tuple(self.batch_prep_config[emb_name]['shape']),
                    mode='r')
        self.logger = logging_utils.get_logger(self.args)

    def __repr__(self):
        return f"Dataset {self.dataset_name}"

    def load_coarse_type_table(self, args, entity_symbols):
        emb_dir = args.data_config.emb_dir
        coarse_type_file = args.data_config.type_prediction.file
        with open(os.path.join(emb_dir, coarse_type_file)) as in_f:
            # take the first type; UNK type is 0
            qid2type = {}
            max_type = 0
            for k, v in json.load(in_f).items():
                if len(v) > 0:
                    qid2type[k] = v[0]+1
                else:
                    qid2type[k] = 0
                max_type = max(max_type, qid2type[k])
        # We assume types are indexed from 0. So, 6 types will have indices 0 - 5. Max type will get 5+1 = 6.
        assert max_type == args.data_config.type_prediction.num_types,\
            f"{args.data_config.type_prediction.num_types} from args.data_config.type_prediction.num_types must match our computed number {max_type}"
        # All qids get unk types
        values = [0 for _ in range(self.num_entities_with_pad_and_nocand)]
        for qid in qid2type:
            if entity_symbols.qid_exists(qid):
                values[entity_symbols.get_eid(qid)] = qid2type[qid]
        # Padded eid gets -1
        values[-1] = -1
        num_types_with_pad = max_type+1
        eid2coarsetype = torch.tensor(values)
        return eid2coarsetype, num_types_with_pad

torch.nn.functional.one_hot torch.from_numpy torch.ones torch.tensor torch.bernoulli torch.distributed.barrier