#!/usr/bin/env python
# coding: utf-8

# In[91]:


import numpy as np

def new_grid(n):
    return np.zeros((n,n), np.int64)

def parse_points(seg):
    start, end = seg.rstrip().split(' -> ')
    x1, y1 = map(int, start.split(','))
    x2, y2 = map(int, end.split(','))
    return x1, y1, x2, y2

def slice_range(x1, y1, x2, y2):
    if y1 > y2:
        y = range(y1, y2-1, -1)
    else:
        y = range(y1, y2+1)
        
    if x1 > x2:
        x = range(x1, x2-1, -1)
    else:
        x = range(x1, x2+1)
    
    return (y, x)
    

def grid_slice(grid, seg, diagonal):
    x1, y1, x2, y2 = parse_points(seg)
    if y1 == y2 or x1 == x2 or diagonal:
        return slice_range(x1, y1, x2, y2)
    else:
        return None

def overlapping(grid, f, diagonal=False):
    for l in f.splitlines():
        s = grid_slice(grid, l, diagonal)
        if s is not None:
            grid[s] += 1
    return np.sum(grid >= 2)


# In[93]:


g = new_grid(10)
test = """0,9 -> 5,9
8,0 -> 0,8
9,4 -> 3,4
2,2 -> 2,1
7,0 -> 7,4
6,4 -> 2,0
0,9 -> 2,9
3,4 -> 1,4
0,0 -> 8,8
5,5 -> 8,2"""
assert overlapping(g, test) == 5
g = new_grid(10)
assert overlapping(g, test, diagonal=True) == 12


# In[50]:


import import_ipynb
import helper
f = helper.read_file('2021_5.txt')
g = new_grid(1000)
overlapping(g, f)


# In[94]:


g = new_grid(1000)
overlapping(g, f, diagonal=True)


# In[ ]: