#!/usr/bin/env python3

import numpy as np
from PIL import Image
import argparse
from pathlib import Path

def gaussian_kernel(size, sigma):
    x = np.arange(-size // 2 + 1, size // 2 + 1)
    y = np.arange(-size // 2 + 1, size // 2 + 1)
    xx, yy = np.meshgrid(x, y)
    kernel = np.exp(-(xx**2 + yy**2) / (2 * sigma**2))
    return kernel / kernel.sum()

def apply_periodic_filter(binary_pattern, kernel):
    from scipy.signal import fftconvolve
    return fftconvolve(binary_pattern, kernel, mode='same')

def generate_blue_noise_void_cluster(width, height, sigma=1.5, iterations=None):
    if iterations is None:
        iterations = width * height

    size = width * height
    pattern = np.zeros((height, width), dtype=np.float32)

    kernel_size = int(6 * sigma)
    if kernel_size % 2 == 0:
        kernel_size += 1
    kernel = gaussian_kernel(kernel_size, sigma)

    initial_pattern = np.random.rand(height, width)

    print(f"Generating {width}x{height} blue noise texture...")
    print(f"Kernel size: {kernel_size}x{kernel_size}, sigma: {sigma}")

    dither_array = np.zeros(size, dtype=np.int32)
    binary_pattern = np.zeros((height, width), dtype=np.float32)

    for i in range(size):
        if i % (size // 10) == 0:
            print(f"Progress: {i}/{size} ({100*i//size}%)")

        filtered = apply_periodic_filter(binary_pattern, kernel)

        if i < size // 2:
            initial_energy = initial_pattern + filtered
            coords = np.unravel_index(np.argmax(initial_energy), initial_energy.shape)
        else:
            coords = np.unravel_index(np.argmin(filtered), filtered.shape)

        dither_array[i] = coords[0] * width + coords[1]
        binary_pattern[coords[0], coords[1]] = 1.0

    print("Converting to threshold map...")

    threshold_map = np.zeros((height, width), dtype=np.float32)
    for rank, pos in enumerate(dither_array):
        y = pos // width
        x = pos % width
        threshold_map[y, x] = rank / size

    print("Done!")
    return threshold_map

def generate_blue_noise_simulated_annealing(width, height, iterations=10000):
    print(f"Generating {width}x{height} blue noise using simulated annealing...")

    pattern = np.random.rand(height, width)

    def energy(pattern):
        fft = np.fft.fft2(pattern)
        power = np.abs(fft) ** 2

        h, w = pattern.shape
        cy, cx = h // 2, w // 2
        y, x = np.ogrid[:h, :w]
        dist = np.sqrt((x - cx)**2 + (y - cy)**2)

        low_freq_mask = dist < min(h, w) * 0.1
        low_freq_energy = np.sum(power * low_freq_mask)

        return low_freq_energy

    current_energy = energy(pattern)
    temperature = 1.0
    cooling_rate = 0.9995

    for i in range(iterations):
        if i % (iterations // 10) == 0:
            print(f"Iteration {i}/{iterations}, Energy: {current_energy:.2f}, Temp: {temperature:.4f}")

        y1, x1 = np.random.randint(0, height), np.random.randint(0, width)
        y2, x2 = np.random.randint(0, height), np.random.randint(0, width)

        pattern[y1, x1], pattern[y2, x2] = pattern[y2, x2], pattern[y1, x1]

        new_energy = energy(pattern)
        delta_energy = new_energy - current_energy

        if delta_energy < 0 or np.random.rand() < np.exp(-delta_energy / temperature):
            current_energy = new_energy
        else:
            pattern[y1, x1], pattern[y2, x2] = pattern[y2, x2], pattern[y1, x1]

        temperature *= cooling_rate

    print("Done!")
    return pattern

def main():
    parser = argparse.ArgumentParser(description='Generate blue noise texture for dithering')
    parser.add_argument('--width', type=int, default=128, help='Texture width (default: 128)')
    parser.add_argument('--height', type=int, default=128, help='Texture height (default: 128)')
    parser.add_argument('--method', choices=['void_cluster', 'annealing'], default='void_cluster',
                        help='Generation method (default: void_cluster)')
    parser.add_argument('--sigma', type=float, default=1.5,
                        help='Gaussian kernel sigma for void_cluster method (default: 1.5)')
    parser.add_argument('--iterations', type=int, default=None,
                        help='Number of iterations (optional)')
    parser.add_argument('--output', type=str, default='../blue_noise.png',
                        help='Output file path (default: ../blue_noise.png)')

    args = parser.parse_args()

    try:
        from scipy.signal import fftconvolve
    except ImportError:
        print("Error: scipy is required for this script.")
        print("Install it with: pip install scipy")
        return

    if args.method == 'void_cluster':
        noise = generate_blue_noise_void_cluster(args.width, args.height, args.sigma, args.iterations)
    else:
        noise = generate_blue_noise_simulated_annealing(args.width, args.height,
                                                         args.iterations or 10000)

    noise_normalized = ((noise - noise.min()) / (noise.max() - noise.min()) * 255).astype(np.uint8)

    img = Image.fromarray(noise_normalized, mode='L')

    output_path = Path(__file__).parent / args.output
    output_path.parent.mkdir(parents=True, exist_ok=True)
    img.save(output_path)

    print(f"\nBlue noise texture saved to: {output_path}")
    print(f"Size: {args.width}x{args.height}")
    print(f"Method: {args.method}")

    fft = np.fft.fft2(noise)
    power_spectrum = np.abs(np.fft.fftshift(fft)) ** 2
    print(f"Power spectrum range: {power_spectrum.min():.2e} - {power_spectrum.max():.2e}")

if __name__ == '__main__':
    main()