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

use aarch64_cpu::asm::barrier;
use core::{
    cell::RefCell,
    fmt::Write,
    panic::PanicInfo,
    sync::atomic::{compiler_fence, Ordering},
};
use half::f16;
use nalgebra::Matrix;
use pim::{
    array::{PimMatrixArena, PimRegion, PimStorage},
    kernel::TEST_KERNEL,
    state::PimState,
    vector::{F16x1, F16x16},
};
use pim_isa::BankMode;
use uart::Uart0;

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

#[no_mangle]
pub extern "C" fn entry() -> ! {
    let mut uart = Uart0;
    let mut pim_state = PimState::new(&TEST_KERNEL);

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

    let mut matrix0 = Matrix::from_data(pim_storage0);
    matrix0.fill_column(0, F16x1(f16::ZERO));
    matrix0.fill_column(1, F16x1(f16::ONE));
    matrix0.fill_column(2, F16x1(f16::PI));
    matrix0.fill_column(3, F16x1(f16::E));
    matrix0.fill_column(4, F16x1(f16::EPSILON));
    matrix0.fill_column(5, F16x1(f16::SQRT_2));
    matrix0.fill_column(6, F16x1(f16::LN_2));
    matrix0.fill_column(7, F16x1(f16::LN_10));

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

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

    // Invalidate and flush array just in case
    pim_matrix_arena0.borrow_mut().invalidate_flush();
    pim_matrix_arena1.borrow_mut().invalidate_flush();
    barrier::dsb(barrier::SY);

    pim_state.set_bank_mode(BankMode::PimAllBank);
    pim_matrix_arena0
        .borrow()
        .execute_instruction_read_dual_bank();
    pim_matrix_arena1
        .borrow()
        .execute_instruction_read_dual_bank();
    pim_matrix_arena0
        .borrow_mut()
        .execute_instruction_write_dual_bank();
    pim_state.set_bank_mode(BankMode::SingleBank);

    pim_matrix_arena0.borrow_mut().invalidate();
    barrier::dsb(barrier::SY);

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

    m5ops::exit();

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

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

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