import numpy as np
import rasterio
from rasterio.transform import from_origin
from rasterio.crs import CRS
import os

def create_synthetic_dem(output_path, width=256, height=256):
    """
    Creates a synthetic DEM (Digital Elevation Model) GeoTIFF.
    Simulates a peak near Mount Everest.
    """
    # 1. Define bounds (centered around Everest: 86.9250, 27.9881)
    # Spanning roughly 0.1 degrees
    lon_min, lat_min = 86.875, 27.938
    pixel_size = 0.0004 # Approx 40m resolution

    # 2. Create synthetic elevation data (Gaussian mountain)
    x = np.linspace(-2, 2, width)
    y = np.linspace(-2, 2, height)
    X, Y = np.meshgrid(x, y)
    
    # Gaussian function for main peak
    Z = np.exp(-(X**2 + Y**2)) * 8848 
    
    # Add some noise/roughness
    noise = np.random.normal(0, 100, (height, width))
    Z = Z + noise
    
    # Ensure minimum elevation is above 0 (e.g., base at 4000m)
    Z = np.maximum(Z, 4000)
    
    # Convert to appropriate data type (int16 or float32)
    Z = Z.astype('int16')

    # 3. Define Transform and CRS
    transform = from_origin(lon_min, lat_min + (height * pixel_size), pixel_size, pixel_size)
    crs = CRS.from_epsg(4326) # WGS84

    # 4. Write to GeoTIFF
    print(f"Writing synthetic DEM to {output_path}...")
    with rasterio.open(
        output_path,
        'w',
        driver='GTiff',
        height=height,
        width=width,
        count=1,
        dtype=Z.dtype,
        crs=crs,
        transform=transform,
    ) as dst:
        dst.write(Z, 1)
    print("Done.")

if __name__ == "__main__":
    # Ensure directories exist
    os.makedirs("data/raw", exist_ok=True)
    
    output_file = "data/raw/synthetic_everest.tif"
    create_synthetic_dem(output_file)