# Prototype - copied from pluto_RayTracingWeekend.jl
# Adapted from [Ray Tracing In One Weekend by Peter Shirley](https://raytracing.github.io/books/RayTracingInOneWeekend.html) and
# [cshenton's Julia implementation](https://github.com/cshenton/RayTracing.jl)"
using Pkg
Pkg.activate(@__DIR__)

"""Next ideas:

- Incorporate ideas from https://discourse.julialang.org/t/why-is-this-raytracer-slow/59176/20
  - @inline to many functions
- save image, e.g. PNG
- The only remaining allocations that appear expensive are of `HitRecord`s (on the stack)
- pre-allocate X paths (ray bundles? not sure what the terminology is) then
  implement versions of hit, scatter, etc. that operate on an entire tensor at once.
  (i.e. efficiently parallelizable with multithreading, on SIMD or GPU)
  - use FieldVector?
- figure out the incorrect look in refraction of negatively scaled sphere
- use parameterizable structs for Vec but keep the same efficiency
- share with community, ask for feedback
- continue watching MIT course
"""

using BenchmarkTools, Images, InteractiveUtils, LinearAlgebra, StaticArrays
Option{T} = Union{Missing, T}
t_col = @SVector[0.4f0, 0.5f0, 0.1f0] # test color

squared_length(v::SVector{3,Float32}) = v ⋅ v

# Before optimization:
#   677-699 ns (41 allocations: 2.77 KiB) # squared_length(v::SVector) = v ⋅ v; @btime squared_length(t_col)
# inside a function:
#   894-906 ns (43 allocations: 2.83 KiB)
#
# don't define `Base.getproperty(vec::SVector{3}, sym::Symbol)`
#   208.527 ns (3 allocations: 80 bytes)
#
# just test `@btime squared_length(t_col)`:
#    12.115 ns (1 allocation: 16 bytes)
# test again after redefining: `Vec3 = SVector{3, Float64}` instead of `Vec3 = SVector{3}`
#    11.966 ns (1 allocation: 16 bytes)
# test again after redefining: `Vec3 = SVector{3, Float32}` instead of `Vec3 = SVector{3, Float64}`
#    13.652 ns (1 allocation: 16 bytes)
#    (SLOWER! I wonder if this would hold with larger vectors or if we become mem-bound)
# replace: squared_length(v::SVector{3,MyFloat})
#     3.256 ns (0 allocations: 0 bytes)
# replace: @btime squared_length($t_col) # notice the $
#     1.152 ns (0 allocations: 0 bytes)
#@btime squared_length($t_col)

@inline near_zero(v::SVector{3,Float32}) = squared_length(v) < 1e-5
#@btime near_zero($t_col) # 1.382 ns (0 allocations: 0 bytes)

# Test images

img = rand(4, 3)

img_rgb = rand(RGB{Float32}, 4, 4)

ex1 = [RGB{Float32}(1,0,0) RGB{Float32}(0,1,0) RGB{Float32}(0,0,1);
       RGB{Float32}(1,1,0) RGB{Float32}(1,1,1) RGB{Float32}(0,0,0)]

ex2 = zeros(RGB{Float32}, 2, 3)
ex2[1,1] = RGB{Float32}(1,0,0)
ex2

# The "Hello World" of graphics
function gradient(nx::Int, ny::Int)
	img = zeros(RGB{Float32}, ny, nx)
	for i in 1:ny, j in 1:nx # Julia is column-major, i.e. iterate 1 column at a time
		x = j; y = (ny-i);
		r = x/nx
		g = y/ny
		b = 0.2
		img[i,j] = RGB{Float32}(r,g,b)
	end
	img
end

gradient(200,100)

rgb(v::SVector{3,Float32}) = RGB{Float32}(v[1], v[2], v[3])
#rgb(v::Vec3) = RGB{MyFloat}(v...)
#@btime rgb($t_col) # 289.914 ns (4 allocations: 80 bytes)

@inline rgb_gamma2(v::SVector{3,Float32}) = RGB{Float32}(sqrt.(v)...)

#@btime rgb_gamma2($t_col) # 454.766 ns (11 allocations: 256 bytes)

struct Ray
	origin::SVector{3, Float32} # Point
	dir::SVector{3, Float32} # Vec3 # direction (unit vector)
end

# interpolates between blue and white
_origin = @SVector[0.0f0,0.0f0,0.0f0]
_v3_minusY = @SVector[0.0f0,-1.0f0,0.0f0]
_t_ray1 = Ray(_origin, _v3_minusY)
#@btime Ray($_origin, $_v3_minusY) # 6.341 ns (0 allocations: 0 bytes)

# equivalent to C++'s ray.at()
@inline point(r::Ray, t::Float32) = r.origin .+ t .* r.dir
#@btime point($_t_ray1, 0.5f0) # 1.392 ns (0 allocations: 0 bytes)

#md"# Chapter 4: Rays, simple camera, and background"

@inline function skycolor(ray::Ray)
	white = @SVector[1.0f0,1.0f0,1.0f0]
	skyblue = @SVector[0.5f0,0.7f0,1.0f0]
	t = 0.5f0(ray.dir[2] + 1.0f0)
    (1.0f0-t)*white + t*skyblue
end
#@btime skycolor($_t_ray1) # 1.402 ns (0 allocations: 0 bytes)

# before optim:
#   1.474 μs (23 allocations: 608 bytes)
# pre-allocate args:
#   1.150 μs (19 allocations: 480 bytes)
# `rgb(skycolor(t_ray1))`
#   1.143 μs (19 allocations: 480 bytes)
# `skycolor(t_ray1)`:
#       636.259 ns (9 allocations: 256 bytes)
#    extract t_white and t_skyblue from skycolor():
#       279.779 ns (5 allocations: 128 bytes)
#    pre-allocate the ray given in argument:
#        51.130 ns (5 allocations: 128 bytes)
#    skycolor!($t_col2, $t_ray1, $t_white, $t_skyblue) using @set to return the value:
#        14.727 ns (1 allocation: 32 bytes)
#    replace all the Color() by @SVector[...], don't use global colors (use local)
#         1.402 ns (0 allocations: 0 bytes)
# restore the rgb(), i.e. `@btime rgb(skycolor($t_ray1))`
#   291.492 ns (4 allocations: 80 bytes)
#@btime rgb(skycolor($t_ray1)) # 291.492 ns (4 allocations: 80 bytes)

# md"""# Random vectors
# C++'s section 8.1"""

@inline random_between(min=0.0f0, max=1.0f0) = rand(Float32)*(max-min) + min # equiv to random_double()
#@btime random_between(50f0, 100f0) # 4.519 ns (0 allocations: 0 bytes)

@inline random_vec3(min::Float32, max::Float32) = @SVector[random_between(min, max) for i ∈ 1:3]

# Before optimization:
#   352.322 ns (6 allocations: 224 bytes)
# Don't use list comprehension:
#   179.684 ns (5 allocations: 112 bytes)
# Return @SVector[] instead of Vec3():
#    12.557 ns (0 allocations: 0 bytes)
#@btime random_vec3(-1.0f0,1.0f0)

@inline random_vec2(min::Float32, max::Float32) = @SVector[random_between(min, max) for i ∈ 1:2]
#@btime random_vec2(-1.0f0,1.0f0) # 8.536 ns (0 allocations: 0 bytes)

@inline function random_vec3_in_sphere() # equiv to random_in_unit_sphere()
	while true
		p = random_vec3(-1f0, 1f0)
		if p⋅p <= 1
			return p
		end
	end
end
#@btime random_vec3_in_sphere() # 46.587 ns (0 allocations: 0 bytes)

squared_length(random_vec3_in_sphere())

"Random unit vector. Equivalent to C++'s `unit_vector(random_in_unit_sphere())`"
@inline random_vec3_on_sphere() = normalize(random_vec3_in_sphere())
#@btime random_vec3_on_sphere()

# Before optim:
#   517.538 ns (12 allocations: 418 bytes)... but random
# After reusing _rand_vec3:
#    92.865 ns (2 allocations: 60 bytes)
# Various speed-ups (don't use Vec3, etc.):
#    52.427 ns (0 allocations: 0 bytes)
random_vec3_on_sphere()
norm(random_vec3_on_sphere())

@inline function random_vec2_in_disk() # equiv to random_in_unit_disk()
	while true
		p = random_vec2(-1f0, 1f0)
		if p⋅p <= 1
			return p
		end
	end
end
#@btime random_vec2_in_disk() # 16.725 ns (0 allocations: 0 bytes)

""" Temporary function to shoot rays through each pixel. Later replaced by `render`

	Args:
		scene: a function that takes a ray, returns the color of any object it hit
"""
function main(nx::Int, ny::Int, scene)
	lower_left_corner = @SVector[-2,-1,-1]
	horizontal = @SVector[4,0,0]
	vertical = @SVector[0f0,2f0,0f0]
	origin = @SVector[0,0,0]

	img = zeros(RGB{Float32}, ny, nx)
	for i in 1:ny, j in 1:nx # Julia is column-major, i.e. iterate 1 column at a time
		u = j/nx
		v = (ny-i)/ny # Y-up!
		ray = Ray(origin, normalize(lower_left_corner + u*horizontal + v*vertical))
		#r = x/nx
		#g = y/ny
		#b = 0.2
		img[i,j] = rgb(scene(ray))
	end
	img
end

# Before optimizations:
#   28.630 ms (520010 allocations: 13.05 MiB)
# After optimizations (skycolor, rand, etc.)
#   10.358 ms (220010 allocations: 5.42 MiB) # at this stage, rgb() is most likely the culprit..
# After replacing Color,Vec3 by @SVector:
#    6.539 ms (80002 allocations: 1.75 MiB)
# Lots more optimizations, including @inline of low-level functions:
# 161.546 μs (     2 allocations: 234.45 KiB)
main(200,100, skycolor)

#md"# Chapter 5: Add a sphere"

@inline function hit_sphere1(center::SVector{3,Float32}, radius::Float32, r::Ray)
	oc = r.origin - center
	a = r.dir ⋅ r.dir
	b = 2oc ⋅ r.dir
	c = (oc ⋅ oc) - radius*radius
	discriminant = b*b - 4a*c
	discriminant > 0
end

@inline function sphere_scene1(r::Ray)
	if hit_sphere1(@SVector[0f0,0f0,-1f0], 0.5f0, r) # sphere of radius 0.5 centered at z=-1
		return @SVector[1f0,0f0,0f0] # red
	else
		skycolor(r)
	end
end

main(200,100,sphere_scene1) # 227.851 μs (2 allocations: 234.45 KiB)

#md"# Chapter 6: Surface normals and multiple objects"

@inline function hit_sphere2(center::SVector{3,Float32}, radius::Float32, r::Ray)
	oc = r.origin - center
	a = r.dir ⋅ r.dir
	b = 2oc ⋅ r.dir
	c = (oc ⋅ oc) - radius*radius
	discriminant = b*b - 4a*c
	if discriminant < 0
		return -1
	else
		return (-b - sqrt(discriminant)) / 2a
	end
end

@inline function sphere_scene2(r::Ray)
	sphere_center = SVector{3,Float32}(0f0,0f0,-1f0)
	t = hit_sphere2(sphere_center, 0.5f0, r) # sphere of radius 0.5 centered at z=-1
	if t > 0f0
		n⃗ = normalize(point(r, t) - sphere_center) # normal vector. typed n\vec
		return 0.5f0n⃗ + SVector{3,Float32}(0.5f0,0.5f0,0.5f0) # remap normal to rgb
	else
		skycolor(r)
	end
end

@btime main(200,100,sphere_scene2)  # 262.286 μs (2 allocations: 234.45 KiB)

"An object that can be hit by Ray"
abstract type Hittable end

"""Materials tell us how rays interact with a surface"""
abstract type Material end

struct NoMaterial <: Material end

const _no_material = NoMaterial()
const _y_up = @SVector[0f0,1f0,0f0]

"Record a hit between a ray and an object's surface"
mutable struct HitRecord
	t::Float32 # vector from the ray's origin to the intersection with a surface.

	# If t==Inf32, there was no hit, and all following values are undefined!
	#
	p::SVector{3,Float32} # point of the intersection between an object's surface and a ray
	n⃗::SVector{3,Float32} # surface's outward normal vector, points towards outside of object?

	# If true, our ray hit from outside to the front of the surface.
	# If false, the ray hit from within.
	front_face::Bool
	mat::Material

	@inline HitRecord() = new(Inf32) # no hit!
	@inline HitRecord(t,p,n⃗,front_face,mat) = new(t,p,n⃗,front_face,mat)
end

struct Sphere <: Hittable
	center::SVector{3,Float32}
	radius::Float32
	mat::Material
end

# md"""The geometry defines an `outside normal`. A HitRecord stores the `local normal`.
# ![Surface normal](https://raytracing.github.io/images/fig-1.06-normal-sides.jpg)
# """

"""Equivalent to `hit_record.set_face_normal()`"""
@inline function ray_to_HitRecord(t, p, outward_n⃗, r_dir::SVector{3,Float32}, mat::Material)
	front_face = r_dir ⋅ outward_n⃗ < 0
	n⃗ = front_face ? outward_n⃗ : -outward_n⃗
	rec = HitRecord(t,p,n⃗,front_face,mat)
end

struct Scatter
	r::Ray
	attenuation::SVector{3,Float32}

	# claforte: TODO: rename to "absorbed?", i.e. not reflected/refracted?
	reflected::Bool # whether the scattered ray was reflected, or fully absorbed
	@inline Scatter(r,a,reflected=true) = new(r,a,reflected)
end

#"Diffuse material"
mutable struct Lambertian<:Material
	albedo::SVector{3,Float32}
end

"""Compute reflection vector for v (pointing to surface) and normal n⃗.

	See [diagram](https://raytracing.github.io/books/RayTracingInOneWeekend.html#metal/mirroredlightreflection)"""
#reflect(v::SVector{3,Float32}, n⃗::SVector{3,Float32}) = normalize(v - (2v⋅n⃗)*n⃗) # claforte: normalize not needed?
@inline reflect(v::SVector{3,Float32}, n⃗::SVector{3,Float32}) = v - (2v⋅n⃗)*n⃗

@assert reflect(@SVector[0.6f0,-0.8f0,0f0], @SVector[0f0,1f0,0f0]) == @SVector[0.6f0,0.8f0,0f00] # diagram's example

"""Create a scattered ray emitted by `mat` from incident Ray `r`.

	Args:
		rec: the HitRecord of the surface from which to scatter the ray.

	Return missing if it's fully absorbed. """
@inline function scatter(mat::Lambertian, r::Ray, rec::HitRecord)::Scatter
	scatter_dir = rec.n⃗ + random_vec3_on_sphere()
	if near_zero(scatter_dir) # Catch degenerate scatter direction
		scatter_dir = rec.n⃗
	else
		scatter_dir = normalize(scatter_dir)
	end
	scattered_r = Ray(rec.p, scatter_dir)
	attenuation = mat.albedo
	return Scatter(scattered_r, attenuation)
end

const _no_hit = HitRecord()

@inline function hit(s::Sphere, r::Ray, tmin::Float32, tmax::Float32)#::Option{HitRecord} # the return type seems to cause excessive allocations
    oc = r.origin - s.center
    a = 1 #r.dir ⋅ r.dir # normalized vector - always 1
    half_b = oc ⋅ r.dir
    c = oc⋅oc - s.radius^2
    discriminant = half_b^2 - a*c
	if discriminant < 0 return _no_hit end
	sqrtd = √discriminant

	# Find the nearest root that lies in the acceptable range
	root = (-half_b - sqrtd) / a
	if root < tmin || tmax < root
		root = (-half_b + sqrtd) / a
		if root < tmin || tmax < root
			return _no_hit
		end
	end

	t = root
	p = point(r, t)
	n⃗ = (p - s.center) / s.radius
	return ray_to_HitRecord(t, p, n⃗, r.dir, s.mat)
end

struct HittableList <: Hittable
    list::Vector{Hittable}
end

#"""Find closest hit between `Ray r` and a list of Hittable objects `h`, within distance `tmin` < `tmax`"""
@inline function hit(hittables::HittableList, r::Ray, tmin::Float32, tmax::Float32)# ::HitRecord
    closest = tmax # closest t so far
    rec = _no_hit
    for h in hittables.list
        temprec = hit(h, r, tmin, closest)
        if temprec !== _no_hit
            rec = temprec
            closest = rec.t # i.e. ignore any further hit > this one's.
        end
    end
    rec
end

@inline color_vec3_in_rgb(v::SVector{3,Float32}) = 0.5normalize(v) + @SVector[0.5f,0.5f,0.5f]

#md"# Metal material"

mutable struct Metal<:Material
	albedo::SVector{3,Float32}
	fuzz::Float32 # how big the sphere used to generate fuzzy reflection rays. 0=none
	@inline Metal(a,f=0.0) = new(a,f)
end

@inline function scatter(mat::Metal, r_in::Ray, rec::HitRecord)::Scatter
	reflected = normalize(reflect(r_in.dir, rec.n⃗) + mat.fuzz*random_vec3_on_sphere())
	Scatter(Ray(rec.p, reflected), mat.albedo)
end

#md"# Scenes"

#"Scene with 2 Lambertian spheres"
function scene_2_spheres()::HittableList
	spheres = Sphere[]

	# small center sphere
	push!(spheres, Sphere(@SVector[0f0,0f0,-1f0], 0.5, Lambertian(@SVector[0.7f0,0.3f0,0.3f0])))

	# ground sphere (planet?)
	push!(spheres, Sphere(@SVector[0f0,-100.5f0,-1f0], 100, Lambertian(@SVector[0.8f0,0.8f0,0.0f0])))
	HittableList(spheres)
end

#"""Scene with 2 Lambertian, 2 Metal spheres.
#
#	See https://raytracing.github.io/images/img-1.11-metal-shiny.png"""
function scene_4_spheres()::HittableList
	scene = scene_2_spheres()

	# left and right Metal spheres
	push!(scene.list, Sphere(@SVector[-1f0,0f0,-1f0], 0.5f0, Metal(@SVector[0.8f0,0.8f0,0.8f0], 0.3f0)))
	push!(scene.list, Sphere(@SVector[1f0,0f0,-1f0], 0.5f0, Metal(@SVector[0.8f0,0.6f0,0.2f0], 0.8f0)))
	return scene
end

#md"""# Camera

# Adapted from C++'s sections 7.2, 11.1 """
mutable struct Camera
	origin::SVector{3,Float32}
	lower_left_corner::SVector{3,Float32}
	horizontal::SVector{3,Float32}
	vertical::SVector{3,Float32}
	u::SVector{3,Float32}
	v::SVector{3,Float32}
	w::SVector{3,Float32}
	lens_radius::Float32
end

"""
	Args:
		vfov: vertical field-of-view in degrees
		aspect_ratio: horizontal/vertical ratio of pixels
      aperture: if 0 - no depth-of-field
"""
function default_camera(lookfrom::SVector{3,Float32}=@SVector[0f0,0f0,0f0], lookat::SVector{3,Float32}=@SVector[0f0,0f0,-1f0],
						vup::SVector{3,Float32}=@SVector[0f0,1f0,0f0], vfov=90.0f0, aspect_ratio=16.0f0/9.0f0,
						aperture=0.0f0, focus_dist=1.0f0)
	viewport_height = 2.0f0 * tand(vfov/2f0)
	viewport_width = aspect_ratio * viewport_height

	w = normalize(lookfrom - lookat)
	u = normalize(vup × w)
	v = w × u

	origin = lookfrom
	horizontal = focus_dist * viewport_width * u
	vertical = focus_dist * viewport_height * v
	lower_left_corner = origin - horizontal/2f0 - vertical/2f0 - focus_dist*w
	lens_radius = aperture/2f0
	Camera(origin, lower_left_corner, horizontal, vertical, u, v, w, lens_radius)
end

default_camera()

#md"# Render

@inline function get_ray(c::Camera, s::Float32, t::Float32)
	rd = SVector{2,Float32}(c.lens_radius * random_vec2_in_disk())
	offset = c.u * rd[1] + c.v * rd[2] #offset = c.u * rd.x + c.v * rd.y
    Ray(c.origin + offset, normalize(c.lower_left_corner + s*c.horizontal +
									 t*c.vertical - c.origin - offset))
end

#@btime get_ray(default_camera(), 0.0f0, 0.0f0)

"""Compute color for a ray, recursively

	Args:
		depth: how many more levels of recursive ray bounces can we still compute?"""
@inline function ray_color(r::Ray, world::HittableList, depth=4)
    if depth <= 0
		return @SVector[0f0,0f0,0f0]
	end

	rec = hit(world, r, 1f-4, Inf32)
    if rec !== _no_hit
		# For debugging, represent vectors as RGB:
		# return color_vec3_in_rgb(rec.p) # show the normalized hit point
		# return color_vec3_in_rgb(rec.n⃗) # show the normal in RGB
		# return color_vec3_in_rgb(rec.p + rec.n⃗)
		# return color_vec3_in_rgb(random_vec3_in_sphere())
		#return color_vec3_in_rgb(rec.n⃗ + random_vec3_in_sphere())

        s = scatter(rec.mat, r, rec)
		if s.reflected
			return s.attenuation .* ray_color(s.r, world, depth-1)
		else
			return @SVector[0f0,0f0,0f0]
		end
    else
        skycolor(r)
    end
end

"""Render an image of `scene` using the specified camera, number of samples.

	Args:
		scene: a HittableList, e.g. a list of spheres
		n_samples: number of samples per pixel, eq. to C++ samples_per_pixel

	Equivalent to C++'s `main` function."""
function render(scene::HittableList, cam::Camera, image_width=400,
				n_samples=1)
	# Image
	aspect_ratio = 16.0f0/9.0f0 # TODO: use cam.aspect_ratio for consistency
	image_height = convert(Int64, floor(image_width / aspect_ratio))

	# Render
	img = zeros(RGB{Float32}, image_height, image_width)
	# Compared to C++, Julia is:
	# 1. column-major, i.e. iterate 1 column at a time, so invert i,j compared to C++
	# 2. 1-based, so no need to subtract 1 from image_width, etc.
	# 3. The array is Y-down, but `v` is Y-up
	for i in 1:image_height, j in 1:image_width
		accum_color = @SVector[0f0,0f0,0f0]
		for s in 1:n_samples
			u = convert(Float32, j/image_width)
			v = convert(Float32, (image_height-i)/image_height) # i is Y-down, v is Y-up!
			if s != 1 # 1st sample is always centered, for 1-sample/pixel
				# claforte: I think the C++ version had a bug, the rand offset was
				# between [0,1] instead of centered at 0, e.g. [-0.5, 0.5].
				u += convert(Float32, (rand()-0.5f0) / image_width)
				v += convert(Float32 ,(rand()-0.5f0) / image_height)
			end
			ray = get_ray(cam, u, v)
			accum_color += ray_color(ray, scene)
		end
		img[i,j] = rgb_gamma2(accum_color / n_samples)
	end
	img
end

# After some optimization:
#  46.506 ms (917106 allocations: 16.33 MiB)
# Using convert(Float32, ...) instead of MyFloat(...):
#  14.137 ms (118583 allocations: 4.14 MiB)
# Don't specify return value of Option{HitRecord} in hit()
#  24.842 ms (527738 allocations: 16.63 MiB)
# Don't initialize unnecessary elements in HitRecord():
#  14.862 ms (118745 allocations: 4.15 MiB)  (but computer was busy...)
# Replace MyFloat by Float32:
#  11.792 ms ( 61551 allocations: 2.88 MiB)
# Remove ::HitRecord return value in remaining hit() method:
#  11.545 ms ( 61654 allocations: 2.88 MiB)
# with mutable HitRecord
#  11.183 ms ( 61678 allocations: 2.88 MiB) (insignificant?)
# @inline tons of stuff:
#   8.129 ms ( 61660 allocations: 2.88 MiB)
render(scene_2_spheres(), default_camera(), 96, 16)

render(scene_4_spheres(), default_camera(), 96, 16)

#md"""# Dielectrics

# from Section 10.2 Snell's Law:
# ![Ray refraction](https://raytracing.github.io/images/fig-1.13-refraction.jpg)

# Refracted angle `sinθ′ = (η/η′)⋅sinθ`, where η (\eta) are the refractive indices.

# Split the parts of the ray into `R′=R′⊥+R′∥` (perpendicular and parallel to n⃗′)."""

# """
# 	Args:
# 		refraction_ratio: incident refraction index divided by refraction index of
# 			hit surface. i.e. η/η′ in the figure above"""
@inline function refract(dir::SVector{3,Float32}, n⃗::SVector{3,Float32}, refraction_ratio::Float32)
	cosθ = min(-dir ⋅ n⃗, 1)
	r_out_perp = refraction_ratio * (dir + cosθ*n⃗)
	r_out_parallel = -√(abs(1-squared_length(r_out_perp))) * n⃗
	normalize(r_out_perp + r_out_parallel)
end

@assert refract(SVector{3,Float32}(0.6,-0.8,0), SVector{3,Float32}(0,1,0), 1.0f0) == SVector{3,Float32}(0.6,-0.8,0) # unchanged

t_refract_widerθ = refract(SVector{3,Float32}(0.6,-0.8,0), SVector{3,Float32}(0,1,0), 2.0f0) # wider angle
@assert isapprox(t_refract_widerθ, SVector{3,Float32}(0.87519, -0.483779, 0.0); atol=1e-3)

t_refract_narrowerθ = refract(SVector{3,Float32}(0.6,-0.8,0), SVector{3,Float32}(0,1,0), 0.5f0) # narrower angle
@assert isapprox(t_refract_narrowerθ, SVector{3,Float32}(0.3, -0.953939, 0.0); atol=1e-3)

mutable struct Dielectric <: Material
	ir::Float32 # index of refraction, i.e. η.
end

@inline function reflectance(cosθ::Float32, refraction_ratio::Float32)
	# Use Schlick's approximation for reflectance.
	# claforte: may be buggy? I'm getting black pixels in the Hollow Glass Sphere...
	r0 = (1f0-refraction_ratio) / (1f0+refraction_ratio)
	r0 = r0^2
	r0 + (1f0-r0)*((1f0-cosθ)^5)
end

@inline function scatter(mat::Dielectric, r_in::Ray, rec::HitRecord)
	attenuation = @SVector[1f0,1f0,1f0]
	refraction_ratio = rec.front_face ? (1.0f0/mat.ir) : mat.ir # i.e. ηᵢ/ηₜ
	cosθ = min(-r_in.dir⋅rec.n⃗, 1.0f0)
	sinθ = √(1.0f0 - cosθ^2)
	cannot_refract = refraction_ratio * sinθ > 1.0
	if cannot_refract || reflectance(cosθ, refraction_ratio) > rand()
		dir = reflect(r_in.dir, rec.n⃗)
	else
		dir = refract(r_in.dir, rec.n⃗, refraction_ratio)
	end
	Scatter(Ray(rec.p, dir), attenuation) # TODO: rename reflected -> !absorbed?
end

#"From C++: Image 15: Glass sphere that sometimes refracts"
@inline function scene_diel_spheres(left_radius=0.5f0)::HittableList # dielectric spheres
	spheres = Sphere[]

	# small center sphere
	push!(spheres, Sphere(SVector{3,Float32}(0f0,0f0,-1f0), 0.5f0, Lambertian(SVector{3,Float32}(0.1f0,0.2f0,0.5f0))))

	# ground sphere (planet?)
	push!(spheres, Sphere(SVector{3,Float32}(0f0,-100.5f0,-1f0), 100f0, Lambertian(SVector{3,Float32}(0.8f0,0.8f0,0.0f0))))

	# left and right spheres.
	# Use a negative radius on the left sphere to create a "thin bubble"
	push!(spheres, Sphere(SVector{3,Float32}(-1f0,0f0,-1f0), left_radius, Dielectric(1.5f0)))
	push!(spheres, Sphere(SVector{3,Float32}( 1f0,0f0,-1f0), 0.5f0, Metal(SVector{3,Float32}(0.8f0,0.6f0,0.2f0), 0.0f0)))
	HittableList(spheres)
end

render(scene_diel_spheres(), default_camera(), 96, 16)
#render(scene_diel_spheres(), default_camera(), 320, 32)

#md"# Positioning camera"

function scene_blue_red_spheres()::HittableList # dielectric spheres
	spheres = Sphere[]
	R = cos(pi/4)
	push!(spheres, Sphere(@SVector[-R,0f0,-1f0], R, Lambertian(@SVector[0f0,0f0,1f0])))
	push!(spheres, Sphere(@SVector[R,0f0,-1f0], R, Lambertian(@SVector[1f0,0f0,0f0])))
	HittableList(spheres)
end

#md"# Spheres with depth-of-field"

#md"# Random spheres"

function scene_random_spheres()::HittableList # dielectric spheres
	spheres = Sphere[]

	# ground
	push!(spheres, Sphere(@SVector[0f0,-1000f0,-1f0], 1000f0,
						  Lambertian(@SVector[0.5f0,0.5f0,0.5f0])))

	for a in -11:10, b in -11:10
		choose_mat = rand()
		center = @SVector[a + 0.9f0*rand(), 0.2f0, b + 0.90f0*rand()]

		# skip spheres too close?
		if norm(center - @SVector[4f0,0.2f0,0f0]) < 0.9f0 continue end

		if choose_mat < 0.8f0
			# diffuse
			albedo = @SVector[rand() for i ∈ 1:3] .* @SVector[rand() for i ∈ 1:3]
			push!(spheres, Sphere(center, 0.2f0, Lambertian(albedo)))
		elseif choose_mat < 0.95f0
			# metal
			albedo = @SVector[random_between(0.5f0,1.0f0) for i ∈ 1:3]
			fuzz = random_between(0.0f0, 5.0f0)
			push!(spheres, Sphere(center, 0.2f0, Metal(albedo, fuzz)))
		else
			# glass
			push!(spheres, Sphere(center, 0.2f0, Dielectric(1.5f0)))
		end
	end

	push!(spheres, Sphere(@SVector[0f0,1f0,0f0], 1.0f0, Dielectric(1.5f0)))
	push!(spheres, Sphere(@SVector[-4f0,1f0,0f0], 1.0f0,
						  Lambertian(@SVector[0.4f0,0.2f0,0.1f0])))
	push!(spheres, Sphere(@SVector[4f0,1f0,0f0], 1.0f0,
						  Metal(@SVector[0.7f0,0.6f0,0.5f0], 0.0f0)))
	HittableList(spheres)
end

scene_random_spheres()

# Hollow Glass sphere using a negative radius
# claforte: getting a weird black halo in the glass sphere... might be due to my
# "fix" for previous black spots, by moving the RecordHit point a bit away from
# the hit surface...
render(scene_diel_spheres(-0.5f0), default_camera(), 96, 16)

render(scene_blue_red_spheres(), default_camera(), 96, 16)

render(scene_diel_spheres(), default_camera(@SVector[-2f0,2f0,1f0], @SVector[0f0,0f0,-1f0],
							 				@SVector[0f0,1f0,0f0], 20.0f0), 96, 16)

t_lookfrom1 = @SVector[13.0f0,2.0f0,3.0f0]
t_lookat1 = @SVector[0.0f0,0.0f0,0.0f0]
t_cam1 = default_camera(t_lookfrom1, t_lookat1, @SVector[0.0f0,1.0f0,0.0f0], 20.0f0, 16.0f0/9.0f0,
                        0.1f0, 10.0f0)

# took ~20s (really rough timing) in REPL, before optimization
# after optimization:
#   880.997 ms (48801164 allocations: 847.10 MiB)
# after switching to Float32 + reducing allocations using rand_vec3!():
#    79.574 ms (  542467 allocations:  14.93 MiB)
# after optimizing skycolor, rand*, probably more stuff I forgot...
#    38.242 ms (  235752 allocations:   6.37 MiB)
# after removing all remaining Color, Vec3, replacing them with @SVector[]...
#    26.790 ms (   91486 allocations: 2.28 MiB)
# Using convert(Float32, ...) instead of MyFloat(...):
#    25.856 ms (   70650 allocations: 1.96 MiB)
# Don't specify return value of Option{HitRecord} in hit()
# Don't initialize unnecessary elements in HitRecord():
#    38.961 ms (   70681 allocations: 1.96 MiB) # WORSE, probably because we're writing a lot more to stack...?
# Replace MyFloat by Float32:
#    36.726 ms (13889 allocations: 724.09 KiB)
# @inline lots of stuff:
#    14.690 ms (13652 allocations: 712.98 KiB)
render(scene_random_spheres(), t_cam1, 96, 1)

# took 5020s in Pluto.jl, before optimizations!
# after lots of optimizations, up to switching to Float32 + reducing allocations using rand_vec3!():
#   10.862032 seconds (70.09 M allocations: 1.858 GiB, 1.53% gc time)
# after optimizing skycolor, rand*, probably more stuff I forgot...
#    4.926 s (25,694,163 allocations: 660.06 MiB)
# after removing all remaining Color, Vec3, replacing them with @SVector[]...
#    3.541 s (9074843 allocations: 195.22 MiB)
# Using convert(Float32, ...) instead of MyFloat(...):
#    3.222 s (2055106 allocations: 88.53 MiB)
# Don't specify return value of Option{HitRecord} in hit()
# Don't initialize unnecessary elements in HitRecord():
#    5.016 s (2056817 allocations: 88.61 MiB) #  WORSE, probably because we're writing a lot more to stack...?
# Replace MyFloat by Float32:
#    5.185 s (1819234 allocations: 83.55 MiB) # Increase of 1% is probably noise
# Remove normalize() in reflect() (that function assumes the inputs are normalized)
#    5.040 s (1832052 allocations: 84.13 MiB)
# @inline lots of stuff:
#    2.110 s (1823044 allocations: 83.72 MiB)
render(scene_random_spheres(), t_cam1, 200, 32)

# After some optimization, took ~5.6 hours:
#   20171.646846 seconds (94.73 G allocations: 2.496 TiB, 1.06% gc time)
# ... however the image looked weird... too blurry
# After removing all remaining Color, Vec3, replacing them with @SVector[]... took ~3.7 hours:
#   13326.770907 seconds (29.82 G allocations: 714.941 GiB, 0.36% gc time)
# Using convert(Float32, ...) instead of MyFloat(...):
# Don't specify return value of Option{HitRecord} in hit()
# Don't initialize unnecessary elements in HitRecord():
# Took ~4.1 hours:
#   14723.339976 seconds (5.45 G allocations: 243.044 GiB, 0.11% gc time) # WORSE, probably because we're writing a lot more to stack...?
# Replace MyFloat by Float32:
# Lots of other optimizations including @inline lots of stuff:
#    6018.101653 seconds (5.39 G allocations: 241.139 GiB, 0.21% gc time)
#render(scene_random_spheres(), t_cam1, 1920, 1000)

t_lookfrom2 = @SVector[3.0f0,3.0f0,2.0f0]
t_lookat2 = @SVector[0.0f0,0.0f0,-1.0f0]
dist_to_focus2 = norm(t_lookfrom2-t_lookat2)
t_cam2 = default_camera(t_lookfrom2, t_lookat2, @SVector[0.0f0,1.0f0,0.0f0], 20.0f0, 16.0f0/9.0f0,
                        2.0f0, dist_to_focus2)

# Before optimization:
#  5.993 s  (193097930 allocations: 11.92 GiB)
# after disabling: `Base.getproperty(vec::SVector{3}, sym::Symbol)`
#  1.001 s  ( 17406437 allocations: 425.87 MiB)
# after forcing Ray and point() to use Float64 instead of AbstractFloat:
#  397.905 ms (6269207 allocations: 201.30 MiB)
# after forcing use of Float32 instead of Float64:
#  487.680 ms (7128113 allocations: 196.89 MiB) # More allocations... something is causing them...
# after optimizing rand_vec3!/rand_vec2! to minimize allocations:
#  423.468 ms (6075725 allocations: 158.92 MiB)
# after optimizing skycolor, rand*, probably more stuff I forgot...
#  217.853 ms (2942272 allocations: 74.82 MiB)
# after removing all remaining Color, Vec3, replacing them with @SVector[]...
#   56.778 ms (1009344 allocations: 20.56 MiB)
# Using convert(Float32, ...) instead of MyFloat(...):
#   23.870 ms (210890 allocations: 8.37 MiB)
# Replace MyFloat by Float32:
#   22.390 ms (153918 allocations: 7.11 MiB)
# Various other changes, e.g. remove unnecessary normalize
#   20.241 ms (153792 allocations: 7.10 MiB)
# @inline lots of stuff:
#   18.065 ms (153849 allocations: 7.10 MiB)
#render(scene_diel_spheres(), t_cam2, 96, 16)

# using Profile
# Profile.clear_malloc_data()
# render(scene_2_spheres(), default_camera(), 96, 16)


# This file is based on LinearAlgebra/test/lu.jl, a part of Julia. License is MIT: https://julialang.org/license

using Test, LinearAlgebra, Random, MatrixFactorizations
using LinearAlgebra: ldiv!, BlasReal, BlasInt, BlasFloat, rdiv!

@testset "UL" begin
    n = 10

    # Split n into 2 parts for tests needing two matrices
    n1 = div(n, 2)
    n2 = 2*n1

    Random.seed!(1234321)

    areal = randn(n,n)/2
    aimg  = randn(n,n)/2
    bimg  = randn(n,2)/2
    breal = randn(n,2)/2
    creal = randn(n)/2
    cimg  = randn(n)/2
    dureal = randn(n-1)/2
    duimg  = randn(n-1)/2
    dlreal = randn(n-1)/2
    dlimg  = randn(n-1)/2
    dreal = randn(n)/2
    dimg  = randn(n)/2

    @testset "UL simple test" begin
        A = randn(n,n)
        U,L = ul(A,Val(false))
        @test U*L ≈ A
        U,L,p = ul(A)
        @test (U*L) ≈ A[p,:]
        P = ul(A).P
        @test P'*U*L ≈ A
        L,U,p = lu(A)
        @test (L*U) ≈ A[p,:] # for context
        P = lu(A).P
        @test P'*L*U ≈ A
    end

    @testset for eltya in (Float32, Float64, ComplexF32, ComplexF64, BigFloat, Int)
        a = eltya == Int ? rand(1:7, n, n) :
            convert(Matrix{eltya}, eltya <: Complex ? complex.(areal, aimg) : areal)
        d = if eltya == Int
            Tridiagonal(rand(1:7, n-1), rand(1:7, n), rand(1:7, n-1))
        elseif eltya <: Complex
            convert(Tridiagonal{eltya}, Tridiagonal(
                complex.(dlreal, dlimg), complex.(dreal, dimg), complex.(dureal, duimg)))
        else
            convert(Tridiagonal{eltya}, Tridiagonal(dlreal, dreal, dureal))
        end
        ε = εa = eps(abs(float(one(eltya))))

        if eltya <: BlasFloat
            @testset "UL factorization for Number" begin
                num = rand(eltya)
                @test (ul(num)...,) == (hcat(one(eltya)), hcat(num), [1])
                @test convert(Array, ul(num)) ≈ eltya[num]
            end
        end
        κ  = cond(a,1)
        κd    = cond(Array(d),1)
        # @testset "Tridiagonal UL" begin
        #     lud   = ul(d)
        #     @test LinearAlgebra.issuccess(lud)
        #     @test ul(lud) == lud
        #     @test_throws ErrorException lud.Z
        #     @test lud.L*lud.U ≈ lud.P*Array(d)
        #     @test lud.L*lud.U ≈ Array(d)[lud.p,:]
        #     @test AbstractArray(lud) ≈ d
        #     @test Array(lud) ≈ d
        #     if eltya != Int
        #         dlu = convert.(eltya, [1, 1])
        #         dia = convert.(eltya, [-2, -2, -2])
        #         tri = Tridiagonal(dlu, dia, dlu)
        #         @test_throws ArgumentError ul!(tri)
        #     end
        # end
        @testset for eltyb in (Float32, Float64, ComplexF32, ComplexF64, Int)
            b  = eltyb == Int ? rand(1:5, n, 2) :
                convert(Matrix{eltyb}, eltyb <: Complex ? complex.(breal, bimg) : breal)
            c  = eltyb == Int ? rand(1:5, n) :
                convert(Vector{eltyb}, eltyb <: Complex ? complex.(creal, cimg) : creal)
            εb = eps(abs(float(one(eltyb))))
            ε  = max(εa,εb)
            @testset "(Automatic) Square UL decomposition" begin
                lua   = factorize(a)
                let Bs = copy(b), Cs = copy(c)
                    for (bb, cc) in ((Bs, Cs), (view(Bs, 1:n, 1), view(Cs, 1:n)))
                        @test norm(a*(lua\bb) - bb, 1) < ε*κ*n*2 # Two because the right hand side has two columns
                        @test norm(a'*(lua'\bb) - bb, 1) < ε*κ*n*2 # Two because the right hand side has two columns
                        @test norm(a'*(lua'\a') - a', 1) < ε*κ*n^2
                        @test norm(a*(lua\cc) - cc, 1) < ε*κ*n # cc is a vector
                        @test norm(a'*(lua'\cc) - cc, 1) < ε*κ*n # cc is a vector
                        @test AbstractArray(lua) ≈ a
                        @test norm(transpose(a)*(transpose(lua)\bb) - bb,1) < ε*κ*n*2 # Two because the right hand side has two columns
                        @test norm(transpose(a)*(transpose(lua)\cc) - cc,1) < ε*κ*n
                    end

                    # Test whether Ax_ldiv_B!(y, UL, x) indeed overwrites y
                    resultT = typeof(oneunit(eltyb) / oneunit(eltya))

                    b_dest = similar(b, resultT)
                    c_dest = similar(c, resultT)

                    ldiv!(b_dest, lua, b)
                    ldiv!(c_dest, lua, c)
                    @test norm(b_dest - lua \ b, 1) < ε*κ*2n
                    @test norm(c_dest - lua \ c, 1) < ε*κ*n

                    ldiv!(b_dest, transpose(lua), b)
                    ldiv!(c_dest, transpose(lua), c)
                    @test norm(b_dest - transpose(lua) \ b, 1) < ε*κ*2n
                    @test norm(c_dest - transpose(lua) \ c, 1) < ε*κ*n

                    ldiv!(b_dest, adjoint(lua), b)
                    ldiv!(c_dest, adjoint(lua), c)
                    @test norm(b_dest - lua' \ b, 1) < ε*κ*2n
                    @test norm(c_dest - lua' \ c, 1) < ε*κ*n

                    if eltyb != Int && !(eltya <: Complex) || eltya <: Complex && eltyb <: Complex
                        p = Matrix(b')
                        q = Matrix(c')
                        p_dest = copy(p)
                        q_dest = copy(q)
                        rdiv!(p_dest, lua)
                        rdiv!(q_dest, lua)
                        @test norm(p_dest - p / lua, 1) < ε*κ*2n
                        @test norm(q_dest - q / lua, 1) < ε*κ*n
                    end
                end
                if eltya <: BlasFloat && eltyb <: BlasFloat
                    e = rand(eltyb,n,n)
                    @test norm(e/lua - e/a,1) < ε*κ*n^2
                end
            end
            # @testset "Tridiagonal UL" begin
            #     lud   = factorize(d)
            #     f = zeros(eltyb, n+1)
            #     @test_throws DimensionMismatch lud\f
            #     @test_throws DimensionMismatch transpose(lud)\f
            #     @test_throws DimensionMismatch lud'\f
            #     @test_throws DimensionMismatch LinearAlgebra.ldiv!(transpose(lud), f)
            #     let Bs = copy(b)
            #         for bb in (Bs, view(Bs, 1:n, 1))
            #             @test norm(d*(lud\bb) - bb, 1) < ε*κd*n*2 # Two because the right hand side has two columns
            #             if eltya <: Real
            #                 @test norm((transpose(lud)\bb) - Array(transpose(d))\bb, 1) < ε*κd*n*2 # Two because the right hand side has two columns
            #                 if eltya != Int && eltyb != Int
            #                     @test norm(LinearAlgebra.ldiv!(transpose(lud), copy(bb)) - Array(transpose(d))\bb, 1) < ε*κd*n*2
            #                 end
            #             end
            #             if eltya <: Complex
            #                 @test norm((lud'\bb) - Array(d')\bb, 1) < ε*κd*n*2 # Two because the right hand side has two columns
            #             end
            #         end
            #     end
            #     if eltya <: BlasFloat && eltyb <: BlasFloat
            #         e = rand(eltyb,n,n)
            #         @test norm(e/lud - e/d,1) < ε*κ*n^2
            #         @test norm((transpose(lud)\e') - Array(transpose(d))\e',1) < ε*κd*n^2
            #         #test singular
            #         du = rand(eltya,n-1)
            #         dl = rand(eltya,n-1)
            #         dd = rand(eltya,n)
            #         dd[1] = zero(eltya)
            #         du[1] = zero(eltya)
            #         dl[1] = zero(eltya)
            #         zT = Tridiagonal(dl,dd,du)
            #         @test !LinearAlgebra.issuccess(ul(zT; check = false))
            #     end
            # end
            # @testset "Thin UL" begin
            #     lua   = @inferred ul(a[:,1:n1])
            #     @test lua.U*lua.L ≈ lua.P'*a[:,1:n1]
            # end
            # @testset "Fat UL" begin
            #     lua   = ul(a[1:n1,:])
            #     @test lua.L*lua.U ≈ lua.P*a[1:n1,:]
            # end
        end

        # @testset "UL of Symmetric/Hermitian" begin
        #     for HS in (Hermitian(a'a), Symmetric(a'a))
        #         luhs = ul(HS)
        #         @test luhs.L*luhs.U ≈ luhs.P*Matrix(HS)
        #     end
        # end
    end

    @testset "Singular matrices" for T in (Float64, ComplexF64)
        A = T[1 2; 0 0]
        @test_throws SingularException ul(A)
        @test_throws SingularException ul!(copy(A))
        @test_throws SingularException ul(A; check = true)
        @test_throws SingularException ul!(copy(A); check = true)
        @test !issuccess(ul(A; check = false))
        @test !issuccess(ul!(copy(A); check = false))
        @test_throws ZeroPivotException ul(A, Val(false))
        @test_throws ZeroPivotException ul!(copy(A), Val(false))
        @test_throws ZeroPivotException ul(A, Val(false); check = true)
        @test_throws ZeroPivotException ul!(copy(A), Val(false); check = true)
        @test !issuccess(ul(A, Val(false); check = false))
        @test !issuccess(ul!(copy(A), Val(false); check = false))
        F = ul(A; check = false)
        @test sprint((io, x) -> show(io, "text/plain", x), F) ==
            "Failed factorization of type $(typeof(F))"
    end

    # @testset "conversion" begin
    #     Random.seed!(3)
    #     a = Tridiagonal(rand(9),rand(10),rand(9))
    #     fa = Array(a)
    #     falu = ul(fa)
    #     alu = ul(a)
    #     falu = convert(typeof(falu),alu)
    #     @test Array(alu) == fa
    #     @test AbstractArray(alu) == fa
    # end

    @testset "Rational Matrices" begin
        ## Integrate in general tests when more linear algebra is implemented in julia
        a = convert(Matrix{Rational{BigInt}}, rand(1:10//1,n,n))/n
        b = rand(1:10,n,2)
        @inferred ul(a)
        lua   = factorize(a)
        l,u,p = lua.L, lua.U, lua.p
        @test l*u ≈ a[p,:]
        @test l[invperm(p),:]*u ≈ a
        @test a*inv(lua) ≈ Matrix(I, n, n)
        let Bs = b
            for b in (Bs, view(Bs, 1:n, 1))
                @test a*(lua\b) ≈ b
            end
        end
        @test @inferred(det(a)) ≈ det(Array{Float64}(a))
    end

    @testset "Rational{BigInt} and BigFloat Hilbert Matrix" begin
        ## Hilbert Matrix (very ill conditioned)
        ## Testing Rational{BigInt} and BigFloat version
        nHilbert = 50
        H = Rational{BigInt}[1//(i+j-1) for i = 1:nHilbert,j = 1:nHilbert]
        Hinv = Rational{BigInt}[(-1)^(i+j)*(i+j-1)*binomial(nHilbert+i-1,nHilbert-j)*
            binomial(nHilbert+j-1,nHilbert-i)*binomial(i+j-2,i-1)^2
            for i = big(1):nHilbert,j=big(1):nHilbert]
        @test inv(H) == Hinv
        setprecision(2^10) do
            @test norm(Array{Float64}(inv(float(H)) - float(Hinv))) < 1e-100
        end
    end

    @testset "logdet" begin
        @test @inferred(logdet(ComplexF32[1.0f0 0.5f0; 0.5f0 -1.0f0])) === 0.22314355f0 + 3.1415927f0im
        @test_throws DomainError logdet([1 1; 1 -1])
    end

    @testset "REPL printing" begin
            bf = IOBuffer()
            show(bf, "text/plain", ul(Matrix(I, 4, 4)))
            seekstart(bf)
            @test String(take!(bf)) == """
UL{Float64,Array{Float64,2},Array{Int64,1}}
U factor:
4×4 Array{Float64,2}:
 1.0  0.0  0.0  0.0
 0.0  1.0  0.0  0.0
 0.0  0.0  1.0  0.0
 0.0  0.0  0.0  1.0
L factor:
4×4 Array{Float64,2}:
 1.0  0.0  0.0  0.0
 0.0  1.0  0.0  0.0
 0.0  0.0  1.0  0.0
 0.0  0.0  0.0  1.0"""
    end

    @testset "propertynames" begin
        names = sort!(collect(string.(Base.propertynames(ul(rand(3,3))))))
        @test names == ["L", "P", "U", "p"]
        allnames = sort!(collect(string.(Base.propertynames(ul(rand(3,3)), true))))
        @test allnames == ["L", "P", "U", "factors", "info", "ipiv", "p"]
    end

    @testset "Issue #30917. Determinant of integer matrix" begin
        @test det([1 1 0 0 1 0 0 0
                1 0 1 0 0 1 0 0
                1 0 0 1 0 0 1 0
                0 1 1 1 0 0 0 0
                0 1 0 0 0 0 1 1
                0 0 1 0 1 0 0 1
                0 0 0 1 1 1 0 0
                0 0 0 0 1 1 0 1]) ≈ 6
    end

    @testset "Issue #33177. No ldiv!(UL, Adjoint)" begin
        A = [1 0; 1 1]
        B = [1 2; 2 8]
        F = ul(B)
        @test (A  / F') * B == A
        @test (A' / F') * B == A'

        a = complex.(randn(2), randn(2))
        @test (a' / F') * B ≈ a'
        @test (transpose(a) / F') * B ≈ transpose(a)

        A = complex.(randn(2, 2), randn(2, 2))
        @test (A' / F') * B ≈ A'
        @test (transpose(A) / F') * B ≈ transpose(A)
    end

    @testset "0x0 matrix" begin
        A = ones(0, 0)
        F = ul(A)
        @test F.U == ones(0, 0)
        @test F.L == ones(0, 0)
        @test F.P == ones(0, 0)
        @test F.p == []
    end

    @testset "more ldiv! methods" begin
        for elty in (Float16, Float64, ComplexF64), transform in (transpose, adjoint)
            A = randn(elty, 5, 5)
            B = randn(elty, 5, 5)
            @test ldiv!(transform(ul(A)), copy(B)) ≈ transform(ul(A)) \ B ≈ transform(A) \ B
            @test ldiv!(transform(ul(A)), transform(copy(B))) ≈ transform(ul(A)) \ transform(B) ≈ transform(A) \ transform(B)
        end
    end

    @testset "more rdiv! methods" begin
        A = randn(5,5)
        B = randn(5,5)
        a = randn(5)
        @test rdiv!(copy(A), ul(B)) ≈ A / B

        for elty in (Float16, Float64, ComplexF64), transform in (transpose, adjoint)
            A = randn(elty, 5, 5)
            C = copy(A)
            B = randn(elty, 5, 5)
            @test rdiv!(transform(A), transform(ul(B))) ≈ transform(C) / transform(ul(B)) ≈ transform(C) / transform(B)
        end
    end

    @testset "conversions" begin
        A = randn(5,5)
        F = ul(A)
        @test ul(F) ≡ F ≡ Factorization{Float64}(F)
        @test UL{Float32}(F).factors ≈ UL{Float32,Matrix{Float32}}(F).factors ≈
                Factorization{Float32}(F).factors ≈ ul(Float32.(A)).factors
        @test copy(F).factors == F.factors
    end

    @testset "det" begin
        A = randn(5,5)
        @test det(A) ≈ det(ul(A))
        @test all(logabsdet(A) .≈ logabsdet(ul(A)))
    end

    @testset "inv" begin
        A = randn(5,5)
        @test inv(ul(A)) ≈ inv(A)
    end

    @testset "MyVector/MyMatrix" begin
        A = randn(5,5)
        b = randn(5)
        B = randn(5,3)
        b̃ = MyVector(copy(b))
        B̃ = MyMatrix(copy(B))
        @test ldiv!(ul(A), b̃) === b̃
        @test b̃ ≈ A \ b
        @test ldiv!(ul(A), B̃) === B̃
        @test B̃ ≈ A \ B
    end

    @testset "Tridiagonal" begin
        n = 10
        A = Tridiagonal(randn(n-1), randn(n), randn(n-1))+10I
        U,L = ul(A, Val(false))
        @test U isa Bidiagonal
        @test L isa Bidiagonal
        Ũ,L̃ = ul(Matrix(A), Val(false))
        @test U ≈ Ũ
        @test L ≈ L̃
        @test U*L ≈ A
    end
end


laptop_joltage(adapters) = maximum(adapters) + 3

parse_adapters(s) = map(x->parse(Int, x), split(s, keepempty=false))

#=
adapter_str = """
16
10
15
5
1
11
7
19
6
12
4
"""
adapters = parse_adapters(adapter_str)
start_jolt = 0
end_jolt = laptop_joltage(adapters)
=#

function diff_dist(adapters, start_jolt, end_jolt)
    v = sort([adapters..., start_jolt, end_jolt])
    diffs = diff(v)
    m = maximum(diffs)
    dist = zeros(m)
    for d in diffs
        dist[d] += 1
    end
    dist
end

function num_paths(adapters, start_jolt, end_jolt)
    v = sort([adapters..., start_jolt, end_jolt])
    n = length(v)
    acc_paths = zeros(n)
    acc_paths[1] = 1
    for i in 1:n, j in i+1:i+3
        if j > n
            continue
        end
        if v[j] - v[i] <= 3
            acc_paths[j] += acc_paths[i]
        end
    end
    last(acc_paths)
end

#=
diff_dist(adapters, start_jolt, end_jolt)
=#

function day10part1(adapters)
    start_jolt = 0
    end_jolt = laptop_joltage(adapters)
    dist = diff_dist(adapters, start_jolt, end_jolt)
    dist[1] * dist[3]
end

function day10part2(adapters)
    start_jolt = 0
    end_jolt = laptop_joltage(adapters)
    num_paths(adapters, start_jolt, end_jolt)
end

function day10(file)
    s = read(file, String)
    adapters = parse_adapters(s)
    (day10part1(adapters), day10part2(adapters))
end

#=
day10("input.txt")
=#


🐙 = mapreduce(line -> parse.(Int8, split(line, "")), hcat, readlines("input.txt"))
💡 = fill(false, size(🐙))
const R = CartesianIndices(🐙)
const If, Il = extrema(R)
const I1 = one(If)

function flash!(🐙, 💡, I)
    💡[I] = true
    for i in max(I-I1, If):min(I+I1, Il)
        🐙[i] += 1
        🐙[i] > 9 && !💡[i] && flash!(🐙, 💡, i)
    end
end

function step!(🐙, 💡)
    fill!(💡, false)
    for I in R; 🐙[I] += 1; end
    for I in R; 🐙[I] > 9 && !💡[I] && flash!(🐙, 💡, I); end
    for I in R; 💡[I] && (🐙[I] = 0); end
end

function solve(🐙, 💡)
    🧮 = 💥 = i = 0
    while !all(💡) || i <= 100
        i += 1
        step!(🐙, 💡)
        🧮 = i <= 100 ? 🧮 + sum(💡) : 🧮
        💥 = all(💡) ? 💥 + i : 💥
    end
    return 🧮, 💥
end

🧮, 💥 = solve(🐙, 💡)
@info "Number of flashes in first 100 steps = $🧮"
@info "First step where all flash = $💥"


using GeometricBase.Config
using GeometricBase.Utils
using GeometricIntegrators.Integrators
using GeometricIntegrators.Tableaus
using GeometricProblems.HarmonicOscillator
using SimpleSolvers
using Test

SimpleSolvers.set_config(:nls_stol_break, 1E3)

using GeometricProblems.HarmonicOscillator: Δt, nt, refx

sode = harmonic_oscillator_sode()

sint = Integrator(sode, TableauLieA(), Δt)
ssol = integrate(sode, sint, nt)
@test relative_maximum_error(ssol.q, refx) < 5E-2

sintc = IntegratorComposition(sode, TableauLieA(), Δt)
ssolc = integrate(sode, sintc, nt)
@test ssol.q == ssolc.q

sint = Integrator(sode, TableauLieB(), Δt)
ssol = integrate(sode, sint, nt)
@test relative_maximum_error(ssol.q, refx) < 5E-2

sintc = IntegratorComposition(sode, TableauLieB(), Δt)
ssolc = integrate(sode, sintc, nt)
@test ssol.q == ssolc.q

sint = Integrator(sode, TableauStrang(), Δt)
ssol = integrate(sode, sint, nt)
@test relative_maximum_error(ssol.q, refx) < 1E-3

sintc = IntegratorComposition(sode, TableauStrang(), Δt)
ssolc = integrate(sode, sintc, nt)
@test ssol.q == ssolc.q

sint = Integrator(sode, TableauStrangA(), Δt)
ssol = integrate(sode, sint, nt)
@test relative_maximum_error(ssol.q, refx) < 1E-3

sintc = IntegratorComposition(sode, TableauStrangA(), Δt)
ssolc = integrate(sode, sintc, nt)
@test ssol.q == ssolc.q

sint = Integrator(sode, TableauStrangB(), Δt)
ssol = integrate(sode, sint, nt)
@test relative_maximum_error(ssol.q, refx) < 1E-3

sintc = IntegratorComposition(sode, TableauStrangB(), Δt)
ssolc = integrate(sode, sintc, nt)
@test ssol.q == ssolc.q

sint1 = Integrator(sode, TableauStrang(), Δt)
ssol1 = integrate(sode, sint, nt)
sint2 = Integrator(sode, TableauStrangA(), Δt)
ssol2 = integrate(sode, sint, nt)
sint3 = Integrator(sode, TableauStrangB(), Δt)
ssol3 = integrate(sode, sint, nt)
@test ssol1.q == ssol2.q
@test ssol1.q ≈  ssol3.q

sint = Integrator(sode, TableauMcLachlan2(), Δt)
ssol = integrate(sode, sint, nt)
@test relative_maximum_error(ssol.q, refx) < 1E-4

sintc = IntegratorComposition(sode, TableauMcLachlan2(), Δt)
ssolc = integrate(sode, sintc, nt)
@test ssol.q == ssolc.q

sint = Integrator(sode, TableauMcLachlan4(), Δt)
ssol = integrate(sode, sint, nt)
@test relative_maximum_error(ssol.q, refx) < 5E-4

sintc = IntegratorComposition(sode, TableauMcLachlan4(), Δt)
ssolc = integrate(sode, sintc, nt)
@test ssol.q == ssolc.q

sint = Integrator(sode, TableauTripleJump(), Δt)
ssol = integrate(sode, sint, nt)
@test relative_maximum_error(ssol.q, refx) < 5E-6

sintc = IntegratorComposition(sode, TableauTripleJump(), Δt)
ssolc = integrate(sode, sintc, nt)
@test ssol.q == ssolc.q

sint = Integrator(sode, TableauSuzukiFractal(), Δt)
ssol = integrate(sode, sint, nt)
@test relative_maximum_error(ssol.q, refx) < 5E-7

sintc = IntegratorComposition(sode, TableauSuzukiFractal(), Δt)
ssolc = integrate(sode, sintc, nt)
@test ssol.q == ssolc.q

DT = eltype(sode)
D  = ndims(sode)

ints_glrk1 = (IntegratorConstructor(DT, D, TableauGauss(1)), IntegratorConstructor(DT, D, TableauGauss(1)))
ints_erk4  = (IntegratorConstructor(DT, D, TableauRK4()), IntegratorConstructor(DT, D, TableauRK4()))

sint = IntegratorComposition(sode, ints_erk4, TableauLieA(), Δt)
ssol = integrate(sode, sint, nt)
# println(relative_maximum_error(ssol.q, refx))
@test relative_maximum_error(ssol.q, refx) < 5E-2

sint = IntegratorComposition(sode, ints_glrk1, TableauStrang(), Δt)
ssol = integrate(sode, sint, nt)
# println(relative_maximum_error(ssol.q, refx))
@test relative_maximum_error(ssol.q, refx) < 1E-3


export make_LN_basis,
       make_LN1N2_basis,
       make_spin1_LSz_basis,
       make_spin1_L_basis,
       get_num_spin1_states

"""
    make_LN_basis(L::Int, N::Int)

Compute Fock basis states with L sites and N particles.
"""
function make_LN_basis(L::Int, N::Int)
    states = Vector{Int}(undef, binomial(L, N))
    tmp = 1
    for i=0:(1<<L)-1
        if count_ones(i) == N
            states[tmp] = i
            tmp += 1
        end
    end
    return states
end

"""
    make_LN1N2_basis(L::Int, N1::Int, N2::Int)

Compute Fock basis states with L sites, N1 particles of type 1 and N2 particles
of type 2 (e.g. spin up and spin down).
"""
function make_LN1N2_basis(L::Int, N1::Int, N2::Int)
    states = Vector{Int}(undef, binomial(L, N1)*binomial(L, N2))
    tmp = 1
    mask1 = 1<<L - 1
    mask2 = 1<<(2L) - 1 - mask1
    for i=0:(1<<(2L))-1
        if count_ones(i&mask1) == N1 && count_ones(i&mask2) == N2
            states[tmp] = i
            tmp += 1
        end
    end
    return states
end

"""
    make_spin1_LSz_basis(L::Int, Sz::Int)

Compute states with L spins s=1 and Sz magnetization.

There can be 3 states per individual site/spin. In the usual notation
these are: |1, 0>, |1, 1>, |1, -1>. Because we can't represent these
states with only one bit of information we map each site to the three
possible states:
* |1, 0> = (00)
* |1, 1> = (01)
* |1,-1> = (10)
For example, a state with three spins, with magnetization -1, 1, 0 in
each site, respectively, is represented as: 100100.
"""
function make_spin1_LSz_basis(L::Int, Sz::Int)
    # Number of states.
    num_states = 0
    l = iseven(L+Sz) ? L : L-1
    # Number of spins with m=1 and m=-1 in a chain of l sites.
    Nu = (l+Sz)>>1
    Nd = (l-Sz)>>1
    while l >= 0 && Nu >= 0 && Nd >= 0
        num_states += binomial(l, Nu)*binomial(L, l)
        l -= 2
        Nu -= 1
        Nd -= 1
    end

    states = zeros(Int, num_states)
    cont = 1
    for i=0:(1<<(2L))-1
        state_Sz = 0
        is_valid = true
        for j=0:L-1
            state_Sz += (i>>(2j))&1
            state_Sz -= (i>>(2j+1))&1
            if (i>>(2j))&1 == 1 && (i>>(2j+1))&1 == 1
                is_valid = false
            end
        end
        if state_Sz == Sz && is_valid
            states[cont] = i
            cont += 1
        end
    end
    return states
end

"""
    get_num_spin1_states(L::Int, Sz::Int)

Compute the number of states with L sites and Sz magnetization.
"""
function get_num_spin1_states(L::Int, Sz::Int)
    return length(make_spin1_LSz_basis(L, Sz))
end

"""
    make_spin1_L_basis(L::Int)

Compute states with L spins s=1.

The states are defined as in make_spin1_LSz_basis.
"""
function make_spin1_L_basis(L::Int)
    num_states = 3^L
    states = zeros(Int, num_states)
    cont = 1
    for i=0:(1<<(2L))-1
        is_valid = true
        for j=0:L-1
            if (i>>(2j))&3 == 3
                is_valid = false
            end
        end
        if is_valid
            states[cont] = i
            cont += 1
        end
    end
    return states
end


push!(LOAD_PATH, pwd())
using AtomicOpt
using Printf
using LinearAlgebra
import Random: seed!, randperm, shuffle
using Distributed
using BSON: @save, @load

function randop(n::Integer, family=:orth)
    if family == :orth
        Q = Matrix(qr(randn(n,n)).Q)
    end
    return Q
end

function randnsparse(n, nnz)
    p = randperm(n)[1:nnz]
    x = zeros(n)
    x[p] = rand(nnz)
    return x
end

struct Signal{opT, atomT}
    n::Int64
    nnz::Int64
    x::Vector{Float64}
    Q::opT
    A::atomT
end

function Signal(n, nnz; family=:orth)
    x = randnsparse(n, nnz)
    Q = randop(n, family)
    λ = norm(x, 1)
    A = λ*(Q*OneBall(n; maxrank = nnz))
    return Signal(n, nnz, x, Q, A)
end

rotate(s::Signal) = s.Q*s.x

function RecordError(m, n, k, nnz, t)
    signals = ntuple(d->Signal(n, nnz), k)
    M = rand(m, n)
    b = M*sum(rotate(s) for s in signals)
    A = mapreduce(s->s.A, ×, signals)
    x, normr = level_set(M, b, A,  α = 0.0, tol = 1e-3, maxIts=10*length(b), logger=0)
    err = 0
    for (i,(x,s)) in enumerate(zip(x, signals))
        err1 = norm(x-s.Q*s.x);
        if err1 > err
            err = err1
        end
    end
    @printf("  m = %3d, n = %3d, k = %1d, nnz = %1d, t = %2d, error = %8.2e\n", m, n, k, nnz, t, err)
    flush(stdout)
    return err
end

function RecordError2(m, n, k, nnz, α, t)
    signals = ntuple(d->Signal(n, nnz), k)
    M = randn(m, n)
    b = M*sum(rotate(s) for s in signals)
    u = randn(m); u .*= α / norm(u)
    b = b + u
    A = mapreduce(s->s.A, ×, signals)
    x, normr = level_set(M, b, A,  α = α^2/2, tol = 1e-3, maxIts=10*length(b), logger=0)
    err = 0
    for (i,(x,s)) in enumerate(zip(x, signals))
        err1 = norm(x-s.Q*s.x);
        if err1 > err
            err = err1
        end
    end
    @printf("  α = %8.2e, t = %3d\n", α, t)
    flush(stdout)
    return err
end

# function AvgError(m, n, k, nnz)
#     out = @distributed (+) for i = 1:100
#         RecordError(m, n, k, nnz)
#     end
#     @printf("  m = %7d, n = %7d, k = %7d, nnz = %7d, error = %8.2e\n", m, n, k, nnz, out/100)
#     flush(stdout)
#     return out/100
# end

# @show RecordError(200, 200, 4, 5, 1)

# m = 125
# n = 200
# k = 3
# nnz = 5
# J = CartesianIndices((31, 31, 50))
# @printf("Start E2!\n")
# E2 = pmap(i->RecordError(500 - (i[1] - 1)*15, 50+(i[2]-1)*15, 2, nnz, i[3]), J)
# @printf("Start E3!\n")
# E3 = pmap(i->RecordError(500 - (i[1] - 1)*15, 50+(i[2]-1)*15, 3, nnz, i[3]), J)
# @printf("Start E4!\n")
# E4 = pmap(i->RecordError(500 - (i[1] - 1)*15, 50+(i[2]-1)*15, 4, nnz, i[3]), J)
# J = CartesianIndices((200, 50))
# E = pmap(i->RecordError2(m, n, k, nnz, 0.01*i[1], i[2]), J)

# # save
# @save "./examples/PhaseTransition/NoisyMeasurement.bson" E
# @save "./examples/PhaseTransition/NoiselessMeasurement.bson" E2 E3 E4

nnz = 3
n = 1000
J = CartesianIndices((19, 9, 10))
E = pmap(i->RecordError(1000 - (i[1] - 1)*50, n, i[2]+1, nnz, i[3]), J)
@save "./examples/PhaseTransition/NoiselessMeasurement2.bson" E


@testset "radial_profile" begin

    X = [0 1 1 1 0
         1 2 2 2 1
         1 2 3 2 1
         1 2 2 2 1
         0 1 1 1 0];

    r, prof = @inferred radial_profile(X);
    @test r == [0, 1, 2, 3]
    @test prof ≈ [3, 2, 1, 0]

    X = [1 2
         2 2]
    r, prof = @inferred radial_profile(X, center=(1, 1))
    @test r == [0, 1]
    @test prof ≈ [1, 2]

    img = randn(rng, 1001, 1001)
    r, prof = @inferred radial_profile(img)
    @test mean(prof) ≈ 0 atol=1e-2
end


### A Pluto.jl notebook ###
# v0.15.1

using Markdown
using InteractiveUtils

# ╔═╡ 7a206e18-4562-469d-b580-f2d0cc2f80df
using Pkg; Pkg.develop(path=".."); Pkg.activate("..")

# ╔═╡ 7c28cf4c-cd0c-11eb-0507-b10f2a883e8c
begin
    using RDatasets
    using LinearAlgebra
    using Plots
    using StatsPlots
    plotly(size=(680, 500))
end

# ╔═╡ 93477f7b-aef4-44e5-ba3b-8285a6261500
md"
# K-Medoids

JingYu Ning
"

# ╔═╡ 64f0d05d-9c4e-4935-bf33-a5ee5835f99a
md"
## IRIS dataset
"

# ╔═╡ 7f59bdb6-c6de-4a25-9ff0-8a81457a4242
iris = dataset("datasets", "iris")

# ╔═╡ 63073482-7add-4234-ab34-6c8068b2f5a6
begin
    cols = names(iris)

    function gen_pots()
        plots = Matrix{Plots.Plot}(undef, 4, 4)
        for i in 1:4, j in 1:4
            if i == j
                plots[i, j] = @df(iris, histogram(cols(i)))
            elseif i < j
                plots[i, j] = @df(iris, histogram2d(
                    cols(i), cols(j),
                    nbins=20, color=:coolwarm, colorbar=false
                ))
            else
                plots[i, j] =  @df(iris, scatter(cols(i), cols(j)))
            end

            (i==1) && (plots[i, j] = plot!(ylabel=cols[j])) # add x label
            (j==4) && (plots[i, j] = plot!(xlabel=cols[i])) # add y label
        end

        return plots
    end
end

# ╔═╡ fde3d98a-be02-47fe-a8c1-293092298659
plot(gen_pots()..., layout=(4, 4), legend=false)

# ╔═╡ 30f8c663-e43b-44e2-9c91-b3e8dd760e5f
md"
## Initial center points
"

# ╔═╡ b1b771a2-3e94-4666-a668-01dee57f9832
md"
**Data size and numbers of classes**
"

# ╔═╡ 52e9ada2-cff4-4fe4-abb0-ec0ebc2d3fa1
begin
    n = size(iris, 1)
    n_classes = 3
end;

# ╔═╡ cbb9a475-36d9-4c9f-a12a-d2cc5c339be7
md"
**Initial three center points**
"

# ╔═╡ ad7bfec2-3fe8-428c-ac92-ad56183cbcdd
# cs_i = rand(1:n, n_classes)
cs_i = [77, 90, 13]

# ╔═╡ 2c6511ca-e070-47ac-8f2e-c1308f491f98
let
    plots = gen_pots()
    for i in 2:4, j in 1:(i-1)
        plots[i, j] = @df iris[cs_i, :] scatter!(plots[i, j], cols(i), cols(j))
    end
    plot(plots..., layout=(4, 4), legend=false)
end

# ╔═╡ bb476b47-ba34-41d7-a32f-fc204f4a328c
md"
## Training
"

# ╔═╡ f558898f-d426-4a36-baed-fa4dd07cf7b2
let
	data = randn(2, 1000)
	data_o = hcat(data, [9, -9])

	function d²(data::Matrix, i::Integer, j::Integer)
		return sum(abs.(data[:, j]-data[:, i]))
	end

	xlim=(-10, 10); ylim=(-10, 10)

	pd = scatter(data[1, :], data[2, :], xlim=xlim, ylim=ylim)
	id = argmin([sum(d²(data, i, j) for j in 1:1000) for i in 1:1000])
	scatter!(pd, [data[1, id]], [data[2, id]])

	po = scatter(data_o[1, :], data_o[2, :], xlim=xlim, ylim=ylim)
	io = argmin([sum(d²(data_o, i, j) for j in 1:1000) for i in 1:1000])
	scatter!(po, [data_o[1, io]], [data_o[2, io]])
	plot(pd, po, legend=false)
end

# ╔═╡ e6f38859-fd04-4dd6-84f3-991fdc72317d
md"
**Manhattan distance between two points**
"

# ╔═╡ 9ce95067-b54f-4a9d-8158-8e25b4a2480f
function d(data::DataFrame, p₁_i::Integer, p₂_i::Integer)
	return sum(abs.(Vector(data[p₁_i, 1:4]) - Vector(data[p₂_i, 1:4])))
end

# ╔═╡ 9cf597b7-9ea2-4c5d-b810-83b5d027d758
begin
    signal = "⚡"

    iris[!, :Predict] = zeros(n)
    for _ in 1:100
        iris[!, :Predict] = map(i->argmin([d(iris, i, ci) for ci in cs_i]), 1:n)

		iris[!, :SumInClassD] = [
			sum((iris[i, :Predict]==iris[j, :Predict] && d(iris, i, j)) for j in 1:n)
			for i in 1:n
		]
		for i in 1:n
			if iris[cs_i[iris[i, :Predict]], :SumInClassD] > iris[i, :SumInClassD]
				cs_i[iris[i, :Predict]] = i
			end
		end
    end
end

# ╔═╡ 7b5ea310-f5c2-4f9d-91b0-2ebae7b8639e
md"
**Prediction**
"

# ╔═╡ e0b73f20-18ba-4dff-9b75-0fe82d072167
let
    signal
    plots = gen_pots()
    for i in 2:4, j in 1:(i-1)
        plots[i, j] = @df iris scatter!(plots[i, j], cols(i), cols(j), group=:Predict)
        plots[i, j] = @df iris[cs_i, :] scatter!(plots[i, j], cols(i), cols(j))
    end
    plot(plots..., layout=(4, 4), legend=false)
end

# ╔═╡ 413aa844-5ce1-4349-a715-b62c8db9a1f8
md"
**Ground truth**
"

# ╔═╡ 1ae9328e-72c1-46e4-ba1e-742205c4bd88
let
    plots = gen_pots()
    for i in 2:4, j in 1:(i-1)
        plots[i, j] = @df iris scatter!(plots[i, j], cols(i), cols(j), group=:Species)
    end
    plot(plots..., layout=(4, 4), legend=false)
end

# ╔═╡ Cell order:
# ╟─93477f7b-aef4-44e5-ba3b-8285a6261500
# ╟─7a206e18-4562-469d-b580-f2d0cc2f80df
# ╠═7c28cf4c-cd0c-11eb-0507-b10f2a883e8c
# ╟─64f0d05d-9c4e-4935-bf33-a5ee5835f99a
# ╠═7f59bdb6-c6de-4a25-9ff0-8a81457a4242
# ╟─63073482-7add-4234-ab34-6c8068b2f5a6
# ╟─fde3d98a-be02-47fe-a8c1-293092298659
# ╟─30f8c663-e43b-44e2-9c91-b3e8dd760e5f
# ╟─b1b771a2-3e94-4666-a668-01dee57f9832
# ╠═52e9ada2-cff4-4fe4-abb0-ec0ebc2d3fa1
# ╟─cbb9a475-36d9-4c9f-a12a-d2cc5c339be7
# ╠═ad7bfec2-3fe8-428c-ac92-ad56183cbcdd
# ╟─2c6511ca-e070-47ac-8f2e-c1308f491f98
# ╟─bb476b47-ba34-41d7-a32f-fc204f4a328c
# ╟─f558898f-d426-4a36-baed-fa4dd07cf7b2
# ╟─e6f38859-fd04-4dd6-84f3-991fdc72317d
# ╠═9ce95067-b54f-4a9d-8158-8e25b4a2480f
# ╠═9cf597b7-9ea2-4c5d-b810-83b5d027d758
# ╟─7b5ea310-f5c2-4f9d-91b0-2ebae7b8639e
# ╟─e0b73f20-18ba-4dff-9b75-0fe82d072167
# ╟─413aa844-5ce1-4349-a715-b62c8db9a1f8
# ╟─1ae9328e-72c1-46e4-ba1e-742205c4bd88


module Geo

include("../src/Jaynes.jl")
using .Jaynes
using Distributions

geo(p::Float64) = rand(:flip, Bernoulli(p)) ? 1 : 1 + rand(:geo, geo, p)

# Define as primitive.
@primitive function logpdf(fn::typeof(geo), p, count)
    return Distributions.logpdf(Geometric(p), count)
end

# Example trace.
ret, cl = simulate(() -> rand(:geo, geo, 0.05))
display(cl.trace)

end # module


##

function _locpoly_alt1!(
  eX::Matrix{T}, w::Vector{T},
  X::Matrix{T}, z::Vector{T}, y::Vector{T},
  z0::T, degree::Int64, kernel::SmoothingKernel{T}) where {T <: AbstractFloat}

  @. w = evaluate(kernel, z, z0)
  _expand_X!(eX, X, z, z0, degree)

  ((eX' * diagm(w)) * eX) \ ((eX' * diagm(w)) * y)
end

function locpoly_alt1(
  X::Matrix{T}, z::Vector{T}, y::Vector{T},
  zgrid::Vector{T},
  degree::Int64,                                     # degree of the polynomial
  kernel::SmoothingKernel{T}=GaussianKernel(one(T))) where {T <: AbstractFloat}

  n, p = size(X)
  ep = p * (degree + 1)
  out = Array{T}(ep, length(zgrid))
  w = Array{T}(n)
  eX = Array{T}(n, ep)

  ind = 0
  for z0 in zgrid
    ind += 1
    out[:, ind] = _locpoly_alt1!(eX, w, X, z, y, z0, degree, kernel)
  end
  out
end

##

function _locpoly_alt!(
  Xt_w_X::Matrix{T}, Xt_w_Y::Vector{T}, w::Vector{T},
  X::Matrix{T}, z::Vector{T}, y::Vector{T},
  z0::T, degree::Int64, kernel::SmoothingKernel{T}) where {T <: AbstractFloat}

  @. w = evaluate(kernel, z, z0)
  _expand_Xt_w_X!(Xt_w_X, w, X, z, z0, degree)
  _expand_Xt_w_Y!(Xt_w_Y, w, X, z, y, z0, degree)

  cholfact!(Xt_w_X) \ Xt_w_Y
end

function locpoly_alt(
  X::Matrix{T}, z::Vector{T}, y::Vector{T},
  zgrid::Vector{T},
  degree::Int64,                                     # degree of the polynomial
  kernel::SmoothingKernel{T}) where {T <: AbstractFloat}

  n, p = size(X)
  ep = p * (degree + 1)
  out = Array{T}(ep, length(zgrid))
  w = Array{T}(n)
  Xt_w_X = Array{T}(ep, ep)
  Xt_w_Y = Array{T}(ep)

  ind = 0
  for z0 in zgrid
    ind += 1
    out[:, ind] = _locpoly_alt!(Xt_w_X, Xt_w_Y, w, X, z, y, z0, degree, kernel)
  end
  out
end

##

function locpolyl1_alt(
  X::Matrix{T}, z::Vector{T}, y::Vector{T},
  zgrid::Vector{T},
  degree::Int64,
  kernel::SmoothingKernel{T},
  λ0::T,
  options::CDOptions=CDOptions()) where {T <: AbstractFloat}

  # construct inner options because we do not want warmStart = false
  # we want to start from the previous iteration since the points
  # on the grid should be close to each other
  opt = CDOptions(options.maxIter, options.optTol, options.randomize, true, options.numSteps)

  n, p = size(X)
  ep = p * (degree + 1)
  out = Array{SparseVector{T, Int64}}(length(zgrid))

  # temporary storage
  w = Array{T}(n)
  wX = Array{T}(n, ep)
  stdX = ones(T, ep)

  ind = 0
  for z0 in zgrid
    ind += 1

    # the following two should update f
    @. w = evaluate(kernel, z, z0)
    _expand_X!(wX, X, z, z0, degree)
    # compute std for each column
    @inbounds for j=1:ep
      v = zero(T)
      @simd for i=1:n
        v += wX[i, j] * wX[i, j] * w[i]
      end
      stdX[j] = sqrt(v / n)
    end

    @. w = sqrt(w)
    f = CDLeastSquaresLoss(diagm(w)*y, diagm(w)*wX)
    g = ProximalBase.ProxL1(λ0, stdX)
    β = ProximalBase.SparseIterate(ep)

    # solve for β
    coordinateDescent!(β, f, g, opt)
    out[ind] = convert(SparseVector, β)
  end
  out
end

##

function _expand_Xt_w_X_alt!(
  Xt_w_X::Matrix{T},
  w::Vector{T}, X::Matrix{T},
  z::Vector{T}, z0::T, degree::Int64) where {T <: AbstractFloat}

  n, p = size(X)
  fill!(Xt_w_X, zero(T))
  @inbounds for j=1:p, jj=0:degree
    col=(j-1)*(degree+1)+1+jj
    for k=j:p
      if k != j
        krange = 0:degree
      else
        krange = jj:degree
      end
      for kk=krange
        row = (k-1)*(degree+1)+1+kk

        # compute Xt_w_X[row, col]
        for i=1:n
          Xt_w_X[row, col] += X[i, j] * (z[i] - z0)^(jj+kk) * X[i, k] * w[i]
        end
        if row != col
          Xt_w_X[col, row] = Xt_w_X[row, col]
        end
      end
    end
  end
  Xt_w_X
end

function genData(n, p)
  X = randn(n, p)
  Z = rand(n)
  ɛ = 0.1 * randn(n)
  Y = zeros(n)
  betaMat = zeros(p, n)

  randBeta = [rand([2,4,6,8]) for j=1:p]
  for i=1:n
    betaMat[:, i] = sin.(randBeta * Z[i])
    Y[i] = dot(betaMat[:, i], X[i, :]) + ɛ[i]
  end
  Y, X, Z, betaMat
end

function genDataIID(n, p)
  X = randn(n, p)
  Z = rand(n)
  ɛ = 0.1 * randn(n)
  Y = zeros(n)
  beta = 0.5+rand(p)

  for i=1:n
    Y[i] = dot(beta, X[i, :]) + ɛ[i]
  end
  Y, X, Z, beta
end

# using BenchmarkTools

####
#
# n, p = 100, 2
# Y, X, Z, betaMat = genData(n, p)
# zgrid = collect(0.01:0.1:0.99)
#
# gk = GaussianKernel(0.2)
# degree = 3
#
# @time o1 = locpoly(X, Z, Y, zgrid, 1, gk)
# @time o2 = locpoly_alt(X, Z, Y, zgrid, 1, gk)
# @time o3 = locpoly_alt1(X, Z, Y, zgrid, 1, gk)
#
# maximum(abs.(o1-o2))
# maximum(abs.(o1-o3))
#
# @benchmark locpoly($X, $Z, $Y, $zgrid, $degree, $gk)
# @benchmark locpoly_alt($X, $Z, $Y, $zgrid, $degree, $gk)
# @benchmark locpoly_alt1($X, $Z, $Y, $zgrid, $degree, $gk)

####

# p = 10
# X = randn(100, p)
# z = rand(100)
# w = zeros(100)
# k = GaussianKernel(0.2)
# @. w = evaluate(k, z, 0.5)
#
#
# degree = 2
# cp = p*(degree+1)
#
# eX = zeros(100, cp)
# _expand_X!(eX, X, z, 0.5, degree)
#
# Xt_w_X = zeros(cp, cp)
# Xt_w_X1 = zeros(cp, cp)
# o1 = _expand_Xt_w_X!(Xt_w_X, w, X, z, 0.5, degree)
# o2 = _expand_Xt_w_X_alt!(Xt_w_X1, w, X, z, 0.5, degree)
# o3 = (eX'*diagm(w))*eX
#
# maximum(abs.(o1 - o2))
# maximum(abs.(o1 - o3))
#
# using BenchmarkTools
# @benchmark _expand_Xt_w_X!($Xt_w_X, $w, $X, $z, 0.5, degree)
# @benchmark _expand_Xt_w_X_alt!($Xt_w_X1, $w, $X, $z, 0.5, degree)

### benchmark local lasso
#
#
# n, s, p = 500, 2, 50
# degree = 2
# gk = GaussianKernel(0.2)
# w = zeros(n)
# sw = zeros(n)
#
#
# Y, X, Z, betaMat = genData(n, s)
# X = [X zeros(n, p-s)]
# z0 = rand()
#
# cp = p*(degree+1)
# eX = zeros(n, cp)
# CoordinateDescent._expand_X!(eX, X, Z, z0, degree)
#
# @. w = evaluate(gk, Z, z0)
# @. sw = sqrt(w)
#
# λ = 0.1
# g = ProximalBase.ProxL1(λ)
# f1 = CDLeastSquaresLoss(diagm(sw) * Y, diagm(sw)*X)
# f2 = CDWeightedLSLoss(Y, X, w)
#
# opt = CDOptions(;maxIter=5000, optTol=1e-8, warmStart=false, randomize=true)
# x1 = convert(ProximalBase.SparseIterate, sprand(p, 0.6))
# x2 = convert(ProximalBase.SparseIterate, sprand(p, 0.6))
#
# @benchmark coordinateDescent!($x1, $f1, $g, $opt)
# @benchmark coordinateDescent!($x2, $f2, $g, $opt)

## benchmark locpolyl1
#
# n, s, p = 500, 2, 50
# degree = 0
# gk = GaussianKernel(0.2)
#
# Y, X, Z, betaT = genData(n, s)
# X = [X zeros(n, p-s)]
# zgrid = collect(0.01:0.1:0.99)
#
# λ0 = 0.01
#
# opt = CDOptions(;randomize=false)
# o1 = locpolyl1(X,Z,Y,zgrid,degree,gk,λ0, opt)
# o2 = locpolyl1_alt(X,Z,Y,zgrid,degree,gk,λ0, opt)
#
# maximum( maximum(abs.(o1[i] - o2[i])) for i=1:length(zgrid) )

#