# coding: utf-8
from __future__ import unicode_literals

import json
import click
from pathlib import Path
from wasabi import Printer
from io import open
from .dframcy import DframCy
from .utils import get_default_columns
from .trainer import DframeTrainer, DframeEvaluator, DframeTrainClassifier

messenger = Printer()
DEFAULT_COLUMNS = ",".join(get_default_columns())


@click.group()
def main():
    pass


@main.command()
@click.option(
    "--input_file", "-i", required=True, type=Path, help="Input text file path."
)
@click.option(
    "--output_file", "-o", required=True, type=Path, help="Output file path/name"
)
@click.option(
    "--convert_type",
    "-t",
    default="csv",
    show_default=True,
    type=str,
    help="Output file format (JSON/CSV)",
)
@click.option(
    "--language_model",
    "-l",
    default="en_core_web_sm",
    show_default=True,
    type=str,
    help="Language model " "to be used.",
)
@click.option(
    "--columns",
    "-c",
    default=DEFAULT_COLUMNS,
    show_default=True,
    type=str,
    help="Annotations to be " "included in dataframe.",
)
@click.option(
    "--separate_entity_frame",
    "-s",
    default=False,
    show_default=True,
    type=bool,
    help="Save separate " "entity dataframe.",
)
def convert(
    input_file,
    output_file,
    convert_type,
    language_model,
    columns,
    separate_entity_frame,
):
    if output_file.is_dir():
        output_file = output_file.joinpath(input_file.stem + "." + str(convert_type))
    if input_file.exists():
        with open(input_file, "r") as infile:
            text = infile.read()
            dframcy = DframCy(language_model)
            doc = dframcy.nlp(text)
            if columns == DEFAULT_COLUMNS:
                annotation_dataframe = dframcy.to_dataframe(
                    doc, separate_entity_dframe=separate_entity_frame
                )
            else:
                annotation_dataframe = dframcy.to_dataframe(
                    doc,
                    columns=columns.split(", "),
                    separate_entity_dframe=separate_entity_frame,
                )
            if separate_entity_frame:
                token_annotation_dataframe, entity_dataframe = annotation_dataframe
            else:
                token_annotation_dataframe = annotation_dataframe
                entity_dataframe = None

            if convert_type == "csv":
                token_annotation_dataframe.to_csv(output_file)
                if separate_entity_frame:
                    entity_output_file = Path(
                        str(output_file).strip(".csv") + "_entity.csv"
                    )
                    entity_dataframe.to_csv(entity_output_file)
            elif convert_type == "json":
                annotation_json = token_annotation_dataframe.to_json(orient="columns")
                with open(output_file, "w") as outfile:
                    json.dump(annotation_json, outfile)
                if separate_entity_frame:
                    entity_output_file = Path(
                        str(output_file).strip(".json") + "_entity.json"
                    )
                    with open(entity_output_file, "w") as ent_outfile:
                        json.dump(entity_dataframe, ent_outfile)
            else:
                messenger.fail(
                    "Unknown output file format '{}'".format(convert_type), exits=-1
                )
    else:
        messenger.fail("input path {} does not exist".format(input_file), exits=-1)


@main.command()
@click.option("--lang", "-l", required=True, type=str, help="Model language.")
@click.option(
    "--output_path",
    "-o",
    required=True,
    type=str,
    help="Output directory to store mode in.",
)
@click.option(
    "--train_path",
    "-t",
    required=True,
    type=str,
    help="Path of CSV containing training data.",
)
@click.option(
    "--dev_path",
    "-d",
    required=True,
    type=str,
    help="Path to CSV containing validation data",
)
@click.option(
    "--debug_data_first",
    "-debug",
    default=True,
    show_default=True,
    type=bool,
    help="Run spaCy's training " "data debugger before " "training",
)
@click.option(
    "--raw_text",
    "-rt",
    default=None,
    show_default=True,
    type=str,
    help="Path to jsonl file with unlabelled " "text documents.",
)
@click.option(
    "--base_model",
    "-b",
    default=None,
    show_default=True,
    type=str,
    help="Name of model to update",
)
@click.option(
    "--pipeline",
    "-p",
    default="tagger,parser,ner",
    show_default=True,
    type=str,
    help="Comma-separated " "names of pipeline " "components",
)
@click.option(
    "--vectors",
    "-v",
    default=None,
    show_default=True,
    type=str,
    help="Model to load vectors from",
)
@click.option(
    "--n_iter",
    "-n",
    default=30,
    show_default=True,
    type=int,
    help="Number of iterations",
)
@click.option(
    "--n_early_stopping",
    "-ne",
    default=None,
    show_default=True,
    type=int,
    help="Maximum number of " "training epochs without " "dev accuracy improvement",
)
@click.option(
    "--n_examples",
    "-ns",
    default=0,
    show_default=True,
    type=int,
    help="Number of examples",
)
@click.option(
    "--use_gpu", "-g", default=-1, show_default=True, type=int, help="Use GPU"
)
@click.option(
    "--version",
    "-v",
    default="0.0.0",
    show_default=True,
    type=str,
    help="Model version",
)
@click.option(
    "--meta_path",
    "-m",
    default=None,
    show_default=True,
    type=Path,
    help="Optional path to meta.json to " "use as base.",
)
@click.option(
    "--init_tok2vec",
    "-t2v",
    default=None,
    show_default=True,
    type=str,
    help="Path to pretrained weights " "for the token-to-vector " "parts of the models",
)
@click.option(
    "--parser_multitasks",
    "-pm",
    default="",
    show_default=True,
    type=str,
    help="Side objectives for parser " "CNN, e.g. 'dep' or 'dep," "tag'",
)
@click.option(
    "--entity_multitasks",
    "-em",
    default="",
    show_default=True,
    type=str,
    help="Side objectives for NER " "CNN, e.g. 'dep' or 'dep," "tag'",
)
@click.option(
    "--noise_level",
    "-n",
    default=0.0,
    show_default=True,
    type=float,
    help="Amount of corruption for data " "augmentation",
)
@click.option(
    "--orth_variant_level",
    "-vl",
    default=0.0,
    show_default=True,
    type=float,
    help="Amount of orthography " "variation for data " "augmentation",
)
@click.option(
    "--eval_beam_widths",
    "-bw",
    default="",
    show_default=True,
    type=str,
    help="Beam widths to evaluate, " "e.g. 4,8",
)
@click.option(
    "--gold_preproc",
    "-G",
    default=False,
    show_default=True,
    type=bool,
    help="Use gold preprocessing",
)
@click.option(
    "--learn_tokens",
    "-T",
    default=False,
    show_default=True,
    type=bool,
    help="Make parser learn " "gold-standard tokenization",
)
@click.option(
    "--textcat_multilabel",
    "-TML",
    default=False,
    show_default=True,
    type=bool,
    help="Textcat classes " "aren't mutually " "exclusive (" "multilabel)",
)
@click.option(
    "--textcat_arch",
    "-ta",
    default="bow",
    show_default=True,
    type=str,
    help="Textcat model architecture",
)
@click.option(
    "--textcat_positive_label",
    "-tpl",
    default=None,
    show_default=True,
    type=str,
    help="Textcat positive " "label for binary " "classes with two " "labels",
)
@click.option(
    "--verbose", "-VV", default=False, show_default=True, type=bool, help="verbosity"
)
def train(
    lang,
    output_path,
    train_path,
    dev_path,
    debug_data_first,
    raw_text,
    base_model,
    pipeline,
    vectors,
    n_iter,
    n_early_stopping,
    n_examples,
    use_gpu,
    version,
    meta_path,
    init_tok2vec,
    parser_multitasks,
    entity_multitasks,
    noise_level,
    orth_variant_level,
    eval_beam_widths,
    gold_preproc,
    learn_tokens,
    textcat_multilabel,
    textcat_arch,
    textcat_positive_label,
    verbose,
):
    dframe_trainer = DframeTrainer(
        lang,
        output_path,
        train_path,
        dev_path,
        debug_data_first,
        raw_text,
        base_model,
        pipeline,
        vectors,
        n_iter,
        n_early_stopping,
        n_examples,
        use_gpu,
        version,
        meta_path,
        init_tok2vec,
        parser_multitasks,
        entity_multitasks,
        noise_level,
        orth_variant_level,
        eval_beam_widths,
        gold_preproc,
        learn_tokens,
        textcat_multilabel,
        textcat_arch,
        textcat_positive_label,
        verbose,
    )
    dframe_trainer.begin_training()


@main.command()
@click.option("--model", "-m", required=True, type=str, help="Model name or path")
@click.option(
    "--data_path",
    "-d",
    required=True,
    type=str,
    help="Path of CSV containing validation data",
)
@click.option(
    "--gpu_id", "-g", default=-1, show_default=True, type=bool, help="Use GPU"
)
@click.option(
    "--gold_preproc",
    "-G",
    default=False,
    show_default=True,
    type=bool,
    help="Use gold preprocessing",
)
@click.option(
    "--displacy_path",
    "-dp",
    default=None,
    show_default=True,
    type=str,
    help="Directory to output rendered " "parses as HTML",
)
@click.option(
    "--displacy_limit",
    "-dl",
    default=25,
    show_default=True,
    type=int,
    help="Limit of parses to render as " "HTML",
)
@click.option(
    "--return_scores",
    "-R",
    default=False,
    show_default=True,
    type=bool,
    help="Return dict containing " "model scores",
)
def evaluate(
    model, data_path, gpu_id, gold_preproc, displacy_path, displacy_limit, return_scores
):
    dframe_evaluator = DframeEvaluator(
        model,
        data_path,
        gpu_id,
        gold_preproc,
        displacy_path,
        displacy_limit,
        return_scores,
    )
    dframe_evaluator.begin_evaluation()


@main.command()
@click.option("--output_path", "-o", required=True, type=str, help="Output model path")
@click.option(
    "--train_path", "-t", required=True, type=str, help="path to training data csv"
)
@click.option(
    "--dev_path",
    "-d",
    required=True,
    type=str,
    help="path to testing/validation data csv",
)
@click.option(
    "--model",
    "-m",
    default=None,
    show_default=True,
    type=str,
    help="language model name",
)
@click.option(
    "--n_iter",
    "-n",
    default=20,
    show_default=True,
    type=int,
    help="Number of training iterations",
)
@click.option(
    "--init_tok2vec",
    "-t2v",
    default=None,
    show_default=True,
    type=Path,
    help="Pretrained tok2vec weights",
)
@click.option(
    "--exclusive_classes",
    "-ec",
    default=False,
    show_default=True,
    type=bool,
    help="classes exclusive",
)
@click.option(
    "--architecture",
    "-a",
    default="ensemble",
    show_default=True,
    type=str,
    help="model architecture",
)
@click.option(
    "--train_split",
    "-s",
    default=0.8,
    show_default=True,
    type=float,
    help="split in case no testing data",
)
def textcat(
    output_path,
    train_path,
    dev_path,
    model,
    n_iter,
    init_tok2vec,
    exclusive_classes,
    architecture,
    train_split,
):
    dframe_textcat_train = DframeTrainClassifier(
        output_path,
        train_path,
        dev_path,
        model,
        n_iter,
        init_tok2vec,
        exclusive_classes,
        architecture,
        train_split,
    )
    dframe_textcat_train.begin_training()


if __name__ == "__main__":
    main()
