feat: generate dynamic schemes

6 files changed, 405 insertions, 8 deletions

diff --git a/src/caelestia/utils/material/__init__.py b/src/caelestia/utils/material/__init__.py
new file mode 100644
index 0000000..8adab1f
--- /dev/null
+++ b/src/caelestia/utils/material/__init__.py
@@ -0,0 +1,52 @@
+import json
+from pathlib import Path
+
+from materialyoucolor.hct import Hct
+
+from caelestia.utils.material.generator import gen_scheme
+from caelestia.utils.material.score import score
+from caelestia.utils.paths import compute_hash, scheme_cache_dir, wallpaper_thumbnail_path
+
+
+def get_score_for_image(image: str, cache_base: Path) -> tuple[list[Hct], list[Hct]]:
+    cache = cache_base / "score.json"
+
+    try:
+        with cache.open("r") as f:
+            return [[Hct.from_int(c) for c in li] for li in json.load(f)]
+    except (IOError, json.JSONDecodeError):
+        pass
+
+    s = score(image)
+
+    cache.parent.mkdir(parents=True, exist_ok=True)
+    with cache.open("w") as f:
+        json.dump([[c.to_int() for c in li] for li in s], f)
+
+    return s
+
+
+def get_colours_for_image(image: str = str(wallpaper_thumbnail_path), scheme=None) -> dict[str, str]:
+    if scheme is None:
+        from caelestia.utils.scheme import get_scheme
+
+        scheme = get_scheme()
+
+    cache_base = scheme_cache_dir / compute_hash(image)
+    cache = (cache_base / scheme.variant / scheme.flavour / scheme.mode).with_suffix(".json")
+
+    try:
+        with cache.open("r") as f:
+            return json.load(f)
+    except (IOError, json.JSONDecodeError):
+        pass
+
+    primaries, colours = get_score_for_image(image, cache_base)
+    i = ["default", "alt1", "alt2"].index(scheme.flavour)
+    scheme = gen_scheme(scheme, primaries[i], colours)
+
+    cache.parent.mkdir(parents=True, exist_ok=True)
+    with cache.open("w") as f:
+        json.dump(scheme, f)
+
+    return scheme
diff --git a/src/caelestia/utils/material/generator.py b/src/caelestia/utils/material/generator.py
new file mode 100755
index 0000000..33ff0e8
--- /dev/null
+++ b/src/caelestia/utils/material/generator.py
@@ -0,0 +1,192 @@
+from materialyoucolor.blend import Blend
+from materialyoucolor.dynamiccolor.material_dynamic_colors import (
+    DynamicScheme,
+    MaterialDynamicColors,
+)
+from materialyoucolor.hct import Hct
+from materialyoucolor.hct.cam16 import Cam16
+from materialyoucolor.scheme.scheme_content import SchemeContent
+from materialyoucolor.scheme.scheme_expressive import SchemeExpressive
+from materialyoucolor.scheme.scheme_fidelity import SchemeFidelity
+from materialyoucolor.scheme.scheme_fruit_salad import SchemeFruitSalad
+from materialyoucolor.scheme.scheme_monochrome import SchemeMonochrome
+from materialyoucolor.scheme.scheme_neutral import SchemeNeutral
+from materialyoucolor.scheme.scheme_rainbow import SchemeRainbow
+from materialyoucolor.scheme.scheme_tonal_spot import SchemeTonalSpot
+from materialyoucolor.scheme.scheme_vibrant import SchemeVibrant
+
+
+def hex_to_hct(hex: str) -> Hct:
+    return Hct.from_int(int(f"0xFF{hex}", 16))
+
+
+light_colours = [
+    hex_to_hct("dc8a78"),
+    hex_to_hct("dd7878"),
+    hex_to_hct("ea76cb"),
+    hex_to_hct("8839ef"),
+    hex_to_hct("d20f39"),
+    hex_to_hct("e64553"),
+    hex_to_hct("fe640b"),
+    hex_to_hct("df8e1d"),
+    hex_to_hct("40a02b"),
+    hex_to_hct("179299"),
+    hex_to_hct("04a5e5"),
+    hex_to_hct("209fb5"),
+    hex_to_hct("1e66f5"),
+    hex_to_hct("7287fd"),
+]
+
+dark_colours = [
+    hex_to_hct("f5e0dc"),
+    hex_to_hct("f2cdcd"),
+    hex_to_hct("f5c2e7"),
+    hex_to_hct("cba6f7"),
+    hex_to_hct("f38ba8"),
+    hex_to_hct("eba0ac"),
+    hex_to_hct("fab387"),
+    hex_to_hct("f9e2af"),
+    hex_to_hct("a6e3a1"),
+    hex_to_hct("94e2d5"),
+    hex_to_hct("89dceb"),
+    hex_to_hct("74c7ec"),
+    hex_to_hct("89b4fa"),
+    hex_to_hct("b4befe"),
+]
+
+colour_names = [
+    "rosewater",
+    "flamingo",
+    "pink",
+    "mauve",
+    "red",
+    "maroon",
+    "peach",
+    "yellow",
+    "green",
+    "teal",
+    "sky",
+    "sapphire",
+    "blue",
+    "lavender",
+    "success",
+    "error",
+]
+
+
+def grayscale(colour: Hct, light: bool) -> None:
+    colour.chroma = 0
+
+
+def mix(a: Hct, b: Hct, w: float) -> Hct:
+    return Hct.from_int(Blend.cam16_ucs(a.to_int(), b.to_int(), w))
+
+
+def harmonize(a: Hct, b: Hct) -> Hct:
+    return Hct.from_int(Blend.harmonize(a.to_int(), b.to_int()))
+
+
+def lighten(colour: Hct, amount: float) -> Hct:
+    diff = (100 - colour.tone) * amount
+    return Hct.from_hct(colour.hue, colour.chroma + diff / 2, colour.tone + diff)
+
+
+def darken(colour: Hct, amount: float) -> Hct:
+    diff = colour.tone * amount
+    return Hct.from_hct(colour.hue, colour.chroma + diff / 2, colour.tone - diff)
+
+
+def distance(colour: Cam16, base: Cam16) -> float:
+    return colour.distance(base)
+
+
+def smart_sort(colours: list[Hct], base: list[Hct]) -> dict[str, Hct]:
+    sorted_colours = [None] * len(colours)
+    distances = {}
+
+    cams = [(c, Cam16.from_int(c.to_int())) for c in colours]
+    base_cams = [Cam16.from_int(c.to_int()) for c in base]
+
+    for colour, cam in cams:
+        dist = [(i, distance(cam, b)) for i, b in enumerate(base_cams)]
+        dist.sort(key=lambda x: x[1])
+        distances[colour] = dist
+
+    for colour in colours:
+        while len(distances[colour]) > 0:
+            i, dist = distances[colour][0]
+
+            if sorted_colours[i] is None:
+                sorted_colours[i] = colour, dist
+                break
+            elif sorted_colours[i][1] > dist:
+                old = sorted_colours[i][0]
+                sorted_colours[i] = colour, dist
+                colour = old
+
+            distances[colour].pop(0)
+
+    return {colour_names[i]: c[0] for i, c in enumerate(sorted_colours)}
+
+
+def get_scheme(scheme: str) -> DynamicScheme:
+    if scheme == "content":
+        return SchemeContent
+    if scheme == "expressive":
+        return SchemeExpressive
+    if scheme == "fidelity":
+        return SchemeFidelity
+    if scheme == "fruitsalad":
+        return SchemeFruitSalad
+    if scheme == "monochrome":
+        return SchemeMonochrome
+    if scheme == "neutral":
+        return SchemeNeutral
+    if scheme == "rainbow":
+        return SchemeRainbow
+    if scheme == "tonalspot":
+        return SchemeTonalSpot
+    return SchemeVibrant
+
+
+def gen_scheme(scheme, primary: Hct, colours: list[Hct]) -> dict[str, str]:
+    light = scheme.mode == "light"
+    base = light_colours if light else dark_colours
+
+    # Sort colours and turn into dict
+    colours = smart_sort(colours, base)
+
+    # Harmonize colours
+    for name, hct in colours.items():
+        if scheme.variant == "monochrome":
+            grayscale(hct, light)
+        else:
+            harmonized = harmonize(hct, primary)
+            colours[name] = darken(harmonized, 0.35) if light else lighten(harmonized, 0.65)
+
+    # Material colours
+    primary_scheme = get_scheme(scheme.variant)(primary, not light, 0)
+    for colour in vars(MaterialDynamicColors).keys():
+        colour_name = getattr(MaterialDynamicColors, colour)
+        if hasattr(colour_name, "get_hct"):
+            colours[colour] = colour_name.get_hct(primary_scheme)
+
+    # FIXME: deprecated stuff
+    colours["text"] = colours["onBackground"]
+    colours["subtext1"] = colours["onSurfaceVariant"]
+    colours["subtext0"] = colours["outline"]
+    colours["overlay2"] = mix(colours["surface"], colours["outline"], 0.86)
+    colours["overlay1"] = mix(colours["surface"], colours["outline"], 0.71)
+    colours["overlay0"] = mix(colours["surface"], colours["outline"], 0.57)
+    colours["surface2"] = mix(colours["surface"], colours["outline"], 0.43)
+    colours["surface1"] = mix(colours["surface"], colours["outline"], 0.29)
+    colours["surface0"] = mix(colours["surface"], colours["outline"], 0.14)
+    colours["base"] = colours["surface"]
+    colours["mantle"] = darken(colours["surface"], 0.03)
+    colours["crust"] = darken(colours["surface"], 0.05)
+    colours["success"] = harmonize(base[8], primary)
+
+    # For debugging
+    print("\n".join(["{}: \x1b[48;2;{};{};{}m   \x1b[0m".format(n, *c.to_rgba()[:3]) for n, c in colours.items()]))
+
+    return {k: hex(v.to_int())[4:] for k, v in colours.items()}
diff --git a/src/caelestia/utils/material/score.py b/src/caelestia/utils/material/score.py
new file mode 100755
index 0000000..da8b062
--- /dev/null
+++ b/src/caelestia/utils/material/score.py
@@ -0,0 +1,129 @@
+#!/usr/bin/env python
+
+import sys
+
+from materialyoucolor.dislike.dislike_analyzer import DislikeAnalyzer
+from materialyoucolor.hct import Hct
+from materialyoucolor.quantize import ImageQuantizeCelebi
+from materialyoucolor.utils.math_utils import difference_degrees, sanitize_degrees_int
+
+
+class Score:
+    TARGET_CHROMA = 48.0
+    WEIGHT_PROPORTION = 0.7
+    WEIGHT_CHROMA_ABOVE = 0.3
+    WEIGHT_CHROMA_BELOW = 0.1
+    CUTOFF_CHROMA = 5.0
+    CUTOFF_EXCITED_PROPORTION = 0.01
+
+    def __init__(self):
+        pass
+
+    @staticmethod
+    def score(colors_to_population: dict) -> tuple[list[Hct], list[Hct]]:
+        desired = 14
+        filter_enabled = True
+        dislike_filter = True
+
+        colors_hct = []
+        hue_population = [0] * 360
+        population_sum = 0
+
+        for rgb, population in colors_to_population.items():
+            hct = Hct.from_int(rgb)
+            colors_hct.append(hct)
+            hue = int(hct.hue)
+            hue_population[hue] += population
+            population_sum += population
+
+        hue_excited_proportions = [0.0] * 360
+
+        for hue in range(360):
+            proportion = hue_population[hue] / population_sum
+            for i in range(hue - 14, hue + 16):
+                neighbor_hue = int(sanitize_degrees_int(i))
+                hue_excited_proportions[neighbor_hue] += proportion
+
+        # Score colours
+        scored_hct = []
+        for hct in colors_hct:
+            hue = int(sanitize_degrees_int(round(hct.hue)))
+            proportion = hue_excited_proportions[hue]
+
+            if filter_enabled and (hct.chroma < Score.CUTOFF_CHROMA or proportion <= Score.CUTOFF_EXCITED_PROPORTION):
+                continue
+
+            proportion_score = proportion * 100.0 * Score.WEIGHT_PROPORTION
+            chroma_weight = Score.WEIGHT_CHROMA_BELOW if hct.chroma < Score.TARGET_CHROMA else Score.WEIGHT_CHROMA_ABOVE
+            chroma_score = (hct.chroma - Score.TARGET_CHROMA) * chroma_weight
+            score = proportion_score + chroma_score
+            scored_hct.append({"hct": hct, "score": score})
+
+        scored_hct.sort(key=lambda x: x["score"], reverse=True)
+
+        # Choose distinct colours
+        chosen_colors = []
+        for difference_degrees_ in range(90, 0, -1):
+            chosen_colors.clear()
+            for item in scored_hct:
+                hct = item["hct"]
+                duplicate_hue = any(
+                    difference_degrees(hct.hue, chosen_hct.hue) < difference_degrees_ for chosen_hct in chosen_colors
+                )
+                if not duplicate_hue:
+                    chosen_colors.append(hct)
+                if len(chosen_colors) >= desired:
+                    break
+            if len(chosen_colors) >= desired:
+                break
+
+        # Get primary colour
+        primary = None
+        for cutoff in range(20, 0, -1):
+            for item in scored_hct:
+                if item["hct"].chroma > cutoff and item["hct"].tone > cutoff * 3:
+                    primary = item["hct"]
+                    break
+            if primary:
+                break
+
+        # Choose distinct primaries
+        chosen_primaries = [primary]
+        for difference_degrees_ in range(90, 14, -1):
+            chosen_primaries = [primary]
+            for item in scored_hct:
+                hct = item["hct"]
+                duplicate_hue = any(
+                    difference_degrees(hct.hue, chosen_hct.hue) < difference_degrees_ for chosen_hct in chosen_primaries
+                )
+                if not duplicate_hue:
+                    chosen_primaries.append(hct)
+                if len(chosen_primaries) >= 3:
+                    break
+            if len(chosen_primaries) >= 3:
+                break
+
+        # Fix disliked colours
+        if dislike_filter:
+            for i, chosen_hct in enumerate(chosen_primaries):
+                chosen_primaries[i] = DislikeAnalyzer.fix_if_disliked(chosen_hct)
+            for i, chosen_hct in enumerate(chosen_colors):
+                chosen_colors[i] = DislikeAnalyzer.fix_if_disliked(chosen_hct)
+
+        return chosen_primaries, chosen_colors
+
+
+def score(image: str) -> tuple[list[Hct], list[Hct]]:
+    return Score.score(ImageQuantizeCelebi(image, 1, 128))
+
+
+if __name__ == "__main__":
+    img = sys.argv[1]
+    mode = sys.argv[2] if len(sys.argv) > 2 else "hex"
+
+    colours = Score.score(ImageQuantizeCelebi(img, 1, 128))
+    for t in colours:
+        if mode != "hex":
+            print("".join(["\x1b[48;2;{};{};{}m   \x1b[0m".format(*c.to_rgba()[:3]) for c in t]))
+        if mode != "swatch":
+            print(" ".join(["{:02X}{:02X}{:02X}".format(*c.to_rgba()[:3]) for c in t]))
diff --git a/src/caelestia/utils/paths.py b/src/caelestia/utils/paths.py
index dfb57d9..6a6e0a8 100644
--- a/src/caelestia/utils/paths.py
+++ b/src/caelestia/utils/paths.py
@@ -1,16 +1,33 @@
+import hashlib
 import os
 from pathlib import Path
 
 config_dir = Path(os.getenv("XDG_CONFIG_HOME", Path.home() / ".config"))
 data_dir = Path(os.getenv("XDG_DATA_HOME", Path.home() / ".local/share"))
 state_dir = Path(os.getenv("XDG_STATE_HOME", Path.home() / ".local/state"))
+cache_dir = Path(os.getenv("XDG_CACHE_HOME", Path.home() / ".cache"))
 
 c_config_dir = config_dir / "caelestia"
 c_data_dir = data_dir / "caelestia"
 c_state_dir = state_dir / "caelestia"
+c_cache_dir = cache_dir / "caelestia"
 
 cli_data_dir = Path(__file__).parent.parent / "data"
 templates_dir = cli_data_dir / "templates"
 
 scheme_path = c_state_dir / "scheme.json"
-scheme_data_path = cli_data_dir / "schemes"
+scheme_data_dir = cli_data_dir / "schemes"
+scheme_cache_dir = c_cache_dir / "schemes"
+
+last_wallpaper_path = c_state_dir / "wallpaper/last.txt"
+wallpaper_thumbnail_path = c_state_dir / "wallpaper/thumbnail.jpg"
+
+
+def compute_hash(path: str) -> str:
+    sha = hashlib.sha256()
+
+    with open(path, "rb") as f:
+        while chunk := f.read(8192):
+            sha.update(chunk)
+
+    return sha.hexdigest()
diff --git a/src/caelestia/utils/scheme.py b/src/caelestia/utils/scheme.py
index c26d2f8..66cd697 100644
--- a/src/caelestia/utils/scheme.py
+++ b/src/caelestia/utils/scheme.py
@@ -1,7 +1,8 @@
 import json
 from pathlib import Path
 
-from caelestia.utils.paths import scheme_data_path, scheme_path
+from caelestia.utils.material import get_colours_for_image
+from caelestia.utils.paths import scheme_data_dir, scheme_path
 
 
 class Scheme:
@@ -89,7 +90,7 @@ class Scheme:
         return self._colours
 
     def get_colours_path(self) -> Path:
-        return (scheme_data_path / self.name / self.flavour / self.mode).with_suffix(".txt")
+        return (scheme_data_dir / self.name / self.flavour / self.mode).with_suffix(".txt")
 
     def save(self) -> None:
         scheme_path.parent.mkdir(parents=True, exist_ok=True)
@@ -118,7 +119,10 @@ class Scheme:
             self._mode = get_scheme_modes()[0]
 
     def _update_colours(self) -> None:
-        self._colours = read_colours_from_file(self.get_colours_path())
+        if self.name == "dynamic":
+            self._colours = get_colours_for_image()
+        else:
+            self._colours = read_colours_from_file(self.get_colours_path())
 
     def __str__(self) -> str:
         return f"Scheme(name={self.name}, flavour={self.flavour}, mode={self.mode}, variant={self.variant})"
@@ -168,7 +172,7 @@ def get_scheme_names() -> list[str]:
     global scheme_names
 
     if scheme_names is None:
-        scheme_names = [f.name for f in scheme_data_path.iterdir() if f.is_dir()]
+        scheme_names = [f.name for f in scheme_data_dir.iterdir() if f.is_dir()]
         scheme_names.append("dynamic")
 
     return scheme_names
@@ -182,7 +186,7 @@ def get_scheme_flavours() -> list[str]:
         if name == "dynamic":
             scheme_flavours = ["default", "alt1", "alt2"]
         else:
-            scheme_flavours = [f.name for f in (scheme_data_path / name).iterdir() if f.is_dir()]
+            scheme_flavours = [f.name for f in (scheme_data_dir / name).iterdir() if f.is_dir()]
 
     return scheme_flavours
 
@@ -192,6 +196,9 @@ def get_scheme_modes() -> list[str]:
 
     if scheme_modes is None:
         scheme = get_scheme()
-        scheme_modes = [f.stem for f in (scheme_data_path / scheme.name / scheme.flavour).iterdir() if f.is_file()]
+        if scheme.name == "dynamic":
+            scheme_modes = ["light", "dark"]
+        else:
+            scheme_modes = [f.stem for f in (scheme_data_dir / scheme.name / scheme.flavour).iterdir() if f.is_file()]
 
     return scheme_modes
diff --git a/src/caelestia/utils/theme.py b/src/caelestia/utils/theme.py
index 3250ea3..d205fb1 100644
--- a/src/caelestia/utils/theme.py
+++ b/src/caelestia/utils/theme.py
@@ -115,7 +115,7 @@ def apply_btop(colours: dict[str, str]) -> None:
 
 def apply_colours(colours: dict[str, str], mode: str) -> None:
     apply_terms(gen_sequences(colours))
-    apply_hypr(gen_conf(colours))
+    apply_hypr(gen_conf(colours))  # FIXME: LAGGY
     apply_discord(gen_scss(colours))
     apply_spicetify(colours, mode)
     apply_fuzzel(colours)