Create a new version of paste:

##CC0 Kaelygon 2025
import math
import numpy as np
from PIL import Image
from dataclasses import dataclass, field
from typing import List, Tuple, Optional
import random

#3D vector operators
def sub_vec3(vec_a,vec_b):
	return [ vec_a[0] - vec_b[0], vec_a[1] - vec_b[1], vec_a[2] - vec_b[2] ]
def add_vec3(vec_a,vec_b):
	return [ vec_a[0] + vec_b[0], vec_a[1] + vec_b[1], vec_a[2] + vec_b[2] ]
def mul_vec3(vec_a,vec_b):
	return [ vec_a[0] * vec_b[0], vec_a[1] * vec_b[1], vec_a[2] * vec_b[2] ]
def div_vec3(vec_a,vec_b):
	return [ vec_a[0] / vec_b[0], vec_a[1] / vec_b[1], vec_a[2] / vec_b[2] ]


#special
def sign_vec3(vec_a):
	return [ 1 if c>=0 else -1 for c in vec_a ]
def lerp_vec3(vec_a,vec_b,a):
   return add_vec3( mul_vec3(vec_a, sub_vec3([1.0]*3,a) ), mul_vec3(vec_b,a) )
def pow_vec3(vec_a,vec_b):
	return [ vec_a[0] ** vec_b[0], vec_a[1] ** vec_b[1], vec_a[2] ** vec_b[2] ]
def spow_vec3(vec_a,vec_b):
	sign = sign_vec3(vec_a)
	return [ sign[0]*(sign[0]*vec_a[0]) ** vec_b[0], sign[1]*(sign[1]*vec_a[1]) ** vec_b[1], sign[2]*(sign[2]*vec_a[2]) ** vec_b[2] ]
def clip_vec3(vec_a,low,hi):
	return [ min(max(vec_a[0],low[0]),hi[0]), min(max(vec_a[1],low[1]),hi[1]), min(max(vec_a[2],low[2]),hi[2]) ]
def lessThan_vec3(vec_a,vec_b):
	return [ vec_a[0]<vec_b[0], vec_a[1]<vec_b[1], vec_a[2]<vec_b[2] ]

#Vector operators
def dot_vec3(a, b):
	product = 0.0
	for i in range(3):
		product+= a[i] * b[i]
	return product

def cross_vec3(a, b):
	c = [a[1]*b[2] - a[2]*b[1],
			a[2]*b[0] - a[0]*b[2],
			a[0]*b[1] - a[1]*b[0]]

	return c

def length_vec3(vec):
	return math.sqrt(vec[0]*vec[0] + vec[1]*vec[1] + vec[2]*vec[2])

def norm_vec3(vec,eps=0.0):
	l = length_vec3(vec)
	if l == eps:
		return [1.0, 0.0, 0.0]
	return [vec[0]/l, vec[1]/l, vec[2]/l]


#other math tools
def math_spow(a,b):
	return -((-a)**b) if a<0 else a**b

### Color conversion
def linearToSrgb(linRGB):
	cutoff = lessThan_vec3(linRGB, [0.0031308]*3)
	gammaAdj = spow_vec3(linRGB, [1.0/2.4]*3 )
	higher = sub_vec3( mul_vec3( [1.055]*3 , gammaAdj ), [0.055]*3 )
	lower = mul_vec3( linRGB, [12.92]*3 )

	return lerp_vec3(higher, lower, cutoff)

def srgbToLinear(sRGB):
	cutoff = lessThan_vec3(sRGB, [0.04045]*3)
	higher = spow_vec3( add_vec3(sRGB, div_vec3([0.055]*3, [1.055]*3) ), [2.4]*3 )
	lower = div_vec3( sRGB, [12.92]*3 )

	return lerp_vec3(higher, lower, cutoff)

def linearToOklab(RGB):
	r,g,b = RGB
	l = 0.4122214708 * r + 0.5363325363 * g + 0.0514459929 * b
	m = 0.2119034982 * r + 0.6806995451 * g + 0.1073969566 * b
	s = 0.0883024619 * r + 0.2817188376 * g + 0.6299787005 * b

	l_ = math_spow(l, 1.0/3.0)
	m_ = math_spow(m, 1.0/3.0)
	s_ = math_spow(s, 1.0/3.0)

	return [
		0.2104542553*l_ + 0.7936177850*m_ - 0.0040720468*s_,
		1.9779984951*l_ - 2.4285922050*m_ + 0.4505937099*s_,
		0.0259040371*l_ + 0.7827717662*m_ - 0.8086757660*s_,
	]

def oklabToLinear(ok):
	L,a,b = ok
	l_ = L + 0.3963377774 * a + 0.2158037573 * b
	m_ = L - 0.1055613458 * a - 0.0638541728 * b
	s_ = L - 0.0894841775 * a - 1.2914855480 * b

	l = l_*l_*l_
	m = m_*m_*m_
	s = s_*s_*s_

	return [
		+4.0767416621 * l - 3.3077115913 * m + 0.2309699292 * s,
		-1.2684380046 * l + 2.6097574011 * m - 0.3413193965 * s,
		-0.0041960863 * l - 0.7034186147 * m + 1.7076147010 * s,
  	]

def srgbToOklab(rgb):
	linRGB = srgbToLinear(rgb)
	oklab = linearToOklab(linRGB)
	return oklab

def oklabToSrgb(ok):
	linRGB = oklabToLinear(ok)
	sRGB = linearToSrgb(linRGB)
	return sRGB


### Color tools ###

def inOklabGamut(ok, eps=1e-7):
	r_lin, g_lin, b_lin = oklabToLinear(ok)

	return (r_lin >= -eps) and (r_lin <= 1+eps) and \
			 (g_lin >= -eps) and (g_lin <= 1+eps) and \
			 (b_lin >= -eps) and (b_lin <= 1+eps)

def isOkSrgbGray(col, threshold = 1.0/255.0):
	r,g,b = oklabToSrgb(col)
	if( abs(r-g) < threshold and 
		abs(g-b) < threshold
	):
		return True
	return False

def srgbToHex(rgb):
	r = int(round(rgb[0] * 255.0))
	g = int(round(rgb[1] * 255.0))
	b = int(round(rgb[2] * 255.0))
	return "#{:02x}{:02x}{:02x}".format(r, g, b)

def oklabGamutVolume(resolution=50):
	
	valid_count = 0
	total_count = resolution ** 3

	for L in range(0,resolution):
		for a in range(0,resolution):
				for b in range(0,resolution):
					if inOklabGamut([
						float(L)/resolution, 
						float(a)/resolution-0.5,
						float(b)/resolution-0.5
					]):
						valid_count += 1
	
	return valid_count / total_count

#Calc color values
def calcOkChroma(col):
	return math.sqrt( col[1]*col[1] + col[2]*col[2] )

def calcSrgbLum(col):
	return srgbToOklab(col)[0]

def separateGrays(ok_array):
	gray_colors = []
	hued_colors = []

	for col in ok_array:
		chroma = calcOkChroma(col)
		if isOkSrgbGray(col):
			gray_colors.append(col)
		else:
			hued_colors.append(col)
	
	return gray_colors, hued_colors


class PointGrid:
	"""
		Store point cloud as a search PointGrid.grid and a 1D PointGrid.cloud
	"""

	def __init__(self, point_radius):
		self.cloud = []
		self.grid = {}
		self.length = 0
		self.cell_size = point_radius

	def key(self, p):
		gx = int(p[0] // self.cell_size)
		gy = int(p[1] // self.cell_size)
		gz = int(p[2] // self.cell_size)
		return (gx,gy,gz)

	def insert(self, p):
		self.length+=1
		k = self.key(p)
		if k not in self.grid:
			self.grid[k] = []
		self.grid[k].append(p)
		self.cloud.append(p)

	def findNearest(self, test_point, point_radius, neighbor_margin=0):
		p0_info = [None,None,None]
		closest=float('inf') 

		neighbor_count=0
		seen = set()
		
		kx, ky, kz = self.key(test_point)
		for dx in (-1, 0, 1):
				for dy in (-1, 0, 1):
					for dz in (-1, 0, 1):
						nk = (kx + dx, ky + dy, kz + dz)

						for p in self.grid.get(nk, ()):
							if p == test_point:
									continue
							pid = id(p) 
							if pid in seen:
									continue
							seen.add(pid)

							delta = sub_vec3(test_point, p)
							dist = length_vec3(delta)

							if dist - neighbor_margin <= point_radius:
									neighbor_count += 1

							if dist < closest:
									closest = dist
									p0_info = [p, delta, dist]
			
		
		return p0_info[0],p0_info[1],p0_info[2],neighbor_count


### Palette generator ###
@dataclass
class PalettePreset:
	reserve_transparent: bool = True

	max_colors: int = None

	gray_count: int = None	#Grayscale color count
	hue_count:  int = None	#Split Hues in this many buckets

	min_sat: float = 0.0 	#min/max ranges are percentages
	max_sat: float = 1.0
	min_lum: float = 0.0
	max_lum: float = 1.0
	packing_fac: float = 1.0 #Packing efficiency
	max_attempts: int = 1024 #After this many max_attempts per point, pointSampler will give up
	
	seed: int = 0 #Seed for pointSampler


class PaletteGenerator:
	"""
		Generate palette where the colors are perceptually evenly spaced out in OKLab colorspace
	"""

	def __init__(self, 
					preset: Optional[PalettePreset] = None, *,
					
					reserve_transparent: Optional[bool] = None,
					hue_count: Optional[int] = None,
					gray_count: Optional[int] = None,
					min_sat: Optional[float] = None,
					max_sat: Optional[float] = None,
					min_lum: Optional[float] = None,
					max_lum: Optional[float] = None,
	):
		self.point_grid = None
		self.point_radius = None

		if preset is None:
			preset = PalettePreset()
	
		self.p = PalettePreset(**{k: getattr(preset, k) for k in preset.__dataclass_fields__})

		if self.p.max_colors:
			self.p.max_colors -= self.p.reserve_transparent
		else:
			self.p.max_colors = 0

	def seed(self, input_seed=None):
		if input_seed:
			random.seed(input_seed)
		if self.p.seed:
			random.seed(self.p.seed)
		else:
			random.random()
 
	def pointSampler(self):
		def rand_vec3():
			return norm_vec3( [random.random()-0.5, random.random()-0.5, random.random()-0.5] )
		def getRandOklab():
			return [random.random(), random.random()-0.5, random.random()-0.5]
		def rand_cloud_point():
			if len(self.point_grid.cloud) and random.random()>0.5:
				return self.point_grid.cloud[ int(random.random()*len(self.point_grid.cloud)) ]
			return getRandOklab()

		origin = [0.5, 0.0, 0.0]
		space_size = math.sqrt(1.0**2 + 1.0**2 + 1.0**2)
		margin = space_size/10000.0
		push_scalar = [self.point_radius]*3

		push_fails = 0
		max_push_fails = 10

		p0 = rand_cloud_point()
		skip_rand = False
		new_point = None

		stale_counter = 0
		attempt_counter = 0

		while len(self.point_grid.cloud) < self.p.max_colors:
			attempt_counter+=1
			stale_counter+=1
			if stale_counter > self.p.max_attempts:
				print("Reached max_attempts to find a new point.")
				break
	
			if not skip_rand:
				#Move to random direction
				rand_vec = rand_vec3()
				move_vec = mul_vec3(rand_vec, [self.point_radius]*3 )
				new_point = add_vec3(p0, move_vec)
			skip_rand = False
	
			if not inOklabGamut(new_point):
				#clip to srgb gamut
				linRGB = oklabToLinear(new_point)
				linRGB = clip_vec3(linRGB,[0.0]*3,[1.0]*3)
				new_point = linearToOklab(linRGB)
	
			col_point, pos_vec, dist, col_count = self.point_grid.findNearest(new_point, self.point_radius, margin)
 
			if col_point:
				delta = dist-self.point_radius
				if abs(delta) > margin:
					#Too close or far: Set 1 radius apart from nearest point
					push_fails+=1
					if push_fails > max_push_fails:
						push_fails = 0
						continue
   
					normal = None
					if push_fails <= max_push_fails//2:
						normal = norm_vec3(pos_vec)
					else:
						normal = rand_vec3() #rand direction
      
					new_point = add_vec3(col_point, mul_vec3(normal, push_scalar)) 
					skip_rand = True
					continue

			self.point_grid.insert(new_point)
			p0=new_point
			stale_counter=0
			push_fails = 0
	
		print("Loop count "+str(attempt_counter))


	def generateGrays(self):
		if self.p.gray_count == None:
			self.p.gray_count = int(round(1.0/(self.point_radius)))
		if self.p.gray_count:
			#Use minimum starting luminosity that second darkest black isn't so close to 0 
			darkest_black = calcSrgbLum([0.499/255,0.499/255,0.499/255]) 
			for i in range(0,self.p.gray_count):
				lum = float(i)/(self.p.gray_count-1)
				scale = ( (self.p.gray_count-i-1)/(self.p.gray_count-1) )
				lum+= darkest_black*scale #Fade that brightest remains 1.0
				new_point = [lum,0,0]
				self.point_grid.insert(new_point)

	def applyColorLimits(self):
		apply_luminosity = self.p.max_lum!=1.0 or self.p.min_lum!=0.0
		apply_saturation = self.p.max_sat!=1.0 or self.p.min_sat!=0.0

		filtered_cloud = []
		for col in self.point_grid.cloud:
			if apply_luminosity:
				lum_width = self.p.max_lum - self.p.min_lum
				col[0] = col[0]*lum_width + self.p.min_lum
		
			if apply_saturation and not isOkSrgbGray(col):
				sat_width = self.p.max_sat - self.p.min_sat
				chroma = calcOkChroma(col)
	
				max_chroma = math.sqrt(0.5**2+0.5**2) 
				rel_sat = chroma / max_chroma
				scaled_sat = (rel_sat * sat_width + self.p.min_sat) * max_chroma

				col_vec = [col[1], col[2]] #Vector a,b 
				col_vec = [col_vec[0]/chroma, col_vec[1]/chroma] #Normalize
				col_vec = [col_vec[0]*scaled_sat, col_vec[1]*scaled_sat] #Scale
				col = [col[0], col_vec[0], col_vec[1]]
				
			filtered_cloud.append(col)
	
		self.point_grid.cloud = filtered_cloud

	def populatePointCloud(self):	
		self.seed()
	
		space_volume = 0.054197416 # print (str(oklabGamutVolume(500))) pre computed
		point_count = self.p.max_colors if self.p.max_colors else self.p.tone_count * self.p.hue_count + self.p.gray_count
		unit_volume = space_volume/(point_count)
		self.point_radius = unit_volume**(1.0/3.0) * self.p.packing_fac 
		print("Using self.point_radius "+str(round(self.point_radius,4)))

		self.point_grid = PointGrid(self.point_radius)

		self.generateGrays()
		self.pointSampler()

		self.applyColorLimits()

		return self.point_grid.cloud
	

	### Palette processing ###

	def paletteToHex(self):
		hex_list = []
		if self.p.reserve_transparent:
				hex_list.append("#00000000")
	
		for col in self.point_grid.cloud:
			srgb_col = oklabToSrgb(col)
			hex_list.append(srgbToHex(srgb_col))
		
		return hex_list

	def paletteToImg(self, hex_list: List[str], filename: str = "palette.png"):
		rgba = []
		if self.p.reserve_transparent:
			rgba.append((0, 0, 0, 0))
	
		for col in self.point_grid.cloud:
			r,g,b = oklabToSrgb(col)
			r = min( max( int(round(r * 255.0)), 0 ), 255) 
			g = min( max( int(round(g * 255.0)), 0 ), 255) 
			b = min( max( int(round(b * 255.0)), 0 ), 255) 
			rgba.append((r, g, b, 255))
		
		arr = np.array([rgba], dtype=np.uint8)
		img = Image.fromarray(arr, mode="RGBA")
		img.save(filename)
		return img

	def sortByLum(self, ok_array):
		return sorted(ok_array, key=lambda x: x[0])

	def sortPalette(self):
		gray_colors, hued_colors = separateGrays(self.point_grid.cloud)
	
		#Place hues in same buckets
		hue_buckets = [[] for _ in range(self.p.hue_count)]
		hue_bucket_width = 2*math.pi * (1.0/self.p.hue_count)

		for col in hued_colors:
			col_hue = math.atan2(col[2], col[1]) + 2* math.pi
			bucket_index = int(col_hue/hue_bucket_width) % self.p.hue_count
			hue_buckets[bucket_index].append(col)

		#Sort hue buckets by luminance
		sorted_hue_buckets = []
		for bucket in hue_buckets:
			sorted_bucket = self.sortByLum(bucket)
			sorted_hue_buckets.append(sorted_bucket)

		#combine colors into single array
		sorted_colors = []

		sorted_grays = self.sortByLum(gray_colors)
		for col in sorted_grays:
			sorted_colors.append(col)

		for bucket in sorted_hue_buckets:
			for col in bucket:
				sorted_colors.append(col)
	
		self.point_grid.cloud = sorted_colors

	#EOF PaletteGenerator


### palette analysis ###

#p0 has type (float3)[L,a,b]
#pair_list has type (float,[],[])[dist, p0, p1]
def printColorPairStats(pair_list, print_count, listName="", calc_error=False):
	def oklabToHex(col):
		return srgbToHex(oklabToSrgb(col))
	
	if len(pair_list) < 2:
		print("Not enough pairs for stats!")
		return
	
	precision = 4

	print(listName+" Closest pairs")
	for pair in pair_list[:print_count]:
		print(str(round(pair[0],precision))+" "+oklabToHex(pair[1])+" "+oklabToHex(pair[2]))

	#print only unseen values
	far_start = max(print_count, len(pair_list)-print_count) 
	if far_start < len(pair_list):
		print(listName+" Farthest pairs")
	for pair in pair_list[far_start:]:
		print(str(round(pair[0],precision))+" "+oklabToHex(pair[1])+" "+oklabToHex(pair[2]))

	#Average
	sumDist = 0.0
	for pair in pair_list:
		sumDist+=pair[0]
	avgDist = sumDist / len(pair_list)

	#Median
	medianDist=0.0
	medianIndex=len(pair_list)//2
	if len(pair_list)%2==0:
		a = medianIndex
		b = max(medianIndex-1,0)
		medianDist = (pair_list[a][0] + pair_list[b][0]) / 2.0
	else:
		medianDist = pair_list[medianIndex][0]

	print(listName+" Avg pair distance: "+str(round(avgDist,precision)))
	print(listName+" Median pair distance: "+str(round(medianDist,precision)))
	print("")

	
	#Compare biggest gap to avg gap
	if calc_error == True:

		hued_pairs = [
			p for p in pair_list
			if not isOkSrgbGray(p[1]) and not isOkSrgbGray(p[2])
		]
		hue_pair_count = len(hued_pairs)
        
		#Average hued only
		sumDist = 0.0
		for pair in hued_pairs:
			sumDist+=pair[0]
		avgHueDist = sumDist / hue_pair_count if hue_pair_count != 0 else 0.0001
	

		#All colors gaps
		all_smallest_gap = pair_list[0][0]
		all_largest_gap = pair_list[-1][0]
	

		#Hued colors gaps
		hued_smallest_gap = hued_pairs[0][0] if hued_pairs else None
		hued_largest_gap = hued_pairs[-1][0] if hued_pairs else None

		allError = abs(1.0 - all_largest_gap/avgDist)
		print("Biggest_gap  to avg_gap delta "+str(round(100*allError,precision))+" %")

		allError = abs(1.0 - all_smallest_gap/avgDist)
		print("Smallest_gap to avg_gap delta "+str(round(100*allError,precision))+" %")

		if hued_largest_gap:
			huedError = abs(1.0 - hued_largest_gap/avgHueDist)
			print("Hued biggest_gap  to avg_gap delta "+str(round(100*huedError,precision))+" %")
		if hued_largest_gap:
			huedError = abs(1.0 - hued_smallest_gap/avgHueDist)
			print("Hued smallest_gap to avg_gap delta "+str(round(100*huedError,precision))+" %")

	return

# Input [[L,a,b], ...]
# Find closest point for every point
# Returns list of point pairs [[dist, p0, p1], ...] sorted by distance
def findClosestPairs_grid(point_grid):
	if len(point_grid.cloud) < 2:
		return []

	dist_pair_array = []

	for p in point_grid.cloud:
		neighbor, delta, dist, neighbor_count = point_grid.findNearest(p, point_grid.cell_size)
		if neighbor is None:
				continue

		if p in neighbor:
				continue

		dist_pair_array.append([dist, p, neighbor])

	dist_pair_array.sort(key=lambda x: x[0])
	return dist_pair_array

def oklabGapStats(point_grid, print_count):
	fullPairArr = findClosestPairs_grid(point_grid)
	printColorPairStats(fullPairArr,print_count,"Full",1)


palette_preset_list = {
	'palTest': PalettePreset(
		reserve_transparent=True,
		gray_count	= None, 
		max_colors	=64, 
		hue_count	=12,
		min_sat		=0.0, 
		max_sat		=1.0, 
		min_lum		=0.0, 
		max_lum		=1.0,
		packing_fac	=1.0,
		max_attempts=1024*64,
		seed=0
	),
}

def printHexList(hex_list, palette_name=""):
	out_string = palette_name + " = ["
	for string in hex_list:
		out_string += "\""+string+"\","
	out_string+="]"
	print(out_string + "\n")

def run_generatePalette():

	palette_name = 'palTest'
	active_preset = palette_preset_list[palette_name]

	palette = PaletteGenerator(preset=active_preset)
	palette.populatePointCloud()

	palette.sortPalette()

	palette_file = 'palette.png'
	palette.paletteToImg(palette_file)

	oklabGapStats(palette.point_grid,3)

	hex_string = palette.paletteToHex()	
	printHexList(hex_string)

	print("Generated "+str(len(palette.point_grid.cloud) + active_preset.reserve_transparent)+" colors to ./"+palette_file)


if __name__ == '__main__':
	linRGB = [1.0,1.0,1.0]
	ok = linearToOklab(linRGB)
	new_srgb = oklabToLinear(ok)
	if not inOklabGamut(ok):
		print("uh oh")
		exit(0)
	
	run_generatePalette()