import json
import math
import os
import textwrap
import traceback
from itertools import groupby
from operator import add
from typing import List
import matplotlib
import difflib

import matplotlib.pyplot as plt
from matplotlib.ticker import EngFormatter

import paraview.simple as pw  # type: ignore

"""Create plots in ParaView using its PythonView funcionality.

The data flow is as follows:
1. main.py calls this (charts.py) create_charts
2. create_charts calls get_post_values_to_plot, which reads the 
visualization config to determine which charts to show by default.
3. A new PythonView is created in ParaView
4. The PythonView script string for the view is set, where:
 - params_to_plot are written, so that the user could change them real-time.
 - This same file is imported from the PythonView script, and add_plots is called.
"""


def get_post_values_to_plot(RESULTS_JSON, VIZ_CONF) -> List[str]:
    """Determine which post values will be displayed by default"""
    params_to_plot = []

    available_parametrs = set(RESULTS_JSON["globalPostValues"].keys())

    all_subplot_params = [i["postValues"] for i in VIZ_CONF["charts"]]
    for subplot_params in all_subplot_params:
        present_params = [i for i in subplot_params if i in available_parametrs]
        params_to_plot.append(present_params)
    return params_to_plot


def create_charts(previous_view, VIZ_CONF, RESULTS_JSON) -> None:
    """Create matplotlib (PythonView) charts in a new layout tab"""
    new_layout = pw.CreateLayout("Time charts")
    pw.RenameLayout("3D View", pw.GetLayoutByName("Layout #1"))
    python_view = pw.CreateView("PythonView")
    # Switch back to the old view, so the charts would
    # not be the first thing the user sees.
    pw.SetActiveView(previous_view)
    new_layout.AssignView(0, python_view)

    subplot_parameters = get_post_values_to_plot(RESULTS_JSON, VIZ_CONF)

    # Everyting in the script string will be revealed to the user.
    # So subplot_parameters is exposed to let the user changed the plotted parameters.
    python_view.Script = """
subplot_parameters = {}
group_domain_values = True

import os, sys
import matplotlib.pyplot as plt
sys.path.append(os.path.join(os.getenv("CENOS_CONFIG"), "paraview"))
import charts

def render(view, width, height): 
    from paraview import python_view
    figure = python_view.matplotlib_figure(width, height)
    
    charts.add_plots(figure, subplot_parameters, group_domain_values)

    return python_view.figure_to_image(figure)
plt.close('all')

    """.format(
        str(subplot_parameters)
    )


### ==================================================================
### Everything below here gets called from the PythonView.
### ==================================================================


def combine_group_values(y_obj: dict) -> None:
    """Combine the values of all grouped domains in y_obj by adding togheter the values"""

    for key, group in groupby(y_obj["domainValues"], lambda x: x["blockGroup"]):
        if key == "none":
            continue

        new_domain = {}
        new_domain["blockGroup"] = "none"
        new_domain["blockName"] = key
        new_domain["values"] = []

        for thing in group:
            if not new_domain["values"]:
                # If list is empty
                new_domain["values"] = thing["values"]
            else:
                # Add the values to the existing ones
                new_domain["values"] = list(
                    map(add, new_domain["values"], thing["values"])
                )

        # Remove the groups which were grouped
        y_obj["domainValues"] = [
            d for d in y_obj["domainValues"] if d["blockGroup"] != key
        ]

        # And add the new domain
        y_obj["domainValues"].append(new_domain)

    return y_obj


def set_defaults() -> None:
    plt.style.use("ggplot")
    matplotlib.rcParams["lines.linewidth"] = 1.5
    matplotlib.rcParams["lines.marker"] = "o"
    matplotlib.rcParams["lines.markersize"] = 3
    matplotlib.rcParams["font.size"] = 12
    matplotlib.rcParams["axes.xmargin"] = 0.01
    matplotlib.rcParams["axes.ymargin"] = 0.01


def set_axes_style(ax) -> None:
    ax.tick_params(axis="both", direction="in", length=4, color="#DDDDDD")
    ax.grid(color="#DDDDDD", linestyle="-", linewidth=0.5)
    ax.tick_params(axis="both", direction="in", length=4, color="#DDDDDD", labelsize=11)
    ax.set_facecolor("#FFFFFF")


def add_plots(figure, plot_params: List[str], group) -> None:
    """This function gets called from the ParaView's PythonView,
    so the results json and case path has to be defined again.

    Args:
        figure: The matplotlib figure created in the PythonView script.
        plot_params: Nested list of globalPostValues to be plotted in each subplot.
        group: Determines if grouped domain values should be added togheter.

    """
    CASE_PATH = os.getcwd()
    results_json_path = os.path.join(CASE_PATH, "results/results.json")
    RESULTS_JSON = json.load(open(results_json_path, encoding="utf-8"))

    post_values = RESULTS_JSON["globalPostValues"]

    subplot_rows = 2
    subplot_columns = math.ceil(len(plot_params) / subplot_rows)
    subplot_index = 1

    set_defaults()

    for subplot_params in plot_params:
        ax = figure.add_subplot(subplot_rows, subplot_columns, subplot_index)
        title = ", ".join([post_values[i]["y"]["name"] for i in subplot_params])
        ax.title.set_text(title)

        try:
            is_multi_param = False
            if len(subplot_params) > 1:
                is_multi_param = True

            for param in subplot_params:

                if param not in post_values.keys():
                    suggestions = difflib.get_close_matches(param, post_values.keys())
                    msg = f'"{param}" was not found in the results.\n'
                    msg += f"Did you mean one of the following:\n"
                    for suggestion in suggestions:
                        msg += "  \u2022" + suggestion

                    raise RuntimeError(msg)

                post_val = post_values[param]

                x_obj = post_val["x"]
                x_values = x_obj["values"]
                ax.set_xlabel(x_obj["name"])
                ax.xaxis.set_major_formatter(EngFormatter(unit=x_obj["units"]))

                y_obj = post_val["y"]
                # ax.set_ylabel(y_obj["name"])
                # Math text must be wrapped in $
                ax.yaxis.set_major_formatter(
                    EngFormatter(unit=f"${y_obj['units']}$", useMathText=True)
                )

                if group:
                    combine_group_values(y_obj)

                for domain in y_obj["domainValues"]:
                    if is_multi_param:
                        legend = domain["blockName"] + " " + y_obj["name"]
                    else:
                        legend = domain["blockName"]

                    y_values = domain["values"]

                    ax.plot(x_values, y_values, label=legend)

                ax.legend(frameon=False)
                set_axes_style(ax)

        except Exception as e:
            trace = textwrap.fill(traceback.format_exc(), width=70)
            ax.annotate(f"Error: {str(e)}\n\n{trace}", (0.03, 0.4))

        subplot_index += 1
    figure.tight_layout()