from sadedegel.dataset import load_raw_corpus, load_sentence_corpus, file_paths
from sadedegel.tokenize import NLTKPunctTokenizer, RegexpSentenceTokenizer, Doc
from sadedegel.bblock.util import flatten, is_eos
from sadedegel.ml import create_model, save_model
from sklearn.model_selection import cross_val_score

import click
from typing import List


@click.group(help="sadedeGel SBD commandline")
def tokenize():
    pass


def iou_eval(name: str, y_true: List[List[str]], y_pred: List[List[str]], paths=None):
    i, u, iou = 0, 0, 0

    for idx, (seq1, seq2) in enumerate(zip(y_true, y_pred)):
        s1 = list(map(lambda x: x.strip(), seq1))
        s2 = list(map(lambda x: x.strip(), seq2))

        i += len(set(s1) & set(s2))
        u += len(set(s1) | set(s2))

        delta = len(set(s1) & set(s2)) / len(set(s1) | set(s2))

        if paths:
            click.secho(f"Lost {1.0 - delta:.2f} because of {paths[idx]}", fg="yellow")

        iou += delta

    click.secho(f"IoU score for {name}")

    click.secho(f"Micro IoU: {i / u:.4f}".rjust(25))
    click.secho(f"Macro IoU: {iou / len(y_true):.4f}".rjust(25))


@tokenize.command()
@click.option("-v", count=True, help="verbosity")
def evaluate(v):
    """Evaluate IoU metric for different SBD algorithms over our stock dataset."""
    click.secho("Loading corpus...")
    raw, sents = load_raw_corpus(False), load_sentence_corpus()

    nltk = NLTKPunctTokenizer()
    reg = RegexpSentenceTokenizer()

    y_pred = [nltk(doc) for doc in raw]
    y_true = [doc['sentences'] for doc in sents]

    iou_eval("NLTKPunctTokenizer", y_true, y_pred, file_paths() if v > 0 else None)

    y_pred = [reg(doc) for doc in raw]

    iou_eval("RegexpSentenceTokenizer", y_true, y_pred, file_paths() if v > 0 else None)

    y_pred = [[s.text for s in Doc(doc).sents] for doc in raw]

    iou_eval("MLBasedTokenizer", y_true, y_pred, file_paths() if v > 0 else None)


@tokenize.command()
def diff():
    """Git like diff tool to compare sentence generated by our tokenizer vs actual list of sentences."""
    click.secho("Loading corpus...")
    raw, sents = load_raw_corpus(False), load_sentence_corpus()

    y_true = [doc['sentences'] for doc in sents]

    y_pred = [Doc(doc).sents for doc in raw]

    paths = file_paths()

    for i in range(len(y_true)):

        if y_true[i] != y_pred[i]:
            click.secho(f"Document {paths[i]}")
            for s_true in y_true[i]:
                if s_true not in y_pred[i]:
                    click.secho(f"+ {s_true}", fg="green")

            click.secho()

        for s_pred in y_pred[i]:
            if s_pred not in y_true[i]:
                click.secho(f"- {s_pred}", fg="red")

        click.secho()
        click.secho()


@tokenize.command()
def build():
    """Build a ML based SBD"""

    raw_corpus = load_raw_corpus(False)
    sent_corpus = load_sentence_corpus(False)

    features = flatten([[span.span_features() for span in Doc(raw).spans] for raw in raw_corpus])
    y = flatten(
        [[is_eos(span, sent['sentences']) for span in Doc(raw).spans] for raw, sent in zip(raw_corpus, sent_corpus)])

    if len(features) != len(y):
        raise Exception(f"Sanity check failed feature list length {len(features)} whereas target list length {len(y)}.")

    sbd_model = create_model()

    scores = cross_val_score(sbd_model, features, y, scoring="f1")

    for i, score in enumerate(scores):
        click.secho(f"Fold {i + 1}: {score:.4f}", fg="yellow")

    sbd_model.fit(features, y)

    click.secho("\nTop 10 Features")
    feature_importance = sbd_model.steps[1][1].feature_importances_
    for idx in list(reversed(feature_importance.argsort()))[:20]:
        click.secho(f"    {sbd_model.steps[0][1].feature_names_[idx]}: {feature_importance[idx]:.4f}", fg="yellow")

    save_model(sbd_model)


if __name__ == '__main__':
    tokenize()
