import numpy as np

# docs, examples: https://plotoptix.rnd.team
from plotoptix import TkOptiX
from plotoptix.utils import make_color
from plotoptix.materials import m_plastic

def trefoil(u, v, r):
    x = r * np.sin(3 * u) / (2 + np.cos(v))
    y = r * (np.sin(u) + 2 * np.sin(2 * u)) / (2 + np.cos(v + np.pi * 2 / 3))
    z = r / 2 * (np.cos(u) - 2 * np.cos(2 * u)) * (2 + np.cos(v)) * (2 + np.cos(v + np.pi * 2 / 3)) / 4
    return np.array([x, y, z], dtype=np.float32)

r = (-np.pi, 3*np.pi)
n = 500

i = np.linspace(r[0], r[1], n)

U, V = np.meshgrid(i, i)
S = np.swapaxes(trefoil(U, V, 5), 0, 2)

rt = TkOptiX(width=1500, height=1800)

rt.set_param(min_accumulation_step=2,
             max_accumulation_frames=500,
             light_shading="Hard")
rt.set_uint("path_seg_range", 6, 15)

rt.setup_material("plastic", m_plastic)

exposure = 0.8; gamma = 2.2
rt.set_float("tonemap_exposure", exposure)
rt.set_float("tonemap_gamma", gamma)
rt.add_postproc("Gamma")

rt.set_background(0)
rt.set_ambient(0.15)

rt.set_surface("surface", S, c=0.94, make_normals=True, mat="plastic")

rt.set_data("plane", geom="Parallelograms",
            pos=[[-100, -14, -100]], u=[200, 0, 0], v=[0, 0, 200],
            c=make_color([0.1, 0.2, 0.3], exposure=exposure, gamma=gamma))

rt.setup_camera("cam1", cam_type="DoF",
                eye=[-50, -7, -15], target=[0, 0, -1], up=[0, 1, 0],
                aperture_radius=0.4, aperture_fract=0.2,
                focal_scale=0.92, fov=35)

rt.setup_light("light1", pos=[-15, 20, 15], color=8, radius=6)

rt.start()