from itertools import chain
import re
filename = 'sample'
input = [list(i.rstrip()) for i in open(filename).readlines()]

def generate_horizontal_cross_sections(search):
    return search

def generate_vertical_cross_sections(search):
    output = []
    for row in range(len(input[0])):
        current = []
        for line in search:
            current.append(line[row])

        output.append(current)

    return output

def generate_diagonal_cross_sections(search):
    modulo = len(search[0])
    concat = ''.join(chain(*search))
    output = []
    for offset in range(len(concat)):
        current = []
        index = offset
        while index % modulo != 0 and index < len(concat):
            current.append(concat[index])
            index += modulo - 1
        try:
            current.append(concat[index])
        except IndexError:
            pass
        output.append(current)

    output = [i for i in output if len(i) >= 4]    
    return output

def generate_antidiagonal_cross_sections(search):
    # start from the points on the right and bottom edges
    # of the array, and then work northwest diagonally
    indexes_to_start = []
    modulo = len(search[0])
    concat = ''.join(chain(*search))

    index = modulo - 1
    while index < len(concat):
        indexes_to_start.append(index)
        index += modulo

    index -= modulo
    while index % modulo != 0:
        indexes_to_start.append(index)
        index -= 1

    indexes_to_start.append(index)

    # work up northwest
    output = []

    for index in indexes_to_start:
        current = []
        while index % modulo != 0 and index > 0:
            current.append(concat[index])
            index -= modulo + 1

        if index > 0: current.append(concat[index])

        output.append(current)

    output = [i for i in output if len(i) >= 4]    
    return output

search_space = [
    *generate_horizontal_cross_sections(input),
    *generate_vertical_cross_sections(input),
    *generate_diagonal_cross_sections(input),
    *generate_antidiagonal_cross_sections(input),
]

xmas = re.compile('XMAS')

total = 0
for item in search_space:
    joined = ''.join(item)
    print(joined)
    results = xmas.findall(joined)
    print(results)
    total += len(results)
    joined = ''.join(reversed(item))
    print(joined)
    results = xmas.findall(joined)
    print(results)
    total += len(results)

print(results)