import logging

import numpy as np

from medcat.utils.matutils import unitvec

from datetime import datetime

import math

import datasets

import random

import copy

def get_all_splits(dataset):

    all_datasets = []

    if 'train' in dataset:

        all_datasets.append(dataset['train'])

    if 'test' in dataset:

        all_datasets.append(dataset['test'])

    if 'valid' in dataset:

        all_datasets.append(dataset['valid'])

    if isinstance(dataset, datasets.arrow_dataset.Dataset):

        # If we have only one, ie no train/test

        all_datasets.append(dataset)

    return all_datasets

def make_example(token, ent_example, token_type='unk', cnt=10**6, time=None, cntx=None):

    out = {'token': token, 'token_type': token_type, 'cnt': cnt, 'time': time}

    if 'context_representation' in ent_example:

        if cntx is None:

            cntx = [0.0 for i in range(len(ent_example['context_representation']))]

        out['context_representation'] = cntx

    return out

def get_duration_separator(separator, start_time, current_time, bucket_size_seconds):

    d_separator = separator

    for i in [1, 7]:

        if (current_time - start_time) >= bucket_size_seconds * i:

            d_separator = f'{separator[0:-1]}-{i}{separator[-1]}'

    return d_separator

def bucket_concepts(examples, bucket_size_seconds=365*24*60*60, separator='<SEP>', duration_separator=False):

    r''' Will bucket concepts into specified bucket_size.

    Args:

        examples

    '''

    for i in range(len(examples['stream'])):

        stream = examples['stream'][i]

        new_stream = []

        _bucket = []

        _tokens = set()

        start_time = -1

        for ent in stream:

            if start_time == -1:

                start_time = ent['time']

            if ent['time'] - start_time >= bucket_size_seconds:

                # Add to stream

                new_stream.extend(_bucket)

                _bucket = []

                _tokens = set()

                if separator is not None:

                    _separator = separator

                    if duration_separator:

                        # This will have different separator for different time spans

                        _separator = get_duration_separator(separator, start_time, ent['time'], bucket_size_seconds)

                    # A separator is +1 of the last token in the stream

                    new_stream.append(make_example(ent_example=ent, token=_separator, token_type='sep', cnt=10**6, time=new_stream[-1]['time']+1))

                # Change start time to current entity time

                start_time = ent['time']

            if ent['token'] not in _tokens:

                _bucket.append(ent)

                _tokens.add(ent['token'])

        if _bucket:

            new_stream.extend(_bucket)

        examples['stream'][i] = new_stream

        new_stream = []

    return examples

def add_position_ids(examples, separators=set()):

    for i in range(len(examples['stream'])):

        stream = examples['stream'][i]

        old_t = None

        cnt = 0

        for ent in stream:

            ent['position_ids'] = cnt

            if ent['token'] in separators:

                cnt += 1

    return examples

def add_age(examples, pt2dob_timestamp, age_prefix='<AGE>', age_suffix=None, age_normalizer=365.25 * 24 * 60 * 60):

    for i in range(len(examples['stream'])):

        stream = examples['stream'][i]

        last_age_added = -1

        new_stream = []

        for ent in stream:

            if examples['patient_id'][i] in pt2dob_timestamp:

                if pt2dob_timestamp is not None:

                    age = int((ent['time'] - pt2dob_timestamp[examples['patient_id'][i]]) / age_normalizer)

                # Age comes a step before the token that caused the change

                if age >= 0 and last_age_added != age:

                    if age_prefix is not None:

                        new_stream.append(make_example(ent_example=ent, token=age_prefix, token_type='age_prefix', cnt=10**6, time=ent['time']))

                    new_stream.append(make_example(ent_example=ent, token=str(age), token_type='age', cnt=10**6, time=ent['time']))

                    last_age_added = age

                    if age_suffix is not None:

                        new_stream.append(make_example(ent_example=ent, token=age_suffix, token_type='age_suffx', cnt=10**6, time=ent['time']))

            new_stream.append(ent)

        examples['stream'][i] = new_stream

        new_stream = []

    return examples

def add_ttd(examples, pt2dod_timestamp, ttd_prefix='<TTD>', ttd_suffix=None, ttd_normalizer=365.25 * 24 * 60 * 60,

            max_ttd=10, ttd_prob=1, max_nttd=10, duplicate_streams=False):

    all_patient_id = []

    all_stream = []

    for i in range(len(examples['stream'])):

        stream = examples['stream'][i]

        last_ttd_added = -1

        new_stream = []

        new_streams = [new_stream]

        n_added_ttds = 0

        for ent in stream:

            if examples['patient_id'][i] in pt2dod_timestamp:

                if n_added_ttds < max_nttd:

                    if random.random() <=  ttd_prob:

                        ttd = int((pt2dod_timestamp[examples['patient_id'][i]] - ent['time']) / ttd_normalizer) + 1

                        if ttd <= max_ttd:

                            if last_ttd_added != ttd:

                                if duplicate_streams:

                                    # At this point we duplicate the first stream fron new_streams (it is the one without TTD always)

                                    new_stream = copy.deepcopy(new_streams[0])

                                    new_streams.append(new_stream)

                                if ttd_prefix is not None:

                                    new_stream.append(make_example(ent_example=ent, token=ttd_prefix, token_type='ttd_prefix', cnt=10**6, time=ent['time']))

                                new_stream.append(make_example(ent_example=ent, token=str(ttd), token_type='ttd', cnt=10**6, time=ent['time']))

                                last_ttd_added = ttd

                                if ttd_suffix is not None:

                                    new_stream.append(make_example(ent_example=ent, token=ttd_suffix, token_type='ttd_suffix', cnt=10**6, time=ent['time']))

                                n_added_ttds += 1

            # append the entity to each stream

            for new_stream in new_streams: new_stream.append(ent)

        if duplicate_streams and len(new_streams) > 1:

            # Remove the first example as it is the base one without time info

            del new_streams[0]

        for new_stream in new_streams:

            all_stream.append(new_stream)

            all_patient_id.append(examples['patient_id'][i])

    examples['patient_id'] = all_patient_id

    examples['stream'] = all_stream

    return examples

def split_stream(examples, max_seq_len=-1):

    if max_seq_len > 0:

        new_streams = []

        new_patient_ids = []

        for ind, stream in enumerate(examples['stream']):

            nparts = math.ceil(len(stream) / max_seq_len)

            for i in range(nparts):

                new_streams.append(stream[i*max_seq_len:(i+1)*max_seq_len])

                new_patient_ids.append(examples['patient_id'][ind])

        examples['stream'] = new_streams

        examples['patient_id'] = new_patient_ids

    return examples

def cleanup_stream(examples, keep_time=True, keep_type=True, keep_position_ids=True, keep_context_representation=True):

    r''' Leave only Tokens and remove the rest from `stream`

    Args:

        examples

        keep_time:

            If set another value will be added to examples that contains the `time` for each

            entity in stream.

        keep_type:

            Same as above

    '''

    if 'token' in examples['stream'][0][0]:

        if keep_time:

            examples['time'] = [[ent['time'] for ent in stream] for stream in examples['stream']]

        if keep_type:

            examples['token_type'] = [[ent['token_type'] for ent in stream] for stream in examples['stream']]

        if keep_position_ids:

            examples['position_ids'] = [[ent['position_ids'] for ent in stream] for stream in examples['stream']]

        if keep_context_representation:

            examples['context_representation'] = [[ent['context_representation'] for ent in stream] for stream in examples['stream']]

        examples['stream'] = [[ent['token'] for ent in stream] for stream in examples['stream']]

    return examples

def add_to_stream(examples, pt2tkn, last=False, prefix=None, unk_tkn='unk', token_type='unk'):

    r''' Add information to the patient stream based on patient_id.

    Args:

        examples

        pt2tkn

        last

        unk_tkn:

            What token will be added if the patient_id is not in pt2tkn

    '''

    for i in range(len(examples['stream'])):

        ent = examples['stream'][i][0]

        if examples['patient_id'][i] in pt2tkn:

            token = pt2tkn.get(examples['patient_id'][i], unk_tkn)

            t_ind = -1 if last else 0 # If -1 means it is the last token, otherwise the first

            to_append = [make_example(ent_example=ent, token=token, cnt=10**6, time=examples['stream'][i][t_ind]['time'], token_type=token_type)]

            if prefix is not None:

                prefix_token = make_example(ent_example=ent, token=prefix, cnt=10**6,

                                            time=examples['stream'][i][t_ind]['time'], token_type="prefix_" + token_type)

                to_append = [prefix_token] + to_append

            if last:

                # Append as last token

                examples['stream'][i] = examples['stream'][i] + to_append

            else:

                examples['stream'][i] = to_append + examples['stream'][i]

    return examples

def remove_tokens_not_in_tokenizer(examples, tokens_to_keep):

    tokens_to_keep = set(tokens_to_keep)

    for i in range(len(examples['stream'])):

        stream = examples['stream'][i]

        new_stream = []

        for ent in stream:

            tkn = ent['token']

            if tkn in tokens_to_keep:

                new_stream.append(ent)

        examples['stream'][i] = new_stream

    return examples

def remove_parents_from_stream(examples, ch2parents, separator=None, separators=None):

    for i in range(len(examples['stream'])):

        stream = examples['stream'][i]

        parents = set()

        new_stream = []

        for ent in stream:

            tkn = ent['token']

            if (separator is not None and tkn == separator) or (separators is not None and tkn in separators):

                # This means we are removing parents only inside of one bucket

                parents = set()

            if tkn in ch2parents:

                # Add only if not in parents

                if tkn not in parents:

                    new_stream.append(ent)

                # Update parents

                parents.update(ch2parents[tkn])

            else:

                new_stream.append(ent)

        examples['stream'][i] = new_stream

    return examples

def get_embeddings_for_tokens(dataset=None, cdb=None, context_type='medium', normalize=True, extra_tokens=['<PAD>'], types=None, concepts=None):

    r''' Given a stream of tokens get the embeddings from MedCAT and make the required maps.

    Args:

        dataset

        cdb

        context_type

        normalize:

            If True the embedding vectors will be normalized

        tkn2type:

            Dictionary mapping from token to type

        types:

            All posible token types (e.g. [T-11, T-12, ...]

        concepts:

            If provided these concepts will also be appened to the tokens and supported by the tokenizer

    Returns:

        embeddings

        tkn2id

        id2tkn

        id2type

        id2type_detailed

    '''

    embeddings = []

    tkn2id = {}

    id2tkn = {}

    def add_tkn(tkn):

        if tkn in cdb.cui2context_vectors and context_type in cdb.cui2context_vectors[tkn]:

            vec = cdb.cui2context_vectors[tkn][context_type]

        else:

            # Token vector is randomly assigned

            vec = np.random.rand(300)

        id2tkn[len(embeddings)] = tkn

        tkn2id[tkn] = len(embeddings)

        vec = unitvec(vec) if normalize else vec

        embeddings.append(vec)

    datasets = get_all_splits(dataset)

    for _dataset in datasets:

        for stream in _dataset['stream']:

            for tkn in stream:

                tkn = str(tkn)

                if tkn not in tkn2id:

                    add_tkn(tkn)

    # Add concepts if they are provided, this is used to build a general

    #tokenizer with all concepts

    if concepts is not None:

        for concept in concepts:

            tkn = str(concept)

            if tkn not in tkn2id:

                add_tkn(tkn)

    # Add named tokens

    for tkn in extra_tokens:

        if tkn not in tkn2id:

            id2tkn[len(embeddings)] = tkn

            tkn2id[tkn] = len(embeddings)

            if tkn != '<PAD>':

                embeddings.append(np.random.rand(len(embeddings[0])))

            else:

                embeddings.append(np.zeros(len(embeddings[0])))

    # Add type tokens

    for tkn in types:

        if tkn not in tkn2id:

            id2tkn[len(embeddings)] = tkn

            tkn2id[tkn] = len(embeddings)

            embeddings.append(np.random.rand(len(embeddings[0])))

    return embeddings, tkn2id, id2tkn

def stream_to_separate_examples(examples):

    r''' Convert a stream to separate examples that can be used to train

    a next concept predictor unable to handle sequences (e.g. random forset). Use with HF datasets map function.

    '''

    out = {}

    out['input_ids'] = [input_ids[0:i+1] for input_ids in examples['input_ids'] for i in range(len(input_ids) - 1)]

    out['labels'] = [input_ids[i+1] for input_ids in examples['input_ids'] for i in range(len(input_ids) - 1)]

    out['labels_all'] = [input_ids[i+1:] for input_ids in examples['input_ids'] for i in range(len(input_ids) - 1)]

    out['patient_id'] = [patient_id for ind, patient_id in enumerate(examples['patient_id']) for _ in range(len(examples['input_ids'][ind]) - 1)]

    return out