//! Project: CMtec CMDR joystick 24
//! Date:    2023-08-01
//! Author:  Christoffer Martinsson
//! Email:   cm@cmtec.se
//! License: Please refer to LICENSE in root directory

#![no_std]
#![no_main]

mod button_matrix;
mod layout;
mod status_led;
mod usb_joystick_device;

use button_matrix::ButtonMatrix;
use core::convert::Infallible;
use cortex_m::delay::Delay;
use dyn_smooth::{DynamicSmootherEcoI32, I32_FRAC_BITS};
use embedded_hal::adc::OneShot;
use embedded_hal::digital::v2::*;
use embedded_hal::timer::CountDown;
use fugit::ExtU32;
use libm::powf;
use panic_halt as _;
use rp2040_hal::{
    adc::Adc,
    gpio::{Function, FunctionConfig, PinId, ValidPinMode},
    pio::StateMachineIndex,
};
use status_led::{StatusMode, Ws2812StatusLed};
use usb_device::class_prelude::*;
use usb_device::prelude::*;
use usb_joystick_device::{JoystickConfig, JoystickReport};
use usbd_human_interface_device::prelude::*;
use waveshare_rp2040_zero::entry;
use waveshare_rp2040_zero::{
    hal::{
        clocks::{init_clocks_and_plls, Clock},
        pac,
        pio::PIOExt,
        timer::Timer,
        watchdog::Watchdog,
        Sio,
    },
    Pins, XOSC_CRYSTAL_FREQ,
};

// Public constants
pub const BUTTON_ROWS: usize = 5;
pub const BUTTON_COLS: usize = 5;
pub const NUMBER_OF_BUTTONS: usize = BUTTON_ROWS * BUTTON_COLS;

pub const AXIS_MIN: u16 = 0;
pub const AXIS_MAX: u16 = 4095;
pub const AXIS_CENTER: u16 = AXIS_MAX / 2;

pub const NBR_OF_GIMBAL_AXIS: usize = 4;
pub const GIMBAL_AXIS_LEFT_X: usize = 0;
pub const GIMBAL_AXIS_LEFT_Y: usize = 1;
pub const GIMBAL_AXIS_RIGHT_X: usize = 2;
pub const GIMBAL_AXIS_RIGHT_Y: usize = 3;

// Analog smoothing settings.
pub const BASE_FREQ: i32 = 2 << I32_FRAC_BITS;
pub const SAMPLE_FREQ: i32 = 1000 << I32_FRAC_BITS;
pub const SENSITIVITY: i32 = (0.01 * ((1 << I32_FRAC_BITS) as f32)) as i32;

// Public types
#[derive(Copy, Clone, Default)]
pub struct Button {
    pub pressed: bool,
    pub fn_mode: u8,
}

#[derive(Copy, Clone)]
pub struct GimbalAxis {
    pub value: u16,
    pub idle_value: u16,
    pub max: u16,
    pub min: u16,
    pub center: u16,
    pub fn_mode: u8,
    pub deadzone: (u16, u16, u16),
    pub expo: f32,
}

impl Default for GimbalAxis {
    fn default() -> Self {
        GimbalAxis {
            value: AXIS_CENTER,
            idle_value: AXIS_CENTER,
            max: AXIS_MAX,
            min: AXIS_MIN,
            center: AXIS_CENTER,
            fn_mode: 0,
            deadzone: (50, 50, 50),
            expo: 0.2,
        }
    }
}

#[entry]
fn main() -> ! {
    // Grab our singleton objects
    let mut pac = pac::Peripherals::take().unwrap();

    // Set up the watchdog driver - needed by the clock setup code
    let mut watchdog = Watchdog::new(pac.WATCHDOG);

    // Configure clocks and PLLs
    let clocks = init_clocks_and_plls(
        XOSC_CRYSTAL_FREQ,
        pac.XOSC,
        pac.CLOCKS,
        pac.PLL_SYS,
        pac.PLL_USB,
        &mut pac.RESETS,
        &mut watchdog,
    )
    .ok()
    .unwrap();

    let core = pac::CorePeripherals::take().unwrap();

    // The single-cycle I/O block controls our GPIO pins
    let sio = Sio::new(pac.SIO);

    // Set the pins to their default state
    let pins = Pins::new(
        pac.IO_BANK0,
        pac.PADS_BANK0,
        sio.gpio_bank0,
        &mut pac.RESETS,
    );

    // Enable adc
    let mut adc = Adc::new(pac.ADC, &mut pac.RESETS);

    // Configure ADC input pins
    // Have not figured out hov to store the adc pins in an array yet
    // TODO: Find a way to store adc pins in an array
    let mut adc_pin_left_x = pins.gp29.into_floating_input();
    let mut adc_pin_left_y = pins.gp28.into_floating_input();
    let mut adc_pin_right_x = pins.gp27.into_floating_input();
    let mut adc_pin_right_y = pins.gp26.into_floating_input();

    // Setting up array with pins connected to button rows
    let button_matrix_row_pins: &[&dyn InputPin<Error = Infallible>; BUTTON_ROWS] = &[
        &pins.gp11.into_pull_up_input(),
        &pins.gp13.into_pull_up_input(),
        &pins.gp9.into_pull_up_input(),
        &pins.gp12.into_pull_up_input(),
        &pins.gp10.into_pull_up_input(),
    ];

    // Setting up array with pins connected to button columns
    let button_matrix_col_pins: &mut [&mut dyn OutputPin<Error = Infallible>; BUTTON_COLS] = &mut [
        &mut pins.gp4.into_push_pull_output(),
        &mut pins.gp5.into_push_pull_output(),
        &mut pins.gp6.into_push_pull_output(),
        &mut pins.gp7.into_push_pull_output(),
        &mut pins.gp8.into_push_pull_output(),
    ];

    // Create button matrix object that scans all buttons
    let mut button_matrix: ButtonMatrix<BUTTON_ROWS, BUTTON_COLS, NUMBER_OF_BUTTONS> =
        ButtonMatrix::new(button_matrix_row_pins, button_matrix_col_pins, 5);

    // Initialize button matrix
    button_matrix.init_pins();

    // Configure USB
    let usb_bus = UsbBusAllocator::new(waveshare_rp2040_zero::hal::usb::UsbBus::new(
        pac.USBCTRL_REGS,
        pac.USBCTRL_DPRAM,
        clocks.usb_clock,
        true,
        &mut pac.RESETS,
    ));

    let mut usb_hid_joystick = UsbHidClassBuilder::new()
        .add_device(JoystickConfig::default())
        .build(&usb_bus);

    let mut usb_dev = UsbDeviceBuilder::new(&usb_bus, UsbVidPid(0x1209, 0x0002))
        .manufacturer("CMtec")
        .product("CMDR Joystick")
        .serial_number("0001")
        .build();

    // Create status LED
    let (mut pio, sm0, _, _, _) = pac.PIO0.split(&mut pac.RESETS);
    let mut status_led = Ws2812StatusLed::new(
        pins.neopixel.into_mode(),
        &mut pio,
        sm0,
        clocks.peripheral_clock.freq(),
    );

    // Create timers/delays
    let timer = Timer::new(pac.TIMER, &mut pac.RESETS);
    let mut delay = Delay::new(core.SYST, clocks.system_clock.freq().to_Hz());

    let mut usb_hid_report_count_down = timer.count_down();
    usb_hid_report_count_down.start(10.millis());

    let mut scan_count_down = timer.count_down();
    scan_count_down.start(1.millis());

    let mut status_led_count_down = timer.count_down();
    status_led_count_down.start(250.millis());

    // Create variable to track modes
    let mut mode: u8 = 0;

    // Create joystick button/axis array
    let mut axis: [GimbalAxis; NBR_OF_GIMBAL_AXIS] = [Default::default(); NBR_OF_GIMBAL_AXIS];
    let mut buttons: [Button; NUMBER_OF_BUTTONS] = [Button::default(); NUMBER_OF_BUTTONS];

    // Set up left gimbal Y axis as full range without return to center spring
    axis[GIMBAL_AXIS_LEFT_Y].idle_value = AXIS_MIN;
    axis[GIMBAL_AXIS_LEFT_Y].deadzone = (50, 0, 50);
    axis[GIMBAL_AXIS_LEFT_Y].expo = 0.0;

    // Manual calibation values
    // TODO: add external EEPROM and make calibration routine
    axis[GIMBAL_AXIS_LEFT_X].center = AXIS_CENTER;
    axis[GIMBAL_AXIS_LEFT_X].max = AXIS_MAX - 450;
    axis[GIMBAL_AXIS_LEFT_X].min = AXIS_MIN + 500;
    axis[GIMBAL_AXIS_LEFT_Y].center = AXIS_CENTER + 105;
    axis[GIMBAL_AXIS_LEFT_Y].max = AXIS_MAX - 250;
    axis[GIMBAL_AXIS_LEFT_Y].min = AXIS_MIN + 500;
    axis[GIMBAL_AXIS_RIGHT_X].center = AXIS_CENTER - 230;
    axis[GIMBAL_AXIS_RIGHT_X].max = AXIS_MAX - 700;
    axis[GIMBAL_AXIS_RIGHT_X].min = AXIS_MIN + 350;
    axis[GIMBAL_AXIS_RIGHT_Y].center = AXIS_CENTER - 68;
    axis[GIMBAL_AXIS_RIGHT_Y].max = AXIS_MAX - 700;
    axis[GIMBAL_AXIS_RIGHT_Y].min = AXIS_MIN + 450;

    // Create dynamic smoother array for gimbal axis
    // TODO: Find a way to store dynamic smoother in the axis struct
    let mut smoother: [DynamicSmootherEcoI32; NBR_OF_GIMBAL_AXIS] = [
        DynamicSmootherEcoI32::new(BASE_FREQ, SAMPLE_FREQ, SENSITIVITY),
        DynamicSmootherEcoI32::new(BASE_FREQ, SAMPLE_FREQ, SENSITIVITY),
        DynamicSmootherEcoI32::new(BASE_FREQ, SAMPLE_FREQ, SENSITIVITY),
        DynamicSmootherEcoI32::new(BASE_FREQ, SAMPLE_FREQ, SENSITIVITY),
    ];

    // Scan matrix to get initial state
    for _ in 0..10 {
        button_matrix.scan_matrix(&mut delay);
    }

    // Fallback way to enter bootloader
    if button_matrix.buttons_pressed()[0] {
        status_led.update(StatusMode::Bootloader);
        let gpio_activity_pin_mask: u32 = 0;
        let disable_interface_mask: u32 = 0;
        rp2040_hal::rom_data::reset_to_usb_boot(gpio_activity_pin_mask, disable_interface_mask);
    }

    loop {
        if status_led_count_down.wait().is_ok() {
            update_status_led(&mut status_led, &mode);
        }

        if usb_hid_report_count_down.wait().is_ok() {
            let pressed_keys = button_matrix.buttons_pressed();

            mode = get_mode(pressed_keys);

            for (index, key) in pressed_keys.iter().enumerate() {
                buttons[index].pressed = *key;
            }

            if button_matrix.buttons_pressed()[0]
                && button_matrix.buttons_pressed()[1]
                && button_matrix.buttons_pressed()[5]
                && button_matrix.buttons_pressed()[6]
                && button_matrix.buttons_pressed()[8]
                && button_matrix.buttons_pressed()[9]
            {
                status_led.update(StatusMode::Bootloader);
                let gpio_activity_pin_mask: u32 = 0;
                let disable_interface_mask: u32 = 0;
                rp2040_hal::rom_data::reset_to_usb_boot(
                    gpio_activity_pin_mask,
                    disable_interface_mask,
                );
            }

            match usb_hid_joystick.device().write_report(&get_joystick_report(
                &mut buttons,
                &mut axis,
                &mode,
            )) {
                Err(UsbHidError::WouldBlock) => {}
                Ok(_) => {}
                Err(e) => {
                    status_led.update(StatusMode::Error);
                    core::panic!("Failed to write joystick report: {:?}", e)
                }
            };
        }

        if scan_count_down.wait().is_ok() {
            button_matrix.scan_matrix(&mut delay);

            // Have not figured out hov to store the adc pins in an array yet
            // so we have to read them one by one
            // TODO: Find a way to store adc pins in an array
            smoother[GIMBAL_AXIS_LEFT_X].tick(adc.read(&mut adc_pin_left_x).unwrap());
            smoother[GIMBAL_AXIS_LEFT_Y].tick(adc.read(&mut adc_pin_left_y).unwrap());
            smoother[GIMBAL_AXIS_RIGHT_X].tick(adc.read(&mut adc_pin_right_x).unwrap());
            smoother[GIMBAL_AXIS_RIGHT_Y].tick(adc.read(&mut adc_pin_right_y).unwrap());

            for (index, item) in axis.iter_mut().enumerate() {
                item.value = calculate_axis_value(
                    smoother[index].value() as u16,
                    item.min,
                    item.max,
                    item.center,
                    item.deadzone,
                    item.expo,
                );
            }
        }

        if usb_dev.poll(&mut [&mut usb_hid_joystick]) {}
    }
}

/// Update status LED colour based on function layer and capslock
///
/// Normal = green (NORMAL)
/// Left Alt mode = blue (GUI LOCK)
/// Error = steady red (ERROR)
///
/// # Arguments
/// * `status_led` - Reference to status LED
/// * `caps_lock_active` - Is capslock active
fn update_status_led<P, SM, I>(status_led: &mut Ws2812StatusLed<P, SM, I>, fn_mode: &u8)
where
    P: PIOExt + FunctionConfig,
    I: PinId,
    Function<P>: ValidPinMode<I>,
    SM: StateMachineIndex,
{
    if *fn_mode & 0x10 == 0x10 {
        status_led.update(StatusMode::Activity);
    } else {
        status_led.update(StatusMode::Normal);
    }
}

/// Get current Fn mode (0, 1, 2 or 3 and alt l/r mode)
/// layout::MAP contains the button types
///
/// # Arguments
///
/// * `pressed_keys` - Array of pressed keys
fn get_mode(pressed_keys: [bool; NUMBER_OF_BUTTONS]) -> u8 {
    // Check how many Fn keys are pressed
    let mut mode: u8 = 0;
    let mut fn_l_active: bool = false;
    let mut fn_r_active: bool = false;
    let mut alt_l_active: bool = false;
    let mut alt_r_active: bool = false;

    for (index, key) in pressed_keys.iter().enumerate() {
        if *key && layout::MAP[0][index] == layout::ButtonType::FnL {
            fn_l_active = true;
        }
        if *key && layout::MAP[0][index] == layout::ButtonType::FnR {
            fn_r_active = true;
        }
        if *key && layout::MAP[0][index] == layout::ButtonType::ModeL {
            alt_l_active = true;
        }
        if *key && layout::MAP[0][index] == layout::ButtonType::ModeR {
            alt_r_active = true;
        }
    }

    if fn_l_active && fn_r_active {
        mode = 3;
    } else if fn_r_active {
        mode = 2;
    } else if fn_l_active {
        mode = 1;
    }

    // Set bit 4 and 5 if alt l/r is active
    if alt_l_active {
        mode |= 0x10;
    }
    if alt_r_active {
        mode |= 0x20;
    }

    mode
}

/// Generate keyboard report based on pressed keys and Fn mode (0, 1 or 2)
/// layout::MAP contains the keycodes for each key in each Fn mode
///
/// # Arguments
///
/// * `matrix_keys` - Array of pressed keys
/// * `axis` - Array of joystick axis values
/// * `fn_mode` - Fn mode (0, 1, 2 or 3)
/// * `alt_l_mode` - Is left alt mode active
/// * `alt_r_mode` - Is right alt mode active
fn get_joystick_report(
    matrix_keys: &mut [Button; NUMBER_OF_BUTTONS],
    axis: &mut [GimbalAxis; 4],
    mode: &u8,
) -> JoystickReport {
    let mut x: u16 = axis[GIMBAL_AXIS_RIGHT_X].value;
    let mut y: u16 = axis[GIMBAL_AXIS_RIGHT_Y].value;
    let z: u16 = axis[GIMBAL_AXIS_LEFT_X].value;
    let mut rx: u16 = AXIS_CENTER;
    let mut ry: u16 = AXIS_CENTER;
    let mut rz: u16 = axis[GIMBAL_AXIS_LEFT_Y].value;

    // Update Fn mode for all axis that are in idle position
    // This is to avoid the Fn mode switching when moving the gimbal
    for item in axis.iter_mut() {
        if item.value == item.idle_value {
            item.fn_mode = mode & 0x0F;
        }
    }

    // Left Alt mode active (bit 4)
    // Full range of left gimbal gives half range of joystick axis (center to max)
    // Left Fn mode = reversed range (center to min)
    if mode & 0x10 == 0x10
        && (axis[GIMBAL_AXIS_LEFT_Y].fn_mode == 0 || axis[GIMBAL_AXIS_LEFT_Y].fn_mode == 2)
    {
        rz = remap(
            axis[GIMBAL_AXIS_LEFT_Y].value,
            AXIS_MIN,
            AXIS_MAX,
            AXIS_CENTER,
            AXIS_MAX,
        );
    } else if mode & 0x10 == 0x10
        && (axis[GIMBAL_AXIS_LEFT_Y].fn_mode == 1 || axis[GIMBAL_AXIS_LEFT_Y].fn_mode == 3)
    {
        rz = AXIS_MAX
            - remap(
                axis[GIMBAL_AXIS_LEFT_Y].value,
                AXIS_MIN,
                AXIS_MAX,
                AXIS_CENTER,
                AXIS_MAX,
            );
    }

    // Right Alt mode active (bit 5)
    // Right gimbal control third joystick axis when right Fn mode is active
    if mode & 0x20 == 0x20
        && (axis[GIMBAL_AXIS_RIGHT_X].fn_mode == 2 || axis[GIMBAL_AXIS_RIGHT_X].fn_mode == 3)
    {
        x = AXIS_CENTER;
        rx = axis[GIMBAL_AXIS_RIGHT_X].value;
    }
    if mode & 0x20 == 0x20
        && (axis[GIMBAL_AXIS_RIGHT_Y].fn_mode == 2 || axis[GIMBAL_AXIS_RIGHT_Y].fn_mode == 3)
    {
        y = AXIS_CENTER;
        ry = axis[GIMBAL_AXIS_RIGHT_Y].value;
    }

    // Set fn mode for all keys taht are in idle position
    // This is to avoid the Fn mode switching when using a button
    for key in matrix_keys.iter_mut() {
        if !key.pressed {
            key.fn_mode = mode & 0x0F;
        }
    }

    // Generate array for all four hat switches with following structure:
    // * bit 1: Up
    // * bit 2: Right
    // * bit 3: Down
    // * bit 4: Left
    // * bit 5: Button
    // * value 0 = not pressed
    // * value 1 = pressed
    let mut hats: [u8; 4] = [0; 4];
    for (index, key) in matrix_keys.iter_mut().enumerate() {
        if key.pressed
            && layout::MAP[key.fn_mode as usize][index] as usize
                >= layout::ButtonType::Hat1U as usize
            && layout::MAP[key.fn_mode as usize][index] as usize
                <= layout::ButtonType::Hat4B as usize
        {
            hats[(layout::MAP[key.fn_mode as usize][index] as usize
                - layout::ButtonType::Hat1U as usize)
                / 5] |= 1
                << ((layout::MAP[key.fn_mode as usize][index] as usize
                    - layout::ButtonType::Hat1U as usize)
                    - (5 * ((layout::MAP[key.fn_mode as usize][index] as usize
                        - layout::ButtonType::Hat1U as usize)
                        / 5)));
        }
    }

    // Convert hat switch data to HID code
    let (hat1, hat_button1) = format_hat_value(hats[0]);
    let (hat2, hat_button2) = format_hat_value(hats[1]);
    let (hat3, hat_button3) = format_hat_value(hats[2]);
    let (hat4, hat_button4) = format_hat_value(hats[3]);

    // Update button state for joystick button 21-24 according to hat button 1-4
    let mut buttons: u32 = (hat_button1 as u32) << 20
        | ((hat_button2 as u32) << 21)
        | ((hat_button3 as u32) << 22)
        | ((hat_button4 as u32) << 23);

    // Update button state for joystick button 1-20
    for (index, key) in matrix_keys.iter_mut().enumerate() {
        if key.pressed
            && layout::MAP[key.fn_mode as usize][index] as usize >= layout::ButtonType::B1 as usize
            && layout::MAP[key.fn_mode as usize][index] as usize <= layout::ButtonType::B20 as usize
        {
            buttons |= 1 << layout::MAP[key.fn_mode as usize][index] as usize;
        }
    }

    JoystickReport {
        x,
        y,
        z,
        rx,
        ry,
        rz,
        hat1,
        hat2,
        hat3,
        hat4,
        buttons,
    }
}

/// Format hat value from 5 switches to USB HID coded value and button state
///
/// # Arguments
/// * `input` - Hat value coded as
///    bit 1-4: direction (U R D L)
///    bit 5: button state
///    0 = not pressed
///    1 = pressed
fn format_hat_value(input: u8) -> (u8, u8) {
    const HAT_CENTER: u8 = 0xf;
    const HAT_UP: u8 = 0;
    const HAT_UP_RIGHT: u8 = 1;
    const HAT_RIGHT: u8 = 2;
    const HAT_DOWN_RIGHT: u8 = 3;
    const HAT_DOWN: u8 = 4;
    const HAT_DOWN_LEFT: u8 = 5;
    const HAT_LEFT: u8 = 6;
    const HAT_UP_LEFT: u8 = 7;

    let direction: u8 = match input & 0x0F {
        1 => HAT_UP,
        2 => HAT_RIGHT,
        3 => HAT_UP_RIGHT,
        4 => HAT_DOWN,
        6 => HAT_DOWN_RIGHT,
        8 => HAT_LEFT,
        12 => HAT_DOWN_LEFT,
        9 => HAT_UP_LEFT,
        _ => HAT_CENTER,
    };

    // Alpine hat switch button filter
    let mut button_state: u8 = 0;
    if input & 0x10 == 0x10 && direction == HAT_CENTER {
        button_state = 1;
    }

    (direction, button_state)
}

/// Calculate value for joystick axis
///
/// # Arguments
/// * `value` - Value to calibrate
/// * `min` - Lower bound of the value's current range
/// * `max` - Upper bound of the value's current range
/// * `center` - Center of the value's current range
/// * `deadzone` - Deadzone of the value's current range (min, center, max)
/// * `expo` - Exponential curve factor
fn calculate_axis_value(
    value: u16,
    min: u16,
    max: u16,
    center: u16,
    deadzone: (u16, u16, u16),
    expo: f32,
) -> u16 {
    let mut calibrated_value = AXIS_CENTER;

    if value > (center + deadzone.1) {
        calibrated_value = remap(
            value,
            center + deadzone.1,
            max - deadzone.2,
            AXIS_CENTER,
            AXIS_MAX,
        );
    } else if value < (center - deadzone.1) {
        calibrated_value = remap(
            value,
            min + deadzone.0,
            center - deadzone.1,
            AXIS_MIN,
            AXIS_CENTER,
        );
    }

    if expo != 0.0 {
        let joystick_x_float = calibrated_value as f32 / AXIS_MAX as f32;
        // Calculate expo using 9th order polynomial function with 0.5 as center point
        let joystick_x_exp: f32 = expo * (0.5 + 256.0 * powf(joystick_x_float - 0.5, 9.0))
            + (1.0 - expo) * joystick_x_float;

        calibrated_value = constrain(
            (joystick_x_exp * AXIS_MAX as f32) as u16,
            AXIS_MIN,
            AXIS_MAX,
        );
    }

    calibrated_value
}

/// Remapping values from one range to another
///
/// # Arguments
/// * `value` - Value to remap
/// * `in_min` - Lower bound of the value's current range
/// * `in_max` - Upper bound of the value's current range
/// * `out_min` - Lower bound of the value's target range
/// * `out_max` - Upper bound of the value's target range
fn remap(value: u16, in_min: u16, in_max: u16, out_min: u16, out_max: u16) -> u16 {
    constrain(
        (value as i64 - in_min as i64) * (out_max as i64 - out_min as i64)
            / (in_max as i64 - in_min as i64)
            + out_min as i64,
        out_min as i64,
        out_max as i64,
    ) as u16
}

/// Constrain a value to a given range
///
/// # Arguments
/// * `value` - Value to constrain
/// * `out_min` - Lower bound of the value's target range
/// * `out_max` - Upper bound of the value's target range
fn constrain<T: PartialOrd>(value: T, out_min: T, out_max: T) -> T {
    if value < out_min {
        out_min
    } else if value > out_max {
        out_max
    } else {
        value
    }
}