import os
import re
from collections import defaultdict
from typing import TYPE_CHECKING, List

import numpy
from typing_extensions import Literal

from nonbonded.library.models.forcefield import Parameter
from nonbonded.library.models.projects import Optimization
from nonbonded.library.models.results import OptimizationResult, TargetResultType
from nonbonded.library.models.targets import OptimizationTarget
from nonbonded.library.plotting.seaborn.utilities import plot_categories
from nonbonded.library.plotting.utilities import (
    combine_target_rmse,
    sort_categories_key,
)
from nonbonded.library.utilities.string import camel_to_kebab_case

if TYPE_CHECKING:
    import simtk.unit
    from openff.toolkit.typing.engines.smirnoff.forcefield import (
        ForceField as OFFForceField,
    )


def _default_parameter_unit(parameter: Parameter) -> "simtk.unit":
    """Returns the default unit for a parameter.

    Parameters
    ----------
    parameter
        The parameter to retrieve the unit of.

    Returns
    -------
        The default parameter unit.
    """

    import simtk.unit

    default_units = {
        "epsilon": simtk.unit.kilocalories_per_mole,
        "rmin_half": simtk.unit.angstrom,
        "sigma": simtk.unit.angstrom,
        "charge_increment1": simtk.unit.elementary_charge,
    }

    return default_units[parameter.attribute_name]


def _extract_parameter_value(
    force_field: "OFFForceField", parameter: Parameter
) -> float:
    """Extracts the value of a parameter from a force field object
    in its default unit.

    Parameters
    ----------
    force_field
        The force field which contains the parameter.
    parameter
        The parameter to extract the value of

    Returns
    -------
        The value of the parameter.
    """

    handler = force_field.get_parameter_handler(parameter.handler_type)

    off_parameter = (
        handler if parameter.smirks is None else handler.parameters[parameter.smirks]
    )

    value = getattr(off_parameter, parameter.attribute_name)

    return value.value_in_unit(_default_parameter_unit(parameter))


def _parameter_attribute_to_title(parameter: Parameter) -> str:
    """Maps a parameter attribute into a human and ``matplotlib`` friendly plot title
    where possible.

    Parameters
    ----------
    parameter
        The parameter which contains the attribute to map.

    Returns
    -------
        A fitting title.
    """

    attribute_symbols = {
        "epsilon": r"\varepsilon",
        "sigma": r"\sigma",
        "rmin_half": r"\dfrac{r_{min}}{2}",
        "charge_increment": "q",
    }

    indexless_attribute = re.sub(r"\d", "", parameter.attribute_name)

    if indexless_attribute not in attribute_symbols:
        return parameter.attribute_name

    if parameter.attribute_name in attribute_symbols:
        return f"${attribute_symbols[parameter.attribute_name]}$"

    symbol = attribute_symbols[indexless_attribute]

    title = re.sub(r"(\d+)", r"_{\1}", parameter.attribute_name)
    title = title.replace(indexless_attribute, symbol)

    return f"${title}$"


def plot_parameter_changes(
    optimizations: List[Optimization],
    optimization_results: List[OptimizationResult],
    mode: Literal["percent", "absolute"],
    show_legend: bool,
    output_directory: str,
    file_type: Literal["png", "pdf"] = "png",
):
    """Plots the change in the force field parameters which have been trained over
    the course of an optimization.

    Parameters
    ----------
    optimizations
        The optimizations which have been performed.
    optimization_results
        The analyzed outputs of the optimizations.
    mode
        The type of change to plot. This may either be absolute changes
        or percent changes.
    show_legend
        Whether to show the legend.
    output_directory
        The directory in which to save the plots.
    file_type
        The file type to use for the plots.
    """
    import pandas
    import seaborn
    from matplotlib import pyplot

    # Determine which parameters were refit
    parameters = defaultdict(dict)

    for index, (optimization, optimization_result) in enumerate(
        zip(optimizations, optimization_results)
    ):

        original_force_field = optimization.force_field.to_openff()
        refit_force_field = optimization_result.refit_force_field.to_openff()

        for parameter in optimization.parameters_to_train:

            parameters[parameter][index] = (
                _extract_parameter_value(original_force_field, parameter),
                _extract_parameter_value(refit_force_field, parameter),
            )

    # Re-shape the parameter values into a plottable pandas frame.
    data_rows = []

    for parameter in parameters:
        for optimization_index in parameters[parameter]:

            original_value, refit_value = parameters[parameter][optimization_index]

            percentage_change = (
                0.0
                if numpy.isclose(original_value, 0.0)
                else (refit_value - original_value) / original_value * 100.0
            )
            absolute_change = refit_value - original_value

            unit = _default_parameter_unit(parameter).get_symbol()

            if unit == "A":
                unit = "Å"

            data_row = {
                "Name": optimizations[optimization_index].name,
                "Handler": parameter.handler_type,
                "SMIRKS": parameter.smirks,
                "Attribute": f"{_parameter_attribute_to_title(parameter)} ({unit})",
                r"$\Delta$": absolute_change,
                r"% $\Delta$": percentage_change,
            }

            data_rows.append(data_row)

    # Plot the data.
    parameter_data = pandas.DataFrame(data_rows)
    palette = seaborn.color_palette(n_colors=len(optimizations))

    header = r"$\Delta$" if mode == "absolute" else r"% $\Delta$"
    handlers = {*parameter_data["Handler"].unique()}

    for handler in handlers:

        handler_data = parameter_data[parameter_data["Handler"] == handler]

        n_attributes = len(handler_data["Attribute"].unique())
        n_smirks = len(handler_data["SMIRKS"].unique())

        plot = seaborn.FacetGrid(
            handler_data,
            row="Attribute",
            height=1.5 * n_attributes,
            aspect=1.0 + 0.1 * (n_smirks - 1),
            sharey=False,
        )
        plot.map_dataframe(plot_categories, "SMIRKS", header, "Name", color=palette)
        plot.set_titles("{row_name}")

        for index in range(n_attributes):

            # Move the x-axis to y = 0
            plot.axes[index, 0].spines["bottom"].set_position(("data", 0.0))
            # Remove the x-axis tick marks.
            plot.axes[index, 0].tick_params(axis="x", which="both", length=0)

        # Move the x-axis labels back under the plot, rather than on the x axis.
        pyplot.setp(
            plot.axes[n_attributes - 1, 0].get_xticklabels(),
            transform=plot.axes[n_attributes - 1, 0].get_xaxis_transform(),
            rotation=90,
        )

        if show_legend:
            plot.add_legend()

        handler_name = camel_to_kebab_case(handler.replace("vdW", "vdw"))

        plot.savefig(
            os.path.join(output_directory, f"{handler_name}-{mode}-change.{file_type}")
        )
        pyplot.close(plot.fig)


def plot_objective_per_iteration(
    optimizations: List[Optimization],
    optimization_results: List[OptimizationResult],
    show_legend: bool,
    output_directory: str,
    file_type: Literal["png", "pdf"] = "png",
):
    """Plots the objective function for each optimizations target as a function of the
    optimization iteration.

    Different optimization are coloured separately while different targets are styled
    with different markers.

    Parameters
    ----------
    optimizations
        The optimizations which have been performed.
    optimization_results
        The analyzed outputs of the optimizations.
    show_legend
        Whether to show the legend.
    output_directory
        The directory in which to save the plots.
    file_type
        The file type to use for the plots.
    """

    import pandas
    import seaborn
    from matplotlib import pyplot

    data_rows = []

    for optimization, optimization_result in zip(optimizations, optimization_results):
        for iteration, iteration_result in optimization_result.target_results.items():
            for target_id, target_result in iteration_result.items():

                data_rows.append(
                    {
                        "Optimization": optimization.name,
                        "Target": target_id,
                        "Iteration": iteration,
                        "Objective Function": target_result.objective_function,
                    }
                )

    plot_data = pandas.DataFrame(data_rows)

    hue = "Optimization"
    style = "Target" if len(plot_data["Target"].unique()) > 1 else None

    colors = seaborn.color_palette(
        "colorblind", n_colors=len(plot_data["Optimization"].unique())
    )

    plot = seaborn.relplot(
        data=plot_data,
        x="Iteration",
        y="Objective Function",
        hue=hue,
        style=style,
        height=4.0,
        legend="full" if show_legend else False,
        palette=colors,
    )
    for target in plot_data["Target"].unique():

        target_data = plot_data[plot_data["Target"] == target]

        seaborn.lineplot(
            data=target_data,
            x="Iteration",
            y="Objective Function",
            hue=hue,
            legend=False,
            estimator=None,
            ci=None,
            palette=colors,
        )

    plot.savefig(os.path.join(output_directory, f"objective-function.{file_type}"))
    pyplot.close(plot.fig)


def plot_target_rmse(
    targets: List[OptimizationTarget],
    target_results: List[TargetResultType],
    target_labels: List[str],
    file_prefix: str,
    output_directory: str,
    file_type: Literal["png", "pdf"] = "png",
):
    """Plots the RMSE for each of the specified targets and for each data type
    contained within the target results.

    Parameters
    ----------
    targets
        The targets which the results were collected for.
    target_results
        The target results to plot.
    target_labels
        The label associated with each target result.
    file_prefix
        The prefix to append to the file name of each produced plot.
    output_directory
        The directory in which to save the plots.
    file_type
        The file type to use for the plots.
    """
    import seaborn
    from matplotlib import pyplot

    plot_data = combine_target_rmse(targets, target_results, target_labels)

    if len(plot_data) == 0:
        return

    # Extract the unique data types (e.g. property types) which will be plotted
    # in separate figures.
    data_types = plot_data["Data Type"].unique()

    # Select a set of colors for each label
    color_palette = (
        ["skyblue", "limegreen"] + []
        if len(target_labels) < 3
        else seaborn.color_palette(n_colors=len(target_labels) - 2)
    )
    colors = {label: color_palette[index] for index, label in enumerate(target_labels)}

    for data_type in data_types:

        # Extract a data frame containing only the data type which should
        # be included in this figure.
        type_plot_data = plot_data[
            (
                plot_data["Data Type"] == data_type
                if data_type is not None
                else plot_data["Data Type"].isnull()
            )
        ]

        categories = sorted(
            type_plot_data["Category"].unique(), key=sort_categories_key, reverse=True
        )
        category_indices = [x * 2 for x in range(1, len(categories) + 1)]

        labels = type_plot_data["Label"].unique()

        with seaborn.axes_style("white"):

            figure, axis = pyplot.subplots(figsize=(5.0, 1.5 + (0.5 * len(categories))))

            # Plot data for each label.
            for label in labels:

                label_plot_data = type_plot_data[type_plot_data["Label"] == label]

                axis.errorbar(
                    label_plot_data["Value"],
                    label_plot_data["Category"]
                    .replace(categories, category_indices)
                    .values,
                    xerr=[
                        label_plot_data["Lower CI"],
                        label_plot_data["Upper CI"],
                    ],
                    color=colors[label],
                    capsize=5.0,
                    capthick=1.5,
                    linestyle="none",
                    marker="o",
                    label=label,
                )

            # Add a simple legend.
            axis.legend()

            # Add title and axis names
            axis.set_yticks(category_indices)
            axis.set_yticklabels(categories)

            axis.set_xlim(left=0.0)
            axis.set_ylim(0.5, (len(categories) + 1) * 2.0 - 0.5)

            axis.set_xlabel("RMSE")

        data_name = "-" if data_type is None else f"-{camel_to_kebab_case(data_type)}-"

        # Save the figure.
        figure.savefig(
            os.path.join(
                output_directory,
                f"{file_prefix}{data_name}rmse.{file_type}",
            ),
            bbox_inches="tight",
        )
        pyplot.close(figure)


def plot_rmse_change(
    optimization: Optimization,
    optimization_result: OptimizationResult,
    output_directory: str,
    file_type: Literal["png", "pdf"] = "png",
):
    """Plots the change in the RMSE between the first and the last iteration of an
    optimization for each optimization target.

    Parameters
    ----------
    optimization
        The optimization which the results were collected for.
    optimization_result
        The analyzed outputs of the optimization.
    output_directory
        The directory in which to save the plots.
    file_type
        The file type to use for the plots.
    """
    if 0 not in optimization_result.target_results:

        raise KeyError(
            "The optimization results must contain the statistics for iteration "
            "0 to plot the change in RMSE values."
        )

    if len(optimization_result.target_results) <= 1:

        raise KeyError(
            "The optimization results must contain the statistics for at least "
            "two iterations to plot the change in RMSE values."
        )

    targets_by_id = {target.id: target for target in optimization.targets}

    final_iteration = sorted(optimization_result.target_results)[-1]

    for target_id, initial_result in optimization_result.target_results[0].items():

        final_result = optimization_result.target_results[final_iteration][target_id]

        plot_target_rmse(
            [targets_by_id[target_id], targets_by_id[target_id]],
            [initial_result, final_result],
            ["Initial", "Final"],
            target_id,
            output_directory,
            file_type,
        )
