icons.py

"""
Weather icon drawing module.
All functions return a pygame.Surface with SRCALPHA.
"""
import math
import pygame


# Color palette
WHITE = (255, 255, 255)
LIGHT_GRAY = (200, 200, 210)
DARK_GRAY = (140, 140, 150)
SUN_YELLOW = (255, 210, 50)
RAIN_BLUE = (100, 180, 255)
SNOW_WHITE = (220, 240, 255)
THUNDER_YELLOW = (255, 230, 50)
FOG_GRAY = (160, 160, 170)
ORANGE = (255, 160, 50)
MOON_YELLOW = (220, 200, 100)


def _cloud_shape(surface, cx, cy, w, h, color):
    """Draw a smooth cloud shape centered at (cx, cy)"""
    r_main = w // 4
    r_left = w // 7
    r_right = w // 5
    r_top = w // 5

    # Bottom fill rect
    pygame.draw.rect(surface, color,
                     (cx - r_main, cy, r_main * 2, r_main),
                     border_radius=r_main // 2)
    # Main body
    pygame.draw.circle(surface, color, (cx, cy), r_main)
    # Left bump
    pygame.draw.circle(surface, color, (cx - r_main + 2, cy + 2), r_left)
    # Right bump
    pygame.draw.circle(surface, color, (cx + r_main - 4, cy - 2), r_right)
    # Top bump
    pygame.draw.circle(surface, color, (cx - r_main // 3, cy - r_main // 2), r_top)
    # Top-right bump
    pygame.draw.circle(surface, color, (cx + r_main // 4, cy - r_main // 2), r_top - 4)


def draw_sun(size=80):
    s = pygame.Surface((size, size), pygame.SRCALPHA)
    cx, cy = size // 2, size // 2
    r = size // 4
    ray_inner = r + 4
    ray_outer = int(size * 0.45)
    pygame.draw.circle(s, SUN_YELLOW, (cx, cy), r)
    for i in range(8):
        angle = i * math.pi / 4 - math.pi / 8
        x1 = cx + int(ray_inner * math.cos(angle))
        y1 = cy + int(ray_inner * math.sin(angle))
        x2 = cx + int(ray_outer * math.cos(angle))
        y2 = cy + int(ray_outer * math.sin(angle))
        pygame.draw.line(s, SUN_YELLOW, (x1, y1), (x2, y2), max(2, size // 20))
    return s


def draw_cloud(size=80, color=LIGHT_GRAY):
    s = pygame.Surface((size, size), pygame.SRCALPHA)
    cx, cy = size // 2, size * 3 // 5
    _cloud_shape(s, cx, cy, size, size, color)
    return s


def draw_partly_cloudy(size=80):
    s = pygame.Surface((size, size), pygame.SRCALPHA)
    # Small sun top-right
    sun_size = size * 55 // 100
    sun = draw_sun(sun_size)
    s.blit(sun, (size - sun_size + 2, -2))
    # Cloud in front
    cx, cy = size * 2 // 5, size * 11 // 16
    _cloud_shape(s, cx, cy, size * 7 // 10, size, LIGHT_GRAY)
    return s


def draw_rain(size=80, heavy=False):
    s = pygame.Surface((size, size), pygame.SRCALPHA)
    cx, cy = size // 2, size * 9 // 20
    _cloud_shape(s, cx, cy, size, size, DARK_GRAY)
    # Rain drops
    drop_count = 5 if heavy else 4
    spacing = size // (drop_count + 1)
    for i in range(drop_count):
        x = (i + 1) * spacing
        y_start = cy + size // 4 + 4
        y_end = y_start + size // 5
        offset = (i % 2) * 4
        pygame.draw.line(s, RAIN_BLUE, (x, y_start + offset), (x - 4, y_end + offset), max(2, size // 25))
    return s


def draw_drizzle(size=80):
    s = pygame.Surface((size, size), pygame.SRCALPHA)
    cx, cy = size // 2, size * 9 // 20
    _cloud_shape(s, cx, cy, size, size, DARK_GRAY)
    # Light drizzle
    for i in range(3):
        x = (i + 1) * size // 4
        y_start = cy + size // 4 + 2
        y_end = y_start + size // 6
        pygame.draw.line(s, RAIN_BLUE, (x, y_start), (x - 3, y_end), max(1, size // 30))
    return s


def draw_thunderstorm(size=80):
    s = pygame.Surface((size, size), pygame.SRCALPHA)
    cx, cy = size // 2, size * 2 // 5
    _cloud_shape(s, cx, cy, size, size, DARK_GRAY)
    # Lightning bolt
    bx = cx - size // 12
    by = cy + size // 5
    bolt = [
        (bx + size // 10, by),
        (bx - size // 14, by + size // 5),
        (bx + size // 18, by + size // 5),
        (bx - size // 10, by + size * 2 // 5),
    ]
    pygame.draw.lines(s, THUNDER_YELLOW, False, bolt, max(2, size // 20))
    return s


def draw_snow(size=80):
    s = pygame.Surface((size, size), pygame.SRCALPHA)
    cx, cy = size // 2, size * 9 // 20
    _cloud_shape(s, cx, cy, size, size, LIGHT_GRAY)
    # Snowflakes
    for i in range(5):
        x = (i + 1) * size // 6
        y = cy + size // 4 + size // 8
        r = max(2, size // 20)
        pygame.draw.circle(s, SNOW_WHITE, (x, y), r)
        if size > 50:
            for angle in range(0, 360, 60):
                rad = math.radians(angle)
                x2 = x + int(r * 2 * math.cos(rad))
                y2 = y + int(r * 2 * math.sin(rad))
                pygame.draw.line(s, SNOW_WHITE, (x, y), (x2, y2), 1)
    return s


def draw_fog(size=80):
    s = pygame.Surface((size, size), pygame.SRCALPHA)
    line_color = (*FOG_GRAY, 200)
    y_positions = [size * 2 // 5, size // 2, size * 3 // 5, size * 7 // 10]
    widths = [size * 4 // 5, size * 3 // 4, size * 4 // 5, size * 2 // 3]
    x_offsets = [size // 10, size // 8, size // 10, size // 6]
    for y, w, xo in zip(y_positions, widths, x_offsets):
        r = size // 20
        pygame.draw.rect(s, line_color,
                         (xo, y - r, w, r * 2),
                         border_radius=r)
    return s


def draw_mist(size=80):
    return draw_fog(size)


def draw_moon(size=80):
    """Crescent moon for 01n."""
    s = pygame.Surface((size, size), pygame.SRCALPHA)
    cx, cy = size // 2, size // 2
    r = size * 5 // 14
    pygame.draw.circle(s, MOON_YELLOW, (cx, cy), r)
    # Cut out an offset circle using alpha multiplication to create the crescent
    offset_x = r * 5 // 8
    offset_y = r // 6
    alpha_mask = pygame.Surface((size, size), pygame.SRCALPHA)
    alpha_mask.fill((255, 255, 255, 255))
    pygame.draw.circle(alpha_mask, (255, 255, 255, 0),
                       (cx + offset_x, cy - offset_y), int(r * 0.88))
    s.blit(alpha_mask, (0, 0), special_flags=pygame.BLEND_RGBA_MULT)
    # Small stars scattered in the background
    import random
    rng = random.Random(42)
    for _ in range(4):
        sx = rng.randint(4, size - 4)
        sy = rng.randint(4, size - 4)
        star_r = max(1, size // 40)
        pygame.draw.circle(s, (255, 255, 220, 160), (sx, sy), star_r)
    return s


def draw_partly_cloudy_night(size=80):
    """Moon peeking behind cloud for 02n."""
    s = pygame.Surface((size, size), pygame.SRCALPHA)
    moon_size = size * 55 // 100
    moon = draw_moon(moon_size)
    s.blit(moon, (size - moon_size + 2, -2))
    cx, cy = size * 2 // 5, size * 11 // 16
    _cloud_shape(s, cx, cy, size * 7 // 10, size, LIGHT_GRAY)
    return s


# OWM icon code → draw function mapping (full 3-char key: prefix + day/night suffix)
ICON_MAP = {
    "01d": lambda sz: draw_sun(sz),
    "01n": lambda sz: draw_moon(sz),
    "02d": lambda sz: draw_partly_cloudy(sz),
    "02n": lambda sz: draw_partly_cloudy_night(sz),
    "03d": lambda sz: draw_cloud(sz),
    "03n": lambda sz: draw_cloud(sz),
    "04d": lambda sz: draw_cloud(sz, DARK_GRAY),
    "04n": lambda sz: draw_cloud(sz, DARK_GRAY),
    "09d": lambda sz: draw_drizzle(sz),
    "09n": lambda sz: draw_drizzle(sz),
    "10d": lambda sz: draw_rain(sz),
    "10n": lambda sz: draw_rain(sz),
    "11d": lambda sz: draw_thunderstorm(sz),
    "11n": lambda sz: draw_thunderstorm(sz),
    "13d": lambda sz: draw_snow(sz),
    "13n": lambda sz: draw_snow(sz),
    "50d": lambda sz: draw_fog(sz),
    "50n": lambda sz: draw_fog(sz),
}

_cache = {}


def get_icon(icon_code: str, size: int = 80) -> pygame.Surface:
    """Get a weather icon surface for the given OWM icon code (e.g. '01d', '04n')."""
    # Normalise to 3-char key; fall back to day variant if suffix missing
    key = icon_code[:3] if len(icon_code) >= 3 else icon_code[:2] + "d"
    cache_key = (key, size)
    if cache_key not in _cache:
        fn = ICON_MAP.get(key) or ICON_MAP.get(key[:2] + "d") or (lambda sz: draw_cloud(sz))
        _cache[cache_key] = fn(size)
    return _cache[cache_key]


def clear_cache():
    _cache.clear()