package filmgrain

import (
	"math"
	"strings"
)

//TODO: transfer function type?

func GetTransferFunction(kind string) *TransferFunction {
	switch strings.ToLower(kind) {
	case "bt470m":
		//gamma22
		return &TransferFunction{
			ToLinear: func(in float64) (linear float64) {
				return math.Pow(in, 2.2)
			},
			FromLinear: func(linear float64) (out float64) {
				return math.Pow(linear, 1/2.2)
			},
			MidTone: 0.18,
		}
	case "bt470bg":
		//gamma28
		return &TransferFunction{
			ToLinear: func(in float64) (linear float64) {
				return math.Pow(in, 2.8)
			},
			FromLinear: func(linear float64) (out float64) {
				return math.Pow(linear, 1/2.8)
			},
			MidTone: 0.18,
		}
	case "bt601", "bt709", "bt2020":
		//TODO: are bt601 and bt709 the same?
		//bt2020. same as bt709, just defined more precise. Can use the same for all
		const beta = 0.018053968510807
		const alpha = 1 + 5.5*beta
		return &TransferFunction{
			ToLinear: func(in float64) (linear float64) {
				if in < 0.081 {
					return in / 4.5
				} else {
					return math.Pow((in+(alpha-1))/alpha, 1/0.45)
				}
			},
			FromLinear: func(linear float64) (out float64) {
				if linear < beta {
					return 4.5 * linear
				} else {
					return alpha*math.Pow(linear, 0.45) - (alpha - 1)
				}
			},
			//TODO: check this for bt2020?
			MidTone: 0.18,
		}
	case "srgb":
		//sRGB
		return &TransferFunction{
			ToLinear: func(in float64) (linear float64) {
				if in <= 0.04045 {
					return in / 12.92
				} else {
					return math.Pow((in+0.055)/1.055, 2.4)
				}
			},
			FromLinear: func(linear float64) (out float64) {
				if linear <= 0.0031308 {
					return 12.92 * linear
				} else {
					return 1.055*math.Pow(linear, 1/2.4) - 0.055
				}
			},
			MidTone: 0.18,
		}
	case "smpte2084", "pq":
		//pq
		const PqM1 = 2610. / 16384
		const PqM2 = 128 * 2523. / 4096
		const PqC1 = 3424. / 4096
		const PqC2 = 32 * 2413. / 4096
		const PqC3 = 32 * 2392. / 4096
		return &TransferFunction{
			ToLinear: func(in float64) (linear float64) {
				pq_pow_inv_m2 := math.Pow(in, 1./PqM2)
				return math.Pow(max(0, pq_pow_inv_m2-PqC1)/(PqC2-PqC3*pq_pow_inv_m2),
					1./PqM1)
			},
			FromLinear: func(linear float64) (out float64) {
				linear_pow_m1 := math.Pow(linear, PqM1)
				return math.Pow((PqC1+PqC2*linear_pow_m1)/(1+PqC3*linear_pow_m1),
					PqM2)
			},
			MidTone: 26. / 1000,
		}
	case "hlg":
		//hlg
		// Note: it is perhaps debatable whether “linear” for HLG should be scene light
		// or display light. Here, it is implemented in terms of display light assuming
		// a nominal peak display luminance of 1000 cd/m², hence the system γ of 1.2. To
		// make it scene light instead, the OOTF (math.Pow(x, 1.2)) and its inverse should
		// be removed from the functions below, and the TransferFunction.MidTone should be replaced
		// with math.Pow(26. / 1000, 1 / 1.2).
		const HlgA = 0.17883277
		const HlgB = 0.28466892
		const HlgC = 0.55991073
		return &TransferFunction{
			ToLinear: func(in float64) (linear float64) {
				// EOTF = OOTF ∘ OETF⁻¹
				if in <= 0.5 {
					linear = in * in / 3
				} else {
					linear = (math.Exp((in-HlgC)/HlgA) + HlgB) / 12
				}
				return math.Pow(linear, 1.2)
			},
			FromLinear: func(linear float64) (out float64) {
				// EOTF⁻¹ = OETF ∘ OOTF⁻¹
				linear = math.Pow(linear, 1./1.2)
				if linear <= (1. / 12) {
					return math.Sqrt(3 * linear)
				} else {
					return HlgA*math.Log(12*linear-HlgB) + HlgC
				}
			},
			MidTone: 26. / 1000,
		}
	}
	return nil
}

type TransferFunction struct {
	ToLinear   func(in float64) (linear float64)
	FromLinear func(linear float64) (out float64)

	// MidTone In linear output light. This would typically be 0.18 for SDR (this matches
	// the definition of Standard Output Sensitivity from ISO 12232:2019), but in
	// HDR, we certainly do not want to consider 18% of the maximum output a
	// “mid-tone”, as it would be e.g. 1800 cd/m² for SMPTE ST 2084 (PQ).
	MidTone float64
}