"""
This example shows various ways of
computing and visualizing the discrete envelope and
orthogonal trajectory of a family of discrete curves.
In the comments, M denotes the Moebius quadric
M  = {[x] in RP^3 | <x, x>_{3, 1} = 0} and M^+
the set of points outside M
M^+ = {[x] in RP^3 | <x, x>_{3, 1} > 0}.

"""
from functools import lru_cache

import bpy

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

import ddg
from ddg.blender.material import material
from ddg.geometry.euclidean_models import moebius_to_projective, projective_to_moebius

# [setup-1]

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

# [setup-3]
euc = ddg.geometry.euclidean(2)
mob = ddg.geometry.euclidean_models.MoebiusModel(2)


# [setup-3]


################################################
# Helper functions
################################################
# [helper-2]
def polarized(fct):
    """
    Create a new function returning the polarized
    objects of the previous function w.r.t. the Moebius quadric.
    """

    @lru_cache(maxsize=128)
    def new_fct(i):
        return mob.absolute.polarize(fct(i))

    return new_fct


def m_points_to_circles(fct):
    """
    Create a new function returning Moebius circles as conic sections,
    while the given function returns Moebius circles as points in M^+.
    """

    @lru_cache(maxsize=128)
    def point_to_circle(i):
        polar_fct = polarized(fct)
        return ddg.geometry.intersect(polar_fct(i), mob.absolute)

    return point_to_circle


# [helper-2]

# [helper-4]
def tangent_edges(fct):
    """Returns a function that returns the edge tangent lines
    for a given function returning Points as vertices of a discrete curve.
    The i'th edge tangent line is the join of fct(i) and fct(i+1).
    """

    @lru_cache(maxsize=128)
    def tangent_edge(i):
        tangent = ddg.geometry.join(fct(i), fct(i + 1))
        return ddg.geometry.orthonormalize_and_center_subspace(
            tangent, np.sum([fct(i).affine_point, fct(i + 1).affine_point], axis=0) / 2
        )

    return tangent_edge


def m_envelope(fct, index=0):
    """
    Returning the envelope of a function of circles represented as points in
    M^+. The envelope is the intersection of consecutive circles. All neighbouring
    circles intersect, therefore there are two envelopes (distinguished by index=0/1).
    A function returning points on M is returned.
    """

    @lru_cache(maxsize=128)
    def envelope_vertex(i):
        # Finding the intersection points of the two circles (in M)
        tangent_fct = tangent_edges(fct)
        polar_tangent_fct = polarized(tangent_fct)
        two_points = ddg.geometry.intersect(mob.absolute, polar_tangent_fct(i))
        mob_points = ddg.geometry.quadric_to_subspaces(two_points)
        plane = ddg.geometry.join(mob.fixed_point, tangent_fct(i))
        # The two points must lie on opposite sides of this plane.
        # We can distinguish them by their signed distance to the plane.
        nl = ddg.geometry.normal_with_level(plane)
        env0 = 0 if np.sign(np.dot(nl, [*mob_points[0].affine_point, 1])) == -1.0 else 1
        env1 = 1 - env0
        env = [mob_points[env0], mob_points[env1]]
        return env[index]

    return envelope_vertex


# [helper-4]

# [helper-5]


def inner_product(x, y):
    """Inner product (+++-)"""
    return np.dot(x[:-1], y[:-1]) - x[-1] * y[-1]


def m_sphere_inversion(point, fix_sphere):
    """
    Applying a sphere inversion in `fix_sphere` to the point `point`.
    `point` is assumed to live in M, `fix_sphere`
    is assumed to be a point in M^+,
    i.e., outside M.
    Returned is the inverted point, living in M.
    """
    inverted_point = ddg.geometry.Point(
        ddg.math.inner_product.reflect(fix_sphere.point, point.point, inner_product)
    )
    return inverted_point


def m_iterative_sphere_inversions(fct, starting_point):
    """
    For a given stating point, apply sphere inversion
    iteratively in a family of spheres given by the function `fct`.
    The spheres are assumed to be given as points in M^+.
    """

    @lru_cache(maxsize=128)
    def new_curve_fct(i):
        if i == 0:
            return starting_point
        else:
            return m_sphere_inversion(new_curve_fct(i - 1), fct(i - 1))

    return new_curve_fct


# [helper-5]

# [helper-6]
def m_homogeneous_lift_fct(fct):
    """Apply the function m_homogeneous_lift
    to a function returning Euclidean spheres."""

    def new_fct(i):
        return m_homogeneous_lift(fct(i))

    return new_fct


def m_homogeneous_lift(sphere):
    """Lift a circle to special homogeneous coordinates
    in R^{3, 1}."""
    c, r = sphere.center.affine_point, sphere.radius
    lift = np.array(
        [*c, (-1 + np.dot(c, c) - r**2) / 2, (1 + np.dot(c, c) - r**2) / 2]
    )
    lift *= 1 / r
    return lift


def m_points_of_similitude_from_homogeneous_lift(fct):
    """
    Returning the circles of similitude for a given family of circles.
    The circles are assumed to be given as homogeneous coordinates in R^{3, 1}.
    The circles of similitude are returned as points in M (in RP^3).
    """

    @lru_cache(maxsize=128)
    def new_fct(i):
        return ddg.geometry.Point(fct(i) + fct(i + 1))

    return new_fct


def m_points_of_anti_similitude_from_homogeneous_lift(fct):
    """
    Returning the circles of anti similitude for a given family of circles.
    The circles are assumed to be given as homogeneous coordinates in R^{3, 1}.
    The circles of anti similitude are returned as points in M (in RP^3).
    """

    @lru_cache(maxsize=128)
    def new_fct(i):
        return ddg.geometry.Point(fct(i) - fct(i + 1))

    return new_fct


# [helper-6]


# [helper-3]
def project_function_of_objects_down(fct):
    """
    Wrapper to apply `moebius_to_projective` to
    a function returning projective objects.
    Returns a new function.
    """

    @lru_cache(maxsize=128)
    def new_fct(i):
        return moebius_to_projective(fct(i))

    return new_fct


# [helper-3]


def function_projective_to_affine(fct):
    """
    Wrapper to apply `.affine_point` to
    a function returning projective points.
    Returns a new function.
    """

    @lru_cache(maxsize=128)
    def new_fct(i):
        return fct(i).affine_point

    return new_fct


# [helper-1]
def function_affine_to_projective(fct):
    """
    Wrapper to apply `ddg.geometry.subspace_from_affine_points` to
    a function returning affine coordinates of points.
    Returns a new function.
    """

    @lru_cache(maxsize=128)
    def new_fct(i):
        return ddg.geometry.subspace_from_affine_points(fct(i))

    return new_fct


# [helper-1]


# [visualization-1]
################################################
# Visualization setup
################################################
orange = material("orange", (0.8, 0.1, 0.036), 0, 0)
red = material("red", (1.0, 0.0, 0.0), 0, 0)
green = material("green", (0.0, 0.5, 0.0), 0, 0)
blue = material("blue", (0.019, 0.052, 0.445), 0, 0)
turquoise = material("turquoise ", (0.02, 0.8, 0.77), 0, 0)
transparency_sphere = 0.6
transparent_sphere = ddg.blender.material.material(
    "transparent_sphere", (0.018, 0.313, 0.656), 0.5, 0.5, alpha=transparency_sphere
)


def visualize_circles(g, domain, name, material=None, collection=None):
    l = []
    for i in range(*domain[0]):
        bobj = ddg.blender.convert(g(i), f"{name}_{i}", material, collection)
        bobj.data.bevel_depth = bevel_circle
        l.append(bobj)

    return l


def visualize_pts(g, domain, name, material=turquoise, collection=None):
    return [
        ddg.blender.convert(
            g(i),
            f"{name}_{i}",
            material,
            collection,
        )
        for i in range(*domain[0])
    ]


def visualize_curve(dnet_fct, domain, name, material=None, collection=None):
    dnet = ddg.nets.DiscreteNet(dnet_fct, domain=domain)
    bobj = ddg.blender.convert(dnet, name, material, collection)
    bobj.data.bevel_depth = bevel_curve
    return bobj


def shift_domain(a, b, domain):
    return [[domain[0][0] + a, domain[0][1] + b]]


# [visualization-1]

################################################
# Example
################################################
for example in ["discrete"]:
    ###############################################
    # Initial data
    ###############################################

    if example == "smooth":
        # Helix spiraling around M
        sampling_curve = np.pi / 50

        r = 1.1

        def parametrization(u):
            return [r * np.sin(u), r * np.cos(u), u]

        smooth_domain = [[-2 * np.pi, 1.5 * np.pi]]
        snet = ddg.nets.SmoothNet(parametrization, domain=smooth_domain)
        dnet = ddg.nets.sample_smooth_net(snet, sampling_curve)

        # Starting points of sphere inversions
        e_start_orth_1_phi = -np.pi / 2
        e_start_orth_2_phi = -np.pi / 2
        e_start_env_2_phi = -np.pi / 2

        # Visualization
        bevel_curve = 0.02
        bevel_circle = 0.003
        circle_sampling = [0.1, 100, "c"]
        m_camera_location = (-11, -10, 3)
        e_camera_location = (0.2, 0.65, 6.2)

    # [example-1]
    elif example == "discrete":
        # Conic section in M^+
        sampling_curve = 10
        h = ddg.geometry.Quadric(np.diag([0.8, 0.8, -1, -1]))
        p = ddg.geometry.hyperplane_from_normal([-1, -1, -5], level=-1)
        snet = ddg.to_smooth_net(ddg.geometry.intersect(h, p))
        dnet = ddg.nets.sample_smooth_net(snet, sampling=[sampling_curve, "t"])

        # Starting points of sphere inversions
        e_start_orth_1_phi = np.pi / 8
        e_start_orth_2_phi = np.pi / 8
        e_start_env_2_phi = np.pi / 8

        # Visualization
        bevel_curve = 0.03
        bevel_circle = 0.015
        circle_sampling = [0.1, 100, "c"]
        m_camera_location = (4, -3.3, 2.5)
        e_camera_location = (-1.2, -1.4, 10)

    else:
        raise Exception(f"{example = } has to be one of 'smooth' or 'discrete'")
    # [example-1]

    # [example-2]
    ###############################################
    # Circles
    ###############################################

    # Moebius #####################################
    m_points_fct = function_affine_to_projective(dnet.fct)  # Points in M^+
    m_circles_fct = m_points_to_circles(m_points_fct)  # Conics in M

    # Euclidean ###################################
    e_circles_fct = project_function_of_objects_down(m_circles_fct)
    # [example-2]

    # [example-3]
    ###############################################
    # Envelopes (as intersection points of circles)
    ###############################################

    # Moebius #####################################
    m_env_fct_0_0 = m_envelope(m_points_fct, index=0)  # Points in M
    m_env_fct_1_0 = m_envelope(m_points_fct, index=1)  # Points in M

    # Euclidean ###################################
    e_env_fct_0_0 = function_projective_to_affine(
        project_function_of_objects_down(m_env_fct_0_0)
    )  # affine coordinates
    e_env_fct_1_0 = function_projective_to_affine(
        project_function_of_objects_down(m_env_fct_1_0)
    )  # affine coordinates
    # [example-3]

    # [example-4]
    ###############################################
    # Orthogonal trajectory (inversion in circle family)
    ###############################################

    # Starting point ##############################
    circle0 = e_circles_fct(0)
    e_start_orth_1 = ddg.math.parametrizations.circle(
        e_start_orth_1_phi, circle0.center.affine_point, circle0.radius
    )  # Affine 3d coordinates
    m_start_orth_1 = projective_to_moebius(
        ddg.geometry.subspace_from_affine_points(e_start_orth_1)
    )  # Point in M

    # Moebius #####################################
    # Apply iterative inversions in the circles
    m_orth_fct_1 = m_iterative_sphere_inversions(
        m_points_fct, m_start_orth_1
    )  # Points in M

    # Euclidean ####################################
    e_orth_fct_1 = function_projective_to_affine(
        project_function_of_objects_down(m_orth_fct_1)
    )  # affine coordinates
    # [example-4]

    # [example-5]
    ###############################################
    # Circles of similitude and anti-similitude
    ###############################################

    # Moebius #####################################
    m_homogeneous_points_fct = m_homogeneous_lift_fct(
        e_circles_fct
    )  # Homogeneous coordinates in R^{3, 1}
    m_points_of_similitude_fct = m_points_of_similitude_from_homogeneous_lift(
        m_homogeneous_points_fct
    )  # Points in M+
    m_circles_of_similitude_fct = m_points_to_circles(
        m_points_of_similitude_fct
    )  # Conics in M
    m_points_of_anti_similitude_fct = m_points_of_anti_similitude_from_homogeneous_lift(
        m_homogeneous_points_fct
    )  # Homogeneous coordinates in R^{3, 1}
    m_circles_of_anti_similitude_fct = m_points_to_circles(
        m_points_of_anti_similitude_fct
    )  # Conics in M

    # Euclidean ####################################
    e_circles_of_similitude_fct = project_function_of_objects_down(
        m_circles_of_similitude_fct
    )  # Circles in R2 in R3
    e_circles_of_anti_similitude_fct = project_function_of_objects_down(
        m_circles_of_anti_similitude_fct
    )  # Circles in R2 in R3
    # [example-5]

    # [example-6]
    ###############################################
    # Envelope (as sphere inversions in circles of anti-similitude)
    ###############################################

    # Starting point ##############################
    e_start_env_2 = ddg.math.parametrizations.circle(
        e_start_env_2_phi, circle0.center.affine_point, circle0.radius
    )  # affine_coordinates
    m_start_env_2 = projective_to_moebius(
        ddg.geometry.subspace_from_affine_points(e_start_env_2)
    )  # Point in M

    # Moebius #####################################
    m_env_fct_2 = m_iterative_sphere_inversions(
        m_points_of_anti_similitude_fct, m_start_env_2
    )  # Points in M

    # Euclidean ####################################
    e_env_fct_2 = function_projective_to_affine(
        project_function_of_objects_down(m_env_fct_2)
    )  # affine coordinates
    # [example-6]

    # [example-7]
    ###############################################
    # Orthogonal trajectory (as sphere inversions in circles of similitude)
    ###############################################

    # Starting point ##############################
    e_start_orth_2 = ddg.math.parametrizations.circle(
        e_start_orth_2_phi, circle0.center.affine_point, circle0.radius
    )  # affine_coordinates
    m_start_orth_2 = projective_to_moebius(
        ddg.geometry.subspace_from_affine_points(e_start_orth_2)
    )  # Point in M

    # Moebius #####################################
    m_orth_fct_2 = m_iterative_sphere_inversions(
        m_points_of_similitude_fct, m_start_orth_2
    )

    # Euclidean ####################################
    e_orth_fct_2 = function_projective_to_affine(
        project_function_of_objects_down(m_orth_fct_2)
    )  # affine coordinates
    # [example-7]

    ################################################
    # Visualization
    ################################################
    # [visualization-2]
    mob_col = ddg.blender.collection.collection(
        f"Moebius_{example}",
        children=[
            f"m_circles_{example}",
            f"m_envelope_1_{example}",
            f"m_orthogonal_trajectory_1_{example}",
            f"m_envelope_2_{example}",
            f"m_orthogonal_trajectory_2_{example}",
        ],
    )

    euc_col = ddg.blender.collection.collection(
        f"Euclidean_{example}",
        children=[
            f"e_circles_{example}",
            f"e_envelope_1_{example}",
            f"e_orthogonal_trajectory_1_{example}",
            f"e_envelope_2_{example}",
            f"e_orthogonal_trajectory_2_{example}",
        ],
    )
    # [visualization-2]

    # [visualization-3]
    ###############################################
    # Moebius
    ###############################################
    mob_curve = ddg.blender.convert(dnet, "mob_curve", orange, mob_col)
    mob_curve.data.bevel_depth = 0.015
    mob_quadric_bobj = ddg.blender.convert(
        mob.absolute,
        "mob_quadric",
        transparent_sphere,
        mob_col,
    )
    ddg.blender.mesh.shade_smooth(mob_quadric_bobj)

    #  dnet.domain = ddg.nets.DiscreteDomain([[*dnet.domain[0]]])
    # # circles
    visualize_circles(
        m_circles_fct,
        domain=shift_domain(0, 1, dnet.domain),
        name="m_circle",
        material=blue,
        collection=mob_col.children[0],
    )

    # envelopes 1 (intersection points)
    visualize_pts(
        m_env_fct_0_0,
        domain=dnet.domain,
        name="m_env_1_point_0",
        collection=mob_col.children[1],
    )
    visualize_pts(
        m_env_fct_1_0,
        domain=dnet.domain,
        name="m_env_1_point_1",
        collection=mob_col.children[1],
    )

    # orthogonal_trajectories 1
    visualize_pts(
        m_orth_fct_1,
        domain=shift_domain(1, 3, dnet.domain),
        name="m_orth_traj_1_point",
        collection=mob_col.children[2],
    )

    # envelopes 2 (circles of anti similitude)
    visualize_circles(
        m_circles_of_anti_similitude_fct,
        domain=dnet.domain,
        name="m_circle_of_anti_similitude",
        material=red,
        collection=mob_col.children[3],
    )
    visualize_pts(
        m_env_fct_2,
        domain=shift_domain(0, 3, dnet.domain),
        name="m_env_2_point",
        collection=mob_col.children[3],
    )

    # orthogonal trajectories 2
    visualize_circles(
        m_circles_of_similitude_fct,
        domain=dnet.domain,
        name="m_circle_of_similitude",
        material=green,
        collection=mob_col.children[4],
    )
    visualize_pts(
        m_orth_fct_2,
        name="m_orth_traj_2_point",
        domain=shift_domain(0, 2, dnet.domain),
        collection=mob_col.children[4],
    )

    # [visualization-3]
    # [visualization-4]
    ###############################################
    # Euclidean
    ###############################################
    # circles
    visualize_circles(
        e_circles_fct,
        domain=shift_domain(0, 1, dnet.domain),
        name="e_circle",
        material=blue,
        collection=euc_col.children[0],
    )

    # envelopes 1 (intersection points)
    visualize_curve(
        e_env_fct_0_0,
        domain=shift_domain(0, 1, dnet.domain),
        name="e_env_1_0",
        material=turquoise,
        collection=euc_col.children[1],
    )
    visualize_curve(
        e_env_fct_1_0,
        domain=shift_domain(0, 1, dnet.domain),
        name="e_env_1_1",
        material=turquoise,
        collection=euc_col.children[1],
    )

    # orthogonal_trajectories 1
    visualize_curve(
        e_orth_fct_1,
        domain=shift_domain(0, 2, dnet.domain),
        name="e_orth_traj_1",
        material=turquoise,
        collection=euc_col.children[2],
    )

    # envelopes 2 (cirlces of anti similitude)
    visualize_circles(
        e_circles_of_anti_similitude_fct,
        domain=dnet.domain,
        name="e_circle_of_anti_similitude",
        material=red,
        collection=euc_col.children[3],
    )
    visualize_circles(
        e_circles_of_similitude_fct,
        domain=dnet.domain,
        name="e_circle_of_similitude",
        material=green,
        collection=euc_col.children[4],
    )
    visualize_curve(
        e_env_fct_2,
        domain=dnet.domain,
        name="e_env_2",
        collection=euc_col.children[3],
        material=turquoise,
    )

    # orthogonal trajectories 2
    visualize_curve(
        e_orth_fct_2,
        domain=shift_domain(0, 1, dnet.domain),
        name="e_orth_traj_2",
        material=turquoise,
        collection=euc_col.children[4],
    )
    # [visualization-4]
    #############################################
    # CAMERAS
    #############################################

    camera_m = ddg.blender.camera.camera(location=m_camera_location, collection=mob_col)
    camera_e = ddg.blender.camera.camera(location=e_camera_location, collection=euc_col)
    ddg.blender.camera.look_at_point(camera_m, (0, 0, 0))
    ddg.blender.camera.look_at_point(camera_e, (0, 0, 0))
    camera_e.rotation_euler = (0, 0, 0)

#############################################
# RENDERING AND LIGHT
#############################################

# Change to cycles rendering engine, higher max_samples size
# enable film transparency, and scale up the resolution for a good quality image
ddg.blender.render.setup_cycles_renderer()
ddg.blender.render.set_film_transparency()

ddg.blender.render.set_world_background(color=(1, 1, 1, 1), strength=0.6)

resolution = 1000
bpy.context.scene.render.resolution_x = resolution
bpy.context.scene.render.resolution_y = resolution

light = ddg.blender.light.light(location=(0, 0, 100), type_="SUN", energy=0.5)