from libs import *
import heapq
import mesh

def intersect_box(ro, rd, c, s):
    m = 1 / rd
    n = m * (ro - c)
    k = np.abs(m) * s
    tn = np.max(-n - k)
    tf = np.min(-n + k)
    if tf > tn and tf > 0:
        return tn, tf
    else:
        return None, None

class WorldRegion(mesh.SquareMesh):
    def __init__(self, size=16):
        super().__init__(refls=[0.5, 0.5, 0.5, 32.0])
        self.grid = np.zeros((size, size, size), dtype=VOXEL_TYPE)
        self.r = size

    def set_voxel(self, pos, vt=0):
        self.grid[*pos] = vt
        if vt == 0:
            self.remesh()
        else:
            bp = bool(vt)
            for d in range(3):
                for s in range(2):
                    if pos[d] == self.r - 1 and s == 0 or pos[d] == 0 and s == 1:
                        continue
                    a = (d + 1) * (-1) ** s
                    o = np.array(axis2offset(a), dtype=np.int64)
                    bc = bool(self.grid[*(pos + o)])
                    if bp ^ bc:
                        a *= bc - bp
                        self.add_sq(a, 1, pos + o / 2, pos + (0 if bp else o))

    def remesh(self):
        offset = vec([0.5, 0.0, 0.0])
        self.faces = []
        prev, curr = 0, 0
        vl = []
        for d in range(3):
            offset = np.array(axis2offset(d + 1))
            for i in range(self.r):
                for j in range(self.r):
                    for k in range(self.r + 1):
                        p = np.array([k, i, j])
                        roll_vec3(p, d)
                        prev, curr = curr, (self.grid[*p] if k < self.r else 0)
                        bp, bc = bool(prev), bool(curr)
                        if bp ^ bc:
                            sq = vec([
                                [0, 1, 1], [0, -1, 1], [0, 1, -1],
                                [0, -1, -1], [0, 1, -1], [0, -1, 1]
                            ])
                            sq /= 2
                            if bp:
                                sq[1:3] = sq[2:0:-1]
                                sq[4:6] = sq[5:3:-1]
                            for v in sq:
                                roll_vec3(v, d)
                            sq += p - offset / 2
                            vl.extend(sq)
                            if bp:
                                p -= offset
                            for _ in range(2):
                                self.faces.append(((bp - bc) * (d + 1), p))
        self.verts = vec(vl)

    def norm_list(self):
        return np.astype(np.repeat([axis2offset(a) for (a, _) in self.faces], 3, 0), FLOAT)

    def raycast(self, ro, rd):
        vr = np.repeat(self.r, 3)
        tn, tf = intersect_box(ro, rd, self.pos + vr / 2 - 0.5, vr / 2)
        if tn is not None:
            ids = []
            for i in range(3):
                k = np.arange(self.r + 1) if rd[i] >= 0 else np.arange(self.r + 1, 0, -1)
                ids.append([(d, i) for d in (k - 0.5 - ro[i]) / rd[i]])
            q, p = None, None
            for (te, i) in heapq.merge(*ids):
                if te < 0 or te < tn or te > tf - 1e-3:
                    continue
                q, p = p, np.astype(np.round(ro + (te + 1e-8) * rd - self.pos), np.int16)
                assert q is None or abs(np.sum(p - q)) == 1
                if self.grid[*p]:
                    sa = np.astype((i + 1) * -np.sign(rd[i]), np.int8)
                    return te, (sa, p)
        return None, None