import json
import random
import math
from pathlib import Path
from PIL import Image, ImageFilter

# Maharashtra approximate bounds
LAT_MIN, LAT_MAX = 15.6, 22.0
LON_MIN, LON_MAX = 72.6, 80.9

# For the PNG, we want high resolution
# width = degrees * pixels_per_degree
PIXELS_PER_DEGREE = 50 # 50 pixels per degree gives ~415x320 image, which is fast to generate
WIDTH = int((LON_MAX - LON_MIN) * PIXELS_PER_DEGREE)
HEIGHT = int((LAT_MAX - LAT_MIN) * PIXELS_PER_DEGREE)

def get_temp_color(temp):
    if temp <= 22: return (0, 0, 0, 0)
    if temp <= 25: return (254, 217, 118, 150)
    if temp <= 30: return (254, 178, 76, 180)
    if temp <= 35: return (253, 141, 60, 200)
    if temp <= 40: return (252, 78, 42, 220)
    if temp <= 43: return (227, 26, 28, 240)
    return (177, 0, 38, 255)

def get_rain_color(rain):
    if rain <= 1: return (0, 0, 0, 0)
    if rain <= 10: return (198, 219, 239, 100)
    if rain <= 40: return (158, 202, 225, 150)
    if rain <= 80: return (107, 174, 214, 180)
    if rain <= 120: return (66, 146, 198, 200)
    if rain <= 160: return (33, 113, 181, 220)
    return (8, 69, 148, 255)

def generate_rasters(out_dir):
    # Create empty images (RGBA)
    temp_img = Image.new('RGBA', (WIDTH, HEIGHT), (0, 0, 0, 0))
    rain_img = Image.new('RGBA', (WIDTH, HEIGHT), (0, 0, 0, 0))
    
    temp_pixels = temp_img.load()
    rain_pixels = rain_img.load()
    
    # Iterate over pixels
    for y in range(HEIGHT):
        # Image Y=0 is TOP (LAT_MAX), Y=HEIGHT is BOTTOM (LAT_MIN)
        lat = LAT_MAX - (y / PIXELS_PER_DEGREE)
        for x in range(WIDTH):
            # Image X=0 is LEFT (LON_MIN), X=WIDTH is RIGHT (LON_MAX)
            lon = LON_MIN + (x / PIXELS_PER_DEGREE)
            
            # Patchy rain simulation using overlapping sine waves (fake Perlin noise)
            patchiness = math.sin(lat * 5.0) * math.cos(lon * 5.0) + math.sin(lat * 2.1 + lon * 3.3)
            
            # Base temperature: hotter inland (east)
            temp_base = 25 + (lon - LON_MIN) * 2.5
            
            # Base rain
            if patchiness > 0.8:
                rain = random.uniform(50, 200) # Heavy patch
            elif patchiness > 0.2:
                rain = random.uniform(5, 50)   # Light patch
            else:
                rain = 0
                
            temp = temp_base + random.uniform(-2, 2)
            temp = max(20, min(temp, 45))
            
            temp_pixels[x, y] = get_temp_color(temp)
            rain_pixels[x, y] = get_rain_color(rain)
            
    # Apply a slight blur so the noise isn't too sharp, creating a smooth raster feel
    temp_img = temp_img.filter(ImageFilter.GaussianBlur(radius=3))
    rain_img = rain_img.filter(ImageFilter.GaussianBlur(radius=5))
            
    temp_img.save(out_dir / "temp_raster.png")
    rain_img.save(out_dir / "rain_raster.png")
    print(f"Saved temp_raster.png and rain_raster.png ({WIDTH}x{HEIGHT}) to {out_dir}")

if __name__ == "__main__":
    out_dir = Path(__file__).parent.parent / "frontend" / "public" / "data"
    out_dir.mkdir(parents=True, exist_ok=True)
    generate_rasters(out_dir)