# GENAI Experiment 6: AI music generation (Markov chain) + WAV, waveform, spectrogram (Jukebox idea)
# ===========================================================================

# # Part 1 — Training melodies (MIDI note numbers, 60 = middle C) and the Markov model

import numpy as np, wave, matplotlib.pyplot as plt
try:
    from IPython.display import Audio, display      # plays audio inside Jupyter / Colab
except ImportError:
    Audio = lambda f: f"(saved audio file: {f})"; display = print
rng = np.random.default_rng(0)

melodies = {
    # C-major, calm, stepwise (classical style)
    "classical": [60, 62, 64, 65, 67, 65, 64, 62, 60, 64, 67, 72, 71, 69, 67, 65, 64, 62, 64, 60,
                  67, 69, 71, 72, 71, 69, 67, 65, 64, 65, 67, 64, 62, 60],
    # C blues scale (C Eb F F# G Bb), jumpy (jazz / blues style)
    "jazz":      [60, 63, 65, 66, 67, 70, 72, 70, 67, 66, 65, 63, 60, 63, 67, 70, 67, 65, 63, 65,
                  66, 67, 63, 60, 70, 72, 75, 72, 70, 67, 63, 60],
}

def train_markov(seq):
    notes = sorted(set(seq)); idx = {n: i for i, n in enumerate(notes)}
    counts = np.ones((len(notes), len(notes))) * 0.01               # tiny smoothing
    for a, b in zip(seq[:-1], seq[1:]):
        counts[idx[a], idx[b]] += 1                                 # count transitions
    probs = counts / counts.sum(1, keepdims=True)                   # P(next | current)
    return notes, probs

models = {g: train_markov(m) for g, m in melodies.items()}
notes, P = models["classical"]
print("classical notes:", notes); print("P(next | 60) =", {n: round(float(p), 2) for n, p in zip(notes, P[0])})

# # Part 2 — Generate a new melody (temperature controls randomness)

def generate_melody(genre, n=24, temperature=1.0, seed=0):
    rng = np.random.default_rng(seed)
    notes, P = models[genre]
    cur = rng.integers(len(notes)); out = [notes[cur]]
    for _ in range(n - 1):
        p = P[cur] ** (1 / temperature); p /= p.sum()               # temperature
        cur = rng.choice(len(notes), p=p); out.append(notes[cur])
    return out

print("classical:", generate_melody("classical"))
print("jazz     :", generate_melody("jazz"))

# # Part 3 — Synthesise audio and save as WAV

SR = 22050
def synth(melody, tempo_bpm=90, swing=False):
    beat = 60 / tempo_bpm; audio = []
    for k, m in enumerate(melody):
        dur = beat * (0.66 if (swing and k % 2 == 0) else (0.33 if swing else 0.5)) * 2   # swing: long-short
        t = np.arange(int(SR * dur)) / SR
        f = 440 * 2 ** ((m - 69) / 12)                                     # MIDI -> Hz
        tone = np.sin(2*np.pi*f*t) + 0.5*np.sin(2*np.pi*2*f*t) + 0.25*np.sin(2*np.pi*3*f*t)   # harmonics
        env = np.minimum(1, t / 0.02) * np.exp(-3 * t)                     # attack + decay envelope
        audio.append(tone * env)
    audio = np.concatenate(audio)
    return 0.8 * audio / np.abs(audio).max()

def save_wav(path, audio):
    with wave.open(path, "wb") as w:
        w.setnchannels(1); w.setsampwidth(2); w.setframerate(SR)
        w.writeframes((audio * 32767).astype(np.int16).tobytes())

settings = {"jazz": dict(tempo_bpm=110, swing=True), "classical": dict(tempo_bpm=80, swing=False)}  # mood: relaxing = slow
files = {}
for genre in ["jazz", "classical"]:
    audio = synth(generate_melody(genre, n=24, temperature=0.9), **settings[genre])
    files[genre] = f"generated_{genre}_sample.wav"; save_wav(files[genre], audio)
    print("saved", files[genre]); display(Audio(files[genre]))

# # Part 4 — Verify the files (not corrupted, audible)

def load_wav(path):
    with wave.open(path, "rb") as w:
        sr, n = w.getframerate(), w.getnframes()
        y = np.frombuffer(w.readframes(n), dtype=np.int16).astype(np.float32) / 32768
    return y, sr

signals = {}
for g, f in files.items():
    y, sr = load_wav(f); signals[g] = y
    rms_db = 20 * np.log10(np.sqrt(np.mean(y**2)) + 1e-12)
    ok = len(y) > sr and np.isfinite(y).all() and rms_db > -40
    print(f"{f}: {sr} Hz, {len(y)/sr:.1f} s, peak {np.abs(y).max():.2f}, RMS {rms_db:.1f} dBFS -> {'OK' if ok else 'PROBLEM'}")

# # Part 5 — Waveform and spectrogram, comparing genres

fig, ax = plt.subplots(2, 2, figsize=(13, 6.5))
for r, (g, y) in enumerate(signals.items()):
    t = np.arange(len(y)) / SR
    ax[r, 0].plot(t, y, lw=0.3); ax[r, 0].set_title(f"{g}: waveform (amplitude vs time)"); ax[r, 0].set_xlabel("time (s)")
    ax[r, 1].specgram(y, NFFT=1024, Fs=SR, noverlap=512, cmap="magma")
    ax[r, 1].set_ylim(0, 3000); ax[r, 1].set_title(f"{g}: spectrogram (frequency vs time)"); ax[r, 1].set_xlabel("time (s)"); ax[r, 1].set_ylabel("Hz")
plt.tight_layout(); plt.show()

for g, y in signals.items():
    spec = np.abs(np.fft.rfft(y)); freqs = np.fft.rfftfreq(len(y), 1 / SR)
    centroid = (freqs * spec).sum() / spec.sum()
    print(f"{g}: duration {len(y)/SR:.1f} s, spectral centroid {centroid:.0f} Hz")

# Exam explanation: the Markov chain learned which note usually follows which in each genre's melody
# and sampled a new melody autoregressively; temperature controls how adventurous it is. In the
# waveform, each note is a burst with a sharp attack and an exponential decay. Jazz has an uneven
# long-short (swing) rhythm, classical a steady, slower one. In the spectrogram, every note is a
# horizontal line (fundamental) with fainter harmonics above it. The jazz blues scale (E♭, F♯, B♭)
# gives different pitch lines from the classical C-major scale.
