use std::collections::HashMap;

use itertools::Itertools;

#[cfg(test)]
use crate::util::read_file;
use crate::{day_solver::DaySolver, util::{Coord, Grid}};

pub struct Day8 {
    map: Grid<char>
}

impl Day8 {

    pub fn create(input: String) -> Self {
        
        Day8 {
            map: Grid::parse(input, None)
        }
    }

    fn get_antennas_per_frequency(&self) -> HashMap<char, Vec<Coord<usize>>> {
        let mut frequency_antennas: HashMap<char, Vec<Coord<usize>>> = HashMap::new();
        for y in 0..self.map.height() {
            for x in 0..self.map.width {
                let cur = self.map.get(x, y);
                let pos: Coord<usize> = Coord::new(x, y);
                if cur != &'.' {
                    if let Some(ant_pos) =  frequency_antennas.get_mut(cur) {
                        ant_pos.push(pos);
                    } else {
                        frequency_antennas.insert(cur.to_owned(), vec![pos]);
                    }
                }

            }
        }
        frequency_antennas
    }

    fn find_antinodes(&self, antennas: &Vec<Coord<usize>>) -> Vec<Coord<usize>> {

        let mut antinodes = Vec::new();

        for a in antennas {
            for b in antennas {
                if a == b {
                    continue;
                }
                let mut a_signed = a.as_signed();
                let dist = a_signed.sub(&b.as_signed());
                antinodes.push(a.to_owned());
                while let Ok(antinode) = TryInto::<Coord<usize>>::try_into(a_signed.add(&dist)) {
                    if self.map.in_bounds(antinode.x, antinode.y) {
                        antinodes.push(antinode);
                    } else {
                        break;
                    }
                    a_signed = antinode.as_signed()
                }
            }
        }

        antinodes
    }
}

impl DaySolver for Day8 {


    fn solve_part1(&mut self) -> String {
        
        let frequency_antennas = self.get_antennas_per_frequency();

        frequency_antennas.iter()
            .flat_map(|(_, antennas)| {
                antennas.iter().flat_map(|a|
                    antennas.iter()
                        .filter_map(|b| if b.to_owned() == a.to_owned() {
                            None
                        } else {
                            let a_signed = a.as_signed();
                            let dist = a_signed.sub(&b.as_signed());
                            Some(a_signed.add(&dist))
                        })
                    )
                })
            .filter_map::<Coord<usize>, _>(|c| c.try_into().ok())
            .filter(|c| self.map.in_bounds(c.x, c.y))
            .unique()
            // .map(|c| { println!("{:?}", c); c })
            .count().to_string()
    }

    fn solve_part2(&mut self) -> String {
        
        let frequency_antennas = self.get_antennas_per_frequency();

        let antinodes = frequency_antennas.iter()
            .flat_map(|(_, antennas)| self.find_antinodes(antennas))
            .unique()
            .collect::<Vec<Coord<usize>>>();

        // let mut full_map = self.map.clone();
        // antinodes.iter().for_each(|n| full_map.set(n.x, n.y, '#'));
        // full_map.print();

        antinodes
            .len().to_string()
    }
}

#[test]
fn test_part1() {
    let mut day = Day8::create(read_file("input/day08_example.txt"));
    assert_eq!("14", day.solve_part1());
}

#[test]
fn test_part1_basic() {
    let sample = "..........
..........
..........
....a.....
..........
.....a....
..........
..........
..........
..........";
    let mut day = Day8::create(sample.to_string());
    assert_eq!("2", day.solve_part1());
}

#[test]
fn test_part2() {
    let mut day = Day8::create(read_file("input/day08_example.txt"));
    assert_eq!("34", day.solve_part2());
}