import numpy as np

import pyopencl as cl
mf = cl.mem_flags

import OpenGL.GL as gl
from OpenGL.arrays import vbo

class MolecularLinks:
    def __init__(self, simulation):
        self.simulation = simulation

    def setup(self):
        if self.simulation.n_bonds > 0:
            self.np_view_link_indices = np.zeros((int(np.sum(self.simulation.np_bond_count)*2),1), dtype=np.uint32)
            curr = 0
            for i in range(self.simulation.n_atoms):
                for j in range(int(self.simulation.np_bond_count[i])):
                    self.np_view_link_indices[2*curr+0] = i
                    self.np_view_link_indices[2*curr+1] = self.simulation.np_bond_indices[self.simulation.max_bonds*i+j][0]
                    curr = curr + 1

            self.cl_view_link_indices = cl.Buffer(self.simulation.context, mf.COPY_HOST_PTR, hostbuf=self.np_view_link_indices)
            self.np_view_link_pos = np.zeros((self.np_view_link_indices.shape[0],4), dtype=np.float32)
            self.gl_view_link_pos = vbo.VBO(data=self.np_view_link_pos, usage=gl.GL_DYNAMIC_DRAW, target=gl.GL_ARRAY_BUFFER)
            self.gl_view_link_pos.bind()
            self.cl_view_link_pos = cl.GLBuffer(self.simulation.context, mf.READ_WRITE, int(self.gl_view_link_pos))

    def update(self):
        pass

    def shutdown(self):
        pass

    def display_decoration(self, uniform):
        if self.simulation.n_bonds > 0:
            cl.enqueue_acquire_gl_objects(self.simulation.queue, [self.cl_view_link_pos])
            self.simulation.program.update_view_links(self.simulation.queue, (int(self.np_view_link_indices.shape[0] / 2),), None, self.simulation.cl_position, self.cl_view_link_indices, self.cl_view_link_pos).wait()
        
            gl.glEnableClientState(gl.GL_VERTEX_ARRAY)
            gl.glUniform3fv(uniform['color'], 1, [0.5,0.5,0.5])
            self.gl_view_link_pos.bind()
            gl.glVertexPointer(4, gl.GL_FLOAT, 0, self.simulation.gl_position)
            gl.glDrawArrays(gl.GL_LINES, 0, self.np_view_link_indices.shape[0])
            gl.glDisableClientState(gl.GL_VERTEX_ARRAY)

    def display_window(self, uniform):
        pass