#!/usr/bin/env python3
import argparse
from compphysutils.parser import readFile
import pyray
from compphysutils.graphics.atom_plot import atom_sizes, atom_colors, wire_color
import numpy
from compphysutils.graphics.plotter import ColorIterator

argparser = argparse.ArgumentParser(prog="plot3dcoords", description="Plots the coordinates and optionally isosurfaces using pyraylib.")
argparser.add_argument("--format", "-f", dest="format", help="Name of the format of the coordinate file.")
argparser.add_argument("--camera_pos", type=float, nargs=3, help="Initial position of the camera. Default : 10 Angström in z direction above the center.", default=False)
argparser.add_argument("--camera_fovy", type=float, help="Field of view for the camera", default=45)
argparser.add_argument("--triangles", nargs="*", default=False, help="Filenames containing column triangulation data for surfaces (isosurfaces) to be plotted.")
argparser.add_argument("--sphere_res", default=30, type=int, help="Resolution of individual spheres - sets rings + sections. [default : 10]")
argparser.add_argument("--tri_colors", default=["255,0,0,120","0,0,255,120"], nargs="+", help="Colors for cycling the isosurfaces, in R,G,B,A format.")
argparser.add_argument("coordinate_file", help="Name of the coordinate file to be processed")
argparser.add_argument("savefile", help="Name of the file into which to save the rendered coordinates")

# Shaders
vertexShader = """#version 330
in vec3 vertexPosition;
in vec4 vertexColor;
in vec3 vertexNormal;

smooth out vec4 fragColor;

uniform mat4 mvp;

uniform vec3 lightDirection;
uniform float ambient;
// Maximum cos for reflection
// The intensity at the max is also reduced
uniform float minCos;

void main(){
	// cos(angle) between normal and direction
	float dotProduct = dot(normalize(lightDirection), normalize(vertexNormal));
	// Intensity together with ambient
	float intensity = clamp(ambient + (clamp(-dotProduct, minCos, 1.0)-minCos)/(1.0-minCos), 0.0, 1.0);

	fragColor = vec4(intensity, intensity, intensity, 1.0) * vertexColor;
	gl_Position = mvp * vec4(vertexPosition, 1.0);
}
"""
fragmentShader = """#version 330
smooth in vec4 fragColor;

uniform vec4 colDiffuse;

out vec4 finalColor;

void main(){
	finalColor = fragColor * colDiffuse;
}
"""

args = argparser.parse_args()
dataset = readFile(args.coordinate_file, args.format)
# Init window and start loading models
pyray.init_window(1680,1050, args.savefile)

atom_models = []
center_coords = [0.0,0.0,0.0]
shader = pyray.load_shader_from_memory(vertexShader, fragmentShader)

light_dir_loc = pyray.get_shader_location(shader, "lightDirection")
ambient_loc = pyray.get_shader_location(shader, "ambient")
minCos_loc = pyray.get_shader_location(shader, "minCos")
#TODO : Set in args
pyray.set_shader_value(shader, light_dir_loc,pyray.Vector3(-1.0, 0.0, 0.0), pyray.SHADER_UNIFORM_VEC3)
pyray.set_shader_value(shader, ambient_loc, pyray.ffi.new("float *", 0.5), pyray.SHADER_UNIFORM_FLOAT)
pyray.set_shader_value(shader, minCos_loc, pyray.ffi.new("float *", 0.4), pyray.SHADER_UNIFORM_FLOAT)
wires = False

for atom_index in range(len(dataset[0])):
    # TODO : Settable rings+sections
    # Also add to set the center
    for i in range(3):
        center_coords[i] += dataset[i][atom_index]
    sphere_mesh = pyray.gen_mesh_sphere(atom_sizes[dataset[3][atom_index].upper()], args.sphere_res, args.sphere_res)
    model = pyray.load_model_from_mesh(sphere_mesh)
    model.materials[0].shader = shader
    # Give shaders to the model
    atom_models.append(model)

for i in range(3):
    center_coords[i] = center_coords[i] / len(dataset[0])

surface_models = []
if args.triangles:
    for surface_index in range(len(args.triangles)):
        triset = readFile(args.triangles[surface_index], "cols", "0 1 2 3 4 5 6 7 8")
        vertices = numpy.zeros(9*len(triset[0]), dtype=numpy.float32)
        normals = numpy.zeros(9*len(triset[0]), dtype=numpy.float32)
        for i in range(len(triset[0])):
            for j in range(9):
                vertices[9*i+j] = triset[j][i]
        #    # Calculate the normal of the face
        #    p1 = numpy.array([triset[0][i],triset[1][i],triset[2][i]])
        #    p2 = numpy.array([triset[3][i],triset[4][i],triset[5][i]])
        #    p3 = numpy.array([triset[6][i],triset[7][i],triset[8][i]])
        #    vec1 = p2 - p1
        #    vec2 = p3 - p1
        #    norm = numpy.cross(vec1, vec2)
        #    norm = norm / numpy.linalg.norm(norm)
        #    # Save the normal to each vertex
        #    for j in range(3):
        #        normals[9*i+3*j] = norm[0]
        #        normals[9*i+3*j+1] = norm[1]
        #        normals[9*i+3*j+2] = norm[2]
        mesh = pyray.Mesh(3*len(triset[0]), len(triset[0]), vertices, normals)
        #mesh = pyray.Mesh(3*len(triset[0]), len(triset[0]), vertices)
        pyray.upload_mesh(mesh, True)
        model = pyray.load_model_from_mesh(mesh)
        model.materials[0].shader = shader
        surface_models.append(model)

# TODO : Allow for custom setting of surface colors
surface_color = ColorIterator(" ".join(args.tri_colors))

camera = pyray.Camera3D()
if not args.camera_pos:
    camera.position = [center_coords[0], center_coords[1], center_coords[2]+10.0]
else:
    camera.position = args.camera_pos
# TODO : Set other camera parameters from the command line
camera.up = [0,1,0]
camera.target = center_coords
camera.fovy = args.camera_fovy
camera.projection = pyray.CAMERA_PERSPECTIVE

# Models uploaded, start the render loop
while not pyray.window_should_close():
    pyray.begin_drawing()
    pyray.clear_background(pyray.RAYWHITE)
    pyray.begin_mode_3d(camera)
    # TODO : This is more flexible, but a single model might be more performant for very large structures
    for atom_index in range(len(dataset[0])):
        # Draw the atom
        pyray.draw_model(atom_models[atom_index],
                         [dataset[0][atom_index],
                          dataset[1][atom_index],
                          dataset[2][atom_index]],
                         1.0,
                         atom_colors[dataset[3][atom_index].upper()]
                         )
        # # Draw the atom wires
        if wires:
            pyray.draw_model_wires(atom_models[atom_index],
                             [dataset[0][atom_index],
                              dataset[1][atom_index],
                              dataset[2][atom_index]],
                             1.0,
                             wire_color
                             )
    # Draw the surfaces
    for surface_index in range(len(surface_models)):
        pyray.draw_model(surface_models[surface_index],
                         [0.0, 0.0, 0.0],
                         1.0,
                         list(map(int, next(surface_color))))
    pyray.end_mode_3d()
    pyray.end_drawing()
    # Check for screenshot
    if pyray.is_key_pressed(pyray.KEY_S):
        pyray.take_screenshot(args.savefile)
    # Updating camera
    if pyray.is_mouse_button_down(pyray.MOUSE_BUTTON_LEFT) or pyray.get_mouse_wheel_move() != 0.0:
        pyray.update_camera(camera, pyray.CAMERA_THIRD_PERSON)

for atom_index in range(len(dataset[0])):
    pyray.unload_model(atom_models[atom_index])
# TODO : Fix the segfault on free
#for surface_index in range(len(surface_models)):
#    pyray.unload_model(surface_models[surface_index])
pyray.close_window()
