import os

import numpy as np
from keras.callbacks import Callback, TensorBoard, EarlyStopping, ModelCheckpoint


def get_callbacks(log_dir=None, valid=(), tensorboard=True, eary_stopping=True):
    """Get callbacks.

    Args:
        log_dir (str): the destination to save logs(for TensorBoard).
        valid (tuple): data for validation.
        tensorboard (bool): Whether to use tensorboard.
        eary_stopping (bool): whether to use early stopping.

    Returns:
        list: list of callbacks
    """
    callbacks = []
    print("Callback log dir {}".format(log_dir))
    if log_dir and tensorboard:
        if not os.path.exists(log_dir):
            print('Successfully made a directory: {}'.format(log_dir))
            os.mkdir(log_dir)
        callbacks.append(TensorBoard(log_dir))

    if valid:
        callbacks.append(F1score(*valid))

    if log_dir:
        if not os.path.exists(log_dir):
            print('Successfully made a directory: {}'.format(log_dir))
            os.mkdir(log_dir)

        file_name = '_'.join(['model_weights', '{epoch:02d}', '{f1:2.2f}']) + '.h5'
        save_callback = ModelCheckpoint(os.path.join(log_dir, file_name),
                                        monitor='f1',
                                        save_weights_only=True)
        callbacks.append(save_callback)

    if eary_stopping:
        callbacks.append(EarlyStopping(monitor='f1', patience=3, mode='max'))

    return callbacks


def get_entities(seq):
    """Gets entities from sequence.

    Args:
        seq (list): sequence of labels.

    Returns:
        list: list of (chunk_type, chunk_start, chunk_end).

    Example:
        >>> seq = ['B-PER', 'I-PER', 'O', 'B-LOC']
        >>> print(get_entities(seq))
        [('PER', 0, 2), ('LOC', 3, 4)]
    """
    i = 0
    chunks = []
    seq = seq + ['O']  # add sentinel
    types = [tag.split('-')[-1] for tag in seq]
    while i < len(seq):
        if seq[i].startswith('B'):
            for j in range(i+1, len(seq)):
                if seq[j].startswith('I') and types[j] == types[i]:
                    continue
                break
            chunks.append((types[i], i, j))
            i = j
        else:
            i += 1
    return chunks


def f1_score(y_true, y_pred, sequence_lengths):
    """Evaluates f1 score.

    Args:
        y_true (list): true labels.
        y_pred (list): predicted labels.
        sequence_lengths (list): sequence lengths.

    Returns:
        float: f1 score.

    Example:
        >>> y_true = []
        >>> y_pred = []
        >>> sequence_lengths = []
        >>> print(f1_score(y_true, y_pred, sequence_lengths))
        0.8
    """
    correct_preds, total_correct, total_preds = 0., 0., 0.
    for lab, lab_pred, length in zip(y_true, y_pred, sequence_lengths):
        lab = lab[:length]
        lab_pred = lab_pred[:length]

        lab_chunks = set(get_entities(lab))
        lab_pred_chunks = set(get_entities(lab_pred))

        correct_preds += len(lab_chunks & lab_pred_chunks)
        total_preds += len(lab_pred_chunks)
        total_correct += len(lab_chunks)

    p = correct_preds / total_preds if correct_preds > 0 else 0
    r = correct_preds / total_correct if correct_preds > 0 else 0
    f1 = 2 * p * r / (p + r) if correct_preds > 0 else 0
    return f1

def __startOfChunk(prevTag, tag, prevTagType, tagType, chunkStart=False):
    if prevTag == 'B' and tag == 'B':
        chunkStart = True
    if prevTag == 'I' and tag == 'B':
        chunkStart = True
    if prevTag == 'O' and tag == 'B':
        chunkStart = True
    if prevTag == 'O' and tag == 'I':
        chunkStart = True

    if prevTag == 'E' and tag == 'E':
        chunkStart = True
    if prevTag == 'E' and tag == 'I':
        chunkStart = True
    if prevTag == 'O' and tag == 'E':
        chunkStart = True
    if prevTag == 'O' and tag == 'I':
        chunkStart = True

    if tag != 'O' and tag != '.' and prevTagType != tagType:
        chunkStart = True
    return chunkStart

def __endOfChunk(prevTag, tag, prevTagType, tagType, chunkEnd=False):
    if prevTag == 'B' and tag == 'B':
        chunkEnd = True
    if prevTag == 'B' and tag == 'O':
        chunkEnd = True
    if prevTag == 'I' and tag == 'B':
        chunkEnd = True
    if prevTag == 'I' and tag == 'O':
        chunkEnd = True

    if prevTag == 'E' and tag == 'E':
        chunkEnd = True
    if prevTag == 'E' and tag == 'I':
        chunkEnd = True
    if prevTag == 'E' and tag == 'O':
        chunkEnd = True
    if prevTag == 'I' and tag == 'O':
        chunkEnd = True

    if prevTag != 'O' and prevTag != '.' and prevTagType != tagType:
        chunkEnd = True
    return chunkEnd

def __splitTagType(tag):
    s = tag.split('-')
    if len(s) > 2 or len(s) == 0:
        raise ValueError('tag format wrong. it must be B-xxx.xxx')
    if len(s) == 1:
        tag = s[0]
        tagType = ""
    else:
        tag = s[0]
        tagType = s[1]
    return tag, tagType

def computeF1ScoreCONLL(correct_slots, pred_slots):

    print("Evaluating F1 score of {} number of predictions ".format(len(correct_slots)))
    correctChunk = {}
    correctChunkCnt = 0
    foundCorrect = {}
    foundCorrectCnt = 0
    foundPred = {}
    foundPredCnt = 0
    correctTags = 0
    tokenCount = 0
    for correct_slot, pred_slot in zip(correct_slots, pred_slots):
        inCorrect = False
        lastCorrectTag = 'O'
        lastCorrectType = ''
        lastPredTag = 'O'
        lastPredType = ''
        for c, p in zip(correct_slot, pred_slot):
            correctTag, correctType = __splitTagType(c)
            predTag, predType = __splitTagType(p)

            if inCorrect == True:
                if __endOfChunk(lastCorrectTag, correctTag, lastCorrectType, correctType) == True and \
                        __endOfChunk(lastPredTag, predTag, lastPredType, predType) == True and \
                        (lastCorrectType == lastPredType):
                    inCorrect = False
                    correctChunkCnt += 1
                    if lastCorrectType in correctChunk:
                        correctChunk[lastCorrectType] += 1
                    else:
                        correctChunk[lastCorrectType] = 1
                elif __endOfChunk(lastCorrectTag, correctTag, lastCorrectType, correctType) != \
                        __endOfChunk(lastPredTag, predTag, lastPredType, predType) or \
                        (correctType != predType):
                    inCorrect = False

            if __startOfChunk(lastCorrectTag, correctTag, lastCorrectType, correctType) == True and \
                    __startOfChunk(lastPredTag, predTag, lastPredType, predType) == True and \
                    (correctType == predType):
                inCorrect = True

            if __startOfChunk(lastCorrectTag, correctTag, lastCorrectType, correctType) == True:
                foundCorrectCnt += 1
                if correctType in foundCorrect:
                    foundCorrect[correctType] += 1
                else:
                    foundCorrect[correctType] = 1

            if __startOfChunk(lastPredTag, predTag, lastPredType, predType) == True:
                foundPredCnt += 1
                if predType in foundPred:
                    foundPred[predType] += 1
                else:
                    foundPred[predType] = 1

            if correctTag == predTag and correctType == predType:
                correctTags += 1

            tokenCount += 1

            lastCorrectTag = correctTag
            lastCorrectType = correctType
            lastPredTag = predTag
            lastPredType = predType

        if inCorrect == True:
            correctChunkCnt += 1
            if lastCorrectType in correctChunk:
                correctChunk[lastCorrectType] += 1
            else:
                correctChunk[lastCorrectType] = 1

    print("foundPredCnt : {}   correctChunkCnt : {}   foundCorrect : {}".format(foundPredCnt, correctChunkCnt, foundCorrectCnt))
    if foundPredCnt > 0:
        precision = 100.0 * correctChunkCnt / foundPredCnt
    else:
        precision = 0

    if foundCorrectCnt > 0:
        recall = 100.0 * correctChunkCnt / foundCorrectCnt
    else:
        recall = 0

    if (precision + recall) > 0:
        f1 = (2 * precision * recall) / (precision + recall)
    else:
        f1 = 0

    return f1, precision, recall

class F1score(Callback):

    def __init__(self, valid_steps, valid_batches, preprocessor=None,verbose = False, mode="dev", raw_data=None, out_file_name=None):
        super(F1score, self).__init__()
        self.valid_steps = valid_steps
        self.valid_batches = valid_batches
        self.p = preprocessor
        self.f1s = []
        self.losses = []
        self.verbose = verbose
        self.mode = mode
        self.raw_data = raw_data
        self.out_file_name = out_file_name
        print("F1score is constructed")


    def on_train_begin(self, logs={}):
        return

    def on_train_end(self, logs={}):
        print(logs)

    def _get_sequence_length(self, seqs):
        seq_length = []
        for seq in seqs:
            #print(type(seq))
            if len(np.argwhere(seq == 0)) > 0 :
                seq_length.append(np.argwhere(seq == 0)[0][0])
            else:
                seq_length.append(len(seq))

        return seq_length

    def on_epoch_end(self, epoch, logs={}):
        correct_preds, total_correct, total_preds = 0., 0., 0.
        print("Evaluating")
        all_y_true = []
        all_y_pred = []

        for i, (data, label) in enumerate(self.valid_batches):
            if i == self.valid_steps:
                break
            y_true = label
            #print(y_true[0])
            y_true = np.argmax(y_true, -1)
            sequence_lengths = data[-1] # shape of (batch_size, 1)
            sequence_lengths = np.reshape(sequence_lengths, (-1,))
            #y_pred = np.asarray(self.model_.predict(data, sequence_lengths))
            #print("EPOCH END : TYPE : {} LEN : {} SHAPE : {} {}".format(type(data), len(data),data[0].shape,data[1].shape))
            y_pred = self.model.predict_on_batch(data)
            y_pred = np.argmax(y_pred, -1)
            #print("LEN of data : {}".format(len(data)))

            #print("TYPE INSIDE ITERATION {}:".format(type(data)))
            #print(data)
            #print("GOLD : {}".format(y_true[0]))
            #print("PREDICTION {}: ".format(y_pred[0]))
            #print("Type : {} Shape : {}".format(type(y_pred), y_pred.shape))
            #print("Sequence length shape : {}".format(sequence_lengths.shape))
            seq_lengths = np.reshape(self._get_sequence_length(y_pred), (-1,))
            #print("Sequence length shape NEW : {}".format(seq_lengths.shape))
            y_pred = [self.p.inverse_transform(y[:l]) for y, l in zip(y_pred, seq_lengths)]
            y_true = [self.p.inverse_transform(y[:l]) for y, l in zip(y_true, seq_lengths)]

            #print(data)
            #print(type(y_true))
            #print(y_true)
            #print(y_pred)
            #exit()
            #words  = [self.p.inverse_id_to_word(y[:l]) for y, l in zip(data[0], seq_lengths)]
            #print(words)
            #for idx, pred in enumerate(y_pred):
            #    for idy, pred_token in enumerate(pred):
            #        f.write("WORD\t{}\t{}\n".format(y_true[idx][idy], y_pred[idx][idy]))
            #    f.write("\n")
            a, b, c = self.count_correct_and_pred(y_true, y_pred, seq_lengths)
            correct_preds += a
            total_preds += b
            total_correct += c
            all_y_true = all_y_true + y_true
            all_y_pred = all_y_pred + y_pred

        f1, _ , _ = computeF1ScoreCONLL(all_y_true, all_y_pred)

        #f1 = self._calc_f1(correct_preds, total_correct, total_preds)


        print(' Evaluation - f1: {:04.2f}'.format(f1))
        self.f1s.append(f1)
        if self.out_file_name is not None :
            with open(self.out_file_name, "w")  as f:
                for i, sent in enumerate(self.raw_data):
                    pred = all_y_pred[i]
                    gold = all_y_true[i]
                    if type(sent[0]) is str :
                        for j, token in enumerate(sent):
                            f.write("{} {} {}\n".format(token, gold[j], pred[j]))
                        f.write("\n")
                    else :
                        for j, token in enumerate(sent[0]):
                            f.write("{} {} {}\n".format(token, gold[j], pred[j]))
                        f.write("\n")

            os.system("perl conlleval.pl < "+str(self.out_file_name))
        logs['f1'] = f1
        #print("On validation {} {}".format(logs['f1'], logs['loss']))
        #self.losses.append(logs['loss'])

        return

    def _calc_f1(self, correct_preds, total_correct, total_preds):

        p = correct_preds / total_preds if correct_preds > 0 else 0
        r = correct_preds / total_correct if correct_preds > 0 else 0
        f1 = 2 * p * r / (p + r) if correct_preds > 0 else 0
        return f1

    def count_correct_and_pred(self, y_true, y_pred, sequence_lengths):
        correct_preds, total_correct, total_preds = 0., 0., 0.
        #print(y_true)
        #print(y_pred)
        for lab, lab_pred, length in zip(y_true, y_pred, sequence_lengths):
            lab = lab[:length]
            lab_pred = lab_pred[:length]

            lab_chunks = set(get_entities(lab))
            lab_pred_chunks = set(get_entities(lab_pred))

            correct_preds += len(lab_chunks & lab_pred_chunks)
            total_preds += len(lab_pred_chunks)
            total_correct += len(lab_chunks)
        return correct_preds, total_correct, total_preds

