use std::collections::HashSet;

use aoc_runner_derive::{aoc as aoc_run, aoc_generator};

use ndarray::prelude::*;

type Forest = Array2<i8>;

#[aoc_generator(day8)]
fn parse_input_day1(input: &str) -> Forest {
    let nrows = input.lines().map(|_| ()).collect::<Vec<_>>().len();
    let ncols = input.lines().next().unwrap().len();
    let shape = (nrows, ncols);
    let mut arr = Array2::zeros(shape);
    for (row, line) in input.lines().enumerate() {
        for (col, tree) in line.chars().enumerate() {
            arr[[row, col]] = tree.to_digit(10).unwrap() as i8;
        }
    }
    arr
}

#[aoc_run(day8, part1)]
fn part1(input: &Forest) -> usize {
    let mut viz = HashSet::new();
    for (ridx, row) in input.rows().into_iter().enumerate() {
        let mut max = -1;
        for (cidx, &tree) in row.as_slice().unwrap().iter().enumerate() {
            if tree > max {
                let addr = (ridx, cidx);
                viz.insert(addr);
                max = tree;
            }
        }
        max = -1;
        for (cidx, &tree) in row.as_slice().unwrap().iter().enumerate().rev() {
            if tree > max {
                viz.insert((ridx, cidx));
                max = tree;
            }
        }
    }

    for (cidx, col) in input.columns().into_iter().enumerate() {
        let mut max = -1;
        for (ridx, &tree) in col
            .as_standard_layout()
            .as_slice()
            .unwrap()
            .iter()
            .enumerate()
        {
            if tree > max {
                viz.insert((ridx, cidx));
                max = tree;
            }
        }
        max = -1;
        for (ridx, &tree) in col
            .as_standard_layout()
            .as_slice()
            .unwrap()
            .iter()
            .enumerate()
            .rev()
        {
            if tree > max {
                viz.insert((ridx, cidx));
                max = tree;
            }
        }
    }

    viz.len()
}

#[aoc_run(day8, part2)]
fn part2(input: &Forest) -> usize {
    let nrows = input.nrows();
    let ncols = input.ncols();
    let mut max = 0;
    for row in 0..nrows {
        for col in 0..ncols {
            let (e, w) = check_span(col, &input.row(row));
            let (n, s) = check_span(row, &input.column(col));
            max = (e * w * n * s).max(max);
        }
    }
    max
}

fn check_span(addr: usize, span: &ArrayView1<i8>) -> (usize, usize) {
    let height = span[addr];
    let (mut a, mut b) = (0, 0);
    //
    let neg = span.slice(s![0..addr]);
    let pos = span.slice(s![addr..]);
    for e in neg.iter().rev() {
        a += 1;
        if *e >= height {
            break;
        }
    }

    for e in pos.iter().skip(1) {
        b += 1;
        if *e >= height {
            break;
        }
    }

    (a, b)
}