use std::fs;
use std::path::Path;
use std::collections::HashMap;
use std::iter::FromIterator;

type Expand = (char,char);
type Rules = HashMap<(char,char),char>;
type Frequency = HashMap<char,u128>;
type FrequencyMap = HashMap<Expand,Frequency>;

/* AOC21 Day 14: https://adventofcode.com/2021/day/14 */
fn main() {
    let input = Path::new("resources").join("input.txt");
    let content = fs::read_to_string(input).expect("Unable to read input file");
    let (template,rules) = content.split_once("\n\n").expect("Malformed input");
    let rules = Rules::from_iter(rules.lines().map(|l| {
        let (from,to) = l.split_once(" -> ").expect("Malformed input");
        let chars = from.chars().collect::<Vec<char>>();
        ((chars[0],chars[1]),to.parse().unwrap())
    }));
    let mut frequency = FrequencyMap::new();
    rules.iter().for_each(|(&(a,b),&c)| {
        let mut map = HashMap::new(); map.insert(c,1);
        frequency.insert((a,b),map);
    });
    (1..10).for_each(|_| frequency = expand(&frequency, &rules));
    println!("Ex1: the result is {}", template_frequency(template, &frequency));
    (0..30).for_each(|_| frequency = expand(&frequency, &rules));
    println!("Ex2: the result is {}", template_frequency(template, &frequency));
}

fn expand(frequency: &FrequencyMap, rules: &Rules) -> FrequencyMap {
    let mut freq2 = FrequencyMap::new();
    rules.iter().for_each(|(&(a,b),&c)| {
        let mut map = merge(frequency.get(&(a,c)).unwrap(), frequency.get(&(c,b)).unwrap());
        *map.entry(c).or_default() += 1;
        *freq2.entry((a,b)).or_default() = map;
    });
    freq2
}

fn template_frequency(template: &str, frequency: &FrequencyMap) -> u128 {
    let mut freq = HashMap::new();
    template.chars().for_each(|c| *freq.entry(c).or_default() += 1);
    freq = template.chars().collect::<Vec<char>>()
                   .as_slice().windows(2)
                   .fold(freq, |acc, w| merge(&acc, frequency.get(&(w[0],w[1])).unwrap()) );
    freq.values().max().unwrap() - freq.values().min().unwrap()
}

fn merge(a: &Frequency, b: &Frequency) -> Frequency {
    let mut map = a.clone();
    b.iter().for_each(|(&k,&v)| *map.entry(k).or_default() += v);
    map
}