#!/usr/bin/env python3
import numpy as np
import matplotlib.pyplot as plt
from alive_progress import alive_bar

img_res_x = 100
img_res_y = 100
total_pixels = img_res_x * img_res_y # so we don't gotta compute it every time

periods = 1
square_x = 0
square_y = 0

#xmin = (-periods * np.pi) + (square_x * np.pi)
#xmax = (periods  * np.pi) + (square_x * np.pi)
#ymin = (-periods * np.pi) + (square_y * np.pi)
#ymax = (periods * np.pi) + (square_y * np.pi)

xmin = -10
xmax = 10
ymin = -10
ymax = 10

escape = 10000
iterations = 255*3
c_x = 2 * np.pi
c_y = 2 * np.pi


image = np.empty([img_res_y, img_res_x])
grid = np.meshgrid(np.linspace(ymin, ymax, img_res_y), np.linspace(xmin, xmax, img_res_x))
print(grid[0].dtype)


class point_charge():
  def __init__(self, x, y, c, mod):
    self.x = x
    self.y = y
    self.c = c
    self.mod = mod
  def get_field(self, to_x, to_y):
    if(self.mod):
      to_x = (to_x % self.mod)
      to_y = (to_y % self.mod)
    return (
        ((self.c * (self.x - to_x)) / ((self.x - to_x)**2 + (self.y - to_y)**2)**1.5),
        ((self.c * (self.y - to_y)) / ((self.x - to_x)**2 + (self.y - to_y)**2)**1.5))

#will remove all the point charge code if it turns out to be good enough to be impliemnted into openCL
#point_charges = [point_charge(-5, -5, 100), point_charge(-5, 5, -100), point_charge(5, 0, 100)]
point_charges = [point_charge(5,5, 100, 10), point_charge(0,0,-100, 0)]


plt.ion()
ax = plt.gca()
fig = plt.gcf()
ax.set_autoscale_on(False)
ax.set_xlim([xmin, xmax])
ax.set_ylim([ymin, ymax])

vector_arrows = None

def show_field():
  global vector_arrows
  grid_f = np.zeros_like(grid)
  for p in point_charges:
    grid_f += p.get_field(grid[0], grid[1])
  #plt.streamplot(grid[0], grid[1], grid_f[0], grid_f[1], density=5)
  vector_arrows = plt.quiver(grid[0], grid[1], grid_f[0], grid_f[1])
  plt.show(block=False)
  plt.pause(.1)


show_field()

timestep = .1
def test_sim():
  particle_grid = np.meshgrid(np.linspace(ymin, ymax, 100), np.linspace(xmin, xmax, 100))
  pos = particle_grid
  acceleration = np.zeros_like(particle_grid)
  velocity = np.zeros_like(particle_grid)
  velocity = [np.ones_like(particle_grid[0]) * 1, np.ones_like(particle_grid[0]) * .5]
  mass = 10
  charge = 1
  particle_plot = ax.plot(velocity[0], velocity[1], 'bo', animated=True)
  #velocity += .1

  background = fig.canvas.copy_from_bbox(ax.bbox)
  ax.draw_artist(vector_arrows)
  fig.canvas.blit(fig.bbox)
  
  while True:
    fig.canvas.restore_region(background)
    field = np.zeros_like(particle_grid)
    # TODO can make this quicker by skipping initilization
    for p in point_charges:
      field += p.get_field(pos[0], pos[1])
    acceleration = ((charge * field) / mass) * timestep
    #print(acceleration)
    velocity += acceleration * timestep
    pos += velocity * timestep

    fig.canvas.restore_region(background)
    particle_plot[0].set_data(pos[0],pos[1])
    ax.draw_artist(particle_plot[0])
    fig.canvas.blit(fig.bbox)
    fig.canvas.flush_events()
    plt.pause(1/60)

    #fig.canvas.draw_idle()
test_sim()

exit(1)



    

#with alive_bar(iterations, bar = 'filling', spinner = 'waves') as bar:
#  for i in range(iterations):
#    next_x = xx / np.sin(yy)
#    yy = yy / np.sin(xx)
#    xx = next_x
#    bar()
#image = np.vstack([xx.ravel(), yy.ravel()])
  

#meshgrid makes things slower as we can't test individual points for breaking to infinity
fractal_test = False
if fractal_test:
  with alive_bar(img_res_y, bar = 'filling', spinner = 'waves') as bar:
    for pix_y, y in enumerate(np.linspace(ymin, ymax, img_res_y)):
      for pix_x, x in enumerate(np.linspace(xmin, xmax, img_res_x)):
        on_x = x
        on_y = y
        for i in range(iterations):
          completed_ratio = (((pix_x * pix_y * 1)) / total_pixels)
          next_x = (completed_ratio * (on_x/np.sin(on_y))) + ((1 - completed_ratio) * on_x/np.tan(on_y))
          on_y = (completed_ratio * (on_y/np.sin(on_x))) + ((1 - completed_ratio) * on_y/np.tan(on_x))
          on_x = next_x
          if on_x**2 + on_y**2 > escape:
            break
        image[pix_y][pix_x] = i
      bar()
else:
  exit()


exit(1)      

plt.style.use('dark_background')
# fuck this shit
fig = plt.figure(frameon=False)
fig.set_size_inches(img_res_x/fig.dpi, img_res_y/fig.dpi)
#fig.set_size_inches(width/height, 1, forward=False)

ax = plt.Axes(fig, [0., 0., 1., 1.])
ax.set_axis_off()
fig.add_axes(ax)

cmap = plt.cm.viridis
cmap.set_bad((0,0,0))
cmap.set_over((0,0,0))
cmap.set_under((0,0,0))

ax.imshow(image, norm="log", aspect="auto", cmap=cmap)
fig.savefig("linear_transform_sin_tan_arnolds_tongue_hotspot.png")
plt.show()