import pygame
import numpy as np
from noise import pnoise3
import math
import random
from math import floor

# --- Constants ---
SCREEN_WIDTH = 1200
SCREEN_HEIGHT = 800
BACKGROUND_COLOR = (135, 206, 235)  # Sky Blue
GROUND_COLOR_LOW = (30, 100, 30)    # Darker forest green for lower terrain
GROUND_COLOR_HIGH = (100, 160, 60)  # Lighter green for higher terrain
GROUND_COLOR_EARTH = (139, 69, 19)  # Brown for lowest terrain
TANK_WIDTH = 40
TANK_HEIGHT = 20
TURRET_WIDTH = 5
TURRET_LENGTH = 25
GRAVITY = 0.2
EXPLOSION_DURATION = 30  # Frames
TANK_MOVE_SPEED = 2
MAX_FUEL = 100
NOISE_UPDATE_INTERVAL = 10  # Update noise every x frames
EXPLOSION_BUFFER = 10  # Extra pixels to update around explosion
MENU_BG_COLOR = (50, 50, 50, 200)  # Semi-transparent dark gray
BUTTON_COLOR = (100, 100, 100)     # Gray for buttons
BUTTON_HOVER_COLOR = (150, 150, 150)  # Lighter gray for hover
TEXT_COLOR = (255, 255, 255)       # White text
WIND_MIN = -0.1  # Minimum wind speed (negative = left)
WIND_MAX = 0.1   # Maximum wind speed (positive = right)
BONUS_ITEM_RADIUS = 10
BONUS_ITEM_POINTS = 50
BONUS_ITEM_COLOR = (255, 215, 0)  # Gold
SMOKE_DURATION = 120  # Frames for smoke effect

# --- Weapon Definitions ---
WEAPONS = [
    {
        "name": "Standard Shell",
        "explosion_radius": 30,
        "crater_multiplier": 1.0,
        "gravity_multiplier": 1.0,
        "wind_multiplier": 1.0,
        "weight": 1.0,
        "warheads": 1,
        "special": None,
        "color": (0, 0, 0),
        "cost": 0,
        "purchased": True
    },
    {
        "name": "Big Bomb",
        "explosion_radius": 70,
        "crater_multiplier": 1.2,
        "gravity_multiplier": 1.2,
        "wind_multiplier": 1.2,
        "weight": 1.5,
        "warheads": 1,
        "special": None,
        "color": (50, 50, 50),
        "cost": 100,
        "purchased": False
    },
    {
        "name": "Digger",
        "explosion_radius": 20,
        "crater_multiplier": 2.5,
        "gravity_multiplier": 1.0,
        "wind_multiplier": 0.8,
        "weight": 1.0,
        "warheads": 1,
        "special": "piercing",
        "color": (139, 69, 19),
        "cost": 80,
        "purchased": False
    },
    {
        "name": "Cluster Bomb",
        "explosion_radius": 25,
        "crater_multiplier": 1.0,
        "gravity_multiplier": 1.0,
        "wind_multiplier": 1.0,
        "weight": 1.2,
        "warheads": 3,
        "special": None,
        "color": (200, 0, 0),
        "cost": 150,
        "purchased": False
    },
    {
        "name": "Bouncer",
        "explosion_radius": 40,
        "crater_multiplier": 1.0,
        "gravity_multiplier": 0.8,
        "wind_multiplier": 1.5,
        "weight": 0.8,
        "warheads": 1,
        "special": "bouncing",
        "color": (0, 200, 0),
        "cost": 120,
        "purchased": False
    },
    {
        "name": "Homing Missile",
        "explosion_radius": 35,
        "crater_multiplier": 1.0,
        "gravity_multiplier": 0.9,
        "wind_multiplier": 0.7,
        "weight": 0.9,
        "warheads": 1,
        "special": "homing",
        "color": (255, 0, 255),
        "cost": 200,
        "purchased": False
    },
    {
        "name": "Fragmentation Shell",
        "explosion_radius": 30,
        "crater_multiplier": 1.0,
        "gravity_multiplier": 1.0,
        "wind_multiplier": 1.0,
        "weight": 1.0,
        "warheads": 1,
        "special": "fragmentation",
        "color": (255, 165, 0),
        "cost": 180,
        "purchased": False
    },
    {
        "name": "Smoke Bomb",
        "explosion_radius": 100,
        "crater_multiplier": 0.0,
        "gravity_multiplier": 1.0,
        "wind_multiplier": 1.0,
        "weight": 1.0,
        "warheads": 1,
        "special": "smoke",
        "color": (128, 128, 128),
        "cost": 90,
        "purchased": False
    }
]

# --- Game Setup ---
pygame.init()
screen = pygame.display.set_mode((SCREEN_WIDTH, SCREEN_HEIGHT))
pygame.display.set_caption("Python Scorched Earth - A/D to Move, W for Weapon, Arrows for Aim, M for Menu")
font = pygame.font.SysFont("consolas", 20)
game_over_font = pygame.font.SysFont("consolas", 72)
menu_font = pygame.font.SysFont("consolas", 30)
clock = pygame.time.Clock()

# --- Perlin Noise Cloud Setup ---
noise_width, noise_height = 200, 150
scale = 50.0
octaves = 4
persistence = 0.5
lacunarity = 2.0
cloud_speed = 0.03
morph_speed = 0.03

x = np.linspace(0, noise_width / scale, noise_width)
y = np.linspace(0, noise_height / scale, noise_height)
x_grid, y_grid = np.meshgrid(x, y)

vectorized_pnoise3 = np.vectorize(
    lambda x, y, z: pnoise3(
        x, y, z,
        octaves=octaves,
        persistence=persistence,
        lacunarity=lacunarity,
        repeatx=int(noise_width / scale),
        repeaty=int(noise_height / scale),
        repeatz=9999999
    )
)

cloud_surface = pygame.Surface((noise_width, noise_height), pygame.SRCALPHA)
cached_noise_array = None
last_noise_update = -NOISE_UPDATE_INTERVAL

def generate_noise(offset_x, time):
    x_shifted = (x_grid + offset_x / scale) % (noise_width / scale)
    z = time * morph_speed
    noise_array = vectorized_pnoise3(x_shifted, y_grid, z)
    return ((noise_array + 0.5) * 255).T

# --- Terrain Texture Setup ---
texture_surface = pygame.Surface((SCREEN_WIDTH, SCREEN_HEIGHT), pygame.SRCALPHA)
texture_scale = 0.1
texture_octaves = 6
texture_persistence = 0.6
texture_lacunarity = 2.0

def generate_terrain_texture():
    texture_array = np.zeros((SCREEN_HEIGHT, SCREEN_WIDTH), dtype=np.float32)
    for y in range(SCREEN_HEIGHT):
        for x in range(SCREEN_WIDTH):
            noise_value = pnoise3(x * texture_scale, y * texture_scale, 0,
                                  octaves=texture_octaves,
                                  persistence=texture_persistence,
                                  lacunarity=texture_lacunarity)
            texture_array[y, x] = (noise_value + 1.0) * 25
    return texture_array

# --- Helper Functions ---
def generate_terrain(width, height):
    terrain = []
    scale = 0.005
    base_height = height * 0.75
    amplitude = 500
    for x in range(width):
        noise_value = pnoise3(x * scale, 0, 0, octaves=4, persistence=0.5, lacunarity=2.0)
        y = base_height + noise_value * amplitude
        terrain.append(int(max(0, min(height, y))))
    return terrain

def get_terrain_y_and_angle(x, terrain_points):
    x = int(max(0, min(len(terrain_points) - 2, x)))
    y1 = terrain_points[x]
    y2 = terrain_points[x+1]
    angle = math.atan2(y1 - y2, 1)
    return y1, angle

def interpolate_color(height, max_height):
    if height > max_height * 0.9:
        t = (height - max_height * 0.9) / (max_height * 0.1)
        r = int(GROUND_COLOR_LOW[0] + t * (GROUND_COLOR_HIGH[0] - GROUND_COLOR_LOW[0]))
        g = int(GROUND_COLOR_LOW[1] + t * (GROUND_COLOR_HIGH[1] - GROUND_COLOR_LOW[1]))
        b = int(GROUND_COLOR_LOW[2] + t * (GROUND_COLOR_HIGH[2] - GROUND_COLOR_LOW[2]))
    else:
        t = height / (max_height * 0.9)
        r = int(GROUND_COLOR_EARTH[0] + t * (GROUND_COLOR_LOW[0] - GROUND_COLOR_EARTH[0]))
        g = int(GROUND_COLOR_EARTH[1] + t * (GROUND_COLOR_LOW[1] - GROUND_COLOR_EARTH[1]))
        b = int(GROUND_COLOR_EARTH[2] + t * (GROUND_COLOR_LOW[2] - GROUND_COLOR_EARTH[2]))
    return (r, g, b)

def generate_terrain_surface(terrain_points, texture_array):
    terrain_surface = pygame.Surface((SCREEN_WIDTH, SCREEN_HEIGHT), pygame.SRCALPHA)
    max_height = max(terrain_points)
    for x in range(SCREEN_WIDTH):
        y_top = int(terrain_points[x])
        for y in range(y_top, SCREEN_HEIGHT):
            base_color = interpolate_color(y_top, max_height)
            texture_value = texture_array[y, x]
            r = int(max(0, min(255, base_color[0] + texture_value * 0.5)))
            g = int(max(0, min(255, base_color[1] + texture_value * 0.5)))
            b = int(max(0, min(255, base_color[2] + texture_value * 0.5)))
            terrain_surface.set_at((x, y), (r, g, b, 255))
    return terrain_surface

def update_terrain_surface(terrain_surface, terrain_points, texture_array, px, radius):
    max_height = max(terrain_points)
    x_start = max(0, px - radius - EXPLOSION_BUFFER)
    x_end = min(SCREEN_WIDTH, px + radius + EXPLOSION_BUFFER)
    y_min = min(terrain_points[x_start:x_end]) if x_end > x_start else terrain_points[px]
    y_min = max(0, y_min - EXPLOSION_BUFFER)
    region_width = x_end - x_start
    region_height = SCREEN_HEIGHT - y_min
    region_surface = pygame.Surface((region_width, region_height), pygame.SRCALPHA)
    for x in range(x_start, x_end):
        y_top = int(terrain_points[x])
        for y in range(y_min, SCREEN_HEIGHT):
            y_region = y - y_min
            if y >= y_top:
                base_color = interpolate_color(y_top, max_height)
                texture_value = texture_array[y, x]
                r = int(max(0, min(255, base_color[0] + texture_value * 0.5)))
                g = int(max(0, min(255, base_color[1] + texture_value * 0.5)))
                b = int(max(0, min(255, base_color[2] + texture_value * 0.5)))
                region_surface.set_at((x - x_start, y_region), (r, g, b, 255))
            else:
                region_surface.set_at((x - x_start, y_region), (0, 0, 0, 0))
    terrain_surface.blit(region_surface, (x_start, y_min))
    return terrain_surface

# --- Bonus Item Class ---
class BonusItem:
    def __init__(self, x, y, points):
        self.x = x
        self.y = y
        self.points = points
        self.radius = BONUS_ITEM_RADIUS
        self.collected = False

    def draw(self, surface):
        if not self.collected:
            pygame.draw.circle(surface, BONUS_ITEM_COLOR, (int(self.x), int(self.y)), self.radius)
            pygame.draw.circle(surface, (255, 255, 255), (int(self.x), int(self.y)), self.radius // 2)

    def check_collision(self, explosion_x, explosion_y, explosion_radius):
        if self.collected:
            return False
        distance = math.sqrt((self.x - explosion_x)**2 + (self.y - explosion_y)**2)
        return distance <= explosion_radius + self.radius

# --- UI Classes ---
class Button:
    def __init__(self, text, x, y, width, height, action=None, weapon_index=None):
        self.text = text
        self.rect = pygame.Rect(x, y, width, height)
        self.hovered = False
        self.action = action
        self.weapon_index = weapon_index

    def draw(self, surface):
        color = BUTTON_HOVER_COLOR if self.hovered else BUTTON_COLOR
        pygame.draw.rect(surface, color, self.rect)
        text_surface = menu_font.render(self.text, True, TEXT_COLOR)
        text_rect = text_surface.get_rect(center=self.rect.center)
        surface.blit(text_surface, text_rect)

    def check_hover(self, mouse_pos):
        self.hovered = self.rect.collidepoint(mouse_pos)
        return self.hovered

    def check_click(self, mouse_pos):
        return self.rect.collidepoint(mouse_pos)

class Menu:
    def __init__(self):
        self.active = False
        self.mode = "main"
        self.buttons = []
        self.init_main_menu()

    def init_main_menu(self):
        self.mode = "main"
        button_width = 300
        button_height = 60
        button_spacing = 20
        start_x = (SCREEN_WIDTH - button_width) // 2
        start_y = (SCREEN_HEIGHT - (4 * button_height + 3 * button_spacing)) // 2
        self.buttons = [
            Button("Resume Game", start_x, start_y, button_width, button_height, action="resume"),
            Button("Restart Game", start_x, start_y + button_height + button_spacing, button_width, button_height, action="restart"),
            Button("Store", start_x, start_y + 2 * (button_height + button_spacing), button_width, button_height, action="store"),
            Button("Quit Game", start_x, start_y + 3 * (button_height + button_spacing), button_width, button_height, action="quit")
        ]

    def init_store_menu(self):
        self.mode = "store"
        button_width = 200
        button_height = 40
        button_spacing = 10
        start_x = (SCREEN_WIDTH - button_width) // 2
        start_y = 150
        self.buttons = [Button("Back", start_x, SCREEN_HEIGHT - button_height - 20, button_width, button_height, action="back")]
        for i, weapon in enumerate(WEAPONS):
            if not weapon["purchased"]:
                self.buttons.append(Button(f"Buy {weapon['name']} ({weapon['cost']} pts)",
                                        start_x, start_y + i * (button_height + button_spacing),
                                        button_width, button_height, action="buy", weapon_index=i))

    def draw(self, surface, players):
        if not self.active:
            return
        overlay = pygame.Surface((SCREEN_WIDTH, SCREEN_HEIGHT), pygame.SRCALPHA)
        overlay.fill(MENU_BG_COLOR)
        surface.blit(overlay, (0, 0))
        title = menu_font.render("Menu" if self.mode == "main" else "Store", True, TEXT_COLOR)
        title_rect = title.get_rect(center=(SCREEN_WIDTH // 2, 100))
        surface.blit(title, title_rect)
        for button in self.buttons:
            button.draw(surface)
        for i, player in enumerate(players):
            points_text = menu_font.render(f"Player {i+1} Points: {player.points}", True, player.color)
            points_rect = points_text.get_rect(topleft=(20, 20 + i * 30))
            surface.blit(points_text, points_rect)

    def handle_event(self, event, terrain_points, terrain_texture_array, terrain_surface, players, game_state, explosions, active_projectile, current_player_index, cloud_offset_x, cached_noise_array, last_noise_update, wind, bonus_items, smoke_effect):
        if not self.active:
            return (terrain_surface, terrain_points, game_state, explosions, active_projectile, current_player_index, cloud_offset_x, cached_noise_array, last_noise_update, wind, bonus_items, smoke_effect)
        if event.type == pygame.MOUSEMOTION:
            mouse_pos = event.pos
            for button in self.buttons:
                button.check_hover(mouse_pos)
        if event.type == pygame.MOUSEBUTTONDOWN and event.button == 1:
            mouse_pos = event.pos
            for button in self.buttons:
                if button.check_click(mouse_pos) and button.action:
                    if self.mode == "main":
                        if button.action == "resume":
                            self.active = False
                        elif button.action == "restart":
                            terrain_points = generate_terrain(SCREEN_WIDTH, SCREEN_HEIGHT)
                            terrain_texture_array = generate_terrain_texture()
                            terrain_surface = generate_terrain_surface(terrain_points, terrain_texture_array)
                            bonus_items = generate_bonus_items(terrain_points)
                            for p in players:
                                p.terrain = terrain_points
                                p.reset()
                                p.score = 0
                                p.points = 0
                                for weapon in WEAPONS:
                                    if weapon["cost"] == 0:
                                        weapon["purchased"] = True
                                    else:
                                        weapon["purchased"] = False
                                p.purchased_weapons = [i for i, w in enumerate(WEAPONS) if w["purchased"]]
                            game_state = "AIMING"
                            explosions = []
                            active_projectile = None
                            current_player_index = 0
                            cloud_offset_x = 0
                            cached_noise_array = None
                            last_noise_update = -NOISE_UPDATE_INTERVAL
                            wind = random.uniform(WIND_MIN, WIND_MAX)
                            smoke_effect = None
                            self.active = False
                            return (terrain_surface, terrain_points, game_state, explosions, active_projectile, current_player_index, cloud_offset_x, cached_noise_array, last_noise_update, wind, bonus_items, smoke_effect)
                        elif button.action == "store":
                            self.init_store_menu()
                        elif button.action == "quit":
                            pygame.event.post(pygame.event.Event(pygame.QUIT))
                    elif self.mode == "store":
                        if button.action == "back":
                            self.init_main_menu()
                        elif button.action == "buy" and button.weapon_index is not None:
                            weapon = WEAPONS[button.weapon_index]
                            if not weapon["purchased"] and players[current_player_index].points >= weapon["cost"]:
                                players[current_player_index].points -= weapon["cost"]
                                weapon["purchased"] = True
                                players[current_player_index].purchased_weapons.append(button.weapon_index)
                                self.init_store_menu()
        return (terrain_surface, terrain_points, game_state, explosions, active_projectile, current_player_index, cloud_offset_x, cached_noise_array, last_noise_update, wind, bonus_items, smoke_effect)

# --- Classes ---
class Tank:
    def __init__(self, x, color, terrain):
        self.initial_x = x
        self.color = color
        self.terrain = terrain
        self.angle = 90
        self.power = 50
        self.is_alive = True
        self.score = 0
        self.points = 0
        self.current_weapon_index = 0
        self.fuel = MAX_FUEL
        self.purchased_weapons = [0]
        self.reset()

    def reset(self):
        self.x = self.initial_x
        self.y, self.surface_angle = get_terrain_y_and_angle(self.x, self.terrain)
        self.rect = pygame.Rect(self.x - TANK_WIDTH // 2, self.y - TANK_HEIGHT, TANK_WIDTH, TANK_HEIGHT)
        self.is_alive = True
        self.reset_fuel()
        if self.x < SCREEN_WIDTH / 2:
            self.angle = 45
        else:
            self.angle = 135

    def reset_fuel(self):
        self.fuel = MAX_FUEL

    def move(self, direction):
        if self.fuel > 0:
            self.x += direction * TANK_MOVE_SPEED
            self.x = max(TANK_WIDTH // 2, min(SCREEN_WIDTH - TANK_WIDTH // 2, self.x))
            self.y, self.surface_angle = get_terrain_y_and_angle(self.x, self.terrain)
            self.rect.centerx = self.x
            self.rect.bottom = self.y
            self.fuel -= 1

    def draw(self, surface):
        if not self.is_alive:
            return
        tank_sprite_surface = pygame.Surface((TANK_WIDTH * 2, TANK_WIDTH * 2), pygame.SRCALPHA)
        sprite_center_x, sprite_center_y = TANK_WIDTH, TANK_WIDTH
        tread_color = (60, 60, 60)
        body_color = self.color
        pygame.draw.circle(tank_sprite_surface, tread_color, (sprite_center_x - TANK_WIDTH // 3, sprite_center_y + TANK_HEIGHT // 4), TANK_HEIGHT // 1.8)
        pygame.draw.circle(tank_sprite_surface, tread_color, (sprite_center_x + TANK_WIDTH // 3, sprite_center_y + TANK_HEIGHT // 4), TANK_HEIGHT // 1.8)
        body_rect = pygame.Rect(0, 0, TANK_WIDTH, TANK_HEIGHT)
        body_rect.center = (sprite_center_x, sprite_center_y)
        pygame.draw.rect(tank_sprite_surface, body_color, body_rect, border_radius=4)
        rotated_surface = pygame.transform.rotate(tank_sprite_surface, math.degrees(self.surface_angle))
        rotated_rect = rotated_surface.get_rect(center=(self.x, self.y - TANK_HEIGHT / 2))
        surface.blit(rotated_surface, rotated_rect.topleft)
        turret_base_x = self.x
        turret_base_y = self.y - TANK_HEIGHT / 2
        turret_end_x = turret_base_x + math.cos(math.radians(self.angle)) * TURRET_LENGTH
        turret_end_y = turret_base_y - math.sin(math.radians(self.angle)) * TURRET_LENGTH
        turret_color = (max(0, self.color[0] - 40), max(0, self.color[1] - 40), max(0, self.color[2] - 40))
        pygame.draw.line(surface, turret_color, (turret_base_x, turret_base_y), (turret_end_x, turret_end_y), TURRET_WIDTH)

    def switch_weapon(self):
        current_idx = self.purchased_weapons.index(self.current_weapon_index) if self.current_weapon_index in self.purchased_weapons else 0
        self.current_weapon_index = self.purchased_weapons[(current_idx + 1) % len(self.purchased_weapons)]

    def fire(self):
        weapon = WEAPONS[self.current_weapon_index]
        angle_rad = math.radians(self.angle)
        start_x = self.x + math.cos(angle_rad) * TURRET_LENGTH
        start_y = self.y - TANK_HEIGHT / 2 - math.sin(angle_rad) * TURRET_LENGTH
        return Projectile(start_x, start_y, self.power, angle_rad, weapon, self)

class Projectile:
    def __init__(self, x, y, power, angle_rad, weapon_type, tank):
        self.x = x
        self.y = y
        self.weapon_type = weapon_type
        self.vx = math.cos(angle_rad) * power * 0.2
        self.vy = -math.sin(angle_rad) * power * 0.2
        self.path = []
        self.time_alive = 0
        self.bounced = False
        self.tank = tank  # Reference to firing tank for homing

    def update(self, wind, terrain_points, players, current_player_index):
        self.time_alive += 1 / 60
        self.vy += GRAVITY * self.weapon_type["gravity_multiplier"] * self.weapon_type["weight"]
        if self.weapon_type["special"] == "homing":
            target = players[1 - current_player_index]
            if target.is_alive:
                dx = target.x - self.x
                dy = target.y - self.y
                distance = math.sqrt(dx**2 + dy**2)
                if distance > 10:
                    speed = math.sqrt(self.vx**2 + self.vy**2)
                    target_vx = dx / distance * speed * 0.1
                    target_vy = dy / distance * speed * 0.1
                    self.vx += (target_vx - self.vx) * 0.05
                    self.vy += (target_vy - self.vy) * 0.05
        self.x += self.vx + wind * self.weapon_type["wind_multiplier"]
        self.y += self.vy
        px, py = int(self.x), int(self.y)
        if self.weapon_type["special"] == "bouncing" and not self.bounced and 0 <= px < SCREEN_WIDTH and py >= terrain_points[px]:
            self.vy = -self.vy * 0.5
            self.vx *= 0.7
            self.bounced = True
        self.path.append((px, py))

    def draw(self, surface):
        if len(self.path) > 2:
            pygame.draw.lines(surface, (255, 255, 0), False, self.path, 2)
        pygame.draw.circle(surface, self.weapon_type["color"], (int(self.x), int(self.y)), 5)

class Explosion:
    def __init__(self, x, y, radius):
        self.x = int(x)
        self.y = int(y)
        self.timer = EXPLOSION_DURATION
        self.max_radius = radius

    def update(self):
        self.timer -= 1
        return self.timer > 0

    def draw(self, surface):
        if self.timer <= 0:
            return
        progress = (EXPLOSION_DURATION - self.timer) / EXPLOSION_DURATION
        current_radius = int(self.max_radius * progress)
        color = (255, 165, 0) if self.timer % 4 < 2 else (255, 69, 0)
        pygame.draw.circle(surface, color, (self.x, self.y), current_radius)

class SmokeEffect:
    def __init__(self, x, y):
        self.x = x
        self.y = y
        self.timer = SMOKE_DURATION
        self.radius = 100

    def update(self):
        self.timer -= 1
        return self.timer > 0

    def draw(self, surface):
        if self.timer <= 0:
            return
        progress = self.timer / SMOKE_DURATION
        alpha = int(100 * progress)
        smoke_surface = pygame.Surface((self.radius * 2, self.radius * 2), pygame.SRCALPHA)
        pygame.draw.circle(smoke_surface, (128, 128, 128, alpha), (self.radius, self.radius), self.radius)
        surface.blit(smoke_surface, (self.x - self.radius, self.y - self.radius))

def generate_bonus_items(terrain_points):
    bonus_items = []
    num_items = random.randint(3, 5)
    for _ in range(num_items):
        x = random.randint(50, SCREEN_WIDTH - 50)
        y = terrain_points[x] - BONUS_ITEM_RADIUS - 5
        bonus_items.append(BonusItem(x, y, BONUS_ITEM_POINTS))
    return bonus_items

def switch_turn():
    global current_player_index, wind
    current_player_index = 1 - current_player_index
    players[current_player_index].reset_fuel()
    wind = random.uniform(WIND_MIN, WIND_MAX)

# --- Game State Variables ---
terrain_points = generate_terrain(SCREEN_WIDTH, SCREEN_HEIGHT)
terrain_texture_array = generate_terrain_texture()
terrain_surface = generate_terrain_surface(terrain_points, terrain_texture_array)
bonus_items = generate_bonus_items(terrain_points)
player1 = Tank(random.randint(50, SCREEN_WIDTH // 2 - 100), (0, 0, 255), terrain_points)
player2 = Tank(random.randint(SCREEN_WIDTH // 2 + 100, SCREEN_WIDTH - 50), (255, 0, 0), terrain_points)
players = [player1, player2]
current_player_index = 0
active_projectile = None
explosions = []
game_state = "AIMING"
cloud_offset_x = 0
frame_count = 0
menu = Menu()
wind = random.uniform(WIND_MIN, WIND_MAX)
smoke_effect = None

# --- Main Game Loop ---
running = True
while running:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False
        if event.type == pygame.KEYDOWN:
            if event.key == pygame.K_m:
                menu.active = not menu.active
                if menu.active:
                    menu.init_main_menu()
            if not menu.active:
                if game_state == "GAMEOVER":
                    if event.key == pygame.K_r:
                        terrain_points = generate_terrain(SCREEN_WIDTH, SCREEN_HEIGHT)
                        terrain_texture_array = generate_terrain_texture()
                        terrain_surface = generate_terrain_surface(terrain_points, terrain_texture_array)
                        bonus_items = generate_bonus_items(terrain_points)
                        for p in players:
                            p.terrain = terrain_points
                            p.reset()
                            p.score = 0
                            p.points = 0
                            for weapon in WEAPONS:
                                if weapon["cost"] == 0:
                                    weapon["purchased"] = True
                                else:
                                    weapon["purchased"] = False
                                p.purchased_weapons = [i for i, w in enumerate(WEAPONS) if w["purchased"]]
                        current_player_index = 0
                        active_projectile = None
                        explosions = []
                        game_state = "AIMING"
                        cloud_offset_x = 0
                        cached_noise_array = None
                        last_noise_update = -NOISE_UPDATE_INTERVAL
                        wind = random.uniform(WIND_MIN, WIND_MAX)
                        smoke_effect = None
                if game_state == "AIMING":
                    current_tank = players[current_player_index]
                    if event.key == pygame.K_LEFT:
                        current_tank.angle = min(180, current_tank.angle + 2)
                    if event.key == pygame.K_RIGHT:
                        current_tank.angle = max(0, current_tank.angle - 2)
                    if event.key == pygame.K_UP:
                        current_tank.power = min(100, current_tank.power + 2)
                    if event.key == pygame.K_DOWN:
                        current_tank.power = max(0, current_tank.power - 2)
                    if event.key == pygame.K_w:
                        current_tank.switch_weapon()
                    if event.key == pygame.K_SPACE:
                        active_projectile = current_tank.fire()
                        game_state = "FIRING"
        if menu.active:
            terrain_surface, terrain_points, game_state, explosions, active_projectile, current_player_index, cloud_offset_x, cached_noise_array, last_noise_update, wind, bonus_items, smoke_effect = menu.handle_event(
                event, terrain_points, terrain_texture_array, terrain_surface, players, game_state, explosions, active_projectile, current_player_index, cloud_offset_x, cached_noise_array, last_noise_update, wind, bonus_items, smoke_effect
            )
            for p in players:
                p.terrain = terrain_points

    if not menu.active:
        if game_state == "AIMING" and current_player_index == 0:
            keys = pygame.key.get_pressed()
            if keys[pygame.K_a]:
                players[0].move(-1)
            if keys[pygame.K_d]:
                players[0].move(1)

        if current_player_index == 1 and game_state == "AIMING":
            ai_tank = players[1]
            player_tank = players[0]
            ai_tank.current_weapon_index = random.choice(ai_tank.purchased_weapons)
            dx = player_tank.x - ai_tank.x
            dy = player_tank.y - ai_tank.y
            if dx > 0:
                ai_tank.angle = random.randint(10, 80)
            else:
                ai_tank.angle = random.randint(100, 170)
            distance = math.sqrt(dx**2 + dy**2)
            ai_tank.power = int(max(20, min(100, distance / 7)))
            active_projectile = ai_tank.fire()
            game_state = "FIRING"

        if game_state == "FIRING" and active_projectile:
            active_projectile.update(wind, terrain_points, players, current_player_index)
            px, py = int(active_projectile.x), int(active_projectile.y)
            weapon_hit = active_projectile.weapon_type
            should_explode = False
            if weapon_hit["special"] == "piercing" and active_projectile.time_alive < 0.5:
                pass
            elif 0 <= px < SCREEN_WIDTH and py >= terrain_points[px]:
                should_explode = True
            if active_projectile:
                for i, tank in enumerate(players):
                    if tank.is_alive and tank.rect.collidepoint(px, py):
                        tank.is_alive = False
                        players[1-i].score += 1
                        should_explode = True
                        break
            if should_explode:
                explosions.append(Explosion(px, py, weapon_hit["explosion_radius"]))
                if weapon_hit["warheads"] > 1:
                    for _ in range(weapon_hit["warheads"] - 1):
                        offset_x = px + random.randint(-20, 20)
                        offset_y = py + random.randint(-20, 20)
                        if 0 <= offset_x < SCREEN_WIDTH:
                            explosions.append(Explosion(offset_x, offset_y, weapon_hit["explosion_radius"] // 2))
                if weapon_hit["special"] == "fragmentation":
                    for _ in range(4):
                        angle = random.uniform(0, 2 * math.pi)
                        power = random.uniform(20, 40)
                        frag_weapon = {
                            "explosion_radius": 15,
                            "crater_multiplier": 0.5,
                            "gravity_multiplier": 1.0,
                            "wind_multiplier": 1.0,
                            "weight": 0.5,
                            "warheads": 1,
                            "special": None,
                            "color": (255, 165, 0)
                        }
                        frag_projectile = Projectile(px, py, power, angle, frag_weapon, active_projectile.tank)
                        explosions.append(Explosion(px, py, frag_weapon["explosion_radius"]))
                if weapon_hit["special"] == "smoke":
                    smoke_effect = SmokeEffect(px, py)
                if weapon_hit["crater_multiplier"] > 0:
                    radius = weapon_hit["explosion_radius"]
                    crater_mult = weapon_hit["crater_multiplier"]
                    for i in range(max(0, px - radius), min(SCREEN_WIDTH, px + radius)):
                        dist = abs(px - i)
                        if dist < radius:
                            crater_depth = math.sqrt(radius**2 - dist**2) * crater_mult
                            terrain_points[i] = int(min(SCREEN_HEIGHT, terrain_points[i] + crater_depth))
                    terrain_surface = update_terrain_surface(terrain_surface, terrain_points, terrain_texture_array, px, radius)
                for p in players:
                    p.terrain = terrain_points
                for item in bonus_items:
                    if item.check_collision(px, py, weapon_hit["explosion_radius"]):
                        item.collected = True
                        players[current_player_index].points += item.points
                active_projectile = None
                game_state = "EXPLODING"
            if active_projectile and not (0 <= px < SCREEN_WIDTH and 0 <= py < SCREEN_HEIGHT + 200):
                active_projectile = None
                game_state = "AIMING"
                switch_turn()

        if game_state == "EXPLODING":
            for explosion in explosions[:]:
                if not explosion.update():
                    explosions.remove(explosion)
            if smoke_effect and not smoke_effect.update():
                smoke_effect = None
            if not explosions:
                for tank in players:
                    if tank.is_alive:
                        tank.y, tank.surface_angle = get_terrain_y_and_angle(tank.x, terrain_points)
                        tank.rect.bottom = tank.y
                if not player1.is_alive or not player2.is_alive:
                    game_state = "GAMEOVER"
                else:
                    game_state = "AIMING"
                    switch_turn()

        cloud_offset_x = (cloud_offset_x + cloud_speed) % (SCREEN_WIDTH * 2)
        time = pygame.time.get_ticks() / 1000.0
        if frame_count - last_noise_update >= NOISE_UPDATE_INTERVAL:
            cached_noise_array = generate_noise(cloud_offset_x, time)
            last_noise_update = frame_count
        if cached_noise_array is not None:
            pixel_array = pygame.surfarray.pixels3d(cloud_surface)
            alpha_array = pygame.surfarray.pixels_alpha(cloud_surface)
            mask = cached_noise_array > 0
            pixel_array[mask] = [255, 255, 255]
            pixel_array[~mask] = [0, 0, 0]
            alpha_array[:] = np.clip((cached_noise_array - 50) * 2, 0, 255).astype(np.uint8)
            del pixel_array
            del alpha_array
            scaled_cloud_surface = pygame.transform.scale(cloud_surface, (SCREEN_WIDTH, SCREEN_HEIGHT))
        else:
            scaled_cloud_surface = pygame.Surface((SCREEN_WIDTH, SCREEN_HEIGHT), pygame.SRCALPHA)

    # Drawing
    screen.fill(BACKGROUND_COLOR)
    screen.blit(scaled_cloud_surface, (0, 0))
    screen.blit(terrain_surface, (0, 0))
    pygame.draw.lines(screen, (20, 80, 20), False, [(x, y) for x, y in enumerate(terrain_points)], 2)

    for item in bonus_items:
        item.draw(screen)
    for player in players:
        player.draw(screen)
    if active_projectile:
        active_projectile.draw(screen)
    for explosion in explosions:
        explosion.draw(screen)
    if smoke_effect:
        smoke_effect.draw(screen)

    if game_state != "GAMEOVER" and not menu.active:
        current_tank = players[current_player_index]
        weapon = WEAPONS[current_tank.current_weapon_index]
        angle_text = font.render(f"Angle: {current_tank.angle:.0f}", True, (0,0,0))
        power_text = font.render(f"Power: {current_tank.power:.0f}", True, (0,0,0))
        weapon_text = font.render(f"Weapon: {weapon['name']}", True, (0,0,0))
        fuel_text = font.render(f"Fuel: {current_tank.fuel:.0f}", True, (0,0,0))
        wind_text = font.render(f"Wind: {wind:.2f}", True, (0,0,0))
        points_text = font.render(f"Points: {current_tank.points}", True, current_tank.color)
        screen.blit(angle_text, (20, 20))
        screen.blit(power_text, (20, 50))
        screen.blit(weapon_text, (20, 80))
        screen.blit(fuel_text, (20, 110))
        screen.blit(wind_text, (20, 140))
        screen.blit(points_text, (20, 170))
        p1_score_text = font.render(f"P1 Score: {player1.score}", True, player1.color)
        p2_score_text = font.render(f"P2 Score: {player2.score}", True, player2.color)
        screen.blit(p1_score_text, (SCREEN_WIDTH - 180, 20))
        screen.blit(p2_score_text, (SCREEN_WIDTH - 180, 50))
    elif game_state == "GAMEOVER" and not menu.active:
        winner_index = -1
        if not player1.is_alive and not player2.is_alive:
            pass
        elif not player2.is_alive:
            winner_index = 0
        elif not player1.is_alive:
            winner_index = 1
        if winner_index == -1:
            win_text_surface = game_over_font.render("MUTUAL DESTRUCTION!", True, (255, 255, 0))
        else:
            winner = players[winner_index]
            win_text_surface = game_over_font.render(f"PLAYER {winner_index+1} WINS!", True, winner.color)
        win_text_rect = win_text_surface.get_rect(center=(SCREEN_WIDTH/2, SCREEN_HEIGHT/2 - 50))
        screen.blit(win_text_surface, win_text_rect)
        reset_text_surface = font.render("Press 'R' to play again or 'M' for Menu", True, (0,0,0))
        reset_text_rect = reset_text_surface.get_rect(center=(SCREEN_WIDTH/2, SCREEN_HEIGHT/2 + 20))
        screen.blit(reset_text_surface, reset_text_rect)

    menu.draw(screen, players)
    pygame.display.flip()
    clock.tick(60)
    frame_count += 1

pygame.quit()