# Copyright (c) 2022 NVIDIA CORPORATION.  All rights reserved.
# NVIDIA CORPORATION and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION is strictly prohibited.

import numpy as np
# import matplotlib.pyplot as plt

import warp as wp
from warp.tests.test_base import *

wp.init()

@wp.kernel
def test_kernel(
    kernel_seed: int,
    int_a: wp.array(dtype=int),
    int_ab: wp.array(dtype=int),
    float_01: wp.array(dtype=float),
    float_ab: wp.array(dtype=float)):

    tid = wp.tid()

    state = wp.rand_init(kernel_seed, tid)

    int_a[tid] = wp.randi(state)
    int_ab[tid] = wp.randi(state, 0, 100)
    float_01[tid] = wp.randf(state)
    float_ab[tid] = wp.randf(state, 0.0, 100.0)

def test_rand(test, device):

    N = 10

    int_a_device = wp.zeros(N, dtype=int, device=device)
    int_a_host = wp.zeros(N, dtype=int, device="cpu")
    int_ab_device = wp.zeros(N, dtype=int, device=device)
    int_ab_host = wp.zeros(N, dtype=int, device="cpu")

    float_01_device = wp.zeros(N, dtype=float, device=device)
    float_01_host = wp.zeros(N, dtype=float, device="cpu")
    float_ab_device = wp.zeros(N, dtype=float, device=device)
    float_ab_host = wp.zeros(N, dtype=float, device="cpu")

    seed = 42

    wp.launch(
        kernel=test_kernel,
        dim=N,
        inputs=[seed, int_a_device, int_ab_device, float_01_device, float_ab_device],
        outputs=[],
        device=device
    )

    wp.copy(int_a_host, int_a_device)
    wp.copy(int_ab_host, int_ab_device)
    wp.copy(float_01_host, float_01_device)
    wp.copy(float_ab_host, float_ab_device)
    wp.synchronize()

    int_a = int_a_host.numpy()
    int_ab = int_ab_host.numpy()
    float_01 = float_01_host.numpy()
    float_ab = float_ab_host.numpy()

    int_a_true = np.array([-575632308, 59537738, 1898992239, 442961864, -1069147335, -478445524, 1803659809, 2122909397, -1888556360, 334603718])
    int_ab_true = np.array([46, 58, 46, 83, 85, 39, 72, 99, 18, 41])
    float_01_true = np.array([0.72961855, 0.86200964, 0.28770837, 0.8187722, 0.186335, 0.6101239, 0.56432086, 0.70428324, 0.64812654, 0.27679986])
    float_ab_true = np.array([96.04259, 73.33809, 63.601555, 38.647305, 71.813896, 64.65809, 77.79791, 46.579605, 94.614456, 91.921814])

    test.assertTrue((int_a == int_a_true).all())
    test.assertTrue((int_ab == int_ab_true).all())

    err = np.max(np.abs(float_01 - float_01_true))
    test.assertTrue(err < 1e-04)

    err = np.max(np.abs(float_ab - float_ab_true))
    test.assertTrue(err < 1e-04)


@wp.kernel
def sample_cdf_kernel(kernel_seed: int, cdf: wp.array(dtype=float), samples: wp.array(dtype=int)):

    tid = wp.tid()
    state = wp.rand_init(kernel_seed, tid)

    samples[tid] = wp.sample_cdf(state, cdf)


def test_sample_cdf(test, device):

    seed = 42
    cdf = np.arange(0.0, 1.0, 0.01, dtype=float)
    cdf = cdf * cdf
    cdf = wp.array(cdf, dtype=float, device=device)
    num_samples = 1000
    samples = wp.zeros(num_samples, dtype=int, device=device)

    wp.launch(kernel=sample_cdf_kernel, dim=num_samples, inputs=[seed, cdf, samples], device=device)

    # histogram should be linear
    # plt.hist(samples.numpy())
    # plt.show()


@wp.kernel
def sampling_kernel(
    kernel_seed: int,
    triangle_samples: wp.array(dtype=wp.vec2),
    square_samples: wp.array(dtype=wp.vec2),
    ring_samples: wp.array(dtype=wp.vec2),
    disk_samples: wp.array(dtype=wp.vec2),
    sphere_surface_samples: wp.array(dtype=wp.vec3),
    sphere_samples: wp.array(dtype=wp.vec3),
    hemisphere_surface_samples: wp.array(dtype=wp.vec3),
    hemisphere_samples: wp.array(dtype=wp.vec3),
    cube_samples: wp.array(dtype=wp.vec3)):

    tid = wp.tid()
    state = wp.rand_init(kernel_seed, tid)

    triangle_samples[tid] = wp.sample_triangle(state)
    ring_samples[tid] = wp.sample_unit_ring(state)
    disk_samples[tid] = wp.sample_unit_disk(state)
    sphere_surface_samples[tid] = wp.sample_unit_sphere_surface(state)
    sphere_samples[tid] = wp.sample_unit_sphere(state)
    hemisphere_surface_samples[tid] = wp.sample_unit_hemisphere_surface(state)
    hemisphere_samples[tid] = wp.sample_unit_hemisphere(state)
    square_samples[tid] = wp.sample_unit_square(state)
    cube_samples[tid] = wp.sample_unit_cube(state)


def test_sampling_methods(test, device):
    
    seed = 42
    num_samples = 100

    triangle_samples = wp.zeros(num_samples, dtype=wp.vec2, device=device)
    square_samples = wp.zeros(num_samples, dtype=wp.vec2, device=device)
    ring_samples = wp.zeros(num_samples, dtype=wp.vec2, device=device)
    disk_samples = wp.zeros(num_samples, dtype=wp.vec2, device=device)
    sphere_surface_samples = wp.zeros(num_samples, dtype=wp.vec3, device=device)
    sphere_samples = wp.zeros(num_samples, dtype=wp.vec3, device=device)
    hemisphere_surface_samples = wp.zeros(num_samples, dtype=wp.vec3, device=device)
    hemisphere_samples = wp.zeros(num_samples, dtype=wp.vec3, device=device)
    cube_samples = wp.zeros(num_samples, dtype=wp.vec3, device=device)

    wp.launch(
        kernel=sampling_kernel,
        dim=num_samples,
        inputs=[seed, triangle_samples, square_samples, ring_samples, disk_samples, sphere_surface_samples, sphere_samples, hemisphere_surface_samples, hemisphere_samples, cube_samples],
        device=device)

    # bounds check
    test.assertTrue((triangle_samples.numpy()[:,0] <= 1.0).all())
    test.assertTrue((triangle_samples.numpy()[:,0] >= 0.0).all())
    test.assertTrue((triangle_samples.numpy()[:,1] >= 0.0).all())
    test.assertTrue((triangle_samples.numpy()[:,1] >= 0.0).all())
    test.assertTrue((square_samples.numpy()[:,0] >= -0.5).all())
    test.assertTrue((square_samples.numpy()[:,0] <= 1.5).all())
    test.assertTrue((square_samples.numpy()[:,1] >= -0.5).all())
    test.assertTrue((square_samples.numpy()[:,1] <= 0.5).all())
    test.assertTrue((cube_samples.numpy()[:,0] >= -0.5).all())
    test.assertTrue((cube_samples.numpy()[:,0] <= 0.5).all())
    test.assertTrue((cube_samples.numpy()[:,1] >= -0.5).all())
    test.assertTrue((cube_samples.numpy()[:,1] <= 0.5).all())
    test.assertTrue((cube_samples.numpy()[:,2] >= -0.5).all())
    test.assertTrue((cube_samples.numpy()[:,2] <= 0.5).all())
    test.assertTrue((hemisphere_surface_samples.numpy()[:,2] >= 0.0).all())
    test.assertTrue((hemisphere_samples.numpy()[:,2] >= 0.0).all())
    test.assertTrue((np.linalg.norm(ring_samples.numpy(), axis=1) <= 1.0 + 1e6).all())
    test.assertTrue((np.linalg.norm(disk_samples.numpy(), axis=1) <= 1.0 + 1e6).all())
    test.assertTrue((np.linalg.norm(sphere_surface_samples.numpy(), axis=1) <= 1.0 + 1e6).all())
    test.assertTrue((np.linalg.norm(sphere_samples.numpy(), axis=1) <= 1.0 + 1e6).all())
    test.assertTrue((np.linalg.norm(hemisphere_surface_samples.numpy(), axis=1) <= 1.0 + 1e6).all())
    test.assertTrue((np.linalg.norm(hemisphere_samples.numpy(), axis=1) <= 1.0 + 1e6).all())


def register(parent):

    devices = wp.get_devices()

    class TestNoise(parent):
        pass

    add_function_test(TestNoise, "test_rand", test_rand, devices=devices)
    add_function_test(TestNoise, "test_sample_cdf", test_sample_cdf, devices=devices)
    add_function_test(TestNoise, "test_sampling_methods", test_sampling_methods, devices=devices)

    return TestNoise

if __name__ == '__main__':
    c = register(unittest.TestCase)
    unittest.main(verbosity=2)