from OpenGL.GL import *
import itertools
import numpy as np
import pygame

FLOAT = np.float32
VOXEL_TYPE = np.uint8
VERTEX_ATTRS = ["vec3 position", "vec4 col", "vec3 normal", "vec4 refls"]
UNIFORMS = {"MVP": ["vert", "mat4", None], "lightPos": ["frag", "vec4", None], "camPos": ["frag", "vec3", None]}

VIEW_DIST = 30
FPS = 60
ARROWS = [pygame.K_RIGHT, pygame.K_LEFT, pygame.K_SPACE, pygame.K_LSHIFT, pygame.K_UP, pygame.K_DOWN]

def vec(l):
    return np.array(l, dtype=FLOAT)

RED = vec([1.0, 0.0, 0.0, 1.0])
YELLOW = vec([1.0, 1.0, 0.0, 1.0])
GREEN = vec([0.0, 1.0, 0.0, 1.0])

def normalise(v):
    norm = np.linalg.norm(v)
    return v / norm if norm > 0 else v

def roll_vec3(l, k):
    match k:
        case 1:
            l[:] = [l[2], l[0], l[1]]
        case 2:
            l[:] = [l[1], l[2], l[0]]

def axis2offset(a, x=1):
    o = [0, 0, x * np.sign(a)]
    roll_vec3(o, abs(a))
    return o

def smerp(ab, w):
    assert all([abs(x) < 2 for x in ab])
    for i in range(w.shape[0]):
        ab = [a + w[i] ** 2 * (3 - 2 * w[i]) * (b - a) for (a, b) in zip(ab[::2], ab[1::2])]
    return ab[0]

# https://iquilezles.org/articles/intersectors/
def intersect_tri(ro, rd, tv):
    e1 = tv[1] - tv[0]
    e2 = tv[2] - tv[0]
    to = ro - tv[0]
    n = np.cross(e1, e2)
    q = np.cross(to, rd)
    d = 1.0 / np.dot(rd, n)
    u = d * np.dot(-q, e2)
    v = d * np.dot(q, e1)
    t = d * np.dot(-n, to)
    if u > 0.0 and v > 0.0 and u + v < 1.0 and t > 1e-3:
        return t

def norm_tri(tv):
    assert tv.shape[0] == 3
    return normalise(np.cross(tv[2] - tv[0], tv[1] - tv[0]))

def update_attr(ent, n, l=None):
    if isinstance(n, int):
        i, name = n, VERTEX_ATTRS[n].split()[1]
    elif isinstance(n, str):
        # https://stackoverflow.com/a/30197797
        i, name = next((i for (i, s) in enumerate(VERTEX_ATTRS) if n in s)), n
    if l is None:
        match name:
            case "position":
                l = ent.verts
            case "col":
                l = ent.mat_list()[0]
            case "normal":
                l = ent.norm_list()
                assert all([abs(np.linalg.norm(v) - 1) < 1e-3 for v in l])
            case "refls":
                l = ent.mat_list()[1]
    if ent.vao:
        glBindBuffer(GL_ARRAY_BUFFER, ent.vbos[i][0])
        if l.shape[0] != ent.vbos[i][1]:
            glBufferData(GL_ARRAY_BUFFER, l.nbytes, l, GL_DYNAMIC_DRAW)
            ent.vbos[i][1] = l.shape[0]
        else:
            glBufferSubData(GL_ARRAY_BUFFER, 0, l.nbytes, l)
        return l