// "Pop Music Generator" - Rust Edition
// David Westgate 2024
// Initial python implementation, tests and some comments from Bart Massey
// https://github.com/pdx-cs-sound/popgen/blob/main/popgen.py

use clap::Parser;
use hound;
use rand::Rng;
use regex::Regex;
use rodio::{buffer::SamplesBuffer, OutputStream, Sink};
use std::{f32::consts::PI, vec};

const C5: &str = "C[5]";

// 11 canonical note names.
const NAMES: [&str; 12] = [
    "C", "Db", "D", "Eb", "E", "F", "Gb", "G", "Ab", "A", "Bb", "B",
];

const NO_OUTPUT: &str = "";

const REGEX_BASE: &str = r"([A-G]b?)(\[([0-8])\])?";

// Relative notes of a major scale.
const MAJOR_SCALE: [u8; 7] = [0, 2, 4, 5, 7, 9, 11];

// Major chord scale tones — one-based.
const MAJOR_CHORD: [u8; 3] = [1, 3, 5];

// Root note offset for each chord in scale tones — one-based.
const CHORD_LOOP: [u8; 4] = [8, 5, 6, 4];

// Given a MIDI key number and an optional number of beats of
// note duration, return a sine wave for that note.
fn make_note(key: i8, n: Option<u8>, beat_samples: u32, samplerate: u16) -> Vec<f32> {
    let num_beats: f32 = n.unwrap_or(1).try_into().unwrap();
    let samplerate_f: f32 = samplerate.try_into().unwrap();
    let beatsamples_f: f32 = beat_samples as f32;
    let key_f: f32 = key.try_into().unwrap();
    let f: f32 = 440f32 * (2f32).powf((key_f - 69f32) / 12f32);
    let b: u32 = (beat_samples as u32) * (n.unwrap_or(1) as u32);
    let b_f: f32 = beatsamples_f * num_beats;

    let cycles: f32 = 2f32 * PI * f * b_f / samplerate_f;
    (0..b)
        .map(|v: u32| ((v as f32) / b_f * cycles).sin())
        .collect::<Vec<f32>>()
}

// Floor division of possibly negative values (GPT generated)
// Intended to match pythons // operator
fn floor_div(a: i8, b: i8) -> i8 {
    let div = a / b;
    let rem = a % b;

    if rem != 0 && (a < 0) != (b < 0) {
        div - 1
    } else {
        div
    }
}

// modulus of possibly negative value (GPT generated)
fn safe_mod_index(value: i8, modulus: usize) -> usize {
    assert!(modulus > 0, "Modulus must be positive");

    // Compute the modulo
    let remainder = value % modulus as i8;

    // Adjust for negative remainders
    let adjusted = if remainder < 0 {
        remainder + modulus as i8
    } else {
        remainder
    };

    // Convert to usize
    adjusted as usize
}

//  Given a scale note with root note 0, return a key offset
//  from the corresponding root MIDI key.
fn note_to_key_offset(note: i8) -> i8 {
    let scale_degree: usize = safe_mod_index(note, 7);
    let major_scale = MAJOR_SCALE[scale_degree];
    (floor_div(note, 7) * 12) + major_scale as i8
}

//  Given a position within a chord, return a scale note
//  offset — zero-based.
fn chord_to_note_offset(posn: i8) -> i8 {
    let chord_posn: usize = safe_mod_index(posn, 3);
    let major_chord = MAJOR_CHORD[chord_posn];
    floor_div(posn, 3) * 7 + (major_chord as i8) - 1
}

// Choose four random notes (one measure) with proper chord offset from the base note
fn pick_notes(chord_root: u8, position: &mut i8, rng: &mut rand::prelude::ThreadRng) -> [i8; 4] {
    let mut notes: [i8; 4] = [0; 4]; //Vec would work too, but less memory efficient
    let mut p: i8 = *position;
    for note in &mut notes {
        let chord_note_offset: i8 = chord_to_note_offset(p);
        let chord_note: i8 = note_to_key_offset((chord_root as i8) + chord_note_offset);
        *note = chord_note;
        if rng.gen::<f32>() > 0.5 {
            p = p + 1;
        } else {
            p = p - 1;
        }
    }

    *position = p;
    notes
}

//  Turn a note name into a corresponding MIDI key number.
//  Format is name with optional bracketed octave, for example
//  "D" or "Eb[5]". Default is octave 4 if no octave is
//  specified.
fn parse_note(note_str: &str, regex: Regex) -> u32 {
    let mut octave: u32 = 4;
    match regex.find(note_str) {
        Some(_) => {}
        None => {
            eprintln!("Invalid Note: {}", note_str);
            std::process::exit(1)
        }
    };
    let chars: std::str::Chars<'_> = note_str.chars();
    let mut it: std::str::Chars<'_> = chars.into_iter();
    let base_note: char = it.next().unwrap();
    let mut flat: bool = false;
    // Logic to Handle all valid cases, like C, C[4], Db, Db[6], etc.
    match it.next() {
        Some(c) => {
            if c == 'b' {
                flat = true;
                match it.next() {
                    Some(c) => {
                        if c == '[' {
                            octave = it.next().unwrap().to_digit(10).unwrap();
                        } else {
                            eprintln!("Invalid Notation: {}", note_str);
                            std::process::exit(1)
                        }
                    }
                    None => {}
                }
            } else if c == '[' {
                octave = it.next().unwrap().to_digit(10).unwrap();
            } else {
                eprintln!("Invalid Notation: {}", note_str);
                std::process::exit(1)
            }
        }
        None => {}
    }

    // Different string format for flat vs not
    let index: usize = match flat {
        true => NAMES.iter().position(|&n| n == format!("{}b", base_note)),
        false => NAMES.iter().position(|&n| n == base_note.to_string()),
    }
    .unwrap();
    let index_32: u32 = index as u32;
    let r: u32 = index_32 + (12 * octave);
    return r;
}

#[derive(Parser, Debug)]
#[command(version, about, long_about = None)]
struct Args {
    #[arg(short, long, default_value_t = 90)]
    bpm: u8,

    #[arg(short, long, default_value_t = 48_000)]
    samplerate: u16,

    #[arg(short, long, default_value_t = C5.to_string())]
    root: String,

    #[arg(long, default_value_t = 2)]
    bassoctave: u32,

    #[arg(long, default_value_t = 5)]
    balance: i8,

    #[arg(short, long, default_value_t = -3)]
    gain: i8,

    #[arg(short, long, default_value_t = NO_OUTPUT.to_string())] // Ideally there is a better way than this
    output: String,
}

// Rodio example https://github.com/RustAudio/rodio/blob/master/examples/basic.rs
fn play(samples: Vec<f32>, samplerate: u16, gain: i8) {
    // TODO: apply gain
    let (_stream, stream_handle) = OutputStream::try_default().unwrap();
    let source: SamplesBuffer<f32> = SamplesBuffer::new(1, samplerate as u32, samples);
    let sink: Sink = Sink::try_new(&stream_handle).unwrap(); //sink makes it easy to keep the program asleep so the whole song plays
    sink.append(source);
    sink.sleep_until_end();
}

fn save(samples: Vec<f32>, samplerate: u16, output: &str, gain: i8) {
    // TODO: apply gain
    let spec: hound::WavSpec = hound::WavSpec {
        channels: 1,
        sample_rate: samplerate as u32,
        bits_per_sample: 32,
        sample_format: hound::SampleFormat::Float,
    };
    let mut writer = hound::WavWriter::create(output, spec).unwrap();
    for s in samples {
        writer.write_sample(s).unwrap();
    }
}

fn main() {
    let args: Args = Args::parse();
    let note_name_regexp: Regex = Regex::new(REGEX_BASE).unwrap();
    let mut rng: rand::prelude::ThreadRng = rand::thread_rng();
    let mut position: i8 = 0;

    // Tempo in beats per minute.
    let bpm: u8 = args.bpm;

    // Audio sample rate in samples per second.
    let samplerate: u16 = args.samplerate;

    let melody_gain = args.balance;

    // MIDI key where melody goes.
    let melody_root: u32 = parse_note(&args.root, note_name_regexp);
    let melody_root_i8: i8 = melody_root.try_into().unwrap();

    // Bass MIDI key is below melody root.
    let bass_root: i8 = ((melody_root as i32) - (12i32 * args.bassoctave as i32))
        .try_into()
        .unwrap();

    // Samples per beat.
    let beat_samples: u32 = ((samplerate as f32) / (f32::from(bpm) / 60f32)).round() as u32;

    let mut sound: Vec<f32> = vec![];

    for chord_root in CHORD_LOOP {
        let notes: [i8; 4] = pick_notes(chord_root - 1, &mut position, &mut rng);

        let mut melody: Vec<Vec<f32>> = vec![];
        for note in notes {
            let melody_note: Vec<f32> =
                make_note(note + melody_root_i8, None, beat_samples, samplerate);
            melody.push(melody_note);
        }

        let bass_note: i8 = note_to_key_offset((chord_root - 1) as i8);
        let bass: Vec<f32> = make_note(bass_note + bass_root, Some(4), beat_samples, samplerate);
        let bass_gain: i8 = 1 - melody_gain;

        // Unlike pythons array arithmetic, here we must unravel and mutate our data structures to apply the gains in the correct place.
        let flattened_melody: Vec<f32> = melody.into_iter().flatten().collect();

        let scaled_melody: Vec<f32> = flattened_melody
            .iter()
            .map(|&x| x * melody_gain as f32)
            .collect();
        let scaled_bass: Vec<f32> = bass.iter().map(|&x| x * bass_gain as f32).collect();

        let paired_notes: Vec<(f32, f32)> = scaled_melody
            .iter()
            .zip(scaled_bass.iter())
            .map(|(&m, &b)| (m, b))
            .collect();

        // Finally, combine back together
        let combined_waveform: Vec<f32> = paired_notes.iter().map(|(m, b)| m + b).collect();

        sound.extend(combined_waveform);
    }
    if args.output.eq(NO_OUTPUT) {
        play(sound, samplerate, args.gain);
    } else {
        save(sound, samplerate, &args.output, args.gain);
    }
}

#[test]
fn test() {
    let note_tests: [(i8, i8); 11] = [
        (-9, -15),
        (-8, -13),
        (-7, -12),
        (-6, -10),
        (-2, -3),
        (-1, -1),
        (0, 0),
        (6, 11),
        (7, 12),
        (8, 14),
        (9, 16),
    ];

    for (n, k) in note_tests {
        let k0 = note_to_key_offset(n);
        assert!(
            k0 == k,
            "note_to_key_offset: for note {}, expected offet {}. Got offset {}",
            n,
            k,
            k0
        );
    }

    let chord_tests: [(i8, i8); 8] = [
        (-3, -7),
        (-2, -5),
        (-1, -3),
        (0, 0),
        (1, 2),
        (2, 4),
        (3, 7),
        (4, 9),
    ];

    for (n, c) in chord_tests {
        let c0 = chord_to_note_offset(n);
        assert!(
            c0 == c,
            "chord_to_note_offset: for chord posn {}, expected scale note {}. Got scale note {}",
            n,
            c,
            c0
        )
    }

    // Random selction of notes to test parse
    let parsed_notes_tests = [
        ("Db", 49),
        ("Db[4]", 49),
        ("C", 48),
        ("C[5]", 60),
        ("A[1]", 21),
        ("F[7]", 89),
        ("Gb[3]", 42),
        ("Gb", 54),
        ("G[1]", 19),
    ];
    let note_name_regexp: Regex = Regex::new(REGEX_BASE).unwrap();
    for (n, v) in parsed_notes_tests {
        let v0 = parse_note(n, note_name_regexp.clone());
        assert!(v0 == v, "parse note: {}, expected {}. Got {}", n, v, v0)
    }

    // Random collection of some spot checks relative to make_note frequency generations
    // make_note(key: i8, n: Option<u8>, beat_samples: u32 , samplerate: u16) -> Vec<f64>
    // TODO:
    //  1) More tests, with a variety of all parameters
    //  2) Improve some sort of rounding/truncating with the long floats
    let make_notes_tests: [(i8, Option<u8>, u32, u16, [f32; 4]); 1] = [(
        48,
        Some(1),
        288000,
        48_000,
        [
            0.0,
            0.017122516431822686,
            0.034240012506541136,
            0.05134746933903016,
        ],
    )];
    for (i, (key, n, beat_samples, samplerate, results)) in make_notes_tests.iter().enumerate() {
        let wave = make_note(*key, *n, *beat_samples, *samplerate);
        for (j, expected_sample) in results.iter().enumerate() {
            let result_sample = wave.get(j).unwrap();
            assert!(
                *expected_sample == *result_sample,
                "Make note: {:?}, expected {}. Got {}",
                make_notes_tests[i],
                *expected_sample,
                *result_sample
            )
        }
    }
}