import numpy as np
from scipy.io.wavfile import write

# Parameters
SAMPLE_RATE = 44100
DURATION = 5.0          # seconds
BPM = 96
BEATS = 8               # 2 bars of 4/4
BEAT_DURATION = 60 / BPM
SECONDS_PER_BEAT = BEAT_DURATION
TOTAL_SAMPLES = int(SAMPLE_RATE * DURATION)

# Time array
t = np.linspace(0, DURATION, TOTAL_SAMPLES, endpoint=False)

# ----- Helper: note frequency -----
def note_freq(note):
    # note like 'C4' -> 261.63
    notes = {'C': 0, 'D': 2, 'E': 4, 'F': 5, 'G': 7, 'A': 9, 'B': 11}
    octave = int(note[1])
    semitone = notes[note[0]] + 12 * (octave - 4)
    return 440.0 * (2 ** (semitone / 12))

# ----- Melody (xylophone-like) -----
def xylophone(freq, duration_sec, amp=0.5):
    samples = int(duration_sec * SAMPLE_RATE)
    env = np.exp(-np.linspace(0, 4, samples))  # quick decay
    wave = np.sin(2 * np.pi * freq * np.linspace(0, duration_sec, samples, endpoint=False))
    return amp * env * wave

melody_notes = ['C4', 'E4', 'G4', 'A4', 'G4', 'E4', 'D4', 'C4']
melody_dur = BEAT_DURATION  # each quarter note
melody = np.zeros(TOTAL_SAMPLES)
for i, note in enumerate(melody_notes):
    start = int(i * BEAT_DURATION * SAMPLE_RATE)
    end = int(start + melody_dur * SAMPLE_RATE)
    seg = xylophone(note_freq(note), melody_dur, amp=0.4)
    melody[start:end] += seg[:end-start]

# ----- Bass (soft sine) -----
bass_notes = ['C2', 'F2', 'G2', 'C2']  # each 2 beats
bass_dur = 2 * BEAT_DURATION
bass = np.zeros(TOTAL_SAMPLES)
for i, note in enumerate(bass_notes):
    start = int(i * bass_dur * SAMPLE_RATE)
    end = int(start + bass_dur * SAMPLE_RATE)
    freq = note_freq(note)
    seg = 0.3 * np.sin(2 * np.pi * freq * np.linspace(0, bass_dur, end-start, endpoint=False))
    bass[start:end] += seg

# ----- Drums (simple) -----
drums = np.zeros(TOTAL_SAMPLES)
# Kick on beats 1,3,5,7 (0-indexed)
kick_beats = [0, 2, 4, 6]
for beat in kick_beats:
    start = int(beat * BEAT_DURATION * SAMPLE_RATE)
    duration = 0.1  # 100 ms
    end = int(start + duration * SAMPLE_RATE)
    if end < TOTAL_SAMPLES:
        # simple kick: decaying sine
        seg = 0.5 * np.exp(-np.linspace(0, 5, end-start)) * np.sin(2 * np.pi * 80 * np.linspace(0, duration, end-start))
        drums[start:end] += seg
# Snare on beats 2,4,6,8 (index 1,3,5,7) but careful with beat 8? We have 8 beats, last snare on beat 8 is at end? Actually beat 8 is index 7, we can include but it may overlap with end. Let's do up to beat 6 (index 5) to avoid edge.
snare_beats = [1, 3, 5]  # 2,4,6
for beat in snare_beats:
    start = int(beat * BEAT_DURATION * SAMPLE_RATE)
    duration = 0.08
    end = int(start + duration * SAMPLE_RATE)
    if end < TOTAL_SAMPLES:
        seg = 0.3 * np.random.normal(0, 1, end-start)  # noise
        drums[start:end] += seg
# Hi-hat every eighth note (0.5 beat)
hihat_beats = np.arange(0, BEATS, 0.5)
for beat in hihat_beats:
    start = int(beat * BEAT_DURATION * SAMPLE_RATE)
    duration = 0.03
    end = int(start + duration * SAMPLE_RATE)
    if end < TOTAL_SAMPLES:
        seg = 0.15 * np.random.normal(0, 1, end-start)
        drums[start:end] += seg

# ----- Mix -----
# Ensure no clipping
mix = melody + bass + drums
mix = mix / np.max(np.abs(mix)) * 0.9  # normalize

# ----- Seamless loop: fade in/out very slightly to avoid clicks -----
# Since we already have a perfect musical loop, we add a tiny 10 ms crossfade at the start/end.
fade_samples = int(0.01 * SAMPLE_RATE)
fade_in = np.linspace(0, 1, fade_samples)
fade_out = np.linspace(1, 0, fade_samples)
mix[:fade_samples] *= fade_in
mix[-fade_samples:] *= fade_out

# Convert to 16-bit PCM
mix_int16 = np.int16(mix * 32767)

# Write WAV (stereo: duplicate mono)
stereo = np.column_stack((mix_int16, mix_int16))
write('garden_loop.wav', SAMPLE_RATE, stereo)

print("✅ garden_loop.wav created – loop it in your video editor!")