#!/usr/bin/env python3

import sys, re
import numpy as np
from scipy.ndimage import rotate
from collections import defaultdict


# copied from stackoverflow
def groups(data, func):
    grouping = defaultdict(list)
    for y in range(len(data)):
        for x in range(len(data[y])):
            grouping[func(x, y)].append(data[y][x])
    return list(map(grouping.get, sorted(grouping)))


def a_cross(puzzle, x, y):
    if puzzle[y][x] != 'A': return False
    elif x == 0 or y == 0 or x == len(puzzle[y])-1 or y == len(puzzle)-1: return False
    elif ((puzzle[y-1][x-1] == 'M' and puzzle[y+1][x+1] == 'S') or (puzzle[y-1][x-1] == 'S' and puzzle[y+1][x+1] == 'M')) and ((puzzle[y+1][x-1] == 'M' and puzzle[y-1][x+1] == 'S') or (puzzle[y+1][x-1] == 'S' and puzzle[y-1][x+1] == 'M')): 
        return True
    else: return False


if __name__ == '__main__':
    puzzle = [line.strip('\n') for line in open(sys.argv[1])]
    words = re.compile(r'(?=XMAS|SAMX)')
    res1 = len(words.findall('D'.join(puzzle)))
    npuzzle = np.array([list(p) for p in puzzle])
    # rows = groups(puzzle, lambda x, y: y)
    #cols 
    res1 += len(words.findall('.'.join(['.'.join([''.join(l) for l in groups(npuzzle, lambda x, y: x)])])))
    #fdiag
    res1 += len(words.findall('.'.join(['.'.join([''.join(l) for l in groups(npuzzle, lambda x, y: x + y)])])))
    #bdiag
    res1 += len(words.findall('.'.join(['.'.join([''.join(l) for l in groups(npuzzle, lambda x, y: x - y)])])))

    res2 = 0
    npuzzle = np.array([list(p) for p in puzzle])
    for y in range(len(npuzzle)):
        for x in range(len(npuzzle[y])):
            res2 += int(a_cross(npuzzle, x, y))

    # challenge 1
    print(f"challenge 1:\n{res1}\n")

    # challenge 2
    print(f"challenge 2:\n{res2}")