import sys

import matplotlib.pyplot as plt
import numpy as np
from util_simuls_regimes import get_sim

from fluidsim import load

from util import (
    compute_kf_kb_ko_keta_kd,
    customize,
    get_paths,
    pos_closest_value,
    save_fig,
)

print(sys.argv)
letter = sys.argv[-1]

if letter not in "DLOWP":
    letter = "P"


sim = get_sim(letter)

coef_compensate = 5 / 3

p_oper = sim.params.oper
kmax = p_oper.coef_dealiasing * p_oper.nx * np.pi / p_oper.Lx


t_start, t_last = sim.output.print_stdout.get_times_start_last()
tmin = t_last - 2.0


if letter == "P":
    fig, ax = plt.subplots()

    temp = sim.output.spectra.load3d_mean(tmin)
    print(temp)
    EK = temp["spectra_E"]
    EA = temp["spectra_A"]
    EKhd = temp["spectra_Khd"]
    EKhr = temp["spectra_Khr"]
    EKh = EKhr + EKhd
    EKz = EK - EKh
    Epolo = EKhd + EKz
    Etoro = EKhr

    k = temp["k"]
    ax.plot(
        k,
        EA * k**coef_compensate,
        "b-",
        label=(r"$E_{A}$"),
    )
    ax.plot(
        k,
        Epolo * k**coef_compensate,
        "m-",
        label=(r"$E_{polo}$"),
    )
    ax.plot(
        k,
        Etoro * k**coef_compensate,
        "r-",
        label=(r"$E_{toro}$"),
    )

    ax.legend(fontsize=10, loc="upper left")

    ax.set_xscale("log")
    ax.set_yscale("log")

    ax.set_xlabel(r"$k$", fontsize=16)
    ax.set_ylabel(r"$E(k) ~ k^{5/3}$", fontsize=16)
    # ax.text(1.1*kb, 1e-10, r"$k_b$", fontsize=16)
    # ax.text(1.1*ko, 1e-10, r"$k_O$", fontsize=16)

    ax.set_xlim([k[1], 0.8 * max(k)])
    ax.set_ylim(bottom=1e-4, top=1e1)

    fig.tight_layout()
    save_fig(fig, f"fig_spectra_slices_regime_{letter}_3D.png")


if True:
    kf, kb, ko, keta, kd = compute_kf_kb_ko_keta_kd(sim, tmin)

    data = sim.output.spectra.load_kzkh_mean(
        tmin, key_to_load=["A", "Khd", "Kz", "Khr"]
    )
    kh = data["kh_spectra"]
    kz = data["kz"]
    delta_kz = kz[1]
    delta_kh = kh[1]

    EA = data["A"]
    EKhd = data["Khd"]
    EKz = data["Kz"]
    EKhr = data["Khr"]
    Epolo = EKhd + EKz
    Etoro = EKhr
    EKh = EKhr + EKhd

    kstxts = ["kfkb", "kbko", "koketa"]
    kss = [0.5 * kb, (kb * ko) ** 0.5, (ko * keta) ** 0.5]

    for n in range(3):

        ks = kss[n]
        kstxt = kstxts[n]

        ikz = pos_closest_value(kz, ks)
        ikh = pos_closest_value(kh, ks)
        print(ks, kz[ikz], ikz)
        print(ks, kh[ikh], ikh)
        EA_kz = EA[:, ikh]
        Epolo_kz = Epolo[:, ikh]
        Etoro_kz = Etoro[:, ikh]
        EKz_kz = EKz[:, ikh]
        EKh_kz = EKh[:, ikh]

        EA_kh = EA[ikz, :]
        Epolo_kh = Epolo[ikz, :]
        Etoro_kh = Etoro[ikz, :]
        EKz_kh = EKz[ikz, :]
        EKh_kh = EKh[ikz, :]

        fig, ax = plt.subplots()  # figsize=(10, 3 * 4.5 / 2))

        if kstxt == "kfkb":
            kstitle = (
                r"$k_h=$"
                + f"{kh[ikh]/kb:.1f}"
                + r"$k_b ~ (--)$"
                + " or "
                + r"$k_z=$"
                + f"{kz[ikz]/kb:.1f}"
                + r"$k_b ~ (-)$"
            )
        elif kstxt == "kbko":
            kstitle = (
                r"$k_h=$"
                + f"{kh[ikh]/ko:.1f}"
                + r"$k_O ~ (--)$"
                + " or "
                + r"$k_z=$"
                + f"{kz[ikz]/ko:.1f}"
                + r"$k_O ~ (-)$"
            )
        else:
            kstitle = (
                r"$k_h=$"
                + f"{kh[ikh]/ko:.1f}"
                + r"$k_O ~ (--)$"
                + " or "
                + r"$k_z=$"
                + f"{kz[ikz]/ko:.1f}"
                + r"$k_O ~ (-)$"
            )

        ax.plot(
            kz,
            EA_kz * kz**coef_compensate,
            "b--",
            label=None,
        )
        ax.plot(
            kz,
            Epolo_kz * kz**coef_compensate,
            "m--",
            label=None,
        )
        ax.plot(
            kz,
            Etoro_kz * kz**coef_compensate,
            "r--",
            label=None,
        )
        """
        ax.plot(
            kz,
            EKz_kz * kz ** coef_compensate,
            "g--",
        )
        ax.plot(
            kz,
            EKh_kz * kz ** coef_compensate,
            "y--",
        )
        """

        ax.plot(
            kh,
            EA_kh * kh**coef_compensate,
            "b-",
            label=(r"$E_A$"),
        )
        ax.plot(
            kh,
            Epolo_kh * kh**coef_compensate,
            "m-",
            label=(r"$E_{polo}$"),
        )
        ax.plot(
            kh,
            Etoro_kh * kh**coef_compensate,
            "r-",
            label=(r"$E_{toro}$"),
        )
        """
        ax.plot(
            kh,
            EKz_kh * kh ** coef_compensate,
            "g-",
            label=(
                r"$E_{z}$"
            ),
        )
        ax.plot(
            kh,
            EKh_kh * kh ** coef_compensate,
            "y-",
            label=(
                r"$E_{h}$"
            ),
        )
        """

        if letter != "P":
            if kstxt == "kfkb":
                x = [delta_kz, kb]
                y = [
                    5e-3 * x[0] ** (-2 + coef_compensate),
                    5e-3 * x[1] ** (-2 + coef_compensate),
                ]
                ax.plot(x, y, "k-")
                ax.text(
                    (x[0] * x[1]) ** 0.5,
                    0.1 * 5e-3 * ((x[0] * x[1]) ** 0.5) ** (-2 + coef_compensate),
                    r"$k_h^{-2}$",
                    fontsize=16,
                )
                ax.legend(fontsize=10, loc="lower left")

            if kstxt == "kbko":
                x = [kb, ko]
                y = [
                    2e-4 * kb ** (-5 / 3 + coef_compensate),
                    2e-4 * ko ** (-5 / 3 + coef_compensate),
                ]
                ax.plot(x, y, "k-")
                ax.text(
                    ks,
                    0.1 * 2e-4 * ks ** (-5 / 3 + coef_compensate),
                    r"$k_h^{-5/3}$",
                    fontsize=16,
                )
                # ax.legend(fontsize=10, loc = "lower left")

            if kstxt == "koketa":
                x = [delta_kz, kb]
                y = [
                    1e-11 * delta_kz ** (1 + coef_compensate),
                    1e-11 * kb ** (1 + coef_compensate),
                ]
                ax.plot(x, y, "k-")
                ax.text(
                    (delta_kz * kb) ** 0.5,
                    0.2
                    * 1e-11
                    * ((delta_kz * kb) ** (0.5 * (1 + coef_compensate))),
                    r"$k_h^1$",
                    fontsize=16,
                )
                # ax.legend(fontsize=10, loc = "upper left")

            x = [delta_kz, ko]
            y = [1e-9 * delta_kz**coef_compensate, 1e-9 * ko**coef_compensate]
            ax.plot(x, y, "k--")
            ax.text(
                (x[0] * x[1]) ** 0.5,
                0.07 * 1e-9 * ((kb * ko) ** (0.5 * coef_compensate)),
                r"$k_z^0$",
                fontsize=16,
            )

            ax.axvline(kb, color="k", linestyle="dotted")
            ax.text(1.1 * kb, 2e-10, r"$k_b$", fontsize=16)
            ax.axvline(ko, color="k", linestyle="dashed")
            ax.text(1.1 * ko, 2e-10, r"$k_O$", fontsize=16)

        # ax.legend(fontsize=10, loc = "lower right")

        ax.set_xscale("log")
        ax.set_yscale("log")

        ax.set_xlabel(r"$k_h, k_z$", fontsize=16)
        ax.set_ylabel(
            r"$E_i ~ k_h^{5/3} ~ (-)$  or $E_i ~ k_z^{5/3} (--)$", fontsize=16
        )
        # ax.text(1.1*kb, 1e-10, r"$k_b$", fontsize=16)
        # ax.text(1.1*ko, 1e-10, r"$k_O$", fontsize=16)

        ax.set_xlim([kh[1], 0.8 * max(kh)])
        ax.set_ylim(bottom=1e-10, top=1e0)

        # plt.legend(loc=2, prop={'size': 10})

        ax.set_title(kstitle, fontsize=12)

        fig.tight_layout()
        save_fig(fig, f"fig_spectra_slices_regime_{letter}_{kstxt}.png")

if __name__ == "__main__":
    plt.show()