# Small Wavefront OBJ/MTL writer shared by the dragon model builders.
import os
import struct
import numpy as np
from scipy import ndimage


class ObjWriterT:
    def __init__(self):
        self.vertices = []
        self.uvs = []
        self.normals = []
        self.normalIndex = {}
        self.lines = []
        self.faceCount = 0
        self.currentMaterial = None
        self.currentObject = None
        self.faces = []

    def vertex(self, p, uv=None):
        self.vertices.append('v %.5f %.5f %.5f' % tuple(p))
        if uv is not None:
            self.uvs.append('vt %.5f %.5f' % tuple(uv))
        return len(self.vertices)

    def point(self, index):
        parts = self.vertices[index - 1].split()
        return np.array([float(parts[1]), float(parts[2]), float(parts[3])])

    def object(self, name):
        self.lines.append('o %s' % name)
        self.currentObject = name
        self.currentMaterial = None

    def material(self, name):
        if name != self.currentMaterial:
            self.lines.append('usemtl %s' % name)
        self.currentMaterial = name

    def normal(self, n):
        key = tuple(round(float(v), 4) for v in n)
        if key not in self.normalIndex:
            self.normals.append('vn %.4f %.4f %.4f' % key)
            self.normalIndex[key] = len(self.normals)
        return self.normalIndex[key]

    def face(self, indices, inside=None, textured=False):
        # Flat normal from the polygon; if an inside point is given the winding is fixed to face away from it.
        pts = [self.point(i) for i in indices]
        n = np.zeros(3)
        for i in range(len(pts)):
            n += np.cross(pts[i], pts[(i + 1) % len(pts)])
        length = np.linalg.norm(n)
        # A zero-area face (collinear points from a traced outline) is kept so the shell stays closed.
        n = n / length if length > 1e-12 else np.array([0.0, 0.0, 1.0])
        if inside is not None and n @ (np.mean(pts, axis=0) - inside) < 0:
            indices = indices[::-1]
            n = -n
        ni = self.normal(n)
        self.faces.append((self.currentObject, self.currentMaterial, list(indices)))
        if textured:
            self.lines.append('f ' + ' '.join('%d/%d/%d' % (i, i, ni) for i in indices))
        else:
            self.lines.append('f ' + ' '.join('%d//%d' % (i, ni) for i in indices))
        self.faceCount += 1

    def write(self, path, mtlName, header):
        out = ['# ' + line for line in header] + ['mtllib %s' % mtlName]
        out += self.vertices + self.uvs + self.normals + self.lines
        open(path, 'w').write('\n'.join(out) + '\n')


    def writeStl(self, path, materials, unitScale, name):
        # Binary STL of the faces using any of the given materials, fan-triangulated, scaled to millimetres.
        triangles = []
        for obj, material, indices in self.faces:
            if material not in materials:
                continue
            pts = [self.point(i) * unitScale for i in indices]
            for i in range(1, len(pts) - 1):
                triangles.append((pts[0], pts[i], pts[i + 1]))
        with open(path, 'wb') as f:
            f.write(('Singe %s' % name).encode('ascii')[:80].ljust(80, b'\0'))
            f.write(struct.pack('<I', len(triangles)))
            for a, b, c in triangles:
                n = np.cross(b - a, c - a)
                length = np.linalg.norm(n)
                n = n / length if length > 1e-12 else np.zeros(3)
                f.write(struct.pack('<12fH', *n, *a, *b, *c, 0))
        return len(triangles)


def pieceCount(writer, voxel, margin=12):
    # Voxelises every object of a writer by ray parity along X on a grid of the given spacing, unions
    # them, and returns (count, [(voxels, objects, (min, max)) per piece]).  Objects may overlap each
    # other, so each is rasterised on its own; the shells within one object must not overlap.
    points = np.array([writer.point(i + 1) for i in range(len(writer.vertices))])
    lower = points.min(axis=0) - margin * voxel
    upper = points.max(axis=0) + margin * voxel
    size = np.ceil((upper - lower) / voxel).astype(int) + 1
    solid = np.zeros((size[1], size[2], size[0]), dtype=bool)
    byObject = {}
    for obj, material, indices in writer.faces:
        byObject.setdefault(obj, []).append(indices)
    voxelsOf = {}
    for obj, faces in byObject.items():
        crossings = np.zeros((size[1], size[2], size[0] + 1), dtype=np.int8)
        for indices in faces:
            pts = [(writer.point(i) - lower) / voxel for i in indices]
            for t in range(1, len(pts) - 1):
                a, b, c = pts[0], pts[t], pts[t + 1]
                det = (b[1] - a[1]) * (c[2] - a[2]) - (c[1] - a[1]) * (b[2] - a[2])
                if abs(det) < 1e-9:
                    continue
                yMin = max(0, int(np.floor(min(a[1], b[1], c[1]))))
                yMax = min(size[1] - 1, int(np.ceil(max(a[1], b[1], c[1]))))
                zMin = max(0, int(np.floor(min(a[2], b[2], c[2]))))
                zMax = min(size[2] - 1, int(np.ceil(max(a[2], b[2], c[2]))))
                if yMin > yMax or zMin > zMax:
                    continue
                # Sample off the half-voxel so outline vertices never sit exactly on a ray.
                gy, gz = np.meshgrid(np.arange(yMin, yMax + 1) + 0.5123, np.arange(zMin, zMax + 1) + 0.5123, indexing='ij')
                l1 = ((b[1] - gy) * (c[2] - gz) - (c[1] - gy) * (b[2] - gz)) / det
                l2 = ((c[1] - gy) * (a[2] - gz) - (a[1] - gy) * (c[2] - gz)) / det
                l3 = 1 - l1 - l2
                inside = (l1 >= 0) & (l2 >= 0) & (l3 >= 0)
                if not inside.any():
                    continue
                x = l1 * a[0] + l2 * b[0] + l3 * c[0]
                xi = np.clip(np.ceil(x).astype(int), 0, size[0])
                ys, zs = np.nonzero(inside)
                np.add.at(crossings, (ys + yMin, zs + zMin, xi[inside]), 1)
        voxelsOf[obj] = (np.cumsum(crossings, axis=2, dtype=np.int8)[:, :, :size[0]] % 2).astype(bool)
        solid |= voxelsOf[obj]
    labels, count = ndimage.label(solid, structure=np.ones((3, 3, 3)))
    members = {}
    for obj, voxels in voxelsOf.items():
        found = labels[voxels]
        if found.size:
            members.setdefault(int(np.bincount(found).argmax()), []).append(obj)
    pieces = []
    for label in range(1, count + 1):
        box = ndimage.find_objects(labels == label)[0]
        low = lower + np.array([box[2].start, box[0].start, box[1].start]) * voxel
        high = lower + np.array([box[2].stop, box[0].stop, box[1].stop]) * voxel
        pieces.append((int((labels == label).sum()), members.get(label, []), (low, high)))
    return count, pieces


def reportPieces(count, pieces):
    print('check: %d piece%s' % (count, '' if count == 1 else 's'))
    for voxels, objects, (low, high) in pieces:
        print('  piece of %d voxels at x %.2f..%.2f y %.2f..%.2f z %.2f..%.2f: %s' % (voxels, low[0], high[0], low[1], high[1], low[2], high[2], ', '.join(objects) or 'no whole object'))


def writeMtl(path, materials, header):
    lines = ['# ' + header]
    for name, m in materials.items():
        lines.append('newmtl %s' % name)
        lines.append('Ka %.3f %.3f %.3f' % tuple(v * 0.2 for v in m['Kd']))
        lines.append('Kd %.3f %.3f %.3f' % m['Kd'])
        lines.append('Ks %.3f %.3f %.3f' % m['Ks'])
        lines.append('Ns %.1f' % m['Ns'])
        lines.append('d %.2f' % m['d'])
        lines.append('Ni %.2f' % m['Ni'])
        lines.append('illum %d' % m['illum'])
        # PBR extension keys, read by util/objToGlb.py and most modern importers.
        lines.append('Pm %.2f' % m['Pm'])
        lines.append('Pr %.2f' % m['Pr'])
        lines.append('')
    open(path, 'w').write('\n'.join(lines) + '\n')


# The three tile materials, shared by every dragon model.
DRAGON_MATERIALS = {
    'dark': {'Kd': (0.30, 0.31, 0.33), 'Ks': (0.70, 0.70, 0.70), 'Ns': 180.0, 'd': 1.0, 'Ni': 1.0, 'illum': 2, 'Pm': 0.8, 'Pr': 0.35},
    'champagne': {'Kd': (0.72, 0.64, 0.54), 'Ks': (0.85, 0.80, 0.70), 'Ns': 120.0, 'd': 1.0, 'Ni': 1.0, 'illum': 2, 'Pm': 0.8, 'Pr': 0.40},
    'glass': {'Kd': (0.88, 0.92, 0.90), 'Ks': (0.95, 0.95, 0.95), 'Ns': 300.0, 'd': 0.55, 'Ni': 1.50, 'illum': 4, 'Pm': 0.0, 'Pr': 0.15},
    # The fourth print colour:  the core, side walls and grout between tiles of the printable dragon.
    'frame': {'Kd': (0.12, 0.11, 0.10), 'Ks': (0.30, 0.30, 0.30), 'Ns': 40.0, 'd': 1.0, 'Ni': 1.0, 'illum': 2, 'Pm': 0.0, 'Pr': 0.7},
    # Backing plate of the printable logotype and plaque; its own body so any filament can be assigned.
    'plate': {'Kd': (0.55, 0.53, 0.50), 'Ks': (0.20, 0.20, 0.20), 'Ns': 30.0, 'd': 1.0, 'Ni': 1.0, 'illum': 2, 'Pm': 0.0, 'Pr': 0.8},
}