"""Calculate the mesh from a 3D SDF"""

import bpy
import numpy as np
import math

import os
import sys

dir = os.path.dirname(bpy.data.filepath)
if not dir in sys.path:
    sys.path.append(dir)

from SDFtoMeshSettings import ADAPTIVE, XMIN, XMAX, YMIN, YMAX, ZMIN, ZMAX
from qef import solve_qef_3d
from utils_3d import V3

def adapt(v0, v1):
    """v0 and v1 are numbers of opposite sign. This returns how far you need to interpolate from v0 to v1 to get to 0."""
    assert (v1 > 0) != (v0 > 0), "v0 and v1 do not have opposite sign"
    if ADAPTIVE:
        return (0 - v0) / (v1 - v0)
    else:
        return 0.5

def dual_contour_3d_find_best_vertex(f, f_normal, x, y, z):
    if not ADAPTIVE:
        return (x+0.5, y+0.5, z+0.5)

    # Evaluate f at each corner
    v = np.empty((2, 2, 2))
    for dx in (0, 1):
        for dy in (0, 1):
            for dz in (0,1):
                v[dx, dy, dz] = f(x + dx, y + dy, z + dz)

    # For each edge, identify where there is a sign change.
    # There are 4 edges along each of the three axes
    changes = []
    for dx in (0, 1):
        for dy in (0, 1):
            if (v[dx, dy, 0] > 0) != (v[dx, dy, 1] > 0):
                changes.append((x + dx, y + dy, z + adapt(v[dx, dy, 0], v[dx, dy, 1])))

    for dx in (0, 1):
        for dz in (0, 1):
            if (v[dx, 0, dz] > 0) != (v[dx, 1, dz] > 0):
                changes.append((x + dx, y + adapt(v[dx, 0, dz], v[dx, 1, dz]), z + dz))

    for dy in (0, 1):
        for dz in (0, 1):
            if (v[0, dy, dz] > 0) != (v[1, dy, dz] > 0):
                changes.append((x + adapt(v[0, dy, dz], v[1, dy, dz]), y + dy, z + dz))

    if len(changes) <= 1:
        return None

    # For each sign change location v[i], we find the normal n[i].
    # The error term we are trying to minimize is sum( dot(x-v[i], n[i]) ^ 2)

    # In other words, minimize || A * x - b || ^2 where A and b are a matrix and vector
    # derived from v and n

    normals = []
    for v in changes:
        n = f_normal(v[0], v[1], v[2])
        normals.append([n.x, n.y, n.z])
        
    return solve_qef_3d(x, y, z, changes, normals)

def swap_quad(quad, shouldSwap = True):
    # It should be a quad
    assert len(quad) == 4
    
    if shouldSwap:
        return (quad[0], quad[3], quad[2], quad[1])
        # return quad
    else:
        return quad

def normal_from_function(f, d=0.01):
    """Given a sufficiently smooth 3d function, f, returns a function approximating of the gradient of f.
    d controls the scale, smaller values are a more accurate approximation."""
    def norm(x, y, z):
        return V3(
            (f(x + d, y, z) - f(x - d, y, z)) / 2 / d,
            (f(x, y + d, z) - f(x, y - d, z)) / 2 / d,
            (f(x, y, z + d) - f(x, y, z - d)) / 2 / d,
        ).normalize()
    return norm

def dual_contour_3d(f, f_normal = None, xmin=XMIN, xmax=XMAX, ymin=YMIN, ymax=YMAX, zmin=ZMIN, zmax=ZMAX):
    """Iterates over a cells of size one between the specified range, and evaluates f and f_normal to produce
        a boundary by Dual Contouring. Returns a Mesh object."""
    
    # Calculate normal from the function itself if not provided
    if f_normal is None:
        f_normal = normal_from_function(f)
        
    # For each cell, find the the best vertex for fitting f
    vert_array = []
    vert_indices = {}
    for x in range(xmin, xmax):
        for y in range(ymin, ymax):
            for z in range(zmin, zmax):
                vert = dual_contour_3d_find_best_vertex(f, f_normal, x, y, z)
                if vert is None:
                    continue
                vert_array.append(vert)
                vert_indices[(x, y, z)] = len(vert_array) - 1

    # For each cell edge, emit an face between the center of the adjacent cells if it is a sign changing edge
    faces = []
    for x in range(xmin, xmax):
        for y in range(ymin, ymax):
            for z in range(ymin, ymax):
                if x > xmin and y > ymin:
                    solid1 = f(x, y, z + 0) > 0
                    solid2 = f(x, y, z + 1) > 0
                    if solid1 != solid2:
                        faces.append(
                            swap_quad(
                                (
                                    vert_indices[(x - 1, y - 1, z)],
                                    vert_indices[(x - 0, y - 1, z)],
                                    vert_indices[(x - 0, y - 0, z)],
                                    vert_indices[(x - 1, y - 0, z)]
                                ),
                                solid2
                            )
                        )
                if x > xmin and z > zmin:
                    solid1 = f(x, y + 0, z) > 0
                    solid2 = f(x, y + 1, z) > 0
                    if solid1 != solid2:
                        faces.append(
                            swap_quad(
                                (
                                    vert_indices[(x - 1, y, z - 1)],
                                    vert_indices[(x - 0, y, z - 1)],
                                    vert_indices[(x - 0, y, z - 0)],
                                    vert_indices[(x - 1, y, z - 0)],
                                ),
                                solid1
                            )
                        )
                if y > ymin and z > zmin:
                    solid1 = f(x + 0, y, z) > 0
                    solid2 = f(x + 1, y, z) > 0
                    if solid1 != solid2:
                        faces.append(
                            swap_quad(
                                (
                                    vert_indices[(x, y - 1, z - 1)],
                                    vert_indices[(x, y - 0, z - 1)],
                                    vert_indices[(x, y - 0, z - 0)],
                                    vert_indices[(x, y - 1, z - 0)],
                                ),
                                solid2
                            )
                        )

    return vert_array, faces
