#!/usr/bin/env python3

import numpy as np
import scipy.ndimage as ndimage

WIDTH = 300
SECTION_HEIGHT = 200

def dither2_diffuse(signal, diffusion_vector, factor=1):
    dithered = np.copy(signal)
    for i in range(len(dithered)):
        dithered[i] = 0 if dithered[i] < 0.5 else 1
        err = signal[i] - dithered[i]
        num_neighbors = len(dithered) - i - 1
        for j in range(min(num_neighbors, len(diffusion_vector))):
            dithered[i+j+1] += err * diffusion_vector[j] * factor
    return np.int8(dithered)

def dither2_noise(signal, make_noise):
    noise = make_noise(len(signal))
    return np.int8(signal + noise)

def white_noise(length):
    return np.random.default_rng().random(length)

def blue_noise(initial_one_fraction=0.1, stdev=1.5):
    def gaussian(arr):
        return np.fft.ifft(ndimage.fourier_gaussian(np.fft.fft(arr), stdev))
    
    def find_largest_void_idx(arr):
        return np.argmax((1-arr) * (1-gaussian(arr)))
    
    def find_tightest_cluster_idx(arr):
        return np.argmax(arr * gaussian(arr))
    
    def make(length):
        if initial_one_fraction >= 0.5:
            raise ValueError("initial_one_fraction must be less than 0.5")
        
        # Initialize the binary pattern with white noise.
        rng = np.random.default_rng()
        ones = max(1, int(length * initial_one_fraction))
        initial_binary_pattern = np.append(np.ones(ones), np.zeros(length-ones))
        rng.shuffle(initial_binary_pattern)
        
        # Break up clusters until the binary pattern is uniform enough.
        converged = False
        while (not converged):
            # The minority pixels are ones. 
            tightest_cluster_idx = find_tightest_cluster_idx(initial_binary_pattern)
            initial_binary_pattern[tightest_cluster_idx] = 0
            largest_void_idx = find_largest_void_idx(initial_binary_pattern)
            if tightest_cluster_idx == largest_void_idx:
                initial_binary_pattern[tightest_cluster_idx] = 1
                converged = True
            else:
                initial_binary_pattern[largest_void_idx] = 1
        
        dither_arr = np.zeros(length)
        
        # Phase I
        prototype_binary_pattern = np.copy(initial_binary_pattern)
        for rank in range(ones-1, -1, -1):
            tightest_cluster_idx = find_tightest_cluster_idx(prototype_binary_pattern)
            prototype_binary_pattern[tightest_cluster_idx] = 0
            dither_arr[tightest_cluster_idx] = rank
        
        # Phase II
        np.copyto(prototype_binary_pattern, initial_binary_pattern)
        rank = ones
        for rank in range(ones, int((length+1)/2)):
            largest_void_idx = find_largest_void_idx(prototype_binary_pattern)
            prototype_binary_pattern[largest_void_idx] = 1
            dither_arr[largest_void_idx] = rank
        
        # Phase III
        # Zeros are the minority pixel now, so invert the pattern to find
        # clusters of them.
        prototype_binary_pattern = 1-prototype_binary_pattern
        for rank in range(int((length+1)/2), length):
            tightest_cluster_idx = find_tightest_cluster_idx(prototype_binary_pattern)
            prototype_binary_pattern[tightest_cluster_idx] = 0
            dither_arr[tightest_cluster_idx] = rank
        
        # Normalize to [0, length)
        return dither_arr/length
    return make
        

if __name__ == '__main__':
    import sys
    stdout = open(sys.__stdout__.fileno(),
        mode=sys.__stdout__.mode,
        buffering=1,
        encoding=sys.__stdout__.encoding,
        errors=sys.__stdout__.errors,
        newline='\n',
        closefd=False)
              
    gradient = np.arange(0, 1, 1/(WIDTH-1))
    diffusion_vector = np.ones(7)
    
    diffused = dither2_diffuse(gradient, diffusion_vector, 1/len(diffusion_vector))
    white_noised = dither2_noise(gradient, white_noise)
    blue_noised = dither2_noise(gradient, blue_noise(stdev=1.9))
    
    sections = [diffused, white_noised, blue_noised]
    print("P1", file=stdout)
    print(WIDTH, SECTION_HEIGHT * len(sections), file=stdout)
    for section in sections:
        for _ in range(SECTION_HEIGHT):
            print(*section, 1, file=stdout)