import ast

import matplotlib.pyplot as plt
import numpy as np
import seaborn as sns

sns.set_theme()
params = {
    "font.family": "Serif",
    "font.serif": "Roman",
    "text.usetex": True,
    "xtick.bottom": False,
    "xtick.labelbottom": False,
    "ytick.left": False,
    "ytick.labelleft": False,
    "axes.grid": False,
    # "figure.autolayout": True
    # "figure.constrained_layout.use": False,
    # "figure.constrained_layout.h_pad":  0.04167,
    # "figure.constrained_layout.w_pad":  0.04167
    # "figure.subplot.wspace": 0,
    # "figure.subplot.hspace": 0
}
plt.rcParams.update(params)


def plot():
    fig, axes = plt.subplots(figsize=(6,6), nrows=3, ncols=3)
    with open("data/color_map.txt") as f:
        lines = f.readlines()
        size = int(lines[0])
        for i, line in enumerate(lines[1:]):
            # Create figures for each plot
            # fig1, ax1 = plt.subplots()
            # fig2, ax2 = plt.subplots()
            # fig3, ax3 = plt.subplots()

            arr = line.strip().split("\t")
            arr = np.asarray(list(map(lambda x: complex(*ast.literal_eval(x)), arr)))
            # Reshape and transpose array
            arr = arr.reshape(size, size).T

            # Plot color maps
            color_map1 = axes[0, i].imshow(
                np.multiply(arr, arr.conj()).real,
                interpolation="nearest",
                cmap=sns.color_palette("mako", as_cmap=True)
            )
            color_map2 = axes[1, i].imshow(
                arr.real,
                interpolation="nearest",
                cmap=sns.color_palette("mako", as_cmap=True)
            )
            color_map3 = axes[2, i].imshow(
                arr.imag,
                interpolation="nearest",
                cmap=sns.color_palette("mako", as_cmap=True)
            )

            # Create color bar
    fig.subplots_adjust(wspace=0, hspace=0)
    # fig.colorbar(color_map1, ax=axes[0, 2])
    # fig.colorbar(color_map2, ax=axes[1, 2])
    # fig.colorbar(color_map3, ax=axes[2, 2])

    # Remove grids
    # for i in range(3):
    #     axes[i, 0].set_xticks([])
    #     axes[i, 0].set_yticks([])
    #     axes[i, 0].
    #     axes[i, 1].set_xticks([])
    #     axes[i, 1].set_yticks([])
    #     axes[i, 2].set_xticks([])
    #     axes[i, 2].set_yticks([])

    
    # axes[0].set_xticks([])
    # axes[0].set_yticks([])
    # axes[1].set_xticks([])
    # axes[].set_yticks([])
    # axes[].set_xticks([])
    # axes[].set_yticks([])

            # # Set labels
            # ax1.set_xlabel("x-axis")
            # ax1.set_ylabel("y-axis")
            # ax2.set_xlabel("x-axis")
            # ax2.set_ylabel("y-axis")
            # ax3.set_xlabel("x-axis")
            # ax3.set_ylabel("y-axis")

            # # Save the figures
    plt.savefig(f"latex/images/color_map_all.pdf", bbox_inches="tight")
            # fig2.savefig(f"latex/images/color_map_{i}_real.pdf", bbox_inches="tight")
            # fig3.savefig(f"latex/images/color_map_{i}_imag.pdf", bbox_inches="tight")

            # # Close figures
            # plt.close(fig1)
            # plt.close(fig2)
            # plt.close(fig3)


if __name__ == "__main__":
    plot()