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


class PinwheelMap:
    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):
        scale = self.scale
        return self.w_scale * np.piecewise(r, [r < int(scale / 2), r >= int(scale / 2) and r < scale,
                                               r >= scale], [1.0, -1.0, 0.0])

    def bound(self, z_i):
        return 1.0 if norm(z_i) < 1.0 else 0.0

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

        if seed != -1:
            np.random.seed(seed)
            self.seed = 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 pinwheel_map_by_id(self, id_x, id_y):
        return self.angle_grid[id_x, id_y]

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

    def tuning(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 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
            if self.bound(z_i) == 1.0:
                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.scale and np.abs(i_y - j_y) < self.scale and vec_i != vec_j:
                        r = np.sqrt((i_x - j_x) * (i_x - j_x) + (i_y - j_y) * (i_y - j_y))
                        d_z = z_j * self.w(r)
                        # print(d_z)
                        d_z_i += d_z
                delta_z_list.append(d_z_i)
            else:
                delta_z_list.append(np.array([0., 0.]))

        for i, vec_i in enumerate(vec_list):
            (i_x, i_y), z_i = vec_i
            if self.bound(z_i) == 1.0:
                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_map(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.pinwheel_map_by_id(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)
        plt.gcf().colorbar(pcm, ax=ax)

    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