import matplotlib.pyplot as plt
from spin_image import SpinImage, FFT
import matplotlib


class Plotter:
    def __init__(self, lat):
        self.lattice = lat
        self.length_2 = 0.05

    def plot_spins(self, si: SpinImage, ax_log=None, ax_lin=None):
        if ax_log:
            t = ax_log.imshow(
                si.img,
                extent=(0, si.length_x, 0, si.length_y),
                norm=matplotlib.colors.LogNorm(vmin=1e-12),
                cmap="viridis",
                origin="lower"
            )
            plt.colorbar(t, ax=ax_log, extend="min")
        if ax_lin:
            t = ax_lin.imshow(
                si.img,
                extent=(0, si.length_x, 0, si.length_y),
                cmap="viridis",
                origin="lower"
            )
            plt.colorbar(t, ax=ax_lin, extend="min")

    def plot_fft(self, fft: FFT, ax_log=None, ax_lin=None, evaluator=None):
        if ax_log:
            if evaluator:
                evaluator.debug(fft)
            t = ax_log.imshow(
                fft.intens,
                extent=fft.extents(),
                norm=matplotlib.colors.LogNorm(),
                cmap="viridis",
                origin="lower"
            )
            plt.colorbar(t, ax=ax_log, extend="min")

        if ax_lin:
            t = ax_lin.imshow(
                fft.intens,
                extent=fft.extents(),
                cmap="viridis",
                origin="lower"
            )
            plt.colorbar(t, ax=ax_lin, extend="min")
            ax_lin.set_xlim(-2, 2)
            ax_lin.set_ylim(-2, 2)