from builtins import str

import matplotlib.pyplot as P
import numpy as np
import tables as T

################################################################
# 3D
################################################################


def pspec_rho(file, path_out, num, color="r", save=False, **kwargs):

    h5f = T.open_file(file, mode="r")
    group = h5f.get_node("/out_" + format(num, "05") + "/d3/rho")
    kbins = h5f.get_node(group, "kbins").read()
    pspec = h5f.get_node(group, "pspec").read()
    h5f.close()

    k = np.sqrt(kbins[1:] * kbins[:-1])
    pspec = (k ** 2) * pspec
    if save:
        P.figure(figsize=(8, 8))
    P.xlabel(r"$k$", fontsize=12)
    P.ylabel(r"$k^2 \mathcal{P}(\rho)$", fontsize=12)
    P.loglog(k, pspec, "-", color=color, **kwargs)
    if save:
        P.savefig(path_out + "pspec_rho_" + format(num, "05") + ".ps")
        P.savefig(path_out + "pspec_rho_" + format(num, "05") + ".jpeg")


def pspec_B(file, path_out, num, color="r", save=False, **kwargs):

    h5f = T.open_file(file, mode="r")
    group = h5f.get_node("/out_" + format(num, "05") + "/d3/B")
    kbins = h5f.get_node(group, "kbins").read()
    pspec = h5f.get_node(group, "pspec").read()
    h5f.close()

    k = np.sqrt(kbins[1:] * kbins[:-1])
    pspec = k ** 2 * pspec
    if save:
        P.figure(figsize=(8, 8))
    P.xlabel(r"$k$", fontsize=12)
    P.ylabel(r"$k^2 \mathcal{P}(B)$", fontsize=12)
    P.loglog(k, pspec, "-", color=color, **kwargs)
    if save:
        P.savefig(path_out + "pspec_B_" + format(num, "05") + ".ps")
        P.savefig(path_out + "pspec_B_" + format(num, "05") + ".jpeg")


def pspec_v(file, path_out, num, color="r", save=False, **kwargs):

    h5f = T.open_file(file, mode="r")
    group = h5f.get_node("/out_" + format(num, "05") + "/d3/v")
    kbins = h5f.get_node(group, "kbins").read()
    pspec = h5f.get_node(group, "pspec").read()
    h5f.close()

    k = np.sqrt(kbins[1:] * kbins[:-1])
    pspec = k ** 4 * pspec
    if save:
        P.figure(figsize=(8, 8))
    P.xlabel(r"$k$", fontsize=12)
    P.ylabel(r"$k^4 \mathcal{P}(v)$", fontsize=12)
    P.loglog(k, pspec, "-", color=color, **kwargs)
    if save:
        P.savefig(path_out + "pspec_v_" + format(num, "05") + ".ps")
        P.savefig(path_out + "pspec_v_" + format(num, "05") + ".jpeg")


def pspec_vz(file, path_out, num, color="r", save=False, **kwargs):

    h5f = T.open_file(file, mode="r")
    group = h5f.get_node("/out_" + format(num, "05") + "/d3/vz")
    kbins = h5f.get_node(group, "kbins").read()
    pspec = h5f.get_node(group, "pspec").read()
    h5f.close()

    k = np.sqrt(kbins[1:] * kbins[:-1])
    pspec = k ** 4 * pspec
    if save:
        P.figure(figsize=(8, 8))
    P.xlabel(r"$k$", fontsize=12)
    P.ylabel(r"$k^4 \mathcal{P}(vz)$", fontsize=12)
    P.loglog(k, pspec, "-", color=color, **kwargs)
    if save:
        P.savefig(path_out + "pspec_vz_" + format(num, "05") + ".ps")
        P.savefig(path_out + "pspec_vz_" + format(num, "05") + ".jpeg")


def pspec_cos(file, path_out, num, color="r", save=False, **kwargs):

    h5f = T.open_file(file, mode="r")
    group = h5f.get_node("/out_" + format(num, "05") + "/d3/cos_vB")
    kbins = h5f.get_node(group, "kbins").read()
    pspec = h5f.get_node(group, "pspec").read()
    h5f.close()

    k = np.sqrt(kbins[1:] * kbins[:-1])
    pspec = k ** 2 * pspec
    if save:
        P.figure(figsize=(8, 8))
    P.xlabel(r"$k$", fontsize=12)
    P.ylabel(r"$k^2 \mathcal{P}(\cos(\vec{v}.\vec{B}))$", fontsize=12)
    P.loglog(k, pspec, "-", color=color, **kwargs)
    if save:
        P.savefig(path_out + "pspec_cos_" + format(num, "05") + ".ps")
        P.savefig(path_out + "pspec_cos_" + format(num, "05") + ".jpeg")


def pspec_logrho(file, path_out, num, color="r", save=False, **kwargs):

    h5f = T.open_file(file, mode="r")
    group = h5f.get_node("/out_" + format(num, "05") + "/d3/logrho")
    kbins = h5f.get_node(group, "kbins").read()
    pspec = h5f.get_node(group, "pspec").read()
    h5f.close()

    k = np.sqrt(kbins[1:] * kbins[:-1])
    pspec = k ** 2 * pspec
    if save:
        P.figure(figsize=(8, 8))
    P.xlabel(r"$k$", fontsize=12)
    P.ylabel(r"$k^2 \mathcal{P}(\log(\rho))$", fontsize=12)
    P.loglog(k, pspec, "-", color=color, **kwargs)
    if save:
        P.savefig(path_out + "pspec_logrho_" + format(num, "05") + ".ps")
        P.savefig(path_out + "pspec_logrho_" + format(num, "05") + ".jpeg")


def pspec_kr(file, path_out, num, color="r", save=False, **kwargs):

    h5f = T.open_file(file, mode="r")
    group = h5f.get_node("/out_" + format(num, "05") + "/d3/kr")
    kbins = h5f.get_node(group, "kbins").read()
    pspec = h5f.get_node(group, "pspec").read()
    h5f.close()

    k = np.sqrt(kbins[1:] * kbins[:-1])
    pspec = (k ** 4) * pspec
    if save:
        P.figure(figsize=(8, 8))
    P.xlabel(r"$k$", fontsize=12)
    P.ylabel(r"$k^4 \mathcal{P}(\rho^{1/3}v)$", fontsize=12)
    P.loglog(k, pspec, "-", color=color, **kwargs)
    if save:
        P.savefig(path_out + "pspec_kr_" + format(num, "05") + ".ps")
        P.savefig(path_out + "pspec_kr_" + format(num, "05") + ".jpeg")


def pspec_vc(file, path_out, num, color="r", save=False, **kwargs):

    h5f = T.open_file(file, mode="r")
    group = h5f.get_node("/out_" + format(num, "05") + "/d3/vc")
    kbins = h5f.get_node(group, "kbins").read()
    pspec = h5f.get_node(group, "pspec").read()
    h5f.close()

    k = np.sqrt(kbins[1:] * kbins[:-1])
    pspec = (k ** 4) * pspec
    if save:
        P.figure(figsize=(8, 8))
    P.xlabel(r"$k$", fontsize=12)
    P.ylabel(r"$k^4 \mathcal{P}(v_c)$", fontsize=12)
    P.loglog(k, pspec, "-", color=color, **kwargs)
    if save:
        P.savefig(path_out + "pspec_vc_" + format(num, "05") + ".ps")
        P.savefig(path_out + "pspec_vc_" + format(num, "05") + ".jpeg")


def pspec_vs(file, path_out, num, color="r", save=False, **kwargs):

    h5f = T.open_file(file, mode="r")
    group = h5f.get_node("/out_" + format(num, "05") + "/d3/vs")
    kbins = h5f.get_node(group, "kbins").read()
    pspec = h5f.get_node(group, "pspec").read()
    h5f.close()

    k = np.sqrt(kbins[1:] * kbins[:-1])
    pspec = (k ** 4) * pspec
    if save:
        P.figure(figsize=(8, 8))
    P.xlabel(r"$k$", fontsize=12)
    P.ylabel(r"$k^4 \mathcal{P}(v_s)$", fontsize=12)
    P.loglog(k, pspec, "-", color=color, **kwargs)
    if save:
        P.savefig(path_out + "pspec_vs_" + format(num, "05") + ".ps")
        P.savefig(path_out + "pspec_vs_" + format(num, "05") + ".jpeg")


###################################################################
# 2D
###################################################################


def pspec_rho_2D(file, path_out, num, plan, color="r", save=False, **kwargs):

    h5f = T.open_file(file, mode="r")
    group = h5f.get_node("/out_" + format(num, "05") + "/d2/rho")
    kbins = h5f.get_node(group, "kbins_" + str(plan)).read()
    pspec = h5f.get_node(group, "pspec_" + str(plan)).read()
    h5f.close()

    k = np.sqrt(kbins[1:] * kbins[:-1])
    pspec = (k ** 1) * pspec
    if save:
        P.figure(figsize=(8, 8))
    P.xlabel(r"$k$", fontsize=12)
    P.ylabel(r"$k \mathcal{P}(\rho)$", fontsize=12)
    P.loglog(k, pspec, "-", color=color, **kwargs)
    if save:
        P.savefig(path_out + "pspec_rho_2D_" + format(num, "05") + ".ps")
        P.savefig(path_out + "pspec_rho_2D_" + format(num, "05") + ".jpeg")


def pspec_B_2D(file, path_out, num, plan, color="r", save=False, **kwargs):

    h5f = T.open_file(file, mode="r")
    group = h5f.get_node("/out_" + format(num, "05") + "/d2/B")
    kbins = h5f.get_node(group, "kbins_" + str(plan)).read()
    pspec = h5f.get_node(group, "pspec_" + str(plan)).read()
    h5f.close()

    k = np.sqrt(kbins[1:] * kbins[:-1])
    pspec = k ** 1 * pspec
    if save:
        P.figure(figsize=(8, 8))
    P.xlabel(r"$k$", fontsize=12)
    P.ylabel(r"$k \mathcal{P}(B)$", fontsize=12)
    P.loglog(k, pspec, "-", color=color, **kwargs)
    if save:
        P.savefig(path_out + "pspec_B_2D_" + format(num, "05") + ".ps")
        P.savefig(path_out + "pspec_B_2D_" + format(num, "05") + ".jpeg")


def pspec_v_2D(file, path_out, num, plan, color="r", save=False, **kwargs):

    h5f = T.open_file(file, mode="r")
    group = h5f.get_node("/out_" + format(num, "05") + "/d2/v")
    kbins = h5f.get_node(group, "kbins_" + str(plan)).read()
    pspec = h5f.get_node(group, "pspec_" + str(plan)).read()
    h5f.close()

    k = np.sqrt(kbins[1:] * kbins[:-1])
    pspec = k ** 3 * pspec
    if save:
        P.figure(figsize=(8, 8))
    P.xlabel(r"$k$", fontsize=12)
    P.ylabel(r"$k^3 \mathcal{P}(v)$", fontsize=12)
    P.loglog(k, pspec, "-", color=color, **kwargs)
    if save:
        P.savefig(path_out + "pspec_v_2D_" + format(num, "05") + ".ps")
        P.savefig(path_out + "pspec_v_2D_" + format(num, "05") + ".jpeg")


def pspec_vz_2D(file, path_out, num, plan, color="r", save=False, **kwargs):

    h5f = T.open_file(file, mode="r")
    group = h5f.get_node("/out_" + format(num, "05") + "/d2/vz")
    kbins = h5f.get_node(group, "kbins_" + str(plan)).read()
    pspec = h5f.get_node(group, "pspec_" + str(plan)).read()
    h5f.close()

    k = np.sqrt(kbins[1:] * kbins[:-1])
    pspec = k ** 3 * pspec
    if save:
        P.figure(figsize=(8, 8))
    P.xlabel(r"$k$", fontsize=12)
    P.ylabel(r"$k^3 \mathcal{P}(v_z)$", fontsize=12)
    P.loglog(k, pspec, "-", color=color, **kwargs)
    if save:
        P.savefig(path_out + "pspec_vz_2D_" + format(num, "05") + ".ps")
        P.savefig(path_out + "pspec_vz_2D_" + format(num, "05") + ".jpeg")


def pspec_cos_2D(file, path_out, num, plan, color="r", save=False, **kwargs):

    h5f = T.open_file(file, mode="r")
    group = h5f.get_node("/out_" + format(num, "05") + "/d2/cos_vB")
    kbins = h5f.get_node(group, "kbins_" + str(plan)).read()
    pspec = h5f.get_node(group, "pspec_" + str(plan)).read()
    h5f.close()

    k = np.sqrt(kbins[1:] * kbins[:-1])
    pspec = k * pspec
    if save:
        P.figure(figsize=(8, 8))
    P.xlabel(r"$k$", fontsize=12)
    P.ylabel(r"$k \mathcal{P}(\cos(\vec{v}.\vec{B}))$", fontsize=12)
    P.loglog(k, pspec, "-", color=color, **kwargs)
    if save:
        P.savefig(path_out + "pspec_cos_2D_" + format(num, "05") + ".ps")
        P.savefig(path_out + "pspec_cos_2D_" + format(num, "05") + ".jpeg")


def pspec_logrho_2D(file, path_out, num, plan, color="r", save=False, **kwargs):

    h5f = T.open_file(file, mode="r")
    group = h5f.get_node("/out_" + format(num, "05") + "/d2/logrho")
    kbins = h5f.get_node(group, "kbins_" + str(plan)).read()
    pspec = h5f.get_node(group, "pspec_" + str(plan)).read()
    h5f.close()

    k = np.sqrt(kbins[1:] * kbins[:-1])
    pspec = k * pspec
    if save:
        P.figure(figsize=(8, 8))
    P.xlabel(r"$k$", fontsize=12)
    P.ylabel(r"$k \mathcal{P}(\log(\rho))$", fontsize=12)
    P.loglog(k, pspec, "-", color=color, **kwargs)
    if save:
        P.savefig(path_out + "pspec_logrho_2D_" + format(num, "05") + ".ps")
        P.savefig(path_out + "pspec_logrho_2D_" + format(num, "05") + ".jpeg")


def pspec_kr_2D(file, path_out, num, plan, color="r", save=False, **kwargs):

    h5f = T.open_file(file, mode="r")
    group = h5f.get_node("/out_" + format(num, "05") + "/d2/kr")
    kbins = h5f.get_node(group, "kbins_" + str(plan)).read()
    pspec = h5f.get_node(group, "pspec_" + str(plan)).read()
    h5f.close()

    k = np.sqrt(kbins[1:] * kbins[:-1])
    pspec = (k ** 3) * pspec
    if save:
        P.figure(figsize=(8, 8))
    P.xlabel(r"$k$", fontsize=12)
    P.ylabel(r"$k^3 \mathcal{P}(\rho^{1/3}v)$", fontsize=12)
    P.loglog(k, pspec, "-", color=color, **kwargs)
    if save:
        P.savefig(path_out + "pspec_kr_2D_" + format(num, "05") + ".ps")
        P.savefig(path_out + "pspec_kr_2D_" + format(num, "05") + ".jpeg")


def pspec_vc_2D(file, path_out, num, plan, color="r", save=False, **kwargs):

    h5f = T.open_file(file, mode="r")
    group = h5f.get_node("/out_" + format(num, "05") + "/d2/vc")
    kbins = h5f.get_node(group, "kbins_" + str(plan)).read()
    pspec = h5f.get_node(group, "pspec_" + str(plan)).read()
    h5f.close()

    k = np.sqrt(kbins[1:] * kbins[:-1])
    pspec = (k ** 3) * pspec
    if save:
        P.figure(figsize=(8, 8))
    P.xlabel(r"$k$", fontsize=12)
    P.ylabel(r"$k^3 \mathcal{P}(v_c)$", fontsize=12)
    P.loglog(k, pspec, "-", color=color, **kwargs)
    if save:
        P.savefig(path_out + "pspec_vc_2D" + format(num, "05") + ".ps")
        P.savefig(path_out + "pspec_vc_2D" + format(num, "05") + ".jpeg")


def pspec_vs_2D(file, path_out, num, plan, color="r", save=False, **kwargs):

    h5f = T.open_file(file, mode="r")
    group = h5f.get_node("/out_" + format(num, "05") + "/d2/vs")
    kbins = h5f.get_node(group, "kbins_" + str(plan)).read()
    pspec = h5f.get_node(group, "pspec_" + str(plan)).read()
    h5f.close()

    k = np.sqrt(kbins[1:] * kbins[:-1])
    pspec = (k ** 3) * pspec
    if save:
        P.figure(figsize=(8, 8))
    P.xlabel(r"$k$", fontsize=12)
    P.ylabel(r"$k^3 \mathcal{P}(v_s)$", fontsize=12)
    P.loglog(k, pspec, "-", color=color, **kwargs)
    if save:
        P.savefig(path_out + "pspec_vs_2D" + format(num, "05") + ".ps")
        P.savefig(path_out + "pspec_vs_2D" + format(num, "05") + ".jpeg")