#!/usr/bin/env python3
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
import pyopencl as cl
from alive_progress import alive_bar
from matplotlib.backend_bases import MouseButton
from PIL import Image

img_res_x = 1000
img_res_y = 1000
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

x_min = (-periods * np.pi) + (square_x * np.pi)
x_max = (periods  * np.pi) + (square_x * np.pi)
y_min = (-periods * np.pi) + (square_y * np.pi)
y_max = (periods * np.pi) + (square_y * np.pi)

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

animation_progres_save = "ani1_less.mp4"
single_frame_save = "asdf.png"

opencl_context = cl.create_some_context(interactive=False)
opencl_queue = cl.CommandQueue(opencl_context)

kernel_src_path = "./kernel.c"

frames = 1
rendered_frames = []

image = np.empty([img_res_x, img_res_y], np.uint32)

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))

mask_path = "mask.png"

bruh = None

def on_click(event):
  global bruh
  split_ratio = 1
  if (event.button is MouseButton.MIDDLE) and (event.inaxes):
      # there's probably a way to set global coordinates;
      # I don't expect this to go anywhere so I don't really care
      on_x = ((event.xdata / img_res_x) * abs(x_max - x_min)) + x_min
      on_y = ((event.ydata / img_res_y) * abs(y_max - y_min)) + y_min

      #I, uh, also don't know the best way to replicate the openCL code automaticly in python
      #so ajust if nessesary
      x_hops = []
      y_hops = []
      for i in range(iterations):
          x_hops.append(((on_x - x_min) / abs(x_max - x_min)) * img_res_x)
          y_hops.append(((on_y - y_min) / abs(y_max - y_min)) * img_res_y)
          next_x = on_x/np.tan(on_y)
          on_y = on_y/np.tan(on_x)
          on_x = next_x
          if on_x**2 + on_y**2 > escape:
              break
      print(y_hops[0])
      print(on_y)
      print("{} hops".format(len(x_hops)))
      if bruh:
          bruh.pop(0).remove()
      bruh = ax.plot(x_hops, y_hops)
      plt.draw()



      print(on_x, on_y)


print("compiling openCL kernel...")
with open(kernel_src_path, 'r') as kernel_src:
  compiled_kernel = cl.Program(opencl_context, kernel_src.read()).build()

encoding_progress = alive_bar(frames, bar = 'filling', spinner = 'waves')

def display_encoder_progress(current_frame: int, total_frames: int):
  print("Encoding: frame {}/{}".format(current_frame, frames))

#open image
mask = np.asarray(Image.open(mask_path).convert("L"), dtype=np.double)
mask.setflags(write=1)
mask /= np.max(mask) # normalize
print(mask.shape)
print(mask.dtype)
#ax.imshow(mask, norm="linear")
#plt.show()

image_buffer = cl.Buffer(opencl_context, cl.mem_flags.WRITE_ONLY, image.nbytes)
mask_buffer = cl.Buffer(opencl_context, cl.mem_flags.READ_ONLY, mask.nbytes)
cl.enqueue_copy(opencl_queue, mask_buffer, mask).wait()
#mask_buffer = cl.array.Array(opencl_context,mask.shape, dtype=np.uint8, data=mask)




# TODO clean this up
print("Rendering {} frames...".format(frames))
if frames > 1:
  with alive_bar(frames, bar = 'filling', spinner = 'waves') as bar_total:
      for frame_i in range(0, frames):
          compiled_kernel.render_frame(opencl_queue, image.shape, None, image_buffer,
                                       np.double(abs(x_max - x_min) / img_res_x),
                                       np.double(abs(y_max - y_min) / img_res_y),
                                       np.double(x_min), np.double(y_min),
                                       np.uint32(iterations), np.uint32(escape),
                                       np.double(frame_i / frames))


          cl.enqueue_copy(opencl_queue, image, image_buffer).wait()
          rendered_frame = ax.imshow(image, norm="log", aspect="auto", cmap=cmap, animated="True")
          rendered_frames.append([rendered_frame])
          bar_total()
  print("Encoding/Saving...")
  ani = animation.ArtistAnimation(fig, rendered_frames, interval=30, blit=True)
  ani.save(animation_progres_save, extra_args=['-preset', 'lossless'], progress_callback=display_encoder_progress, codec="h264_nvenc")
else:
  compiled_kernel.render_frame(opencl_queue, image.shape, None, image_buffer, mask_buffer,
                               np.double(abs(x_max - x_min) / img_res_x),
                               np.double(abs(y_max - y_min) / img_res_y),
                               np.double(x_min), np.double(y_min),
                               np.uint32(iterations), np.uint32(escape),
                               np.double(1))
  cl.enqueue_copy(opencl_queue, image, image_buffer).wait()
  ax.imshow(image, norm="log", aspect="auto", cmap=cmap)
  fig.savefig(single_frame_save)
  plt.autoscale(False)
  plt.connect('motion_notify_event', on_click)
  plt.show()
