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

# In[67]:


def next_generation(fs):
    fs = {k-1:v for k, v in fs.items()}
    fs[6] += fs[-1]
    fs[8] = fs[-1]
    fs[-1] = 0
    return fs

def nth_generation(fs, n):
    for _ in range(n):
        fs = next_generation(fs)
    return fs

def fish(l):
    fs = {i:0 for i in range(9)}
    for d in l:
        fs[d] = fs[d] + 1
    return fs


# In[71]:


fs = fish([3,4,3,1,2])
fs = nth_generation(fs, 80)
sum(fs.values())


# In[74]:


import import_ipynb
import helper
s = helper.read_file('2021_6.txt')


# In[75]:


fs = map(int, s.rstrip().split(','))
fs = fish(fs)
fs = nth_generation(fs, 80)
sum(fs.values())


# In[76]:


fs = map(int, s.rstrip().split(','))
fs = fish(fs)
fs = nth_generation(fs, 256)
sum(fs.values())


# In[ ]: