"""
Post Process DA
----------------------------------

Tools to process data after sixdb has calculated the
da. Includes functions for extracting data from database
as well as plotting of DA polar plots.
"""
import logging
from pathlib import Path
from typing import Any, Tuple, Iterable

import numpy as np
import pandas as pd
from generic_parser import DotDict
from matplotlib import pyplot as plt
from matplotlib import rcParams, lines as mlines
from scipy.interpolate import interp1d
from tfs import TfsDataFrame, write_tfs

from pylhc_submitter.constants.autosix import (
    get_tfs_da_path,
    get_tfs_da_seed_stats_path,
    get_tfs_da_angle_stats_path,
    get_autosix_results_path,
    HEADER_NTOTAL,
    HEADER_INFO,
    HEADER_HINT,
    MEAN,
    STD,
    MIN,
    MAX,
    N,
    SEED,
    ANGLE,
    ALOST1,
    ALOST2,
    AMP,
)
from pylhc_submitter.sixdesk_tools.extract_data_from_db import extract_da_data

LOG = logging.getLogger(__name__)

DA_COLUMNS = (ALOST1, ALOST2)
INFO = (
    "Statistics over the N={n:d} {over:s} per {per:s}. "
    "The N-Columns indicate how many non-zero DA values were used."
)
HINT = "{param:s} {val:} is the respective value calculated over all other {param:s}s."
OVER_WHICH = {SEED: "angles", ANGLE: "seeds"}

# Fixed plot styles ---
COLOR_MEAN = "red"
COLOR_SEED = "grey"
COLOR_LIM = "black"
COLOR_FILL = "blue"
ALPHA_SEED = 0.5
ALPHA_FILL = 0.2
ALPHA_FILL_STD = 0.2


def post_process_da(jobname: str, basedir: Path):
    """ Post process the DA results into dataframes and DA plots. """
    LOG.info("Post-Processing Sixdesk Results.")
    df_da, df_angle, df_seed = create_da_tfs(jobname, basedir)
    create_polar_plots(jobname, basedir, df_da, df_angle)
    LOG.info("Post-Processing finished.")


# Data Analysis ----------------------------------------------------------------


def create_da_tfs(jobname: str, basedir: Path) -> Tuple[TfsDataFrame, TfsDataFrame, TfsDataFrame]:
    """ Extracts data from db into dataframes, and writes and returns them.

    Args:
        jobname (str): Name of the Job
        basedir (Path): SixDesk Basefolder Location
    """
    LOG.info("Gathering DA data into tfs-files.")
    df_da = extract_da_data(jobname, basedir)
    df_angle = _create_stats_df(df_da, ANGLE)
    df_seed = _create_stats_df(df_da, SEED, global_index=0)

    write_tfs(get_tfs_da_path(jobname, basedir), df_da)
    write_tfs(get_tfs_da_angle_stats_path(jobname, basedir), df_angle, save_index=ANGLE)
    write_tfs(get_tfs_da_seed_stats_path(jobname, basedir), df_seed, save_index=SEED)
    return df_da, df_angle, df_seed


def _create_stats_df(df: pd.DataFrame, parameter: str, global_index: Any = None) -> TfsDataFrame:
    """ Calculates the stats over a given parameter.
    Note: Could be refactored to use `group_by`.

    Args:
        df (DataFrame): DataFrame containing the DA information over all seeds.
        parameter (str): The parameter over which we want to average, i.e.
                         SEED or ANGLE.
        global_index (Any): identifier to use as a global index, i.e. the statistics
                            over all entries are stored here. (e.g. '0' for SEEDs)

    """
    operation_map = DotDict({MEAN: np.mean, STD: np.std, MIN: np.min, MAX: np.max})

    pre_index = [] if global_index is None else [global_index]
    index = sorted(set(df[parameter]))
    n_total = sum(df[parameter] == index[0])

    df_stats = TfsDataFrame(
        index=pre_index + index,
        columns=[f"{fun}{al}" for al in DA_COLUMNS for fun in list(operation_map.keys()) + [N]],
    )
    df_stats.headers[HEADER_INFO] = INFO.format(
        over=OVER_WHICH[parameter], per=parameter.lower(), n=n_total
    )
    df_stats.headers[HEADER_NTOTAL] = n_total

    for col_da in DA_COLUMNS:
        for idx in index:
            mask = (df[parameter] == idx) & (df[col_da] != 0)
            df_stats.loc[idx, f"{N}{col_da}"] = sum(mask)
            for name, operation in operation_map.items():
                df_stats.loc[idx, f"{name}{col_da}"] = operation(df.loc[mask, col_da])
            for name, operation in operation_map.get_subdict([MIN, MAX]).items():
                df_stats.loc[idx, f"{name}{AMP}"] = operation(df.loc[mask, f"{name}{AMP}"])

        if global_index is not None:
            # Note: could be done over df_stats for MEAN, MIN and MAX, but not STD
            mask = df[col_da] != 0
            df_stats.loc[global_index, f"{N}{col_da}"] = sum(mask)

            # Global MEAN, MIN, MAX Dynamic Aperture
            for name, operation in operation_map.get_subdict([MEAN, MIN, MAX, STD]).items():
                df_stats.loc[global_index, f"{name}{col_da}"] = operation(
                    df.loc[mask, col_da]
                )

            # Global MIN, MAX Amplitudes
            for name, operation in operation_map.get_subdict([MIN, MAX]).items():
                df_stats.loc[global_index, f"{name}{AMP}"] = operation(
                    df.loc[mask, f"{name}{AMP}"]  # min(MINA) and max(MAXA)
                )

            df_stats.headers[HEADER_HINT] = HINT.format(param=parameter, val=global_index)

    return df_stats


# Single Plots -----------------------------------------------------------------


def create_polar_plots(jobname: str, basedir: Path, df_da: TfsDataFrame, df_angles: TfsDataFrame):
    """ Plotting loop over da-methods and wrapper so save plots.

    Args:
        jobname (str): Name of the Job
        basedir (Path): SixDesk Basefolder Location
        df_da (TfsDataFrame): Full DA analysis result.
        df_angles (TfsDataFrame): Dataframe with the statistics (min, max, mean) per angle
    """
    LOG.info("Creating Polar Plots.")
    outdir_path = get_autosix_results_path(jobname, basedir)
    for da_col in DA_COLUMNS:
        fig = plot_polar(df_angles, da_col, jobname, df_da)
        fig.tight_layout(), fig.tight_layout()
        fig.savefig(outdir_path / fig.canvas.get_default_filename())

    # plt.show()


def plot_polar(
    df_angles: TfsDataFrame,
    da_col: str = ALOST2,
    jobname: str = "",
    df_da: TfsDataFrame = None,
    **kwargs,
) -> plt.Figure:
    """Create Polar Plot for DA analysis data.

    Keyword arguments are all optional.

    Args:
        df_angles (TfsDataFrame): Dataframe with the statistics (min, max, mean) per angle
        da_col (str): DA-Column name from sixdesk analysis to be used , e.g. ``ALOST2``.
                      (optional, default: ``ALOST2``)
        jobname (str): Name of the job. Used in window title only (optional).
        df_da (TfsDataFrame): Full DA analysis result. If given, plots
                              the individual DA results per seed. (optional)

    Keyword Arguments:
        interpolated (bool): If true, uses interpolation to plot the lines curved
        fill (bool): If true, fills the area between min and max with light blue
        angle_ticks (Iterable[numeric]): Positions in degree of the angle ticks (and lines)
        amplitude ticks (Iterable[numeric]): Positions of the amplitude ticks.

    Returns:
        Figure of the polar plot.
    """
    interpolated: bool = kwargs.pop("interpolated", True)
    fill: bool = kwargs.pop("fill", df_da is None)
    angle_ticks: Iterable[np.numeric] = kwargs.pop("angle_ticks", None)
    amplitude_ticks: Iterable[np.numeric] = kwargs.pop("amplitude_ticks", None)

    if "lines.marker" not in kwargs:
        kwargs["lines.marker"] = "None"
    fig, ax = plt.subplots(nrows=1, ncols=1, subplot_kw={"projection": "polar"})
    fig.canvas.manager.set_window_title(f"{jobname} polar plot for {da_col}")

    angles = np.deg2rad(df_angles.index)
    da_min, da_mean, da_max, da_std = (
        df_angles[f"{name}{da_col}"] for name in (MIN, MEAN, MAX, STD)
    )

    seed_h, seed_l = _plot_seeds(ax, df_da, da_col, interpolated)
    if interpolated:
        mean_h, max_h = _plot_interpolated(ax, angles, da_min, da_mean, da_max, da_std, fill)
    else:
        mean_h, max_h = _plot_straight(ax, angles, da_min, da_mean, da_max, da_std, fill)

    ax.set_thetamin(0)
    ax.set_thetamax(90)
    ax.set_rlim([0, None])
    ax.set_xlabel(r"DA$_{x}~[\sigma_{nominal}]$", labelpad=15)
    ax.set_ylabel(r"DA$_{y}~[\sigma_{nominal}]$", labelpad=20)

    if angle_ticks is not None:
        ax.set_xticks(np.deg2rad(angle_ticks))

    if amplitude_ticks is not None:
        ax.set_yticks(amplitude_ticks)
    ax.tick_params(labelright=True, labelleft=True)

    ax.legend(
        loc="upper right",
        bbox_to_anchor=(0.9, 0.95),
        bbox_transform=fig.transFigure,
        # frameon=False,
        handles=seed_h + [mean_h, max_h],
        labels=seed_l + ["Mean DA", "Limits"],
        ncol=1,
    )

    return fig


def _plot_seeds(ax, df_da: TfsDataFrame, da_col: str, interpolated: bool) -> Tuple[list, list]:
    """Add the Seed lines to the polar plots, if df_da is given.

    Args:
        ax: Axes to plot in
        df_da: DataFrame with DA information
        da_col: Dynamic Aperture column (ALOST1 or ALOST2)
        interpolated: If true, the lines will be curved.

    Returns:
        Tuple of list of one line handle and a list of a single label
    """
    if df_da is not None:
        seed_h = None
        for seed in sorted(set(df_da[SEED])):
            seed_mask = df_da[SEED] == seed
            angles = np.deg2rad(df_da.loc[seed_mask, ANGLE])
            da_data = df_da.loc[seed_mask, da_col]
            da_data.loc[da_data == 0] = np.NaN
            if interpolated:
                seed_h, _, _ = _interpolated_line(
                    ax,
                    angles,
                    da_data,
                    c=COLOR_SEED,
                    ls="-",
                    label=f"Seed {seed:d}",
                    alpha=ALPHA_SEED,
                )
            else:
                (seed_h,) = ax.plot(
                    angles, da_data, c=COLOR_SEED, ls="-", label=f"Seed {seed:d}", alpha=ALPHA_SEED
                )
        return [seed_h], ["DA per Seed"]
    return [], []


def _plot_interpolated(ax, angles, da_min, da_mean, da_max, da_std, fill):
    """Plot interpolated DA lines and areas."""
    _, _, ip_min = _interpolated_line(ax, angles, da_min, c=COLOR_LIM, ls="--", label="Minimum DA")
    max_h, ip_x, ip_max = _interpolated_line(
        ax, angles, da_max, c=COLOR_LIM, ls="--", label="Maximum DA"
    )
    if fill:
        ax.fill_between(ip_x, ip_min, ip_max, color=COLOR_FILL, alpha=ALPHA_FILL)
        _, ip_std_min = _interpolated_coords(angles, da_mean - da_std)
        _, ip_std_max = _interpolated_coords(angles, da_mean + da_std)
        ax.fill_between(ip_x, ip_std_min, ip_std_max, color=COLOR_FILL, alpha=ALPHA_FILL_STD)

    mean_h, _, _ = _interpolated_line(ax, angles, da_mean, c=COLOR_MEAN, ls="-", label="Mean DA")
    return mean_h, max_h


def _plot_straight(ax, angles, da_min, da_mean, da_max, da_std, fill):
    """Plot straight DA lines and areas."""
    (_,) = ax.plot(angles, da_min, c=COLOR_LIM, ls="--", label="Minimum DA")
    (max_h,) = ax.plot(angles, da_max, c=COLOR_LIM, ls="--", label="Maximum DA")
    if fill:
        ax.fill_between(
            angles,
            da_min.astype(float),
            da_max.astype(float),  # weird conversion to obj otherwise
            color=COLOR_FILL,
            alpha=ALPHA_FILL,
        )
        ax.fill_between(
            angles,
            (da_mean - da_std).astype(float),
            (da_mean + da_std).astype(float),
            color=COLOR_FILL,
            alpha=ALPHA_FILL_STD,
        )
    (mean_h,) = ax.plot(angles, da_mean, c=COLOR_MEAN, ls="-", label="Mean DA")
    return mean_h, max_h


def _interpolated_line(ax, x, y, npoints: int = 100, **kwargs):
    """Plot a line that interpolates linearly between points.
    Useful for polar plots with sparse points."""
    ls = kwargs.pop("linestyle", kwargs.pop("ls", rcParams["lines.linestyle"]))
    marker = kwargs.pop("marker", rcParams["lines.marker"])
    label = kwargs.pop("label")

    ip_x, ip_y = _interpolated_coords(x, y, npoints)
    (line_h,) = ax.plot(ip_x, ip_y, marker="None", ls=ls, label=f"_{label}_line", **kwargs)

    if marker.lower() not in ["none", ""]:
        ax.plot(x, y, ls="None", marker=marker, label=f"_{label}_markers", **kwargs)

    # fake handle for legend
    handle = mlines.Line2D([], [], color=line_h.get_color(), ls=ls, marker=marker, label=label)
    return handle, ip_x, ip_y


def _interpolated_coords(x, y, npoints: int = 100):
    """ Do linear interpolation between points. """
    ip_x = np.linspace(min(x), max(x), npoints)
    ip_y = interp1d(x, y)(ip_x)
    return ip_x, ip_y
