import re
import sys
import argparse
import os
import stat
from math import *
from dataclasses import dataclass
from typing import List, Union, Tuple
import shutil
from pathlib import Path
import random

from PIL import Image, ImageDraw
import numpy as np

def read_image(filename):
    rgb_img = Image.open(filename).convert('RGB')
    pixel_array = np.array(rgb_img, dtype=np.float32) / 255.0
    return pixel_array

# We only care about color distribution
def read_and_compress_image(filename):
    img = Image.open(filename).convert('RGB')
    width = min(img.width, 50)
    height = min(img.height, 50)
    img = img.resize((width, height))
    pixel_array = np.array(img, dtype=np.float32) / 255.0
    return pixel_array

def conv_color(color: List[float]) -> Tuple[int, int, int, int]:
    return (int(round(255 * color[0])), int(round(255 * color[1])), int(round(255 * color[2])), 255)

def create_color_row_image(colors, output_filename, rect_width = 5, rect_height = 3):
    num_colors = len(colors)
    assert num_colors > 0
    img_width = num_colors * rect_width + (num_colors - 1)
    img_height = rect_height
    new_image = Image.new('RGBA', (img_width, img_height))
    pixels = new_image.load()

    for X, color in enumerate(colors):
        for i in range(rect_width):
            for j in range(rect_height):
                pixels[X + X * rect_width + i, j] = conv_color(color)

    new_image.save(output_filename)
    print(f"Image saved as {output_filename}")

def euclidian_distance(c1: List[float], c2: List[float]):
    return sqrt(sum(map(lambda i: (c1[i] - c2[i]) * (c1[i] - c2[i]), range(3))))

@dataclass
class ColorDistancePrice:
    def apply(self, c1: List[float], c2: List[float]) -> float:
        pass

    def show_image_quality(self, c2: List[float], image):
        new_image = Image.new('RGB', (image.shape[1], image.shape[0]))
        pixels = new_image.load()
        for y in range(image.shape[0]):
            for x in range(image.shape[1]):
                cost: float = self.apply(image[y, x], c2)
                cc: int = int(round(255 * cost))
                pixels[x, y] = (min(cc, 255), max(0, min(cc - 255, 255)), 0)
        new_image.show()


@dataclass
class OmegaPrice(ColorDistancePrice):
    d: float

    def apply(self, c1: List[float], c2: List[float]) -> float:
        e = euclidian_distance(c1, c2)
        if e <= self.d:
            return pow(self.d - e, 2)
        return 0

@dataclass
class EtaPrice(ColorDistancePrice):
    d: float

    def apply(self, c1: List[float], c2: List[float]) -> float:
        e = euclidian_distance(c1, c2)
        if e <= self.d:
            return 1 - e / self.d
        return 1 - exp(self.d - e)

@dataclass
class AlphaPrice(ColorDistancePrice):
    dd: float
    dp: float
    def apply(self, c1: List[float], c2: List[float]) -> float:
        e = euclidian_distance(c1, c2)
        if e < self.dd:
            return 0
        return self.dp + pow(e / sqrt(3), 3)

@dataclass
class ColorDistanceRule:
    i: int
    j: int
    func: ColorDistancePrice

def annealing(T: float, cost_delta: float) -> bool:
    if cost_delta < 0:
        return True
    P = exp(-(cost_delta / T))
    return random.random() < P

class ColorSchemeSearchConf:
    def __init__(self, image100: np.ndarray, n, pair_rules: List[ColorDistanceRule], image_pixel_price: ColorDistancePrice):
        self.n = n
        self.pair_price = [[[] for i in range(m)] for m in range(n) ]
        for rule in pair_rules:
            self.pair_price[max(rule.i, rule.j)][min(rule.i, rule.j)].append(rule.func)
        self.image_cost_c = len(pair_rules)
        self.image_pixel_price = image_pixel_price
        self.image100 = image100

    def drift(self, sample_pos: List[Tuple[int, int]], coef: float) -> List[Tuple[int, int]]:
        w: int = self.image100.shape[1]
        h: int = self.image100.shape[0]
        def drift_point(pos: Tuple[int, int]):
            return (
                max(0, min(pos[0] + int(round(coef * (random.random() - .5))), w - 1)),
                max(0, min(pos[1] + int(round(coef * (random.random() - .5))), h - 1))
            )

        return list(map(drift_point, sample_pos))

    def sample_image(self, sample_positions: List[Tuple[int, int]]) -> List[List[float]]:
        return list(map(lambda pos: self.image100[pos[1], pos[0]], sample_positions))

    def evaluate_colorscheme(self, colorscheme: List[List[float]]):
        w: int = self.image100.shape[1]
        h: int = self.image100.shape[0]
        assert len(colorscheme) == self.n
        C_with_image = 0
        for y in range(h):
            for x in range(w):
                c1 = self.image100[x, y]
                for i in range(self.n):
                    C_with_image += self.image_pixel_price.apply(c1, colorscheme[i])
        C_with_image *= (self.image_cost_c / w / h)
        C_with_each_other = 0
        for i, func_arr_arr in enumerate(self.pair_price):
            for j, func_arr in enumerate(func_arr_arr):
                for func in func_arr:
                    a = func.apply(colorscheme[i], colorscheme[j])
                    C_with_each_other += a
        return C_with_image + C_with_each_other

    def find_best_colorscheme(self, steps: int, initial_temperature: float):
        positions = [(0, 0) for i in range(self.n)]
        colors: List[List[float]] = self.sample_image(positions)
        cost = self.evaluate_colorscheme(colors)
        T = initial_temperature
        for i in range(steps):
            print(f"Temperature {T}. Cost: {cost}. Positions: {positions}")
            new_positions = self.drift(positions, 200 * (1 - exp(-T)))
            new_colorscheme = self.sample_image(new_positions)
            new_cost = self.evaluate_colorscheme(new_colorscheme)
            if annealing(T, new_cost - cost):
                positions = new_positions
                colors = new_colorscheme
                cost = new_cost
            T *= (1 - 10/steps)
        return colors

def color_to_hex(clr: List[float]) -> str:
    return f"#{int(round(255 * clr[0])):02X}{int(round(255 * clr[1])):02X}{int(round(255 * clr[2])):02X}"


# One element per monitor
# swaybar: background
sway_color_bar_background = 0
# swaybar: statusline
sway_color_bar_statusline = 1
#swaybar: focused_background
sway_color_bar_focused_background = 2
#swaybar: focused_statusbar
sway_color_bar_focused_statusline = 3
#swaybar: focused_workspace
sway_color_bar_workspace_focused_border = 4
sway_color_bar_workspace_focused_background = 5
sway_color_bar_workspace_focused_text = 6
#swaybar: active_workspace
sway_color_bar_workspace_active_border = 7
sway_color_bar_workspace_active_background = 8
sway_color_bar_workspace_active_text = 9
#swaybar: inactive_workspace
sway_color_bar_workspace_inactive_border = 10
sway_color_bar_workspace_inactive_background = 11
sway_color_bar_workspace_inactive_text = 12
#swaybar: urgent_workspace
sway_color_bar_workspace_urgent_border = 13
sway_color_bar_workspace_urgent_background = 14
sway_color_bar_workspace_urgent_text = 15
#swaybar: binding_mode
sway_color_bar_mode_indicator_border = 16
sway_color_bar_mode_indicator_background = 17
sway_color_bar_mode_indicator_text = 18
# n when constructing my ColorSchemeSearchConf
sway_color_n = 19

def sway_monitor_image(scheme, output_filename):
    lwksp = 20
    gap = 5
    ltxt = 10

    ls = (lwksp + gap * 3 + ltxt)
    wksp_gap = (lwksp - ltxt) / 2
    img = Image.new('RGBA', (5 * lwksp + 6 * gap, 2 * ls))
    draw = ImageDraw.Draw(img)
    def draw_section(i, ho):
        draw.rectangle([(0, i * ls), (5 * lwksp + 6 * gap, i * ls + ls)], fill=conv_color(scheme[ho]))
        draw.rectangle([(gap, i * ls + 2 * gap + lwksp), (5 * lwksp + 5 * gap, i * ls + 2 * gap + lwksp + ltxt)], fill=conv_color(scheme[ho+1]))
        def draw_wksp(j, brdr):
            draw.rectangle([((j+1) * gap + j * lwksp, ls * i + gap), ((j+1) * gap + j * lwksp + lwksp, ls * i + gap + lwksp)], fill=conv_color(scheme[brdr + 1]), outline=conv_color(scheme[brdr]))
            draw.rectangle([((j+1) * gap + j * lwksp + wksp_gap, ls * i + gap + wksp_gap),
                            ((j+1) * gap + j * lwksp + wksp_gap + ltxt, ls * i + gap + wksp_gap + ltxt)],
                           fill=conv_color(scheme[brdr + 2]))
        draw_wksp(0, sway_color_bar_workspace_focused_border)
        draw_wksp(1, sway_color_bar_workspace_active_border)
        draw_wksp(2, sway_color_bar_workspace_inactive_border)
        draw_wksp(3, sway_color_bar_workspace_urgent_border)
        draw_wksp(4, sway_color_bar_mode_indicator_border)

    draw_section(0, sway_color_bar_background)
    draw_section(1, sway_color_bar_focused_background)
    img.save(output_filename)
    print(f"Image saved as {output_filename}")


def get_sway_search_conf(image100: np.ndarray) -> ColorSchemeSearchConf:
    def wild_trio(first_ind):
        return [
            ColorDistanceRule(first_ind + 0, first_ind + 1, OmegaPrice(0.2)),
            ColorDistanceRule(first_ind + 1, first_ind + 2, OmegaPrice(0.6)),
            ColorDistanceRule(first_ind + 0, first_ind + 2, OmegaPrice(0.5)),
        ]

    rules = [
        ColorDistanceRule(sway_color_bar_background, sway_color_bar_statusline, OmegaPrice(10)),
        ColorDistanceRule(sway_color_bar_focused_background, sway_color_bar_focused_statusline, OmegaPrice(10)),
        ColorDistanceRule(sway_color_bar_background, sway_color_bar_focused_background, OmegaPrice(0.2)),
        ColorDistanceRule(sway_color_bar_statusline, sway_color_bar_focused_statusline, OmegaPrice(0.2)),

        ColorDistanceRule(sway_color_bar_mode_indicator_background, sway_color_bar_workspace_urgent_background, EtaPrice(0.025)),

        ColorDistanceRule(sway_color_bar_workspace_focused_background, sway_color_bar_workspace_active_background, OmegaPrice(0.4)),
        ColorDistanceRule(sway_color_bar_workspace_focused_background, sway_color_bar_workspace_inactive_background, OmegaPrice(0.4)),
        ColorDistanceRule(sway_color_bar_workspace_active_background, sway_color_bar_workspace_inactive_background, OmegaPrice(0.1)),

        ColorDistanceRule(sway_color_bar_background, sway_color_bar_workspace_inactive_background, OmegaPrice(0.07)),
        ColorDistanceRule(sway_color_bar_focused_background, sway_color_bar_workspace_inactive_background, OmegaPrice(0.07)),
        ColorDistanceRule(sway_color_bar_background, sway_color_bar_workspace_active_background, OmegaPrice(0.07)),
        ColorDistanceRule(sway_color_bar_focused_background, sway_color_bar_workspace_active_background, OmegaPrice(0.07)),
        ColorDistanceRule(sway_color_bar_background, sway_color_bar_workspace_focused_background, OmegaPrice(0.07)),
        ColorDistanceRule(sway_color_bar_focused_background, sway_color_bar_workspace_focused_background,OmegaPrice(0.07)),
    ]
    rules += wild_trio(sway_color_bar_workspace_focused_border)
    rules += wild_trio(sway_color_bar_workspace_active_border)
    rules += wild_trio(sway_color_bar_workspace_inactive_border)
    rules += wild_trio(sway_color_bar_workspace_urgent_border)
    rules += wild_trio(sway_color_bar_mode_indicator_border)
    return ColorSchemeSearchConf(image100, sway_color_n, rules, AlphaPrice(0.1, 0.5))


def extract_colors_for_sway_complete_io(source_path, dest_colorscheme_path, dest_thumbnail_path):
    image100 = read_and_compress_image(source_path)
    print(f"Image shape: {image100.shape}", file=sys.stderr)
    searcher = get_sway_search_conf(image100)
    colorscheme = searcher.find_best_colorscheme(20, 1000)

    with open(dest_colorscheme_path, 'a') as f:
        for i in range(searcher.n):
            print(color_to_hex(colorscheme[i]), file=f)

    sway_monitor_image(colorscheme, dest_thumbnail_path)


def compile_a_directory(source_dir, target_txt_dir, target_thumbnail_path):
    image_extensions = {'.jpg', '.jpeg', '.png', '.webp'}
    source_path = Path(source_dir)
    target_txt_path = Path(target_txt_dir)
    target_thumbnail_path = Path(target_thumbnail_path)

    # Walk through all directories and files
    for root, dirs, files in os.walk(source_dir):
        root_path = Path(root)

        for file in files:
            file_path = root_path / file
            file_ext = file_path.suffix.lower()
            if file_ext in image_extensions:
                # Create corresponding target directory and colorscheme
                relative_path = file_path.relative_to(source_path)
                target_txt_dir_path = target_txt_path / relative_path.parent
                target_txt_dir_path.mkdir(parents=True, exist_ok=True)
                target_txt_file_path = target_txt_dir_path / f"{file_path.stem}.txt"
                target_thumbnail_dir_path = target_thumbnail_path / relative_path.parent
                target_thumbnail_dir_path.mkdir(parents=True, exist_ok=True)
                target_thumbnail_file_path = target_thumbnail_dir_path / f"{file_path.stem}.png"
                print(str(file_path), "To", str(target_txt_file_path), "And", str(target_thumbnail_file_path))
                extract_colors_for_sway_complete_io(str(file_path), str(target_txt_file_path), str(target_thumbnail_file_path))

if __name__ == "__main__":
    # parser_arg = argparse.ArgumentParser(description="Extract colors from image")
    # parser_arg.add_argument("input_file", help="Path to the input image file")
    # parser_arg.add_argument("color_file", help="Path for color blocks output image", nargs='?', default=None)
    # parser_arg.add_argument("monitor_file", help="Path for monitor thumbnail output image", nargs='?', default=None)
    # args = parser_arg.parse_args()
    parser_arg = argparse.ArgumentParser(description="Extract colors from a set of images")
    parser_arg.add_argument("source_dir", help="Path to directory with input files")
    parser_arg.add_argument("dest_txt_dir", help="Path to output directory for colorscheme specs")
    parser_arg.add_argument("dest_thumbnail_dir", help="Path to output directory for thumbnail of sway theme")
    args = parser_arg.parse_args()
    compile_a_directory(args.source_dir, args.dest_txt_dir, args.dest_thumbnail_dir)