use std::collections::HashMap;

use burn::tensor::Tensor;
use perf_event::{
    events::{Event, Hardware},
    Builder, Counter,
};

use crate::freq::FrequencyKHZ;
use std::sync::{Arc, RwLock};

use crate::energy::Estimator;
use crate::model::BurnBackend;

pub struct PerfEstimator {
    registry: HashMap<u64, Counters>,
    model: crate::model::Net<BurnBackend>,
    device: <BurnBackend as burn::prelude::Backend>::Device,
    shared_cpu_current_frequencies: Arc<RwLock<Vec<FrequencyKHZ>>>,
}

impl PerfEstimator {
    pub fn new(shared_cpu_current_frequencies: Arc<RwLock<Vec<FrequencyKHZ>>>) -> Self {
        let model = crate::model::load_model("perf.pt");
        Self {
            registry: Default::default(),
            model,
            device: Default::default(),
            shared_cpu_current_frequencies,
        }
    }
}

struct Counters {
    counters: Vec<Counter>,
    old_time: u64,
    old_total_energy: f64,
    cpu: i32,
}

static EVENT_TYPES_P: &[Event] = &[
    Event::Hardware(Hardware::BRANCH_MISSES),
    Event::Hardware(Hardware::CACHE_MISSES),
    Event::Hardware(Hardware::CACHE_REFERENCES),
    Event::Hardware(Hardware::CPU_CYCLES),
    Event::Hardware(Hardware::INSTRUCTIONS),
    #[cfg(not(feature = "amd"))]
    Event::Hardware(Hardware::REF_CPU_CYCLES),
];

impl Estimator for PerfEstimator {
    fn start_trace(&mut self, pid: u64, cpu: i32) -> Result<(), ()> {
        let counters: Result<Vec<_>, _> = EVENT_TYPES_P
            .iter()
            .map(|kind| {
                Builder::new()
                    .kind(kind.clone())
                    .observe_pid(pid as i32)
                    .inherit_thread(true)
                    .build()
            })
            .collect();

        let mut counters = match counters {
            Ok(counters) => counters,
            Err(e) => {
                eprintln!(
                    "Failed to create performance counter for PID {}: {}",
                    pid, e
                );
                return Err(());
            }
        };

        for counter in counters.iter_mut() {
            if let Err(e) = counter.enable() {
                eprintln!("Failed to enable performance counters: {}", e);
                return Err(());
            }
            if let Err(e) = counter.reset() {
                eprintln!("Failed to reset performance counters: {}", e);
                return Err(());
            }
        }

        let old_time = counters[0].read_count_and_time().unwrap().time_running;
        let counters = Counters {
            counters,
            old_time,
            old_total_energy: 0.,
            cpu,
        };
        self.registry.insert(pid, counters);
        Ok(())
    }

    fn stop_trace(&mut self, pid: u64) {
        self.registry.remove(&pid);
    }

    fn update_information(&mut self, pid: u64, cpu: i32) {
        if let Some(info) = self.registry.get_mut(&pid) {
            info.cpu = cpu;
        } else {
            eprintln!("Tried to update an unknown task")
        }
    }

    fn read_consumption(&mut self, pid: u64) -> Option<f64> {
        let Some(counters) = self.registry.get_mut(&pid) else {
            println!("did not find counters for {pid}");
            return None;
        };

        let time_running_ns = counters.counters[0]
            .read_count_and_time()
            .unwrap()
            .time_running;
        if time_running_ns - counters.old_time == 0 {
            return None;
        }
        let correction_factor = 10_000_000. / (time_running_ns - counters.old_time) as f64;
        counters.old_time = time_running_ns;

        let mut values = vec![
            (self.shared_cpu_current_frequencies.read().unwrap()[counters.cpu as usize] / 1000)
                as f64,
        ];
        for ty in counters.counters.iter_mut() {
            let count: u64 = ty.read().unwrap();
            values.push((count as f64) * correction_factor);
        }
        #[cfg(feature = "amd")]
        values.push(values[4]);

        let result = self
            .model
            .forward(Tensor::from_floats(&values.as_slice()[0..], &self.device));

        let energy = result.into_scalar() as f64;
        counters.old_total_energy += energy / correction_factor;
        for counter in counters.counters.iter_mut() {
            counter.reset().unwrap();
        }
        Some(energy / correction_factor)
    }
}