#![feature(generic_const_exprs)]
#![no_std]
#![no_main]

extern crate alloc;

use aarch64_cpu::asm::barrier;
use alloc::{boxed::Box, rc::Rc};
use core::{
    cell::RefCell,
    fmt::Write,
    mem::MaybeUninit,
    panic::PanicInfo,
    sync::atomic::{compiler_fence, Ordering},
};
use embedded_alloc::Heap;
use half::f16;
use nalgebra::Matrix;
use pim::{
    array::{DummyArray, PimMatrixArena, PimStorage, NUMBER_OF_BANKS},
    kernel::{execute_matrix_multiply_rowwise, MATRIX_MUL},
    state::PimState,
    vector::{F16x1, F16x16},
};
use uart::Uart0;

mod boot;
mod critical_section;
mod m5ops;
mod pim;
mod uart;

#[global_allocator]
static PIM_ALLOC: Heap = Heap::empty();

const PIM_ARENA_SIZE: usize = 0x2000000;

#[link_section = ".pim_data"]
static mut PIM_ARENA: [MaybeUninit<u8>; PIM_ARENA_SIZE] = [MaybeUninit::uninit(); PIM_ARENA_SIZE];

#[no_mangle]
pub extern "C" fn entry() -> ! {
    unsafe {
        PIM_ALLOC.init(PIM_ARENA.as_ptr() as usize, PIM_ARENA_SIZE);
    }

    let mut uart = Uart0;
    let mut pim_state = PimState::new(&MATRIX_MUL);
    pim_state.set_kernel();

    let pim_matrix_arena0 = Rc::new(RefCell::new(PimMatrixArena(
        [[[F16x16::default(); NUMBER_OF_BANKS]; 8]; 8],
    )));
    let pim_matrix_arena1 = Rc::new(RefCell::new(PimMatrixArena(
        [[[F16x16::default(); NUMBER_OF_BANKS]; 8]; 8],
    )));
    let pim_matrix_arena2 = Rc::new(RefCell::new(PimMatrixArena(
        [[[F16x16::default(); NUMBER_OF_BANKS]; 8]; 8],
    )));
    let pim_storage0 = PimStorage {
        arena: &pim_matrix_arena0,
        index: 0,
    };
    let pim_storage1 = PimStorage {
        arena: &pim_matrix_arena1,
        index: 0,
    };
    let pim_storage2 = PimStorage {
        arena: &pim_matrix_arena2,
        index: 0,
    };

    writeln!(
        &mut uart,
        "arena0: {:?}\narena1: {:?}\narena2: {:?}",
        core::ptr::addr_of!(*pim_matrix_arena0.borrow()),
        core::ptr::addr_of!(*pim_matrix_arena1.borrow()),
        core::ptr::addr_of!(*pim_matrix_arena2.borrow()),
    )
    .unwrap();

    let mut matrix0 = Matrix::from_data(pim_storage0);
    matrix0.fill_lower_triangle(F16x1(f16::ONE), 0);

    let mut matrix1 = Matrix::from_data(pim_storage1);
    matrix1.fill_lower_triangle(F16x1(f16::ONE), 0);

    let matrix2 = Matrix::from_data(pim_storage2);

    writeln!(&mut uart, "{matrix0} * {matrix1}\n=").unwrap();

    let mut dummy_array = Box::new(DummyArray([F16x16::default(); NUMBER_OF_BANKS]));
    barrier::dsb(barrier::SY);

    execute_matrix_multiply_rowwise(
        &mut pim_state,
        &mut pim_matrix_arena0.borrow_mut(),
        &mut pim_matrix_arena1.borrow_mut(),
        &mut pim_matrix_arena2.borrow_mut(),
        dummy_array.as_mut(),
    );

    writeln!(&mut uart, "{matrix2}").unwrap();

    m5ops::exit();

    loop {
        compiler_fence(Ordering::SeqCst);
    }
}

#[panic_handler]
fn panic(info: &PanicInfo) -> ! {
    writeln!(Uart0, "{info}").unwrap();

    loop {
        compiler_fence(Ordering::SeqCst);
    }
}