"""
Command-line tool for starting an explainerdashboard from a particular directory
"""
import os
import webbrowser
from pathlib import Path
from importlib import import_module
from copy import deepcopy
import pickle
import oyaml as yaml

import pandas as pd
import joblib
import click

import waitress

from explainerdashboard import *
from explainerdashboard.explainers import BaseExplainer
from explainerdashboard.dashboards import ExplainerDashboard


explainer_ascii = """

 _____ ___ __| |__ _(_)_ _  ___ _ _ __| |__ _ __| |_ | |__  ___  __ _ _ _ __| |
/ -_) \ / '_ \ / _` | | ' \/ -_) '_/ _` / _` (_-< ' \| '_ \/ _ \/ _` | '_/ _` |
\___/_\_\ .__/_\__,_|_|_||_\___|_| \__,_\__,_/__/_||_|_.__/\___/\__,_|_| \__,_|
        |_| 

"""


def build_explainer(explainer_config):
    if isinstance(explainer_config, (Path, str)) and str(explainer_config).endswith(".yaml"):
        config = yaml.safe_load(open(str(explainer_config), "r"))
    elif isinstance(explainer_config, dict):
        config = explainer_config
    assert 'explainer' in config, \
        "Please pass a proper explainer.yaml config file that starts with `explainer:`!"
    config = explainer_config['explainer']

    print(f"explainerdashboard ===> Loading model from {config['modelfile']}")
    model = pickle.load(open(config['modelfile'], "rb"))
    
    print(f"explainerdashboard ===> Loading data from {config['datafile']}")
    if str(config['datafile']).endswith('.csv'):
        df = pd.read_csv(config['datafile'])
    elif str(config['datafile']).endswith('.parquet'):
        df = pd.read_parquet(config['datafile'])
    else:
        raise ValueError("datafile should either be a .csv or .parquet!")

    print(f"explainerdashboard ===> Using column {config['data_target']} to generate X, y ")
    target_col = config['data_target']
    X = df.drop(target_col, axis=1)
    y = df[target_col]

    if config['data_index'] is not None:
        print(f"explainerdashboard ===> Generating index from column {config['data_index']}")
        assert config['data_index'] in X.columns, \
            (f"Cannot find data_index column ({config['data_index']})"
             f" in datafile ({config['datafile']})!"
              "Please set it to the proper index column name, or set it to null")
        X = X.set_index(config['data_index'])
        
    params = config['params']

    if config['explainer_type'] == "classifier":
        print(f"explainerdashboard ===> Generating ClassifierExplainer...")
        explainer = ClassifierExplainer(model, X, y, **params)
    elif config['explainer_type'] == "regression":
        print(f"explainerdashboard ===> Generating RegressionExplainer...")
        explainer = ClassifierExplainer(model, X, y, **params)
    return explainer


def build_and_dump_explainer(explainer_config, dashboard_config=None):
    explainer = build_explainer(explainer_config)

    click.echo(f"explainerdashboard ===> Calculating properties by building Dashboard...")
    if dashboard_config is not None:
        ExplainerDashboard.from_config(explainer, dashboard_config)
    elif Path(explainer_config['explainer']['dashboard_yaml']).exists():
        click.echo(f"explainerdashboard ===> Calculating properties by building Dashboard from {explainer_config['explainer']['dashboard_yaml']}...")
        dashboard_config = yaml.safe_load(open(str(explainer_config['explainer']['dashboard_yaml']), "r"))
        ExplainerDashboard.from_config(explainer, dashboard_config)
    else:
        click.echo(f"explainerdashboard ===> Calculating all properties")
        explainer.calculate_properties()

    click.echo(f"explainerdashboard ===> Saving explainer to {explainer_config['explainer']['explainerfile']}...")
    if (dashboard_config is not None and 
        explainer_config['explainer']['explainerfile'] != dashboard_config['dashboard']['explainerfile']):
        click.echo(f"explainerdashboard ===> Warning explainerfile in explainer config and dashboard config do not match!")
    explainer.dump(explainer_config['explainer']['explainerfile'])
    return


def launch_dashboard_from_pkl(explainer_filepath, no_browser, port, no_dashboard=False):
    explainer = BaseExplainer.from_file(explainer_filepath)

    if port is None: 
        click.echo(f"explainerdashboard ===> Setting port to 8050, override with e.g. --port 8051")
        port = 8050

    db = ExplainerDashboard(explainer, port=port)
    
    if not no_browser and not os.environ.get("WERKZEUG_RUN_MAIN"):
        webbrowser.open_new(f"http://127.0.0.1:{port}/")

    if not no_dashboard:
        waitress.serve(db.flask_server(), host='0.0.0.0', port=port)
    return


def launch_dashboard_from_yaml(dashboard_config, no_browser, port, no_dashboard=False):
    if isinstance(dashboard_config, (Path, str)) and str(dashboard_config).endswith(".yaml"):
        config = yaml.safe_load(open(str(dashboard_config), "r"))
    elif isinstance(dashboard_config, dict):
        config = dashboard_config
    else:
        raise ValueError(f"dashboard_config should either be a .yaml filepath or a dict!")

    if not Path(config['dashboard']['explainerfile']).exists():
        click.echo(f"explainerdashboard ===> {config['dashboard']['explainerfile']} does not exist!")
        click.echo(f"explainerdashboard ===> first generate {config['dashboard']['explainerfile']} with explainerdashboard build")
        return

    click.echo(f"explainerdashboard ===> Building dashboard from {config['dashboard']['explainerfile']}")
    
    db = ExplainerDashboard.from_config(config)

    if port is None: 
        port =  config['dashboard']['params']['port'] 
        if port is None:
            port = 8050
        click.echo(f"explainerdashboard ===> Setting port to {port}, override with e.g. --port 8051")

    if not no_browser and not os.environ.get("WERKZEUG_RUN_MAIN"):
        click.echo(f"explainerdashboard ===> launching browser at {f'http://localhost:{port}/'}")
        webbrowser.open_new(f"http://localhost:{port}/")
    
    click.echo(f"explainerdashboard ===> Starting dashboard:")
    if not no_dashboard:
        waitress.serve(db.flask_server(), host='0.0.0.0', port=port)
    return


@click.group()
@click.pass_context
def explainerdashboard_cli(ctx):
    """
    explainerdashboard CLI tool. Used to launch an explainerdashboard from 
    the commandline. 

    \b
    explainerdashboard run
    ----------------------  

    Run explainerdashboard and start browser directly from command line.

    \b
    Example use:
        explainerdashboard run explainer.joblib
        explainerdashboard run dashboard.yaml
        explainerdashboard run dashboard.yaml --no-browser --port 8051
        explainerdashboard run --help

    If you pass an explainer.joblib file, will launch the full default dashboard.
    Generate this file with explainer.dump("explainer.joblib")

    If you pass a dashboard.yaml file, will launch a fully configured 
    explainerdashboard. Generate dashboard.yaml with 
    ExplainerDashboard.to_yaml('dashboard.yaml')

    If no argument given, searches for either dashboard.yaml or 
    explainer.joblib, in that order, so if you keep that naming convention 
    you can simply start with:

    \b
        explainerdashboard run

    \b
    explainerdashboard build
    ------------------------

    Build and store an explainer object, based on explainer.yaml file, that indicates
    where to find stored model (e.g. model.pkl), stored datafile (e.g. data.csv),
    and other explainer parameters.

    \b
    Example use:
        explainerdashboard build explainer.yaml
        explainerdashboard build explainer.yaml dashboard.yaml
        explainerdashboard build --help

    If given a second dashboard.yaml argument, will use that dashboard 
    configuration to calculate necessary properties for that specific dashboard
    configuration before storing to disk. Otherwise will use dashboard_yaml
    parameter in explainer.yaml to find configuration, or alternatively
    simply calculate all properties. 

    explainer.yaml file can be generated with explainer.to_yaml("explainer.yaml")

    If no argument given, searches for explainer.yaml, so if you keep that 
    naming convention you can simply start the build with:

    \b
        explainerdashboard build

    """


@explainerdashboard_cli.command(help="run dashboard and open browser")
@click.pass_context
@click.argument("explainer_filepath", nargs=1, required=False)
@click.option("--no-browser", "-nb", "no_browser", is_flag=True,
                 help="Launch a dashboard, but do not launch a browser.")
@click.option("--port", "-p", "port", default=None,
                help="port to run dashboard on defaults.")
def run(ctx, explainer_filepath, no_browser, port):
    click.echo(explainer_ascii)
    if explainer_filepath is None:
        if (Path().cwd() / "dashboard.yaml").exists():
            explainer_filepath = Path().cwd() / "dashboard.yaml"
        elif (Path().cwd() / "explainer.joblib").exists():
            explainer_filepath = Path().cwd() / "explainer.joblib"
        else:
            click.echo("No argument given and could find neither a "
                    "dashboard.yaml nor a explainer.joblib. Aborting.")
            return

    if (str(explainer_filepath).endswith(".joblib") or
        str(explainer_filepath).endswith(".pkl") or
        str(explainer_filepath).endswith(".pickle") or
        str(explainer_filepath).endswith(".dill")):
        launch_dashboard_from_pkl(explainer_filepath, no_browser, port)
        return
    elif str(explainer_filepath).endswith(".yaml"):
        launch_dashboard_from_yaml(explainer_filepath, no_browser, port)
    else:
        click.echo("Please pass a proper argument to explainerdashboard run"
                    "(i.e. either an explainer.joblib or a dashboard.yaml)")
    return


@explainerdashboard_cli.command(help="build and save explainer object")
@click.pass_context
@click.argument("explainer_filepath", nargs=1, required=False)
@click.argument("dashboard_filepath", nargs=1, required=False)
def build(ctx, explainer_filepath, dashboard_filepath):
    click.echo(explainer_ascii)
    if explainer_filepath is None:
        if (Path().cwd() / "explainer.yaml").exists():
            explainer_filepath = Path().cwd() / "explainer.yaml"
        else:
            click.echo("No argument given to explainerdashboard build and "
                        "could not find an explainer.yaml. Aborting.")
            return

    if str(explainer_filepath).endswith(".yaml"):
        explainer_config = yaml.safe_load(open(str(explainer_filepath), "r"))

        click.echo(f"explainerdashboard ===> Building {explainer_config['explainer']['explainerfile']}")

        if (dashboard_filepath is not None and 
            str(dashboard_filepath).endswith(".yaml")
            and Path(dashboard_filepath).exists()):
            click.echo(f"explainerdashboard ===> Using {dashboard_filepath} to calculate explainer properties")
            dashboard_config = yaml.safe_load(open(str(dashboard_filepath), "r"))
        else:
            dashboard_config = None

        print(f"explainerdashboard ===> Building {explainer_config['explainer']['explainerfile']}")
        build_and_dump_explainer(explainer_config, dashboard_config)
        print(f"explainerdashboard ===> Build finished!")
        return 

@explainerdashboard_cli.command(help="run without launching dashboard")
@click.pass_context
@click.argument("explainer_filepath", nargs=1, required=True)
@click.option("--port", "-p", "port", default=None,
                help="port to run dashboard on defaults.")
def test(ctx, explainer_filepath, port):
    if (str(explainer_filepath).endswith(".joblib") or
        str(explainer_filepath).endswith(".pkl") or
        str(explainer_filepath).endswith(".pickle") or
        str(explainer_filepath).endswith(".dill")):
        launch_dashboard_from_pkl(explainer_filepath, 
            no_browser=True, port=port, no_dashboard=True)
        return

    elif str(explainer_filepath).endswith(".yaml"):
        launch_dashboard_from_yaml(explainer_filepath, 
            no_browser=True, port=port, no_dashboard=True)
        return
    else:
        raise ValueError("Please pass a proper argument "
                    "(i.e. .joblib, .pkl, .dill or .yaml)!")

if __name__ =="__main__":
    explainerdashboard_cli()