#![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::{Const, Matrix, Matrix2, SMatrixView};
use pim::{
    array::{PimMatrixArena, PimRegion, PimStorage},
    // array::PimMatrix,
    // array::{BankArray, ComputeArray},
    kernel::TEST_KERNEL,
    matrix::{F16x1, F16x16},
    state::PimState,
};
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 mut pim_matrix_arena = RefCell::new(PimMatrixArena([[F16x16::default(); 8]; 8]));
    let pim_storage0 = PimStorage {
        arena: &pim_matrix_arena,
        index: 0,
    };
    let pim_storage1 = PimStorage {
        arena: &pim_matrix_arena,
        index: 1,
    };
    let pim_storage2 = PimStorage {
        arena: &pim_matrix_arena,
        index: 2,
    };
    let mut matrix0 = Matrix::from_data(pim_storage0);
    let mut matrix1 = Matrix::from_data(pim_storage1);
    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_lower_triangle(F16x1(f16::ONE), 0);
    // matrix1.fill_upper_triangle(F16x1(f16::from_f32(2.0)), 0);
    writeln!(&mut uart, "{matrix0} * 2\n=").unwrap();

    // let mut compute_array: ComputeArray<3> = ComputeArray([
    //     BankArray([F16x16([f16::from_f32(0.1); 16]); 32]),
    //     BankArray([f16::from_f32(0.2); 512]),
    //     BankArray([f16::from_f32(0.3); 512]),
    // ]);
    // let dummy_array = BankArray::default();

    // writeln!(
    //     &mut uart,
    //     "PIM array is at {:x?}",
    //     core::ptr::addr_of!(compute_array)
    // )
    // .unwrap();

    // writeln!(
    //     &mut uart,
    //     "BankArray0: [{:?}, ...]\nBankArray1: [{:?}, ...]\nBankArray2: [{:?}, ...]",
    //     compute_array.0[0].0[0], compute_array.0[1].0[0], compute_array.0[2].0[0]
    // )
    // .unwrap();

    // writeln!(&mut uart, "MAC: BankArray2 += BankArray0 * BankArray1",).unwrap();

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

    pim_state.set_bank_mode(BankMode::PimAllBank);
    pim_matrix_arena
        .borrow_mut()
        .execute_instruction_read_dual_bank();
    pim_matrix_arena
        .borrow_mut()
        .execute_instruction_read_dual_bank();
    pim_matrix_arena
        .borrow_mut()
        .execute_instruction_write_dual_bank();
    pim_state.set_bank_mode(BankMode::SingleBank);

    pim_matrix_arena.borrow_mut().invalidate();
    // compute_array.invalidate();
    barrier::dsb(barrier::SY);

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

    // writeln!(
    //     &mut uart,
    //     "BankArray2: [{:?}, ...]",
    //     compute_array.0[2].0[0]
    // )
    // .unwrap();

    // writeln!(&mut uart, "ComputeArray:\n{:#?}", compute_array.0[2]).unwrap();

    m5ops::exit();

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

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

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