use crate::parties::SimulatedParties;
use crate::server::SimulatedServer;
use rand_distr::Pareto;
use std::io::Write;
use termcolor::{Color, ColorChoice, ColorSpec, StandardStream, WriteColor};
mod datasets;
mod oracle;
mod parties;
mod probability;
mod server;
use clap::Clap;
use tracing::event;

use crate::datasets::{CsvDataset, TemporalDataset};
use crate::oracle::Oracle;
use tracing::Level;
use tracing_subscriber;
use tracing_subscriber::FmtSubscriber;

#[derive(Clap)]
#[clap(version = "1.0", author = "Sarah Jamie Lewis <sarah@openprivacy.ca>")]
struct Opts {
    /// the number of parties to simulate
    #[clap(short, long, default_value = "10")]
    num_parties: usize,

    /// dataset
    #[clap(short, long, default_value = "")]
    dataset: String,

    /// samples per round
    #[clap(short, long, default_value = "10")]
    samples_per_round: usize,

    /// minimum false positive rate
    #[clap(long, default_value = "1")]
    min_p: usize,

    /// maximum false positive rate
    #[clap(long, default_value = "8")]
    max_p: usize,

    #[clap(short, long)]
    trace: bool,

    #[clap(short, long, default_value = "0")]
    prob_entangled: f64,

    #[clap(long, default_value = "0")]
    sample: usize,

    #[clap(long, default_value = "0")]
    skip: usize,
}

fn main() {
    let opts: Opts = Opts::parse();
    let level = match opts.trace {
        true => Level::TRACE,
        _ => Level::INFO,
    };

    let mut oracle = Oracle::new();
    let mut max = 0;
    let subscriber = FmtSubscriber::default();
    tracing::subscriber::with_default(subscriber, || {
        let mut rng = rand::thread_rng();
        let mut server = SimulatedServer::new();

        if opts.dataset == "none" {
            let simulated_parties = SimulatedParties::new_simulation(opts.num_parties);
            {
                event!(Level::INFO, "Generating {} Parties and registering them with the server", opts.num_parties);
                simulated_parties.register_with_server(&mut server, &mut rng, opts.min_p, opts.max_p, &mut oracle);
                server.finalize();
            }

            let pareto = Pareto::new(1.0, 1.0).unwrap();

            {
                event!(Level::INFO, "Simulating message sends using {} samples from a pareto distribution...", opts.samples_per_round);
                (0..opts.samples_per_round).for_each(|_i| simulated_parties.sample_traffic(&mut server, &mut rng, pareto, opts.prob_entangled, &mut oracle));
            }
        } else {
            let dataset = CsvDataset::load_dataset(opts.dataset.as_str());
            event!(Level::INFO, "Registering parties from {} which containts {}", opts.dataset, dataset.num_parties());
            dataset.register_with_server(&mut server, &mut rng, opts.min_p, opts.max_p, &mut oracle);
            server.finalize();
            event!(Level::INFO, "Playing back {} events from {}", dataset.num_records(), opts.dataset);
            let mut sample = dataset.num_records();
            if opts.sample != 0 {
                sample = opts.sample;
            }
            dataset.playthough_traffic(&mut server, &mut oracle, opts.sample, opts.skip, opts.prob_entangled, &mut rng);
        }

        {
            let (server_oracle, round_stats, party_stats) = server.statistics();
            let if_uniform = (round_stats.num_messages as f64) / (round_stats.num_registered_parties as f64);
            event!(
                Level::INFO,
                "Round had {}  messages, send to {} parties ({:.2} is uniform distribution)",
                round_stats.num_messages,
                round_stats.num_registered_parties,
                if_uniform
            );
            let mut stdout = StandardStream::stdout(ColorChoice::Always);
            for (index, (party, stats)) in party_stats.iter().enumerate() {
                if stats.trivial_breaks > 0 || (stats.observed_messages > 2 && stats.observed_skew > 2.0) {
                    stdout.set_color(ColorSpec::new().set_fg(Some(Color::Red))).unwrap();
                } else {
                    stdout.set_color(ColorSpec::new().set_fg(Some(Color::Green))).unwrap();
                }
                writeln!(
                    &mut stdout,
                    "({}) Party {} | Ideal: {:>8.2} {:>8.2} | Observed: {:>8.2} ({:>6.2})  | Skew: {:>8.2} ({:>6.2}) | Trivial Attributions this Round: {}",
                    index,
                    party,
                    stats.ideal_rate,
                    stats.ideal_rate * (round_stats.num_messages as f64),
                    stats.observed_messages,
                    100.0 * stats.observed_rate,
                    stats.observed_skew_messages,
                    stats.observed_skew,
                    stats.trivial_breaks
                )
                .unwrap();
            }
            let max = server_oracle.compile_to_dot("server_event.dot", true, false, 1.0);
            println!("Max from server: {}", max);
            oracle.compile_to_dot("actual_events.dot", false, false, 1.0);

            //server_oracle.compile_to_dot("server_event_inverse.dot", true, true,max);
        }
    });
}