import numpy

from mako.template import Template
from mako.lookup import TemplateLookup

from pathlib import Path

from simulation         import Lattice, Geometry
from symbolic.generator import LBM

import symbolic.D3Q27 as D3Q27

from OpenGL.GL   import *
from OpenGL.GLUT import *

from OpenGL.GL import shaders

from geometry.box import Box

from utility.projection import Projection, Rotation
from utility.opengl     import MomentsTexture
from utility.mouse      import MouseDragMonitor, MouseScrollMonitor

lattice_x = 200
lattice_y = 64
lattice_z = 64

updates_per_frame = 50

inflow = 0.01
relaxation_time = 0.51

lbm = LBM(D3Q27)

boundary = Template("""
    if ( m == 2 ) {
        u_0 = 0.0;
        u_1 = 0.0;
        u_2 = 0.0;
    }
    if ( m == 3 ) {
        u_0 = min(time/5000.0 * ${inflow}, ${inflow});
        u_1 = 0.0;
        u_2 = 0.0;
    }
    if ( m == 4 ) {
        rho = 1.0;
    }
""").render(
    inflow = inflow
)

channel = """
float sdf(vec3 v) {
    return add(
        add(
            ssub(
                box(translate(v, v3(25,center.y,center.z)), v3(5,32,32)),
                add(
                    sphere(translate(v, v3(20,0.5*center.y,1.5*center.z)), 10),
                    sphere(translate(v, v3(30,1.5*center.y,0.5*center.z)), 10)
                ),
                2
            ),
            ssub(
                box(translate(v, v3(85,center.y,center.z)), v3(5,32,32)),
                add(
                    sphere(translate(v, v3(90,1.5*center.y,1.5*center.z)), 10),
                    sphere(translate(v, v3(80,0.5*center.y,0.5*center.z)), 10)
                ),
                2
            )
        ),
        ssub(
            box(translate(v, v3(145,center.y,center.z)), v3(5,32,32)),
            add(
                cylinder(rotate_y(translate(v, v3(145,1.5*center.y,0.5*center.z)), 1), 10, 10),
                cylinder(rotate_y(translate(v, v3(145,0.5*center.y,1.5*center.z)), -1), 10, 10)
            ),
            2
        )
    );
}

float sdf_bounding(vec3 v) {
    return add(
        add(
            box(translate(v, v3(25,center.y,center.z)), v3(5,32,32)),
            box(translate(v, v3(85,center.y,center.z)), v3(5,32,32))
        ),
        box(translate(v, v3(145,center.y,center.z)), v3(5,32,32))
    );
}
"""

def glut_window(fullscreen = False):
    glutInit(sys.argv)
    glutSetOption(GLUT_MULTISAMPLE, 8)
    glutInitDisplayMode(GLUT_RGBA | GLUT_DOUBLE | GLUT_MULTISAMPLE)

    if fullscreen:
        window = glutEnterGameMode()
    else:
        glutInitWindowSize(800, 500)
        glutInitWindowPosition(0, 0)
        window = glutCreateWindow("LBM")

    return window

window = glut_window(fullscreen = False)

vertex_shader = shaders.compileShader("""
#version 430

layout (location=0) in vec4 vertex;

uniform mat4 projection;
uniform mat4 rotation;

void main() {
    gl_Position = projection * rotation * vertex;
}""", GL_VERTEX_SHADER)

fragment_shader = shaders.compileShader("""
#version 430

in vec3 color;

layout(location = 0) out vec4 frag_color;

void main(){
    frag_color = vec4(vec3(0.5), 1.0);
}""", GL_FRAGMENT_SHADER)

raycast_vertex_shader = shaders.compileShader("""
#version 430

layout (location=0) in vec4 vertex;
                   out vec3 frag_pos;

uniform mat4 projection;
uniform mat4 rotation;

void main() {
    gl_Position = projection * rotation * vertex;
    frag_pos    = vertex.xyz;
}""", GL_VERTEX_SHADER)

raycast_fragment_shader = shaders.compileShader(Template("""
#version 430

#define EPSILON 1e-1
#define RAYMARCH_STEPS 64
#define OBSTACLE_STEPS 16

in vec3 frag_pos;

uniform vec4 camera_pos;

uniform sampler3D moments;

out vec4 result;

const vec3 cuboid = vec3(${size_x}, ${size_y}, ${size_z});
const vec3 center = vec3(${size_x/2.5}, ${size_y/2}, ${size_z/2});

vec3 v3(float x, float y, float z) {
    return vec3(x,y,z);
}

vec2 v2(float x, float y) {
    return vec2(x,y);
}

vec3 fabs(vec3 x) {
    return abs(x);
}

float fabs(float x) {
    return abs(x);
}

<%include file="template/sdf.lib.glsl.mako"/>

${sdf_source}

vec3 sdf_normal(vec3 v) {
    return normalize(vec3(
        sdf(vec3(v.x + EPSILON, v.y, v.z)) - sdf(vec3(v.x - EPSILON, v.y, v.z)),
        sdf(vec3(v.x, v.y + EPSILON, v.z)) - sdf(vec3(v.x, v.y - EPSILON, v.z)),
        sdf(vec3(v.x, v.y, v.z + EPSILON)) - sdf(vec3(v.x, v.y, v.z - EPSILON))
    ));
}

vec3 unit(vec3 v) {
    return v / cuboid;
}

vec3 palette(float x) {
    return mix(
        vec3(0.251, 0.498, 0.498),
        vec3(0.502, 0.082, 0.082),
        x
    );
}

vec3 getVelocityColorAt(vec3 pos) {
    const vec4 data = texture(moments, unit(pos));
    return palette(length(data.yzw) / ${2*inflow});
}

float distanceToLattice(vec3 v) {
  return box(v, vec3(${size_x},${size_y},${size_z}));
}

float maxRayLength(vec3 origin, vec3 ray) {
    return max(1.0, ${max_ray_length} - distanceToLattice(origin + ${max_ray_length}*ray));
}

vec4 trace_obstacle(vec3 origin, vec3 ray, float delta) {
    vec3 sample_pos = origin;
    float ray_dist = 0.0;

    for (int i = 0; i < OBSTACLE_STEPS; ++i) {
        const float sdf_dist = sdf(sample_pos);
        ray_dist += sdf_dist;

        if (ray_dist > delta) {
            return vec4(0.0);
        }

        if (abs(sdf_dist) < EPSILON) {
            const vec3 color = vec3(0.5);
            const vec3 n = normalize(sdf_normal(sample_pos));
            return vec4(abs(dot(n, ray)) * color, 1.0);
        } else {
            sample_pos = origin + ray_dist*ray;
        }
    }

    return vec4(0.0);
}

vec3 trace(vec3 pos, vec3 ray) {
    const float depth = maxRayLength(pos, ray);
    const float delta = depth / RAYMARCH_STEPS;
    const float gamma = 0.5 / RAYMARCH_STEPS;

    vec3 color = vec3(0.0);
    vec3 sample_pos = pos;

    for (int i=0; i < RAYMARCH_STEPS; ++i) {
        sample_pos += delta*ray;

        if (sdf_bounding(sample_pos) > delta) {
            color += gamma * getVelocityColorAt(sample_pos);
        } else {
            const vec4 obstacle_color = trace_obstacle(sample_pos, ray, delta);
            if (obstacle_color.w == 1.0) {
                return color + obstacle_color.xyz;
            } else {
                color += gamma * getVelocityColorAt(sample_pos);
            }
        }
    }

    return color;
}

void main(){
    const vec3 ray = normalize(frag_pos - camera_pos.xyz);

    result = vec4(trace(frag_pos, ray), 1.0);
}
""", lookup = TemplateLookup(directories = [
    Path(__file__).parent
])).render(
    size_x = lattice_x,
    size_y = lattice_y,
    size_z = lattice_z,
    inflow = inflow,
    max_ray_length = max(lattice_x,lattice_y,lattice_z)**3,
    sdf_source = channel
), GL_FRAGMENT_SHADER)

domain_program   = shaders.compileProgram(vertex_shader, fragment_shader)
domain_projection_id = shaders.glGetUniformLocation(domain_program, 'projection')
domain_rotation_id   = shaders.glGetUniformLocation(domain_program, 'rotation')

raycast_program = shaders.compileProgram(raycast_vertex_shader, raycast_fragment_shader)
raycast_projection_id       = shaders.glGetUniformLocation(raycast_program, 'projection')
raycast_rotation_id         = shaders.glGetUniformLocation(raycast_program, 'rotation')
raycast_camera_pos_id = shaders.glGetUniformLocation(raycast_program, 'camera_pos')

lattice = Lattice(
    descriptor   = D3Q27,
    geometry     = Geometry(lattice_x, lattice_y, lattice_z),
    moments      = lbm.moments(optimize = True),
    collide      = lbm.bgk(f_eq = lbm.equilibrium(), tau = relaxation_time, optimize = True),
    boundary_src = boundary,
    opengl       = True
)

lattice.setup_channel_with_sdf_obstacle(channel)

moments_texture = MomentsTexture(lattice, include_materials = False)

projection = Projection(distance = 2*lattice_x)
rotation = Rotation([-0.5*lattice_x, -0.5*lattice_y, -0.5*lattice_z])

cube_vertices, cube_edges = lattice.geometry.wireframe()

def on_display():
    for i in range(0,updates_per_frame):
        lattice.evolve()

    moments_texture.collect()

    glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)

    glEnable(GL_DEPTH_TEST)
    glDepthFunc(GL_LESS)

    camera_pos = numpy.matmul([0,-projection.distance,0,1], rotation.get_inverse())

    shaders.glUseProgram(domain_program)
    glUniformMatrix4fv(domain_projection_id, 1, False, numpy.ascontiguousarray(projection.get()))
    glUniformMatrix4fv(domain_rotation_id,   1, False, numpy.ascontiguousarray(rotation.get()))
    glLineWidth(3)
    glBegin(GL_LINES)
    for i, j in cube_edges:
        glVertex(cube_vertices[i])
        glVertex(cube_vertices[j])
    glEnd()

    shaders.glUseProgram(raycast_program)
    glUniformMatrix4fv(raycast_projection_id, 1, False, numpy.ascontiguousarray(projection.get()))
    glUniformMatrix4fv(raycast_rotation_id,   1, False, numpy.ascontiguousarray(rotation.get()))
    glUniform4fv(raycast_camera_pos_id, 1, camera_pos)
    moments_texture.bind()
    Box(0,lattice.geometry.size_x,0,lattice.geometry.size_y,0,lattice.geometry.size_z).draw()

    glutSwapBuffers()

mouse_monitors = [
    MouseDragMonitor(GLUT_LEFT_BUTTON,  lambda dx, dy: rotation.update(0.005*dy, 0.005*dx)),
    MouseDragMonitor(GLUT_RIGHT_BUTTON, lambda dx, dy: rotation.shift(0.25*dx, 0.25*dy)),
    MouseScrollMonitor(lambda zoom: projection.update_distance(5*zoom))
]

def on_timer(t):
    glutTimerFunc(t, on_timer, t)
    glutPostRedisplay()

glutDisplayFunc(on_display)
glutReshapeFunc(lambda w, h: projection.update_ratio(w, h))
glutMouseFunc(lambda *args: list(map(lambda m: m.on_mouse(*args), mouse_monitors)))
glutMotionFunc(lambda *args: list(map(lambda m: m.on_mouse_move(*args), mouse_monitors)))
glutTimerFunc(30, on_timer, 30)

glutMainLoop()