program advent04
    implicit none

    character, dimension(146, 146) :: grid
    integer :: row, col, part1, part2
    real :: start, end

    call cpu_time(start)
    grid = load("data/04.txt", 140, 140)
    part1 = 0
    part2 = 0


    do col = 4, 143
        do row = 4, 143
            if (grid(row, col) == 'X') then
                part1 = part1 + count_xmas(row, col)
            end if

            if (grid(row, col) == 'A') then
                part2 = part2 + count_mas(row, col)
            end if
        end do
    end do

    call cpu_time(end)

    print *, "Part 1:", part1
    print *, "Part 2:", part2
    print *, "Execution time:", end - start

    contains

    function load(path, n_rows, n_cols) result(grid)
        implicit none

        character(*), intent(in) :: path
        integer, intent(in) :: n_rows, n_cols
        integer :: handle, row, col
        character, dimension(:, :), allocatable :: grid
        character(n_cols) :: line

        allocate(grid(n_rows + 6, n_cols + 6))
        grid = '.'

        open(newunit=handle, file=path, status="old", action="read")
        do row = 4, n_rows + 3
            read(handle, *) line
            do col = 1, n_cols
                grid(row, col + 3) = line(col:col)
            end do
        end do

        close(handle)
    end function load

    integer function count_xmas(row, col) result(count)
        implicit none

        integer, intent(in) :: row, col
        integer :: i
        integer(8) :: prod
        integer(8), dimension(8) :: primes
        character, dimension(7, 7) :: test_grid, window
        integer(8), dimension(7, 7) :: prime_mask, matches, matches_prime

        test_grid = reshape( &
            [&
                'S', '.', '.', 'S', '.', '.', 'S', &
                '.', 'A', '.', 'A', '.', 'A', '.', &
                '.', '.', 'M', 'M', 'M', '.', '.', &
                'S', 'A', 'M', 'X', 'M', 'A', 'S', &
                '.', '.', 'M', 'M', 'M', '.', '.', &
                '.', 'A', '.', 'A', '.', 'A', '.', &
                'S', '.', '.', 'S', '.', '.', 'S' &
            ], &
            shape(test_grid) &
        )

        primes = [2, 3, 5, 7, 11, 13, 17, 19]

        prime_mask = reshape( &
            [ &
                2,  1,  1,  3,  1,  1,  5, &
                1,  2,  1,  3,  1,  5,  1, &
                1,  1,  2,  3,  5,  1,  1, &
                19, 19, 19, 1,  7,  7,  7, &
                1,  1,  17, 13, 11, 1,  1, &
                1,  17,  1, 13, 1,  11, 1, &
                17,  1,  1, 13, 1,  1,  11 &
            ], &
            shape(prime_mask) &
        )

        window = grid(row - 3:row + 3, col - 3:col + 3)
        matches = logical_to_int64(window == test_grid)
        matches_prime = matches * prime_mask
        prod = product(zero_to_one(matches_prime))

        count = 0
        do i = 1, 8
            if (mod(prod, primes(i) ** 3) == 0) then
                count = count + 1
            end if
        end do
    end function count_xmas

    integer function count_mas(row, col) result(count)
        implicit none

        integer, intent(in) :: row, col
        integer :: i
        character, dimension(3, 3) :: window, t1, t2, t3, t4

        t1 = reshape( &
            [ &
                'M', '.', 'S', &
                '.', 'A', '.', &
                'M', '.', 'S' &
            ], &
            shape(t1) &
        )
        t2 = t1(3:1:-1, :) ! flip t1 top-to-bottom
        t3 = transpose(t1) ! swap t1 rows for columns
        t4 = t3(:, 3:1:-1) ! flip t3 lef-to-right

        window = grid(row - 1:row + 1, col - 1:col + 1)
        if ( &
            count_matches(window, t1) == 5 &
            .or. count_matches(window, t2) == 5 &
            .or. count_matches(window, t3) == 5 &
            .or. count_matches(window, t4) == 5 &
        ) then
            count = 1
        else
            count = 0
        end if
    end function count_mas

    integer function count_matches(a1, a2) result(matches)
        implicit none

        character, dimension(:, :) :: a1, a2

        matches = count(a1 == a2)
    end function count_matches

    elemental integer(8) function logical_to_int64(b) result(i)
        implicit none

        logical, intent(in) :: b

        if (b) then
            i = 1
        else
            i = 0
        end if
    end function logical_to_int64

    elemental integer(8) function zero_to_one(x) result(y)
        implicit none

        integer(8), intent(in) :: x

        if (x == 0) then
            y = 1
        else
            y = x
        end if
    end function zero_to_one
end program advent04