use std::{collections::HashMap, str::FromStr};

use itertools::Itertools;

#[derive(Clone)]
pub struct Polymer {
    pub first: char,
    pub last: char,
    pub map: HashMap<(char, char), usize>,
}

impl FromStr for Polymer {
    type Err = &'static str;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        let mut map = HashMap::new();
        for neighbours in s.chars().tuple_windows::<(char, char)>() {
            map.entry(neighbours)
                .and_modify(|count| *count += 1)
                .or_insert(1);
        }
        let first = s.chars().next().unwrap();
        let last = s.chars().last().unwrap();
        Ok(Self { first, last, map })
    }
}

#[derive(Copy, Clone)]
pub struct Rule {
    pub neighbours: (char, char),
    pub separator: char,
}

pub struct Rules {
    pub map: HashMap<(char, char), char>,
}

impl TryFrom<Vec<Rule>> for Rules {
    type Error = &'static str;

    fn try_from(rule_vec: Vec<Rule>) -> Result<Self, Self::Error> {
        if rule_vec.is_empty() {
            return Err("rule set can't be empty");
        }
        let mut rules = Rules {
            map: Default::default(),
        };
        for rule in rule_vec {
            if rules.map.contains_key(&rule.neighbours) {
                return Err("ambigous rule set");
            }
            rules.map.insert(rule.neighbours, rule.separator);
        }
        Ok(rules)
    }
}

impl Polymer {
    pub fn apply(&mut self, rules: &Rules) {
        let mut new_map = HashMap::new();

        for (neighbours, count) in self.map.iter() {
            if let Some(separator) = rules.map.get(neighbours) {
                new_map
                    .entry((neighbours.0, *separator))
                    .and_modify(|old_count| *old_count += count)
                    .or_insert(*count);
                new_map
                    .entry((*separator, neighbours.1))
                    .and_modify(|old_count| *old_count += count)
                    .or_insert(*count);
            } else {
                new_map
                    .entry(*neighbours)
                    .and_modify(|old_count| *old_count += count)
                    .or_insert(*count);
            }
        }
        self.map = new_map
    }
    pub fn element_counts(&self) -> HashMap<char, usize> {
        let mut map = HashMap::new();
        map.insert(self.first, 1);
        map.insert(self.last, 1);
        for ((left, right), count) in self.map.iter() {
            map.entry(*left)
                .and_modify(|old_count| *old_count += count)
                .or_insert(*count);
            map.entry(*right)
                .and_modify(|old_count| *old_count += count)
                .or_insert(*count);
        }
        map.iter()
            .map(|(element, count)| (*element, count / 2))
            .collect()
    }
}

pub struct Input {
    pub polymer: Polymer,
    pub rules: Rules,
}