import pandas as pd
import numpy as np
from shutil import rmtree
from matplotlib.axes import Axes
from concurrent.futures import ProcessPoolExecutor, as_completed
from typing import cast

from ..exec import N_CORES, SimJob, print_process_msg
from ..analysis import SimType, collect_avg_analyses, collect_run_time_series

from .utils import (
    LINE_STYLE,
    FILTERS,
    create_standard_figure,
    create_colorbar_figure,
    plot_horizontal_bands,
    plot_errorbar,
    plot_errorband,
    set_heatmap_colorbar,
    get_heatmap_norm,
    plot_main_heatmap,
    plot_side_heatmap,
    get_strat_eval,
    get_optimal_strat_phe_0,
    plot_dist_phe_0_lims,
    get_dist_avg_strat_phe_0,
    plot_expected_values,
    plot_extinct_times,
    plot_time_series,
    plot_colored_errorbar,
    plot_colored_curve,
    interpolate_values,
    interpolate_extinct_rates,
    plot_avg_strat_phe_0,
)


def make_param_plots(param: str, df: pd.DataFrame, job: SimJob) -> None:
    param_df = FILTERS[param](df, job).sort_values(param)
    if len(param_df) < 2:
        return

    two_panels = param == "strat_phe_0_i"
    fig_0, ax_0 = create_standard_figure(param, "avg_growth_rate")
    fig_1, ax_1 = create_standard_figure(param, "extinct_rate")
    fig_2, ax_2 = create_standard_figure("avg_growth_rate", "extinct_rate")
    fig_3, axs_3 = create_colorbar_figure(param, "norm_n_agents", two_panels)
    fig_4, axs_4 = create_colorbar_figure(param, "avg_strat_phe_0", two_panels)
    fig_5, ax_5 = create_standard_figure(param, "avg_strat_phe_0")
    fig_6, ax_6 = create_standard_figure(param, "dist_phe_0")
    fig_7, ax_7 = create_standard_figure(param, "std_dev_growth_rate")
    fig_8, ax_8 = create_standard_figure(param, "avg_birth_rate")
    fig_9, ax_9 = create_standard_figure("avg_growth_rate", "std_dev_growth_rate")

    if param in ["prob_mut", "n_agents_i"]:
        for ax in [axs_3[0], axs_4[0]]:
            ax.set_xscale("log")

    for sim_type, group_df in param_df.groupby("sim_type"):

        def plot_mean_and_uncertainty(
            ax: Axes, y_col: str, y_span_col: str | None
        ) -> None:
            if param == "strat_phe_0_i" and sim_type == SimType.RANDOM:
                if y_span_col is None:
                    plot_horizontal_bands(ax, group_df, (y_col, "mean"), (y_col, "sem"))
                else:
                    plot_horizontal_bands(
                        ax, group_df, (y_col, "mean"), (y_span_col, "mean")
                    )
            else:
                plot_errorbar(ax, group_df, param, y_col, False)
                if y_span_col is not None:
                    plot_errorband(ax, group_df, param, y_col, y_span_col)

        plot_mean_and_uncertainty(ax_0, "avg_growth_rate", None)
        plot_mean_and_uncertainty(ax_1, "extinct_rate", None)

        plot_errorbar(ax_2, group_df, "avg_growth_rate", "extinct_rate", True)

        if param == "strat_phe_0_i":
            if sim_type == SimType.FIXED:
                plot_main_heatmap(
                    fig_3, axs_3[0], axs_3[2], group_df, param, "dist_n_agents"
                )
            elif sim_type == SimType.EVOL:
                plot_main_heatmap(
                    fig_4, axs_4[0], axs_4[2], group_df, param, "dist_avg_strat_phe_0"
                )
            elif sim_type == SimType.RANDOM:
                plot_side_heatmap(axs_3[1], group_df, "dist_n_agents")
                plot_side_heatmap(axs_4[1], group_df, "dist_avg_strat_phe_0")
        else:
            plot_main_heatmap(
                fig_3, axs_3[0], axs_3[1], group_df, param, "dist_n_agents"
            )
            plot_main_heatmap(
                fig_4, axs_4[0], axs_4[1], group_df, param, "dist_avg_strat_phe_0"
            )

        plot_mean_and_uncertainty(ax_5, "avg_strat_phe_0", "std_dev_strat_phe")
        plot_mean_and_uncertainty(ax_6, "dist_phe_0", None)
        plot_mean_and_uncertainty(ax_7, "std_dev_growth_rate", None)
        plot_mean_and_uncertainty(ax_8, "avg_birth_rate", None)

        plot_errorbar(ax_9, group_df, "avg_growth_rate", "std_dev_growth_rate", True)

    if param == "strat_phe_0_i":
        plot_dist_phe_0_lims(ax_6, param_df, job)

    if param == "prob_mut":
        fixed_i_df = FILTERS["fixed_i"](df, job).sort_values("strat_phe_0_i")
        plot_errorbar(ax_2, fixed_i_df, "avg_growth_rate", "extinct_rate", True)
        plot_errorbar(ax_9, fixed_i_df, "avg_growth_rate", "std_dev_growth_rate", True)
        plot_expected_values(ax_0, df, job, param, "avg_growth_rate")
        plot_expected_values(ax_1, df, job, param, "extinct_rate")

    max_avg_growth = get_optimal_strat_phe_0(df, job, "avg_growth_rate", "max")
    min_extinct = get_optimal_strat_phe_0(df, job, "extinct_rate", "min")
    if param == "strat_phe_0_i":
        for ax in [ax_0, ax_1, ax_7]:
            ax.axvline(max_avg_growth, ls="--", **LINE_STYLE)
            ax.axvline(min_extinct, ls=":", **LINE_STYLE)
    for ax in axs_4[:-1] + [ax_5]:
        ax.axhline(max_avg_growth, ls="--", **LINE_STYLE)
        ax.axhline(min_extinct, ls=":", **LINE_STYLE)

    if param in ["prob_mut", "n_agents_i"]:
        for ax in [ax_0, ax_1, ax_5, ax_6, ax_7, ax_8]:
            ax.set_xscale("log")
    for ax in [ax_1, ax_2]:
        ax.set_yscale("log")

    for ax in [ax_0, ax_1, ax_2, ax_5, ax_6, ax_7, ax_8, ax_9]:
        ax.legend()

    fig_dir = job.base_dir / "plots" / param
    fig_dir.mkdir(parents=True, exist_ok=True)

    fig_0.savefig(fig_dir / "avg_growth_rate.pdf")
    fig_1.savefig(fig_dir / "extinct_rate.pdf")
    fig_2.savefig(fig_dir / "rates.pdf")
    fig_3.savefig(fig_dir / "dist_n_agents.pdf")
    fig_4.savefig(fig_dir / "dist_avg_strat_phe_0.pdf")
    fig_5.savefig(fig_dir / "avg_strat_phe_0.pdf")
    fig_6.savefig(fig_dir / "dist_phe_0.pdf")
    fig_7.savefig(fig_dir / "std_dev_growth_rate.pdf")
    fig_8.savefig(fig_dir / "avg_birth_rate.pdf")
    fig_9.savefig(fig_dir / "growth_rates.pdf")

    print_process_msg(f"made '{param}' plots")


def make_time_series_plots(df: pd.DataFrame, job: SimJob) -> None:
    fig_dir = job.base_dir / "plots" / "time_series"
    fig_dir.mkdir(parents=True, exist_ok=True)

    for y_col in ["n_agents", "n_extinct", "avg_strat_phe_0", "dist_phe_0"]:
        fig, ax = create_standard_figure("time", y_col)
        y_span_col = "std_dev_strat_phe" if y_col == "avg_strat_phe_0" else None
        plot_extinct_times(ax, df)
        plot_time_series(ax, df, y_col, y_span_col)
        ax.margins(0.0)
        ax.legend()
        fig.savefig(fig_dir / f"{y_col}.pdf")

    print_process_msg("made 'time_series' plots")


def make_fixed_plots(df: pd.DataFrame, job: SimJob) -> None:
    fixed_df = FILTERS["fixed"](df, job).sort_values("strat_phe_0_i")
    if fixed_df["n_agents_i"].nunique() < 2:
        return

    fig_0, axs_0 = create_colorbar_figure("strat_phe_0_i", "avg_growth_rate", False)
    fig_1, axs_1 = create_colorbar_figure("strat_phe_0_i", "extinct_rate", False)
    fig_2, axs_2 = create_colorbar_figure("strat_phe_0_i", "dist_phe_0", False)
    fig_3, axs_3 = create_colorbar_figure("strat_phe_0_i", "std_dev_growth_rate", False)
    fig_4, axs_4 = create_colorbar_figure("strat_phe_0_i", "avg_birth_rate", False)

    fig_5, axs_5 = create_colorbar_figure("n_agents_i", "extinct_rate", False)

    fig_6, axs_6 = create_colorbar_figure(
        "avg_strat_phe_0", "dist_avg_strat_phe_0", False
    )
    fig_7, axs_7 = create_colorbar_figure(
        "avg_strat_phe_0", "dist_avg_strat_phe_0", False
    )

    fig_8, axs_8 = create_colorbar_figure("n_agents_i", "prob_mut", False)
    fig_9, axs_9 = create_colorbar_figure("n_agents_i", "prob_mut", False)

    log_norm = get_heatmap_norm(
        "log", fixed_df["n_agents_i"].min() / 2, fixed_df["n_agents_i"].max() * 2
    )

    for n_agents_i, group_df in fixed_df.groupby("n_agents_i"):

        def plot_mean_and_uncertainty(ax: Axes, y_col: str) -> None:
            value = cast(float, n_agents_i)
            plot_colored_errorbar(ax, group_df, "strat_phe_0_i", y_col, log_norm, value)

        plot_mean_and_uncertainty(axs_0[0], "avg_growth_rate")
        plot_mean_and_uncertainty(axs_1[0], "extinct_rate")
        plot_mean_and_uncertainty(axs_2[0], "dist_phe_0")
        plot_mean_and_uncertainty(axs_3[0], "std_dev_growth_rate")
        plot_mean_and_uncertainty(axs_4[0], "avg_birth_rate")

        y = group_df["n_agents_i"] ** (
            -group_df[("avg_growth_rate", "mean")]
            / (4 * (group_df[("std_dev_growth_rate", "mean")] ** 2))
        )
        y *= float(group_df[("extinct_rate", "mean")].iloc[-1] / y.iloc[-1])
        axs_1[0].errorbar(group_df["strat_phe_0_i"], y, ls=":", **LINE_STYLE)

    for fig, axs in zip(
        [fig_0, fig_1, fig_2, fig_3, fig_4], [axs_0, axs_1, axs_2, axs_3, axs_4]
    ):
        set_heatmap_colorbar(fig, axs[1], "n_agents_i", log_norm)

    fixed_df = fixed_df.sort_values("n_agents_i")

    linear_norm = get_heatmap_norm("linear", 0, 1)

    for strat_phe_0_i, group_df in fixed_df.groupby("strat_phe_0_i"):
        value = cast(float, strat_phe_0_i)
        plot_colored_errorbar(
            axs_5[0], group_df, "n_agents_i", "extinct_rate", linear_norm, value
        )

    set_heatmap_colorbar(fig_5, axs_5[1], "strat_phe_0_i", linear_norm)

    for ax in [axs_5[0], axs_8[0], axs_9[0]]:
        ax.set_xscale("log")
    for ax in [axs_1[0], axs_5[0], axs_8[0], axs_9[0]]:
        ax.set_yscale("log")
    for ax in (axs_1[0], axs_5[0]):
        ax.set_ylim(1e-9, 1e-1)

    fixed_df = fixed_df.sort_values("strat_phe_0_i")
    max_N_df = fixed_df[fixed_df["n_agents_i"] == fixed_df["n_agents_i"].max()]
    avg_growth_rate = interpolate_values(axs_0[0], max_N_df, "avg_growth_rate")
    avg_birth_rate = interpolate_values(axs_4[0], max_N_df, "avg_birth_rate")

    random_df = FILTERS["random"](df, job).sort_values("n_agents_i")
    extinct_rates = interpolate_extinct_rates(axs_5[0], df, job)
    s_exp = get_strat_eval()

    log_norm = get_heatmap_norm(
        "log", random_df["n_agents_i"].min() / 2, random_df["n_agents_i"].max() * 2
    )

    avg_s_exp, avg_s = [], []
    for idx, (n_agents_i, group_df) in enumerate(random_df.groupby("n_agents_i")):
        extinct_rate = extinct_rates[idx]
        value = cast(float, n_agents_i)

        avg_s_exp_row, avg_s_row = [], []
        for prob_mut, subgroup_df in group_df.groupby("prob_mut"):
            log_p_s_exp: np.ndarray = (n_agents_i * avg_growth_rate) - np.log(
                extinct_rate + prob_mut * avg_birth_rate
            )
            p_s_exp = np.exp(log_p_s_exp - np.max(log_p_s_exp))
            p_s_exp /= np.sum(p_s_exp) / len(s_exp)
            avg_s_exp_row.append((s_exp * p_s_exp / len(s_exp)).sum())

            s, p_s = get_dist_avg_strat_phe_0(subgroup_df)
            s = np.array(s)
            p_s = np.array(p_s).ravel()
            avg_s_row.append((s * p_s / len(s)).sum())

            if prob_mut == job.config["model"]["prob_mut"]:
                plot_colored_curve(axs_6[0], s_exp, p_s_exp, log_norm, value)
                plot_colored_curve(axs_7[0], s, p_s, log_norm, value)

        avg_s_exp.append(avg_s_exp_row)
        avg_s.append(avg_s_row)

    for fig, axs in zip([fig_6, fig_7], [axs_6, axs_7]):
        set_heatmap_colorbar(fig, axs[1], "n_agents_i", log_norm)

    plot_avg_strat_phe_0(fig_8, axs_8[0], axs_8[1], random_df, np.array(avg_s_exp))
    plot_avg_strat_phe_0(fig_9, axs_9[0], axs_9[1], random_df, np.array(avg_s))

    fig_dir = job.base_dir / "plots" / "fixed"
    fig_dir.mkdir(parents=True, exist_ok=True)

    fig_0.savefig(fig_dir / "avg_growth_rate.pdf")
    fig_1.savefig(fig_dir / "extinct_rate.pdf")
    fig_2.savefig(fig_dir / "dist_phe_0.pdf")
    fig_3.savefig(fig_dir / "std_dev_growth_rate.pdf")
    fig_4.savefig(fig_dir / "avg_birth_rate.pdf")
    fig_5.savefig(fig_dir / "extinct_rate_scaling.pdf")
    fig_6.savefig(fig_dir / "exp_dist_avg_strat_phe_0.pdf")
    fig_7.savefig(fig_dir / "dist_avg_strat_phe_0.pdf")
    fig_8.savefig(fig_dir / "exp_avg_strat_phe_0.pdf")
    fig_9.savefig(fig_dir / "avg_strat_phe_0.pdf")

    print_process_msg("made 'fixed' plots")


def plot_sim_jobs(sim_jobs: list[SimJob]) -> None:
    avg_analyses = collect_avg_analyses(sim_jobs)
    job = sim_jobs[0]
    run_time_series = collect_run_time_series(job, 0)

    rmtree(job.base_dir / "plots", ignore_errors=True)

    with ProcessPoolExecutor(max_workers=N_CORES) as pool:
        futures = [
            pool.submit(make_param_plots, "strat_phe_0_i", avg_analyses, job),
            pool.submit(make_param_plots, "prob_mut", avg_analyses, job),
            pool.submit(make_param_plots, "n_agents_i", avg_analyses, job),
            pool.submit(make_time_series_plots, run_time_series, job),
            pool.submit(make_fixed_plots, avg_analyses, job),
        ]

        for future in as_completed(futures):
            future.result()
