import noise

import numpy as np

import matplotlib.pyplot as plt

import random as rng

from matplotlib.patches import Ellipse

class Pickle:

    def __init__(self,x,y,r_in,r_out,phi):
        self.x = x
        self.y = y
        self.r_in = r_in
        self.r_out = r_out
        self.phi = phi

    def get_p1(self):
        x1 = self.x - self.r_in / 2.0 * np.cos(self.phi)
        y1 = self.y - self.r_in / 2.0 * np.sin(self.phi)
        return x1,y1

    def get_p2(self):
        x2 = self.x + self.r_in / 2.0 * np.cos(self.phi)
        y2 = self.y + self.r_in / 2.0 * np.sin(self.phi)
        return x2,y2

    def get_ellipse(self):
        c = self.r_in+self.r_out

        # Compute ellipse parameters
        a = c / 2  # Semimajor axis
        b = np.sqrt(a ** 2 - (self.r_in/2.0) ** 2)  # Semiminor axis
        ell = Ellipse((self.x, self.y), 2*a, 2*b,
                     angle=self.phi*180./np.pi, linewidth=1, fill=False, zorder=2)
        return ell

grid_x = 30
grid_y = 30
max_x = 8
max_y = 8


rng.seed(1)
pickle_r_in = 10.0
pickle_r_out = 2.0
pickle_n = 20

x_rng_min = pickle_r_out+0.5*pickle_r_in
x_rng_max = grid_x-pickle_r_out-0.5*pickle_r_in
y_rng_min = pickle_r_out+0.5*pickle_r_in
y_rng_max = grid_y-pickle_r_out-0.5*pickle_r_in

# Start with random pickle distribution
pickles = []
for n in range(pickle_n):
    pickles.append(Pickle(rng.uniform(x_rng_min,x_rng_max), \
                          rng.uniform(y_rng_min,y_rng_max), \
                          pickle_r_in,pickle_r_out,rng.random()*2.*np.pi))

x = np.linspace(0,max_x, grid_x)
y = np.linspace(0,max_y, grid_y)

head_dir_preference = np.zeros((grid_x, grid_y))

for x_idx, xx in enumerate(x):
    for y_idx, yy in enumerate(y):
        head_dir_preference[x_idx, y_idx] = noise.pnoise2(xx, yy, octaves=2,  )

# for pickle in pickles:
#     # print(pickle.phi)
#     x1,y1 = pickle.get_p1()
#     x2, y2 = pickle.get_p2()
#     # print('x,y: ',pickle.x,pickle.y)
#     # print('x1,y1,x2,y2: ',x1,y1,x2,y2)
#     # print('As Int x1,y1,x2,y2: ', int(x1),int(y1), int(x2),int(y2))
#     head_dir_preference[int(x1),int(y1)] = 2.0
#     head_dir_preference[int(x2),int(y2)] = 2.0
def plotting():
    fig = plt.figure(0)
    ax = fig.add_subplot(111)
    # plt.set_cmap('twilight')
    ax.imshow(head_dir_preference.T, vmin=-1, vmax=1, origin='lower')
    # plt.colorbar()
    for pickle in pickles:
        ell = pickle.get_ellipse()
        # print(ell)
        ax.add_artist(ell)
    plt.show()

plotting()

n_iterations = 10000
f_scale = 0.1

for i in range(n_iterations):

    for p_i in pickles:
        for p_j in pickles:
            if p_i is not p_j:
                f_p_i = f_scale*1/((p_i.x-p_j.x)**2+(p_i.y-p_j.y)**2)**2
                # better would be simultaniously
                if p_i.x + f_p_i * (p_i.x - p_j.x) < grid_x-pickle_r_out \
                    and p_i.x + f_p_i * (p_i.x - p_j.x) > pickle_r_out:
                    p_i.x += f_p_i * (p_i.x - p_j.x)
                if p_i.y + f_p_i * (p_i.y - p_j.y) < grid_y - pickle_r_out \
                        and p_i.y + f_p_i * (p_i.y - p_j.y) > pickle_r_out:
                    p_i.y += f_p_i * (p_i.y - p_j.y)
plotting()
