use std::any::Any;
use std::sync::Arc;
use std::fs::{write, create_dir_all};
use rand::rngs::ThreadRng;
use nalgebra::{DVector, dvector, dmatrix};

use modppl::prelude::*;

mod pointed_model;
use pointed_model::types_2d::Bounds;
use pointed_model::{PointedModel, DriftProposal};

mod dyngenfns;
use dyngenfns::{pointed_2d_model, pointed_2d_drift_proposal};
use dyngenfns::{hierarchical_model, read_coeffs,
    hierarchical_drift_proposal,
    add_or_remove_param_proposal
};


#[test]
fn test_metropolis_hastings_handcoded() -> std::io::Result<()> {
    create_dir_all("../data")?;

    const NUM_ITERS: u32 = 25000;

    let model = PointedModel { obs_cov: dmatrix![1., -3./5.; -3./5., 2.] };
    let proposal = DriftProposal { drift_cov: dmatrix![0.25, 0.; 0., 0.25] };
    let bounds = Bounds { xmin: -5., xmax: 5., ymin: -5., ymax: 5. };
    let obs = dvector![0., 0.];

    let constraints = (None, Some(obs));

    let (mut trace, _) = model.generate(bounds, constraints);
    for iter in 0..NUM_ITERS {
        dbg!(iter);
        let (new_trace, accepted) = modppl::mh(&model, trace, &proposal, ());
        dbg!(accepted);
        trace = new_trace;
        let data = trace.data.0.clone().unwrap();
        let json = format!("[{},{}]", data[0], data[1]);
        write(format!("../data/mh_trace_{}.json", iter), json)?;
    }
    
    Ok(())
}


#[test]
pub fn test_metropolis_hastings_dyngenfn() -> std::io::Result<()> {
    create_dir_all("../data")?;

    const NUM_ITERS: u32 = 25000;

    let bounds = Bounds { xmin: -5., xmax: 5., ymin: -5., ymax: 5. };
    let obs = dvector![0., 0.];

    let mut observations = Trie::new();
    observations.observe("obs", Arc::new(obs) as Arc<dyn Any + Send + Sync>);

    let mut trace = pointed_2d_model.generate((bounds, dmatrix![1., -3./5.; -3./5., 2.]), observations).0;
    for iter in 0..NUM_ITERS {
        dbg!(iter);
        let (new_trace, accepted) = mh(&pointed_2d_model, trace, &pointed_2d_drift_proposal, dmatrix![0.25, 0.; 0., 0.25]);
        dbg!(accepted);
        trace = new_trace;
        let data = trace.data.read::<DVector<f64>>("latent");
        let json = format!("[{},{}]", data[0], data[1]);
        write(format!("../data/mh_trace_{}.json", iter), json)?;
    }

    Ok(())
}

#[test]
pub fn test_metropolis_hastings_hierarchical() -> std::io::Result<()> {
    create_dir_all("../data")?;

    let mut rng = ThreadRng::default();

    let xs = vec![-5.,-4.,-3.,-2.,-1.,0.,1.,2.,3.,4.,5.];

    let mut observations = DynTrie::new();
    let (a, b, c) = (0.3, 0.4, 0.5);
    let ys = xs.iter().map(|x|
        a + b*x + c*x*x + normal.random(&mut rng, (0., 0.1))
    ).collect::<Vec<f64>>();
    write("../data/hierarchical_data.json", format!("[{:?}, {:?}]", xs, ys))?;
    ys.into_iter().enumerate().for_each(|(i, y)| { observations.observe(&format!("(y, {})", i), Arc::new(y)); });

    let mut trace = hierarchical_model.generate(xs, observations).0;
    let mut all_coeffs = vec![];
    for _ in 0..100 {
        let (new_trace, _) = mh(&hierarchical_model, trace, &add_or_remove_param_proposal, ());
        trace = new_trace;
        all_coeffs.push(read_coeffs(&trace));
        for _ in 0..3 {
            let (new_trace, _) = mh(&hierarchical_model, trace, &hierarchical_drift_proposal, 0.1);
            trace = new_trace;
            all_coeffs.push(read_coeffs(&trace));
        }
        for _ in 0..10 {
            let (new_trace, _) = mh(&hierarchical_model, trace, &hierarchical_drift_proposal, 0.01);
            trace = new_trace;
            all_coeffs.push(read_coeffs(&trace));
        }
        write("../data/hierarchical_model.json", format!("{:?}", all_coeffs))?;
    }
    Ok(())
}