from libs import *
from raycast import intersect_box
import heapq
import mesh

class WorldRegion(mesh.SquareMesh):
    def __init__(self, size=0):
        super().__init__(refls=[0.5, 0.5, 0.5, 32.0])
        size = size or conf["region"]
        self.grid = np.zeros((size, size, size), dtype=np.uint8)
        self.r = size
        self.dirty = False

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

    def remesh(self):
        if not np.sum(self.grid):
            return
        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)
        self.dirty = False

    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):
                if rd[i] >= 0:
                    k = np.arange(self.r + 1)
                else:
                    k = np.arange(self.r + 1, 0, -1)
                dists = (self.pos[i] + k - 0.5 - ro[i]) / rd[i]
                ids.append([(d, i) for d in dists])
            for (te, i) in heapq.merge(*ids):
                if te < 0 or te < tn or te > tf - 1e-3:
                    continue
                p = ro + (te + 1e-4) * rd - self.pos
                p = np.astype(np.round(p), np.int16)
                if self.grid[*p]:
                    sa = np.astype((i + 1) * -np.sign(rd[i]), np.int8)
                    return te, (sa, p)
        return None, None