#!/usr/bin/env python3
"""
build-track-v02.py

Reconstruction of the Python process used to build the final 64-bar UK garage
artefact. The raw source samples are not included in the static website. To run
this script locally, place the private sample pack at ./Samples.zip or provide
an extracted ./Samples directory containing the WAV files referenced below.

Outputs:
- ukg-v02-final.wav
- ukg-v02-final.mid
- metadata.json
"""

from pathlib import Path
import math
import random
import shutil
import zipfile

import mido
import numpy as np
import soundfile as sf
from scipy.signal import butter, lfilter, resample_poly


PROJECT_NAME = "ukg-v02"
TEMPO_BPM = 140
BARS = 64
SR = 44100
BEATS_PER_BAR = 4
SECONDS_PER_BEAT = 60 / TEMPO_BPM
BAR_SAMPLES = int(round(BEATS_PER_BAR * SECONDS_PER_BEAT * SR))
TOTAL_SAMPLES = BAR_SAMPLES * BARS
PPQ = 480
TICKS_PER_BAR = PPQ * 4

RANDOM_SEED = 7
random.seed(RANDOM_SEED)


def prepare_samples():
    """Extract Samples.zip if needed and return the sample directory."""
    here = Path(".")
    sample_dir = here / "Samples"
    zip_path = here / "Samples.zip"

    if sample_dir.exists():
        return sample_dir

    if zip_path.exists():
        extract_root = here / "_samples_extracted"
        if extract_root.exists():
            shutil.rmtree(extract_root)
        with zipfile.ZipFile(zip_path, "r") as z:
            z.extractall(extract_root)
        nested = extract_root / "Samples"
        return nested if nested.exists() else extract_root

    raise FileNotFoundError(
        "No raw samples found. Place Samples.zip or a Samples/ folder next to this script."
    )


def load_audio(path, target_sr=SR):
    data, source_sr = sf.read(path, always_2d=True, dtype="float32")
    if source_sr != target_sr:
        divisor = math.gcd(source_sr, target_sr)
        up = target_sr // divisor
        down = source_sr // divisor
        data = np.stack(
            [resample_poly(data[:, channel], up, down).astype(np.float32)
             for channel in range(data.shape[1])],
            axis=1,
        )
    if data.shape[1] == 1:
        data = np.repeat(data, 2, axis=1)
    return data[:, :2]


def fade(signal, fade_in=0.0, fade_out=0.0):
    signal = signal.copy()
    if fade_in:
        length = min(len(signal), int(fade_in * SR))
        signal[:length] *= np.linspace(0, 1, length)[:, None]
    if fade_out:
        length = min(len(signal), int(fade_out * SR))
        signal[-length:] *= np.linspace(1, 0, length)[:, None]
    return signal


def normalize(signal, peak=1.0):
    max_value = np.max(np.abs(signal))
    return signal if max_value == 0 else signal * (peak / max_value)


def filter_audio(signal, cutoff, mode):
    b, a = butter(2, cutoff / (SR / 2), mode)
    return lfilter(b, a, signal, axis=0).astype(np.float32)


def high_pass(signal, cutoff):
    return filter_audio(signal, cutoff, "highpass")


def low_pass(signal, cutoff):
    return filter_audio(signal, cutoff, "lowpass")


def fit_to_length(signal, samples):
    if len(signal) >= samples:
        return signal[:samples]
    repeats = int(np.ceil(samples / len(signal)))
    return np.tile(signal, (repeats, 1))[:samples]


def bar_start(bar):
    return bar * BAR_SAMPLES


def step_start(bar, step):
    base = int(round((bar * 16 + step) * SECONDS_PER_BEAT / 4 * SR))
    if step % 2 == 1:
        base += int(round(0.08 * (SECONDS_PER_BEAT / 2) * SR))
    return base


def midi_tick(bar, step):
    tick = bar * TICKS_PER_BAR + step * (PPQ // 4)
    if step % 2 == 1:
        tick += int(0.08 * (PPQ / 2))
    return tick


def add_audio(destination, start, signal, gain=1.0, pan=0.0):
    if start >= len(destination):
        return
    signal = signal.astype(np.float32).copy()
    left = math.cos((pan + 1) * math.pi / 4)
    right = math.sin((pan + 1) * math.pi / 4)
    signal[:, 0] *= left
    signal[:, 1] *= right
    end = min(len(destination), start + len(signal))
    if end > start:
        destination[start:end] += signal[:end - start] * gain


def add_midi(events, note, tick, duration=60, velocity=90):
    events.append((tick, "on", note, velocity))
    events.append((tick + duration, "off", note, 0))


def build():
    sample_dir = prepare_samples()
    source = {path.name: load_audio(path) for path in sample_dir.glob("*.wav")}

    mix = np.zeros((TOTAL_SAMPLES + SR * 2, 2), dtype=np.float32)
    events = []

    kick_hit = fade(normalize(source["Kick Fm.wav"][: int(0.32 * SR)], 0.9), fade_out=0.08)
    snare_hit = fade(normalize(source["Snare Fm.wav"][: int(0.42 * SR)], 0.85), fade_out=0.10)
    hat_source = high_pass(source["Hats Fm.wav"], 3500)
    hat_hit = fade(normalize(hat_source[: int(0.12 * SR)], 0.55), fade_out=0.05)
    open_hat = fade(normalize(hat_source[: int(0.38 * SR)], 0.50), fade_out=0.18)

    beat_loop = fade(normalize(fit_to_length(source["Beat.wav"], 2 * BAR_SAMPLES), 0.85), 0.01, 0.01)
    pad_loop = fade(normalize(fit_to_length(source["Gmin Pad.wav"], 4 * BAR_SAMPLES), 0.80), 0.03, 0.08)
    moody_loop = fade(normalize(fit_to_length(source["Moody 140.wav"], 4 * BAR_SAMPLES), 0.80), 0.02, 0.08)
    sevens_loop = fade(normalize(fit_to_length(source["Sevens 140.wav"], 4 * BAR_SAMPLES), 0.80), 0.02, 0.08)
    clang_loop = fade(normalize(fit_to_length(source["Fm Clang.wav"], 8 * BAR_SAMPLES), 0.75), 0.05, 0.20)
    bam = fade(normalize(source["Bam.wav"][: 2 * BAR_SAMPLES], 0.80), fade_out=0.15)
    plink = fade(normalize(source["G6 Plinkity.wav"][: 2 * BAR_SAMPLES], 0.80), fade_out=0.12)
    bap = fade(normalize(source["G3 Bap de liddle la.wav"][: 4 * BAR_SAMPLES], 0.80), 0.02, 0.12)

    c3 = fade(low_pass(normalize(fit_to_length(source["C3.wav"], 16 * BAR_SAMPLES), 0.75), 5000), 0.10, 0.30)
    c4 = fade(low_pass(normalize(fit_to_length(source["C4.wav"], 16 * BAR_SAMPLES), 0.75), 7000), 0.10, 0.30)
    c5 = fade(high_pass(normalize(fit_to_length(source["C5.wav"], 16 * BAR_SAMPLES), 0.75), 180), 0.10, 0.30)

    for bar in range(0, BARS, 4):
        phrase = bar // 4
        if phrase in [0, 1, 2, 3, 8, 9, 10, 11, 12, 13, 14, 15]:
            add_audio(mix, bar_start(bar), pad_loop, gain=0.14 if bar < 32 else 0.18, pan=-0.08)
        if phrase in [2, 3, 6, 7, 10, 11, 14, 15]:
            add_audio(mix, bar_start(bar), sevens_loop, gain=0.10, pan=0.20)
        if phrase in [4, 5, 6, 7, 12, 13, 14, 15]:
            add_audio(mix, bar_start(bar), moody_loop, gain=0.17, pan=-0.18)

    for bar, signal, gain, pan in [(0, c3, 0.08, -0.10), (16, c4, 0.06, 0.12), (32, c3, 0.09, -0.14), (48, c5, 0.055, 0.18)]:
        add_audio(mix, bar_start(bar), signal, gain=gain, pan=pan)

    for bar in [8, 24, 40, 56]:
        add_audio(mix, bar_start(bar), clang_loop, gain=0.09 if bar != 56 else 0.13, pan=0.05)

    for bar in list(range(16, 24, 2)) + list(range(32, 40, 2)) + list(range(48, 64, 2)):
        add_audio(mix, bar_start(bar), high_pass(beat_loop, 120), gain=0.18 if bar < 48 else 0.24)

    for bar in range(BARS):
        bar_number = bar + 1
        phrase_8 = (bar // 8) % 8
        within_8 = bar % 8
        within_4 = bar % 4

        kick_steps = {0, 10}
        kick_steps.add(6 if phrase_8 % 2 == 0 else 7)
        if within_4 in (1, 2):
            kick_steps.add(13)
        if within_8 in (3, 7):
            kick_steps.discard(10)
            kick_steps.update([9, 12])
        if bar_number % 16 in (12, 16):
            kick_steps.add(15)
        if bar_number in (32, 64):
            kick_steps.update([14, 15])

        for step in sorted(kick_steps):
            velocity = 105 if step == 0 else random.choice([76, 82, 88, 94])
            add_audio(mix, step_start(bar, step), kick_hit, gain=0.65 * (velocity / 100))
            add_midi(events, 36, midi_tick(bar, step), 70, velocity)

        for step in [4, 12]:
            velocity = 96 if step == 4 else 102
            add_audio(mix, step_start(bar, step), snare_hit, gain=0.48 * (velocity / 100), pan=0.02)
            add_midi(events, 38, midi_tick(bar, step), 70, velocity)

        for step in range(16):
            if step in (4, 12) and random.random() < 0.2:
                continue
            velocity = 42 + (12 if step % 4 == 0 else 0) + (9 if step % 2 else 0) + random.randint(-5, 5)
            add_audio(mix, step_start(bar, step), hat_hit, gain=0.18 * (velocity / 60), pan=random.uniform(-0.12, 0.12))
            add_midi(events, 42, midi_tick(bar, step), 36, max(20, min(80, velocity)))

        open_hat_steps = [6, 14]
        if phrase_8 % 3 == 1:
            open_hat_steps = [2, 10, 14]
        if within_4 == 3:
            open_hat_steps.append(15)
        for step in open_hat_steps:
            if random.random() < 0.80:
                add_audio(mix, step_start(bar, step), open_hat, gain=0.16, pan=random.uniform(-0.2, 0.2))
                add_midi(events, 46, midi_tick(bar, step), 100, 58)

        for step in [3, 7, 11, 15]:
            probability = 0.28 + (0.16 if within_4 == 3 else 0) + (0.18 if within_8 == 7 else 0)
            if random.random() < probability:
                ghost = fade(snare_hit[: int(0.12 * SR)] * 0.5, fade_out=0.04)
                add_audio(mix, step_start(bar, step), ghost, gain=0.14, pan=random.uniform(-0.08, 0.08))
                add_midi(events, 37, midi_tick(bar, step), 45, random.randint(24, 42))

        if within_4 == 3:
            fill_steps = [13, 15] if bar_number % 8 else [10, 11, 13, 15]
            if bar_number % 16 == 0:
                fill_steps = [8, 10, 11, 13, 14, 15]
            for index, step in enumerate(fill_steps):
                add_audio(mix, step_start(bar, step), snare_hit[: int(0.20 * SR)], gain=0.22 + 0.04 * index)
                add_midi(events, 38 if index % 2 else 37, midi_tick(bar, step), 45, 55 + index * 8)

        if bar_number % 8 == 1 and bar != 0:
            add_audio(mix, bar_start(bar), bam, gain=0.13, pan=-0.20)
            add_midi(events, 49, bar * TICKS_PER_BAR, 240, 72)
        if bar_number % 16 == 9:
            add_audio(mix, bar_start(bar), plink, gain=0.12, pan=0.25)
        if bar_number % 16 == 1 and bar != 0:
            add_audio(mix, bar_start(bar), bap, gain=0.10)

    for bar in [3, 7, 11, 15, 23, 31, 39, 47, 55, 63]:
        tail = plink[: int(0.42 * SR)][::-1] * np.linspace(0, 1, int(0.42 * SR))[:, None]
        add_audio(mix, bar_start(bar + 1) - len(tail), tail, gain=0.045, pan=0.18)

    final = mix[:TOTAL_SAMPLES]
    final = high_pass(final, 25)
    final = np.tanh(final * 1.15)
    final = final / (np.max(np.abs(final)) or 1) * 0.92

    out_dir = Path("dist")
    out_dir.mkdir(exist_ok=True)
    sf.write(out_dir / f"{PROJECT_NAME}-final.wav", final, SR)

    midi = mido.MidiFile(type=1, ticks_per_beat=PPQ)
    track = mido.MidiTrack()
    midi.tracks.append(track)
    track.append(mido.MetaMessage("track_name", name="64-bar UKG drums revised for uploaded samples", time=0))
    track.append(mido.MetaMessage("set_tempo", tempo=mido.bpm2tempo(TEMPO_BPM), time=0))
    track.append(mido.MetaMessage("time_signature", numerator=4, denominator=4, time=0))

    events.sort(key=lambda event: (event[0], 0 if event[1] == "off" else 1))
    last_tick = 0
    for tick, event_type, note, velocity in events:
        delta = max(0, tick - last_tick)
        last_tick = tick
        if event_type == "on":
            track.append(mido.Message("note_on", channel=9, note=note, velocity=int(velocity), time=delta))
        else:
            track.append(mido.Message("note_off", channel=9, note=note, velocity=0, time=delta))
    track.append(mido.MetaMessage("end_of_track", time=max(0, BARS * TICKS_PER_BAR - last_tick)))
    midi.save(out_dir / f"{PROJECT_NAME}-final.mid")

    metadata = {
        "project": PROJECT_NAME,
        "tempo_bpm": TEMPO_BPM,
        "bars": BARS,
        "sample_rate": SR,
        "raw_samples_included": False,
        "method": "Python-generated arrangement using private local samples, MIDI event scheduling, sample chopping, phrase-level variation, and offline audio rendering."
    }
    (out_dir / "metadata.json").write_text(__import__("json").dumps(metadata, indent=2))


if __name__ == "__main__":
    build()