import matrix
import shared
from pprint import pprint
from dataclasses import dataclass
from typing import Set, Tuple
from functools import cached_property
from mpl_toolkits.mplot3d import Axes3D
import numpy as np
import matplotlib.pyplot as plt




@dataclass
class Cube:
    x: int
    y: int
    z: int
    neighbor_coords: Set[Tuple[int,int,int]] = None

    @cached_property
    def xyz(self):
        return (self.x,self.y,self.z)
    
    def set_neighbor_coords(self):
        self.neighbor_coords = get_neighbors(self.x,self.y,self.z)

def get_neighbors(x,y,z):
        # Generate the six neighbor_coords
        # Look at a die 6 to the left, 2 on top, one on right
        offsets = (
           (1,0,0), # 1 
           (0,0,1), # 2 
           (0,-1,0),# 3 
           (0,1,0), # 4 
           (0,0,-1),# 5 
           (-1,0,0) # 6 
        )
        return set([(x+o[0], y+o[1], z+o[2]) for o in offsets])

def get_flat_neighbors(x,y,z):
        # Generate the six neighbor_coords
        # Look at a die 6 to the left, 2 on top, one on right 
        # ignore top and bottom
        offsets = (
           (1,0,0), # 1 
           #(0,0,1), # 2 
           (0,-1,0),# 3 
           (0,1,0), # 4 
           #(0,0,-1),# 5 
           (-1,0,0) # 6 
        )
        return set([(x+o[0], y+o[1], z+o[2]) for o in offsets])
        
def get_flat_neighbors_from(x,y,z, f):
    return get_flat_neighbors(x,y,z).intersection(f)


# @shared.profile
def part1(rows):
    cubes = {}
    for idx, (x,y,z) in enumerate(rows):
        cubes[(x,y,z)] = {'n':get_neighbors(x,y,z), 'not_me': tuple(map(tuple,rows[:idx] + rows[idx+1:]))}
    potential = len(cubes) * 6
    for coords, cube in cubes.items():
        for other in cube['not_me']:
            if other in cube['n']:
                potential -= 1
    print(potential)



def surface_area(rows):
    cubes = {}
    maxX,maxY,maxZ = 0,0,0
    for idx, (x,y,z) in enumerate(rows):
        cubes[(x,y,z)] = {'n':get_neighbors(x,y,z), 'not_me': set(map(tuple,rows[:idx] + rows[idx+1:]))}
        maxX = max(maxX,x)
        maxY = max(maxY,y)
        maxZ = max(maxZ,z)
    _cubes = frozenset(cubes.keys())

    potential = len(cubes) * 6
    for coords, cube in cubes.items():
        for other in cube['not_me']:
            if other in cube['n']:
                potential -= 1
    return potential, maxX,maxY,maxZ, cubes, _cubes


def part2(rows):
    potential, maxX, maxY,maxZ,cubes,_cubes = surface_area(rows)


    air = set()
    for x in range(0,maxX+1):
        for y in range(0,maxY+1):
            for z in range(0,maxZ+1):
                if (x,y,z) not in _cubes:
                    air.add((x,y,z))
    air -= _cubes # Remove all lava from air
    all_air_count = len(air)
    all_lava_count = len(_cubes)
    possible_count = (maxX+1)*(maxY+1)*(maxZ+1)

    print("all_air: ", all_air_count)
    print("all_lava:", all_lava_count)
    print("possible:", possible_count)
    print(maxZ,maxY,maxX)
    print()
    #shared.render_cubes(maxX+2,maxY+2,maxZ+1, [x for x in _cubes])

    air_nx = {}
    for a in air:
        neighbors = get_neighbors(*a)
        air_nx[a] = {
                'lava':neighbors & _cubes or None,
                'air': neighbors & air or None
                }

    # loop row by row
    inside = []
    for z in range(0,maxZ+2):
        seen = set()
        seen.add((0,0,z))
        mx = matrix.matrix_of_size(maxX+2,maxY+2)

        for x,y in matrix.spiral_generator(maxX+2, maxY+2):
            xyz = x,y,z
            mx[y][x] = "#"
            if xyz in _cubes:
                continue
            ns = get_flat_neighbors(x,y,z)
            for neigh in ns:
                if neigh in seen:
                    seen.add(xyz)

        this_level = set(x for x in air_in_row(air, z))
        for (x,y,z) in seen:
            mx[y][x] = 0

        for x,y,z in this_level:
            if (x,y,z) not in seen:
                mx[y][x] = "o"
        print(matrix.ppmx(mx, pad=False,space=False))
        print()

        #print(len(this_level), len(seen))
        #inside.extend([x for x in this_level if x not in seen])
        #print(seen)
        #shared.render_cubes(maxX+2,maxY+2,maxZ+1, [x for x in seen])
        #shared.render_cubes(maxX+2,maxY+2,maxZ+1, [x for x in this_level-seen])
    ##print()
    ##print(inside)
    ##print()
    #actually_inside = []
    #for i in inside:
    #    ns = get_neighbors(*i)
    #    # check for surrounded 100% by rock
    #    rock_count = 0
    #    for n in ns:
    #        if n in _cubes:
    #            rock_count +=1
    #    if rock_count == 6:
    #        print("in rock")
    #        actually_inside.append(i)
    #        continue
    #    #check for surrounded 100% by air
    #    air_count = 0
    #    for n in ns:
    #        if n in air:
    #            air_count +=1
    #            print("in air")
    #            break
    #        actually_inside.append(i)
    ##print(actually_inside)
    #tot, _, _, _, _, _ = surface_area(actually_inside)
    #print(tot)
    #print(potential - tot)

def air_in_row(air, z):
    x = []
    for a in air:
        if a[2] == z:
            x.append(a)
    return x



def main():
    rows = [list(map(int,row.split(","))) for row in shared.load_rows(18)]
    with shared.elapsed_timer() as elapsed:
        part1(rows)
        print("🕒", elapsed())
    #rows = [map(int,row.split(",")) for row in shared.load_rows(18)]
    with shared.elapsed_timer() as elapsed:
        part2(rows)
        print("🕒", elapsed())


if __name__ == "__main__":
    main()