### Package import ###

import numpy as np
from scipy import signal, fft
from scipy.io.wavfile import write
import pyaudio
from PIL import Image
import matplotlib.pyplot as plt
import os
from pathlib import Path
import math


### Musical helpers ###

# Escala pentatónica mayor en MIDI (C mayor pentatónica como patrón)
PENTATONIC = [60, 62, 64, 67, 69]  # C, D, E, G, A

def midi_to_freq(midi_note):
    return 440.0 * (2 ** ((midi_note - 69) / 12.0))

def freq_to_midi(freq):
    return 69 + 12 * math.log2(freq / 440.0)

def build_scale(base_midi=60, octaves=3):
    """
    Construye una lista de notas en varias octavas
    a partir de la escala pentatónica anterior.
    """
    notes = []
    for o in range(octaves):
        for n in PENTATONIC:
            notes.append(n + 12 * o + (base_midi - 60))
    return notes

def quantize_freq_to_scale(freq, scale_freqs):
    """
    Busca la frecuencia de la escala más cercana a la que has introducido.
    """
    diffs = [abs(freq - f) for f in scale_freqs]
    idx = int(np.argmin(diffs))
    return scale_freqs[idx]

def melodic_sequence_from_image(image_data, sample_rate, note_duration_sec, waveform_type, scale_freqs):
    """
    Construye una secuencia de notas usando la imagen como secuenciador.
    Cada bloque de píxeles elige un índice en scale_freqs.
    """
    img = np.asarray(image_data, dtype=np.float32).reshape(-1)

    # Normalizamos a 0..1
    img_norm = (img - img.min()) / (img.max() - img.min() + 1e-9)

    # Tamaño de bloque en píxeles (cuantos píxeles por nota)
    block_size = 256  # prueba 64, 128, 512 para variar densidad de notas
    num_blocks = len(img_norm) // block_size
    if num_blocks == 0:
        return np.array([], dtype=np.float32)

    blocks = img_norm[:num_blocks * block_size].reshape(num_blocks, block_size)

    note_len_samples = int(note_duration_sec * sample_rate)
    seq = []

    for b in blocks:
        val = b.mean()  # brillo medio del bloque
        idx = int(val * (len(scale_freqs) - 1))
        freq = scale_freqs[idx]

        if waveform_type == 1:
            note = sine_gen(sample_rate, freq, note_len_samples)
        elif waveform_type == 2:
            note = square_gen(sample_rate, freq, note_len_samples)
        else:
            note = saw_gen(sample_rate, freq, note_len_samples)

        fade_audio(note)
        seq.append(note)

    seq_signal = np.concatenate(seq)
    return seq_signal


### Defining functions ###

def my_fft(data):
    data_fft = fft.rfft(data)
    positions = np.linspace(10, 22050, 22040)
    plt.semilogx(positions, np.absolute(data_fft[10:22050]))
    plt.xlabel('Frequency')
    plt.ylabel('Amplitude')
    plt.title('FFT')
    plt.show()

def raster_scan(image, bi_sel):
    width, height = image.size
    pixel_values = []

    for y in range(height):
        # even / odd rows
        left_to_right = (y + 2) % 2 == 0 or ((y + 2) % 2 == 1 and bi_sel == 2)

        if left_to_right:
            xs = range(width)
        else:
            xs = range(width - 1, -1, -1)

        for x in xs:
            pixel_value = image.getpixel((x, y))
            pixel_values.append(pixel_value)

    return pixel_values

def rgb_to_luma(pixel_values):
    """
    Convierte RGB a luminancia perceptual (mantiene 'huella' del color).
    Y devuelve un array 1D de floats.
    """
    out = []
    for p in pixel_values:
        if isinstance(p, tuple):
            if len(p) >= 3:
                r, g, b = p[:3]
            else:
                r = g = b = p[0]
        else:
            r = g = b = p
        # Luminancia ITU-R BT.709
        l = 0.2126 * r + 0.7152 * g + 0.0722 * b
        out.append(l)
    return np.array(out, dtype=np.float32)

def bidirectional_select():
    while True:
        try:
            bidir = int(input(
"""
-------------------------------------------------
Please select bidirectional (L->R then R->L)
or unidirectional (just L->R) raster scanning
-------------------------------------------------

1. Bidirectional
2. Unidirectional

-------------------------------------------------
Select 1/2
-------------------------------------------------

"""))
        except ValueError:
            print("Please enter either 1 or 2.")
            continue

        if 1 <= bidir <= 2:
            break
        else:
            print("Please enter either 1 or 2.")
            continue

    return bidir

def mode_select():
    while True:
        try:
            mode = int(input(
"""
-------------------------------------------------
Please select synthesis mode
-------------------------------------------------

1. Drone (single sustained note modulated by image)
2. Melodic (sequence of notes from image)

-------------------------------------------------
Select 1/2
-------------------------------------------------

"""))
        except ValueError:
            print("Please enter either 1 or 2.")
            continue

        if 1 <= mode <= 2:
            break
        else:
            print("Please enter either 1 or 2.")
            continue

    return mode

def waveform_select():
    while True:
        try:
            waveform = int(input(
"""
-------------------------------------------------
Please select a waveform for the oscillator
-------------------------------------------------

1. Sine wave
2. Square wave
3. Sawtooth wave

-------------------------------------------------
Select 1/2/3
-------------------------------------------------

"""))
        except ValueError:
            print("Please enter an integer between 1 and 3.")
            continue

        if waveform > 3 or waveform < 1:
            print("Please enter an integer between 1 and 3.")
            continue
        else:
            break

    return waveform

def frequency_select():
    while True:
        try:
            osc_freq = int(input(
"""
-------------------------------------------------
Please input a frequency value in Hertz for the
oscillator between 10Hz and 22,000Hz.
(This will define the key/tonality)
-------------------------------------------------

"""))
        except ValueError:
            print("Please enter an integer between 10 and 22,000.")
            continue

        if osc_freq > 22000 or osc_freq < 10:
            print("Please enter an integer between 10 and 22,000.")
            continue
        else:
            break

    return osc_freq

def normalise(data, max_min):
    data = np.asarray(data, dtype=np.float32).reshape(-1)
    min_value = np.min(data)
    max_value = np.max(data)
    normalised_data = ((2 * (data - min_value) / (max_value - min_value + 1e-9)) - 1) * max_min
    print(f"""MaxValue: {np.max(normalised_data)}
MinValue: {np.min(normalised_data)}""")
    return normalised_data

def fade_audio(unfaded_data):
    fade_length = 2400
    if len(unfaded_data) < fade_length * 2:
        fade_length = len(unfaded_data) // 2
    if fade_length <= 0:
        return

    fade_multiplier = np.linspace(0, 1, fade_length)
    for i in range(fade_length):
        unfaded_data[i] *= fade_multiplier[i]
        unfaded_data[-1 - i] *= fade_multiplier[i]

def my_plot(data, plot_name):
    positions = list(range(len(data)))
    plt.plot(positions, data)
    plt.xlabel('Sample No.')
    plt.ylabel('Amplitude')
    plt.title(f'{plot_name}')
    plt.show()

def sine_gen(audio_rate, f, l):
    t = np.linspace(0, l / audio_rate, l, dtype=np.float32)
    y = np.sin(2 * np.pi * f * t)
    return y

def square_gen(audio_rate, f, l):
    t = np.linspace(0, l / audio_rate, l, dtype=np.float32)
    y = signal.square(2 * np.pi * f * t)
    return y

def saw_gen(audio_rate, f, l):
    t = np.linspace(0, l / audio_rate, l, dtype=np.float32)
    y = signal.sawtooth(2 * np.pi * f * t)
    return y


### MAIN ###

try:
    dir_path = os.path.dirname(os.path.realpath(__file__))
    sample_rate = 48000
    channels = 1
    format = pyaudio.paFloat32
    NOTE_DURATION_SEC = 0.18   # duración de cada nota en modo melódico
    MAX_IMG_DIM = 512          # límite de tamaño imagen para no irnos a la luna

    # --- Carga de imagen ---

    img = Image.open(f'{dir_path}/input.png')
    width, height = img.size

    # Redimensionamos si es muy grande para que la duración no sea eterna
    if max(width, height) > MAX_IMG_DIM:
        scale = MAX_IMG_DIM / max(width, height)
        new_w = int(width * scale)
        new_h = int(height * scale)
        img = img.resize((new_w, new_h), Image.BICUBIC)
        width, height = img.size
        print(f"Image resized to {width}x{height} for performance.")

    print(img.format)
    print(img.mode)
    print(img.size)

    img.show()

    # --- Scan de imagen ---

    bi_select = bidirectional_select()
    pixel_values = raster_scan(img, bi_select)

    # RGB -> luminancia (mantiene info de color pero en 1 canal)
    luma = rgb_to_luma(pixel_values)

    # Normalizar a [-0.95, 0.95]
    image_data = normalise(luma, 0.95)
    fade_audio(image_data)

    # Duración en segundos: 1 píxel = 1 muestra
    img_duration_sec = len(image_data) / sample_rate
    print(f"Image scan duration ≈ {img_duration_sec:.2f} seconds")

    # --- Selección de modo, forma de onda, frecuencia ---

    mode = mode_select()
    wf = waveform_select()
    osc_f_input = frequency_select()

    # Escala basada en la frecuencia introducida
    midi_base = int(round(freq_to_midi(osc_f_input)))
    SCALE_MIDI = build_scale(base_midi=midi_base, octaves=3)
    SCALE_FREQS = [midi_to_freq(m) for m in SCALE_MIDI]

    osc_f_quant = quantize_freq_to_scale(osc_f_input, SCALE_FREQS)
    print(f"Quantized base frequency to scale: {osc_f_quant:.2f} Hz")

    # --- Construcción del oscilador según modo ---

    if mode == 1:
        # DRONE: mismo número de muestras que image_data
        osc_len = len(image_data)

        if wf == 1:
            osc_signal = sine_gen(sample_rate, osc_f_quant, osc_len)
            wf_type = "Sine_Drone"
        elif wf == 2:
            osc_signal = square_gen(sample_rate, osc_f_quant, osc_len)
            wf_type = "Square_Drone"
        else:
            osc_signal = saw_gen(sample_rate, osc_f_quant, osc_len)
            wf_type = "Saw_Drone"

        fade_audio(osc_signal)

        # Aquí viene la magia: modulamos la amplitud del oscilador con la imagen
        audio_signal = osc_signal * image_data

    else:
        # MELODIC: la imagen decide una secuencia de notas
        osc_signal = melodic_sequence_from_image(image_data,
                                                 sample_rate,
                                                 NOTE_DURATION_SEC,
                                                 wf,
                                                 SCALE_FREQS)
        wf_type = {1: "Sine_Melodic", 2: "Square_Melodic", 3: "Saw_Melodic"}[wf]

        # En modo melódico, audio_signal es directamente la secuencia
        audio_signal = osc_signal

    fade_audio(audio_signal)

    # --- Audio setup ---

    p = pyaudio.PyAudio()

    stream = p.open(format=format,
                    channels=channels,
                    rate=sample_rate,
                    output=True)

    # Escuchamos primero el "barrido" puro de la imagen (ruido textural)
    stream.write(image_data.astype(np.float32).tobytes())
    my_plot(image_data, "Image Scan (Luma) Waveform")
    my_fft(image_data)

    # Luego el oscilador (drone o secuencia)
    stream.write(osc_signal.astype(np.float32).tobytes())
    my_plot(osc_signal, f"Oscillator ({wf_type})")
    my_fft(osc_signal)

    # Y luego el sonido final (oscilador modulado por la imagen)
    stream.write(audio_signal.astype(np.float32).tobytes())
    my_plot(audio_signal, "Final Audio (Image-modulated)")
    my_fft(audio_signal)

    # --- Menú interactivo ---

    quit_program = False

    while not quit_program:
        try:
            select = int(input(
f"""
-------------------------------------------------
Play again, export audio, or quit.
-------------------------------------------------

1. Play final audio again
2. Export final audio
3. Quit

-------------------------------------------------
Select 1/2/3
-------------------------------------------------

"""))
        except ValueError:
            print("Please enter 1, 2 or 3.")
            continue

        if select == 1:
            stream.write(audio_signal.astype(np.float32).tobytes())
            my_plot(audio_signal, "Final Audio (Image-modulated)")
            my_fft(audio_signal)

        elif select == 2:
            export_dir = f"{dir_path}/WAV Exports"
            os.makedirs(export_dir, exist_ok=True)
            export_name = f"{Path(img.filename).stem}_{wf_type}_{int(osc_f_quant)}Hz.wav"
            export_path = f"{export_dir}/{export_name}"
            write(export_path, sample_rate, audio_signal.astype(np.float32))
            print(f"File saved as {export_name}")

        elif select == 3:
            quit_program = True
            print("Shutting down...")

        else:
            print("Please enter 1, 2 or 3.")

except KeyboardInterrupt:
    print("\nCtrl+C detected. Exiting gracefully...")

finally:
    # Cleanup de audio
    try:
        if 'stream' in locals() and stream is not None:
            stream.stop_stream()
            stream.close()
        if 'p' in locals() and p is not None:
            p.terminate()
    except Exception:
        pass