#!/usr/bin/env python3
import argparse
import json
import math
import struct
import wave
from pathlib import Path

import librosa
import numpy as np
from PIL import Image, ImageDraw

NOTE_NAMES = ['C', 'C#', 'D', 'D#', 'E', 'F', 'F#', 'G', 'G#', 'A', 'A#', 'B']


def vlq(value):
    value = int(max(0, value))
    buffer = value & 0x7f
    out = bytearray()
    while value := value >> 7:
        buffer <<= 8
        buffer |= (value & 0x7f) | 0x80
    while True:
        out.append(buffer & 0xff)
        if buffer & 0x80:
            buffer >>= 8
        else:
            break
    return bytes(out)


def write_midi(path, notes, bpm=105, ticks_per_beat=480):
    tempo = round(60_000_000 / bpm)
    events = []
    for note in notes:
        start = round(note['start'] * bpm / 60 * ticks_per_beat)
        end = round(note['end'] * bpm / 60 * ticks_per_beat)
        pitch = int(note['midi'])
        velocity = int(note['velocity'])
        events.append((start, 1, bytes([0x90, pitch, velocity])))
        events.append((end, 0, bytes([0x80, pitch, 0])))
    events.sort(key=lambda x: (x[0], x[1]))
    track = bytearray()
    track += vlq(0) + b'\xff\x51\x03' + tempo.to_bytes(3, 'big')
    previous = 0
    for tick, _, event in events:
        track += vlq(tick - previous) + event
        previous = tick
    track += vlq(0) + b'\xff\x2f\x00'
    header = b'MThd' + struct.pack('>IHHH', 6, 0, 1, ticks_per_beat)
    body = b'MTrk' + struct.pack('>I', len(track)) + track
    path.write_bytes(header + body)


def analyze_audio(source, bpm):
    y, sr = librosa.load(source, sr=22050, mono=True)
    hop = 256
    onset_env = librosa.onset.onset_strength(y=y, sr=sr, hop_length=hop)
    onset_frames = librosa.onset.onset_detect(
        onset_envelope=onset_env,
        sr=sr,
        hop_length=hop,
        backtrack=True,
        delta=0.05,
        wait=1,
        pre_max=3,
        post_max=3,
        pre_avg=3,
        post_avg=3,
    )
    onset_times = librosa.frames_to_time(onset_frames, sr=sr, hop_length=hop)
    f0, voiced, probability = librosa.pyin(
        y,
        fmin=librosa.note_to_hz('C2'),
        fmax=librosa.note_to_hz('C7'),
        sr=sr,
        frame_length=2048,
        hop_length=hop,
    )
    f0_times = librosa.frames_to_time(np.arange(len(f0)), sr=sr, hop_length=hop)
    rms = librosa.feature.rms(y=y, frame_length=2048, hop_length=hop)[0]
    stft = np.abs(librosa.stft(y, n_fft=4096, hop_length=hop))
    frequencies = librosa.fft_frequencies(sr=sr, n_fft=4096)
    notes = []
    for index, start in enumerate(onset_times):
        if index + 1 < len(onset_times):
            end = float(onset_times[index + 1])
        else:
            end = float(librosa.get_duration(y=y, sr=sr))
        if end - start < 0.04:
            continue
        left = int(np.searchsorted(f0_times, start))
        right = int(np.searchsorted(f0_times, min(end, start + 0.35)))
        right = max(left + 1, right)
        pitch_slice = f0[left:right]
        prob_slice = probability[left:right]
        valid = np.isfinite(pitch_slice) & (prob_slice > 0.45)
        pitch_hz = None
        if np.any(valid):
            pitch_hz = float(np.median(pitch_slice[valid]))
        if pitch_hz is None:
            frame_start = max(0, int(round(start * sr / hop)))
            frame_end = min(stft.shape[1], max(frame_start + 1, int(round(min(end, start + 0.35) * sr / hop))))
            spectrum = np.median(stft[:, frame_start:frame_end], axis=1)
            spectrum[:max(1, int(librosa.note_to_hz('C2') / (sr / 4096)))] = 0
            peak_bins = librosa.util.peak_pick(spectrum, pre_max=3, post_max=3, pre_avg=3, post_avg=3, delta=0.05, wait=1)
            if len(peak_bins):
                peak = int(peak_bins[np.argmax(spectrum[peak_bins])])
                pitch_hz = float(frequencies[peak])
        if pitch_hz is None or pitch_hz <= 0:
            continue
        midi_float = float(librosa.hz_to_midi(pitch_hz))
        midi = int(np.clip(round(midi_float), 0, 127))
        frame = min(len(rms) - 1, max(0, int(round(start * sr / hop))))
        local_rms = float(rms[frame])
        velocity = int(np.clip(round(35 + local_rms * 180), 1, 127))
        notes.append({
            'index': len(notes),
            'start': round(float(start), 4),
            'end': round(float(end), 4),
            'duration': round(float(end - start), 4),
            'pitch_hz': round(pitch_hz, 3),
            'midi_float': round(midi_float, 3),
            'midi': midi,
            'note': f'{NOTE_NAMES[midi % 12]}{midi // 12 - 1}',
            'velocity': velocity,
        })
    return {
        'source': str(source),
        'sample_rate': sr,
        'duration': round(float(librosa.get_duration(y=y, sr=sr)), 4),
        'bpm': bpm,
        'onsets_detected': len(onset_times),
        'notes': notes,
    }


def render_svg(path, analysis, width=1600, height=900):
    notes = analysis['notes']
    duration = analysis['duration']
    pitches = [n['midi'] for n in notes]
    low = min(pitches) - 3
    high = max(pitches) + 3
    left, right, top, bottom = 100, width - 60, 60, height - 100
    def x(t): return left + t / max(duration, 0.001) * (right - left)
    def y(p): return bottom - (p - low) / max(high - low, 1) * (bottom - top)
    parts = [f'<svg xmlns="http://www.w3.org/2000/svg" width="{width}" height="{height}" viewBox="0 0 {width} {height}">', '<rect width="100%" height="100%" fill="#10141c"/>']
    for p in range(low, high + 1):
        yy = y(p)
        color = '#263142' if p % 12 else '#46546d'
        parts.append(f'<line x1="{left}" y1="{yy:.1f}" x2="{right}" y2="{yy:.1f}" stroke="{color}" stroke-width="1"/>')
        if p % 12 == 0:
            parts.append(f'<text x="20" y="{yy + 5:.1f}" fill="#aeb9ca" font-size="18">{NOTE_NAMES[p % 12]}{p // 12 - 1}</text>')
    if notes:
        points = ' '.join(f'{x(n["start"]):.1f},{y(n["midi"]):.1f}' for n in notes)
        parts.append(f'<polyline points="{points}" fill="none" stroke="#8de1cc" stroke-width="5" opacity="0.65"/>')
    for n in notes:
        xx = x(n['start'])
        yy = y(n['midi'])
        ww = max(10, x(n['end']) - xx)
        radius = 8 + n['velocity'] / 10
        parts.append(f'<rect x="{xx:.1f}" y="{yy - radius:.1f}" width="{ww:.1f}" height="{2 * radius:.1f}" rx="{radius:.1f}" fill="#f2bd68" opacity="0.88"/>')
    parts.append(f'<text x="{left}" y="35" fill="#f4f0e8" font-size="26">MIDI Landscape — {Path(analysis["source"]).name}</text>')
    parts.append(f'<text x="{left}" y="{height - 35}" fill="#aeb9ca" font-size="18">time →     pitch ↑     notes: {len(notes)}</text>')
    parts.append('</svg>')
    path.write_text('\n'.join(parts))


def render_png(path, analysis, width=1600, height=900):
    notes = analysis['notes']
    duration = analysis['duration']
    pitches = [n['midi'] for n in notes]
    low = min(pitches) - 3
    high = max(pitches) + 3
    left, right, top, bottom = 100, width - 60, 60, height - 100
    image = Image.new('RGB', (width, height), '#10141c')
    draw = ImageDraw.Draw(image)
    def x(t): return left + t / max(duration, 0.001) * (right - left)
    def y(p): return bottom - (p - low) / max(high - low, 1) * (bottom - top)
    for p in range(low, high + 1):
        yy = y(p)
        color = '#263142' if p % 12 else '#46546d'
        draw.line((left, yy, right, yy), fill=color, width=1)
        if p % 12 == 0:
            draw.text((20, yy - 8), f'{NOTE_NAMES[p % 12]}{p // 12 - 1}', fill='#aeb9ca')
    points = [(x(n['start']), y(n['midi'])) for n in notes]
    if len(points) > 1:
        draw.line(points, fill='#8de1cc', width=5)
    for n in notes:
        xx, yy = x(n['start']), y(n['midi'])
        ww = max(10, x(n['end']) - xx)
        radius = 8 + n['velocity'] / 10
        draw.rounded_rectangle((xx, yy - radius, xx + ww, yy + radius), radius=radius, fill='#f2bd68')
    draw.text((left, 20), f'MIDI Landscape — {Path(analysis["source"]).name}', fill='#f4f0e8')
    draw.text((left, height - 35), f'time →     pitch ↑     notes: {len(notes)}', fill='#aeb9ca')
    image.save(path)


def write_stl(path, analysis, width=120.0, depth=34.0, base_height=2.0):
    notes = analysis['notes']
    duration = analysis['duration']
    pitches = [n['midi'] for n in notes]
    low = min(pitches)
    high = max(pitches)
    height = 70.0
    radius_min = 1.8
    radius_max = 4.0
    triangles = []
    def tri(a, b, c): triangles.append((a, b, c))
    def box(x0, x1, y0, y1, z0, z1):
        v = [(x0,y0,z0),(x1,y0,z0),(x1,y1,z0),(x0,y1,z0),(x0,y0,z1),(x1,y0,z1),(x1,y1,z1),(x0,y1,z1)]
        for a,b,c in [(0,2,1),(0,3,2),(4,5,6),(4,6,7),(0,1,5),(0,5,4),(1,2,6),(1,6,5),(2,3,7),(2,7,6),(3,0,4),(3,4,7)]: tri(v[a],v[b],v[c])
    box(0, width, 0, depth, 0, base_height)
    for n in notes:
        cx = n['start'] / max(duration, 0.001) * width
        cy = depth / 2
        z = base_height + 4 + (n['midi'] - low) / max(high - low, 1) * height
        r = radius_min + (n['velocity'] / 127) * (radius_max - radius_min)
        segments = 16
        top = (cx, cy, z)
        bottom = (cx, cy, base_height)
        for i in range(segments):
            a0 = 2 * math.pi * i / segments
            a1 = 2 * math.pi * (i + 1) / segments
            p0 = (cx + r * math.cos(a0), cy + r * math.sin(a0), base_height)
            p1 = (cx + r * math.cos(a1), cy + r * math.sin(a1), base_height)
            q0 = (cx + r * math.cos(a0), cy + r * math.sin(a0), z)
            q1 = (cx + r * math.cos(a1), cy + r * math.sin(a1), z)
            tri(bottom, p1, p0)
            tri(top, q0, q1)
            tri(p0, p1, q1)
            tri(p0, q1, q0)
    with path.open('w') as f:
        f.write('solid midi_landscape\n')
        for a,b,c in triangles:
            ux,uy,uz = b[0]-a[0],b[1]-a[1],b[2]-a[2]
            vx,vy,vz = c[0]-a[0],c[1]-a[1],c[2]-a[2]
            nx,ny,nz = uy*vz-uz*vy, uz*vx-ux*vz, ux*vy-uy*vx
            norm = math.sqrt(nx*nx+ny*ny+nz*nz) or 1
            f.write(f' facet normal {nx/norm:.6f} {ny/norm:.6f} {nz/norm:.6f}\n  outer loop\n')
            for v in (a,b,c): f.write(f'   vertex {v[0]:.6f} {v[1]:.6f} {v[2]:.6f}\n')
            f.write('  endloop\n endfacet\n')
        f.write('endsolid midi_landscape\n')
    return {'triangles': len(triangles), 'width_mm': width, 'depth_mm': depth, 'height_mm': round(base_height + 4 + height, 3)}


def validate_stl(path):
    vertices = []
    edge_counts = {}
    with path.open() as f:
        for line in f:
            if line.strip().startswith('vertex'):
                v = tuple(round(float(x), 6) for x in line.split()[1:])
                vertices.append(v)
    for i in range(0, len(vertices), 3):
        face = vertices[i:i+3]
        if len(face) != 3: return {'valid': False, 'reason': 'incomplete facet'}
        for a,b in ((face[0],face[1]),(face[1],face[2]),(face[2],face[0])):
            edge = tuple(sorted((a,b)))
            edge_counts[edge] = edge_counts.get(edge, 0) + 1
    counts = list(edge_counts.values())
    return {'valid': all(math.isfinite(v) for point in vertices for v in point) and all(c == 2 for c in counts), 'facets': len(vertices)//3, 'unique_vertices': len(set(vertices)), 'boundary_edges': sum(c != 2 for c in counts), 'reason': 'closed manifold' if all(c == 2 for c in counts) else 'open or non-manifold edges'}


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('source')
    parser.add_argument('--out-dir', default='/home/workspace/ghojualamanchu/Projects/audio-sonar/midi-geometry')
    parser.add_argument('--bpm', type=float, default=105)
    args = parser.parse_args()
    out = Path(args.out_dir)
    out.mkdir(parents=True, exist_ok=True)
    analysis = analyze_audio(args.source, args.bpm)
    stem = Path(args.source).stem
    midi_path = out / f'{stem}.mid'
    json_path = out / f'{stem}.json'
    svg_path = out / f'{stem}.svg'
    png_path = out / f'{stem}.png'
    stl_path = out / f'{stem}.stl'
    write_midi(midi_path, analysis['notes'], args.bpm)
    render_svg(svg_path, analysis)
    render_png(png_path, analysis)
    mesh = write_stl(stl_path, analysis)
    validation = validate_stl(stl_path)
    analysis['outputs'] = { 'midi': str(midi_path), 'json': str(json_path), 'svg': str(svg_path), 'png': str(png_path), 'stl': str(stl_path), 'mesh': mesh, 'validation': validation }
    json_path.write_text(json.dumps(analysis, indent=2) + '\n')
    print(json.dumps(analysis, indent=2))
    if not validation['valid']:
        raise SystemExit('STL validation failed')

if __name__ == '__main__':
    main()