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

sns.set_theme()

def plot():
    with open("data/color_map.txt") as f:
        lines = f.readlines();
        size = int(lines[0])
        for i, line in enumerate(lines[1:]):
            fig, ax = plt.subplots()
            arr = line.strip().split("\t")
            arr = np.asarray(list(map(lambda x: ((a := complex(*ast.literal_eval(x)))*a.conjugate()).real, arr)))

            arr = arr.reshape(size,size)

            ax.imshow(arr.T, interpolation="nearest")
            plt.savefig(f"latex/images/color_map_{i}.pdf")
            plt.close(fig)


if __name__ == "__main__":
    plot()