#![no_std]
#![no_main]

// Several lines of code are adapted or copied from Microbit-spectrum, by Bart Massey. These will be noted
// https://github.com/BartMassey/microbit-spectrum/blob/main/src/main.rs

use core::f32::consts::PI;
use cortex_m_rt::entry;
use critical_section_lock_mut::LockMut;
use embedded_hal::digital::OutputPin;
use libm::{floorf, log2f};
use microbit::{
    board::Board,
    display::nonblocking::{BitImage, Display},
    hal::{
        gpio::{p0::P0_05, Floating, Input},
        saadc::SaadcConfig,
        timer::Timer,
        Saadc,
    },
    pac::TIMER1,
};
use microfft::real::rfft_64 as rfft;
use num_complex::Complex;
use num_traits::Float;
use ordered_float::NotNan;
use panic_rtt_target as _;
use rtt_target::{rprintln, rtt_init_print};

// Microbit-spectrum
const FFT_WIDTH: usize = 64;
/// ADC is signed; it is configured for a sample precision of 10 bits. Must match below.
const ADC_MAX: usize = 512;

const NOTES: [&str; 12] = [
    "A", "B♭", "B", "C", "D♭", "D", "E♭", "E", "F", "G♭", "G", "A♭",
];
static DISPLAY: LockMut<Display<TIMER1>> = LockMut::new();
const DEFAULT_BITMAP: [[u8; 5]; 5] = [
    [9, 0, 0, 0, 9],
    [0, 9, 0, 9, 0],
    [0, 0, 9, 0, 0],
    [0, 9, 0, 9, 0],
    [9, 0, 0, 9, 9],
];
// Bit representations of the 7 note, A to G in order
const NOTE_IMAGE_BITMAPS: [[[u8; 5]; 5]; 7] = [
    [
        [0, 9, 9, 0, 0],
        [9, 0, 0, 9, 0],
        [9, 9, 9, 9, 0],
        [9, 0, 0, 9, 0],
        [9, 0, 0, 9, 0],
    ],
    [
        [9, 9, 9, 0, 0],
        [9, 0, 0, 9, 0],
        [9, 9, 9, 0, 0],
        [9, 0, 0, 9, 0],
        [9, 9, 9, 0, 0],
    ],
    [
        [0, 9, 9, 9, 0],
        [9, 0, 0, 0, 0],
        [9, 0, 0, 0, 0],
        [9, 0, 0, 0, 0],
        [0, 9, 9, 9, 0],
    ],
    [
        [9, 9, 9, 0, 0],
        [9, 0, 0, 9, 0],
        [9, 0, 0, 9, 0],
        [9, 0, 0, 9, 0],
        [9, 9, 9, 0, 0],
    ],
    [
        [9, 9, 9, 9, 0],
        [9, 0, 0, 0, 0],
        [9, 9, 9, 0, 0],
        [9, 0, 0, 0, 0],
        [9, 9, 9, 9, 0],
    ],
    [
        [9, 9, 9, 9, 0],
        [9, 0, 0, 0, 0],
        [9, 9, 9, 0, 0],
        [9, 0, 0, 0, 0],
        [9, 0, 0, 0, 0],
    ],
    [
        [0, 9, 9, 9, 0],
        [9, 0, 0, 0, 0],
        [9, 0, 9, 9, 0],
        [9, 0, 0, 9, 0],
        [0, 9, 9, 9, 0],
    ],
];

fn read_samples(saadc: &mut Saadc, mic_in: &mut P0_05<Input<Floating>>) -> [i16; FFT_WIDTH] {
    let mut result = [0; FFT_WIDTH];
    for r in &mut result {
        *r = saadc.read_channel(mic_in).unwrap();
    }
    result
}

// Microbit-spectrum - read function
fn samples_to_freq(
    samples: [i16; FFT_WIDTH],
    window: [f32; FFT_WIDTH],
) -> [Complex<f32>; FFT_WIDTH / 2] {
    const fn nn(f: f32) -> NotNan<f32> {
        unsafe { NotNan::new_unchecked(f) }
    }
    let mut sample_buf = [0.0f32; FFT_WIDTH];
    for (i, s) in sample_buf.iter_mut().enumerate() {
        *s = samples[i].abs() as f32 / ADC_MAX as f32
    }
    let dc = sample_buf.iter().copied().sum::<f32>() / sample_buf.len() as f32;
    let peak = sample_buf
        .iter()
        .copied()
        .map(nn)
        .max()
        .unwrap()
        .into_inner();
    for (i, s) in sample_buf.iter_mut().enumerate() {
        *s = window[i] * (*s - dc) / peak;
    }

    // Take the FFT and process the result.
    let freqs: &mut [Complex<f32>; FFT_WIDTH / 2] = rfft(&mut sample_buf);
    *freqs
}

fn get_dominant_freq(_freqs: [Complex<f32>; FFT_WIDTH / 2]) -> f32 {
    // todo!();
    0.0
}

fn get_note(freq: f32) -> (usize, f32) {
    if freq > 0.0 {
        let note_num = log2f(freq / 440.0) * 12.0;
        let octave = floorf(4.0 + note_num / 12.0);
        let rel_note = (note_num % 12.0) as usize;
        return (rel_note, octave);
    }
    (0, 0.0)
}

fn get_image(noctive: (usize, f32)) -> [[u8; 5]; 5] {
    let (note, _octive) = noctive; // TODO - Add octive and flat modifier to display
    if note > 0 {
        NOTE_IMAGE_BITMAPS[note]
    } else {
        DEFAULT_BITMAP
    }
}

#[entry]
fn init() -> ! {
    rtt_init_print!();
    let board: Board = Board::take().unwrap();
    let mut timer = Timer::new(board.TIMER0);

    let display: Display<TIMER1> = Display::new(board.TIMER1, board.display_pins);
    DISPLAY.init(display);

    let mut mic_in = board.microphone_pins.mic_in;
    let mut mic_run = board.microphone_pins.mic_run;

    let _ = mic_run.set_high();

    let saadc_config: SaadcConfig = SaadcConfig::default();
    let mut saadc: Saadc = Saadc::new(board.ADC, saadc_config);


    let mut window = [0.0f32; FFT_WIDTH];
    // Use a Hann window, which is easy to compute and reasonably good. Source: microbit-spectrum
    for (n, w) in window.iter_mut().enumerate() {
        let s: f32 = (PI * n as f32 / (FFT_WIDTH - 1) as f32).sin();
        *w = s * s;
    }

    loop {
        let samples: [i16; FFT_WIDTH] = read_samples(&mut saadc, &mut mic_in);
        let freqs: [Complex<f32>; FFT_WIDTH / 2] = samples_to_freq(samples, window);
        let dominant_freq: f32 = get_dominant_freq(freqs);
        let noctive: (usize, f32) = get_note(dominant_freq); //note + octive tuple
        let bitmap: [[u8; 5]; 5] = get_image(noctive);
        DISPLAY.with_lock(|display: &mut Display<TIMER1>| {
            display.show(&BitImage::new(&bitmap)); // XXX Not the most memory efficient - Improve
        });
        rprintln!("{}", NOTES[noctive.0]);
        timer.delay(400);
    }
}