"""
Be aware that the example, at the end, contains settings for internal snapshot testing.
Comment or remove these to apply the actual (rendering) settings of the example.
"""

import bpy
import numpy as np

import ddg
from ddg.blender.camera import camera
from ddg.blender.light import light, look_at_point
from ddg.blender.material import material
from ddg.blender.render import (
    set_film_transparency,
    set_world_background,
    setup_cycles_renderer,
)

# [setup]
#################
# Initial setup #
#################

ddg.blender.scene.clear()

collection_mathematical = ddg.blender.collection.collection(
    "theoretic_caustics", children=["cardioid", "nephroid"]
)
collection_physical = ddg.blender.collection.collection("real_caustics")
# [setup]


# [helper-functions]
####################
# Helper functions #
####################


def reflect_ray(ray_segment, mirror_line):
    """Reflect a ray in a line in R^2.

    Parameters
    ----------
    ray_segment : tuple of numpy.ndarray and numpy.ndarray
        Two points on the ray in affine coordinates.
    mirror_line : ddg.geometry.Subspace
        A line in R^2.

    Returns
    -------
    mirrored_ray_origin : numpy.ndarray
        The intersection of the ray with the mirror_line in affine coordinates.
        This is origin of the mirrored ray.
    mirrored_ray_direction : numpy.ndarray
        The direction of the mirrored ray, a vector in R^2.
    """
    line = ddg.geometry.subspace_from_affine_points(*ray_segment)
    reflection = ddg.geometry.euclidean_models.ProjectiveModel.reflect_in_hyperplane(
        line, mirror_line
    )
    p, q = reflection.affine_points  # p is the intersection of line and mirror_line
    d = q - p
    return p, d / np.linalg.norm(d)


def emitted_ray_segment(p, light, length):
    """A segment from the ray from light to p containing p.

    Parameters
    ----------
    p : numpy.ndarray
        A point in R^2 in affine coordinates.
    light : numpy.ndarray
        A point in RP^2 in projective coordinates. The light can be at infinity.
    length : float
        The length of the segment if the line is at infinity.

    Returns
    -------
    p, q
        q is a numpy.ndarray representing the affine coordinates of a point on
        the ray. If the light is at infinity, |q - p| == length and q is
        between p and the light. Otherwise q represents the affine coordinates
        of the light.
    """
    if light[-1] == 0:
        d = light[:-1]
        d = d / np.linalg.norm(d)
        return p, p + length * d
    else:
        return p, light[:-1] / light[-1]


def reflected_ray_segment(emitted_ray_segment, tangent, construction):
    """A segment from the reflected emitted ray.

    Parameters
    ----------
    emitted_ray_segment : tuple of numpy.ndarray and numpy.ndarray of shape (2,)
        Two points on the emitted ray.
    tangent : ddg.geometry.Subspace
        The tangent on which the emitted ray is reflected.
    construction : {"cardioid", "nephroid"}
        This determines the length of the reflected ray segment. They are
        chosen such that the picture looks pretty.

    Returns
    -------
    p, q : numpy.ndarray of shape (2,)
        Two points on the reflected emitted ray.
    """
    a, b = emitted_ray_segment
    p, d = reflect_ray(emitted_ray_segment, tangent)
    if construction == "cardioid":
        return p, p - np.linalg.norm(a - b) * d
    elif construction == "nephroid":
        return p, p - d
    else:
        raise ValueError(f"{construction = } must be one of 'cardioid' or 'nephroid'")


def envelope(reflected_ray_segments, periodicity):
    """The envelope of a curve.

    Let r_1, ..., r_n be the rays reflected at the curve. This function returns
    the intersections of

        - (r_1, r_2), ..., (r_{n-2}, r_{n-1}) if periodicity is False and
        - (r_n, r_1), (r_1, r_2), ..., (r_{n-1}, r_n}) if periodicity is True.

    Parameters
    ----------
    reflected_ray_segments : sequence of pair numpy.ndarrays of shape (2,)
        These are the light rays that are reflected by the curve.
        Each ray is encoded by a pair of a point p on the ray and the direction
        d of the ray.
    periodicity : bool
        The periodicity of the curve.

    Returns
    -------
    list of numpy.ndarray of shape (2,)
    """
    reflected_rays = [
        ddg.geometry.subspace_from_affine_points(p, d)
        for p, d in reflected_ray_segments
    ]
    intersections = [
        ddg.geometry.intersect(r, s).affine_point
        for r, s in zip(np.roll(reflected_rays, 1), reflected_rays)
    ]
    if periodicity:
        return intersections
    else:
        return intersections[1:-1]


# [helper-functions]


# [parameters]
###########################
# Initial values and data #
###########################


example = "discrete"
if example == "smooth":
    num = 80
    thick_line = 0.02
    thin_line = 0.005
elif example == "discrete":
    num = 16
    thick_line = 0.02
    thin_line = 0.002
else:
    raise Exception(f"{example = } has to be one of 'discrete' or 'smooth'")


####################
# Parameterization #
####################
# We define a parametrization for a curve.


def circle(u):
    return np.column_stack([np.sin(u), np.cos(u)])


# [parameters]

# [visualization]
#################
# Visualization #
#################


orange = material("orange", (0.8, 0.1, 0.036), 0, 0)
red = material("red", (0.8, 0.01, 0.036), 0, 0)
blue = material("blue", (0.019, 0.052, 0.445), 0, 0)
turquoise = material("turquoise", (0.02, 0.8, 0.77), 0, 0)
dark_green = material("dark_green", (0.0, 0.015, 0.0), 1, 0.5)


# [visualization]


# [theoretical-construction]
###########################################
# Creating theoretical caustic in Blender #
###########################################


for construction in ["cardioid", "nephroid"]:
    # Setup
    if construction == "nephroid":
        light_origin = np.array([1, 0, 0])  # projective coordinates
        interval = (-np.pi, 0)
        sampled_interval = np.linspace(*interval, num)
        periodicity = False
        emitted_ray_segment_length = 5
    elif construction == "cardioid":
        light_origin = np.array([1, 0, 1])  # projective coordinates
        interval = (-np.pi, np.pi)
        sampled_interval = np.linspace(*interval, num, False)
        periodicity = True
        emitted_ray_segment_length = 1
    else:
        raise Exception()

    samples = circle(sampled_interval)
    curve = ddg.arrays.Curve(samples, periodicity)

    # Create caustic as envelope of reflected light rays
    edge_midpoints = [
        (p + q) / 2 for p, q in zip(samples, np.roll(samples, -1, axis=0))
    ]
    tangents = [
        ddg.geometry.subspace_from_affine_points(p, q)
        for p, q in zip(samples, np.roll(samples, -1, axis=0))
    ]
    emitted_ray_segments = [
        emitted_ray_segment(m, light_origin, emitted_ray_segment_length)
        for m in edge_midpoints
    ]
    reflected_ray_segments = [
        reflected_ray_segment(irs, t, construction)
        for irs, t in zip(emitted_ray_segments, tangents)
    ]
    caustic_ = envelope(reflected_ray_segments, periodicity)

    # Visualize objects
    collection = collection_mathematical.children[construction]
    curve_bobj = ddg.blender.convert(curve, f"curve {construction}", orange, collection)
    curve_bobj.data.bevel_depth = thick_line

    for i, ray_segment in enumerate(emitted_ray_segments):
        ray_bobj = ddg.blender.convert(
            ddg.arrays.line_segment_from_points(*ray_segment),
            f"emitted ray {i} {construction}",
            blue,
            collection,
        )
        ray_bobj.data.bevel_depth = thin_line

    for i, ray_segment in enumerate(reflected_ray_segments):
        ray_bobj = ddg.blender.convert(
            ddg.arrays.line_segment_from_points(*ray_segment),
            f"reflected ray {i} {construction}",
            turquoise,
            collection,
        )
        ray_bobj.data.bevel_depth = thin_line

    caustic_bobj = ddg.blender.convert(
        ddg.arrays.Curve(caustic_, periodicity),
        f"caustic {construction}",
        red,
        collection,
    )
    caustic_bobj.data.bevel_depth = thick_line

# Add camera
camera_theoretical = camera(
    type_="ORTHO", location=(0, 0, 15), collection=collection_mathematical
)
camera_theoretical.rotation_euler.z = np.pi

# Rendering setup
set_film_transparency()
# [theoretical-construction]


# [physical-construction]
#########################################
# Creating real caustic in Blender
#########################################

# Curve. Additionally, in Blender a Material consisting of a
# Glossy node with color=(0.8, 0.1, 0.036) and roughness=0.1 was added
discrete_curve_real_bobj = ddg.blender.convert(
    curve, "Discrete Curve Real", collection=collection_physical
)
discrete_curve_real_bobj.scale = (1, 1, 20)

# Plane
plane = ddg.geometry.coordinate_hyperplane(3)
bobj_plane = ddg.blender.convert(plane, "plane", dark_green, collection_physical)

# Cycles settings
setup_cycles_renderer(samples=2048, device="GPU")
set_world_background(color=(0.5, 0.5, 0.5, 1), strength=1)

# Overall caustic settings
bpy.context.scene.cycles.blur_glossy = 0
bpy.context.scene.cycles.sample_clamp_indirect = 0

# Light
sun = light(energy=30, collection=collection_physical)
look_at_point(sun, (-5, 0, -1))
sun.data.angle = 0.001
sun.data.cycles.max_bounces = 1

# Camera
camera_physical = camera(location=(1.95, 0, 2.5), collection=collection_physical)
camera_physical.rotation_euler = (np.pi / 4, 0, np.pi / 2)
bpy.context.scene.render.resolution_x = 2000
bpy.context.scene.render.resolution_y = 2000
# [physical-construction]