// Based on: https://bottosson.github.io/posts/oklab/
//
// OkLab first defines a transform from xyz, and multiplies by this matrix:
// [+0.8189330101, +0.3618667424, -0.1288597137]
// [+0.0329845436, +0.9293118715, +0.0361456387]
// [+0.0482003018, +0.2643662691, +0.633851707 ]
//
// Wikipedia lists this matrix for converting from linear sRGB to D65 CIE XYZ, so
// I'm considering it canonical:
// [+0.4124, +0.3576, +0.1805]
// [+0.2126, +0.7152, +0.0722]
// [+0.0193, +0.1192, +0.9505]
//
// Combined, we get:
// [+0.41217385032507, +0.5362974607032, +0.05146302925248]
// [+0.21187214048845, +0.6807476834212, +0.10740645682645]
// [+0.08831541121808, +0.2818663070584, +0.63026344660742]
//
// And its inverse:
// [+4.0767584135565013494237930518854,    -3.3072279873944731418619352916485,  +0.230721459944885632473018834049 ]
// [-1.2681810851624033989047813181437,    +2.6092932102856398573991970933594,  -0.3411121165477535569679616041822]
// [-0.0040984077180314400491332639337372, -0.70350366010241732765095902557887, +1.7068604529788013559365593912662]
//
// After a non-linear transformation (cube root), OkLab applies a second matrix:
// [+0.2104542553, +0.793617785,  -0.0040720468]
// [+1.9779984951, -2.428592205,  +0.4505937099]
// [+0.0259040371, +0.7827717662, -0.808675766 ]
//
// And its inverse:
// [+0.99999999845051981426207542502031, +0.39633779217376785682345989261573,  +0.21580375806075880342314146183004 ]
// [+1.0000000088817607767160752456705,  -0.10556134232365634941095687705472,  -0.063854174771705903405254198817796]
// [+1.0000000546724109177012928651534,  -0.089484182094965759689052745863391, -1.2914855378640917399489287529148  ]
//
// It should be noted that the numbers in the first column were supposed to be 1, and other implementations skip
// multiplying by that column altogether. I however am insane, and shall continue barging forward.

package oklab

import (
	"image/color"
	"math"
)

type Color struct {
	Lightness, ChromaA, ChromaB, A float64
}

func linearize(c uint32) float64 {
	l := float64(c) / 0xffff

	if l <= 0.039285714285714285714285714285714 {
		return l / 12.923210180787861094641554898407
	}

	return math.Pow((l+0.055)/1.055, 2.4)
}

func delinearize(l float64) uint32 {
	switch {
	case l <= 0:
		return 0
	case l <= 0.0030399346397784299969770436366690:
		return uint32(l*846922.57919793247683733430026710 + 0.5)
	case l >= 1:
		return 0xffff
	default:
		return uint32(69139.425*math.Pow(l, 1/2.4) - 3603.925)
	}
}

// FromNRGBA create an OkLab color from non-pre-multiplied RGBA.
func FromNRGBA(r, g, b, a uint32) Color {
	rLin, gLin, bLin := linearize(r), linearize(g), linearize(b)

	l := math.Cbrt(0.41217385032507*rLin + 0.5362974607032*gLin + 0.05146302925248*bLin)
	m := math.Cbrt(0.21187214048845*rLin + 0.6807476834212*gLin + 0.10740645682645*bLin)
	s := math.Cbrt(0.08831541121808*rLin + 0.2818663070584*gLin + 0.63026344660742*bLin)

	return Color{
		0.2104542553*l + 0.793617785*m - 0.0040720468*s,
		1.9779984951*l - 2.428592205*m + 0.4505937099*s,
		0.0259040371*l + 0.7827717662*m - 0.808675766*s,
		float64(a) / float64(0xffff),
	}
}

// FromRGBA create an OkLab color from pre-multiplied RGBA.
func FromRGBA(r, g, b, a uint32) Color {
	switch a {
	case 0xffff:
		// do nothing.
	case 0:
		// completely transparent, color information was lost.
		// pretend it was gray.
		r, g, b = 0x7fff, 0x7fff, 0x7fff
	default:
		// un-premultiply rgb.
		//
		// Note that I'm rounding up here, which is the opposite of what the NRGBA/NRGBA64 colors do,
		// which may be a bug as RGBA64->NRGBA64->RGBA64 is lossy.
		r = (r*0xffff + a - 1) / a
		g = (g*0xffff + a - 1) / a
		b = (b*0xffff + a - 1) / a
	}

	return FromNRGBA(r, g, b, a)
}

// NRGBAColor represent a color that can give us non pre-multiplied RGBA components.
//
// This isn't a standard interface, but we implement it and check for it regardless.
type NRGBAColor interface {
	color.Color
	NRGBA() (r, g, b, a uint32)
}

// FromColor converts an arbitrary color type to an OKLab [Color], with special handling
// for [color.NRGBA], [color.NRGBA64], and anything implementing the [NRGBAColor] interface
// to preserve transparent colors.
func FromColor(c color.Color) Color {
	switch c := c.(type) {
	case Color:
		return c

	// Special handling for [color.NRGBA] and [color.NRGBA64]
	case color.NRGBA:
		return FromNRGBA(uint32(c.R)*0x101, uint32(c.G)*0x101, uint32(c.B)*0x101, uint32(c.A)*0x101)

	case color.NRGBA64:
		return FromNRGBA(uint32(c.R), uint32(c.G), uint32(c.B), uint32(c.A))

	case NRGBAColor:
		return FromNRGBA(c.NRGBA())

	default:
		return FromRGBA(c.RGBA())
	}
}

func cube(v float64) float64 {
	return v * v * v
}

// NRGBA converts to non-premultiplied RGBA.
func (c Color) NRGBA() (r, g, b, a uint32) {
	l := cube(0.99999999845051981426207542502031*c.Lightness + 0.39633779217376785682345989261573*c.ChromaA + 0.21580375806075880342314146183004*c.ChromaB)
	m := cube(1.0000000088817607767160752456705*c.Lightness - 0.10556134232365634941095687705472*c.ChromaA - 0.063854174771705903405254198817796*c.ChromaB)
	s := cube(1.0000000546724109177012928651534*c.Lightness - 0.089484182094965759689052745863391*c.ChromaA - 1.2914855378640917399489287529148*c.ChromaB)

	r = delinearize(4.0767584135565013494237930518854*l - 3.3072279873944731418619352916485*m + 0.230721459944885632473018834049*s)
	g = delinearize(-1.2681810851624033989047813181437*l + 2.6092932102856398573991970933594*m - 0.3411121165477535569679616041822*s)
	b = delinearize(-0.0040984077180314400491332639337372*l - 0.70350366010241732765095902557887*m + 1.7068604529788013559365593912662*s)
	a = uint32(c.A*0xffff + 0.5)

	return
}

// RGBA converts to premultiplied RGBA.
func (c Color) RGBA() (r, g, b, a uint32) {
	r, g, b, a = c.NRGBA()

	r = r * a / 0xffff
	g = g * a / 0xffff
	b = b * a / 0xffff

	return
}

func sqr(a float64) float64 {
	return a * a
}

// Distance returns the perceptual distance between two colors.
//
// This controls for possible transparency. Two different but otherwise fully transparent colors would have a distance of zero.
func Distance(a, b Color) float64 {
	L1, a1, b1 := a.Lightness*a.A, a.ChromaA*a.A, a.ChromaB*a.A
	L2, a2, b2 := b.Lightness*b.A, b.ChromaA*b.A, b.ChromaB*b.A
	dL := L1 - L2
	da := a1 - a2
	db := b1 - b2
	dA := a.A - b.A
	return math.Sqrt(max(sqr(dL), sqr(dL+dA)) + max(sqr(da), sqr(da+dA)) + max(sqr(db), sqr(db+dA)))
}

func okLabModel(c color.Color) color.Color {
	return FromColor(c)
}

var (
	// A color model for converting arbitrary colors to OKLab.
	//
	// Wraps the [FromColor] function, returning a [color.Color] interface rather than the [Color] type.
	Model = color.ModelFunc(okLabModel)

	// Type assertions.
	_ NRGBAColor = Color{}
)