#!/usr/bin/env python3
import numpy as np
import matplotlib.pyplot as plt
import utils
import pandas as pd
import itertools
import analytical


plt.style.use("../physrev.mplstyle")


def load_data(N, which):
    """Load data from file.

    Returns (avMM2, avMM3, delta, na_range, avg_see, avg_traces)
    """
    # avMM2 = np.load(f"analytic-trace2-N{N}.npy")
    # avMM3 = np.load(f"analytic-trace3-N{N}.npy")

    if which == 'entropy':
        file_name = f"/home/felipe/git-projects/syk-nonergodic/syk/data/entropy/N{N}/entropy.npz"
    elif which == 'moments':
        file_name = f"/home/felipe/git-projects/syk-nonergodic/syk/data/moments/N{N}/moments.npz"


    loaded = np.load(file_name)


    if N < 13:
        delta, avg = (
            loaded["delta"],
            loaded["avg_see"],
        )

    else:
        delta, avg = (
            loaded["delta"],
            loaded["avg"],
        )
    return delta, avg


def plot_entropy_na(neff, delta, avg_entropy):
    na_range = np.arange(1, neff//2+1)
    if neff == 12:
        delta_idx_list = [4, 7, 10]
    elif neff == 14:
        delta_idx_list = [4, 6, 8]
    else:
        delta_idx_list = [8, 12, 14]

    print('newp', delta[delta_idx_list])

    page_result = [analytical.page_entropy(neff, na) for na in na_range]
    # tobias_result = [analytical.tobias_entropy(neff, na) for na in na_range]

    fig, ax = plt.subplots()

    for delta_idx in delta_idx_list:
        ax.plot(
            na_range,
            avg_entropy[delta_idx, :],
            "+",
            linestyle="none",
            label=rf"$\delta = ${delta[delta_idx]}",
        )

    ax.plot(na_range, page_result, label="Page", color='black')
    # ax.plot(na_range, tobias_result, label="NEE")

    ax.set(xlabel='$N_A$',
           ylabel='$S_{EE}$',)
    ax.grid()
    ax.legend()

    fig.savefig(f"entropy-na-N{neff+1}.pdf")
    plt.clf()

    return None


def plot_entropy_n(n_range):
    fraction = 4
    n_range = np.asarray(n_range)
    neff_range = n_range - 1
    fig, ax = plt.subplots()

    y = np.zeros((len(n_range), 3))

    for i, n in enumerate(n_range):
        neff = n - 1
        na_range = np.arange(1, neff//2+1)
        na = neff/fraction
        remainder = neff%fraction

        if neff == 12:
            delta_idx_list = [4, 7, 10]
        elif neff == 14:
            delta_idx_list = [4, 6, 8]
        else:
            delta_idx_list = [8, 12, 14]

        if remainder:
            na_idx = np.flatnonzero(np.logical_or(na_range == np.floor(na), na_range == np.ceil(na)))
            delta, avg_entropy = load_data(n, which='entropy')
            print(n, na_range[na_idx], delta[delta_idx_list])
            y[i] = avg_entropy[delta_idx_list][:, na_idx].mean(axis=1)
        else:
            na_idx = np.flatnonzero(na_range == na)
            delta, avg_entropy = load_data(n, which='entropy')
            print(n, na_range[na_idx], delta[delta_idx_list])
            y[i] = avg_entropy[delta_idx_list][:, na_idx].reshape(-1)


    # x = np.arange(2, 20)
    x = n_range
    page_n = np.array([analytical.page_entropy(n, (n-1)/fraction) for n in x])
    page_neff = np.array([analytical.page_entropy(n-1, (n-1)/fraction) for n in x])
    ax.plot(x, page_n, '-', color='black', label='$Page N$')
    ax.plot(x, page_neff, '--', color='black', label='$Page N_{eff}$')

    for i, delta in enumerate([0.3, 1, 3]):
        ax.plot(n_range, y[:, i], "x", label=rf"$\delta = {delta}$")

    ax.set(xticks=n_range, xlim=[n_range[0]-2, n_range[-1]+2])
    ax.grid()
    handles, labels = plt.gca().get_legend_handles_labels()
    by_label = dict(zip(labels, handles))
    ax.legend(by_label.values(), by_label.keys())

    fig.savefig(f"entropy.pdf")
    plt.clf()

    return None


def plot_entropy_delta(neff, delta, avg_see):
    """Plot entropy as function of disorder.

    One plot for each size N.
    All N_A on same plot.
    """
    markercycle = itertools.cycle("+x^so")
    fig, ax = plt.subplots()
    namax = neff // 2
    for na in range(1, namax+1):
        page = analytical.page_entropy(neff, na)
        ax.plot(
            delta,
            avg_see[:, na-1] / page,
            marker=next(markercycle),
            markerfacecolor="none",
            linestyle="none",
            color=(0.6, 0, 0),
            label=fr"$N_A={na+1}$",
        )

    ax.set(
        title=fr"$N={neff+1}$",
        xlabel=r"$\delta$",
        xscale='log',
        xlim=[0.01, 200],
        ylabel=r"$S_{EE}/S_{Page}$",
        yticks=np.arange(0, 1.1, 0.1),
    )
    ax.legend(loc=3)
    ax.grid()
    fig.savefig(f"entropy-delta-N{neff+1}.pdf")
    plt.clf()

    return None


def plot_trace_moment_na(neff, m, delta, avg_traces):
    na_range = np.arange(1, neff//2+1)

    if neff == 12:
        idx_list = [1, 2, 3]
    elif neff == 14:
        idx_list = [4, 6, 8]

    page_result = [analytical.page_moment(neff, na, m) for na in na_range]
    tobias_result = [analytical.tobias_moment(neff, na, m) for na in na_range]
    fig, ax = plt.subplots()

    for delta_idx in idx_list:
        ax.plot(
            na_range,
            avg_traces[delta_idx, :, m - 2],
            "+",
            linestyle="none",
            label=rf"$\delta = ${delta[delta_idx]}",
        )

    ax.plot(na_range, page_result, label="Page")
    ax.plot(na_range, tobias_result, label="NEE")

    ax.set(xlabel=r"$N_A$",
           ylabel=r"$\langle \mathrm{tr} \rho_A^{%d}\rangle$" % m,
           yscale='log'
    )
    ax.grid()
    ax.legend(loc='lower left')

    fig.savefig(f"trace{m}-na-N{neff+1}.pdf")
    plt.clf()
    return None


def plot_trace_moment_delta(neff, m, delta, avg_traces):
    na_range = np.arange(1, neff//2 + 1)
    plot_ipr = True
    plot_kl = True
    fig, ax, = plt.subplots(nrows=1+plot_ipr+plot_kl, ncols=1, figsize=(8, 18))
    if plot_ipr:
        data_path = f'/home/felipe/git-projects/syk-nonergodic/syk/data/iq_vs_delta/masaki_iq_vs_delta_N{neff+1}.npz'
        loaded = np.load(data_path)
        delta_iq = np.sort(loaded['delta'])
        iqd = np.sort(loaded['IQD'])
        # std = loaded['std']
        ax[0].plot(delta_iq, iqd, '+-', color=(0.6, 0, 0))
        ax[0].set(
            title=r"$N={{{0}}}$".format(neff + 1),
            xscale='log', yscale='log', ylabel="IPR", xlabel=r"$\delta$", xlim=[0.01, 100])
        ax[0].grid()

    if plot_kl:
        data_path = f'/home/felipe/git-projects/syk-nonergodic/syk/data/kl_vs_delta/masaki_kl_vs_delta_N{neff+1}.npz'
        loaded = np.load(data_path)
        delta_kl = loaded['delta']
        kl_ps = loaded['kl_ps']
        kl_wd = loaded['kl_wd']
        ax[2].plot(delta_kl, kl_ps, marker='s', markerfacecolor='none', linestyle='--', color=(0.3, 0.3, 0.3), label='Poisson')
        ax[2].plot(delta_kl, kl_wd, marker='+', markerfacecolor='none', linestyle=':', color=(0.6, 0, 0), label='Wigner-Dyson')
        ax[2].set(xscale='log', xlabel=r"$\delta$", ylabel="KL", xlim=[0.01, 100])
        ax[2].grid()
        ax[2].legend()

    for na_idx, na in enumerate(na_range):
        ones = np.ones_like(delta)
        page_result = analytical.page_moment(neff, na, m) * ones
        tobias_result = analytical.tobias_moment(neff, na, m) * ones

        p = ax[1].plot(
            delta,
            avg_traces[:, na_idx, m-2],
            # color=(0.6, 0, 0),
            marker="+",
            linestyle="-",
            label=fr"$N_A={na}$",
        )

        clr = p[0].get_color()

        ax[1].plot(delta, page_result, linestyle="-", color=clr,
                # label=fr"Page, $N_A={na}$"
        )
        ax[1].plot(delta, tobias_result, linestyle="--", color=clr,
                # label=fr"Analytical, $N_A={na}$"
        )

    ax[1].set(
        xlabel=r"$\delta$",
        ylabel=r"$\langle \mathrm{{tr}} \rho_A^{{{0}}}\rangle$".format(m),
        xscale="log",
        yscale="log",
        xlim=[0.01, 100]
    )

    ax[1].legend(bbox_to_anchor=(1.04,1), loc="upper left")
    ax[1].grid()

    fig.savefig(f"trace{m}-delta-N{neff+1}.pdf", bbox_inches="tight")
    plt.clf()

    return None


if __name__ == "__main__":
    N_LIST = [5, 7, 9, 11, 13, 15]
    # for N in N_LIST:
        # neff = N - 1

        # delta, avg = load_data(N, which='moments')

        # m = 2
        # plot_trace_moment_delta(neff, m, delta, avg)

        # plot_trace_moment_na(neff, m, delta, avg)

        # m = 3
        # plot_trace_moment_delta(neff, m, delta, avg)

        # plot_trace_moment_na(neff, m, delta, avg)

        ####################

        # delta, avg = load_data(N, which='entropy')

        # plot_entropy_delta(neff, delta, avg)

        # plot_entropy_na(neff, delta, avg)

    plot_entropy_n(N_LIST)

#!/usr/bin/env python3
import numpy as np
import matplotlib.pyplot as plt
import utils
import pandas as pd
import itertools
import analytical


plt.style.use("../physrev.mplstyle")


def load_data(N, which):
    """Load data from file.

    Returns (avMM2, avMM3, delta, na_range, avg_see, avg_traces)
    """
    # avMM2 = np.load(f"analytic-trace2-N{N}.npy")
    # avMM3 = np.load(f"analytic-trace3-N{N}.npy")

    if which == 'entropy':
        file_name = f"/home/felipe/git-projects/syk-nonergodic/syk/data/entropy/N{N}/entropy.npz"
    elif which == 'moments':
        file_name = f"/home/felipe/git-projects/syk-nonergodic/syk/data/moments/N{N}/moments.npz"


    loaded = np.load(file_name)


    if N < 13:
        delta, avg = (
            loaded["delta"],
            loaded["avg_see"],
        )

    else:
        delta, avg = (
            loaded["delta"],
            loaded["avg"],
        )
    return delta, avg


def plot_entropy_na(neff, delta, avg_entropy):
    na_range = np.arange(1, neff//2+1)
    if neff == 12:
        delta_idx_list = [4, 7, 10]
    elif neff == 14:
        delta_idx_list = [4, 6, 8]
    else:
        delta_idx_list = [8, 12, 14]

    print('newp', delta[delta_idx_list])

    page_result = [analytical.page_entropy(neff, na) for na in na_range]
    # tobias_result = [analytical.tobias_entropy(neff, na) for na in na_range]

    fig, ax = plt.subplots()

    for delta_idx in delta_idx_list:
        ax.plot(
            na_range,
            avg_entropy[delta_idx, :],
            "+",
            linestyle="none",
            label=rf"$\delta = ${delta[delta_idx]}",
        )

    ax.plot(na_range, page_result, label="Page", color='black')
    # ax.plot(na_range, tobias_result, label="NEE")

    ax.set(xlabel='$N_A$',
           ylabel='$S_{EE}$',)
    ax.grid()
    ax.legend()

    fig.savefig(f"entropy-na-N{neff+1}.pdf")
    plt.clf()

    return None


def plot_entropy_n(n_range):
    fraction = 4
    n_range = np.asarray(n_range)
    neff_range = n_range - 1
    fig, ax = plt.subplots()

    y = np.zeros((len(n_range), 3))

    for i, n in enumerate(n_range):
        neff = n - 1
        na_range = np.arange(1, neff//2+1)
        na = neff/fraction
        remainder = neff%fraction

        if neff == 12:
            delta_idx_list = [4, 7, 10]
        elif neff == 14:
            delta_idx_list = [4, 6, 8]
        else:
            delta_idx_list = [8, 12, 14]

        if remainder:
            na_idx = np.flatnonzero(np.logical_or(na_range == np.floor(na), na_range == np.ceil(na)))
            delta, avg_entropy = load_data(n, which='entropy')
            print(n, na_range[na_idx], delta[delta_idx_list])
            y[i] = avg_entropy[delta_idx_list][:, na_idx].mean(axis=1)
        else:
            na_idx = np.flatnonzero(na_range == na)
            delta, avg_entropy = load_data(n, which='entropy')
            print(n, na_range[na_idx], delta[delta_idx_list])
            y[i] = avg_entropy[delta_idx_list][:, na_idx].reshape(-1)


    # x = np.arange(2, 20)
    x = n_range
    page_n = np.array([analytical.page_entropy(n, (n-1)/fraction) for n in x])
    page_neff = np.array([analytical.page_entropy(n-1, (n-1)/fraction) for n in x])
    ax.plot(x, page_n, '-', color='black', label='$Page N$')
    ax.plot(x, page_neff, '--', color='black', label='$Page N_{eff}$')

    for i, delta in enumerate([0.3, 1, 3]):
        ax.plot(n_range, y[:, i], "x", label=rf"$\delta = {delta}$")

    ax.set(xticks=n_range, xlim=[n_range[0]-2, n_range[-1]+2])
    ax.grid()
    handles, labels = plt.gca().get_legend_handles_labels()
    by_label = dict(zip(labels, handles))
    ax.legend(by_label.values(), by_label.keys())

    fig.savefig(f"entropy.pdf")
    plt.clf()

    return None


def plot_entropy_delta(neff, delta, avg_see):
    """Plot entropy as function of disorder.

    One plot for each size N.
    All N_A on same plot.
    """
    markercycle = itertools.cycle("+x^so")
    fig, ax = plt.subplots()
    namax = neff // 2
    for na in range(1, namax+1):
        page = analytical.page_entropy(neff, na)
        ax.plot(
            delta,
            avg_see[:, na-1] / page,
            marker=next(markercycle),
            markerfacecolor="none",
            linestyle="none",
            color=(0.6, 0, 0),
            label=fr"$N_A={na+1}$",
        )

    ax.set(
        title=fr"$N={neff+1}$",
        xlabel=r"$\delta$",
        xscale='log',
        xlim=[0.01, 200],
        ylabel=r"$S_{EE}/S_{Page}$",
        yticks=np.arange(0, 1.1, 0.1),
    )
    ax.legend(loc=3)
    ax.grid()
    fig.savefig(f"entropy-delta-N{neff+1}.pdf")
    plt.clf()

    return None


def plot_trace_moment_na(neff, m, delta, avg_traces):
    na_range = np.arange(1, neff//2+1)

    if neff == 12:
        idx_list = [1, 2, 3]
    elif neff == 14:
        idx_list = [4, 6, 8]

    page_result = [analytical.page_moment(neff, na, m) for na in na_range]
    tobias_result = [analytical.tobias_moment(neff, na, m) for na in na_range]
    fig, ax = plt.subplots()

    for delta_idx in idx_list:
        ax.plot(
            na_range,
            avg_traces[delta_idx, :, m - 2],
            "+",
            linestyle="none",
            label=rf"$\delta = ${delta[delta_idx]}",
        )

    ax.plot(na_range, page_result, label="Page")
    ax.plot(na_range, tobias_result, label="NEE")

    ax.set(xlabel=r"$N_A$",
           ylabel=r"$\langle \mathrm{tr} \rho_A^{%d}\rangle$" % m,
           yscale='log'
    )
    ax.grid()
    ax.legend(loc='lower left')

    fig.savefig(f"trace{m}-na-N{neff+1}.pdf")
    plt.clf()
    return None


def plot_trace_moment_delta(neff, m, delta, avg_traces):
    na_range = np.arange(1, neff//2 + 1)
    plot_ipr = True
    plot_kl = True
    fig, ax, = plt.subplots(nrows=1+plot_ipr+plot_kl, ncols=1, figsize=(8, 18))
    if plot_ipr:
        data_path = f'/home/felipe/git-projects/syk-nonergodic/syk/data/iq_vs_delta/masaki_iq_vs_delta_N{neff+1}.npz'
        loaded = np.load(data_path)
        delta_iq = np.sort(loaded['delta'])
        iqd = np.sort(loaded['IQD'])
        # std = loaded['std']
        ax[0].plot(delta_iq, iqd, '+-', color=(0.6, 0, 0))
        ax[0].set(
            title=r"$N={{{0}}}$".format(neff + 1),
            xscale='log', yscale='log', ylabel="IPR", xlabel=r"$\delta$", xlim=[0.01, 100])
        ax[0].grid()

    if plot_kl:
        data_path = f'/home/felipe/git-projects/syk-nonergodic/syk/data/kl_vs_delta/masaki_kl_vs_delta_N{neff+1}.npz'
        loaded = np.load(data_path)
        delta_kl = loaded['delta']
        kl_ps = loaded['kl_ps']
        kl_wd = loaded['kl_wd']
        ax[2].plot(delta_kl, kl_ps, marker='s', markerfacecolor='none', linestyle='--', color=(0.3, 0.3, 0.3), label='Poisson')
        ax[2].plot(delta_kl, kl_wd, marker='+', markerfacecolor='none', linestyle=':', color=(0.6, 0, 0), label='Wigner-Dyson')
        ax[2].set(xscale='log', xlabel=r"$\delta$", ylabel="KL", xlim=[0.01, 100])
        ax[2].grid()
        ax[2].legend()

    for na_idx, na in enumerate(na_range):
        ones = np.ones_like(delta)
        page_result = analytical.page_moment(neff, na, m) * ones
        tobias_result = analytical.tobias_moment(neff, na, m) * ones

        p = ax[1].plot(
            delta,
            avg_traces[:, na_idx, m-2],
            # color=(0.6, 0, 0),
            marker="+",
            linestyle="-",
            label=fr"$N_A={na}$",
        )

        clr = p[0].get_color()

        ax[1].plot(delta, page_result, linestyle="-", color=clr,
                # label=fr"Page, $N_A={na}$"
        )
        ax[1].plot(delta, tobias_result, linestyle="--", color=clr,
                # label=fr"Analytical, $N_A={na}$"
        )

    ax[1].set(
        xlabel=r"$\delta$",
        ylabel=r"$\langle \mathrm{{tr}} \rho_A^{{{0}}}\rangle$".format(m),
        xscale="log",
        yscale="log",
        xlim=[0.01, 100]
    )

    ax[1].legend(bbox_to_anchor=(1.04,1), loc="upper left")
    ax[1].grid()

    fig.savefig(f"trace{m}-delta-N{neff+1}.pdf", bbox_inches="tight")
    plt.clf()

    return None


if __name__ == "__main__":
    N_LIST = [5, 7, 9, 11, 13, 15]
    # for N in N_LIST:
        # neff = N - 1

        # delta, avg = load_data(N, which='moments')

        # m = 2
        # plot_trace_moment_delta(neff, m, delta, avg)

        # plot_trace_moment_na(neff, m, delta, avg)

        # m = 3
        # plot_trace_moment_delta(neff, m, delta, avg)

        # plot_trace_moment_na(neff, m, delta, avg)

        ####################

        # delta, avg = load_data(N, which='entropy')

        # plot_entropy_delta(neff, delta, avg)

        # plot_entropy_na(neff, delta, avg)

    plot_entropy_n(N_LIST)

#!/usr/bin/env python3 Load data from file.      Returns (avMM2, avMM3, delta, na_range, avg_see, avg_traces) # avMM2 = np.load(f"analytic-trace2-N{N}.npy") # avMM3 = np.load(f"analytic-trace3-N{N}.npy") # tobias_result = [analytical.tobias_entropy(neff, na) for na in na_range] # ax.plot(na_range, tobias_result, label="NEE") # x = np.arange(2, 20) Plot entropy as function of disorder.      One plot for each size N.     All N_A on same plot. # std = loaded['std'] # color=(0.6, 0, 0), # label=fr"Page, $N_A={na}$" # label=fr"Analytical, $N_A={na}$" # for N in N_LIST: # neff = N - 1 # delta, avg = load_data(N, which='moments') # m = 2 # plot_trace_moment_delta(neff, m, delta, avg) # plot_trace_moment_na(neff, m, delta, avg) # m = 3 # plot_trace_moment_delta(neff, m, delta, avg) # plot_trace_moment_na(neff, m, delta, avg) #################### # delta, avg = load_data(N, which='entropy') # plot_entropy_delta(neff, delta, avg) # plot_entropy_na(neff, delta, avg)