from functools import lru_cache

import bpy
import numpy as np

import ddg

# [helper-functions]
##########################################
#  HELPER FUNCTIONS
##########################################

euc_3 = ddg.geometry.euclidean(3)


def circles(fct):
    """
    Returns a function that returns the discrete curvature circle at given index
    of the given discrete curve fct.
    The i'th curvature circle is given by the unique circle through the
    points fct(i-1), fct(i), fct(i+1).
    Domain: from 1 to n-1.
    """

    @lru_cache(maxsize=128)
    def circle(i):
        c, r, n = ddg.math.euclidean.circle_through_three_points(
            fct(i - 1), fct(i), fct(i + 1), atol=1e-17
        )
        euc_circle = euc_3.sphere_from_affine_point_and_normals(c, r, n)
        return euc_3.sphere_to_quadric(euc_circle)

    return circle


def inverse_project_fct(fct, sphere):
    """
    Returns a function that returns the projection of
    a given discrete curve fct to the quadric `sphere`.
    Each vertex of the discrete curve is inverse stereographically
    projected to `sphere`.
    """

    @lru_cache(maxsize=128)
    def inverse_project_v(i):
        p = ddg.geometry.Point(ddg.math.projective.homogenize(fct(i)))
        return ddg.geometry.inverse_stereographic_project(p, quadric=sphere)

    return inverse_project_v


def m_circles(g, sphere):
    """
    Assumes, that the function g returns Points on the quadric `sphere`.
    Returns a function that returns the Moebius circles.
    The i'th Moebius circle is given by the intersection of the sphere with
    the plane spanned by the points fct(i-1), fct(i), fct(i+1).
    Domain: from 1 to n-1.
    """

    @lru_cache(maxsize=128)
    def m_circle(i):
        plane = ddg.geometry.join(g(i - 1), g(i), g(i + 1))
        return ddg.geometry.meet(plane, sphere)

    return m_circle


def project_fct(g, plane):
    """
    Assumes, that the function fct returns projective objects
    in RP3.
    Returns a function that projects the objects to the plane, resulting
    in planar objects of the same type.
    Domain: from 0 to n.
    """

    @lru_cache(maxsize=128)
    def project_obj(i):
        proj_on_plane = ddg.geometry.stereographic_project(g(i), hyperplane=plane)
        return proj_on_plane

    return project_obj


# [helper-functions]

# [setup]
#############################################
# CLEAR AND ADD COLLECTION FOR STATIC OBJECTS
#############################################

# Animated objects will be linked to their own collections
ddg.blender.scene.clear(remove_collections=True)
static_coll = ddg.blender.collection.collection("static")
animated_coll = ddg.blender.collection.collection("animated objects")

##########################################
# EUC PLANE AND MOB SPHERE
##########################################

plane_pt, plane_directions = [0, 0, -1], [[1, 0, 0], [0, 1, 0]]
plane = ddg.geometry.subspace_from_affine_points_and_directions(
    plane_pt, plane_directions
)
north_pole = ddg.geometry.euclidean_models.MoebiusModel(2).fixed_point
sphere = ddg.geometry.Quadric(np.diag([1, 1, 1, -1]))
# [setup]

# [curve-and-functions]
##########################################
# PLANE CURVE
##########################################
# (-a,0) and (a,0) are  the centers of circles which generate cardioid
a = 0.3
# translation of cardioid
b = [2.3, 0]


def f(t):
    return np.array(
        [
            2 * a * (1 - np.cos(t)) * np.cos(t) + b[0],
            2 * a * (1 - np.cos(t)) * np.sin(t) + b[1],
            -1,
        ]
    )


samples = 20  # 24 * 6 - 1
sampling_curve = [samples, "t"]
snet = ddg.nets.SmoothNet(f, [[0, 2 * np.pi, True]])
dnet = ddg.nets.sample_smooth_net(snet, sampling=sampling_curve)

##########################################
# FUNCTION CALLS
##########################################

# circles_fct(i) returns i'th curvature circle (as Quadric)
circles_fct = circles(dnet.fct)


# centers_fct(i) returns the center of the i'th curvature circle (as Point/Subspace)
@lru_cache(maxsize=128)
def centers_fct(i):
    circle = euc_3.quadric_to_sphere(circles_fct(i))
    return ddg.geometry.Point(circle.center.point)


# spherical_fct(i) returns the inverse stereographic
# projection of the i'th vertex (as Point/Subspace)
spherical_fct = inverse_project_fct(dnet.fct, sphere)
# m_circles_fct(i) returns the i'th curvature circle on the Moebius sphere (as Quadric)
m_circles_fct = m_circles(spherical_fct, sphere)


# m_centers_fct(i) returns polar point of the plane corresponding
# to the i'th curvature circle (as Point/Subspace)
@lru_cache(maxsize=128)
def m_centers_fct(i):
    return sphere.polarize(m_circles_fct(i).subspace)


# Function returning the stereographic projection of the projective objects.
projected_vertices_fct = project_fct(spherical_fct, plane)
projected_m_circles_fct = project_fct(m_circles_fct, plane)
projected_m_centers_fct = project_fct(m_centers_fct, plane)
# [curve-and-functions]

# [visualization]
##########################################
#  VISUALIZATION
##########################################

# Materials
orange = ddg.blender.material.material("orange", (1, 0.026, 0))
blue = ddg.blender.material.material("blue", (0.019, 0.052, 0.445))
lightgreen = ddg.blender.material.material("lightgreen", (0, 0.42, 0.078), 0, 0)
transparent_cone = ddg.blender.material.material(
    "transparent_cone", (0.2, 0.154, 0.154), alpha=0.3
)
transparency_sphere = 0.8
transparent_sphere = ddg.blender.material.material(
    "transparent_sphere", (0, 0.168, 0.6), 0.5, alpha=transparency_sphere
)
transparent_plane = ddg.blender.material.material(
    "transparent_plane", (0, 0.05, 0.005), 0.4
)
transparent_m_planes = ddg.blender.material.material(
    "transparent_m_planes", (0.06, 0.06, 0.06), alpha=0.9
)

# Distinguish weather to use settings for a smooth
# or a discrete version of the animation.
# For more than 40 samples we interpret the
# curve as a smooth curve.
discrete_animation = samples < 40

# Sampling and size settings for objects to visualize
plane_size = (4.5, 4.5)

bevel_curvature_circles = 0.01 if discrete_animation else 0.008
bevel_curve = 0.01

sampling_general = [0.1, 100, "c"]
sampling_polar_plane = 1
point_size = 0.02
# [visualization]

# [animation-1]
##########################################
#  ANIMATION
##########################################
# Animation 1: Euclidean curvature circles
# Animation 2: Euclidean curvature circles with tracking of the evolute
# Animation 3: Fade in of the Moebius sphere
# Animation 4: Projection of planar vertices to the sphere
# Animation 5: Moebius circles as intersections of the sphere
# with planes through three consecutive points
# Animation 6: Project circles and polar points back to the plane
# Animation 7: Trace curvature circles on the plane and on the sphere

# For internal snapshot testing set animation_no=6
# For internal testing the max_samples are set low.
# For a qualitative good picture increase these, e.g. to 512
animation_no = 6
output_dir = "/tmp/"
samples = 8


# Static objects
# ----------------

# Plane
plane_bobj = ddg.blender.convert(
    plane,
    "plane",
    material=transparent_plane,
    collection=static_coll,
)

# Curve
bobj = ddg.blender.convert(
    dnet,
    "discrete curve",
    material=orange,
    collection=static_coll,
)
bobj.data.bevel_depth = bevel_curve

# Moebius sphere
if animation_no > 2:
    sphere_bobj = ddg.blender.convert(
        sphere,
        "sphere",
        material=transparent_sphere,
        collection=static_coll,
    )
    ddg.blender.mesh.shade_smooth(sphere_bobj)


# Fade in of the Moebius sphere
if animation_no == 3:
    frames_fade_in_sphere = 4 * 24
    ddg.blender.animation.animate_opacity(
        transparent_sphere,
        0,
        frames_fade_in_sphere,
        frames_fade_in_sphere + 1,
        frames_fade_in_sphere + 1,
        opacity_outer=0,
        opacity_inner=transparency_sphere,
    )

# Moebius curve
# discrete: set of Points, smooth: DiscreteNet -> Blender Curve object
if animation_no > 4:
    if discrete_animation:
        for i in range(samples):
            ddg.blender.vertices(
                spherical_fct(i),
                f"spherical point {i}",
                radius=point_size,
                material=orange,
                collection=static_coll,
            )
    else:

        def spherical_fct_affine(i):
            return spherical_fct(i).affine_point

        bobj_spherical_curve = ddg.blender.convert(
            ddg.nets.DiscreteNet(spherical_fct_affine, [[0, samples]]),
            "spherical curve",
            material=orange,
            collection=static_coll,
        )
        bobj_spherical_curve.bevel_depth = bevel_curve


# [animation-1]

# [animation-2]
def blender_objects(i):
    ddg.blender.collection.clear([animated_coll], deep=True)
    pts = [
        ddg.geometry.subspace_from_affine_points(p)
        for p in [dnet(i - 1), dnet(i), dnet(i + 1)]
    ]

    ##########################################
    #  PLANAR
    ##########################################

    # Vertices of the curve
    if animation_no in [1, 2, 4, 5] and discrete_animation:
        pts_tmp = [pts[1]] if animation_no == 4 else pts
        for i, p in enumerate(pts_tmp):
            ddg.blender.vertices(
                p,
                "animated planar point",
                radius=point_size,
                material=orange,
                collection=animated_coll,
            )

    # Curvature circles
    if animation_no in [1, 2, 6, 7]:
        bobj_circle = ddg.blender.convert(
            circles_fct(i),
            f"animated circle {i}",
            material=blue,
            collection=animated_coll,
        )
        bobj_circle.data.bevel_depth = bevel_curvature_circles
        bobj_circle

    # Curvature circe center and evolute trace
    if animation_no in [2, 6, 7]:
        # Curvature circle center
        ddg.blender.vertices(
            centers_fct(i),
            f"animated center {i}",
            radius=point_size,
            material=lightgreen,
            collection=animated_coll,
        )

        # Evolute curve
        def centers_fct_affine(i):
            return centers_fct(i).affine_point

        bobj_evolute = ddg.blender.convert(
            ddg.nets.DiscreteNet(centers_fct_affine, [[0, i]]),
            f"evolute {i}",
            material=lightgreen,
            collection=animated_coll,
        )
        bobj_evolute.data.bevel_depth = bevel_curve

    #########################################
    #  PROJECTION TO THE SPHERE
    ##########################################

    # Projection to sphere visualized as line
    if animation_no in [4, 5, 6]:
        pts_tmp = [pts[1]] if animation_no == 4 else pts
        # Line as join of the North Pole and curve's vertex
        if animation_no == 4 or discrete_animation:
            for j, p in enumerate(pts_tmp):
                line_segment = ddg.arrays.line_segment_from_points(
                    north_pole.affine_point, p.affine_point
                )
                bobj_spherical_projection = ddg.blender.convert(
                    line_segment,
                    f"spherical projection {i}_{j}",
                    material=orange,
                    collection=animated_coll,
                )
                bobj_spherical_projection.data.bevel_depth = bevel_curvature_circles

    # Track of Moebius curve, discrete: set of Points, smooth: DiscreteNet
    if animation_no in [4]:
        if discrete_animation:
            ddg.blender.vertices(
                spherical_fct(i),
                f"spherical point {i}",
                radius=point_size,
                material=orange,
                collection=animated_coll,
            )
        else:

            def spherical_fct_affine(i):
                return spherical_fct(i).affine_point

            bobj_spherical_curve = ddg.blender.convert(
                ddg.nets.DiscreteNet(spherical_fct_affine, [[0, i]]),
                f"spherical curve {i}",
                material=orange,
                collection=animated_coll,
            )
            bobj_spherical_curve.bevel_depth = bevel_curve

    ##########################################
    #  MOEBIUS PLANES AND CIRCLES
    ##########################################
    # Catch if a circle degenerates in a cusp of the curve
    non_degenerate = ddg.geometry.signatures.Signature(2, 1, 0)
    if m_circles_fct(i).signature() == non_degenerate:
        # Moebius planes
        if animation_no in [5]:
            center = ddg.geometry.quadric_to_euclidean_sphere(m_circles_fct(i)).center
            plane = m_circles_fct(i).subspace.orthonormalize_and_center(center=center)
            ddg.blender.convert(
                plane,
                f"Moebius plane {i}",
                material=transparent_m_planes,
                collection=animated_coll,
            )

        # Moebius curvature circle
        if animation_no in [5, 6, 7]:
            # Track in animation_no = 7
            bobj_m_circle = ddg.blender.convert(
                m_circles_fct(i),
                f"Moebius curvature circle {i}",
                material=blue,
                collection=animated_coll,
            )
            bobj_m_circle.data.bevel_depth = bevel_curvature_circles
            bobj_m_circle

        # Moebius polar point
        if animation_no in [5, 6, 7]:
            # Track in animation_no = 7
            bobj_m_center = ddg.blender.vertices(
                m_centers_fct(i),
                f"Moebius polar point {i}",
                radius=point_size,
                material=lightgreen,
                collection=animated_coll,
            )
            bobj_m_center

    ##########################################
    #  PROJECT BACK
    ##########################################
    # Projection cone and line
    if animation_no in [6]:
        cone = ddg.geometry.join(north_pole, circles_fct(i))

        # Projection cone
        bobj_projection_cone = ddg.blender.convert(
            cone,
            f"projection cone {i}",
            material=transparent_cone,
            collection=animated_coll,
        )
        ddg.blender.mesh.shade_smooth(bobj_projection_cone)

        # Cut cone on one side
        with ddg.blender.bmesh.bmesh_from_mesh(bobj_projection_cone.data) as bm:
            ddg.blender.bmesh.cut_bounding_box(
                bm, np.array([3, 3, 3]), np.array([3, 0, 0])
            )

        # Projection line
        line_segment = ddg.arrays.line_segment_from_points(
            north_pole.affine_point, centers_fct(i).affine_point
        )
        bobj_projection_line = ddg.blender.convert(
            line_segment,
            f"projection line {i}",
            material=lightgreen,
            collection=animated_coll,
        )
        bobj_projection_line.data.bevel_depth = bevel_curve


# [animation-2]

# [animation-3]
##########################################
# Animation Setup
##########################################

ddg.blender.animation.clear_animation_data(bpy.context.scene)

ddg.blender.props.add_props_with_callback(blender_objects, ("i"), 0)

ddg.blender.animation.set_keyframe(bpy.context.scene, 0, "i", 0)
ddg.blender.animation.set_keyframe(bpy.context.scene, samples, "i", samples)

bpy.context.scene.frame_end = samples if animation_no != 3 else frames_fade_in_sphere
bpy.context.scene.frame_start = 0
bpy.context.scene.frame_current = 0

# [animation-3]

# [rendering-setup]
##########################################
# RENDERING SETUP
##########################################

# Add light and camera
cam = ddg.blender.camera.camera(location=(5.1, -2, 1.5))
ddg.blender.camera.look_at_point(cam, (0, 0.5, -1))
light = ddg.blender.light.light(location=(4, 0, 20), energy=1)
ddg.blender.light.look_at_point(light, (5, 0, 0))
light.data.angle = 0.17

# Setup rendering
ddg.blender.render.setup_cycles_renderer(samples=samples)
ddg.blender.render.set_world_background((0.7, 0.7, 0.7, 1), 1)
bpy.context.scene.render.resolution_x = 1500
bpy.context.scene.render.resolution_y = 1500
ddg.blender.render.set_film_transparency()
bpy.context.scene.view_settings.view_transform = "Standard"
ddg.blender.render.set_render_output_images(
    output_dir,
    time=False,
    file_format="PNG",
    alpha=True,
)

# [rendering-setup]