//! Boolean network implementation.

extern crate rand;

use rand::{
    prelude::*,
    distributions::Bernoulli
};

use std::{
    fs::File,
    io::{
        BufReader,
        BufWriter,
        Read,
        Write
    }
};

use crate::{
    bool3funcs::U128_BOOL_3_FUNCS,
    rwbytes::{
        ReadBytes,
        ReadBytesVec,
        WriteBytes
    }
};

type TCount = u32;

const IN_BITS: usize = 3;
type TFunc = u8; // "u(1<<(IN_BITS))"

#[derive(Clone)]
struct Boolon {
    inputs: Vec<TCount>, // indices of Boolons whose values are inputs to this one's input tree
    funcs: Vec<Vec<TFunc>> // i-th bit is value of boolean function when inputs' values as bin-digits make i (equivalently, it is an index for U128_BOOL_3_FUNC array) - for each level of input tree, from highest to lowest
}

impl Boolon {
    pub fn new_rand(rng: &mut ThreadRng, height: usize, size: usize, rands: usize) -> Self {
        let mut funcs = Vec::<Vec<TFunc>>::with_capacity(height);

        let mut leaves: usize = 1;
        for _ in 0..height {
            funcs.push(
                (0..leaves).into_iter().map(|_|
                    rng.gen::<TFunc>()
                ).collect()
            );
            leaves *= IN_BITS;
        }

        let mut inputs = Vec::<TCount>::with_capacity(leaves);
        inputs.extend((0..leaves).into_iter().map(|_|
            rng.gen_range(0, size + rands) as TCount
        ));

        Self{inputs, funcs}
    }

    pub fn new_inherit(rng: &mut ThreadRng, parent: &Self, size: usize, rands: usize, mutat_prob: f64) -> Self {
        match rng.gen_bool(mutat_prob) {
            false => parent.clone(), // Inheriting a boolon entirely seems to make "learning" faster...
            true => {
                let height = parent.funcs.len();
                Self::new_rand(rng, height, size, rands)
            }
        }
    }

    pub fn read<T: Read>(reader: &mut T, height: usize) -> Result<Self, String> {
        let leaves = IN_BITS.pow(height as u32);

        let inputs = TCount::read_bytes_vec(reader, leaves)?;

        let mut funcs = Vec::<Vec<TFunc>>::with_capacity(height);
        let mut leaves: usize = 1;
        for _ in 0..height {
            funcs.push(TFunc::read_bytes_vec(reader, leaves)?);
            leaves *= IN_BITS;
        }

        Ok(Self{inputs, funcs})
    }

    pub fn write<T: Write>(&self, writer: &mut T) -> Result<(), String> {
        self.inputs.write_bytes(writer)?;
        self.funcs.write_bytes(writer)?;
        Ok(())
    }

}

pub struct Boolnet {
    cycles: usize, // at the end of cycle, outputs of last layer become inputs of first layer
    rands: usize, // boolons with random output in each layer
    boolons: Vec<Vec<Boolon>>,

    age: u64,

    score: f64
}

impl Boolnet {
    pub fn new_rand(height: usize, size: usize, rands: usize, layers: usize, cycles: usize) -> Self {
        let mut rng = thread_rng();

        let mut boolons = Vec::<Vec<Boolon>>::with_capacity(layers);
        for _ in 0..layers {
            let mut lr = Vec::<Boolon>::with_capacity(size);
            for _ in 0..size {
                lr.push(Boolon::new_rand(&mut rng, height, size, rands));
            }
            boolons.push(lr);
        }

        let age: u64 = 0;

        let score = std::f64::NEG_INFINITY;

        Self{cycles, rands, boolons, age, score}
    }

    pub fn new_inherit(mut rng: &mut ThreadRng, parents: Vec<&Self>, par_sw_prob: f64, mutat_prob: f64) -> Self {
        let p0 = parents[0];

        let (cycles, layers, rands, size) = (p0.cycles, p0.boolons.len(), p0.rands, p0.boolons[0].len());

        let mut boolons = Vec::<Vec<Boolon>>::with_capacity(layers);
        for i in 0..layers {
            let mut lr = Vec::<Boolon>::with_capacity(size);
            let mut p = parents[rng.gen_range(0, parents.len())];
            let switch: Vec<bool> = Bernoulli::new(par_sw_prob).unwrap()
                .sample_iter(&mut rng)
                .take(size)
                .collect();
            for j in 0..size {
                if switch[j] { // switch parent whose boolons are inherited
                    p = parents[rng.gen_range(0, parents.len())];
                }
                lr.push(Boolon::new_inherit(rng, & p.boolons[i][j], size, rands, mutat_prob));
            }
            boolons.push(lr);
        }

        let age: u64 = 0;

        let score = std::f64::NEG_INFINITY;

        Self{cycles, rands, boolons, age, score}
    }

    pub fn load(filepath: &str) -> Result<Self, String> {
        // "Binary deserialization"... TODO: compactify
        let mut reader = BufReader::new(File::open(filepath).map_err(|e| e.to_string())?);

        let age = u64::read_bytes(&mut reader)?;

        let height = TCount::read_bytes(&mut reader)? as usize;
        let size = TCount::read_bytes(&mut reader)? as usize;
        let rands = TCount::read_bytes(&mut reader)? as usize;
        let layers = TCount::read_bytes(&mut reader)? as usize;
        let cycles = TCount::read_bytes(&mut reader)? as usize;

        let mut boolons = Vec::<Vec<Boolon>>::with_capacity(layers);
        for _ in 0..layers {
            let mut lr = Vec::<Boolon>::with_capacity(size);
            for _ in 0..size {
                lr.push(Boolon::read(&mut reader, height)?);
            }
            boolons.push(lr);
        }

        let score = std::f64::NEG_INFINITY;

        Ok(Self{cycles, rands, boolons, age, score})
    }

    pub fn run_batch(&self, rng: &mut ThreadRng, input_batch: &[u128], output_batch: &mut [u128]) {
        let (layers, rands, size, height) = (self.boolons.len(), self.rands, self.boolons[0].len(), self.boolons[0][0].funcs.len());

        let mut values = vec![0u128; size + rands]; // i-th value corresp. to i-th bit of bitstring, j-th bit of value corresp. to j-th bitstring in "batch"

        // Copy input to values of first boolons
        values[..input_batch.len()].copy_from_slice(input_batch);

        // Add outputs of "random" (fictitious) boolons
        rng.fill(&mut values[size..]);

        let mut tree_values = Vec::<Vec<u128>>::with_capacity(height);
        let mut leaves: usize = 1;
        for _ in 0..height {
            tree_values.push(vec![0u128; leaves]);
            leaves *= IN_BITS;
        }

        for _ in 0..self.cycles {
            for i in 0..layers {
                let mut next_values = vec![0u128; size + rands];

                let boolons = & self.boolons[i];
                for j in 0..size {
                    let b = & boolons[j];
                    let (inputs, funcs) = (& b.inputs, & b.funcs);

                    // From previous layer's boolons' values to highest level of input tree
                    let level_funcs = & funcs[height - 1];
                    let mut offs: usize = 0;
                    for k in 0..tree_values[height - 1].len() {
                        tree_values[height - 1][k] = U128_BOOL_3_FUNCS[level_funcs[k] as usize]( // eval for 128 "bit flows" (batch) simultaneously
                            values[inputs[offs] as usize],
                            values[inputs[offs + 1] as usize],
                            values[inputs[offs + 2] as usize]
                        );
                        offs += IN_BITS; // offs = IN_BITS * k
                    }

                    for level in (0..(height - 1)).rev() { // From higher level of input tree to lower one
                        let level_funcs = & funcs[level];
                        let mut offs: usize = 0;
                        for k in 0..tree_values[level].len() {
                            tree_values[level][k] = U128_BOOL_3_FUNCS[level_funcs[k] as usize]( // eval for 128 "bit flows" (batch) simultaneously
                                tree_values[level + 1][offs],
                                tree_values[level + 1][offs + 1],
                                tree_values[level + 1][offs + 2]
                            );
                            offs += IN_BITS; // offs = IN_BITS * k
                        }
                    }

                    next_values[j] =  tree_values[0][0];
                }

                // Add outputs of "random" (fictitious) boolons
                rng.fill(&mut next_values[size..]);

                values = next_values;
            }
        }

        // Copy values of last-before-random boolons to output
        let l = output_batch.len();
        output_batch.copy_from_slice(& values[(size - l)..size]);
    }

    pub fn run_single(&self, rng: &mut ThreadRng, input: &[u8], output: &mut [u8]) {
        // Explode input bytes to lowest bits of input batch
        let mut inp_batch = vec![0u128; input.len() << 3];
        for i in 0..inp_batch.len() {
            inp_batch[i] = ((input[i >> 3] >> (i & 7)) & 1) as u128;
        }

        // Run batch
        let mut outp_batch = vec![0u128; output.len() << 3];
        self.run_batch(rng, &inp_batch, &mut outp_batch);

        // Implode lowest bits of output batch to output bytes
        for i in 0..output.len() {
            output[i] = 0;
        }
        for i in 0..outp_batch.len() {
            output[i >> 3] |= ((outp_batch[i] & 1) as u8) << (i & 7);
        }
    }

    pub fn age(&self) -> u64 {
        self.age
    }

    pub fn inc_age(&mut self) {
        self.age += 1;
    }

    pub fn height(&self) -> usize {
        self.boolons[0][0].funcs.len()
    }

    pub fn size(&self) -> usize {
        self.boolons[0].len()
    }

    pub fn rands(&self) -> usize {
        self.rands
    }

    pub fn layers(&self) -> usize {
        self.boolons.len()
    }

    pub fn cycles(&self) -> usize {
        self.cycles
    }

    pub fn set_cycles(&mut self, cycles: usize) {
        self.cycles = cycles;
        self.score = std::f64::NEG_INFINITY;
    }

    pub fn score(&self) -> f64 {
        self.score
    }

    pub fn set_score(&mut self, score: f64) {
        self.score = score;
    }

    pub fn save(&self, filepath: &str) -> Result<(), String> {
        // "Binary serialization" ... TODO: compactify
        let mut writer = BufWriter::new(File::create(filepath).map_err(|e| e.to_string())?);

        self.age.write_bytes(&mut writer)?; // age

        (self.boolons[0][0].funcs.len() as TCount).write_bytes(&mut writer)?; // height
        (self.boolons[0].len() as TCount).write_bytes(&mut writer)?; // size
        (self.rands as TCount).write_bytes(&mut writer)?; // rands
        (self.boolons.len() as TCount).write_bytes(&mut writer)?; // layers
        (self.cycles as TCount).write_bytes(&mut writer)?; // cycles

        for i in 0..self.boolons.len() {
            for j in 0..self.boolons[i].len() {
                self.boolons[i][j].write(&mut writer)?;
            }
        }

        writer.flush().map_err(|e| e.to_string())?;

        Ok(())
    }

}
