import numpy as np
from opensimplex import OpenSimplex
from settings import WorldSettings
from numba import njit

class ClimateGenerator:
    def __init__(self, height_map: np.ndarray, settings: WorldSettings):
        self.height_map = height_map
        self.settings = settings
        self.noise = OpenSimplex(seed=settings.get_valid_seed())
        
        # Палитра биомов
        self.biome_colors = {
            'water': (0.2, 0.4, 0.8, 1.0),
            'ice': (0.95, 0.95, 0.98, 0.8),
            'tundra': (0.4, 0.6, 0.9, 0.9),
            'taiga': (0.1, 0.3, 0.15, 1.0),
            'forest': (0.3, 0.6, 0.3, 1.0),
            'desert': (0.9, 0.6, 0.2, 1.0),
            'steppe': (0.8, 0.75, 0.5, 1.0),
            'swamp': (0.4, 0.35, 0.2, 0.9),
            'jungle': (0.1, 0.5, 0.1, 1.0),
            'savanna': (0.8, 0.7, 0.3, 1.0)
        }

    def generate_biome_map(self, enabled_biomes=None) -> np.ndarray:
        """Генерация карты биомов с учетом выбранных биомов"""
        if enabled_biomes is None:
            enabled_biomes = list(self.biome_colors.keys())
            
        height, width = self.height_map.shape
        
        # Генерация шумов
        temp_noise = self._generate_noise(width, height, 40)
        moist_noise = self._generate_noise(width, height, 50)
        
        biome_map = np.zeros((height, width, 4), dtype=np.float32)
        self._fill_biome_map(
            biome_map, 
            self.height_map,
            temp_noise,
            moist_noise,
            enabled_biomes
        )
        return biome_map

    def _generate_noise(self, width, height, scale):
        """Генерация многооктавного шума"""
        noise = np.zeros((height, width))
        for octave in [1.0, 0.5, 0.25]:
            for y in range(height):
                for x in range(width):
                    nx = x / (scale * octave)
                    ny = y / (scale * octave)
                    noise[y, x] += self.noise.noise2(nx, ny) * octave
        return (noise - noise.min()) / (noise.max() - noise.min())

    @staticmethod
    @njit
    def _fill_biome_map(biome_map, height_map, temp_noise, moist_noise, enabled_biomes):
        height, width = height_map.shape
        water_level = 0.05
        
        for y in range(height):
            latitude = abs(y/height - 0.5) * 2  # 0 на экваторе, 1 на полюсах
            
            for x in range(width):
                h = height_map[y,x]
                
                # Вода (всегда включена, без прозрачности)
                if h < water_level:
                    biome_map[y,x] = (0.2, 0.4, 0.8, 1.0)
                    continue
                
                # Климатические параметры
                temp = (0.7 - latitude * 0.5) * (1.0 - h * 0.5) * (0.8 + temp_noise[y,x] * 0.4)
                moist = moist_noise[y,x]
                
                # Расчёт прозрачности: сильнее зависит от высоты
                alpha = 1.0 - h**2 * 0.7  # h^2 делает прозрачность резче на вершинах
                
                # Выбор биома
                if temp < 0.3:  # Холодные зоны
                    if h > 0.8 and 'ice' in enabled_biomes:
                        biome = (0.95, 0.95, 0.98, alpha * 0.7)  # Лёд прозрачнее
                    elif moist > 0.65 and 'tundra' in enabled_biomes:
                        biome = (0.4, 0.6, 0.9, alpha)
                    elif 'taiga' in enabled_biomes:
                        biome = (0.1, 0.3, 0.15, alpha)
                    else:
                        biome = (0.3, 0.6, 0.3, alpha)
                
                elif temp < 0.6:  # Умеренные зоны
                    if moist > 0.7 and 'swamp' in enabled_biomes:
                        biome = (0.4, 0.35, 0.2, alpha)
                    elif moist > 0.5 and 'forest' in enabled_biomes:
                        biome = (0.3, 0.6, 0.3, alpha)
                    elif moist > 0.3 and 'steppe' in enabled_biomes:
                        biome = (0.8, 0.75, 0.5, alpha)
                    elif 'desert' in enabled_biomes:
                        biome = (0.9, 0.6, 0.2, alpha)
                    else:
                        biome = (0.3, 0.6, 0.3, alpha)
                
                else:  # Теплые зоны
                    if moist < 0.3 and 'desert' in enabled_biomes:
                        biome = (0.9, 0.6, 0.2, alpha)
                    elif moist < 0.5 and 'savanna' in enabled_biomes:
                        biome = (0.8, 0.7, 0.3, alpha)
                    elif moist > 0.7 and 'jungle' in enabled_biomes:
                        biome = (0.1, 0.5, 0.1, alpha)
                    elif 'forest' in enabled_biomes:
                        biome = (0.3, 0.6, 0.3, alpha)
                    else:
                        biome = (0.8, 0.75, 0.5, alpha)
                
                biome_map[y,x] = biome