import math
from dataclasses import dataclass
from itertools import combinations

import numpy as np
from PIL import Image
from scipy.cluster import vq

# https://en.wikipedia.org/wiki/SRGB#Transformation
linearize_srgb = np.vectorize(
  lambda v: (v / 12.92) if v <= 0.04045 else (((v + 0.055) / 1.055) ** 2.4)
)
delinearize_lrgb = np.vectorize(
  lambda v: (v * 12.92) if v <= 0.0031308 else ((v ** (1 / 2.4)) * 1.055 - 0.055)
)
# https://mina86.com/2019/srgb-xyz-matrix/
RGB_TO_XYZ = np.array([
  [33786752 / 81924984, 29295110 / 81924984, 14783675 / 81924984],
  [8710647 / 40962492, 29295110 / 40962492, 2956735 / 40962492],
  [4751262 / 245774952, 29295110 / 245774952, 233582065 / 245774952],
])
XYZ_TO_RGB = [
  [4277208 / 1319795, -2028932 / 1319795, -658032 / 1319795],
  [-70985202 / 73237775, 137391598 / 73237775, 3043398 / 73237775],
  [164508 / 2956735, -603196 / 2956735, 3125652 / 2956735],
]

# https://bottosson.github.io/posts/oklab/
XYZ_TO_LMS = np.array([
  [0.8189330101, 0.3618667424, -0.1288597137],
  [0.0329845436, 0.9293118715, 0.0361456387],
  [0.0482003018, 0.2643662691, 0.6338517070],
])
RGB_TO_LMS = XYZ_TO_LMS @ RGB_TO_XYZ
LMS_TO_RGB = np.linalg.inv(RGB_TO_LMS)
LMS_TO_OKLAB = np.array([
  [0.2104542553, 0.7936177850, -0.0040720468],
  [1.9779984951, -2.4285922050, 0.4505937099],
  [0.0259040371, 0.7827717662, -0.8086757660],
])
OKLAB_TO_LMS = np.linalg.inv(LMS_TO_OKLAB)


def oklab2hex(pixel: np.array) -> str:
  # no need for a vectorized version, this is only for providing the mean hex
  return "#" + "".join(f"{int(x * 255):02X}" for x in delinearize_lrgb(((pixel @ OKLAB_TO_LMS.T) ** 3) @ LMS_TO_RGB.T))


def srgb2oklab(pixels: np.array) -> np.array:
  return (linearize_srgb(pixels / 255) @ RGB_TO_LMS.T) ** (1 / 3) @ LMS_TO_OKLAB.T


@dataclass
class Stats:
  size: int
  proportion: int
  variance: float
  stddev: float
  hex: str
  Lbar: float
  abar: float
  bbar: float
  Cbar: float
  hbar: float
  Lhat: float
  ahat: float
  bhat: float


def calc_statistics(pixels: np.array, total_size=None) -> Stats:
  # mean pixel of the image, (L-bar, a-bar, b-bar)
  mean = pixels.mean(axis=0)
  # square each component
  squared = pixels ** 2
  # Euclidean norm squared by summing squared components
  sqnorms = squared.sum(axis=1)
  # mean pixel of normalized image, (L-hat, a-hat, b-hat)
  tilt = (pixels / np.sqrt(sqnorms)[:, np.newaxis]).mean(axis=0)
  # variance = mean(||p||^2) - ||mean(p)||^2
  variance = sqnorms.mean(axis=0) - sum(mean ** 2)
  # chroma^2 = a^2 + b^2
  chroma = np.sqrt(squared[:, 1:].sum(axis=1))
  # hue = atan2(b, a), but we need a circular mean
  # https://en.wikipedia.org/wiki/Circular_mean#Definition
  # cos(atan2(b, a)) = a / sqrt(a^2 + b^2) = a / chroma
  # sin(atan2(b, a)) = b / sqrt(a^2 + b^2) = b / chroma
  hue = math.atan2(*(pixels[:, [2, 1]] / chroma[:, np.newaxis]).mean(axis=0))
  return Stats(
    size=len(pixels),
    proportion=1 if total_size is None else (len(pixels) / total_size),
    variance=variance,
    stddev=math.sqrt(variance),
    hex=oklab2hex(mean),
    Lbar=mean[0],
    abar=mean[1],
    bbar=mean[2],
    Cbar=chroma.mean(axis=0),
    hbar=(hue * 180 / math.pi) % 360,
    Lhat=tilt[0],
    ahat=tilt[1],
    bhat=tilt[2],
  )


def find_clusters(pixels: np.array, cluster_attempts=5, seed=0) -> list[Stats]:
  means, labels = max(
    (
      # Try k = 2, 3, and 4, and try a few times for each
      vq.kmeans2(pixels.astype(float), k, minit="++", seed=seed + i)
      for k in (2, 3, 4)
      for i in range(cluster_attempts)
    ),
    key=lambda c:
      # Evaluate clustering by seeing the average distance in the ab-plane
      # between the centers. Maximizing this means the clusters are highly
      # distinct, which gives a sense of which k was best.
      (np.array([m1 - m2 for m1, m2 in combinations(c[0][:, 1:], 2)]) ** 2)
        .sum(axis=1)
        .mean(axis=0)
  )
  return [calc_statistics(pixels[labels == i], len(pixels)) for i in range(len(means))]


def get_pixels(img: Image.Image) -> np.array:
  rgb = []
  for fr in range(getattr(img, "n_frames", 1)):
    img.seek(fr)
    rgb += [
      [r, g, b]
      for r, g, b, a in img.convert("RGBA").getdata()
      if a > 0 and (r, g, b) != (0, 0, 0)
    ]
  return srgb2oklab(np.array(rgb))


if __name__ == "__main__":
  print("TODO")