import os
import numpy as np
import rasterio
from rasterio.plot import show
from scipy.ndimage import median_filter
import matplotlib.pyplot as plt

def load_and_preprocess_sar(image_path):
    """
    Loads a preprocessed Sentinel-1 SAR GeoTIFF slice 
    and applies secondary spatial filtering for flood analysis.
    """
    with rasterio.open(image_path) as src:
        meta = src.meta.copy()
        # Read VV polarization band (standard for flood mapping)
        sar_data = src.read(1).astype(np.float32)
        transform = src.transform
        crs = src.crs
        
    # Convert DN to decibels if not already scaled: dB = 10 * log10(intensity)
    # Assuming input is linear power, convert safely
    with np.errstate(divide='ignore', invalid='ignore'):
        sar_db = 10.0 * np.log10(sar_data)
        sar_db[np.isinf(sar_db)] = np.nan
        
    # Apply a median filter to suppress remaining speckle texture
    filtered_db = median_filter(np.nan_to_num(sar_db, nan=0.0), size=5)
    
    return filtered_db, meta, transform, crs

# Example invocation path placeholders
# pre_path = "data/sentinel1_pre_flood.tif"
# post_path = "data/sentinel1_post_flood.tif"

from skimage.filters import threshold_otsu

def extract_flood_extent(pre_db, post_db, threshold_offset=-3.0):
    """
    Performs change detection between pre- and post-flood SAR imagery
    and applies Otsu thresholding to isolate flooded pixels.
    """
    # Calculate backscatter difference (Post - Pre)
    # Significant drops in dB indicate newly emerged water surfaces
    diff_image = post_db - pre_db
    
    # Valid pixels only
    valid_mask = ~np.isnan(diff_image) & (post_db != 0)
    valid_diff = diff_image[valid_mask]
    
    if valid_diff.size == 0:
        raise ValueError("No valid overlapping pixels found in SAR rasters.")
        
    # Calculate automated threshold using scikit-image Otsu method
    try:
        thresh = threshold_otsu(valid_diff)
    except Exception:
        thresh = -3.0 # Fallback empirical threshold
        
    print(f"Calculated optimal change detection threshold: {thresh:.2f} dB")
    
    # Binary classification: areas with drop greater than threshold are potential floods
    flood_mask = np.zeros_input = np.zeros(diff_image.shape, dtype=np.uint8)
    flood_mask[(diff_image <= thresh) & valid_mask] = 1
    
    return flood_mask, diff_image

def mask_permanent_water(flood_mask, permanent_water_raster_path):
    """
    Removes permanent water bodies from the candidate flood mask 
    to retain strictly temporary flood inundation zones.
    """
    with rasterio.open(permanent_water_raster_path) as perm_src:
        perm_water = perm_src.read(1)
        
    # Ensure binary mask format (1 = permanent water, 0 = land)
    perm_water_binary = np.where(perm_water > 0, 1, 0)
    
    # Exclude permanent water from the dynamic flood mask
    isolated_flood_mask = np.where((flood_mask == 1) & (perm_water_binary == 0), 1, 0)
    
    return isolated_flood_mask

def export_flood_vector(flood_mask, transform, crs, output_geojson_path):
    """
    Vectorizes the raster flood inundation mask into standard GeoJSON polygons.
    """
    import geopandas as gpd
    from rasterio.features import shapes
    
    shape_generator = shapes(
        flood_mask.astype(np.uint8), 
        mask=(flood_mask == 1), 
        transform=transform
    )
    
    records = []
    for geom, value in shape_generator:
        records.append({
            "type": "Feature",
            "geometry": geom,
            "properties": {"grid_code": int(value)}
        })
        
    if records:
        gdf = gpd.GeoDataFrame.from_features(records, crs=crs)
        gdf.to_file(output_geojson_path, driver="GeoJSON")
        print(f"Successfully exported flood polygons to {output_geojson_path}")
    else:
        print("No flood features detected above threshold.")