use crate::types::*;
// use rand_xoshiro::{SplitMix64, rand_core::RngCore};
use crate::stats::{CumulativeStats, CurrentStats};
use reservoir_sampling::unweighted::core::l;
use rand_pcg::Pcg64;
use rand::RngCore;
use rand::seq::SliceRandom;

pub struct Simulation {
    pub done_init: bool,
    pub now: TimeCount,
    pub quorums: Vec<Quorum>,
    pub regions: Vec<Region>,
    pub rand: Pcg64,
    pub lg_regions_per_quorum: RegionCount,
    pub num_region_mask: RegionCount,
    pub cur_stats: CurrentStats,
    pub cum_stats: CumulativeStats,
    pub k: NodeCount,
    pub g: NodeCount,
}

impl Simulation {
    fn rand_region(&mut self) -> RegionCount {
        (self.rand.next_u64() as RegionCount) & self.num_region_mask
    }

    pub fn random_shuffle_regions(&mut self, n:usize) -> Vec<u32> {
        let mut vec: Vec<u32> = (0..n as u32).collect();
        vec.shuffle(&mut self.rand);
        vec
    }

    // Reservoir-sampling to pick which nodes to kick out for CCR
    fn pick_honest_malicious_nodes_to_kick_out(&mut self, m: NodeCount, n: NodeCount, n_bad: NodeCount) -> NodeCount {
        let mut selected_indices: Vec<usize> = vec![0usize; m as usize];
        // Picks m indices out of 0..n using the optimized L reservoir sampling algo.
        l(0usize..n, selected_indices.as_mut_slice(), &mut self.rand);
        // First 0..n_bad - 1 indices out of total are treated as malicious
        let num_bad = selected_indices.iter().fold(0, |bad_count, &index| if index < n_bad { bad_count + 1 } else { bad_count });
        num_bad
    }

    // Kick-out nodes as per CCR.
    fn kick_out_nodes(&mut self, region: RegionCount) -> (bool, usize) {
        // This func will always be called after at least 1 malicious node is inserted into the quorum.
        // TODO Check: We only want to chuck out k*g'/g nodes out of the *existing* nodes in the quorum.
        // So effectively, the no. of malicious nodes is decremented by 1 here.
        let num_malicious = self.regions[region].num_malicious - 1;
        let current_region_size = num_malicious + self.regions[region].num_honest;
        let number_to_kick_out : usize = (self.k as f64 * current_region_size as f64/ self.g as f64).round() as usize;
        // println!(
        //     "KICKING OUT: ***Region {}*** had {} honest {} malicious nodes = total of {} nodes.",
        //     region,
        //     self.regions[region].num_honest,
        //     num_malicious,
        //     current_region_size
        // );
        let quorum = region >> self.lg_regions_per_quorum;

        if number_to_kick_out > 0
        {
            let num_malicious_to_kick_out = self.pick_honest_malicious_nodes_to_kick_out(number_to_kick_out, current_region_size, num_malicious);
            let num_honest_to_kick_out = number_to_kick_out - num_malicious_to_kick_out;
            // println!("KICKING OUT: We choose to kick out {} honest {} malicious nodes = total of {} nodes.",
            //          num_honest_to_kick_out,
            //          num_malicious_to_kick_out,
            //          number_to_kick_out
            // );

            self.regions[region].num_malicious -= num_malicious_to_kick_out;
            self.regions[region].num_honest -= num_honest_to_kick_out;
            self.quorums[quorum].tot_malicious -= num_malicious_to_kick_out;
            self.quorums[quorum].tot_honest -= num_honest_to_kick_out;

            // Re-insert each node that was kicked out earlier, into new quorums, while maintaining
            // honest/malicious status.
            for _ctr in 0..num_honest_to_kick_out {
                let secondary_join_region = self.rand_region();
                // Don't need to check the return value as it'll continue to be true (only honest nodes were inserted)
                self.insert(false, secondary_join_region, false);
                // println!(
                //     "KICKING OUT: honest node {} to region {}",
                //     ctr, secondary_join_region
                // );
            }
            for ctr in 0..num_malicious_to_kick_out {
                let secondary_join_region = self.rand_region();
                let below_bmax = self.insert(true, secondary_join_region, false);
                if !below_bmax
                {
                    return (below_bmax, ctr + 1);
                }
                // println!(
                //     "KICKING OUT: malicious node {} to region {}",
                //     ctr, secondary_join_region
                // );
            }
            return (true, num_malicious_to_kick_out)
        }
        (true, 0)
    }

    // Insert a new node into a given region in the DHT.
    fn insert(&mut self, is_malicious: bool, region: RegionCount, is_primary_join: bool) -> bool {
        let quorum = region >> self.lg_regions_per_quorum;

        // Insert the new node into that region.
        // Also update honest/malicious counts for region + quorum.
        if is_malicious {
            self.regions[region].num_malicious += 1;
            self.quorums[quorum].tot_malicious += 1;
        } else {
            self.regions[region].num_honest += 1;
            self.quorums[quorum].tot_honest += 1;
        }
        if is_primary_join {
            self.regions[region].num_nodes_since_last_primary_join = 0
        } else {
            self.regions[region].num_nodes_since_last_primary_join += 1
        }
        if self.done_init {
            return self.cur_stats.update(quorum, &self.quorums[quorum], false);
            // if broke_b_max {
            //     print!("FAILED INVARIANT b_0 < 1/3 while re-inserting cuckoo-ed out nodes");
            // }
        }
        true
    }

    pub fn init(&mut self, num_honest: NodeCount, num_malicious: NodeCount) -> (usize, bool, bool, bool, usize) {
        for _ in 0..num_honest {
            // The original honest nodes are simply "mapped" to random locations -
            let target_region = self.rand_region();
            // Don't need to check the return value as it'll always be true (b_0 max won't be broken as only honest nodes were inserted)
            self.insert(false, target_region, false);
        }
        for ctr  in 0..num_malicious {
            let (inserted, inserted_below_bmax, rejoined_below_bmax, number_rejoined) = self.cuckoo_insert();
            if !inserted || !inserted_below_bmax || !rejoined_below_bmax {
                return (ctr, inserted, inserted_below_bmax, rejoined_below_bmax, number_rejoined );
            }
        }
        self.done_init = true;
        self.cur_stats.dirty = true;
        self.collect_stats();
        self.update_cumstats(true);
        (num_malicious, true, true, true, 0)
    }

    // Insert a new node into a random region in the DHT, then
    // evict a fraction of nodes in that region into random
    // regions in the DHT (but don't evict nodes from the regions
    // _those_ nodes land in).
    pub fn cuckoo_insert(&mut self) -> (bool, bool, bool, usize) {
        let regions_ordering = self.random_shuffle_regions(self.regions.len());
        // println!("Random permutation of target regions: {:?}", regions_ordering);

        for ctr in 0..regions_ordering.len() {
            // Pick a random region to put the new node into. Also get the quorum for that region.
            let region = regions_ordering[ctr] as usize;
            // println!(
            //     "TRYING CUCKOO INSERT: target region {} quorum {} no. of nodes since last 1ary join {} k {}",
            //     region, quorum, self.regions[region].num_nodes_since_last_primary_join, self.k
            // );
            if self.regions[region].num_nodes_since_last_primary_join >= self.k - 1 {
                // println!("------ SUCCESSFUL CUCKOO INSERT: target region {}", region);
                let below_bmax_insert_own = self.insert(true, region, true); // True for primary join
                let (below_bmax_insert_kicked_out, no_of_malicious_kicked_out) = self.kick_out_nodes(region);

                // Update the age of the region that was emptied out.
                // self.quorums[quorum].tot_last_join += self.now - last_join;
                self.now += 1;
                // Cuckoo-ing after initialization: update stats for the quorum with the cuckood-out region
                if self.done_init {
                    // as it's being updated in the collect_stats call essentially.
                    // self.cur_stats.update(quorum, &self.quorums[quorum], false);
                    self.collect_stats();
                    self.update_cumstats(false);
                }
                return (true, below_bmax_insert_own, below_bmax_insert_kicked_out, no_of_malicious_kicked_out)
            }
        }
        (false, true, true, 0)
    }

    // Remove an existing malicious node from the quorum that has the lowest fraction of faulty nodes
    pub fn move_malicious(&mut self) -> (bool, bool, bool, usize) {
        let find_min_b_0_region = || {
            let compute_b_0 = |r: Region| {
                (r.num_malicious as f64) / (r.num_honest as f64 + r.num_malicious as f64)
            };

            let b_0_array: Vec<f64> = self.regions.iter().map(|&q| compute_b_0(q)).collect();
            // eprintln!("b_0 array for");
            // let mut ctr: i8 = 0;
            // for &b_0 in &b_0_array {
            //     print!("{} {},", ctr, b_0);
            //     ctr = ctr + 1;
            // }
            // println!("\n");
            let (min_b_0_region, _) = b_0_array.iter().enumerate().fold((0, 1.0), |(min_index, min_val), (index, &val)| if val < min_val && val > 0.0 {
                // eprintln!(
                //     "Replacing {} index {} with current val {} index {}",
                //     min_val, min_index, val, index
                // );
                (index, val)
            } else {
                (min_index, min_val)
            });

            min_b_0_region
        };

        // Pick quorum with least fraction of byz nodes
        let min_b_0_region = find_min_b_0_region();
        let min_b_0_quorum = min_b_0_region >> self.lg_regions_per_quorum;
        // eprintln!(
        //     "MOVE MALICIOUS out of region {} which has {} honest {} malicious nodes",
        //     min_b_0_region,
        //     self.regions[min_b_0_region].num_honest,
        //     self.regions[min_b_0_region].num_malicious
        // );
        self.regions[min_b_0_region].num_malicious -= 1;
        self.quorums[min_b_0_quorum].tot_malicious -= 1;
        // Don't need to check return value as b_0 for min_b_0_quorum decreases.
        self.cur_stats
            .update(min_b_0_quorum, &self.quorums[min_b_0_quorum], false);

        // Insert it back into the DHT
        return self.cuckoo_insert();
    }

    pub fn collect_stats(&mut self) {
        if self.cur_stats.dirty {
            for (i, q) in self.quorums.iter().enumerate() {
                self.cur_stats.update(i, q, i==0);
            }
            self.cur_stats.dirty = false;
        }
    }

    pub fn update_cumstats(&mut self, force: bool) {
        let stat = &self.cur_stats;
        if force || stat.min_tot_nodes < self.cum_stats.min_tot_nodes {
            self.cum_stats.min_tot_nodes = stat.min_tot_nodes;
        }
        if force || stat.max_tot_nodes > self.cum_stats.max_tot_nodes {
            self.cum_stats.max_tot_nodes = stat.max_tot_nodes;
        }
        if force || stat.min_tot_honest < self.cum_stats.min_tot_honest {
            self.cum_stats.min_tot_honest = stat.min_tot_honest;
        }
        if force || stat.max_tot_honest > self.cum_stats.max_tot_honest {
            self.cum_stats.max_tot_honest = stat.max_tot_honest;
        }
        if force || stat.min_tot_malicious < self.cum_stats.min_tot_malicious {
            self.cum_stats.min_tot_malicious = stat.min_tot_malicious;
        }
        if force || stat.max_tot_malicious > self.cum_stats.max_tot_malicious {
            self.cum_stats.max_tot_malicious = stat.max_tot_malicious;
        }
/*        let min_age = (self.now<<self.lg_regions_per_quorum) - stat.max_tot_last_join;
        let max_age = (self.now<<self.lg_regions_per_quorum) - stat.min_tot_last_join;
        if force || min_age < self.cum_stats.min_age {
            self.cum_stats.min_age = min_age;
        }
        if force || max_age > self.cum_stats.max_age {
            self.cum_stats.max_age = max_age;
        }
*/        if force || stat.min_b_0 < self.cum_stats.min_b_0 {
            self.cum_stats.min_b_0 = stat.min_b_0;
        }
        if force || stat.max_b_0 > self.cum_stats.max_b_0 {
            self.cum_stats.max_b_0 = stat.max_b_0;

        }
    }




}

