Обнови код, что нужно сделать:
Добавить разные типы юнитов.
Сделать более сложное поведение солдат, например, искать врагов, двигаться к ним.
Добавить базу и ограничение на создание юнитов.
Добавить интерфейс выбора действий.
Оптимизировать обновления для большого поля.
Добавить искусственный интеллект для управления одной из команд.
```python
import pygame
import sys
import random

# Настройки
CELL_SIZE = 15  # размер клетки в пикселях
GRID_WIDTH = 40
GRID_HEIGHT = 40
SCREEN_WIDTH = CELL_SIZE * GRID_WIDTH
SCREEN_HEIGHT = CELL_SIZE * GRID_HEIGHT

FPS = 10  # скорость обновления игры

# Цвета
WHITE = (255, 255, 255)
GRAY = (200, 200, 200)
RED = (220, 50, 50)
BLUE = (50, 50, 220)
BLACK = (0, 0, 0)

# Команды
TEAM_NONE = 0
TEAM_RED = 1
TEAM_BLUE = 2

class Soldier:
    def __init__(self, x, y, team):
        self.x = x
        self.y = y
        self.team = team
        self.hp = 3  # здоровье

    def is_alive(self):
        return self.hp > 0

    def move(self, grid):
        # Пример простого передвижения: случайный шаг в соседнюю клетку, если там пусто
        directions = [(-1,0),(1,0),(0,-1),(0,1)]
        random.shuffle(directions)
        for dx, dy in directions:
            nx, ny = self.x + dx, self.y + dy
            if 0 <= nx < GRID_WIDTH and 0 <= ny < GRID_HEIGHT and grid[ny][nx] is None:
                grid[self.y][self.x] = None
                self.x, self.y = nx, ny
                grid[self.y][self.x] = self
                break

    def attack(self, grid):
        # Атакуем врагов в соседних клетках
        directions = [(-1,0),(1,0),(0,-1),(0,1)]
        for dx, dy in directions:
            nx, ny = self.x + dx, self.y + dy
            if 0 <= nx < GRID_WIDTH and 0 <= ny < GRID_HEIGHT:
                target = grid[ny][nx]
                if target is not None and target.team != self.team:
                    target.hp -= 1
                    # Можно добавить эффект или звук
                    if target.hp <= 0:
                        grid[ny][nx] = None
                    break  # атакуем одного врага за ход

def draw_grid(screen):
    for x in range(0, SCREEN_WIDTH, CELL_SIZE):
        pygame.draw.line(screen, GRAY, (x,0), (x, SCREEN_HEIGHT))
    for y in range(0, SCREEN_HEIGHT, CELL_SIZE):
        pygame.draw.line(screen, GRAY, (0,y), (SCREEN_WIDTH, y))

def draw_soldiers(screen, grid):
    for row in grid:
        for soldier in row:
            if soldier is not None:
                color = RED if soldier.team == TEAM_RED else BLUE
                rect = pygame.Rect(soldier.x*CELL_SIZE, soldier.y*CELL_SIZE, CELL_SIZE, CELL_SIZE)
                pygame.draw.rect(screen, color, rect)
                # Можно нарисовать HP
                font = pygame.font.SysFont(None, 16)
                hp_text = font.render(str(soldier.hp), True, WHITE)
                screen.blit(hp_text, (soldier.x*CELL_SIZE+2, soldier.y*CELL_SIZE+2))

def main():
    pygame.init()
    screen = pygame.display.set_mode((SCREEN_WIDTH, SCREEN_HEIGHT))
    pygame.display.set_caption("Клеточная битва")
    clock = pygame.time.Clock()

    grid = [[None for _ in range(GRID_WIDTH)] for _ in range(GRID_HEIGHT)]

    current_team = TEAM_RED  # команда, которой мы ставим солдат

    font = pygame.font.SysFont(None, 24)
    instructions = font.render("LMB - создать красного, RMB - создать синего", True, BLACK)

    while True:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                sys.exit()

            elif event.type == pygame.MOUSEBUTTONDOWN:
                mx, my = pygame.mouse.get_pos()
                grid_x = mx // CELL_SIZE
                grid_y = my // CELL_SIZE

                if grid[grid_y][grid_x] is None:
                    if event.button == 1:
                        # ЛКМ - красный солдат
                        soldier = Soldier(grid_x, grid_y, TEAM_RED)
                        grid[grid_y][grid_x] = soldier
                    elif event.button == 3:
                        # ПКМ - синий солдат
                        soldier = Soldier(grid_x, grid_y, TEAM_BLUE)
                        grid[grid_y][grid_x] = soldier

        # Обновляем состояние всех солдат
        # Чтобы избежать проблем при изменении массива во время обхода,
        # сначала собираем всех солдат в список
        soldiers = [soldier for row in grid for soldier in row if soldier is not None]

        # Сначала солдаты ходят
        for soldier in soldiers:
            if soldier.is_alive():
                soldier.move(grid)

        # Затем солдаты атакуют
        for soldier in soldiers:
            if soldier.is_alive():
                soldier.attack(grid)

        # Рисуем всё
        screen.fill(WHITE)
        draw_grid(screen)
        draw_soldiers(screen, grid)
        screen.blit(instructions, (5, 5))
        pygame.display.flip()

        clock.tick(FPS)

if __name__ == "__main__":
    main()
```