"""
main executable for viewing result files from the trajectory metric apps
author: Michael Grupp

This file is part of evo (github.com/MichaelGrupp/evo).

evo is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

evo is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with evo.  If not, see <http://www.gnu.org/licenses/>.
"""

import argparse
import logging
import sys

import pandas as pd

from evo.tools import log, user, pandas_bridge
from evo.tools.settings import SETTINGS

logger = logging.getLogger(__name__)

SEP = "-" * 80  # separator line

CONFLICT_TEMPLATE = """
Mismatching titles - risk of aggregating data from different metrics. Conflict:

<<<<<<< {first_file}
{first_title}
=======
{mismatching_title}
>>>>>>> {mismatching_file}

Only the first one will be used as the title!"""


def run(args: argparse.Namespace) -> None:

    pd.options.display.width = 80
    pd.options.display.max_colwidth = 20

    log.configure_logging(
        args.verbose, args.silent, args.debug, local_logfile=args.logfile
    )
    if args.debug:
        import pprint

        arg_dict = {arg: getattr(args, arg) for arg in vars(args)}
        logger.debug(f"main_parser config:\n{pprint.pformat(arg_dict)}\n")

    df = pandas_bridge.load_results_as_dataframe(
        args.result_files, args.use_filenames, args.merge
    )

    keys = df.columns.values.tolist()
    if SETTINGS.plot_usetex:
        keys = [key.replace("_", "\\_") for key in keys]
        df.columns = keys
    duplicates = [x for x in keys if keys.count(x) > 1]
    if duplicates:
        logger.error(
            f"Values of 'est_name' must be unique - duplicates: {', '.join(duplicates)}\n"
            f"Try using the --use_filenames option to use filenames "
            f"for labeling instead."
        )
        sys.exit(1)

    # derive a common index type if possible - preferably timestamps
    common_index = None
    time_indices = ["timestamps", "seconds_from_start", "sec_from_start"]
    if args.use_rel_time:
        del time_indices[0]
    for idx in time_indices:
        if idx not in df.loc["np_arrays"].index:
            continue
        if df.loc["np_arrays", idx].isnull().values.any():
            continue
        else:
            common_index = idx
            break

    # build error_df (raw values) according to common_index
    if common_index is None:
        # use a non-timestamp index
        error_df = pd.DataFrame(
            df.loc["np_arrays", "error_array"].tolist(), index=keys
        ).T
    else:
        error_df = pd.DataFrame()
        for key in keys:
            new_error_df = pd.DataFrame(
                {key: df.loc["np_arrays", "error_array"][key]},
                index=df.loc["np_arrays", common_index][key],
            )
            duplicates = new_error_df.index.duplicated(keep="first")
            if any(duplicates):
                logger.warning(
                    f"duplicate indices in error array of {key} - "
                    f"keeping only first occurrence of duplicates"
                )
                new_error_df = new_error_df[~duplicates]  # type: ignore
            error_df = pd.concat([error_df, new_error_df], axis=1)
        error_df.sort_index(inplace=True)

    # check titles
    if args.ignore_title:
        first_title = ""
    else:
        first_title = df.loc["info", "title"].iloc[0]
    first_file = args.result_files[0]
    if not args.no_warnings and not args.ignore_title:
        checks = df.loc["info", "title"] != first_title
        for i, differs in enumerate(checks):
            if not differs:
                continue
            else:
                mismatching_title = df.loc["info", "title"].iloc[i]
                mismatching_file = args.result_files[i]
                logger.debug(SEP)
                logger.warning(
                    CONFLICT_TEMPLATE.format(
                        first_file=first_file,
                        first_title=first_title,
                        mismatching_title=mismatching_title,
                        mismatching_file=mismatching_file,
                    )
                )
                if not user.confirm(
                    "You can use --ignore_title to just aggregate data.\n"
                    "Go on anyway? - enter 'y' or any other key to exit"
                ):
                    sys.exit()

    logger.debug(SEP)
    logger.debug(
        "Aggregated dataframe:\n%s",
        df.to_string(line_width=80, max_colwidth=40),
    )

    # show a statistics overview
    logger.debug(SEP)
    if not args.ignore_title:
        logger.info("\n" + first_title + "\n\n")
    logger.info(df.loc["stats"].T.to_string(line_width=80) + "\n")

    if args.rerun:
        send_to_rerun(args, df, keys, common_index, first_title)

    if args.save_table:
        logger.debug(SEP)
        if SETTINGS.table_export_data.lower() == "error_array":
            data = error_df
        elif SETTINGS.table_export_data.lower() in ("info", "stats"):
            data = df.loc[SETTINGS.table_export_data.lower()]
        else:
            raise ValueError(
                f"unsupported export data specifier: {SETTINGS.table_export_data}"
            )
        pandas_bridge.save_df_as_table(
            data, args.save_table, confirm_overwrite=not args.no_warnings
        )

    if args.plot or args.save_plot:
        # check if data has NaN "holes" due to different indices
        inconsistent = error_df.isnull().values.any()
        if (
            inconsistent
            and common_index != "timestamps"
            and not args.no_warnings
        ):
            logger.debug(SEP)
            logger.warning(
                "Data lengths/indices are not consistent, "
                "raw value plot might not be correctly aligned"
            )

        from evo.tools import plot
        import matplotlib.pyplot as plt
        import seaborn as sns
        import math

        # use default plot settings
        figsize = (SETTINGS.plot_figsize[0], SETTINGS.plot_figsize[1])
        use_cmap = SETTINGS.plot_multi_cmap.lower() != "none"
        colormap = SETTINGS.plot_multi_cmap if use_cmap else None
        linestyles = (
            ["-o" for x in args.result_files] if args.plot_markers else None
        )

        # labels according to first dataset
        if (
            "xlabel" in df.loc["info"].index
            and not df.loc["info", "xlabel"].isnull().values.any()
        ):
            index_label = df.loc["info", "xlabel"].iloc[0]
        else:
            index_label = "$t$ (s)" if common_index else "index"
        metric_label = df.loc["info", "label"].iloc[0]

        plot_collection = plot.PlotCollection(first_title)
        # raw value plot
        fig_raw = plt.figure(figsize=figsize)
        # handle NaNs from concat() above
        error_df.interpolate(method="index", limit_area="inside").plot(
            ax=fig_raw.gca(),
            colormap=colormap,
            style=linestyles,
            title=first_title,
            alpha=SETTINGS.plot_trajectory_alpha,
            legend=SETTINGS.plot_show_legend,
        )
        plt.xlabel(index_label)
        plt.ylabel(metric_label)
        plot_collection.add_figure("raw", fig_raw)

        # statistics plot
        if SETTINGS.plot_statistics:
            fig_stats = plt.figure(figsize=figsize)
            include = df.loc["stats"].index.isin(SETTINGS.plot_statistics)
            if any(include):
                df.loc["stats"][include].plot(
                    kind="barh",
                    ax=fig_stats.gca(),
                    colormap=colormap,
                    stacked=False,
                    legend=SETTINGS.plot_show_legend,
                )
                plt.xlabel(metric_label)
                plot_collection.add_figure("stats", fig_stats)

        # grid of distribution plots
        raw_tidy = pd.melt(
            error_df,
            value_vars=list(error_df.columns.values),
            var_name="estimate",
            value_name=metric_label,
        )
        col_wrap = (
            2
            if len(args.result_files) <= 2
            else math.ceil(len(args.result_files) / 2.0)
        )
        # FacetGrid uses tight_layout internally, which clashes with the
        # constrained_layout we set globally in evo.tools.plot.
        import warnings

        with plt.rc_context({"figure.constrained_layout.use": False}):
            dist_grid = sns.FacetGrid(
                raw_tidy, col="estimate", col_wrap=col_wrap
            )
            # TODO: see issue #98
            with warnings.catch_warnings():
                warnings.simplefilter("ignore")
                dist_grid.map(sns.distplot, metric_label)  # fits=stats.gamma
        plot_collection.add_figure("histogram", dist_grid.fig)

        # box plot
        fig_box = plt.figure(figsize=figsize)
        ax = sns.boxplot(
            x=raw_tidy["estimate"], y=raw_tidy[metric_label], ax=fig_box.gca()
        )
        # ax.set_xticklabels(labels=[item.get_text() for item in ax.get_xticklabels()], rotation=30)
        plot_collection.add_figure("box_plot", fig_box)

        # violin plot
        fig_violin = plt.figure(figsize=figsize)
        ax = sns.violinplot(
            x=raw_tidy["estimate"],
            y=raw_tidy[metric_label],
            ax=fig_violin.gca(),
        )
        # ax.set_xticklabels(labels=[item.get_text() for item in ax.get_xticklabels()], rotation=30)
        plot_collection.add_figure("violin_histogram", fig_violin)

        if args.plot:
            plot_collection.show()
        if args.save_plot:
            logger.debug(SEP)
            plot_collection.export(
                args.save_plot, confirm_overwrite=not args.no_warnings
            )


def send_to_rerun(
    args: argparse.Namespace,
    df: pd.DataFrame,
    keys: list,
    common_index: str | None,
    first_title: str,
) -> None:
    try:
        import pyarrow as pa
        import rerun as rr
        import rerun.blueprint as rrb
        from rerun.experimental import ViewerClient
    except ImportError:
        logger.error(
            "Optional dependency rerun-sdk is not installed. "
            "Install it with: pip install rerun-sdk"
        )
        sys.exit(1)

    from evo.tools.plot import color_cycle
    from evo.tools import rerun_bridge as revo

    evo_app_name = "evo_res"

    logger.debug(SEP)
    logger.debug("Sending data to Rerun.")
    rr.init(evo_app_name, recording_id=args.rerun_rec_id)
    client: ViewerClient = revo.connect_or_spawn()

    # Send a combined stats table.
    stats_table = df.loc["stats"].T.join(df.loc["info"].T).reset_index()
    stats_table.rename(columns={"index": "result"}, inplace=True)
    client.send_table(
        f"{evo_app_name} stats",
        pa.Table.from_pandas(stats_table).to_batches(),
    )

    # Blueprint layout.
    has_time_index = common_index is not None
    timeline = revo.TIMELINE if has_time_index else revo.INDEX_TIMELINE
    time_range = rrb.VisibleTimeRange(
        timeline=timeline,
        start=rrb.TimeRangeBoundary.infinite(),
        end=rrb.TimeRangeBoundary.cursor_relative(
            seconds=0.0 if has_time_index else 0
        ),
    )
    rr.send_blueprint(
        rrb.Blueprint(
            rrb.Grid(
                contents=[
                    rrb.TimeSeriesView(
                        name="Results",
                        time_ranges=time_range,
                        plot_legend=rrb.Corner2D.RightTop,
                    ),
                    rrb.Grid(
                        contents=[
                            rrb.BarChartView(
                                name=f"Statistics {key}",
                                origin=f"/{evo_app_name}/statistics/{key}",
                                plot_legend=rrb.PlotLegend(
                                    None, visible=False
                                ),
                            )
                            for key in keys
                        ],
                        grid_columns=len(keys),
                    ),
                ],
                grid_columns=1,
                row_shares=[2.0, 1.0],
            ),
            rrb.SelectionPanel(expanded=False),
            rrb.TimePanel(expanded=False),
        )
    )

    # Send error time series and statistics per result.
    colors = color_cycle()
    for i, key in enumerate(keys):
        color_rgba = colors[i % len(colors)]
        error_array = df.loc["np_arrays", "error_array"][key]

        timestamps = None
        if common_index == "timestamps":
            timestamps = df.loc["np_arrays", "timestamps"][key]
        elif common_index in ("seconds_from_start", "sec_from_start"):
            timestamps = df.loc["np_arrays", common_index][key]

        revo.send_scalars(
            entity_path=f"{evo_app_name}/errors/{key}",
            scalars=error_array,
            timestamps=timestamps,
            color=revo.Color(static=color_rgba),
            labelname=key,
        )

        stats = {
            stat: float(df.loc["stats"][key][stat])
            for stat in SETTINGS.plot_statistics
            if stat in df.loc["stats"].index
        }
        revo.send_statistics_bar_chart(
            entity_path=f"{evo_app_name}/statistics/{key}",
            stats=stats,
        )