use std::collections::HashMap;
use std::error::Error;
use std::fs::File;
use std::io::{self, BufRead};
use std::vec::Vec;

fn recurse_pair(
    a: char,
    b: char,
    depth: usize,
    rules: &HashMap<(char, char), char>,
    cache: &mut HashMap<(char, char, usize), HashMap<char, usize>>,
) -> HashMap<char, usize> {
    if cache.contains_key(&(a, b, depth)) {
        return cache[&(a, b, depth)].clone();
    }

    let middle = rules.get(&(a, b)).unwrap();

    let mut res = HashMap::new();
    res.insert(*middle, 1);

    if depth == 1 {
        return res;
    }

    let res1 = recurse_pair(a, *middle, depth - 1, rules, cache);
    let res2 = recurse_pair(*middle, b, depth - 1, rules, cache);

    for (c, count) in res1.iter() {
        let count2 = res.get(c).unwrap_or_else(|| &0).clone();
        res.insert(*c, count + count2);
    }

    for (c, count) in res2.iter() {
        let count2 = res.get(c).unwrap_or_else(|| &0).clone();
        res.insert(*c, count + count2);
    }

    cache.insert((a, b, depth), res.clone());

    res
}

fn compute_dist(
    polymer_template: &Vec<char>,
    depth: usize,
    rules: &HashMap<(char, char), char>,
    cache: &mut HashMap<(char, char, usize), HashMap<char, usize>>,
) -> usize {
    let mut dist: HashMap<char, usize> = HashMap::new();
    for c in polymer_template.iter() {
        let count = dist.get(c).unwrap_or_else(|| &0).clone();
        dist.insert(*c, count + 1);
    }

    for pair in polymer_template.windows(2) {
        let res = recurse_pair(pair[0], pair[1], depth, &rules, cache);
        for (c, count) in res.iter() {
            let count2 = dist.get(c).unwrap_or_else(|| &0).clone();
            dist.insert(*c, count + count2);
        }
    }

    println!("{:?}", dist);

    let mut counts: Vec<&usize> = dist.values().collect();
    counts.sort();

    *counts.last().unwrap() - *counts.first().unwrap()
}

fn main() -> Result<(), Box<dyn Error>> {
    let file = File::open("inputs/day14.txt")?;
    let mut lines = io::BufReader::new(file).lines();

    let polymer_template: Vec<char> = lines.next().unwrap().unwrap().chars().collect();
    lines.next();

    let mut rules: HashMap<(char, char), char> = HashMap::new();
    while let Some(Ok(line)) = lines.next() {
        if line == "" {
            break;
        }

        let mut parts = line.split(" -> ");
        let pair = parts.next().unwrap();
        let new = parts.next().unwrap();

        rules.insert(
            (pair.chars().nth(0).unwrap(), pair.chars().nth(1).unwrap()),
            new.chars().nth(0).unwrap(),
        );
    }

    let mut cache: HashMap<(char, char, usize), HashMap<char, usize>> = HashMap::new();

    let answer1 = compute_dist(&polymer_template, 10, &rules, &mut cache);

    println!("Answer1: {}", answer1);

    let answer2 = compute_dist(&polymer_template, 40, &rules, &mut cache);

    println!("Answer2: {}", answer2);

    Ok(())
}