// Fit to a ring and plot the results
//
// To run:
//   cargo run --release --example -- <OPTIONS>
//
// Example:
//   cargo run --release --example -- -n 1000 --nstates 1
//

#[cfg(feature = "formats")]
mod requires_formats {
    use clap::Parser;
    use rand::{Rng, SeedableRng};
    use rand_distr::Uniform;
    use std::io;
    use tempfile::NamedTempFile;

    use lace::prelude::*;

    #[derive(Debug, Parser)]
    struct Opt {
        /// The number of data to generate
        #[clap(short, default_value = "1000")]
        n: usize,
        /// The number of states in the Engine
        #[clap(long, default_value = "8")]
        nstates: usize,
        /// Scales the ring up or down from its default range, which is (-1, 1)
        #[clap(long, default_value = "1.0")]
        scale: f64,
        /// Width of the ring
        #[clap(long, default_value = "0.2")]
        width: f64,
    }

    fn gen_ring<R: Rng, W: io::Write>(
        n: usize,
        scale: f64,
        width: f64,
        fileout: &mut W,
        rng: &mut R,
    ) -> io::Result<(Vec<f64>, Vec<f64>)> {
        let unif = Uniform::new(-1.0, 1.0);

        let mut n_collected: usize = 0;
        let mut xs = Vec::with_capacity(n);
        let mut ys = Vec::with_capacity(n);

        writeln!(fileout, "ID,x,y")?;
        while n_collected < n {
            let x: f64 = rng.sample(unif) * scale;
            let y: f64 = rng.sample(unif) * scale;
            let r = (x * x + y * y).sqrt();
            if (1.0 - width) * scale <= r && r <= 1.0 * scale {
                writeln!(fileout, "{},{},{}", n_collected, x, y)?;
                xs.push(x);
                ys.push(y);
                n_collected += 1;
            }
        }
        Ok((xs, ys))
    }

    fn plot(
        xs_in: Vec<f64>,
        ys_in: Vec<f64>,
        xs_sim: Vec<f64>,
        ys_sim: Vec<f64>,
    ) {
        use plotly::common::Mode;
        use plotly::layout::{GridPattern, Layout, LayoutGrid};
        use plotly::{Plot, Scatter};

        let trace1 =
            Scatter::new(xs_in, ys_in).name("Input").mode(Mode::Markers);

        let trace2 = Scatter::new(xs_sim, ys_sim)
            .mode(Mode::Markers)
            .name("Simulated")
            .x_axis("x2")
            .y_axis("y2");

        let mut plot = Plot::new();
        plot.add_trace(trace1);
        plot.add_trace(trace2);

        let layout = Layout::new().grid(
            LayoutGrid::new()
                .rows(1)
                .columns(2)
                .pattern(GridPattern::Independent),
        );

        plot.set_layout(layout);
        plot.show();
    }

    pub fn main() {
        let opt = Opt::parse();

        // generate csv data
        println!("Generating data");
        let mut rng = rand::thread_rng();
        let mut f = NamedTempFile::new().unwrap();
        let (xs_in, ys_in) =
            gen_ring(opt.n, opt.scale, opt.width, &mut f, &mut rng).unwrap();
        println!("Data written to {:?}", f.path());

        // generate codebook
        println!("Generating codebook");
        let codebook = lace_codebook::data::codebook_from_csv(
            f.path(),
            None,
            None,
            None,
            false,
        )
        .unwrap();

        // generate engine
        println!("Constructing Engine");
        let mut engine = Engine::new(
            opt.nstates,
            codebook,
            DataSource::Csv(f.path().into()),
            0,
            rand_xoshiro::Xoshiro256Plus::from_entropy(),
        )
        .unwrap();

        // run engine with default config. At the time of April 20 2021, the default
        // reassignment kernels for both rows and column is Gibbs.
        println!("Running Engine");
        engine.run(1_000).unwrap();

        println!("Simulating data");
        let mut xs_sim = Vec::with_capacity(opt.n);
        let mut ys_sim = Vec::with_capacity(opt.n);
        engine
            .simulate(&[0, 1], &Given::<usize>::Nothing, opt.n, None, &mut rng)
            .unwrap()
            .drain(..)
            .for_each(|xy| {
                xs_sim.push(xy[0].to_f64_opt().unwrap());
                ys_sim.push(xy[1].to_f64_opt().unwrap());
            });

        engine
            .states
            .iter()
            .for_each(|state| print!("{} ", state.views[0].asgn().n_cats));

        println!("\nPlotting");
        plot(xs_in, ys_in, xs_sim, ys_sim);
        println!("Done");
    }
}

#[cfg(feature = "formats")]
fn main() {
    requires_formats::main()
}

#[cfg(not(feature = "formats"))]
fn main() {
    eprintln!("Please enable the 'formats' feature to use this example.")
}