use std::fs;
use std::mem;
use std::iter::Sum;
use std::ops::Add;
use std::path::Path;

#[derive(Clone,Debug,PartialEq)]
enum BinaryTree {
    Leaf(u8),
    Pair(Box<BinaryTree>, Box<BinaryTree>)
}

impl BinaryTree {
    fn pair(a: u8, b: u8) -> BinaryTree {
        BinaryTree::Pair(Box::new(BinaryTree::Leaf(a)), Box::new(BinaryTree::Leaf(b)))    
    }

    fn is_leaf(&self) -> bool {
        match self {
            BinaryTree::Leaf(_) => true,
            _                   => false
        }
    }

    fn is_simple_pair(&self) -> bool {
        if let BinaryTree::Pair(l,r) = self {
            l.is_leaf() && r.is_leaf()
        } else {
            false
        }    
    }

    fn get_simple_pair(&self) -> Option<(u8,u8)> {
        if let BinaryTree::Pair(l,r) = self {
            if let BinaryTree::Leaf(l) = **l {
                if let BinaryTree::Leaf(r) = **r {
                    return Some((l,r))
                }
            }
        }
        None
    }

    fn leftmost(&mut self) -> &mut BinaryTree {
        match self {
            BinaryTree::Leaf(_) => self,
            BinaryTree::Pair(l,_) => l.leftmost(),
        }
    }
    fn rightmost(&mut self) -> &mut BinaryTree {
        match self {
            BinaryTree::Leaf(_) => self,
            BinaryTree::Pair(_,r) => r.rightmost(),
        }
    }

    fn reduce(&mut self) {
        while self.explode().or_else(|| self.split()).is_some() { }
    }

    fn explode(&mut self) -> Option<BinaryTree> {
        self.explode_aux(0,None,None)
    }

    fn explode_aux(&mut self, depth: usize, left: Option<&mut BinaryTree>, right: Option<&mut BinaryTree>) -> Option<BinaryTree> {
        if depth > 3 && self.is_simple_pair() {
            let (l1,r1) = self.get_simple_pair().unwrap();
            if let Some(BinaryTree::Leaf(l2)) = left {
                *(left.unwrap()) = BinaryTree::Leaf(l1+*l2)    
            }
            if let Some(BinaryTree::Leaf(r2)) = right {
                *(right.unwrap()) = BinaryTree::Leaf(r1+*r2)
            }
            Some(mem::replace(self, BinaryTree::Leaf(0)))
        } else {
            match self {
                BinaryTree::Pair(l,r) => 
                    l.explode_aux(depth+1, left, Some(r.leftmost()))
                     .or_else(|| r.explode_aux(depth+1, Some(l.rightmost()), right)),
                BinaryTree::Leaf(_) => None
            }
        }
    }

    fn split(&mut self) -> Option<BinaryTree> {
        match self {
            BinaryTree::Leaf(l) if *l > 9 => {
                let n = *l;
                Some(mem::replace(self, BinaryTree::pair(n/2, n-n/2)))
            },
            BinaryTree::Pair(l,r) => {
                l.split().or_else(|| r.split())
            }
            _ => None
        }
    }

    fn magnitude(&self) -> u32 {
        match self {
            BinaryTree::Leaf(n) => *n as u32,
            BinaryTree::Pair(l,r) => 3*l.magnitude() + 2*r.magnitude()
        }
    }

    fn parse_number(bytes: &[u8]) -> (u8,usize) {
        let mut num = 0;
        let mut size = 0;
        while bytes[size].is_ascii_digit() {
            num = 10*num + (bytes[size]-b'0');
            size += 1;
        }
        (num, size)
    }

    fn parse(s: &[u8]) -> (BinaryTree,usize) {
        match s[0] {
            b'[' => {
                let (left, lsize) = BinaryTree::parse(&s[1..]);
                let (right, rsize) = BinaryTree::parse(&s[1+lsize+1..]);
                (BinaryTree::Pair(Box::new(left), Box::new(right)), lsize+rsize+3)
            },
            d if d.is_ascii_digit() => {
                let (num, size) = Self::parse_number(s);
                (BinaryTree::Leaf(num), size)
            },
            _ => unreachable!()
        }
    }
}

impl From<&str> for BinaryTree {
    fn from(s: &str) -> BinaryTree {
        BinaryTree::parse(s.as_bytes()).0
    }
}

impl Add for BinaryTree {
    type Output = Self;
    fn add(self, other: Self) -> Self {
        let mut pair = BinaryTree::Pair(Box::new(self),Box::new(other));
        pair.reduce();
        pair
    }
}

impl Sum for BinaryTree{
    fn sum<I>(iter: I) -> BinaryTree
    where
        I: Iterator<Item = BinaryTree> {
        iter.reduce(|a,b| a+b).unwrap()
    }
}

/* AOC21 Day 18: https://adventofcode.com/2021/day/18 */
fn main() {
    let input = Path::new("resources").join("input.txt");
    let content = fs::read_to_string(input).expect("Unable to read input file");
    println!("Ex1: the result is {}", content.lines().map(|l| BinaryTree::from(l)).sum::<BinaryTree>().magnitude());
    let nums: Vec<BinaryTree> = content.lines().map(|l| BinaryTree::from(l)).collect();
    println!("Ex2: the result is {}", nums.iter().flat_map(|t1| nums.iter().map(move |t2| (t1.clone()+t2.clone()).magnitude())).max().unwrap());
}

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

    #[test]
    fn parse_pair() {
        let result = BinaryTree::Pair(
            Box::new(BinaryTree::Leaf(1)),
            Box::new(BinaryTree::Leaf(2))
        );
        assert_eq!(result, BinaryTree::from("[1,2]"));
    }

    #[test]
    fn parse_left_subpair() {
        let result = BinaryTree::Pair(
            Box::new(BinaryTree::Pair(
                Box::new(BinaryTree::Leaf(1)),
                Box::new(BinaryTree::Leaf(2))
            )),
            Box::new(BinaryTree::Leaf(3))
        );
        assert_eq!(result, BinaryTree::from("[[1,2],3]"));
    }

    #[test]
    fn parse_right_subpair() {
        let result = BinaryTree::Pair(
            Box::new(BinaryTree::Leaf(9)),
            Box::new(BinaryTree::Pair(
                Box::new(BinaryTree::Leaf(8)),
                Box::new(BinaryTree::Leaf(7))
            ))
        );
        assert_eq!(result, BinaryTree::from("[9,[8,7]]"));
    }

    #[test]
    fn parse_subpair() {
        let result = BinaryTree::Pair(
            Box::new(BinaryTree::Pair(
                Box::new(BinaryTree::Leaf(1)),
                Box::new(BinaryTree::Leaf(9))
            )),
            Box::new(BinaryTree::Pair(
                Box::new(BinaryTree::Leaf(8)),
                Box::new(BinaryTree::Leaf(5))
            ))
        );
        assert_eq!(result, BinaryTree::from("[[1,9],[8,5]]"))
    }

    #[test]
    fn parse_complex1() {
        let result = BinaryTree::Pair(
            Box::new(BinaryTree::Pair(
                Box::new(BinaryTree::Pair(
                    Box::new(BinaryTree::Pair(
                        Box::new(BinaryTree::Leaf(1)),
                        Box::new(BinaryTree::Leaf(2))
                    )),
                    Box::new(BinaryTree::Pair(
                        Box::new(BinaryTree::Leaf(3)),
                        Box::new(BinaryTree::Leaf(4))
                    ))
                )),
                Box::new(BinaryTree::Pair(
                    Box::new(BinaryTree::Pair(
                        Box::new(BinaryTree::Leaf(5)),
                        Box::new(BinaryTree::Leaf(6))
                    )),
                    Box::new(BinaryTree::Pair(
                        Box::new(BinaryTree::Leaf(7)),
                        Box::new(BinaryTree::Leaf(8))
                    )),
                ))
            )),
            Box::new(BinaryTree::Leaf(9))
        );
        assert_eq!(result, BinaryTree::from("[[[[1,2],[3,4]],[[5,6],[7,8]]],9]"))
    }

    #[test]
    fn parse_complex2() {
        let result = BinaryTree::Pair(
            Box::new(BinaryTree::Pair(
                Box::new(BinaryTree::Pair(
                    Box::new(BinaryTree::Leaf(9)),
                    Box::new(BinaryTree::Pair(
                        Box::new(BinaryTree::Leaf(3)),
                        Box::new(BinaryTree::Leaf(8))
                    ))
                )),
                Box::new(BinaryTree::Pair(
                    Box::new(BinaryTree::Pair(
                        Box::new(BinaryTree::Leaf(0)),
                        Box::new(BinaryTree::Leaf(9))
                    )),
                    Box::new(BinaryTree::Leaf(6))
                ))
            )),
            Box::new(BinaryTree::Pair(
                Box::new(BinaryTree::Pair(
                    Box::new(BinaryTree::Pair(
                        Box::new(BinaryTree::Leaf(3)),
                        Box::new(BinaryTree::Leaf(7))
                    )),
                    Box::new(BinaryTree::Pair(
                        Box::new(BinaryTree::Leaf(4)),
                        Box::new(BinaryTree::Leaf(9))
                    )),
                )),
                Box::new(BinaryTree::Leaf(3)),
            )),
        );
        assert_eq!(result, BinaryTree::from("[[[9,[3,8]],[[0,9],6]],[[[3,7],[4,9]],3]]"))
    }

    #[test]
    fn parse_complex3() {
        let result = BinaryTree::Pair(
            Box::new(BinaryTree::Pair(
                Box::new(BinaryTree::Pair(
                    Box::new(BinaryTree::Pair(
                        Box::new(BinaryTree::Leaf(1)),
                        Box::new(BinaryTree::Leaf(3))
                    )),
                    Box::new(BinaryTree::Pair(
                        Box::new(BinaryTree::Leaf(5)),
                        Box::new(BinaryTree::Leaf(3))
                    ))
                )),
                Box::new(BinaryTree::Pair(
                    Box::new(BinaryTree::Pair(
                        Box::new(BinaryTree::Leaf(1)),
                        Box::new(BinaryTree::Leaf(3))
                    )),
                    Box::new(BinaryTree::Pair(
                        Box::new(BinaryTree::Leaf(8)),
                        Box::new(BinaryTree::Leaf(7))
                    ))
                )),
            )),
            Box::new(BinaryTree::Pair(
                Box::new(BinaryTree::Pair(
                    Box::new(BinaryTree::Pair(
                        Box::new(BinaryTree::Leaf(4)),
                        Box::new(BinaryTree::Leaf(9))
                    )),
                    Box::new(BinaryTree::Pair(
                        Box::new(BinaryTree::Leaf(6)),
                        Box::new(BinaryTree::Leaf(9))
                    ))
                )),
                Box::new(BinaryTree::Pair(
                    Box::new(BinaryTree::Pair(
                        Box::new(BinaryTree::Leaf(8)),
                        Box::new(BinaryTree::Leaf(2))
                    )),
                    Box::new(BinaryTree::Pair(
                        Box::new(BinaryTree::Leaf(7)),
                        Box::new(BinaryTree::Leaf(3))
                    ))
                )),
            )),
        );
        assert_eq!(result, BinaryTree::from("[[[[1,3],[5,3]],[[1,3],[8,7]]],[[[4,9],[6,9]],[[8,2],[7,3]]]]"))
    }

    #[test]
    fn explode_left() {
        let mut num = BinaryTree::from("[[[[[9,8],1],2],3],4]");
        num.explode();
        assert_eq!(BinaryTree::from("[[[[0,9],2],3],4]"), num)
    }

    #[test]
    fn explode_right() {
        let mut num = BinaryTree::from("[7,[6,[5,[4,[3,2]]]]]");
        num.explode();
        assert_eq!(BinaryTree::from("[7,[6,[5,[7,0]]]]"), num)
    }

    #[test]
    fn explode_left_right() {
        let mut num = BinaryTree::from("[[6,[5,[4,[3,2]]]],1]");
        num.explode();
        assert_eq!(BinaryTree::from("[[6,[5,[7,0]]],3]"), num)
    }

    #[test]
    fn explode_leftmost() {
        let mut num = BinaryTree::from("[[3,[2,[1,[7,3]]]],[6,[5,[4,[3,2]]]]]");
        num.explode();
        assert_eq!(BinaryTree::from("[[3,[2,[8,0]]],[9,[5,[4,[3,2]]]]]"), num)
    }

    #[test]
    fn sum() {
        assert_eq!(
            BinaryTree::from("[[[[0,7],4],[[7,8],[6,0]]],[8,1]]"),
            BinaryTree::from("[[[[4,3],4],4],[7,[[8,4],9]]]") + BinaryTree::from("[1,1]")
        )
    }

    #[test]
    fn sum_array() {
        let array = [
            "[[[0,[4,5]],[0,0]],[[[4,5],[2,6]],[9,5]]]",
            "[7,[[[3,7],[4,3]],[[6,3],[8,8]]]]",
            "[[2,[[0,8],[3,4]]],[[[6,7],1],[7,[1,6]]]]",
            "[[[[2,4],7],[6,[0,5]]],[[[6,8],[2,8]],[[2,1],[4,5]]]]",
            "[7,[5,[[3,8],[1,4]]]]",
            "[[2,[2,2]],[8,[8,1]]]",
            "[2,9]",
            "[1,[[[9,3],9],[[9,0],[0,7]]]]",
            "[[[5,[7,4]],7],1]",
            "[[[[4,2],2],6],[8,7]]"
        ];
        assert_eq!(
            BinaryTree::from("[[[[8,7],[7,7]],[[8,6],[7,7]]],[[[0,7],[6,6]],[8,7]]]"),
            array.iter().map(|&s| BinaryTree::from(s)).sum()
        )
    }
    #[test]

    fn sum_array_with_magnitude() {
        let array = [
            "[[[0,[5,8]],[[1,7],[9,6]]],[[4,[1,2]],[[1,4],2]]]",
            "[[[5,[2,8]],4],[5,[[9,9],0]]]",
            "[6,[[[6,2],[5,6]],[[7,6],[4,7]]]]",
            "[[[6,[0,7]],[0,9]],[4,[9,[9,0]]]]",
            "[[[7,[6,4]],[3,[1,3]]],[[[5,5],1],9]]",
            "[[6,[[7,3],[3,2]]],[[[3,8],[5,7]],4]]",
            "[[[[5,4],[7,7]],8],[[8,3],8]]",
            "[[9,3],[[9,9],[6,[4,9]]]]",
            "[[2,[[7,7],7]],[[5,8],[[9,3],[0,2]]]]",
            "[[[[5,2],5],[8,[3,7]]],[[5,[7,5]],[4,4]]]"
        ];
        let sum = array.iter().map(|&s| BinaryTree::from(s)).sum();
        assert_eq!(BinaryTree::from("[[[[6,6],[7,6]],[[7,7],[7,0]]],[[[7,7],[7,7]],[[7,8],[9,9]]]]"), sum);
        assert_eq!(4140, sum.magnitude())
    }
}