import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as anim
from matplotlib.colors import ListedColormap

np.random.seed(0)

# Цвета состояний
colors = [
    (139/255, 69/255, 19/255),    # 0 - Земля (коричневый)
    (144/255, 238/255, 144/255),  # 1 - Редкая трава (светло-зелёный)
    (34/255, 139/255, 34/255),    # 2 - Обычная трава (зелёный)
    (0/255, 100/255, 0/255),      # 3 - Густая трава (тёмно-зелёный)
    (0/255, 0/255, 255/255),      # 4 - Вода (синий)
    (169/255, 169/255, 169/255),  # 5 - Камень (серый)
]
cmap = ListedColormap(colors)

# Веса состояний
state_weights = {
    0: {  # Земля
        0:  0.001,  # Земля ~ Земля
        1:  0.3,    # Земля ~ Редкая трава
        2:  0.6,    # Земля ~ Обычная трава
        3:  0.9,    # Земля ~ Густая трава
        4:  1.5,    # Земля ~ Вода
        5: -0.5,    # Земля ~ Камень
    },
    1: {  # Редкая трава
        0:  0.001,  # Редкая трава ~ Земля
        1:  0.15,   # Редкая трава ~ Редкая трава
        2:  0.3,    # Редкая трава ~ Обычная трава
        3:  0.45,   # Редкая трава ~ Густая трава
        4:  0.75,   # Редкая трава ~ Вода
        5: -1.0,    # Редкая трава ~ Камень
    },
    2: {  # Обычная трава
        0:  0.001,  # Обычная трава ~ Земля
        1:  0.1,    # Обычная трава ~ Редкая трава
        2:  0.2,    # Обычная трава ~ Обычная трава
        3:  0.3,    # Обычная трава ~ Густая трава
        4:  0.5,    # Обычная трава ~ Вода
        5: -1.5,    # Обычная трава ~ Камень
    },
}

# Размер сетки
grid_size = 100

# Начальная сетка: вся земля
grid = np.zeros((grid_size, grid_size), dtype=int)

# Генерация воды и камней на сетке
for _ in range(np.random.randint(2, 6)):  # Случайные области воды
    x0, y0 = np.random.randint(0, grid_size, size=2)
    w, h = np.random.randint(2, 21, size=2)
    for i in range(w):
        for j in range(h):
            x, y = (x0 + i) % grid_size, (y0 + j) % grid_size
            grid[x, y] = 4  # Вода
for _ in range(np.random.randint(5, 21)):  # Случайные области камней
    x0, y0 = np.random.randint(0, grid_size, size=2)
    w, h = np.random.randint(2, 5, size=2)
    for i in range(w):
        for j in range(h):
            x, y = (x0 + i) % grid_size, (y0 + j) % grid_size
            grid[x, y] = 5  # Камень

# Генерация матрицы весов
def generate_weight_matrix(radius):
    size = 2 * radius + 1
    weight_matrix = np.zeros((size, size), dtype=float)
    for x in range(-radius, radius + 1):
        for y in range(-radius, radius + 1):
            weight_matrix[x + radius, y + radius] = (radius + 1 - abs(x)) * (radius + 1 - abs(y))
    return weight_matrix

# Нормализация матрицы весов
def normalize_matrix(matrix):
    return matrix / matrix.sum()

# Вычисление влияния для ячейки (x, y) на основе её соседей
def calculate_influence_for_cell(grid, weight_matrix, x, y):
    influence = 0
    height, width = grid.shape
    
    weight_height, weight_width = weight_matrix.shape
    
    half_weight_height = weight_height // 2
    half_weight_width = weight_width // 2

    xy_state = grid[x, y]
    
    for i in range(weight_height):
        for j in range(weight_width):
            dx = i - half_weight_height
            dy = j - half_weight_width
            
            nx = (x + dx) % height
            ny = (y + dy) % width
            
            nxny_state = grid[nx, ny]
            
            influence += state_weights[xy_state][nxny_state] * weight_matrix[i, j]
    
    return influence

# Обработка подгруппы ячеек
def process_cells(cells, grid, new_grid):
    for x, y in cells:
        if grid[x, y] >= 0 and grid[x, y] < 3:
            growth_probability = calculate_influence_for_cell(grid, weight_matrix, x, y)
            if np.random.rand() < growth_probability:
                new_grid[x, y] += 1

# Генерация матрицы весов
radius = 2
weight_matrix = normalize_matrix(generate_weight_matrix(radius))

# Подготовка ячеек для обработки
all_cells = [(x, y) for x in range(grid_size) for y in range(grid_size)]

# Количество подгрупп
num_chunks = 20

# Перемешивание и разбиение на подгруппы
def shuffle_and_split_cells(all_cells, num_chunks):
    np.random.shuffle(all_cells)
    return np.array_split(all_cells, num_chunks)

# Изначальное разбиение
chunks = shuffle_and_split_cells(all_cells, num_chunks)
chunk_index = 0
new_grid = grid.copy()

# Анимация
fig, ax = plt.subplots(figsize=(6, 6))
im = ax.imshow(grid, cmap=cmap, vmin=0, vmax=5)
ax.axis("off")

def update(frame):
    global grid, chunk_index, new_grid, chunks

    process_cells(chunks[chunk_index], grid, new_grid)
    chunk_index += 1
    
    if chunk_index >= num_chunks:
        grid[:] = new_grid
        chunks = shuffle_and_split_cells(all_cells, num_chunks)
        chunk_index = 0

    im.set_array(new_grid)
    return [im]

ani = anim.FuncAnimation(fig, update, frames=100, interval=10, blit=True)

plt.show()