dotfiles/home-modules/wallpaper-to-colors.nix

{ pkgs, ... }@all: with all;
{
  home.packages = [
    (pkgs.writers.writePython3Bin "change_colors_json" {
      libraries = with pkgs.python3Packages; [ numpy pillow scikit-learn ];
      flakeIgnore = [ "E302" "E305" "E226" "E501" ];
    } /*python */ ''
      from colorsys import hls_to_rgb, rgb_to_hls
      import json
      import sys
      import subprocess
      from time import sleep
      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 = 1):
              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 = 0.6
          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)

      def hyprpicker() -> Color:
          ret = str(subprocess.run(["bash", "-c", "${pkgs.hyprpicker}/bin/hyprpicker -n -f rgb | tail -n 1"], capture_output=True).stdout)[2:-3]
          return Color([int(c) for c in ret.split(" ")])

      if __name__ == "__main__":
          if sys.argv[1] == "img":
              img = sys.argv[2]
              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)
          elif sys.argv[1] == "manual":
              accent = ensure_color(c=hyprpicker(), alter_sat=False)
              sleep(0.1)
              secondary = ensure_color(c=hyprpicker(), alter_sat=True)
              sleep(0.1)
              tertiary = ensure_color(c=hyprpicker(), alter_sat=False)

          weird = alter_hue(ilist=accent, hue=80)
          special = alter_hue(ilist=accent, hue=180)
          foreground = alter_l(accent, 0.9)
          background = alter_l(accent, 0.05)

          with open("${hyper.home}/nix-dots/flake-modules/colors.json", "r") as f:
              full_d = json.load(f)

          full_d['${hyper.host}'] = {
              "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("${hyper.home}/nix-dots/flake-modules/colors.json", "w") as f:
              f.write(json.dumps(full_d, indent=4))
    '')
  ];
}