use itertools::Itertools;
use rayon::prelude::*;
use std::cell::RefCell;
use std::collections::HashMap;
use std::fs::File;
use std::io::{BufRead, BufReader, Lines};
use std::time::{Duration, Instant};

// BOILERPLATE
type InputIter = Lines<BufReader<File>>;

fn get_input() -> InputIter {
    let f = File::open("input").unwrap();
    let br = BufReader::new(f);
    br.lines()
}

fn main() {
    let start = Instant::now();
    let ans1 = problem1(get_input());
    let duration = start.elapsed();
    println!("Problem 1 solution: {} [{}s]", ans1, duration.as_secs_f64());

    let start = Instant::now();
    let ans2 = problem2(get_input());
    let duration = start.elapsed();
    println!("Problem 2 solution: {} [{}s]", ans2, duration.as_secs_f64());
}

// PARSE

#[derive(Debug)]
struct SpringRow<'a> {
    springs: Vec<char>,
    spans: Vec<usize>,
    cache: RefCell<HashMap<(usize, &'a [usize], usize), u64>>,
}

impl<'a> From<&str> for SpringRow<'a> {
    fn from(s: &str) -> Self {
        let (springs_s, spans_s) = s.split_once(' ').unwrap();
        SpringRow {
            springs: springs_s.chars().collect(),
            spans: spans_s.split(',').map(|x| x.parse().unwrap()).collect(),
            cache: RefCell::new(HashMap::new()),
        }
    }
}

impl<'a> SpringRow<'a> {
    fn possible_arrangements(&self) -> u64 {
        let mut permutations = vec![self.springs.clone()];
        while let Some(cur) = permutations.iter().position(|perm| perm.contains(&'?')) {
            let perm = permutations.remove(cur);
            let rep_pos = perm.iter().position(|c| *c == '?').unwrap();

            let mut new_elem = perm.clone();

            new_elem[rep_pos] = '.';
            permutations.push(new_elem.clone());
            new_elem[rep_pos] = '#';
            permutations.push(new_elem);
        }

        permutations
            .iter()
            .filter_map(|perm| {
                let groups: Vec<usize> = perm
                    .iter()
                    .group_by(|c| **c == '#')
                    .into_iter()
                    .filter_map(|(key, group)| if key { Some(group.count()) } else { None })
                    .collect();
                if groups == self.spans {
                    Some(())
                } else {
                    None
                }
            })
            .count() as u64
    }
    fn memoize_arrangements_r(
        &self,
        idx: usize,
        spans_left: &'a [usize],
        cur_span_left: usize,
    ) -> u64 {
        if let Some(value) = self.cache.borrow().get(&(idx, spans_left, cur_span_left)) {
            return *value;
        }
        let result = self.possible_arrangements_r(idx, spans_left, cur_span_left);
        self.cache
            .borrow_mut()
            .insert((idx, spans_left, cur_span_left), result);
        result
    }
    fn possible_arrangements_r(
        &self,
        idx: usize,
        spans_left: &'a [usize],
        cur_span_left: usize,
    ) -> u64 {
        if cur_span_left == 0 && spans_left.len() == 1 {
            // everything is consumed. if we're at the end, we good
            // if we have stuff left, as long as it doesn't contain a fixed broken spring we good
            if idx == self.springs.len() || !self.springs[idx..].contains(&'#') {
                // println!("{:?} valid", self.springs);
                return 1;
            }
            // something invalid
            return 0;
        }
        if idx >= self.springs.len() {
            // reached the end of the string without consuming all our spans
            // println!(
            //     "{:?} invalid idx: {} spans_left: {:?}",
            //     self.springs, idx, spans_left
            // );
            return 0;
        }
        // println!(
        //     "call idx: {} str: {:?} spans_left: {:?} cur_span_left: {}",
        //     idx,
        //     &self.springs[idx..],
        //     spans_left,
        //     cur_span_left
        // );
        return if cur_span_left == 0 {
            match self.springs[idx] {
                '#' => 0, // a broken spring after we have completed our span is invalid
                // a space or a ? (becomes a space) starts a new span
                '.' | '?' => self.memoize_arrangements_r(idx + 1, &spans_left[1..], spans_left[1]),
                _ => panic!("invalid spring"),
            }
        } else {
            // we are mid span
            match self.springs[idx] {
                // cur_span_left != 0, so we still have work to do with broken springs
                '#' => self.memoize_arrangements_r(idx + 1, spans_left, cur_span_left - 1),
                '.' => {
                    if spans_left[0] == cur_span_left {
                        self.memoize_arrangements_r(idx + 1, spans_left, cur_span_left)
                    } else {
                        0
                    }
                }
                '?' => {
                    if spans_left[0] == cur_span_left {
                        self.memoize_arrangements_r(idx + 1, spans_left, cur_span_left - 1)
                            + self.memoize_arrangements_r(idx + 1, spans_left, cur_span_left)
                    } else {
                        self.memoize_arrangements_r(idx + 1, spans_left, cur_span_left - 1)
                    }
                }
                _ => panic!("invalid spring"),
            }
        };
    }
}

// PROBLEM 1 solution

fn problem1<T: BufRead>(input: Lines<T>) -> u64 {
    let rows: Vec<SpringRow> = input
        .map(|row| SpringRow::from(row.unwrap().as_str()))
        .collect();

    rows.iter()
        .map(|row| row.possible_arrangements_r(0, &row.spans, row.spans[0]))
        .sum()
}

// PROBLEM 2 solution
fn problem2<T: BufRead>(input: Lines<T>) -> u64 {
    let mut rows: Vec<SpringRow> = input
        .map(|row| SpringRow::from(row.unwrap().as_str()))
        .collect();

    for row in &mut rows {
        let orig_slice = &row.springs.clone();
        for _i in 0..4 {
            row.springs.push('?');
            row.springs.extend_from_slice(orig_slice);
        }

        let orig_span_slice = &row.spans.clone();
        for _i in 0..4 {
            row.spans.extend_from_slice(orig_span_slice);
        }
    }
    rows.iter()
        .map(|row| row.possible_arrangements_r(0, &row.spans, row.spans[0]))
        .sum()
}

#[cfg(test)]
mod tests {
    use crate::*;
    use std::io::Cursor;

    const EXAMPLE: &str = &"???.### 1,1,3
.??..??...?##. 1,1,3
?#?#?#?#?#?#?#? 1,3,1,6
????.#...#... 4,1,1
????.######..#####. 1,6,5
?###???????? 3,2,1";

    #[test]
    fn problem1_example() {
        let c = Cursor::new(EXAMPLE);
        assert_eq!(problem1(c.lines()), 21);
    }

    #[test]
    fn problem2_example() {
        let c = Cursor::new(EXAMPLE);
        assert_eq!(problem2(c.lines()), 525152);
    }
}