from tqdm import tqdm
from scipy.ndimage import uniform_filter1d
from scipy.io.wavfile import write
import math, cv2, sys, getopt, wave
import numpy as np

def unwrap_img(input_file):
    ring_img = cv2.imread(input_file)

    size = ring_img.shape[0]
    outer_radius = size // 2
    inner_radius = 0

    unwrapped_img = cv2.warpPolar(
        ring_img,
        (size, int(size * math.pi)),
        (outer_radius, outer_radius),
        outer_radius,
        flags=0
    )

    rotated_unwrapped = cv2.rotate(unwrapped_img, cv2.ROTATE_90_COUNTERCLOCKWISE)
    cropped_unwrapped = rotated_unwrapped[inner_radius:, :]

    inverted_unwrapped = cv2.bitwise_not(cropped_unwrapped)

    return inverted_unwrapped


def image_to_audio(input_img, out_wav, duration_seconds, sample_rate, vertical_res, amp_threshold):
    max_freq = 10000
    min_freq = 50

    duration_seconds = float(duration_seconds)

    # Downsample image vertically to reduce number of frequencies
    input_img = input_img[::vertical_res]  # Use one row every 20 pixels
    height, width, _ = input_img.shape

    num_samples = int(sample_rate * duration_seconds)
    freqs = np.logspace(np.log10(min_freq), np.log10(max_freq), height)[::-1]

    brightness = np.mean(input_img / 255.0, axis=2)
    amplitudes = np.where(brightness >= 0.1, brightness, 0)

    samples = np.zeros(num_samples, dtype=np.float32)
    chunk_size = 10000

    for start in tqdm(range(0, num_samples, chunk_size)):
        end = min(start + chunk_size, num_samples)
        t = np.linspace(start / sample_rate, end / sample_rate, end - start)

        pixel_xs = (t * width / duration_seconds).astype(int)
        pixel_xs = np.clip(pixel_xs, 0, width - 1)

        amp_per_sample = amplitudes[:, pixel_xs]
        amp_per_sample[amp_per_sample < amp_threshold] = 0
        amp_per_sample = uniform_filter1d(amp_per_sample, size=7, axis=1)  # Smooth
        
        phases = 2 * np.pi * freqs[:, None] * t[None, :]
        active = np.count_nonzero(amp_per_sample, axis=0)
        chunk = np.sum(amp_per_sample * np.sin(phases), axis=0)

        # Normalize chunk by number of active oscillators to avoid "snow"
        chunk = np.where(active > 0, chunk / active, 0)
        samples[start:end] = chunk

    # Final normalization
    samples /= np.max(np.abs(samples) + 1e-8)

    # Convert to int16 for wav
    wav_samples = (samples * 32767).astype(np.int16)
    write(out_wav, sample_rate, wav_samples)
    print(f"Audio saved to {out_wav}")


if __name__ == '__main__':
    input_file = ''
    output_file = ''
    duration = 10

    try:
        opts, args = getopt.getopt(sys.argv[1:], "hi:o:d:yRes:thresh:")
    except getopt.GetoptError:
        print('error: img_to_freq.py -i <input_picture> -o <output_sound> -d <audio_duration>')

    for opt, arg in opts:
        if opt == '-h':
            print('error: img_to_freq.py -i <input_picture> -o <output_sound> -d <audio_duration>')
            sys.exit(2)
        elif opt == '-i':
            input_file = arg
        elif opt == '-o':
            output_file = arg
        elif opt == '-d':
            duration = arg

    unwrapped_img = unwrap_img(input_file)
    cv2.imwrite('output.png', unwrapped_img)

    image_to_audio(unwrapped_img, output_file, duration, 44100, 20, 0.1)