"""Example: Stereographic projection of a tetrahedron and cross ratio.

Consists of the following steps:

- Construct a regular tetrahedron circumscribed by a sphere.
- Stereographically project the tetrahedron into a plane.
- Check that the complex cross-ratio of the images of the four vertices equals
  exp(±i*pi/3).
"""

#############################################
# Setup
#############################################

# [setup-1]
# Import necessary libraries and modules
import bpy
import numpy as np

import ddg
from ddg.nets import compose

# [setup-1]

# [setup-2]
# Clear existing objects and materials in Blender scene
ddg.blender.object.clear()
ddg.blender.material.clear()
# [setup-2]


#############################################
# Constants and Functions
#############################################


# [constants]
# Define a unit sphere as a quadric
UNIT_SPHERE = ddg.geometry.Quadric(np.diag([1, 1, 1, -1]))
# Define the north pole as a point on the unit sphere
NORTH_POLE = ddg.geometry.subspace_from_affine_points([0, 0, 1])
# [constants]


# [functions-1]
# Define a function to construct a tetrahedron inscribed in a sphere
def tetrahedron_in_sphere(vertex):
    """Return a tetrahedron inscribed in the unit sphere
    based on the position of one vertex.

    Parameters
    ----------
    vertex : array_like of shape (3,)
        The position of one vertex of the tetrahedron.
        It determines the whole tetrahedron. Vector will be normalized.

    Returns
    -------
    ddg.halfedge.Surface
        Tetrahedron as Halfedge object.
    """
    # This tetrahedron is already regular.
    tetrahedron = ddg.halfedge.tetrahedron(co_attr=None)
    tetrahedron.verts.add_attribute("co")
    a, b, c, d = list(tetrahedron.verts)
    a.co = np.array((0, 0, 0))
    b.co = np.array((1, 1, 0))
    c.co = np.array((0, 1, 1))
    d.co = np.array((1, 0, 1))
    random_vertex = ddg.halfedge.some_vertex(tetrahedron)

    # Translate and scale the tetrahedron so it's inscribed in the unit sphere
    barycenter = sum(v.co for v in tetrahedron.verts) / len(tetrahedron.verts)
    for v in tetrahedron.verts:
        v.co = v.co - barycenter
    norm = np.linalg.norm(random_vertex.co)
    for v in tetrahedron.verts:
        v.co = v.co / norm

    # Rotate the tetrahedron to match the provided vertex
    F = ddg.math.euclidean.rotation_from_to(random_vertex.co, vertex, homogeneous=False)
    for v in tetrahedron.verts:
        v.co = F @ v.co

    return tetrahedron


# [functions-1]


# [functions-2]
# Define a function to stereographically project a tetrahedron
def project_tetrahedron(tetrahedron):
    """Stereographically project a tetrahedron.

    Parameters
    ----------
    tetrahedron : Surface
        Halfedge object representing the tetrahedron.

    Returns
    -------
    points_projected : list of complex
        Points in the plane x3 = 0. A point [x, y, 0] is represented by complex(x, y).
        The north pole is projected to complex('inf').

        This is also the representation used in the ddg.math.complex module, so
        the cross ratio can be computed directly from the points in this list.

    edges_projected : list of ddg.nets.SmoothCurve
        List of projected edges.
    """
    # Project each vertex of the tetrahedron to the plane x3 = 0
    points_projected = [
        ddg.geometry.stereographic_project(
            ddg.geometry.subspace_from_affine_points(v.co)
        )
        for v in tetrahedron.verts
    ]
    for i, p in enumerate(points_projected):
        if p == float("inf"):
            points_projected[i] = complex("inf")
        else:
            p = p.affine_point
            points_projected[i] = complex(p[0], p[1])

    # Construct edge list as a list of lists of vertices
    edges = ddg.halfedge.single_edges(tetrahedron)
    edges_projected = [project_edge_to_sphere(e) for e in edges]

    # Restrict the domain to exclude the north pole because
    # that would be projected to infinity
    for e in edges_projected:
        if ddg.geometry.subspace_from_affine_points(e(0)) == NORTH_POLE:
            e.domain = ddg.nets.SmoothDomain([[0.1, 1]])
        elif ddg.geometry.subspace_from_affine_points(e(1)) == NORTH_POLE:
            e.domain = ddg.nets.SmoothDomain([[0, 0.9]])

    # Transform the edges to match the projection
    def trafo(p):
        p = ddg.geometry.subspace_from_affine_points(p)
        p_projected = ddg.geometry.stereographic_project(p)
        return p_projected.affine_point

    edges_projected_transformed = [compose(trafo, e) for e in edges_projected]

    return points_projected, edges_projected_transformed


# [functions-2]


# [functions-3]
# Define a function to project an edge of a halfedge object to the unit sphere
def project_edge_to_sphere(edge):
    """Project an edge of a halfedge object to the unit sphere.

    Parameters
    ----------
    edge : Edge
        Edge of a halfedge object

    Returns
    -------
    SmoothCurve
    """
    p1 = edge.tail.co
    p2 = edge.head.co

    # Create a smooth curve representing the projected edge
    edge_projected = ddg.nets.SmoothNet(lambda t: t * p1 + (1 - t) * p2, [[0, 1]])
    edge_projected = compose(lambda x: x / np.linalg.norm(x), edge_projected)
    return edge_projected


# [functions-3]


#############################################
# Calculations
#############################################

# [calculations-1]
# Define the coordinates of one vertex of the tetrahedron (will be normalized)
my_vertex = [1, -2, 2]
# Create a tetrahedron inscribed in the unit sphere based on the provided vertex
my_tetrahedron = tetrahedron_in_sphere(my_vertex)
# Project the edges of the tetrahedron to the unit sphere
edges_on_sphere = [
    project_edge_to_sphere(e) for e in ddg.halfedge.single_edges(my_tetrahedron)
]
# Stereographically project the tetrahedron into the plane
my_points_projected, my_edges_projected = project_tetrahedron(my_tetrahedron)
# [calculations-1]

# [calculations-2]
# Calculate the complex cross-ratio of the projected points
# and check if it equals exp(±i*pi/3)
ccr = ddg.math.complex.cr(*my_points_projected)
assert ddg.nonexact.isclose(
    ccr, np.exp(np.pi / 3 * complex(0, 1))
) or ddg.nonexact.isclose(ccr, np.exp(-np.pi / 3 * complex(0, 1)))
# [calculations-2]


#############################################
# Visualization
#############################################

# [visualization-1]
# Define colors for visualization
BLUE = 1 * np.array((0.02, 0, 1))
GREEN = 0.3 * np.array((0, 1, 0.2))
YELLOW = 1 * np.array((1, 0.5, 0))

# Create and set up camera
camera = ddg.blender.camera.camera(location=(3, 4, 3))
ddg.blender.camera.look_at_point(camera, (0.5, -0.5, 0))
bpy.context.scene.camera = camera
# [visualization-1]

# [visualization-2]
# Create materials and collections for visualization
mat_tet = ddg.blender.material.material(color=BLUE, name="mat_tet")
mat_tet_on_sphere = ddg.blender.material.material(
    color=GREEN, name="mat_tet_on_sphere", alpha=0.5
)
mat_tet_proj = ddg.blender.material.material(color=YELLOW, name="mat_tet_proj")
mat_plane = ddg.blender.material.material(
    color=(0.5, 0.5, 0.5), alpha=0.25, name="mat_plane"
)
# Set blend_method for transparency to work with Eevee render engine
mat_plane.blend_method = "BLEND"

# Create collections to organize objects
tetrahedron_on_sphere_coll = ddg.blender.collection.collection("Tetrahedron on sphere")
tetrahedron_projected_coll = ddg.blender.collection.collection("Projected tetrahedron")
# [visualization-2]

# [visualization-3]
# Visualize the inscribed tetrahedron
tetrahedrom_mesh = ddg.arrays.from_half_edge_surface(my_tetrahedron)
ddg.blender.vertices(
    ddg.arrays.Points(tetrahedrom_mesh.points),
    "tetrahedron_points",
    radius=0.04,
    material=mat_tet,
)
bobj = ddg.blender.edges(tetrahedrom_mesh, "tetrahedron_edges", material=mat_tet)
bobj.data.bevel_depth = 0.015

# Visualize edges on the unit sphere
for i, edge_smooth_net in enumerate(edges_on_sphere):
    edge_discrete_net = ddg.nets.sample_smooth_net(
        edge_smooth_net, sampling=[0.05, 60, "c"]
    )

    bobj = ddg.blender.convert(
        edge_discrete_net, f"spherical edge {i}", collection=tetrahedron_on_sphere_coll
    )
    ddg.blender.material.set_material(bobj, mat_tet_on_sphere)
    bobj.data.bevel_depth = 0.01

# Visualize projected points and edges
for i, point in enumerate(my_points_projected):
    if point != complex("inf"):
        bobj = ddg.blender.vertices(
            ddg.geometry.subspace_from_affine_points([point.real, point.imag, 0]),
            f"projected point {i}",
            radius=0.04,
            material=mat_tet_proj,
            collection=tetrahedron_projected_coll,
        )
# NOTE: No useful sphere radius

for i, edge_smooth_net in enumerate(my_edges_projected):
    edge_discrete_net = ddg.nets.sample_smooth_net(edge_smooth_net, [0.05, 60, "c"])

    bobj = ddg.blender.convert(
        edge_discrete_net, f"projected edge {i}", collection=tetrahedron_projected_coll
    )
    ddg.blender.material.set_material(bobj, mat_tet_proj)
    bobj.data.bevel_depth = 0.015
# [visualization-3]

# [visualization-4]
# Create a plane in the xy plane
xy_plane = ddg.geometry.subspace_from_affine_points([0, 0, 0], [1, 0, 0], [0, 1, 0])
xy_plane = xy_plane.orthonormalize()
bobj = ddg.blender.convert(xy_plane, "xy plane")
ddg.blender.material.set_material(bobj, mat_plane)
# [visualization-4]

# [visualization-5]
# Configure Blender settings for rendering
# Set world color to pure white and (emission?) strength to 0.5
bpy.data.worlds["World"].node_tree.nodes["Background"].inputs[0].default_value = (
    1,
    1,
    1,
    1,
)
bpy.data.worlds["World"].node_tree.nodes["Background"].inputs[1].default_value = 0.5
# Make world background transparent
bpy.data.scenes["Scene"].render.film_transparent = True
# Set render engine to Eevee for faster rendering
bpy.data.scenes["Scene"].render.engine = "BLENDER_EEVEE"
# [visualization-5]


# from testing.tests.examples.blender.snapshot import opt_in  # noqa: E402
#
# opt_in()