import numpy as np
from matplotlib import cm

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

b = 8000   # number of curves
n = 80     # nodes per curve
dt = 0.06  # nodes distance

inp = np.zeros((b, 3, 4), dtype=np.float32)
xyz = np.stack((
    np.random.normal(loc=0, scale=1.0, size=b),
    np.random.normal(loc=0, scale=1.0, size=b),
    np.random.normal(loc=0, scale=1.0, size=b))).T
for c in range(b):
    mag = np.linalg.norm(xyz[c])
    xyz[c] *= np.sqrt(mag) / mag

ofs = 50 * np.random.rand(3)
for c in range(b):
    inp[c,:,:3] = xyz[c]
    inp[c,:,3] = ofs            # sync the 4'th dim of the noise

pos = np.zeros((b, n, 3), dtype=np.float32)
col = np.zeros((b, n, 3), dtype=np.float32)
r = np.zeros((b, n), dtype=np.float32)

rnd = simplex(inp)

rprev = np.copy(rnd)
for t in range(n):
    rt = 2.0 * (t+1) / (n+2) - 1
    rt = 1 - rt*rt
    r[:,t] = 0.07 * rt * rt
    for c in range(b):
        mag = np.linalg.norm(rnd[c])
        r[c,t] *= 0.2 + 0.8 * mag
        
        rnd[c] *= (dt/mag)      # normalize and scale the step size
        inp[c,:,:3] += rnd[c]   # step in the field direction
        pos[c,t] = inp[c,0,:3]
        
        fi = (1/(dt*dt)) * np.dot(rnd[c], rprev[c])
        cc = cm.get_cmap("bone")(np.power(2*fi-1,19))[:3]
        col[c,t] = make_color(cc)
        
    rprev = np.copy(rnd)
    rnd = simplex(inp, rnd)     # noise at the next pos

rt = TkOptiX(start_now=False)
rt.set_param(
    min_accumulation_step=4,
    max_accumulation_frames=500,
    rt_timeout=100000,
    light_shading="Soft"
)
rt.set_uint("path_seg_range", 4, 10)
    
exposure = 1.2; gamma = 2.2
rt.set_float("tonemap_exposure", exposure)
rt.set_float("tonemap_gamma", gamma)
rt.set_float("denoiser_blend", 0.25)
rt.add_postproc("Denoiser")
    
m_metallic["VarFloat"]["base_roughness"] = 0.004
rt.setup_material("metal", m_metallic)
rt.setup_material("plastic", m_plastic)
    
rt.setup_camera("dof_cam", eye=[0, 0, 12], target=[0, 0, 0], fov=40, focal_scale=0.86, cam_type="DoF")

rt.setup_light("l1", pos=[8, -3, 13], color=1.5*np.array([0.99, 0.97, 0.93]), radius=5)
rt.setup_light("l2", pos=[-17, -7, 5], u=[0, 0, -10], v=[0, 14, 0], color=1*np.array([0.25, 0.28, 0.35]), light_type="Parallelogram")

for c in range(b):
    if np.random.uniform() < 0.1:
        rt.set_data("c"+str(c), pos=pos[c], r=2.1*r[c], c=col[c], geom="BezierChain", mat="plastic")
    else:
        rt.set_data("c"+str(c), pos=pos[c], r=0.33*r[c], c=[0.94, 0.93, 0.9], geom="BezierChain", mat="metal")

rt.show()