# [setup]
# From discrete Weierstrass data (orthogonal circle pattern)
# to an S-isothermic minimal surface
import os

import bpy
import numpy as np

import ddg

# Clear the Blender scene
ddg.blender.scene.clear(deep=True)

# geo = "euc"
geo = "lor"

if geo == "euc":
    us = ddg.geometry.Quadric(np.diag([1.0, 1.0, 1.0, -1.0]))

    def dot(u, v):
        return u[0] * v[0] + u[1] * v[1] + u[2] * v[2]


if geo == "lor":
    us = ddg.geometry.Quadric(np.diag([1.0, 1.0, -1.0, 1.0]))

    def dot(u, v):
        return u[0] * v[0] + u[1] * v[1] - u[2] * v[2]


# Point of stereographic projection
N = ddg.geometry.Point([0, 0, -1.0, 1.0])
# [setup]


# [setup2]
# ================================================================
# Load data
# ================================================================
dir = os.path.dirname(os.path.realpath(__file__))
combinatorics = ddg.conversion.halfedge.obj.obj_to_hds(
    f"{dir}/weierstrass_schwarz_p.obj"
)
# [setup2]


# [prepare]
# ================================================================
# Color vertices
# ================================================================
# Divide vertices in black and white vertices and the white ones again, to
# distinguish which will become spheres and which will become circles.
ddg.halfedge.bicolor_vertices(combinatorics, colors=["w", "b"])
white_verts = [v for v in combinatorics.verts if v.color == "w"]
black_verts = [v for v in combinatorics.verts if v.color == "b"]

verts = sorted(list(combinatorics.verts), key=lambda v: v.index)
white_net = ddg.halfedge.diagonals(combinatorics)[0]
assert verts[ddg.halfedge.some_vertex(white_net).old_index].color == "w"

ddg.halfedge.bicolor_vertices(white_net, colors=["wc", "ws"])
combinatorics.verts.add_attribute("white_type", None)
for v in white_net.verts:
    verts[v.old_index].white_type = "wc" if v.color == "wc" else "ws"
# [prepare]

# [to_koebe]
# ================================================================
# Project to Koebe net
# ================================================================
# Construct the central extension of a Koebe net: black vertices are
# the points of tangency, they lie in 'us'. White 'ws' vertices are
# sphere centers outside of 'us' and white 'wc' are circle centers inside 'us'.
combinatorics.verts.add_attribute("koebe_co")

for b in black_verts:
    proj = ddg.geometry.inverse_stereographic_project(
        ddg.geometry.Point([*b.co, 1], atol=1e-15, rtol=0.0), us, N
    )
    b.koebe_co = proj.affine_point


# For vertices on the corner we can not take the least square subspace since
# they only have two neighbours. We treat them seperately
ddg.halfedge.set_minimal_valency_attr(
    combinatorics, "verts", "is_corner", [True, False]
)

corner_verts = [
    v for v in combinatorics.verts if v.is_corner is True and v.color == "w"
]
no_corner_verts = [
    v for v in combinatorics.verts if v.is_corner is False and v.color == "w"
]

for w in no_corner_verts:
    neighbors = [e.head for e in ddg.halfedge.out_edges(w)]
    plane = ddg.geometry.least_square_subspace(
        [ddg.geometry.Point([*b.koebe_co, 1]) for b in neighbors], k=2
    )
    w.koebe_co = us.polarize(plane).affine_point
    if w.white_type == "wc":
        w.koebe_co = w.koebe_co / np.abs(dot(w.koebe_co, w.koebe_co))

# Intersect tangent lines to obtain corner vertices
for w in corner_verts:
    assert w.white_type == "ws"
    lines = []
    for e in ddg.halfedge.out_edges(w):
        if not ddg.halfedge.is_boundary_edge(e):
            v2 = e.opp.pre.tail
        else:
            assert ddg.halfedge.is_boundary_edge(e)
            v2 = e.nex.head
        lines.append(
            ddg.geometry.subspace_from_affine_points(
                e.head.koebe_co, v2.koebe_co, atol=1e-5, rtol=1e-5
            )
        )
    pt = ddg.geometry.intersect(*lines)
    w.koebe_co = pt.affine_point

# [to_koebe]

# [dualize]
# ================================================================
# Christoffel dual
# ================================================================
# Set attribute storing edges of the Christoffel dual (with global scaling)
ddg.halfedge.bicolor_edges(combinatorics, "color", [1, -1])
combinatorics.edges.add_attribute("new_edge")

for e in combinatorics.edges:
    diff = e.head.koebe_co - e.tail.koebe_co
    normsq = dot(diff, diff)
    e.new_edge = (1 / normsq) * diff * e.color / 70

# Define boundary vertices where the integration of the one-form stops proceeding
combinatorics.verts.add_attribute("bdy")
for v in combinatorics.verts:
    if ddg.halfedge.is_boundary_vertex(v) and np.isclose(np.linalg.norm(v.co), 1.0):
        v.bdy = True

# Construct Christoffel dual
ddg.halfedge.verify_closure_one_form(combinatorics, "new_edge", tol=1e-3)
combinatorics = ddg.halfedge.integrate_one_form(
    combinatorics,
    ddg.halfedge.some_vertex(combinatorics),
    np.array([5.0, 1.0, 1.5]),
    "dual_co",
    "new_edge",
    boundary_attr="bdy",
)
# [dualize]

# [vis]
# ================================================================
# Visualize
# ================================================================
primal = ddg.halfedge.diagonals(ddg.halfedge.diagonals(combinatorics)[0])[1]

black = ddg.blender.material.material(color=(0.0, 0.0, 0.0), name="black")
white = ddg.blender.material.material(color=(1.0, 1.0, 1.0), name="white")

ddg.blender.edges(combinatorics, "weierstrass", radius=0.005, material=black)

bobj_us = ddg.blender.convert(
    ddg.arrays.from_quadric(us, (300, 0.025), (300, 0.025)),
    "unit_sphere",
    material=black,
)
with ddg.blender.bmesh.bmesh_from_mesh(bobj_us.data) as bm:
    ddg.blender.bmesh.cut_bounding_box(
        bm, np.array([np.inf, np.inf, 3]), np.array([0, 0, 2])
    )

ddg.blender.convert(
    ddg.arrays.from_half_edge_surface(combinatorics, "koebe_co"),
    "koebe_primal",
    material=white,
)
ddg.blender.edges(
    ddg.arrays.from_half_edge_surface(combinatorics, "koebe_co"),
    "koebe_edges",
    radius=0.005,
    material=black,
)

# ddg.blender.convert(
#    ddg.arrays.from_half_edge_surface(primal, "dual_co"), "max_primal", material=white
# )

ddg.blender.edges(
    ddg.arrays.from_half_edge_surface(primal, "dual_co"),
    "max_primal_edges",
    radius=0.005,
    material=black,
)
# [vis]

# [vis2]
# ================================================================
# Visualize face circles
# ================================================================
# For an orthonormal basis (u,w) of a planar face we can use a circle
# parameterization function to compute the incircle
# of the face and convert it to a disc.

for v in combinatorics.verts:
    if v.white_type == "wc" and not ddg.halfedge.is_boundary_vertex(v):
        e = next(ddg.halfedge.out_edges(v))
        u = e.head.dual_co - v.dual_co
        w = e.nex.head.dual_co - e.head.dual_co

        u /= np.sqrt(dot(u, u))
        w /= np.sqrt(dot(w, w))

        r = np.mean(
            [
                np.sqrt(dot(v.dual_co - e.head.dual_co, v.dual_co - e.head.dual_co))
                for e in ddg.halfedge.out_edges(v)
            ]
        )

        def circle(theta):
            return v.dual_co + r * np.cos(theta) * u + r * np.sin(theta) * w

        cos = [circle(t) for t in np.arange(0, 2 * np.pi, 0.1)]

        disc = ddg.halfedge.Surface()
        disc.add_face(len(cos))
        disc.verts.add_attribute("co")
        for vert in disc.verts:
            setattr(vert, "co", cos[vert.index])

        ddg.blender.convert(disc, f"face_circle_{v.index}", material=black)
# [vis2]

# Export OBJ (optional)
ddg.conversion.obj.halfedge.hds_to_obj(
    combinatorics,
    filename=f"{dir}/schwarz_p_fundamental_piece_{geo}.obj",
    co_attr="dual_co",
)

# Add light and camera
cam = ddg.blender.camera.camera(location=(2.3, 15, 2))
bpy.context.scene.camera = cam
ddg.blender.camera.look_at_point(cam, (2.3, 0, 1))
light = ddg.blender.light.light(location=(0, 0, 0), energy=1)
ddg.blender.light.look_at_point(light, (0, 0, 10))

# Setup rendering
ddg.blender.render.setup_eevee_renderer(taa_render_samples=12)
ddg.blender.render.set_world_background((1, 1, 1, 1), 1)
bpy.context.scene.render.resolution_x = 2000
bpy.context.scene.render.resolution_y = 1000
bpy.context.scene.view_settings.view_transform = "Standard"