extern crate num;

use std::cmp::Ordering;
use std::collections::HashMap;
use std::hash::{Hash, Hasher};

use num::{One, Zero};

mod node_defs;
mod node_ordering;

use node_defs::{Container, container, NodeData, Sum, Product, Power, Primitive};

impl<T: NodeData> Sum<T> {
    fn negate(&self, env: &mut Environment<T>) -> Container<Sum<T>> {
        env.make_sum(!self.minus, &self.constant, &self.terms)
    }
}

impl<T: NodeData> Product<T> {
    fn negate(&self, env: &mut Environment<T>) -> Container<Product<T>> {
        env.make_product(&-self.coefficient.clone(), &self.powers).unwrap()
    }
}

type SelfMap<T> = HashMap<T, Container<T>>;

struct Environment<T: NodeData> {
    sum: SelfMap<Sum<T>>,
    product: SelfMap<Product<T>>,
    power: SelfMap<Power<T>>,
    primitive: SelfMap<Primitive<T>>,
    sum_derivative: HashMap<(Sum<T>, T::Parameter), Container<Sum<T>>>,
    product_derivative: HashMap<(Product<T>, T::Parameter), Container<Sum<T>>>,
    power_derivative: HashMap<(Power<T>, T::Parameter), Container<Sum<T>>>,
    primitive_derivative: HashMap<(Primitive<T>, T::Parameter), Container<Sum<T>>>,
}

impl<T: NodeData> Sum<T> {
    fn at_zero(&self) -> T::Constant {
        if self.minus {
            -self.constant.clone()
        } else {
            self.constant.clone()
        }
    }

    fn adjust_constant(&self,
                       constant: &T::Constant,
                       env: &mut Environment<T>)
                       -> Container<Sum<T>> {
        env.make_sum(self.minus,
                     &(if self.minus {
                           -constant.clone()
                       } else {
                           constant.clone()
                       }),
                     &self.terms)
    }

    fn conditional_negate(&self, env: &mut Environment<T>) -> Vec<Container<Product<T>>> {
        if self.minus {
            env.negate_products(&self.terms)
        } else {
            self.terms
                .iter()
                .map(|p| p.insert(env))
                .collect::<Vec<_>>()
        }
    }
}

impl<T: NodeData> Environment<T> {
    fn make_sum(&mut self,
                minus: bool,
                constant: &T::Constant,
                terms: &Vec<Container<Product<T>>>)
                -> Container<Sum<T>> {
        if minus && terms.len() == 0 {
            self.make_sum(false, &-constant.clone(), terms)
        } else {
            Sum::new(minus, constant.clone(), terms.clone()).insert(self)
        }
    }

    fn make_product(&mut self,
                    coefficient: &T::Constant,
                    powers: &Vec<Container<Power<T>>>)
                    -> Option<Container<Product<T>>> {
        if coefficient.clone() == T::Constant::zero() || powers.len() == 0 {
            None
        } else {
            Some(Product::new(coefficient.clone(), powers.clone()).insert(self))
        }
    }

    fn make_power(&mut self,
                  primitive: &Container<Primitive<T>>,
                  exponent: &T::Exponent)
                  -> Option<Container<Power<T>>> {
        if exponent.clone() <= T::Exponent::zero() {
            None
        } else {
            Some(Power::new(exponent.clone(), primitive.clone()).insert(self))
        }
    }

    fn make_system_variable(&mut self, index: &T::SystemVariable) -> Container<Primitive<T>> {
        Primitive::SystemVariable(index.clone()).insert(self)
    }

    fn make_parameter(&mut self, name: &T::Parameter) -> Container<Primitive<T>> {
        Primitive::Parameter(name.clone()).insert(self)
    }

    fn make_input(&mut self, index: &T::Input) -> Container<Primitive<T>> {
        Primitive::Input(index.clone()).insert(self)
    }

    fn make_sigmoid(&mut self, minus: bool, sum: &Container<Sum<T>>) -> Container<Primitive<T>> {
        if sum.minus || (sum.terms.len() == 0 && sum.constant < T::Constant::zero()) {
            let sum = sum.negate(self);
            self.make_sigmoid(!minus, &sum)
        } else {
            Primitive::Sigmoid(minus, sum.clone()).insert(self)
        }
    }

    fn negate_products(&mut self,
                       terms: &Vec<Container<Product<T>>>)
                       -> Vec<Container<Product<T>>> {
        terms.iter().map(|p| p.negate(self)).collect::<Vec<_>>()
    }
    fn number(&mut self, number: &T::Constant) -> Container<Sum<T>> {
        self.make_sum(false, number, &vec![])
    }

    fn add<A: Node<T>, B: Node<T>>(&mut self,
                                   left: &Container<A>,
                                   right: &Container<B>)
                                   -> Container<Sum<T>> {
        let left_sum = left.as_sum(self);
        let right_sum = right.as_sum(self);
        let constant = left_sum.at_zero() + right_sum.at_zero();
        if left_sum.terms.len() == 0 {
            return right_sum.adjust_constant(&constant, self);
        }
        if right_sum.terms.len() == 0 {
            return left_sum.adjust_constant(&constant, self);
        }

        // Begin port of weird logic
        let left_terms = left_sum.conditional_negate(self);
        let right_terms = right_sum.conditional_negate(self);
        let left_length = left_terms.len();
        let right_length = right_terms.len();

        let mut terms = Vec::with_capacity(left_length + right_length);
        let mut left_index = 0;
        let mut right_index = 0;

        while left_index < left_length && right_index < right_length {
            let ref left_term = left_terms[left_index];
            let ref right_term = right_terms[right_length];
            match left_term.fuzzy_cmp(right_term.as_ref()) {
                Ordering::Less => {
                    terms.push(left_term.clone());
                    left_index += 1;
                }
                Ordering::Equal => {
                    let coefficient = left_term.coefficient.clone() +
                                      right_term.coefficient.clone();
                    if coefficient != T::Constant::zero() {
                        terms.push(self.make_product(&coefficient, &left_term.powers).unwrap());
                    }
                    left_index += 1;
                    right_index += 1;
                }
                Ordering::Greater => {
                    terms.push(right_term.clone());
                    right_index += 1;
                }
            }
            if left_index == left_length {
                terms.extend_from_slice(&right_terms[..].split_at(right_index).1);
            }
            if right_index == right_length {
                terms.extend_from_slice(&left_terms[..].split_at(left_index).1);
            }
        }

        terms.shrink_to_fit(); // Why not.

        if terms.len() == 0 || terms[0].coefficient >= T::Constant::zero() {
            self.make_sum(false, &constant, &terms)
        } else {
            let terms = self.negate_products(&terms);
            self.make_sum(true, &constant, &terms)
        }
    }

    fn multiply<A: Node<T>, B: Node<T>>(&mut self,
                                        left: &Container<A>,
                                        right: &Container<B>)
                                        -> Container<Sum<T>> {
        let left_sum = left.as_sum(self);
        let right_sum = right.as_sum(self);
        let minus = left_sum.minus != right_sum.minus;

        let first_constant = left_sum.constant.clone() * right_sum.constant.clone();
        let first = self.make_sum(minus, &first_constant, &vec![]);

        let mut outer_terms = Vec::with_capacity(right_sum.terms.len());
        if left_sum.constant != T::Constant::zero() {
            for term in &right_sum.terms {
                let outer_coefficient = left_sum.constant.clone() * term.coefficient.clone();
                let outer_term = self.make_product(&outer_coefficient, &term.powers).unwrap();
                if left_sum.constant < T::Constant::zero() {
                    outer_terms.push(outer_term.negate(self));
                } else {
                    outer_terms.push(outer_term);
                }
            }
        }
        outer_terms.shrink_to_fit();
        let outer = self.make_sum((left_sum.constant < T::Constant::zero()) != minus,
                                  &T::Constant::zero(),
                                  &outer_terms);

        // Should be possible to turn these into functions.
        let mut inner_terms = Vec::with_capacity(left_sum.terms.len());
        if right_sum.constant != T::Constant::zero() {
            for term in &left_sum.terms {
                let inner_coefficient = right_sum.constant.clone() * term.coefficient.clone();
                let inner_term = self.make_product(&inner_coefficient, &term.powers).unwrap();
                if right_sum.constant < T::Constant::zero() {
                    inner_terms.push(inner_term.negate(self));
                } else {
                    inner_terms.push(inner_term);
                }
            }
        }
        inner_terms.shrink_to_fit();
        let inner = self.make_sum((right_sum.constant < T::Constant::zero()) != minus,
                                  &T::Constant::zero(),
                                  &inner_terms);

        let mut last = self.number(&T::Constant::zero());
        for left_term in &left_sum.terms {
            for right_term in &right_sum.terms {
                let combined_terms = self._combine_terms(&left_term.powers, &right_term.powers);
                let last_coefficient = left_term.coefficient.clone() *
                                       right_term.coefficient.clone();
                let product = self.make_product(&last_coefficient, &combined_terms).unwrap();
                last = self.add(&last, &product);
            }
        }

        last = self.add(&last, &inner);
        last = self.add(&last, &outer);
        self.add(&last, &first)
    }

    fn _combine_terms(&mut self,
                      left: &Vec<Container<Power<T>>>,
                      right: &Vec<Container<Power<T>>>)
                      -> Vec<Container<Power<T>>> {
        let left_length = left.len();
        let right_length = right.len();

        let mut powers = Vec::with_capacity(left_length + right_length);
        let mut left_index = 0;
        let mut right_index = 0;

        while left_index < left_length && right_index < right_length {
            let ref left_power = left[left_index];
            let ref right_power = right[right_length];
            match left_power.fuzzy_cmp(&right_power) {
                Ordering::Less => {
                    powers.push(left_power.clone());
                    left_index += 1;
                }
                Ordering::Equal => {
                    let exponent = left_power.exponent.clone() + right_power.exponent.clone();
                    powers.push(self.make_power(&left_power.primitive, &exponent).unwrap());
                    left_index += 1;
                    right_index += 1;
                }
                Ordering::Greater => {
                    powers.push(right_power.clone());
                    right_index += 1;
                }
            }
            if left_index == left_length {
                powers.extend_from_slice(&right[..].split_at(right_index).1);
            }
            if right_index == right_length {
                powers.extend_from_slice(&left[..].split_at(left_index).1);
            }
        }

        powers.shrink_to_fit(); // Why not.

        powers

    }
}

trait Summable<T: NodeData> {
    fn as_sum(&self, env: &mut Environment<T>) -> Container<Sum<T>>;
}

trait Productable<T: NodeData> {
    fn as_product(&self, env: &mut Environment<T>) -> Container<Product<T>>;
}

fn _as_sum<T: NodeData>(productable: &Productable<T>,
                        env: &mut Environment<T>)
                        -> Container<Sum<T>> {
    let product = productable.as_product(env);
    if product.coefficient < T::Constant::zero() {
        let product = product.negate(env);
        env.make_sum(true, &T::Constant::zero(), &vec![product])
    } else {
        env.make_sum(false, &T::Constant::zero(), &vec![product])
    }
}

impl<T: NodeData> Summable<T> for Product<T> {
    fn as_sum(&self, env: &mut Environment<T>) -> Container<Sum<T>> {
        _as_sum(self, env)
    }
}

impl<T: NodeData> Summable<T> for Power<T> {
    fn as_sum(&self, env: &mut Environment<T>) -> Container<Sum<T>> {
        _as_sum(self, env)
    }
}

impl<T: NodeData> Summable<T> for Primitive<T> {
    fn as_sum(&self, env: &mut Environment<T>) -> Container<Sum<T>> {
        _as_sum(self, env)
    }
}

trait Node<T: NodeData>: Summable<T> {
    fn partial_derivative(&self,
                          variable: &T::Parameter,
                          env: &mut Environment<T>)
                          -> Container<Sum<T>>;

    fn fancy_clone(&self, env: &mut Environment<T>) -> Self;

    fn get_storage<'a>(&self, env: &'a mut Environment<T>) -> &'a mut SelfMap<Self>
        where Self: Sized;

    fn insert(&self, env: &mut Environment<T>) -> Container<Self>
        where Self: Sized + Clone + Eq + Hash
    {
        if !self.get_storage(env).contains_key(self) {
            let key = self.fancy_clone(env);
            let value = key.clone();
            self.get_storage(env).insert(key, container(value));
        }
        match self.get_storage(env).get(self) {
                Some(node) => node,
                None => unreachable!(),
            }
            .clone()
    }
}

impl<T: NodeData> Summable<T> for Sum<T> {
    fn as_sum(&self, env: &mut Environment<T>) -> Container<Sum<T>> {
        self.insert(env)
    }
}

impl<T: NodeData> Node<T> for Sum<T> {
    fn partial_derivative(&self,
                          variable: &T::Parameter,
                          env: &mut Environment<T>)
                          -> Container<Sum<T>> {
        if !env.sum_derivative.contains_key(&(self.clone(), variable.clone())) {
            let key = (self.fancy_clone(env), variable.clone());
            let mut result = env.number(&T::Constant::zero());
            for term in &self.terms {
                let partial = term.partial_derivative(variable, env);
                result = env.add(&result, &partial);
            }
            if self.minus {
                result = result.negate(env);
            }
            env.sum_derivative.insert(key, result);
        }
        match env.sum_derivative.get(&(self.clone(), variable.clone())) {
            Some(derivative) => derivative.clone(),
            None => unreachable!(),
        }
    }

    fn fancy_clone(&self, env: &mut Environment<T>) -> Sum<T> {
        Sum {
            pre_hash: self.pre_hash.clone(),
            minus: self.minus.clone(),
            constant: self.constant.clone(),
            terms: self.terms
                .iter()
                .map(|p| p.insert(env))
                .collect::<Vec<_>>(),
        }
    }

    fn get_storage<'a>(&self, env: &'a mut Environment<T>) -> &'a mut SelfMap<Sum<T>> {
        &mut env.sum
    }
}

impl<T: NodeData> Productable<T> for Product<T> {
    fn as_product(&self, env: &mut Environment<T>) -> Container<Product<T>> {
        self.insert(env)
    }
}

impl<T: NodeData> Node<T> for Product<T> {
    fn partial_derivative(&self,
                          variable: &T::Parameter,
                          env: &mut Environment<T>)
                          -> Container<Sum<T>> {
        if !env.product_derivative.contains_key(&(self.clone(), variable.clone())) {
            let key = (self.fancy_clone(env), variable.clone());
            let result = if self.powers.len() == 1 {
                self.powers[0].partial_derivative(variable, env)
            } else {
                let mut result = env.number(&T::Constant::zero());
                unimplemented!()
                /*for (index, focus) in self.powers.iter().enumerate() {
                    let remainder = env.make_product(T::Constant::one(), )
                }*/
            };
            env.product_derivative.insert(key, result);
        }
        match env.product_derivative.get(&(self.clone(), variable.clone())) {
            Some(derivative) => derivative.clone(),
            None => unreachable!(),
        }
    }

    fn fancy_clone(&self, env: &mut Environment<T>) -> Product<T> {
        Product {
            pre_hash: self.pre_hash.clone(),
            coefficient: self.coefficient.clone(),
            powers: self.powers
                .iter()
                .map(|p| p.insert(env))
                .collect::<Vec<_>>(),
        }
    }

    fn get_storage<'a>(&self, env: &'a mut Environment<T>) -> &'a mut SelfMap<Product<T>> {
        &mut env.product
    }
}

impl<T: NodeData> Productable<T> for Power<T> {
    fn as_product(&self, env: &mut Environment<T>) -> Container<Product<T>> {
        let power = self.insert(env);
        env.make_product(&T::Constant::one(), &vec![power]).unwrap()
    }
}

impl<T: NodeData> Node<T> for Power<T> {
    fn partial_derivative(&self,
                          variable: &T::Parameter,
                          env: &mut Environment<T>)
                          -> Container<Sum<T>> {
        unimplemented!();
    }

    fn fancy_clone(&self, env: &mut Environment<T>) -> Power<T> {
        Power {
            exponent: self.exponent.clone(),
            primitive: self.primitive.insert(env),
        }
    }

    fn get_storage<'a>(&self, env: &'a mut Environment<T>) -> &'a mut SelfMap<Power<T>> {
        &mut env.power
    }
}

impl<T: NodeData> Primitive<T> {
    fn as_power(&self, env: &mut Environment<T>) -> Container<Power<T>> {
        let primitive = self.insert(env);
        env.make_power(&primitive, &T::Exponent::one()).unwrap()
    }
}

impl<T: NodeData> Productable<T> for Primitive<T> {
    fn as_product(&self, env: &mut Environment<T>) -> Container<Product<T>> {
        match self {
            &Primitive::Sigmoid(true, ref sum) => {
                let sigmoid = env.make_sigmoid(false, sum);
                let power = sigmoid.as_power(env);
                env.make_product(&-T::Constant::one(), &vec![power]).unwrap()
            }
            _ => {
                let power = self.as_power(env);
                power.as_product(env)
            }
        }
    }
}

impl<T: NodeData> Node<T> for Primitive<T> {
    fn partial_derivative(&self,
                          variable: &T::Parameter,
                          env: &mut Environment<T>)
                          -> Container<Sum<T>> {
        unimplemented!();
    }

    fn fancy_clone(&self, env: &mut Environment<T>) -> Primitive<T> {
        match self {
            &Primitive::Sigmoid(minus, ref a) => Primitive::Sigmoid(minus, a.insert(env)),
            _ => self.clone(),
        }
    }

    fn get_storage<'a>(&self, env: &'a mut Environment<T>) -> &'a mut SelfMap<Primitive<T>> {
        &mut env.primitive
    }
}

#[cfg(test)]
mod tests {
    #[test]
    fn it_works() {}
}