import numpy as np
from raytrace_2D import Disk
import csv
import colorspace
from rich.progress import Progress, track
from PIL import Image
from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor, as_completed

class Ray:
	def __init__(self, origin, direction, intensity):
		self.origin = np.array(origin)
		self.direction = np.array(direction) / np.linalg.norm(direction)
		self.intensity = intensity

center = [0,0]
r = 1
# n = 1.3
# disk = Disk(center, r, n)
N = 100000
min_intensity = 0.00001
max_ray = 100*N

stack = []
result = []
points = []

# def init():
# 	stack = []
# 	for x in np.linspace(-r + dx, r-dx, N):
# 		stack.append(Ray([x, 2*r], [0,-1], 10*np.abs(x)))
# 	return stack

def reflection_and_refraction(ray:Ray, intersection_point, normal, n):
	# print("reflection/refraction at the point:", intersection_point)
	k_n = np.dot(ray.direction, normal)
	if k_n > 0:
		n1,n2 = n,1
	else:
		n1, n2 = 1,n
	kp = ray.direction - k_n*normal
	k2p = n1*kp/n2
	if np.dot(k2p, k2p) >= 1:
		ray1 = Ray(intersection_point, ray.direction-2*k_n*normal, ray.intensity)
		return ray1,None
	else:
		k2n = np.sqrt(1 - np.dot(k2p, k2p))*np.sign(k_n)
		Rs = (n1*k_n - n2*k2n)*(n1*k_n - n2*k2n)/((n1*k_n + n2*k2n)*(n1*k_n + n2*k2n))
		Rp = (n1*k2n - n2*k_n)*(n1*k2n - n2*k_n)/((n1*k2n + n2*k_n)*(n1*k2n + n2*k_n))
		R = (Rs+Rp)/2
		ray1 = Ray(intersection_point, ray.direction-2*k_n*normal, ray.intensity*R)
		ray2 = Ray(intersection_point, k2n*normal+k2p, ray.intensity*(1-R))
		return ray1,ray2

def trace(disk:Disk, stack:list, max_ray:int):
	ray_count = 0
	result = []
	while stack and ray_count < max_ray:
		ray = stack.pop()
		ray_count += 1
		if ray is None:
			continue
		if isinstance(ray, Ray):
			if ray.intensity < min_intensity:
				continue
			direction = ray.direction
			t,intersection_point = disk.find_intersection(ray)
			if intersection_point is not None:
				points.append(np.concatenate((ray.origin, intersection_point,[ray.intensity])))
				normal = disk.get_normal(intersection_point)
				stack.extend(reflection_and_refraction(ray, intersection_point, normal, disk.refractive_index))
			else:
				points.append(np.concatenate((ray.origin, ray.origin+ray.direction,[ray.intensity])))
				if direction[1] > 0:
					result.append([np.arccos(direction[1]), ray.intensity*direction[1]])
	# print("ray trace finish, total ray count:",ray_count)
	return result

def water_density(t):
	data_t = np.linspace(0,101,102)
	data_rho = np.array([999.79,999.84,999.88,999.92,999.97,999.92,999.88,999.84,999.79,999.72,999.65,999.55,999.44,999.32,999.19,999.05,998.90,998.74,998.56,998.36,998.16,997.96,997.74,997.50,997.25,996.99,996.74,996.48,996.20,995.90,995.59,995.29,995.97,994.65,994.33,994.98,993.64,993.28,992.93,992.55,992.17,991.78,991.39,990.99,990.57,990.16,989.75,989.33,988.90,988.45,988.04,987.55,987.06,986.58,986.19,985.70,985.22,984.73,984.67,983.18,983.13,982.70,982.22,981.64,981.06,980.48,979.91,979.33,978.85,978.28,977.70,977.13,976.56,975.99,975.41,974.84,974.27,973.70,973.14,972.47,971.81,971.25,970.59,969.93,969.27,968.61,967.96,967.30,966.65,965.99,965.34,964.69,964.04,963.39,962.64,961.90,961.26,960.52,959.78,959.04,958.40,957.67])
	return np.interp(t, data_t, data_rho)

def water_refraction_index(t, wavelength):
	rhobar = 0.999974950 * (1 - (t - 3.983035)*(t - 3.983035)*(t + 301.797)/(522528.9*(t+69.34881)))
	# print(rhobar)
	a0 = 0.244257733
	a1 = 0.00974634476
	a2 = -0.00373234996
	a3 = 0.000268678472
	a4 = 0.0015892057
	a5 = 0.00245934259
	a6 = 0.90070492
	a7 = -0.0166626219
	Tstar = 273.15
	Tbar = 1 + t/Tstar
	lambdastar = 589
	lambdabar = wavelength/lambdastar
	lambda_IR = 5.432937
	lambda_UV = 0.229202
	C = a0 + a1*rhobar + a2*Tbar + a3*lambdabar*lambdabar*Tbar + a4/(lambdabar*lambdabar) + a5/((lambdabar+lambda_UV)*(lambdabar-lambda_UV)) + a6/((lambdabar + lambda_IR)*(lambdabar - lambda_IR)) + a7*rhobar*rhobar
	C *= rhobar
	n = np.sqrt((2*C+1)/(1-C))
	return n

def sun_spectral(wavelength):
	return 1e16/(np.power(wavelength,5)*(np.exp(6.62607015e6*2.99792458/(wavelength*1.380649*(5250+273.15)))-1))

def modified_trace(wavelength, temp, center, r, N, n_theta, d_theta, min_intensity, max_ray, max_angle):
	n = water_refraction_index(temp, wavelength)
	disk = Disk(center, r, n)
	stack = []
	for u in np.linspace(0, 1, N, endpoint=False):
		stack.append(Ray([r*np.sqrt(u), 2*r], [0,-1], 1))
	ray_count = 0
	XYZ = colorspace.wavelength2XYZ(wavelength)*sun_spectral(wavelength)
	own_angle_XYZ = np.zeros((n_theta, 3))
	while stack and ray_count < max_ray:
		ray = stack.pop()
		ray_count += 1
		if ray is None:
			continue
		if isinstance(ray, Ray):
			if ray.intensity < min_intensity:
				continue
			direction = ray.direction
			t,intersection_point = disk.find_intersection(ray)
			if intersection_point is not None:
				points.append(np.concatenate((ray.origin, intersection_point,[ray.intensity])))
				normal = disk.get_normal(intersection_point)
				stack.extend(reflection_and_refraction(ray, intersection_point, normal, disk.refractive_index))
			else:
				points.append(np.concatenate((ray.origin, ray.origin+ray.direction,[ray.intensity])))
				if direction[1] > 0:
					angle = np.arccos(direction[1])
					if angle < max_angle:
						own_angle_XYZ[int(angle/d_theta)] += XYZ*ray.intensity*direction[1]
	return own_angle_XYZ

def rainbow(n_theta, max_theta, temp):
	angles, d_theta = np.linspace(0,max_theta,n_theta, retstep=True)

	angle_XYZ = np.zeros((n_theta, 3))
	angle_sRGB = np.zeros((n_theta, 3))

	num_wavelength = len(colorspace.data_wavelength)

	with Progress() as progress:
		# 创建一个进度条
		task = progress.add_task("[green]Simulating...", total=num_wavelength)

		# 使用 ThreadPoolExecutor 并行执行
		with ProcessPoolExecutor() as executor:
			futures = [executor.submit(modified_trace, wavelength, temp, center, r, N, n_theta, d_theta, min_intensity, max_ray, max_theta) for wavelength in colorspace.data_wavelength]

			for future in as_completed(futures):
				result = future.result()
				angle_XYZ += result
				progress.update(task, advance=1)

	for i,XYZ in enumerate(angle_XYZ):
		angle_sRGB[i] = colorspace.XYZ2RGB(XYZ)

	maxRGB = np.max(angle_sRGB)
	angle_sRGB/= maxRGB
	angle_sRGB = np.clip(angle_sRGB, 0, 1)
	colorspace.vectorized_gamma_correct(angle_sRGB)
	return angles,angle_sRGB

def bin_find(x, list, start):
	l,r = 0, len(list)-1

	if x < list[l]:
		return 0
	if x >= list[r]:
		return r

	if list[start] <= x < list[start + 1]:
		return start
	elif list[start] < x:
		l = start
	else:
		r = start

	while l <= r:
		mid = l + (r-l)//2
		if list[mid] <= x < list[mid + 1]:
			return mid
		elif list[mid] < x:
			l = mid + 1
		else:
			r = mid - 1
	return -1

def draw_column(i, w, h, radius, angle_sRGB):
		index = 0
		column = np.zeros((h, 3), dtype=np.int8)
		for j in range(h):
			r = np.linalg.norm(np.array([i,j]) - [w/2, 0])
			index = bin_find(r, radius, index)
			ratio = (r - radius[index])/(radius[index+1] - radius[index])
			column[j] = ((1-ratio)*angle_sRGB[index]+ratio*angle_sRGB[index+1])*255
		return i,column

def take_picture(angles, angle_sRGB, w, h, distance, filename="image.png"):
	image_array = np.zeros((w, h, 3), dtype=np.uint8)
	radius = distance * np.tan(angles)

	with Progress() as progress:
		# 创建一个进度条
		task2 = progress.add_task("[green]Rendering...", total=w)

		# 使用 ThreadPoolExecutor 并行执行
		with ProcessPoolExecutor() as executor:
			futures2 = [executor.submit(draw_column, i, w,h,radius,angle_sRGB) for i in range(w)]

			for future in as_completed(futures2):
				ind,column = future.result()
				image_array[ind] = column
				progress.update(task2, advance=1)
	image = Image.fromarray(image_array, "RGB")
	image.save(filename)

if __name__=="__main__":
	w = 7680 *2
	h = 4320 *2
	dis = 2400 *2
	max_angle = 1.1*np.arctan(np.sqrt(1+(h*h+w*w/4)/(dis*dis)))
	angles,sRGB=rainbow(10000, max_angle, 10)
	np.savez_compressed("saved.npz", a=angles, b=sRGB)
	# loaded = np.load("saved.npz")
	# angles, sRGB = loaded['a'], loaded['b']
	take_picture(angles,sRGB, w, h, dis, "image16k.png")