use na::{DMatrix, DVector};

use pairing::{PrimeField, Field};
use pairing::bls12_381::{Fq, FqRepr};

pub fn u64_to_field<F: PrimeField>(x: u64) -> F {
    F::from_repr(
        <F::Repr as From<u64>>::from(x)
    ).unwrap()
}


#[derive(Debug, Clone, PartialEq)]
pub enum Gate<'a> {
    Attribute(&'a [u8]),
    Threshold {t: usize,  children: Vec<Gate<'a>>},
}

impl<'a> Gate<'a> {
    fn children(self) -> Vec<Gate<'a>> {
        match self {
            Gate::Threshold {children, ..} => children,
            Gate::Attribute (_) => Vec::new(),
        }
    }

    fn label(self) -> &'a [u8] {
        match self {
            Gate::Threshold {..} => b"",
            Gate::Attribute (l)  => l,
        }
    }

    fn lsss(&self) -> DMatrix<u64> {
        match self {
            &Gate::Attribute (_)  => DMatrix::from_row_slice(1, 1, &[1u64]),
            &Gate::Threshold {t, ref children} => {
                let m = children.len();
                DMatrix::from_fn(m, t, |i, j| {
                    ((i as u64) + 1).pow(j as u32)
                })
            }
        }
    }

    fn rec_msp(ast: AS<'a>)  -> AS<'a>{
        // find the first threshold-tree string rather than a simple attribute
        let at = ast.rho.iter().enumerate().skip_while(|&(_, g)| {
            match *g {
                Gate::Attribute (_) => true,
                Gate::Threshold {..} => false,
            }
        }).next();

        // if none is found, simply return ast, we're done.
        // otherwise, [cont.]
        if let Some((i, g)) = at {
            let mg = g.lsss();

            // refresh rho with the resolved attribute
            let mut rho = ast.rho.clone();
            rho.remove(i);
            // XXX. lots of cloning for the safety god.
            // perhaps the fact that we're inside a class can spare us
            // the copy of all the gates?
            for c in g.clone().children().iter() {
                rho.insert(i, c.clone());
            }

            let (m1, d1) = ast.m.shape();
            let (m2, d2) = mg.shape();

            // insert the lsss matrix into m
            let w = ast.m.row(i).kronecker(&mg.column(0));
            let mg = mg.remove_column(0);
            let m = ast.m.remove_row(i);

            let m = DMatrix::from_fn(
                m1 + m2 - 1,
                d1 + d2 - 1,
                |row, col| {
                    if row < i && col < d1           { m[(row, col)] }
                    else if row < i && col >= d1     { 0 }
                    else if row < i+m2 && col < d1   { w[(row-i, col)] }
                    else if row < i+m2 && col >= d1  { mg[(row-i, col-d1)] }
                    else if row >= i+m2 && col < d1  { m[(row-m2, col)] }
                    else                             { 0 }
                }
            );
            Self::rec_msp(AS {m, rho})
        } else {
            ast.clone()
        }
    }

    fn msp(self) -> AS<'a> {
        let rho = vec![self];
        let m = DMatrix::from_row_slice(1, 1, &[1u64]);
        let ast = Self::rec_msp(AS {m, rho});
        ast
    }

}

#[derive(Clone)]
struct AS<'a> {
    m: DMatrix<u64>,
    rho: Vec<Gate<'a>>,
}

impl<'a> AS<'a> {
    fn gen_shares<F: PrimeField>(&self) -> DVector<F> {
        let m = DMatrix::<F>::from_fn(self.m.nrows(), self.m.ncols(),
                                  |row, col| { u64_to_field(self.m[(row, col)]) }
        );
        DVector::from_fn(self.m.nrows(),
                         |row, col| { F::one() }
        )
    }

}

mod tc {
    use std::str;
    use std::str::FromStr;
    use nom::{digit, alphanumeric};
    use msp::Gate;

    named!(parens <Gate>, ws!(delimited!(tag!("("), threshold_gate, tag!(")") )) );

    named!(threshold_gate <Gate>,
           do_parse!(
               t: threshold >>
               children: attributes >>
               (Gate::Threshold {t, children})
           )
    );

    named!(attributes < Vec<Gate<'a>> >,
           many1!(preceded!(tag!(","), parse))
    );

    named!(threshold <usize>,
           map_res!(map_res!(ws!(digit), str::from_utf8),
                FromStr::from_str)
    );


    named!(attribute <Gate>,
           map!(ws!(alphanumeric), Gate::Attribute)
    );

    named!(pub parse <Gate>,
           alt!(attribute | parens)
    );
}


mod bc {
    use nom::{IResult, alphanumeric};
    use msp::Gate;

    named!(parens<Gate>, ws!(delimited!( tag!("("), expr, tag!(")") )) );

    named!(attribute <Gate>,
           map!(ws!(alphanumeric), Gate::Attribute)
    );

    named!(gate <Gate>,
           alt!(attribute | parens)
    );

    named!(term <Gate<'a>>,
           do_parse!(
               init: gate >>
               res:  fold_many0!(
                   pair!(alt!(tag!("^") | tag!("∧") | tag!("and")), gate),
                   init,
                   |acc, (op, g): (&[u8], Gate<'a>)| {
                       let  v: Vec<Gate<'a>> = vec![acc, g];
                       Gate::Threshold {t: 2usize, children: v}
                   }
               ) >>
               (res)
           )
    );

    named!(expr <Gate<'a>>,
           do_parse!(
               init: term >>
               res:  fold_many0!(
                   pair!(alt!(tag!("v") | tag!("or") | tag!("∨")), term),
                   init,
                   |acc, (op, g): (&[u8], Gate<'a>)| {
                       let  v: Vec<Gate<'a>> = vec![acc, g];
                       Gate::Threshold {t: 2usize, children: v}
                   }
               ) >>
               (res)
           )
    );

    // XXX. this needs to be changed into a result type.
    // Also note that making instead "expr" pub will fuck up everything
    // because the "lifetime is not defined" - not sure I really understood why.
    pub fn parse(i: &[u8]) -> IResult<&[u8], Gate, u32> {
        expr(i)
    }
}


#[test]
fn test_tc_parse() {
    // see eprint 2010/374
    let policy = b"(2, (2, (3, I, J, K, L), G, H), (2, A, B, C), (2, D, E, F))";
    let (_, pp) = tc::parse(policy).unwrap();
    let msp = pp.msp();
    assert_eq!(msp.rho.len(), 12);
    let expected = DMatrix::from_row_slice(12, 2 + (2-1) + (3-1) + (2-1) + (2-1),
                                           &[1, 1, 1, 0, 0, 0, 0,
                                             1, 1, 2, 0, 0, 0, 0,
                                             1, 1, 3, 0, 0, 0, 0,
                                             1, 2, 0, 1, 0, 0, 0,
                                             1, 2, 0, 2, 0, 0, 0,
                                             1, 2, 0, 3, 0, 0, 0,
                                             1, 3, 0, 0, 1, 0, 0,
                                             1, 3, 0, 0, 2, 0, 0,
                                             1, 3, 0, 0, 3, 1, 1,
                                             1, 3, 0, 0, 3, 2, 4,
                                             1, 3, 0, 0, 3, 3, 9,
                                             1, 3, 0, 0, 3, 4, 16,
                                           ]
    );
    assert_eq!(expected, msp.m)
}



use std::collections::VecDeque;

pub struct Bfs<'a> {
    stack: VecDeque<Gate<'a>>
}

impl<'a> Bfs<'a> {
    fn new(g: Gate<'a>) -> Self {
        let mut stack = VecDeque::new();
        stack.push_back(g);
        Bfs { stack }
    }

}

impl<'a> Iterator for Bfs<'a> {
    type Item = Gate<'a>;

    fn next(&mut self) -> Option<Self::Item> {
        let g = self.stack.pop_front()?;
        let children = g.clone().children();
        self.stack.extend(children.iter().cloned());
        Some(g)
    }

}

#[test]
fn test_bc_parse() {
    let policy = b"intern and topsecret or xkeyscore";
    let (_, gate) = bc::parse(policy).unwrap();
    let children = gate.children();
    assert_eq!(children.len(), 2);
}


#[test]
fn test_bfs() {
    let policy = b"1 and (3 and 4)";
    let (_, gate) = bc::parse(policy).unwrap();
    assert_eq!(gate.clone().children().len(), 2);
    let mut bfs = Bfs::new(gate);
    let root = bfs.next();
    let expected = Gate::Attribute(b"1");
    assert_eq!(Some(expected), bfs.next());
    assert_eq!(bfs.next().unwrap().children().len(), 2);
}

#[test]
fn test_matrix_field() {
    let policy = b"1 and 2";
    let (_, gate) = bc::parse(policy).unwrap();
    let msp = gate.msp();
    let m = msp.m;
    let mm = DVector::<Fq>::from_fn(m.ncols(),
                              |row, col| {
                                  u64_to_field(m[(row, col)])
                              }
    );
    let v  = DVector::from_fn(m.ncols(), |row, col| { 1u64 });
    let vv = DVector::from_fn(m.ncols(), |row, col| { Fq::one() });
    let d = mm.dot(&vv);

}