import pygame
import random
import math

# ===== CONFIG =====
WIDTH, HEIGHT = 1000, 700
MAP_WIDTH, MAP_HEIGHT = 120, 120
TILE_SIZE = 8
MIN_ZOOM, MAX_ZOOM = 0.5, 18
CIV_COUNT = 3

# ===== INIT =====
pygame.init()
screen = pygame.display.set_mode((WIDTH, HEIGHT))
clock = pygame.time.Clock()
pygame.display.set_caption("Enhanced Civilization Sim")

FPS = 60

# ===== CAMERA & ZOOM =====
camera_x, camera_y = 0, 0
zoom = 1.0

# ===== TERRAIN GENERATION =====
def smooth_noise(w, h, scale=6):
    base = [[random.random() for _ in range(w // scale + 2)] for _ in range(h // scale + 2)]
    def interpolate(x, y):
        ix, fx = int(x), x - int(x)
        iy, fy = int(y), y - int(y)
        v1 = base[iy][ix]
        v2 = base[iy][ix + 1]
        v3 = base[iy + 1][ix]
        v4 = base[iy + 1][ix + 1]
        i1 = v1 * (1 - fx) + v2 * fx
        i2 = v3 * (1 - fx) + v4 * fx
        return i1 * (1 - fy) + i2 * fy
    return [[interpolate(x / scale, y / scale) for x in range(w)] for y in range(h)]

height_map = smooth_noise(MAP_WIDTH, MAP_HEIGHT, scale=8)

# ===== RESOURCE MAP =====
resource_map = [[0 for _ in range(MAP_WIDTH)] for _ in range(MAP_HEIGHT)]
def generate_resources():
    for y in range(MAP_HEIGHT):
        for x in range(MAP_WIDTH):
            h = height_map[y][x]
            if 0.2 < h < 0.6:
                resource_map[y][x] = 1 if random.random() < 0.15 else 0
            elif h < 0.25:
                resource_map[y][x] = 2 if random.random() < 0.5 else 0
generate_resources()

# ===== ANIMALS =====
animals = set()
for _ in range(300):
    while True:
        ax, ay = random.randint(0, MAP_WIDTH - 1), random.randint(0, MAP_HEIGHT - 1)
        if 0.3 < height_map[ay][ax] < 0.6:
            animals.add((ax, ay))
            break

# ===== BIOME & RESOURCE COLORS =====
def biome_color(h):
    if h < 0.2: return (0, 90, 130)
    elif h < 0.25: return (30, 144, 255)
    elif h < 0.35: return (230, 220, 170)
    elif h < 0.6: return (80, 160, 60)
    elif h < 0.8: return (34, 100, 34)
    else: return (100, 100, 100)

def resource_color(r):
    if r == 1: return (0, 110, 0)
    if r == 2: return (0, 180, 255)
    return None

# ===== DISCOVERED MAP =====
discovered_map = [[False for _ in range(MAP_WIDTH)] for _ in range(MAP_HEIGHT)]

# ===== COLOR DISTANCE FOR DIPLOMACY =====
def color_distance(c1, c2):
    return sum((a - b) ** 2 for a, b in zip(c1, c2)) ** 0.5

# ===== CIVILIZATION CLASS =====
class Civilization:
    def __init__(self, x, y, color, name):
        self.x, self.y = x, y
        self.color = color
        self.name = name
        self.population = 50
        self.claimed = set([(x, y)])
        self.units = [Unit(x, y, color) for _ in range(3)]
        self.expansion_timer = 0
        self.resources = {'wood': 10, 'water': 10, 'food': 15}
        self.mines = set()
        self.allies = set()
        self.enemies = set()

    def update(self, others):
        self.expansion_timer += 1
        self.diplomacy(others)
        if self.expansion_timer >= 60:
            self.expansion_timer = 0
            self.try_expand()
        for unit in self.units:
            unit.update()
            ux, uy = int(unit.x), int(unit.y)
            for dy in range(-3, 4):
                for dx in range(-3, 4):
                    nx, ny = ux + dx, uy + dy
                    if 0 <= nx < MAP_WIDTH and 0 <= ny < MAP_HEIGHT:
                        if dx*dx + dy*dy <= 9:
                            discovered_map[ny][nx] = True
        self.check_food()

    def diplomacy(self, others):
        for other in others:
            if other == self:
                continue
            dist = color_distance(self.color, other.color)
            if dist < 100:
                self.allies.add(other)
            elif dist > 200:
                self.enemies.add(other)

    def try_expand(self):
        if self.resources['wood'] < 2 or self.resources['water'] < 2:
            return
        frontier = list(self.claimed)
        random.shuffle(frontier)
        for (tx, ty) in frontier[:8]:
            for dx, dy in [(-1,0),(1,0),(0,-1),(0,1)]:
                nx, ny = tx + dx, ty + dy
                if 0 <= nx < MAP_WIDTH and 0 <= ny < MAP_HEIGHT:
                    if height_map[ny][nx] > 0.25 and (nx, ny) not in self.claimed:
                        res = resource_map[ny][nx]
                        if res in (1, 2) or random.random() < 0.1:
                            self.claimed.add((nx, ny))
                            if res == 1:
                                self.resources['wood'] += 2
                                resource_map[ny][nx] = 0
                            elif res == 2:
                                self.resources['water'] += 2
                            self.resources['food'] += 1
                            self.population += 1
                            if (nx, ny) in animals:
                                animals.remove((nx, ny))
                                self.resources['food'] += 5
                            if random.random() < 0.25:
                                self.units.append(Unit(nx, ny, self.color))
                            if random.random() < 0.1:
                                self.mines.add((nx, ny))
                            return

    def check_food(self):
        if self.resources['food'] >= self.population:
            if random.random() < 0.05:
                self.population += 100
                self.resources['food'] -= self.population // 2
        else:
            self.population -= 1
            if self.population < 1:
                self.population = 1

    def draw(self, surf):
        for (x, y) in self.claimed:
            px = (x * TILE_SIZE - camera_x) * zoom
            py = (y * TILE_SIZE - camera_y) * zoom
            s = TILE_SIZE * zoom
            claim_color = (*self.color[:3], 60)
            claim_surf = pygame.Surface((s, s), pygame.SRCALPHA)
            claim_surf.fill(claim_color)
            surf.blit(claim_surf, (px, py))
            pygame.draw.rect(surf, self.color, (px, py, s, s), width=1)
            # Inside the Civilization.draw method
            px = (self.x * TILE_SIZE - camera_x) * zoom
            py = (self.y * TILE_SIZE - camera_y) * zoom
            pygame.draw.circle(surf, (139, 69, 19), (int(px), int(py)), int(5 * zoom))

        for u in self.units:
            u.draw(surf)
        for mx, my in self.mines:
            mx_px = (mx * TILE_SIZE - camera_x) * zoom
            my_py = (my * TILE_SIZE - camera_y) * zoom
            pygame.draw.rect(surf, (80, 80, 80), (mx_px + 2, my_py + 2, TILE_SIZE * zoom - 4, TILE_SIZE * zoom - 4))

# ===== UNIT CLASS =====
class Unit:
    def __init__(self, x, y, color):
        self.x, self.y = x + random.uniform(-0.5, 0.5), y + random.uniform(-0.5, 0.5)
        self.color = color
        self.speed = 0.05

    def update(self):
        dx, dy = random.uniform(-1, 1), random.uniform(-1, 1)
        self.x += dx * self.speed
        self.y += dy * self.speed
        self.x = max(0, min(self.x, MAP_WIDTH - 1))
        self.y = max(0, min(self.y, MAP_HEIGHT - 1))

    def draw(self, surf):
        px = (self.x * TILE_SIZE - camera_x) * zoom
        py = (self.y * TILE_SIZE - camera_y) * zoom
        pygame.draw.circle(surf, self.color, (int(px), int(py)), int(2 * zoom))

# ===== DRAW ANIMALS =====
def draw_animals(surf):
    for (x, y) in animals:
        px = (x * TILE_SIZE - camera_x) * zoom
        py = (y * TILE_SIZE - camera_y) * zoom
        r = max(1, int(2 * zoom))
        pygame.draw.circle(surf, (139, 69, 19), (int(px), int(py)), r)

# ===== DRAW MAP =====
def draw_map():
    for y in range(MAP_HEIGHT):
        for x in range(MAP_WIDTH):
            if not discovered_map[y][x]:
                continue
            h = height_map[y][x]
            r = resource_map[y][x]
            px = (x * TILE_SIZE - camera_x) * zoom
            py = (y * TILE_SIZE - camera_y) * zoom
            s = TILE_SIZE * zoom
            pygame.draw.rect(screen, biome_color(h), (px, py, s, s))
            if r:
                pygame.draw.circle(screen, resource_color(r), (px + s//2, py + s//2), max(1, int(1.5 * zoom)))

# ===== CIVILIZATION SETUP =====
civilizations = []
for _ in range(5):
    while True:
        x, y = random.randint(0, MAP_WIDTH-1), random.randint(0, MAP_HEIGHT-1)
        if height_map[y][x] > 0.3:
            civ_color = tuple(random.randint(50, 255) for _ in range(3))
            civ = Civilization(x, y, civ_color, f"Civ{_}")
            civilizations.append(civ)
            break

# ===== MAIN LOOP =====
running = True
while running:
    clock.tick(FPS)
    for event in pygame.event.get():
        if event.type == pygame.QUIT: running = False
        elif event.type == pygame.MOUSEBUTTONDOWN:
            if event.button == 4:  # Scroll up
                mx, my = pygame.mouse.get_pos()
                world_x = camera_x + mx / zoom
                world_y = camera_y + my / zoom
                zoom = min(MAX_ZOOM, zoom * 1.1)
                camera_x = world_x - mx / zoom
                camera_y = world_y - my / zoom
            elif event.button == 5:  # Scroll down
                mx, my = pygame.mouse.get_pos()
                world_x = camera_x + mx / zoom
                world_y = camera_y + my / zoom
                zoom = max(MIN_ZOOM, zoom / 1.1)
                camera_x = world_x - mx / zoom
                camera_y = world_y - my / zoom


    keys = pygame.key.get_pressed()
    move_speed = 10 / zoom
    if keys[pygame.K_a]: camera_x -= move_speed
    if keys[pygame.K_d]: camera_x += move_speed
    if keys[pygame.K_w]: camera_y -= move_speed
    if keys[pygame.K_s]: camera_y += move_speed

    screen.fill((0, 0, 0))
    draw_map()
    draw_animals(screen)
    for civ in civilizations:
        civ.update(civilizations)
        civ.draw(screen)

    pygame.display.flip()

pygame.quit()