import matplotlib.pyplot as plt
import numpy as np
from numpy.linalg import norm


class OrientationMap:
    A = 2.0
    B = 1.0
    lambda_1 = 2.0
    lambda_2 = 1.0
    max_length = 5.0
    w_scale = 0.01

    # def w(self,r):
    #     return self.A*np.exp(-self.lambda_1*r**2)-self.B*np.exp(-self.lambda_2*r**2)
    def w(self, r):
        corr_len = self.corr_len
        return self.w_scale * np.piecewise(r, [r < int(corr_len / 2), r >= int(corr_len / 2) and r < corr_len,
                                               r >= corr_len], [1.0, -1.0, 0.0])

    # def bound(self,z_i):
    #     return self.max_length - norm(z_i)
    def bound(self, z_i):
        return 1.0 if norm(z_i) < 1.0 else 0.0

    def __init__(self,x_dim,y_dim, corr_len, sheet_x=0, sheet_y=0, rnd_seed = -1):
        self.x_dim = x_dim
        self.y_dim = y_dim
        self.corr_len = int((corr_len/sheet_x)*x_dim) # in number of neurons?!?
        self.sheet_x = sheet_x
        self.sheet_y = sheet_y

        if rnd_seed != -1:
            np.random.seed(rnd_seed)
            self.rnd_seed = rnd_seed

        self.angle_grid = np.random.rand(x_dim,y_dim)*2.0*np.pi-np.pi
        self.vec_grid = np.ndarray((x_dim,y_dim,2))

        for idx in range(x_dim):
            for idy in range(y_dim):
                phi = self.angle_grid[idx, idy]
                vec = [0.01 * np.cos(phi), 0.01 * np.sin(phi)]
                self.vec_grid[idx, idy] = vec

    # def __init__(self,filename):
    #     self.angle_grid = np.load("../noise_maps/orientation_map.npy")

    def orientation_map_by_id(self,id_x,id_y):
        return self.angle_grid[id_x,id_y]

    def orientation_map(self,x,y):
        id_x = int(x * (self.x_dim - 1) / self.sheet_x)
        id_y = int(y * (self.y_dim - 1) / self.sheet_y)
        return self.angle_grid[id_x, id_y]

    # def dist(self,v_1,v_2):
    #     return np.norm(v_1 - v_2)

    def iterate(self):
        vec_grid = self.vec_grid
        vec_list = list([((i, j), vec_grid[i, j]) for i in range(vec_grid.shape[0]) for j in range(vec_grid.shape[1])])
        delta_z_list = []
        for vec_i in vec_list:
            (i_x, i_y), z_i = vec_i
            # print(i_x,i_y,z_i)
            d_z_i = np.array([0., 0.])
            for vec_j in vec_list:
                (j_x, j_y), z_j = vec_j
                if np.abs(i_x - j_x) < self.corr_len and np.abs(i_y - j_y) < self.corr_len and vec_i != vec_j:
                    r = np.sqrt((i_x - j_x) ** 2 + (i_y - j_y) ** 2)
                    d_z = z_j * self.w(r) * self.bound(z_i)
                    # print(d_z)
                    d_z_i += d_z
            delta_z_list.append(d_z_i)
            # print(d_z_i)
        for i, vec_i in enumerate(vec_list):
            (i_x, i_y), z_i = vec_i
            vec_grid[i_x, i_y] += delta_z_list[i]

    def improve(self, iterations):
        for i in range(iterations):
            self.iterate()
            print("Iteration {} out of {}".format(i + 1, iterations))
        vec_list = list([((i, j), self.vec_grid[i, j]) for i in range(self.vec_grid.shape[0]) for j in
                         range(self.vec_grid.shape[1])])
        for vec_i in vec_list:
            (i_x, i_y), z_i = vec_i
            self.angle_grid[i_x, i_y] = np.arctan2(z_i[1], z_i[0])

    def plot(self, ax=None):
        if ax is None:
            fig, ax = plt.subplots(1, 1)

        X, Y = self.get_meshgrid_of_neuron_positions()

        Z = np.zeros((self.x_dim, self.y_dim))

        # For correctly displaying ticks
        for y_idx in range(Z.shape[1]):
            for x_idx in range(Z.shape[0]):
                o_map_val = self.orientation_map(x_idx, y_idx)
                Z[x_idx, y_idx] = o_map_val
        if ax is None:
            fig = plt.figure()
            ax = fig.add_subplot(111)
        plt.set_cmap('twilight')
        pcm = ax.pcolormesh(X, Y, Z.T)
        for y_idx in range(Z.shape[1]):
            for x_idx in range(Z.shape[0]):
                vec = self.vec_grid[x_idx, y_idx]
                ax.quiver(x_idx, y_idx, vec[0], vec[1], scale=0.02, units="x", scale_units="x", width=0.09)
        plt.gcf().colorbar(pcm, ax=ax)

    def load_orientation_map(self):
        # try:
        self.angle_grid = np.load('../noise_maps/seeded_orientation_maps/map_{}x{}_pixels_{}_corr_pix_{}_seed.npy'.format(self.x_dim, self.y_dim, self.corr_len,self.rnd_seed))
        # except:
        #     print('No such map exists yet, generate and save it')

    def save_orientation_map(self):
        np.save('../noise_maps/map_{}x{}_pixels_{}_corr_pix.npy'.format(map.x_dim, map.y_dim, map.corr_len),
                map.angle_grid)

    def get_meshgrid_of_neuron_positions(self):
        x_max = self.sheet_x
        d_x = x_max / self.x_dim
        y_max = self.sheet_y
        d_y = y_max / self.y_dim
        X, Y = np.meshgrid(np.arange(0, x_max + d_x, d_x) - d_x / 2., np.arange(0, y_max + d_y, d_y) - d_y / 2.)
        return X, Y


    def get_neuron_positions(self):
        X, Y = self.get_meshgrid_of_neuron_positions()
        return list(zip(X.flatten(), Y.flatten()))

    def get_neuron_tunings(self):
        return [self.orientation_map_float(x,y ) for x,y in self.get_neuron_positions()]

