use burn::prelude::*;
use burn::module::{AutodiffModule, ModuleVisitor, Param};
use burn::tensor::backend::AutodiffBackend;

struct SanityChecker {
    max_param: f32,
    param_count: usize,
}

impl SanityChecker {
    pub fn new() -> Self {
        Self {
            max_param: 0.0,
            param_count: 0,
        }
    }
}

impl<B: Backend> ModuleVisitor<B> for SanityChecker {
    fn visit_int<const D: usize>(&mut self, _param: &Param<Tensor<B, D, Int>>) {}
    fn visit_bool<const D: usize>(&mut self, _param: &Param<Tensor<B, D, Bool>>) {}
    fn visit_float<const D: usize>(&mut self, param: &Param<Tensor<B, D, Float>>) {
        let tensor = param.val();
        let max = tensor.clone().abs().max().into_scalar().elem::<f32>();
        if max > self.max_param {
            self.max_param = max;
        }
        let mut prod = 1;
        let dims = tensor.dims();
        for i in 0..D {
            prod *= dims[i];
        }
        self.param_count += prod;
    }
}

pub fn foo<B: AutodiffBackend>(model: impl AutodiffModule<B>) {
    loop {
        let mut sanity_checker = SanityChecker::new();
        model.visit(&mut sanity_checker);
        
        println!("max_param: {}, param_count: {}",
            sanity_checker.max_param,
            sanity_checker.param_count,
        );
    }
}