extern crate build_a_smart_contract_in_rust;

use wasm_bindgen_test::*;

wasm_bindgen_test::wasm_bindgen_test_configure!(run_in_browser);

#[wasm_bindgen_test]
fn initialise_is_public() {
    use build_a_smart_contract_in_rust::initialise;
    let _t = initialise();
    assert!(true, "This test should compile to pass");
}

#[wasm_bindgen_test]
fn initialise_returns_result() {
    use build_a_smart_contract_in_rust::initialise;

    let res = initialise();
    let _v = res.is_ok();
    assert!(true, "This test should compile to pass");
}

#[wasm_bindgen_test]
fn initialise_returns_new_context() {
    use build_a_smart_contract_in_rust::{initialise, Context};

    let res = initialise();
    match res {
        Ok(v) => {
            let _context: Context = v.into_serde().expect("initialise returns a value serialisable into a Context.");
        },
        Err(_e) => {
            assert!(false, "initialise should not have returned an Err variant");
        }
    };
}

#[wasm_bindgen_test]
fn context_has_base_account() {
    use build_a_smart_contract_in_rust::{initialise, Context, Account};

    let res = initialise();
    match res {
        Ok(v) => {
            let context: Context = v.into_serde().expect("initialise returns a value serialisable into a Context.");
            let _account: Account = context.base_account;
        },
        Err(_e) => {
            assert!(false, "initialise should not have returned an Err variant");
        }
    };
}

#[wasm_bindgen_test]
fn base_account_has_total_clicks() {
    use build_a_smart_contract_in_rust::{initialise, Context, Account};

    let res = initialise();
    match res {
        Ok(v) => {
            let context: Context = v.into_serde().expect("initialise returns a value serialisable into a Context.");
            let account: Account = context.base_account;

            let _total_clicks: u64 = account.total_clicks;
        },
        Err(_e) => {
            assert!(false, "initialise should not have returned an Err variant");
        }
    };
}

#[wasm_bindgen_test]
fn base_account_has_clickers() {
    use build_a_smart_contract_in_rust::{initialise, Context, Account};

    let res = initialise();
    match res {
        Ok(v) => {
            let context: Context = v.into_serde().expect("initialise returns a value serialisable into a Context.");
            let account: Account = context.base_account;

            let _clickers: Vec<String> = account.clickers;
        },
        Err(_e) => {
            assert!(false, "initialise should not have returned an Err variant");
        }
    };
}

#[wasm_bindgen_test]
fn total_clicks_is_initialised_with_0() {
    use build_a_smart_contract_in_rust::{initialise, Context, Account};

    let res = initialise();
    match res {
        Ok(v) => {
            let context: Context = v.into_serde().expect("initialise returns a value serialisable into a Context.");
            let account: Account = context.base_account;

            let total_clicks: u64 = account.total_clicks;
            assert_eq!(total_clicks, 0);
        },
        Err(_e) => {
            assert!(false, "initialise should not have returned an Err variant");
        }
    };
}

#[wasm_bindgen_test]
fn clickers_is_initialised_with_empty_vec() {
    use build_a_smart_contract_in_rust::{initialise, Context, Account};

    let res = initialise();
    match res {
        Ok(v) => {
            let context: Context = v.into_serde().expect("initialise returns a value serialisable into a Context.");
            let account: Account = context.base_account;

            let clickers: Vec<String> = account.clickers;
            let com: Vec<String> = vec![];
            assert_eq!(clickers, com);
        },
        Err(_e) => {
            assert!(false, "initialise should not have returned an Err variant");
        }
    };
}


extern crate build_a_smart_contract_in_rust;

use build_a_smart_contract_in_rust::{Account, Context};
use wasm_bindgen::{JsValue, JsError};
use wasm_bindgen_test::*;

wasm_bindgen_test::wasm_bindgen_test_configure!(run_in_browser);

#[wasm_bindgen_test]
fn is_public() {
    use build_a_smart_contract_in_rust::set_click;
    let _t = set_click;
    assert!(true, "This test should just compile to pass");
}

#[wasm_bindgen_test]
fn returns_result() {
    use build_a_smart_contract_in_rust::set_click;

    let res = set_click(js_con(&fix_context()), String::from("quincy"));
    let _v = res.is_ok();
    assert!(true, "This test should just compile to pass");
}

#[wasm_bindgen_test]
fn returns_context() {
    use build_a_smart_contract_in_rust::set_click;

    let res = set_click(js_con(&fix_context()), String::from("Ahmad"));
    match res {
        Ok(v) => {
            let _context: Context = v.into_serde().expect("set_click returns a value serialisable into a Context.");
        },
        Err(_e) => {
            assert!(false, "set_click should not have returned an Err variant");
        }
    };
}

#[wasm_bindgen_test]
fn pushes_clicker_and_increments_total_clicks() {
    use build_a_smart_contract_in_rust::set_click;

    let mut context = fix_context();
    let mut res = set_click(js_con(&context), "shaun".to_string());
    context = js_to_con(&res);
    assert_eq!(context.base_account.total_clicks, 11);
    assert_eq!(context.base_account.clickers.len(), 3);
    assert_eq!(context.base_account.clickers[2], "shaun".to_string());

    res = set_click(js_con(&context), "quincy".to_string());
    assert_eq!(js_to_con(&res).base_account.total_clicks, 12);
    assert_eq!(js_to_con(&res).base_account.clickers.len(), 4);
    assert_eq!(js_to_con(&res).base_account.clickers[3], "quincy".to_string());
    match res {
        Ok(v) => {
            let _context: Context = v.into_serde().expect("set_click returns a value serialisable into a Context.");
        },
        Err(_e) => {
            assert!(false, "set_click should not have returned an Err variant");
        }
    };
}

fn fix_context() -> Context {
    let acc = Account {
        total_clicks: 10,
        clickers: vec![String::from("tom"), String::from("shaun")]
    };
    let con = Context {
        base_account: acc
    };
    con
}

fn js_con(context: &Context) -> JsValue {
    JsValue::from_serde(context).expect("can turn context into JsValue")
}

fn js_to_con(js: &Result<JsValue, JsError>) -> Context {
    match js {
        Ok(v) => v.into_serde().expect("can deserialise JsValue into Context"),
        Err(_e) => panic!("set_click should not have returned an Err variant")
    }
}


//! # Block
//!
//! A block is a piece of data that is stored in the blockchain.

use serde::{Deserialize, Serialize};

use crate::node::Node;

/// The block added to the chain of the blockchain.
///
/// **Note:** This is a reference type, and does not contain any implementations.
#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct Block {
    pub id: u64,
    pub hash: String,
    pub previous_hash: String,
    pub timestamp: u64, // TODO: Consider removing if testing too difficult
    pub data: Vec<Node>,
    pub nonce: u64,
    pub next_miner: String,
    pub next_validators: Vec<String>,
}


//! # Chain
//!
//! A chain represents the main data of the blockchain, and is passed in full between Nodes.

use rand::Rng;
use serde::{Deserialize, Serialize};

use crate::{block::Block, calculate_hash, hash_to_binary, node::Node, DIFFICULTY_PREFIX};

/// The chain consists of the immutable `chain` data, and the changing `network` data.
#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct Chain {
    pub chain: Vec<Block>,
    pub network: Vec<String>,
}

impl Chain {
    /// Creates a `Chain` with empty `chain` and `network` data.
    pub fn new() -> Self {
        Self {
            chain: vec![],
            network: vec![],
        }
    }

    /// Returns the last block of the current `Chain`, if it exists. Otherwise, returns `None`.
    ///
    /// # Examples
    ///
    /// ```
    /// let mut chain = Chain::new();
    /// let last_block = chain.last_block();
    /// assert!(last_block.is_none());
    /// ```
    ///
    /// If chain is initialised by Node:
    ///
    /// ```javascript
    /// const chain = initialise("node_1");
    /// assert.equal(chain.chain.length, 1);
    /// ```
    ///
    /// ```
    /// assert_eq!(chain.last_block().unwrap().id, 0);
    /// ```
    pub fn get_last_block(&self) -> Option<Block> {
        match self.chain.last() {
            Some(block) => Some(block.clone()),
            None => None,
        }
    }

    /// Applies weighted, random selection to all `Node`s in the `Chain` to determine the `next_miner` for the next block.
    ///
    /// **Note:** Defaults to returning `"Camper"`, if no `Node`s are present in the `Chain`.
    ///
    /// # Examples
    ///
    /// ```
    /// let chain = Chain::new();
    /// let next_miner = chain.get_next_miner();
    /// assert_eq!(next_miner, String::from("Camper"));
    /// ```
    pub fn get_next_miner(&self) -> String {
        let mut nodes: Vec<&Node> = self.get_nodes();

        nodes.sort_by(|a, b| a.weight_as_miner().cmp(&b.weight_as_miner()));
        let cumulative_weight = nodes
            .iter()
            .fold(0, |acc, node| acc + node.weight_as_miner());
        let cumulative_weights = nodes
            .iter()
            .map(|node| node.weight_as_miner() as f64 / cumulative_weight as f64)
            .collect::<Vec<f64>>();

        let rand_num = rand::thread_rng().gen::<f64>();
        let mut ind = 0;
        for (i, c_m) in cumulative_weights.iter().enumerate() {
            ind = i;
            if c_m > &rand_num {
                break;
            }
        }
        let next_miner = match nodes.get(ind) {
            Some(node) => node.name.clone(),
            None => "Camper".to_string(),
        };

        next_miner
    }

    /// Applies weighted, random selection to all `Node`s in the `Chain` to determine the `next_validators` for the next block.
    pub fn get_next_validators(&self, next_miner: &String) -> Vec<String> {
        let num_of_nodes = self.network.len();
        let next_miner_reputation = match self.get_node_by_name(next_miner) {
            Some(node) => node.reputation,
            None => 0,
        };

        let mut nodes = self.get_nodes();

        let mut max_reputation = 0;
        for node in nodes.iter() {
            if node.reputation > max_reputation {
                max_reputation = node.reputation;
            }
        }

        // The greater the reputation, the fewer validators are needed
        let num_needed_validators = (max_reputation - next_miner_reputation) as usize;
        let num_needed_validators = if num_needed_validators > num_of_nodes {
            num_of_nodes
        } else {
            num_needed_validators + 1
        };

        nodes.sort_by(|a, b| a.weight_as_validator().cmp(&b.weight_as_validator()));

        let cumulative_weight = nodes
            .iter()
            .fold(0, |acc, node| acc + node.weight_as_validator());
        let cumulative_weights = nodes
            .iter()
            .map(|node| node.weight_as_validator() as f64 / cumulative_weight as f64)
            .collect::<Vec<f64>>();

        println!(
            "Cumulative Weight: {:?}\nCumulative Weights: {:?}\n\n",
            cumulative_weight, cumulative_weights
        );
        let mut next_validators = vec![];
        for _ in 0..num_needed_validators {
            let rand_num = rand::thread_rng().gen::<f64>();
            let mut ind = 0;
            for (i, c_m) in cumulative_weights.iter().enumerate() {
                if c_m > &rand_num {
                    break;
                }
                ind = i;
            }
            let validator = match nodes.get(ind) {
                Some(node) => node.name.clone(),
                None => "Camper".to_string(),
            };
            next_validators.push(validator);
        }
        next_validators
    }

    /// Returns the `Node` with the given `name` if it exists in the `Chain`. Otherwise, returns `None`.
    ///
    /// # Examples
    ///
    /// ```
    /// let mut chain = Chain::new();
    /// assert!(chain.get_node_by_name("Camper").is_none());
    /// ```
    pub fn get_node_by_name(&self, name: &str) -> Option<&Node> {
        // Search Chain data in reverse
        for block in self.chain.iter().rev() {
            for node in block.data.iter() {
                if node.name == name {
                    return Some(node);
                }
            }
        }
        None
    }

    /// Returns a `Vec` of all `Node`s in the `Chain`.
    pub fn get_nodes(&self) -> Vec<&Node> {
        let mut nodes = vec![];
        for block in self.chain.iter().rev() {
            for node in block.data.iter() {
                // If node.name is not in nodes, add it
                if !nodes.iter().any(|n: &&Node| n.name == node.name) {
                    nodes.push(node);
                }
            }
        }
        nodes
    }

    /// Mines the given `data` into a new `Block` on the `Chain`.
    pub fn mine_block(&mut self, data: Vec<Node>) {
        println!("\nMining Block...");
        let mut nonce = 0;

        let id = self.chain.len() as u64;
        let next_miner = self.get_next_miner();
        let next_validators = self.get_next_validators(&next_miner);
        let previous_hash = match self.get_last_block() {
            Some(block) => block.hash.clone(),
            None => "".to_string(),
        };
        let timestamp = chrono::Utc::now().timestamp() as u64;
        loop {
            if nonce % 100_000 == 0 {
                println!("Trying nonce: {}", nonce);
            }

            let hash = calculate_hash(
                &data,
                id,
                &next_miner,
                &next_validators,
                nonce,
                &previous_hash,
                timestamp,
            );
            let bin_hash = hash_to_binary(&hash);
            if bin_hash.starts_with(DIFFICULTY_PREFIX) {
                let new_block = Block {
                    id,
                    hash: bin_hash,
                    previous_hash,
                    timestamp,
                    data,
                    nonce,
                    next_miner,
                    next_validators,
                };
                self.chain.push(new_block);
                break;
            }
            nonce += 1;
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn new_chain_returns_empty_vec() {
        let chain: Chain = Chain::new();
        assert_eq!(chain.chain.len(), 0);
    }
    #[test]
    fn get_last_block_returns_none_when_chain_is_empty() {
        let chain: Chain = Chain::new();
        assert!(chain.get_last_block().is_none());
    }
    #[test]
    fn get_last_block_returns_last_block_when_chain_is_not_empty() {
        let chain = _fixture_chain();
        assert!(chain.get_last_block().is_some());
    }
    #[test]
    fn get_next_miner_returns_different_miner_when_chain_is_not_empty() {
        let chain = _fixture_chain();
        // run 100 times, break if miner is different
        let mut i = 0;
        let a = loop {
            let previous_miner = chain.get_last_block().unwrap().next_miner.clone();
            let next_miner = chain.get_next_miner();
            if previous_miner != next_miner {
                break true;
            }
            if i == 99 {
                break false;
            }
            i += 1;
        };
        assert!(a);
    }
    #[test]
    fn get_next_validators_returns_different_validators_when_chain_is_not_empty() {
        let chain = _fixture_chain();
        // run 100 times, break if validators are different
        let mut i = 0;
        let a = loop {
            let previous_validators = chain.get_last_block().unwrap().next_validators.clone();
            let next_validators = chain.get_next_validators(&chain.get_next_miner());
            if previous_validators != next_validators {
                break true;
            }
            if i == 99 {
                break false;
            }
            i += 1;
        };
        assert!(a);
    }
    #[test]
    fn get_node_by_name_returns_none_when_node_is_not_in_chain() {
        let chain = _fixture_chain();
        assert!(chain.get_node_by_name("node_not_in_chain").is_none());
    }
    #[test]
    fn get_node_by_name_returns_node_when_node_is_in_chain() {
        let chain = _fixture_chain();
        assert!(chain.get_node_by_name("Camper").is_some());
    }
    #[test]
    fn get_nodes_returns_all_nodes_in_chain() {
        let chain = _fixture_chain();
        let nodes = chain.get_nodes();
        assert_eq!(nodes.len(), 4);
    }
    #[test]
    fn mine_block_does_not_panic() {
        let mut chain = _fixture_chain();
        chain.mine_block(vec![Node::new("Shaun")]);
        assert_eq!(chain.chain.len(), 3);
    }

    fn _fixture_chain<'a>() -> Chain {
        let mut chain = Chain::new();
        let mut camper = Node::new("Camper");
        camper.reputation = 1;
        camper.tokens = 10;
        camper.staked = 5;
        let mut tom = Node::new("Tom");
        tom.reputation = 2;
        tom.tokens = 20;
        tom.staked = 10;
        let mut mrugesh = Node::new("Mrugesh");
        mrugesh.reputation = 4;
        mrugesh.tokens = 100;
        mrugesh.staked = 80;
        let mut ahmad = Node::new("Ahmad");
        ahmad.reputation = 3;
        ahmad.tokens = 30;
        ahmad.staked = 22;

        let data = vec![camper, tom, mrugesh];
        chain.network = vec![
            String::from("Camper"),
            String::from("Tom"),
            String::from("Mrugesh"),
        ];
        chain.mine_block(data);
        let data = vec![ahmad];
        chain.network.push(String::from("Ahmad"));
        chain.mine_block(data);
        chain
    }
}


//! # Blockchain
//!
//! `blockchain` is a WASM module for handling a Proof of Stake blockchain.

pub mod block;
pub mod chain;
pub mod node;

use block::Block;
use chain::Chain;
use node::Node;

use rand::Rng;
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use wasm_bindgen::prelude::*;

/// Increasing the number of leading zeros in the hash of a block increases the difficulty of mining a block.
pub static DIFFICULTY_PREFIX: &str = "0";

/// Events that can be emitted in the `event` field of a `Transaction`.
#[derive(Serialize, Deserialize, Debug, Clone)]
pub enum Events {
    BuyRack,
    BlockInvalidated,
    Stake,
    SubmitTask,
    Unstake,
    UpdateChain,
}

/// A transaction describes the change which needs to be mined into a block. The transaction is associated with the `name` of a Node.
#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct Transaction {
    pub event: Events,
    pub name: String,
}

/// The current state of the Node calling the API.
#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct NodeState {
    pub chain: Chain,
    pub transactions: Vec<Transaction>,
    pub task_valid: bool,
}

/// The structure of the response from the Blockchain API.
#[derive(Serialize, Deserialize, Debug)]
pub struct Res {
    pub chain: Chain,
    pub errors: Vec<String>,
}

/// Mines the next block onto the given chain passed in the `node_state` argument.
///
/// # Examples
///
/// ```js
/// const nodeState = {
///   chain: {
///     chain: [],
///     network: ["node_1", "node_2"],
///   },
///   transactions: [
///     {
///       event: "UpdateChain",
///       name: "node_2"
///     }
///   ],
///   task_valid: true
/// };
/// const result = handle_mine(nodeState);
/// ```
#[wasm_bindgen]
pub fn handle_mine(node_state: JsValue) -> Result<JsValue, JsError> {
    let node_state: NodeState = node_state.into_serde()?;
    let mut chain = node_state.chain;
    let mut errors: Vec<String> = vec![];

    let mut unique_nodes: Vec<Node> = vec![];
    for transaction in node_state.transactions.iter() {
        if let Some(node) = chain.get_node_by_name(&transaction.name) {
            if !unique_nodes.iter().any(|n| n.name == node.name) {
                let node = node.clone();
                unique_nodes.push(node);
            }
        }
    }

    for transaction in node_state.transactions.iter() {
        if let Some(node) = unique_nodes.iter_mut().find(|n| n.name == transaction.name) {
            match transaction.event {
                Events::Unstake => {
                    if node.can_unstake() {
                        node.staked -= 1;
                    } else {
                        errors.push(format!("{} cannot unstake", node.name));
                    }
                }
                Events::Stake => {
                    if node.can_stake() {
                        node.staked += 1;
                    } else {
                        errors.push(format!("{} cannot stake", node.name));
                    }
                }
                Events::BuyRack => {
                    if node.can_buy_rack() {
                        node.racks += 1;
                        node.tokens -= 10;
                    } else {
                        errors.push(format!("{} cannot buy rack", node.name));
                    }
                }
                Events::BlockInvalidated => {
                    if node.can_punish() {
                        if node.tokens == node.staked {
                            node.staked -= 1;
                        }
                        node.tokens -= 1;
                        node.reputation -= 1;
                    } else {
                        errors.push(format!("{} cannot be punished", node.name));
                    }
                }
                Events::SubmitTask => {
                    if node_state.task_valid {
                        node.tokens += 1;
                        // More staked tokens == higher chance of reputation increasing
                        if rand::thread_rng().gen_range(node.staked..=node.tokens) == node.tokens {
                            node.reputation += 1;
                        }
                    } else {
                        if node.can_punish() {
                            if node.tokens == node.staked {
                                node.staked -= 1;
                            }
                            node.tokens -= 1;
                            node.reputation -= 1;
                        } else {
                            errors.push(format!("{} cannot be punished", node.name));
                        }
                    }
                }
                _ => {
                    errors.push(format!("{} transacted with an invalid event", node.name));
                }
            };
        } else {
            match transaction.event {
                Events::UpdateChain => {
                    // Add node to chain
                    unique_nodes.push(Node::new(&transaction.name));
                }
                _ => {
                    errors.push(format!("{} not found in chain", transaction.name));
                }
            };
        }
    }

    if errors.len() == node_state.transactions.len() || unique_nodes.len() == 0 {
        return Err(JsError::new("Invalid transactions. No change in chain"));
    }
    chain.mine_block(unique_nodes);
    // let response = Res { chain, errors };
    Ok(JsValue::from_serde(&(chain, errors))?)
}

/// Validates whether the provided `chain` argument is valid for the latest two blocks in the chain.
///
/// # Examples
///
/// ```js
/// const chain = {
///   chain: [
///    {
///     id: 0,
///     hash: "01101101",
///     previous_hash: "genesis",
///     timestamp: 123456789,
///     data: [
///       {
///         name: "node_1",
///         tokens: 20,
///         staked: 0,
///         reputation: 1,
///         racks: 0
///       }],
///     nonce: 0,
///     next_miner: "node_1",
///     next_validators: ["node_2"]
///    }],
///   network: ["node_1", "node_2"]
/// };
/// const isChainValid = handle_validate(chain);
/// assert.equal(isChainValid, true);
/// ```
///
/// # Errors
///
/// If `chain` argument is not deserialisable into type `Chain`, a `JsError` is thrown.
#[wasm_bindgen]
pub fn handle_validate(chain: JsValue) -> Result<bool, JsError> {
    let chain: Chain = chain.into_serde()?;
    if let Some(previous_block) = chain.chain.get(chain.chain.len() - 2) {
        let last_block: Block = match chain.get_last_block() {
            Some(block) => block,
            None => return Err(JsError::new("Chain is empty")),
        };
        Ok(Node::validate_block(&last_block, previous_block))
    } else {
        Err(JsError::new("Chain is too short"))
    }
}

/// Initialise a new blockchain, and returns the corresponding chain.
/// This is only to be called by the first Node starting the network.
#[wasm_bindgen]
pub fn initialise(name: String) -> Result<JsValue, JsError> {
    let mut chain: Chain = Chain::new();

    // Create and mine genesis block
    let data = vec![Node::new(&name)];
    chain.mine_block(data);

    Ok(JsValue::from_serde(&chain)?)
}

/// Takes a hash slice, and returns the binary representation.
pub fn hash_to_binary(hash: &[u8]) -> String {
    let mut res: String = String::default();
    for c in hash {
        res.push_str(&format!("{:b}", c));
    }
    res
}

/// Uses `Sha256` to calculate the hash from a `serde_json::Value` of the input arguments.
pub fn calculate_hash(
    data: &Vec<Node>,
    id: u64,
    next_miner: &String,
    next_validators: &Vec<String>,
    nonce: u64,
    previous_hash: &str,
    timestamp: u64,
) -> Vec<u8> {
    let data = serde_json::json!({
        "id": id,
        "previous_hash": previous_hash,
        "data": data,
        "timestamp": timestamp,
        "nonce": nonce,
        "next_miner": next_miner,
        "next_validators": next_validators,
    });
    let mut hasher = Sha256::new();
    hasher.update(data.to_string().as_bytes());
    hasher.finalize().as_slice().to_owned()
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn difficulty_is_not_too_high() {
        assert!(DIFFICULTY_PREFIX.len() <= 3);
    }
    #[test]
    fn calculate_hash_works() {
        let data = vec![Node::new("test")];
        let hash = calculate_hash(
            &data,
            1,
            &"test".to_string(),
            &vec!["test".to_string()],
            1,
            &"test".to_string(),
            1,
        );
        assert_eq!(hash.len(), 32);
    }
    #[test]
    fn hash_to_binary_works() {
        let hash = vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15];
        let hash_str = hash_to_binary(&hash);
        assert_eq!(hash_str.len(), 50);
    }
}


//! # Node
//!
//! A node is a piece of data that is stored in a `BLock` on the blockchain.

use serde::{Deserialize, Serialize};

use crate::{block::Block, calculate_hash, hash_to_binary, DIFFICULTY_PREFIX};

#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct Node {
    pub name: String,
    pub racks: u64,
    pub reputation: u64,
    pub staked: u64,
    pub tokens: u64,
}

impl Node {
    /// Creates a new Node with the given name.
    ///
    /// # Examples
    ///
    /// ```
    /// let node = Node::new("Camper");
    /// assert_eq!(node.name, "Camper");
    /// ```
    pub fn new(name: &str) -> Self {
        Self {
            name: name.to_string(),
            staked: 0,
            tokens: 20,
            racks: 0,
            reputation: 1,
        }
    }

    /// Check if Node can afford a server rack
    pub fn can_buy_rack(&self) -> bool {
        self.tokens - self.staked >= 10
    }
    /// Check if a Node can stake, by checking if it has any unstaked tokens
    pub fn can_stake(&self) -> bool {
        self.tokens > self.staked
    }
    /// Check if a Node can unstake, by checking if it has any staked tokens
    pub fn can_unstake(&self) -> bool {
        self.staked > 0
    }
    /// Check if a Node can be punished, by checking if it has any tokens and reputation
    pub fn can_punish(&self) -> bool {
        self.tokens > 0 && self.reputation > 0
    }
    /// Calculates the miner weight of Node
    pub fn weight_as_miner(&self) -> u64 {
        self.reputation * self.staked
    }
    /// Calculates the validator weight of Node
    pub fn weight_as_validator(&self) -> u64 {
        self.reputation
    }

    /// Validates if two adjacent blocks have been correctly mined
    pub fn validate_block(block: &Block, previous_block: &Block) -> bool {
        if block.previous_hash != previous_block.hash {
            println!("block with id: {} has wrong previous hash", block.id);
            return false;
        } else if !&block.hash.starts_with(DIFFICULTY_PREFIX) {
            println!("block with id: {} has invalid difficulty", block.id);
            return false;
        } else if block.id != previous_block.id + 1 {
            println!(
                "block with id: {} is not the next block after the latest: {}",
                block.id, previous_block.id
            );
            return false;
        } else if hash_to_binary(&calculate_hash(
            &block.data,
            block.id,
            &block.next_miner,
            &block.next_validators,
            block.nonce,
            &block.previous_hash,
            block.timestamp,
        )) != block.hash
        {
            println!("block with id: {} has invalid hash", block.id);
            return false;
        }
        true
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn new_always_creates_same_node() {
        let node1 = Node::new("node1");
        let node2 = Node::new("node1");
        assert_eq!(node1.name, node2.name);
        assert_eq!(node1.staked, node2.staked);
        assert_eq!(node1.tokens, node2.tokens);
        assert_eq!(node1.racks, node2.racks);
        assert_eq!(node1.reputation, node2.reputation);
    }
    #[test]
    fn cannot_buy_rack_when_all_staked() {
        let node = _fixture_nodes().0;
        assert!(!node.can_buy_rack());
    }
    #[test]
    fn can_buy_rack_when_all_unstaked() {
        let node = _fixture_nodes().1;
        assert!(node.can_buy_rack());
    }
    #[test]
    fn can_buy_rack_when_no_rep() {
        let node = _fixture_nodes().2;
        assert!(node.can_buy_rack());
    }
    #[test]
    fn cannot_buy_rack_when_no_tokens() {
        let node = _fixture_nodes().3;
        assert!(!node.can_buy_rack());
    }
    #[test]
    fn cannot_stake_when_all_staked() {
        let node = _fixture_nodes().0;
        assert!(!node.can_stake());
    }
    #[test]
    fn can_stake_when_all_unstaked() {
        let node = _fixture_nodes().1;
        assert!(node.can_stake());
    }
    #[test]
    fn can_stake_when_no_rep() {
        let node = _fixture_nodes().2;
        assert!(node.can_stake());
    }
    #[test]
    fn cannot_stake_when_no_tokens() {
        let node = _fixture_nodes().3;
        assert!(!node.can_stake());
    }
    #[test]
    fn can_unstake_when_all_staked() {
        let node = _fixture_nodes().0;
        assert!(node.can_unstake());
    }
    #[test]
    fn cannot_unstake_when_all_unstaked() {
        let node = _fixture_nodes().1;
        assert!(!node.can_unstake());
    }
    #[test]
    fn can_unstake_when_no_rep() {
        let node = _fixture_nodes().2;
        assert!(node.can_unstake());
    }
    #[test]
    fn cannot_unstake_when_no_tokens() {
        let node = _fixture_nodes().3;
        assert!(!node.can_unstake());
    }
    #[test]
    fn all_staked_miner_weight() {
        let node = _fixture_nodes().0;
        assert_eq!(node.weight_as_miner(), 800);
    }
    #[test]
    fn all_unstaked_miner_weight() {
        let node = _fixture_nodes().1;
        assert_eq!(node.weight_as_miner(), 0);
    }
    #[test]
    fn no_rep_miner_weight() {
        let node = _fixture_nodes().2;
        assert_eq!(node.weight_as_miner(), 0);
    }
    #[test]
    fn no_tokens_miner_weight() {
        let node = _fixture_nodes().3;
        assert_eq!(node.weight_as_miner(), 0);
    }
    #[test]
    fn all_staked_validator_weight() {
        let node = _fixture_nodes().0;
        assert_eq!(node.weight_as_validator(), 8);
    }
    #[test]
    fn all_unstaked_validator_weight() {
        let node = _fixture_nodes().1;
        assert_eq!(node.weight_as_validator(), 8);
    }
    #[test]
    fn no_rep_validator_weight() {
        let node = _fixture_nodes().2;
        assert_eq!(node.weight_as_validator(), 0);
    }
    #[test]
    fn no_tokens_validator_weight() {
        let node = _fixture_nodes().3;
        assert_eq!(node.weight_as_validator(), 1);
    }
    #[test]
    fn invalidate_block_unequal_previous_hash() {
        let previous_block = _fixture_blocks().0;
        let mut block = _fixture_blocks().1;
        block.previous_hash = block.previous_hash.replace("1", "0");
        assert!(!Node::validate_block(&block, &previous_block));
    }
    #[test]
    fn invalidate_block_hash_not_start_with_difficulty() {
        let previous_block = _fixture_blocks().0;
        let mut block = _fixture_blocks().1;
        block.hash = block.previous_hash.replace("0", "1"); //"011111101111000110011110001110011001111011011010011011000101110010100101001001111001110001011010111000010011000010000100101000011010111110001110010011110101011000011101011110011001110010111001011011011111111010110100000".to_string();
        assert!(!Node::validate_block(&block, &previous_block));
    }
    #[test]
    fn invalidate_block_id_not_incremented() {
        let previous_block = _fixture_blocks().0;
        let mut block = _fixture_blocks().1;
        block.id = previous_block.id;
        assert!(!Node::validate_block(&block, &previous_block));
    }
    #[test]
    fn invalidate_encoded_hash_not_correct() {
        let previous_block = _fixture_blocks().0;
        let mut block = _fixture_blocks().1;
        block.hash = "001111101111000110011110001110011001111011011010011011000101110010100101001001111001110001011010111000010011000010000100101000011010111110001110010011110101011000011101011110011001110010111001011011011111111010110100000".to_string();
        assert!(!Node::validate_block(&block, &previous_block));
    }
    #[test]
    fn validate_block_correct() {
        let previous_block = _fixture_blocks().0;
        let block = _fixture_blocks().1;
        assert!(Node::validate_block(&block, &previous_block));
    }

    fn _fixture_nodes() -> (Node, Node, Node, Node) {
        let all_staked = Node {
            name: "Shaun".to_string(),
            staked: 100,
            tokens: 100,
            racks: 5,
            reputation: 8,
        };
        let all_unstaked = Node {
            name: "Tom".to_string(),
            staked: 0,
            tokens: 100,
            racks: 5,
            reputation: 8,
        };
        let no_rep = Node {
            name: "Ahmad".to_string(),
            staked: 80,
            tokens: 100,
            racks: 5,
            reputation: 0,
        };
        let no_tokens = Node {
            name: "Quincy".to_string(),
            staked: 0,
            tokens: 0,
            racks: 1,
            reputation: 1,
        };
        (all_staked, all_unstaked, no_rep, no_tokens)
    }
    fn _fixture_blocks() -> (Block, Block) {
        let id = 0;
        let previous_hash = String::new();
        let data = vec![
            Node {
                name: "Camper".to_string(),
                staked: 0,
                tokens: 10,
                racks: 1,
                reputation: 1,
            },
            Node {
                name: "Tom".to_string(),
                staked: 0,
                tokens: 10,
                racks: 1,
                reputation: 0,
            },
            Node {
                name: "Mrugesh".to_string(),
                staked: 0,
                tokens: 10,
                racks: 1,
                reputation: 0,
            },
        ];
        let nonce = 83;
        let next_miner = "Camper".to_string();
        let next_validators = vec![
            "Camper".to_string(),
            "Camper".to_string(),
            "Camper".to_string(),
        ];
        let timestamp = 1648987026;
        let hash = hash_to_binary(&calculate_hash(
            &data,
            id,
            &next_miner,
            &next_validators,
            nonce,
            &previous_hash,
            timestamp,
        ));
        let genesis = Block {
            id,
            hash: hash.clone(),
            previous_hash,
            timestamp,
            data,
            nonce,
            next_miner,
            next_validators,
        };

        // ----------------------------
        let id = 1;
        let previous_hash = hash;
        let timestamp = 1648996144;
        let data = vec![Node {
            name: "Ahmad".to_string(),
            staked: 0,
            tokens: 10,
            racks: 1,
            reputation: 0,
        }];
        let nonce = 306;
        let next_miner = "Mrugesh".to_string();
        let next_validators = vec!["Mrugesh".to_string(), "Mrugesh".to_string()];
        let hash = hash_to_binary(&calculate_hash(
            &data,
            id,
            &next_miner,
            &next_validators,
            nonce,
            &previous_hash,
            timestamp,
        ));
        let block = Block {
            id,
            hash,
            previous_hash,
            timestamp,
            data,
            nonce,
            next_miner,
            next_validators,
        };

        (genesis, block)
    }
}


extern crate blockchain;

use blockchain::{chain::Chain, handle_mine, node::Node, Events, NodeState, Transaction};
use wasm_bindgen::JsValue;
use wasm_bindgen_test::*;
use web_sys::{console, ErrorEvent};

wasm_bindgen_test::wasm_bindgen_test_configure!(run_in_browser);

#[wasm_bindgen_test]
fn stake() {
    let data = vec![Node::new("Camper")];
    let mut fix_node_state = fix(data);
    fix_node_state.transactions[0].event = Events::Stake;
    fix_node_state.transactions[0].name = "Camper".to_string();

    let (chain, _) = mine(fix_node_state).expect("result to be chain");
    assert_eq!(chain.get_last_block().unwrap().data[0].staked, 1);
}

#[wasm_bindgen_test]
fn buy_rack() {
    let data = vec![Node::new("Camper")];
    let mut fix_node_state = fix(data);
    fix_node_state.transactions[0].event = Events::BuyRack;
    let (chain, _) = mine(fix_node_state).expect("result to be chain");
    assert_eq!(chain.get_last_block().unwrap().data[0].racks, 1);
    assert_eq!(chain.get_last_block().unwrap().data[0].tokens, 10);
}

#[wasm_bindgen_test]
fn block_invalidated() {
    let data = vec![Node::new("Camper")];
    let mut fix_node_state = fix(data);
    fix_node_state.transactions[0].event = Events::BlockInvalidated;
    let (chain, _) = mine(fix_node_state).expect("result to be chain");
    assert_eq!(chain.get_last_block().unwrap().data[0].racks, 0);
    assert_eq!(chain.get_last_block().unwrap().data[0].tokens, 19);
    assert_eq!(chain.get_last_block().unwrap().data[0].reputation, 0);
}

#[wasm_bindgen_test]
fn unstake() {
    let mut camper = Node::new("Camper");
    camper.staked = 1;
    let data = vec![camper];
    let mut fix_node_state = fix(data);
    fix_node_state.transactions[0].event = Events::Unstake;
    let (chain, _) = mine(fix_node_state).expect("result to be chain");
    assert_eq!(chain.get_last_block().unwrap().data[0].staked, 0);
}

#[wasm_bindgen_test]
fn update_chain() {
    let mut camper = Node::new("Camper");
    camper.staked = 1;
    let data = vec![camper];
    let mut fix_node_state = fix(data);
    fix_node_state.transactions[0].event = Events::UpdateChain;
    fix_node_state.transactions[0].name = "Tom".to_string();
    let (chain, _) = mine(fix_node_state).expect("result to be chain");
    assert!(chain.get_node_by_name("Tom").is_some());
    assert_eq!(chain.get_nodes().len(), 2);
}

#[wasm_bindgen_test]
fn all_invalid_buy_rack() {
    let mut camper = Node::new("Camper");
    camper.staked = 11;
    let data = vec![camper];
    let mut fix_node_state = fix(data);
    fix_node_state.transactions[0].event = Events::BuyRack;
    let chain_res = mine(fix_node_state);
    assert!(chain_res.is_err());
    if let Err(e) = chain_res {
        // Get the error message
        let error_message = ErrorEvent::from(e);
        assert_eq!(
            error_message.message(),
            "Invalid transactions. No change in chain"
        );
    }
}

#[wasm_bindgen_test]
fn all_invalid_unstake() {
    let mut camper = Node::new("Camper");
    camper.staked = 0;
    let data = vec![camper];
    let mut fix_node_state = fix(data);
    fix_node_state.transactions[0].event = Events::Unstake;
    let chain_res = mine(fix_node_state);
    assert!(chain_res.is_err());
    if let Err(e) = chain_res {
        // Get the error message
        let error_message = ErrorEvent::from(e);
        assert_eq!(
            error_message.message(),
            "Invalid transactions. No change in chain"
        );
    }
}

#[wasm_bindgen_test]
fn all_invalid_stake() {
    let mut camper = Node::new("Camper");
    camper.staked = 20;
    let data = vec![camper];
    let mut fix_node_state = fix(data);
    fix_node_state.transactions[0].event = Events::Stake;
    let chain_res = mine(fix_node_state);
    assert!(chain_res.is_err());
    if let Err(e) = chain_res {
        // Get the error message
        let error_message = ErrorEvent::from(e);
        assert_eq!(
            error_message.message(),
            "Invalid transactions. No change in chain"
        );
    }
}

#[wasm_bindgen_test]
fn all_invalid_punish() {
    let mut camper = Node::new("Camper");
    camper.tokens = 0;
    let data = vec![camper];
    let mut fix_node_state = fix(data);
    fix_node_state.transactions[0].event = Events::BlockInvalidated;
    let chain_res = mine(fix_node_state);
    assert!(chain_res.is_err());
    if let Err(e) = chain_res {
        // Get the error message
        let error_message = ErrorEvent::from(e);
        assert_eq!(
            error_message.message(),
            "Invalid transactions. No change in chain"
        );
    }
}

#[wasm_bindgen_test]
fn all_invalid_find_node() {
    let mut camper = Node::new("Camper");
    camper.staked = 1;
    let data = vec![camper];
    let mut fix_node_state = fix(data);
    fix_node_state.transactions[0].event = Events::Stake;
    fix_node_state.transactions[0].name = "Tom".to_string();
    let chain_res = mine(fix_node_state);
    console::log_1(&format!("{:?}", chain_res).into());
    assert!(chain_res.is_err());
    if let Err(e) = chain_res {
        // Get the error message
        let error_message = ErrorEvent::from(e);
        assert_eq!(
            error_message.message(),
            "Invalid transactions. No change in chain"
        );
    }
}

#[wasm_bindgen_test]
fn stake_multiple_tokens() {
    let camper = Node::new("Camper");
    let data = vec![camper];
    let mut fix_node_state = fix(data);
    fix_node_state.transactions.push(Transaction {
        name: "Camper".to_string(),
        event: Events::Stake,
    });
    fix_node_state.transactions.push(Transaction {
        name: "Camper".to_string(),
        event: Events::Stake,
    });
    let (chain, _) = mine(fix_node_state).expect("result to be chain");
    assert_eq!(chain.get_node_by_name("Camper").unwrap().staked, 2);
}

#[wasm_bindgen_test]
fn one_invalid_transaction() {
    let camper = Node::new("Camper");
    let data = vec![camper];
    let mut fix_node_state = fix(data);
    fix_node_state.transactions.push(Transaction {
        name: "Camper".to_string(),
        event: Events::Stake,
    });
    fix_node_state.transactions.push(Transaction {
        name: "Tom".to_string(),
        event: Events::Stake,
    });
    let (chain, errors) = mine(fix_node_state).expect("result to be chain");
    assert_eq!(chain.get_node_by_name("Camper").unwrap().staked, 1);
    assert_eq!(errors[0], "Tom not found in chain");
}

#[wasm_bindgen_test]
fn submit_task_lots_staked() {
    let mut camper = Node::new("Camper");
    camper.staked = 15;
    let data = vec![camper];
    let mut fix_node_state = fix(data);
    fix_node_state.transactions.push(Transaction {
        name: "Camper".to_string(),
        event: Events::Stake,
    });
    fix_node_state.transactions.push(Transaction {
        name: "Camper".to_string(),
        event: Events::Stake,
    });
    // mine until reputation increases
    let mut c = 0;
    let a = loop {
        let (chain, _) = mine(fix_node_state.clone()).expect("result to be chain");
        if chain.get_node_by_name("Camper").unwrap().reputation == 2 {
            break true;
        }
        if c >= 25 {
            break false;
        }
        c += 1;
    };
    assert!(a, "probability-based test - may fail");
}

#[wasm_bindgen_test]
fn submit_task_few_staked() {
    let mut camper = Node::new("Camper");
    camper.staked = 1;
    let data = vec![camper];
    let mut fix_node_state = fix(data);
    fix_node_state.transactions.push(Transaction {
        name: "Camper".to_string(),
        event: Events::Stake,
    });
    fix_node_state.transactions.push(Transaction {
        name: "Camper".to_string(),
        event: Events::Stake,
    });
    // mine until reputation increases
    let mut c = 0;
    let mut l = 0;
    let a = loop {
        let (chain, _) = mine(fix_node_state.clone()).expect("result to be chain");
        if chain.get_node_by_name("Camper").unwrap().reputation == 1 {
            c += 1;
        }
        if c >= 10 {
            break true;
        }
        if l >= 25 {
            break false;
        }
        l += 1;
    };
    assert!(a, "probability-based test - may fail");
}

#[wasm_bindgen_test]
fn submit_task_rep_chance() {
    let camper = Node::new("Camper");
    let data = vec![camper];
    let mut fix_node_state = fix(data);
    fix_node_state.transactions.push(Transaction {
        name: "Camper".to_string(),
        event: Events::Stake,
    });
    fix_node_state.transactions.push(Transaction {
        name: "Camper".to_string(),
        event: Events::Stake,
    });
    let (chain, _) = mine(fix_node_state).expect("result to be chain");
    assert_eq!(chain.get_node_by_name("Camper").unwrap().tokens, 21);
}

fn fix(data: Vec<Node>) -> NodeState {
    let node_vec_str = serde_json::to_string(&data).unwrap();
    let fix_node_state = format!(
        r#"{{
      "chain": {{
        "chain": [
          {{
            "id": 0,
            "hash": "00110101",
            "previous_hash": "",
            "timestamp": 123456789,
            "data": {},
            "nonce": 123,
            "next_miner": "Camper",
            "next_validators": ["Camper"]
          }}
        ],
        "network": ["Camper"]
      }},
      "transactions": [
        {{
          "event": "SubmitTask",
          "name": "Camper"
        }}
      ],
      "task_valid": true
    }}"#,
        node_vec_str
    );
    serde_json::from_str(&fix_node_state).unwrap()
}

fn mine(fix_node_state: NodeState) -> Result<(Chain, Vec<String>), JsValue> {
    let node_state = JsValue::from_serde(&fix_node_state).unwrap();
    let res = handle_mine(node_state);
    let response = match res {
        Ok(v) => match v.into_serde() {
            Ok(v) => v,
            Err(e) => {
                console::log_1(&format!("{:?}", e).into());
                panic!("could not serde response");
            }
        },
        Err(e) => {
            // Error is converted into a JsValue to make use of Debug trait
            return Err(JsValue::from(e));
        }
    };
    Ok(response)
}


extern crate gym;

use gym::*;

fn main() {
	println!("**********************************");

	let mut client = GymClient::new("http://localhost:5000".to_string());
	//println!("already running environments:\n{:?}\n", client.get_envs().unwrap());

	let mut env = match client.make("CartPole-v0") {
		Ok(env) => env,
		Err(msg) => panic!("Could not make environment because of error:\n{}", msg)
	};

	//println!("observation space:\n{:?}\n", env.observation_space());
	//println!("action space:\n{:?}\n", env.action_space());

	let _ = env.monitor_start("/tmp/random-agent-results".to_string(), true, false);
	for ep in 0..10 {
		let mut tot_reward = 0.;
		let _ = env.reset();
		loop {
			let action = env.action_space().sample();
			let state = env.step(action, true).unwrap();
			assert_eq!(state.observation.len(), env.observation_space().sample().len());
			tot_reward += state.reward;

			if state.done {
				break;
			}
		}
		println!("Finished episode {} with total reward {}", ep, tot_reward);
	}
	let _ = env.monitor_stop();

	println!("**********************************");
}


extern crate serde_json;
extern crate hyper;
extern crate rand;

use std::collections::BTreeMap;
use std::io::Read;

use serde_json::Value;
use serde_json::value::{ToJson, from_value};

use hyper::client::Client;
use hyper::header::Headers;

use rand::{thread_rng, Rng};

pub type GymResult<T> = Result<T, hyper::Error>;

#[derive(Debug, Clone)]
pub enum Space {
	DISCRETE{n: u64},
	BOX{shape: Vec<u64>, high: Vec<f64>, low: Vec<f64>},
	TUPLE{spaces: Vec<Box<Space>>}
}

impl Space {
	fn from_json(info: &Value) -> Space {
		match info.find("name").unwrap().as_str().unwrap() {
			"Discrete" => {
				let n = info.find("n").unwrap().as_u64().unwrap();
				Space::DISCRETE{n: n}
			},
			"Box" => {
				let shape = info.find("shape").unwrap().as_array().unwrap()
								.into_iter().map(|x| x.as_u64().unwrap())
								.collect::<Vec<_>>();

				let high = info.find("high").unwrap().as_array().unwrap()
							   .into_iter().map(|x| x.as_f64().unwrap())
							   .collect::<Vec<_>>();

				let low = info.find("low").unwrap().as_array().unwrap()
							  .into_iter().map(|x| x.as_f64().unwrap())
							  .collect::<Vec<_>>();

				Space::BOX{shape: shape, high: high, low: low}
			},
			"Tuple" => panic!("Parsing for Tuple spaces is not yet implemented"),
			e @ _ => panic!("Unrecognized space name: {}", e)
		}
	}
	pub fn sample(&self) -> Vec<f64> {
		let mut rng = thread_rng();
		match *self {
			Space::DISCRETE{n} => {
				vec![(rng.gen::<u64>()%n) as f64]
			},
			Space::BOX{ref shape, ref high, ref low} => {
				let mut ret = Vec::with_capacity(shape.iter().map(|x| *x as usize).product());
				let mut index = 0;
				for &i in shape {
					for _ in 0..i {
						ret.push(rng.gen_range(low[index], high[index]));
						index += 1;
					}
				}
				ret
			},
			Space::TUPLE{ref spaces} => {
				let mut ret = Vec::new();
				for space in spaces {
					ret.extend(space.sample());
				}
				ret
			}
		}
	}
}

#[allow(dead_code)]
#[derive(Debug)]
pub struct State {
	pub observation:	Vec<f64>,
	pub reward:			f64,
	pub done:			bool,
	pub info:			Value,
}

#[allow(dead_code)]
pub struct Environment {
	client:			GymClient,
	instance_id:	String,
	act_space:		Space,
	obs_space:		Space,
}

impl Environment {
	pub fn action_space<'a>(&'a self) -> &'a Space {
		&self.act_space
	}
	pub fn observation_space<'a>(&'a self) -> &'a Space {
		&self.obs_space
	}
	pub fn reset(&mut self) -> GymResult<Vec<f64>> {
		let path = "/v1/envs/".to_string() + &self.instance_id + "/reset/";
		let observation = try!(self.client.post(path, Value::Null));

		let ret: Vec<_> = observation.find("observation").unwrap().as_array().unwrap()
									 .into_iter().map(|x| x.as_f64().unwrap())
									 .collect();
		Ok(ret)
	}
	pub fn step(&mut self, action: Vec<f64>, render: bool) -> GymResult<State> {
		let mut req = BTreeMap::new();
		req.insert("render", Value::Bool(render));
		match self.act_space {
			Space::DISCRETE{..} => {
				assert_eq!(action.len(), 1);
				req.insert("action", Value::U64(action[0] as u64));
			},
			Space::BOX{ref shape, ..} => {
				assert_eq!(action.len(), shape[0] as usize);
				req.insert("action", action.to_json());
			},
			Space::TUPLE{..} => panic!("Actions for Tuple spaces not implemented yet")
		}

		let path = "/v1/envs/".to_string() + &self.instance_id + "/step/";
		let state = try!(self.client.post(path, req.to_json()));

		Ok(State {
			observation: from_value(state.find("observation").unwrap().clone()).unwrap(),
			reward: state.find("reward").unwrap().as_f64().unwrap(),
			done: state.find("done").unwrap().as_bool().unwrap(),
			info: state.find("info").unwrap().clone()
		})
	}
	pub fn monitor_start(&mut self, directory: String, force: bool, resume: bool) -> GymResult<()> {
		let mut req = BTreeMap::new();
		req.insert("directory", Value::String(directory));
		req.insert("force", Value::Bool(force));
		req.insert("resume", Value::Bool(resume));

		let path = "/v1/envs/".to_string() + &self.instance_id + "/monitor/start/";
		try!(self.client.post(path, req.to_json()));
		Ok(())
	}
	pub fn monitor_stop(&mut self) -> GymResult<()> {
		let path = "/v1/envs/".to_string() + &self.instance_id + "/monitor/close/";
		try!(self.client.post(path, Value::Null));
		Ok(())
	}
	pub fn upload(&mut self, training_dir: String, api_key: String, algorithm_id: String) -> GymResult<()> {
		let mut req = BTreeMap::new();
		req.insert("training_dir", training_dir);
		req.insert("api_key", api_key);
		req.insert("algorithm_id", algorithm_id);

		try!(self.client.post("/v1/upload/".to_string(), req.to_json()));
		Ok(())
	}
}

pub struct GymClient {
	address:	String,
	handle:		Client,
	headers:	Headers,
}

impl GymClient {
    pub fn new(addr: String) -> GymClient {
		let mut headers = Headers::new();
		headers.set_raw("Content-Type", vec![b"application/json".to_vec()]);

    	GymClient {
    		address: addr,
    		handle: Client::new(),
    		headers: headers
    	}
    }
    pub fn make(mut self, env_id: &str) -> GymResult<Environment> {
    	let mut req: BTreeMap<&str, &str> = BTreeMap::new();
    	req.insert("env_id", env_id);

    	let instance_id = try!(self.post("/v1/envs/".to_string(), req.to_json()));
    	let instance_id = match instance_id.find("instance_id") {
    		Some(id) => id.as_str().unwrap(),
    		None => panic!("Unrecognized environment id: {}", env_id)
    	};

    	let obs_space = try!(self.get("/v1/envs/".to_string() + instance_id + "/observation_space/"));
    	let act_space = try!(self.get("/v1/envs/".to_string() + instance_id + "/action_space/"));

    	Ok(Environment {
    		client: self,
    		instance_id: instance_id.to_string(),
    		act_space: Space::from_json(act_space.find("info").unwrap()),
    		obs_space: Space::from_json(obs_space.find("info").unwrap())})
    }
    pub fn get_envs(&mut self) -> GymResult<BTreeMap<String, String>> {
    	let json = try!(self.get("/v1/envs/".to_string()));
    	Ok(from_value(json.find("all_envs").unwrap().clone()).unwrap())
    }

    fn post(&mut self, route: String, request: Value) -> GymResult<Value> {
    	let url = self.address.clone() + &route;
    	let mut resp = try!(self.handle.post(&url)
    							  	   .body(&request.to_string())
    							  	   .headers(self.headers.clone())
    							  	   .send());

    	let mut json = String::new();
    	let _ = resp.read_to_string(&mut json);

	    Ok(serde_json::from_str(&json).unwrap_or(Value::Null))
    }
    fn get(&mut self, route: String) -> GymResult<Value> {
    	let url = self.address.clone() + &route;
    	let mut resp = try!(self.handle.get(&url)
    							  	   .send());
    	let mut json = String::new();
    	let _ = resp.read_to_string(&mut json);

		Ok(serde_json::from_str(&json).unwrap_or(Value::Null))
    }
}


extern crate gym;

use gym::*;

#[test]
fn test_space_sample() {
	let discrete_space = Space::DISCRETE{n: 15};
	for _ in 0..10 {
		let sample = discrete_space.sample();
		assert!(sample.len() == 1 && 0. <= sample[0] && sample[0] < 15.);
	}

	let box_space = Space::BOX{
		shape: vec![5],
		high: vec![1., 2., 3., 4., 5.],
		low: vec![-1., -2., -3., -4., -5.]
	};
	for _ in 0..10 {
		let sample = box_space.sample();
		assert_eq!(sample.len(), 5);
		for i in 0..5 {
			let bound = (i+1) as f64;
			assert!(-bound <= sample[i] && sample[i] <= bound);
		}
	}

	let tuple_space = Space::TUPLE{spaces: vec![Box::new(discrete_space), Box::new(box_space)]};
	for _ in 0..10 {
		let sample = tuple_space.sample();
		assert_eq!(sample.len(), 6);
		assert!(0. <= sample[0] && sample[0] <= 15.);
		for i in 1..6 {
			let bound = i as f64;
			assert!(-bound <= sample[i] && sample[i] <= bound);
		}
	}
}


// This check is new and seems buggy (possibly with PyO3 interaction)
#![allow(clippy::borrow_deref_ref)]

use std::collections::HashSet;
use std::thread;

use fancy_regex::Regex;
use pyo3::exceptions;
use pyo3::prelude::*;
use pyo3::types::{PyBytes, PyList, PyTuple};
use pyo3::PyResult;
use rustc_hash::FxHashMap as HashMap;

fn _byte_pair_merge<T>(
    piece: &[u8],
    ranks: &HashMap<Vec<u8>, usize>,
    f: impl Fn(std::ops::Range<usize>) -> T,
) -> Vec<T> {
    // This is a vector of (start, rank).
    // The rank is of the byte pair starting at position start.
    // The rank of the last item in the vector is not a valid value.
    let mut parts: Vec<(usize, usize)> = (0..piece.len() + 1).map(|i| (i, usize::MAX)).collect();

    let get_rank = {
        #[inline(always)]
        |parts: &Vec<(usize, usize)>, start_idx: usize, skip: usize| {
            if (start_idx + skip + 2) < parts.len() {
                ranks
                    .get(&piece[parts[start_idx].0..parts[start_idx + skip + 2].0])
                    .copied()
            } else {
                None
            }
        }
    };

    // We look up the ranks once in the beginning and iteratively update
    // them during each merge, which reduces the number of rank lookups.
    for i in 0..parts.len() - 2 {
        match get_rank(&parts, i, 0) {
            Some(rank) => {
                // usize::MAX is a sentinel value and cannot be a valid rank
                debug_assert!(rank != usize::MAX);
                parts[i].1 = rank;
            }
            None => {
                continue;
            }
        };
    }

    // If you have n parts and m merges, this does O(mn) work.
    // We could do something with a heap and do O(m log n) work.
    // It is important to consider that n is often small (<100), and as such
    // the cache-locality benefits outweigh the algorithmic complexity downsides
    // of the `parts` vector data structure above.

    // Note that we hash bytes, not token pairs. As long as we train BPE the way we
    // currently do, this is equivalent. An easy way to break this would be to decouple
    // merge priority from token index or to prevent specific token merges.
    loop {
        if parts.len() == 1 {
            break;
        }

        // usize::MAX is a sentinel rank value allowing us to
        // take the min more quickly
        let mut min_rank: (usize, usize) = (usize::MAX, 0);
        for (i, &(_, rank)) in parts[..parts.len() - 1].iter().enumerate() {
            if rank < min_rank.0 {
                min_rank = (rank, i);
            }
        }

        if min_rank.0 != usize::MAX {
            let i = min_rank.1;

            // NOTE: We are about to remove parts[i + 1]. We do not do it
            // yet because there are cache-locality benefits to updating
            // parts[i] and parts[i-1] before removing, which could thrash
            // the cache. Thus, we update the rank calculation by skipping over
            // parts[i + 1], by invoking `get_rank!` with `skip = 1`.
            parts[i].1 = get_rank(&parts, i, 1).unwrap_or(usize::MAX);
            if i > 0 {
                parts[i - 1].1 = get_rank(&parts, i - 1, 1).unwrap_or(usize::MAX);
            }

            parts.remove(i + 1);
        } else {
            break;
        }
    }
    let mut out: Vec<T> = Vec::with_capacity(parts.len() - 1);
    for i in 0..parts.len() - 1 {
        out.push(f(parts[i].0..parts[i + 1].0));
    }
    out
}

pub fn byte_pair_encode(piece: &[u8], ranks: &HashMap<Vec<u8>, usize>) -> Vec<usize> {
    if piece.len() == 1 {
        return vec![ranks[piece]];
    }
    _byte_pair_merge(piece, ranks, |p| ranks[&piece[p.start..p.end]])
}

pub fn byte_pair_split<'a>(piece: &'a [u8], ranks: &HashMap<Vec<u8>, usize>) -> Vec<&'a [u8]> {
    if piece.len() == 1 {
        return vec![piece];
    }
    _byte_pair_merge(piece, ranks, |p| &piece[p.start..p.end])
}

// Various performance notes:
//
// Regex
// =====
// Most of the time is spent in regex. The easiest way to speed this up is by using less fancy
// regex features. For instance, using a regex parse-able by `regex` crate is 3x faster than
// the usual regex we use.
//
// However, given that we're using a regex parse-able by `regex`, there isn't much difference
// between using the `regex` crate and using the `fancy_regex` crate.
//
// There is an important interaction between threading, `regex` and `fancy_regex`.
// When using `fancy_regex`, we hit `regex.find_at`. It turns out that this causes contention on
// some mutable scratch space inside of `regex`. This absolutely kills performance. When using plain
// old `regex`, we don't hit this, because `find_iter` has a different code path.
// Related: https://github.com/rust-lang/regex/blob/master/PERFORMANCE.md
// Anyway, the way we get around this is with having a (mostly) thread local clone of the regex for
// each thread.
//
// Threading
// =========
// I tried using `rayon`. It wasn't really faster than using Python threads and releasing the GIL.
// So goodbye `rayon`! Let thread count etc be in control of our Python users.
//
// Caching
// =======
// The reference tokeniser has an lru cache over the equivalent of `byte_pair_encode`.
// Originally, we had one too! Without it, we were only vaguely faster than Python.
// I used an RWLock to protect the cache. This didn't seem to hurt single threaded performance
// noticeably, but it did affect multi-threaded performance. Weirdly, it seemed to affect
// multi-threaded performance even when I only had readers (maybed I messed something up?).
// Anyway, I realised that we could get rid of the cache, if we treat the set of tokens as a cache!
// These are exactly the set or merges that are likely to be hot. And now we don't have to think
// about interior mutability, memory use, or cloning.
//
// Hashing
// =======
// We use FxHashMap instead of the standard HashMap. This is maybe like a 5-10% win?
// The current implementation ends up doing a lot of hashing of bytes. In theory, this could be made
// to be hashing of two-tuples of ints, which looks like it may also be a couple percent faster.

use std::num::NonZeroU64;
pub struct FakeThreadId(NonZeroU64);

fn hash_current_thread() -> usize {
    // It's easier to use unsafe than to use nightly. Rust has this nice u64 thread id counter
    // that works great for our use case of avoiding collisions in our array. Unfortunately,
    // it's private. However, there are only so many ways you can layout a u64, so just transmute
    // https://github.com/rust-lang/rust/issues/67939
    const _: [u8; 8] = [0; std::mem::size_of::<std::thread::ThreadId>()];
    const _: [u8; 8] = [0; std::mem::size_of::<FakeThreadId>()];
    let x = unsafe {
        std::mem::transmute::<std::thread::ThreadId, FakeThreadId>(thread::current().id()).0
    };
    u64::from(x) as usize
}

const MAX_NUM_THREADS: usize = 128;
#[pyclass]
struct CoreBPE {
    encoder: HashMap<Vec<u8>, usize>,
    special_tokens_encoder: HashMap<String, usize>,
    decoder: HashMap<usize, Vec<u8>>,
    special_tokens_decoder: HashMap<usize, Vec<u8>>,
    regex_tls: Vec<Regex>,
    special_regex_tls: Vec<Regex>,
    sorted_token_bytes: Vec<Vec<u8>>,
}

impl CoreBPE {
    fn _get_tl_regex(&self) -> &Regex {
        // See performance notes above for what this is about
        // It's also a little janky, please make a better version of it!
        // However, it's nice that this doesn't leak memory to short-lived threads
        &self.regex_tls[hash_current_thread() % MAX_NUM_THREADS]
    }

    fn _get_tl_special_regex(&self) -> &Regex {
        &self.special_regex_tls[hash_current_thread() % MAX_NUM_THREADS]
    }

    fn _decode_native(&self, tokens: &[usize]) -> Vec<u8> {
        let mut ret = Vec::with_capacity(tokens.len() * 2);
        for token in tokens {
            let token_bytes = self
                .decoder
                .get(token)
                .unwrap_or_else(|| &self.special_tokens_decoder[token]);
            ret.extend(token_bytes);
        }
        ret
    }

    fn _encode_ordinary_native(&self, text: &str) -> Vec<usize> {
        // This is the core of the encoding logic; the other functions in here
        // just make things complicated :-)
        let regex = self._get_tl_regex();
        let mut ret = vec![];
        for mat in regex.find_iter(text) {
            let piece = mat.unwrap().as_str().as_bytes();
            if let Some(token) = self.encoder.get(piece) {
                ret.push(*token);
                continue;
            }
            ret.extend(&byte_pair_encode(piece, &self.encoder));
        }
        ret
    }

    fn _encode_native(&self, text: &str, allowed_special: &HashSet<&str>) -> (Vec<usize>, usize) {
        let special_regex = self._get_tl_special_regex();
        let regex = self._get_tl_regex();
        let mut ret = vec![];

        let mut start = 0;
        let mut last_piece_token_len = 0;
        loop {
            let mut next_special;
            let mut start_find = start;
            loop {
                // Find the next allowed special token, if any
                next_special = special_regex.find_from_pos(text, start_find).unwrap();
                match next_special {
                    Some(m) => {
                        if allowed_special.contains(&text[m.start()..m.end()]) {
                            break;
                        }
                        start_find = m.start() + 1;
                    }
                    None => break,
                }
            }
            let end = next_special.map_or(text.len(), |m| m.start());

            // Okay, here we go, compare this logic to _encode_ordinary_native
            for mat in regex.find_iter(&text[start..end]) {
                let piece = mat.unwrap().as_str().as_bytes();
                if let Some(token) = self.encoder.get(piece) {
                    last_piece_token_len = 1;
                    ret.push(*token);
                    continue;
                }
                let tokens = byte_pair_encode(piece, &self.encoder);
                last_piece_token_len = tokens.len();
                ret.extend(&tokens);
            }

            match next_special {
                // And here we push the special token
                Some(m) => {
                    let piece = m.as_str();
                    let token = self.special_tokens_encoder[piece];
                    ret.push(token);
                    start = m.end();
                    last_piece_token_len = 0;
                }
                None => break,
            }
        }

        // last_piece_token_len is how many tokens came from the last regex split. This is used
        // for determining unstable tokens, since you can't merge across (stable) regex splits
        (ret, last_piece_token_len)
    }

    fn _increase_last_piece_token_len(
        &self,
        tokens: Vec<usize>,
        mut last_piece_token_len: usize,
    ) -> (Vec<usize>, usize) {
        // Unfortunately, the locations where our regex splits can be unstable.
        // For the purposes of determining unstable tokens, unstable regex splitting
        // is only a problem if a split that was present disappears, since this can
        // lead to merging of tokens otherwise thought to be stable.
        // cl100k_base makes our life hard by including the \s*[\r\n]+
        // pattern. This can e.g. cause "\n" + " " to become "\n \n".
        // Here is a quick and dirty fix:
        {
            let token_is_all_space = |token| {
                self.decoder
                    .get(token)
                    .map(|token_bytes| {
                        token_bytes
                            .iter()
                            .rev()
                            .all(|&b| [b' ', b'\n', b'\t'].contains(&b))
                    })
                    .unwrap_or(false)
            };
            if last_piece_token_len > 0
                && token_is_all_space(&tokens[tokens.len() - last_piece_token_len])
            {
                while (last_piece_token_len < tokens.len())
                    && token_is_all_space(&tokens[tokens.len() - last_piece_token_len - 1])
                {
                    last_piece_token_len += 1;
                }
            }
        }
        debug_assert!(last_piece_token_len <= tokens.len());

        (tokens, last_piece_token_len)
    }

    fn _encode_unstable_native(
        &self,
        text: &str,
        allowed_special: &HashSet<&str>,
    ) -> (Vec<usize>, HashSet<Vec<usize>>) {
        let (tokens, last_piece_token_len) = self._encode_native(text, allowed_special);
        if last_piece_token_len == 0 {
            // If last_piece_token_len is zero, the last token was a special token and we have
            // no unstable bytes
            return (tokens, HashSet::new());
        }
        let (mut tokens, last_piece_token_len) =
            self._increase_last_piece_token_len(tokens, last_piece_token_len);

        let unstable_bytes = self._decode_native(&tokens[tokens.len() - last_piece_token_len..]);
        tokens.truncate(tokens.len() - last_piece_token_len);

        // TODO: we should try harder to find additional stable tokens
        // This would reduce the amount of retokenising when determining completions
        // Refer to the logic in an older version of this file

        let mut completions = HashSet::new();
        if unstable_bytes.is_empty() {
            return (tokens, completions);
        }

        // This is the easy bit. Just find all single tokens that start with unstable_bytes
        // (including tokens that exactly match unstable_bytes)
        // Separating this from the loop below helps with performance in a common case.
        let mut point = self
            .sorted_token_bytes
            .partition_point(|x| x.as_slice() < unstable_bytes.as_slice());
        while point < self.sorted_token_bytes.len()
            && self.sorted_token_bytes[point].starts_with(&unstable_bytes)
        {
            completions.insert(vec![
                self.encoder[self.sorted_token_bytes[point].as_slice()],
            ]);
            point += 1;
        }

        // Now apply even more brute force. At every (other) possible position for the straddling
        // token, concatenate additional bytes from that token (if any) to unstable_bytes,
        // and retokenise the whole thing and see what we get.
        for i in 1..unstable_bytes.len() {
            let prefix = &unstable_bytes[..i];
            let suffix = &unstable_bytes[i..];
            let mut point = self
                .sorted_token_bytes
                .partition_point(|x| x.as_slice() < suffix);
            // TODO: Perf optimisation if suffix starts with " "?
            while point < self.sorted_token_bytes.len()
                && self.sorted_token_bytes[point].starts_with(suffix)
            {
                let possibility = [prefix, self.sorted_token_bytes[point].as_slice()].concat();
                let encoded = match std::str::from_utf8(&possibility) {
                    // Morally, this is byte_pair_encode(&possibility, &self.encoder)
                    // But we might have introduced a regex split which would prevent merges.
                    // (particularly possible in the presence of unstable regex splits)
                    // So convert to UTF-8 and do regex splitting.
                    // E.g. with cl100k_base "  !" gets split to " " + " !",
                    // but byte_pair_encode("  !") != byte_pair_encode(" ")
                    Ok(s) => self._encode_ordinary_native(s),

                    // Technically, whether or not this arm is correct depends on whether there
                    // would be a regex split before the UTF-8 truncation point.
                    // Probably niche enough that no one will ever notice (after all, people didn't
                    // notice all the big holes in the previous unstable token implementation)
                    Err(_) => byte_pair_encode(&possibility, &self.encoder),
                    // Something like the following is intriguing but incorrect:
                    // Err(e) => self._encode_ordinary_native(unsafe {
                    //     std::str::from_utf8_unchecked(&possibility[..e.valid_up_to()])
                    // }),
                };
                let mut seq = Vec::new();
                let mut seq_len = 0;
                for token in encoded {
                    seq.push(token);
                    seq_len += self.decoder[&token].len();
                    if seq_len >= unstable_bytes.len() {
                        break;
                    }
                }
                completions.insert(seq);
                point += 1;
            }
        }

        // This is also not straightforward. While we generally assume that regex splits are stable,
        // unfortunately, they are not. That is, if adding bytes were to make a split appear in
        // unstable_bytes, this could make tokens possible which our logic would otherwise think
        // would be merged.
        // For example, with gpt2, the use of \s+(?!\S) means that "\n\n" could
        // develop a split, e.g. "\n\n0" splits into "\n"+"\n"+"0", making "\n" a possible token.
        // Here is a quick and dirty fix:
        // This isn't right if we ever remove \s+(?!\S)
        if unstable_bytes.len() > 1 {
            let last_decoded = bstr::decode_last_utf8(unstable_bytes.as_slice());
            if unstable_bytes.len() - last_decoded.1 > 0
                && last_decoded.0.map_or(false, |c| c.is_whitespace())
            {
                let mut reencoded = byte_pair_encode(
                    &unstable_bytes[..unstable_bytes.len() - last_decoded.1],
                    &self.encoder,
                );
                reencoded.extend(byte_pair_encode(
                    &unstable_bytes[unstable_bytes.len() - last_decoded.1..],
                    &self.encoder,
                ));
                completions.insert(reencoded);
            }
        }

        (tokens, completions)
    }
}

#[pymethods]
impl CoreBPE {
    #[new]
    fn new(
        encoder: HashMap<Vec<u8>, usize>,
        special_tokens_encoder: HashMap<String, usize>,
        pattern: &str,
    ) -> PyResult<Self> {
        let regex = Regex::new(pattern)
            .map_err(|e| PyErr::new::<exceptions::PyValueError, _>(e.to_string()))?;

        let special_regex = {
            let _parts = special_tokens_encoder
                .keys()
                .map(|s| fancy_regex::escape(s))
                .collect::<Vec<_>>();
            Regex::new(&_parts.join("|"))
                .map_err(|e| PyErr::new::<exceptions::PyValueError, _>(e.to_string()))?
        };

        let decoder: HashMap<usize, Vec<u8>> =
            encoder.iter().map(|(k, v)| (*v, k.clone())).collect();

        assert!(
            encoder.len() == decoder.len(),
            "Encoder and decoder must be of equal length; maybe you had duplicate token indices in your encoder?"
        );

        let special_tokens_decoder: HashMap<usize, Vec<u8>> = special_tokens_encoder
            .iter()
            .map(|(k, v)| (*v, k.as_bytes().to_vec()))
            .collect();

        // Clone because I don't know how to tell Rust I'm not going to change the map
        let mut sorted_token_bytes: Vec<Vec<u8>> = encoder.keys().cloned().collect();
        sorted_token_bytes.sort();

        Ok(CoreBPE {
            encoder,
            special_tokens_encoder,
            decoder,
            special_tokens_decoder,
            regex_tls: (0..MAX_NUM_THREADS).map(|_| regex.clone()).collect(),
            special_regex_tls: (0..MAX_NUM_THREADS)
                .map(|_| special_regex.clone())
                .collect(),
            sorted_token_bytes,
        })
    }

    // ====================
    // Encoding
    // ====================

    fn encode_ordinary(&self, py: Python, text: &str) -> Vec<usize> {
        py.allow_threads(|| self._encode_ordinary_native(text))
    }

    fn encode(&self, py: Python, text: &str, allowed_special: HashSet<&str>) -> Vec<usize> {
        py.allow_threads(|| self._encode_native(text, &allowed_special).0)
    }

    fn _encode_bytes(&self, py: Python, bytes: &[u8]) -> Vec<usize> {
        py.allow_threads(|| {
            match std::str::from_utf8(bytes) {
                Ok(text) => self._encode_ordinary_native(text),
                Err(e) => {
                    let text = unsafe { std::str::from_utf8_unchecked(&bytes[..e.valid_up_to()]) };
                    let (tokens, last_piece_token_len) = self._encode_native(text, &HashSet::new());
                    let (mut tokens, last_piece_token_len) =
                        self._increase_last_piece_token_len(tokens, last_piece_token_len);
                    if !tokens.is_empty() && last_piece_token_len > 0 {
                        // Lop off the tokens from the last piece and run BPE on the remaining bytes
                        // Somewhat niche, but this may not be correct if we'd have had a regex
                        // split between the valid UTF-8 and the invalid bytes, which is why this
                        // method is private
                        let mut unstable_bytes =
                            self._decode_native(&tokens[tokens.len() - last_piece_token_len..]);
                        unstable_bytes.extend_from_slice(&bytes[e.valid_up_to()..]);

                        tokens.truncate(tokens.len() - last_piece_token_len);
                        tokens.extend(byte_pair_encode(&unstable_bytes, &self.encoder));
                    }
                    tokens
                }
            }
        })
    }

    fn encode_with_unstable(
        &self,
        py: Python,
        text: &str,
        allowed_special: HashSet<&str>,
    ) -> Py<PyTuple> {
        let (tokens, completions) =
            py.allow_threads(|| self._encode_unstable_native(text, &allowed_special));
        let py_completions =
            PyList::new(py, completions.iter().map(|seq| PyList::new(py, &seq[..])));
        (tokens, py_completions).into_py(py)
    }

    fn encode_single_token(&self, piece: &[u8]) -> PyResult<usize> {
        if let Some(token) = self.encoder.get(piece).copied() {
            return Ok(token);
        }
        if let Ok(piece_str) = std::str::from_utf8(piece) {
            if let Some(token) = self.special_tokens_encoder.get(piece_str).copied() {
                return Ok(token);
            }
        }
        Err(PyErr::new::<exceptions::PyKeyError, _>(piece.to_owned()))
    }

    fn encode_single_piece(&self, piece: &[u8]) -> Vec<usize> {
        if let Some(token) = self.encoder.get(piece) {
            return vec![*token];
        }
        byte_pair_encode(piece, &self.encoder)
    }

    // ====================
    // Decoding
    // ====================

    fn decode_bytes(&self, py: Python, tokens: Vec<usize>) -> Py<PyBytes> {
        let bytes = py.allow_threads(|| self._decode_native(&tokens));
        PyBytes::new(py, &bytes).into()
    }

    fn decode_single_token_bytes(&self, py: Python, token: usize) -> PyResult<Py<PyBytes>> {
        if let Some(bytes) = self.decoder.get(&token) {
            return Ok(PyBytes::new(py, bytes).into());
        }
        if let Some(bytes) = self.special_tokens_decoder.get(&token) {
            return Ok(PyBytes::new(py, bytes).into());
        }
        Err(PyErr::new::<exceptions::PyKeyError, _>(token.to_string()))
    }

    // ====================
    // Miscellaneous
    // ====================

    fn token_byte_values(&self, py: Python) -> Vec<Py<PyBytes>> {
        self.sorted_token_bytes
            .iter()
            .map(|x| PyBytes::new(py, x).into())
            .collect()
    }
}

#[pymodule]
fn _tiktoken(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add_class::<CoreBPE>()?;
    Ok(())
}

#[cfg(test)]
mod tests {
    use rustc_hash::FxHashMap as HashMap;

    use crate::byte_pair_split;

    #[test]
    fn very_simple_test() {
        let mut ranks = HashMap::default();
        ranks.insert(b"ab".to_vec(), 1);
        ranks.insert(b"cd".to_vec(), 2);

        let res = byte_pair_split(b"abcd", &ranks);
        assert_eq!(res, vec![b"ab", b"cd"]);
    }
}


#[macro_use]
extern crate lazy_static;
pub mod backtracking;
pub mod big_integer;
pub mod ciphers;
pub mod compression;
pub mod conversions;
pub mod data_structures;
pub mod dynamic_programming;
pub mod general;
pub mod graph;
pub mod math;
pub mod navigation;
pub mod searching;
pub mod sorting;
pub mod string;

#[cfg(test)]
mod tests {
    use super::sorting;
    #[test]
    fn quick_sort() {
        //descending
        let mut ve1 = vec![6, 5, 4, 3, 2, 1];
        sorting::quick_sort(&mut ve1);

        assert!(sorting::is_sorted(&ve1));

        //pre-sorted
        let mut ve2 = vec![1, 2, 3, 4, 5, 6];
        sorting::quick_sort(&mut ve2);

        assert!(sorting::is_sorted(&ve2));
    }
}


/*
    In this problem, we want to determine all possible combinations of k
    numbers out of 1 ... n. We use backtracking to solve this problem.
    Time complexity: O(C(n,k)) which is O(n choose k) = O((n!/(k! * (n - k)!)))

    generate_all_combinations(n=4, k=2) => [[1, 2], [1, 3], [1, 4], [2, 3], [2, 4], [3, 4]]
*/
pub fn generate_all_combinations(n: i32, k: i32) -> Vec<Vec<i32>> {
    let mut result = vec![];
    create_all_state(1, n, k, &mut vec![], &mut result);

    result
}

fn create_all_state(
    increment: i32,
    total_number: i32,
    level: i32,
    current_list: &mut Vec<i32>,
    total_list: &mut Vec<Vec<i32>>,
) {
    if level == 0 {
        total_list.push(current_list.clone());
        return;
    }

    for i in increment..(total_number - level + 2) {
        current_list.push(i);
        create_all_state(i + 1, total_number, level - 1, current_list, total_list);
        current_list.pop();
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_output() {
        let expected_res = vec![
            vec![1, 2],
            vec![1, 3],
            vec![1, 4],
            vec![2, 3],
            vec![2, 4],
            vec![3, 4],
        ];

        let res = generate_all_combinations(4, 2);

        assert_eq!(expected_res, res);
    }

    #[test]
    fn test_empty() {
        let expected_res: Vec<Vec<i32>> = vec![vec![]];

        let res = generate_all_combinations(0, 0);

        assert_eq!(expected_res, res);
    }
}


mod all_combination_of_size_k;
mod sudoku;

pub use all_combination_of_size_k::generate_all_combinations;
pub use sudoku::Sudoku;


/*
    A Rust implementation of Sudoku solver using Backtracking.
    GeeksForGeeks: https://www.geeksforgeeks.org/sudoku-backtracking-7/
*/

pub struct Sudoku {
    board: [[u8; 9]; 9],
}

impl Sudoku {
    pub fn new(board: [[u8; 9]; 9]) -> Sudoku {
        Sudoku { board }
    }

    fn find_empty_cell(&self) -> Option<(usize, usize)> {
        // Find a empty cell in the board (returns (-1, -1) if all cells are filled)
        for i in 0..9 {
            for j in 0..9 {
                if self.board[i][j] == 0 {
                    return Some((i, j));
                }
            }
        }

        None
    }

    fn check(&self, index_tuple: (usize, usize), value: u8) -> bool {
        let (y, x) = index_tuple;

        // checks if the value to be added in the board is an acceptable value for the cell

        // checking through the row
        for i in 0..9 {
            if self.board[i][x] == value {
                return false;
            }
        }
        // checking through the column
        for i in 0..9 {
            if self.board[y][i] == value {
                return false;
            }
        }

        // checking through the 3x3 block of the cell
        let sec_row = y / 3;
        let sec_col = x / 3;

        for i in (sec_row * 3)..(sec_row * 3 + 3) {
            for j in (sec_col * 3)..(sec_col * 3 + 3) {
                if self.board[i][j] == value {
                    return false;
                }
            }
        }

        true
    }

    pub fn solve(&mut self) -> bool {
        let empty_cell = self.find_empty_cell();

        if let Some((y, x)) = empty_cell {
            for val in 1..10 {
                if self.check((y, x), val) {
                    self.board[y][x] = val;
                    if self.solve() {
                        return true;
                    }
                    // backtracking if the board cannot be solved using current configuration
                    self.board[y][x] = 0
                }
            }
        } else {
            // if the board is complete
            return true;
        }

        // returning false the board cannot be solved using current configuration
        false
    }

    pub fn print_board(&self) {
        // helper function to display board

        let print_3_by_1 = |arr: Vec<u8>, last: bool| {
            let str = arr
                .iter()
                .map(|n| n.to_string())
                .collect::<Vec<String>>()
                .join(", ");

            if last {
                println!("{str}",);
            } else {
                print!("{str} | ",);
            }
        };

        for i in 0..9 {
            if i % 3 == 0 && i != 0 {
                println!("- - - - - - - - - - - - - -")
            }

            print_3_by_1(self.board[i][0..3].to_vec(), false);
            print_3_by_1(self.board[i][3..6].to_vec(), false);
            print_3_by_1(self.board[i][6..9].to_vec(), true);
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_sudoku_correct() {
        let board: [[u8; 9]; 9] = [
            [3, 0, 6, 5, 0, 8, 4, 0, 0],
            [5, 2, 0, 0, 0, 0, 0, 0, 0],
            [0, 8, 7, 0, 0, 0, 0, 3, 1],
            [0, 0, 3, 0, 1, 0, 0, 8, 0],
            [9, 0, 0, 8, 6, 3, 0, 0, 5],
            [0, 5, 0, 0, 9, 0, 6, 0, 0],
            [1, 3, 0, 0, 0, 0, 2, 5, 0],
            [0, 0, 0, 0, 0, 0, 0, 7, 4],
            [0, 0, 5, 2, 0, 6, 3, 0, 0],
        ];

        let board_result = [
            [3, 1, 6, 5, 7, 8, 4, 9, 2],
            [5, 2, 9, 1, 3, 4, 7, 6, 8],
            [4, 8, 7, 6, 2, 9, 5, 3, 1],
            [2, 6, 3, 4, 1, 5, 9, 8, 7],
            [9, 7, 4, 8, 6, 3, 1, 2, 5],
            [8, 5, 1, 7, 9, 2, 6, 4, 3],
            [1, 3, 8, 9, 4, 7, 2, 5, 6],
            [6, 9, 2, 3, 5, 1, 8, 7, 4],
            [7, 4, 5, 2, 8, 6, 3, 1, 9],
        ];

        let mut sudoku = Sudoku::new(board);
        let is_solved = sudoku.solve();

        assert!(is_solved);
        assert_eq!(sudoku.board, board_result);
    }

    #[test]
    fn test_sudoku_incorrect() {
        let board: [[u8; 9]; 9] = [
            [6, 0, 3, 5, 0, 8, 4, 0, 0],
            [5, 2, 0, 0, 0, 0, 0, 0, 0],
            [0, 8, 7, 0, 0, 0, 0, 3, 1],
            [0, 0, 3, 0, 1, 0, 0, 8, 0],
            [9, 0, 0, 8, 6, 3, 0, 0, 5],
            [0, 5, 0, 0, 9, 0, 6, 0, 0],
            [1, 3, 0, 0, 0, 0, 2, 5, 0],
            [0, 0, 0, 0, 0, 0, 0, 7, 4],
            [0, 0, 5, 2, 0, 6, 3, 0, 0],
        ];

        let mut sudoku = Sudoku::new(board);
        let is_solved = sudoku.solve();

        assert!(!is_solved);
    }
}


// Algorithm created by Peter Borwein in 1985
// https://doi.org/10.1016/0196-6774(85)90006-9

use crate::math::sieve_of_eratosthenes;
use num_bigint::BigUint;
use num_traits::One;
use std::collections::BTreeMap;

/// Calculate the sum of n / p^i with integer division for all values of i
fn index(p: usize, n: usize) -> usize {
    let mut index = 0;
    let mut i = 1;
    let mut quot = n / p;

    while quot > 0 {
        index += quot;
        i += 1;
        quot = n / p.pow(i);
    }

    index
}

/// Calculate the factorial with time complexity O(log(log(n)) * M(n * log(n))) where M(n) is the time complexity of multiplying two n-digit numbers together.
pub fn fast_factorial(n: usize) -> BigUint {
    if n < 2 {
        return BigUint::one();
    }

    // get list of primes that will be factors of n!
    let primes = sieve_of_eratosthenes(n);

    let mut p_indeces = BTreeMap::new();

    // Map the primes with their index
    primes.into_iter().for_each(|p| {
        p_indeces.insert(p, index(p, n));
    });

    let max_bits = p_indeces.get(&2).unwrap().next_power_of_two().ilog2();

    // Create a Vec of 1's
    let mut a = Vec::with_capacity(max_bits as usize);
    a.resize(max_bits as usize, BigUint::one());

    // For every prime p, multiply a[i] by p if the ith bit of p's index is 1
    for (p, i) in p_indeces.into_iter() {
        let mut bit = 1usize;
        while bit.ilog2() < max_bits {
            if (bit & i) > 0 {
                a[bit.ilog2() as usize] *= p;
            }

            bit <<= 1;
        }
    }

    a.into_iter()
        .enumerate()
        .map(|(i, a_i)| a_i.pow(2u32.pow(i as u32))) // raise every a[i] to the 2^ith power
        .product() // we get our answer by multiplying the result
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::big_integer::hello_bigmath::factorial;

    #[test]
    fn fact() {
        assert_eq!(fast_factorial(30), factorial(30));
        assert_eq!(fast_factorial(52), factorial(52));
        assert_eq!(fast_factorial(100), factorial(100));
        assert_eq!(fast_factorial(1000), factorial(1000));
        assert_eq!(fast_factorial(5000), factorial(5000));
    }
}


use num_bigint::BigUint;
use num_traits::One;

// A trivial example of using Big integer math

pub fn factorial(num: u32) -> BigUint {
    let mut result: BigUint = One::one();
    for i in 1..=num {
        result *= i;
    }
    result
}

#[cfg(test)]
mod tests {
    use std::str::FromStr;

    use super::*;

    #[test]
    fn basic_factorial() {
        assert_eq!(factorial(10), BigUint::from_str("3628800").unwrap());
        assert_eq!(
            factorial(50),
            BigUint::from_str("30414093201713378043612608166064768844377641568960512000000000000")
                .unwrap()
        );
    }
}


#![cfg(feature = "big-math")]

mod fast_factorial;
mod hello_bigmath;
mod poly1305;

pub use self::fast_factorial::fast_factorial;
pub use self::hello_bigmath::factorial;
pub use self::poly1305::Poly1305;


use num_bigint::BigUint;
use num_traits::Num;
use num_traits::Zero;

macro_rules! hex_uint {
    ($a:literal) => {
        BigUint::from_str_radix($a, 16).unwrap()
    };
}

/**
 * Poly1305 Message Authentication Code:
 * This implementation is based on RFC8439.
 * Note that the Big Integer library we are using may not be suitable for
 * cryptographic applications due to non constant time operations.
*/
pub struct Poly1305 {
    p: BigUint,
    r: BigUint,
    s: BigUint,
    /// The accumulator
    pub acc: BigUint,
}

impl Default for Poly1305 {
    fn default() -> Self {
        Self::new()
    }
}

impl Poly1305 {
    pub fn new() -> Self {
        Poly1305 {
            p: hex_uint!("3fffffffffffffffffffffffffffffffb"), // 2^130 - 5
            r: Zero::zero(),
            s: Zero::zero(),
            acc: Zero::zero(),
        }
    }
    pub fn clamp_r(&mut self) {
        self.r &= hex_uint!("0ffffffc0ffffffc0ffffffc0fffffff");
    }
    pub fn set_key(&mut self, key: &[u8; 32]) {
        self.r = BigUint::from_bytes_le(&key[..16]);
        self.s = BigUint::from_bytes_le(&key[16..]);
        self.clamp_r();
    }
    /// process a 16-byte-long message block. If message is not long enough,
    /// fill the `msg` array with zeros, but set `msg_bytes` to the original
    /// chunk length in bytes. See `basic_tv1` for example usage.
    pub fn add_msg(&mut self, msg: &[u8; 16], msg_bytes: u64) {
        let mut n = BigUint::from_bytes_le(msg);
        n.set_bit(msg_bytes * 8, true);
        self.acc += n;
        self.acc *= &self.r;
        self.acc %= &self.p;
    }
    /// The result is guaranteed to be 16 bytes long
    pub fn get_tag(&self) -> Vec<u8> {
        let result = &self.acc + &self.s;
        let mut bytes = result.to_bytes_le();
        bytes.resize(16, 0);
        bytes
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::fmt::Write;
    fn get_tag_hex(tag: &[u8]) -> String {
        let mut result = String::new();
        for &x in tag {
            write!(result, "{x:02x}").unwrap();
        }
        result
    }
    #[test]
    fn basic_tv1() {
        let mut mac = Poly1305::new();
        let key: [u8; 32] = [
            0x85, 0xd6, 0xbe, 0x78, 0x57, 0x55, 0x6d, 0x33, 0x7f, 0x44, 0x52, 0xfe, 0x42, 0xd5,
            0x06, 0xa8, 0x01, 0x03, 0x80, 0x8a, 0xfb, 0x0d, 0xb2, 0xfd, 0x4a, 0xbf, 0xf6, 0xaf,
            0x41, 0x49, 0xf5, 0x1b,
        ];
        let mut tmp_buffer = [0_u8; 16];
        mac.set_key(&key);
        mac.add_msg(b"Cryptographic Fo", 16);
        mac.add_msg(b"rum Research Gro", 16);
        tmp_buffer[..2].copy_from_slice(b"up");
        mac.add_msg(&tmp_buffer, 2);
        let result = mac.get_tag();
        assert_eq!(
            get_tag_hex(result.as_slice()),
            "a8061dc1305136c6c22b8baf0c0127a9"
        );
    }
}


pub fn another_rot13(text: &str) -> String {
    let input = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz";
    let output = "NOPQRSTUVWXYZABCDEFGHIJKLMnopqrstuvwxyzabcdefghijklm";
    text.chars()
        .map(|c| match input.find(c) {
            Some(i) => output.chars().nth(i).unwrap(),
            None => c,
        })
        .collect()
}

#[cfg(test)]
mod tests {
    // Note this useful idiom: importing names from outer (for mod tests) scope.
    use super::*;

    #[test]
    fn test_simple() {
        assert_eq!(another_rot13("ABCzyx"), "NOPmlk");
    }

    #[test]
    fn test_every_alphabet_with_space() {
        assert_eq!(
            another_rot13("The quick brown fox jumps over the lazy dog"),
            "Gur dhvpx oebja sbk whzcf bire gur ynml qbt"
        );
    }

    #[test]
    fn test_non_alphabet() {
        assert_eq!(another_rot13("ðŸŽƒ Jack-o'-lantern"), "ðŸŽƒ Wnpx-b'-ynagrea");
    }
}