{ pkgs-unstable, ... }:
{
  home.packages = with pkgs-unstable; [
    (writers.writePython3Bin "change_colors_json" {
      libraries = with python3Packages; [ numpy pillow scikit-learn ];
      flakeIgnore = [ "E302" "E305" "E226" "E501" ];
    } /*python */ ''
from colorsys import hls_to_rgb, rgb_to_hls
import json
import sys
from typing import Literal, cast
from numpy.typing import NDArray
from sklearn.cluster import KMeans
import numpy as np
from PIL import Image

def fc(c: int) -> str:
    assert c < 256
    s = str(hex(c))[2:]
    if c < 16:
        return "0" + s
    elif len(s) == 1:
        return s + s
    else:
        return s

class Color(object):
    def __init__(self, rgb: tuple[int, ...], frequency: float):
        assert len(rgb) == 3, "RGB values must be a tuple of length 3"
        self.rgb = cast(tuple[int, int, int], rgb)
        self.freq: float = frequency

    def __lt__(self, other: "Color") -> bool:
        return self.freq < other.freq

    @property
    def hls(self) -> tuple[float, float, float]:
        return rgb_to_hls(r=self.rgb[0] / 255, g=self.rgb[1] / 255, b=self.rgb[2] / 255)

    @property
    def luminance(self) -> float:
        return np.dot(np.array([0.2126, 0.7152, 0.0722]), self.rgb)

def k_means_extraction(arr: NDArray[float], height: int, width: int, palette_size: int) -> list[Color]:
    arr = np.reshape(arr, (width * height, -1))
    model = KMeans(n_clusters=palette_size, n_init="auto", init="k-means++", random_state=2024)
    labels = model.fit_predict(arr)
    palette = np.array(model.cluster_centers_, dtype=int)
    color_count = np.bincount(labels)
    color_frequency = color_count / float(np.sum(color_count))
    colors = []
    for color, freq in zip(palette, color_frequency):
        colors.append(Color(color, freq))
    return colors


class Palette:
    def __init__(self, colors: list[Color]):
        self.colors = colors
        self.frequencies = [c.freq for c in colors]

    def __getitem__(self, item: int) -> Color:
        return self.colors[item]

    def __len__(self) -> int:
        return self.number_of_colors

def ensure_color(c: Color, alter_sat: bool) -> list[int]:
    hue, lum, sat = c.hls
    if alter_sat:
        new_sat = min((sat**0.5) + 0.4, 1)
    else:
        new_sat = sat
    new_lum = max(lum, 0.5)
    r, g, b = hls_to_rgb(h=hue, l=new_lum, s=new_sat)
    return [int(r*255), int(g*255), int(b*255)]

def list_to_hex(ilist: list[int]) -> str:
    return f"#{fc(ilist[0])}{fc(ilist[1])}{fc(ilist[2])}"

def alter_hue(ilist: list[int], hue: int) -> list[int]:
    assert hue >= 0 and hue <= 360
    r, g, b = ilist
    h, l, s = rgb_to_hls((r/255), (g/255), (b/255))
    new_hue = (((h*360) + hue) % 360) / 360
    r, g, b = hls_to_rgb(h=new_hue, l=l, s=s)
    return [int(r*255), int(g*255), int(b*255)]

def alter_l(ilist: list[int], l_in_1_0: float) -> list[int]:
    assert l_in_1_0 >= 0 and l_in_1_0 <= 1
    r, g, b = ilist
    h, _, s = rgb_to_hls((r/255), (g/255), (b/255))
    r, g, b = hls_to_rgb(h=h, l=l_in_1_0, s=s)
    return [int(r*255), int(g*255), int(b*255)]

def extract_colors(
    image: str,
    palette_size: int = 5,
    resize: bool = True,
    sort_mode: Literal["luminance", "frequency"] | None = None,
) -> Palette:

    img = Image.open(image).convert("RGB")

    # open the image
    img = img.resize((256, 256))
    width, height = img.size
    arr = np.asarray(img)

    colors = k_means_extraction(arr, height, width, palette_size)

    if sort_mode == "luminance":
        colors.sort(key=lambda c: c.luminance, reverse=False)
    else:
        colors.sort(reverse=True)

    return Palette(colors)
if __name__ == "__main__":
    img = sys.argv[1]

    palette = extract_colors(image=img, palette_size=3)

    accent = ensure_color(c=palette[0], alter_sat=False)
    secondary = ensure_color(c=palette[1], alter_sat=True)
    tertiary = ensure_color(c=palette[2], alter_sat=False)
    weird = alter_hue(ilist=secondary, hue=180)
    special = alter_hue(ilist=accent, hue=180)
    foreground = alter_l(accent, 0.9)
    background = alter_l(accent, 0.1)

    d = {
        "base": {
            "foreground": list_to_hex(foreground),
            "background": list_to_hex(background)
        },
        "to_alter": {
            "accent": list_to_hex(accent),
            "secondary": list_to_hex(secondary),
            "tertiary": list_to_hex(tertiary),
            "special": list_to_hex(special),
            "weird": list_to_hex(weird)
        }
    }

    with open("/home/nx2/nix-dots/flake-modules/colors.json", "w") as f:
        f.write(json.dumps(d, indent=4))
'')
  ];
}