import numpy as np
import euler_python_my as euler
import matplotlib.pylab as plt


class EulerImagery:
    def __init__(self, outcoord, outdata, configs):
        self.xi = outcoord['easting'].data
        self.yi = outcoord['northing'].data
        self.outdata = outdata
        self.configs = configs
        self.est_classic = []

    def handle_data(self):
        area = self.get_coord_area()
        shape = self.outdata.shape
        xx = self.xi.repeat(shape[0])
        xx = xx.reshape(shape[1], shape[0]).T
        yy = self.yi.repeat(shape[1])
        yy = yy.reshape(shape)
        zi = np.ones_like(self.outdata) * 20
        zz = zi.reshape(self.outdata.size)
        zz = zz.reshape(shape)
        np.nan_to_num(self.outdata, nan=-100, copy=False)
        for SI in self.configs.get("SI_VET"):
            classic = euler.euler_deconv(self.outdata, xx, yy, zz, shape, area,
                                         SI, self.configs.get("WINSIZE"), self.configs.get("FILT"))
            self.est_classic.append(classic)
        plot_classic(self.outdata, self.est_classic, self.xi, self.yi)

    def get_coord_area(self):
        x_min = int(self.xi.min())
        x_max = int(self.xi.max()) + 1
        y_min = int(self.yi.min())
        y_max = int(self.yi.max()) + 1
        return [y_min, y_max, x_min, x_max]


def plot_classic(data, est_classic, xi, yi):
    '''
    Classic plot of the depth and base level estimates for all SIs
    '''

    '''
    Figure 4 source-position (depth) estimates
    '''
    vet_title = ["(a)", "(b)", "(c)", "(d)"]

    minz = np.min(-10.)
    maxz = np.max(0.)
    levelsz = np.linspace(minz, maxz, 11)

    # depth plots
    # plt.figure(figsize=(12, 8.5))
    # for i in range(4):
    #     plt.subplot(2, 2, i + 1)
    #     plt.title(vet_title[i], fontsize=14, loc='center', y=-0.27)
    #     plt.pcolormesh(xi, yi, data, vmin=-250, vmax=50)
    #     ax = plt.gca()
    #     ax.set_ylabel('Northing (m)', fontsize=14)
    #     ax.set_xlabel('Easting (m)', fontsize=14)
    #     ax.tick_params(labelsize=13)
    #     scat = plt.scatter(est_classic[i][:, 0],
    #                        est_classic[i][:, 1], s=40,
    #                        c=(est_classic[i][:, 2]),
    #                        cmap='terrain_r', vmin=minz, vmax=maxz,
    #                        edgecolors='k')
    #     cbar = plt.colorbar(scat, ticks=levelsz, pad=0.01, shrink=1,
    #                         format='%0.1f')
    #     cbar.set_label('$\^z_o$ (m)', labelpad=-18, y=-0.03, rotation=0,
    #                    fontsize=13)
    #     cbar.ax.tick_params(labelsize=13)
    #     ax.set_xlim(np.min(xi), np.max(xi))
    #     ax.set_ylim(np.min(yi), np.max(yi))
    #
    # plt.subplots_adjust(wspace=0.15, hspace=0.32)
    # plt.show()
    # plt.savefig('figures/FIG4.png',bbox_inches='tight', dpi = 600)
    # plt.close('all')

    '''
    Figure 7 base level-position (base-level) estimates
    '''
    plt.figure(figsize=(12, 8.5))
    for i in range(4):

        if i == 0:
            # base level estimates for SI = 0 have higher amplitude
            minb = np.min(-20)
            maxb = np.max(20)
            levelsb = np.linspace(minb, maxb, 7)

            plt.subplot(2, 2, i + 1)
            plt.title(vet_title[i], fontsize=14, loc='center',
                      y=-0.27)
            plt.pcolormesh(xi, yi, data, vmin=-250, vmax=50)
            ax = plt.gca()
            ax.set_ylabel('Northing (m)', fontsize=14)
            ax.set_xlabel('Easting (m)', fontsize=14)
            ax.tick_params(labelsize=13)
            scat = plt.scatter(est_classic[i][:, 0],
                               est_classic[i][:, 1], s=40,
                               c=(est_classic[i][:, 3]),
                               cmap='jet', vmin=minb,
                               vmax=maxb, edgecolors='k')
            cbar = plt.colorbar(scat, ticks=levelsb, pad=0.01, shrink=1,
                                format='%d')
            cbar.set_label('$\^b$ (nT)', labelpad=-18, y=-0.03, rotation=0,
                           fontsize=13)
            cbar.ax.tick_params(labelsize=13)
            ax.set_xlim(np.min(xi), np.max(xi))
            ax.set_ylim(np.min(yi), np.max(yi))


        else:
            minb = np.min(-20)
            maxb = np.max(20)
            levelsb = np.linspace(minb, maxb, 7)

            plt.subplot(2, 2, i + 1)
            plt.title(vet_title[i], fontsize=14, loc='center',
                      y=-0.27)
            plt.pcolormesh(xi, yi, data, vmin=-250, vmax=50)
            ax = plt.gca()
            ax.set_ylabel('Northing (m)', fontsize=14)
            ax.set_xlabel('Easting (m)', fontsize=14)
            ax.tick_params(labelsize=13)
            scat = plt.scatter(est_classic[i][:, 0],
                               est_classic[i][:, 1], s=40,
                               c=(est_classic[i][:, 3]), cmap='jet',
                               vmin=minb, vmax=maxb, edgecolors='k')
            cbar = plt.colorbar(scat, ticks=levelsb, pad=0.01, shrink=1,
                                format='%d')
            cbar.set_label('$\^b$ (nT)', labelpad=-18, y=-0.03, rotation=0,
                           fontsize=13)
            cbar.ax.tick_params(labelsize=13)
            ax.set_xlim(np.min(xi), np.max(xi))
            ax.set_ylim(np.min(yi), np.max(yi))

    plt.subplots_adjust(wspace=0.15, hspace=0.32)
    plt.show()