"""
gee_codegen.py — Chronostasis GEE Code Generator
==================================================
Generates ready-to-run Google Earth Engine JavaScript code
for SAR flood analysis of any of the 15 supported Indian river basins.

Users paste the output into code.earthengine.google.com and get:
  - Sentinel-1 SAR VV composite (2022-2024 monsoon)
  - Flood extent detection via backscatter threshold
  - CHIRPS rainfall overlay
  - HydroSHEDS flow accumulation
  - Risk zone classification (high/moderate/low)
  - Export-ready flood masks as GeoTIFF
"""

from typing import Dict, Any

# Basin bounding boxes for GEE geometry
BASIN_BOUNDS = {
    "brahmaputra": {"minLon": 89.5,  "maxLon": 96.0,  "minLat": 24.5, "maxLat": 28.5},
    "ganga":       {"minLon": 78.0,  "maxLon": 89.0,  "minLat": 23.0, "maxLat": 28.0},
    "mahanadi":    {"minLon": 80.5,  "maxLon": 87.5,  "minLat": 18.5, "maxLat": 23.5},
    "krishna":     {"minLon": 73.5,  "maxLon": 81.5,  "minLat": 14.5, "maxLat": 19.5},
    "godavari":    {"minLon": 73.5,  "maxLon": 83.5,  "minLat": 15.0, "maxLat": 22.0},
    "narmada":     {"minLon": 72.5,  "maxLon": 81.5,  "minLat": 21.0, "maxLat": 24.5},
    "tapti":       {"minLon": 72.5,  "maxLon": 78.5,  "minLat": 19.5, "maxLat": 22.5},
    "cauvery":     {"minLon": 75.0,  "maxLon": 80.5,  "minLat": 10.0, "maxLat": 15.0},
    "damodar":     {"minLon": 84.5,  "maxLon": 88.5,  "minLat": 22.0, "maxLat": 24.5},
    "sabarmati":   {"minLon": 71.5,  "maxLon": 74.5,  "minLat": 21.5, "maxLat": 24.5},
    "mahi":        {"minLon": 72.5,  "maxLon": 74.5,  "minLat": 21.5, "maxLat": 24.5},
    "baitarani":   {"minLon": 84.5,  "maxLon": 87.5,  "minLat": 20.0, "maxLat": 23.0},
    "subarnarekha":{"minLon": 85.0,  "maxLon": 88.0,  "minLat": 21.5, "maxLat": 23.5},
    "indus":       {"minLon": 72.5,  "maxLon": 77.5,  "minLat": 29.0, "maxLat": 33.0},
    "luni":        {"minLon": 70.5,  "maxLon": 75.5,  "minLat": 24.5, "maxLat": 28.5},
}


def generate_gee_code(region_id: str, region_data: Dict[str, Any],
                       year: int = 2022) -> str:
    """
    Generates complete GEE JavaScript for SAR flood analysis of a river basin.

    Args:
        region_id:   basin identifier (e.g. 'brahmaputra')
        region_data: full region dict from tasks.REGIONS
        year:        analysis year (2022, 2023, or 2024)

    Returns:
        Complete GEE JavaScript string ready to paste into code.earthengine.google.com
    """
    bounds      = BASIN_BOUNDS.get(region_id, BASIN_BOUNDS["brahmaputra"])
    name        = region_data["name"]
    river       = region_data["river"]
    state       = region_data["state"]
    sar_thresh  = region_data.get("sar_threshold_db", -16)
    peak_year   = region_data["peak_year"]
    accuracy    = region_data["accuracy_pct"]
    chronic_km2 = region_data["chronic_km2"]
    fa          = region_data["flood_areas"]
    flood_km2   = fa.get(year, fa.get(str(year), 0))
    districts   = region_data["chronic_districts"]
    risk_zones  = region_data["risk_zones_km2"]
    rainfall    = region_data["peak_rainfall_mm"]
    pop         = region_data["chronic_pop"]

    # All 3 years for comparison
    years_str = ", ".join(str(y) for y in [2022, 2023, 2024])
    flood_areas_str = " | ".join(
        f"{y}: {fa.get(y, fa.get(str(y), 0))} km²" for y in [2022, 2023, 2024])

    code = f"""// ============================================================
// CHRONOSTASIS — SAR Flood Intelligence
// Basin:  {name}
// River:  {river}
// State:  {state}
// Year:   {year}
// Generated by Chronostasis OpenEnv v2.0
// https://huggingface.co/spaces/LunaAmagi/chronostasis
// ============================================================
// Model accuracy: {accuracy}%
// Flood extents: {flood_areas_str}
// Peak flood year: {peak_year}
// Chronic inundation area: {chronic_km2} km²
// Population at risk: {pop:,}
// Peak rainfall: {rainfall} mm
// ============================================================

// ── 1. STUDY AREA ────────────────────────────────────────────
var basin = ee.Geometry.Rectangle([
  {bounds['minLon']}, {bounds['minLat']},
  {bounds['maxLon']}, {bounds['maxLat']}
]);

Map.centerObject(basin, 7);
Map.setOptions('HYBRID');

print('=== CHRONOSTASIS: {name} ===');
print('Analysis year: {year}');
print('SAR threshold: {sar_thresh} dB');
print('Expected flood extent: {flood_km2:,.0f} km²');

// ── 2. SENTINEL-1 SAR VV — MONSOON COMPOSITES ─────────────────
// Load SAR data for monsoon season (June–September) for all years
var sarBase = ee.ImageCollection('COPERNICUS/S1_GRD')
  .filterBounds(basin)
  .filter(ee.Filter.eq('instrumentMode', 'IW'))
  .filter(ee.Filter.listContains('transmitterReceiverPolarisation', 'VV'))
  .filter(ee.Filter.eq('orbitProperties_pass', 'DESCENDING'))
  .select('VV');

// Per-year monsoon composites
var getSAR = function(yr) {{
  return sarBase
    .filterDate(yr + '-06-01', yr + '-09-30')
    .median()
    .clip(basin);
}};

var sar2022 = getSAR('2022');
var sar2023 = getSAR('2023');
var sar2024 = getSAR('2024');
var sarYear = getSAR('{year}');  // Selected analysis year

print('SAR composite loaded for year {year}');

// ── 3. PRE-FLOOD BASELINE (dry season reference) ───────────────
var sarDry = sarBase
  .filterDate('{year - 1}-11-01', '{year}-03-31')
  .median()
  .clip(basin);

// ── 4. PERMANENT WATER MASK (Landsat 8 NDWI) ──────────────────
var landsat = ee.ImageCollection('LANDSAT/LC08/C02/T1_L2')
  .filterBounds(basin)
  .filterDate('2020-01-01', '2022-05-31')
  .filter(ee.Filter.lt('CLOUD_COVER', 20))
  .map(function(img) {{
    var nir   = img.select('SR_B5').multiply(0.0000275).add(-0.2);
    var green = img.select('SR_B3').multiply(0.0000275).add(-0.2);
    return img.addBands(nir.subtract(green).divide(nir.add(green))
             .rename('NDWI'));
  }});

var permWater = landsat.select('NDWI').median().gt(0.3).clip(basin);
print('Permanent water mask computed');

// ── 5. FLOOD DETECTION (SAR backscatter threshold) ────────────
// Threshold: {sar_thresh} dB — calibrated for {name}
var sarThreshold = {sar_thresh};

var floodRaw = sarYear.lt(sarThreshold)
  .and(sarDry.gt(sarThreshold))   // changed from dry = new flood
  .and(permWater.not());          // exclude permanent water

// Morphological cleaning — remove speckle
var floodMask = floodRaw
  .focal_min({{radius: 1, kernelType: 'square', units: 'pixels'}})
  .focal_max({{radius: 1, kernelType: 'square', units: 'pixels'}});

print('Flood mask computed');

// ── 6. COMPUTE FLOOD EXTENT ────────────────────────────────────
var pixelArea  = ee.Image.pixelArea().divide(1e6); // km²
var floodArea  = floodMask.multiply(pixelArea).reduceRegion({{
  reducer:  ee.Reducer.sum(),
  geometry: basin,
  scale:    30,
  maxPixels: 1e10,
}});
print('Computed flood extent (km²):', floodArea);

// ── 7. YEAR-ON-YEAR COMPARISON ─────────────────────────────────
var flood2022 = sar2022.lt(sarThreshold).and(permWater.not())
  .focal_min({{radius:1,kernelType:'square',units:'pixels'}})
  .focal_max({{radius:1,kernelType:'square',units:'pixels'}});
var flood2023 = sar2023.lt(sarThreshold).and(permWater.not())
  .focal_min({{radius:1,kernelType:'square',units:'pixels'}})
  .focal_max({{radius:1,kernelType:'square',units:'pixels'}});
var flood2024 = sar2024.lt(sarThreshold).and(permWater.not())
  .focal_min({{radius:1,kernelType:'square',units:'pixels'}})
  .focal_max({{radius:1,kernelType:'square',units:'pixels'}});

// Chronic inundation — flooded in ALL 3 years
var chronicFlood = flood2022.and(flood2023).and(flood2024);
var chronicArea  = chronicFlood.multiply(pixelArea).reduceRegion({{
  reducer:  ee.Reducer.sum(),
  geometry: basin,
  scale:    30,
  maxPixels: 1e10,
}});
print('Chronic inundation area (km²) — expected {chronic_km2} km²:', chronicArea);

// ── 8. CHIRPS RAINFALL OVERLAY ────────────────────────────────
var chirps = ee.ImageCollection('UCSB-CHG/CHIRPS/DAILY')
  .filterBounds(basin)
  .filterDate('{year}-06-01', '{year}-09-30')
  .sum()
  .clip(basin);

print('CHIRPS total monsoon rainfall loaded ({year})');

// ── 9. TERRAIN — DEM + SLOPE + FLOW ACCUMULATION ─────────────
var dem   = ee.Image('USGS/SRTMGL1_003').clip(basin);
var slope = ee.Terrain.slope(dem);

var hydro = ee.Image('WWF/HydroSHEDS/15ACC').clip(basin);
var flowAcc = hydro.select('b1');

// ── 10. MULTI-FACTOR RISK MODEL ────────────────────────────────
// Combines: flood frequency + rainfall + low elevation + high flow
var floodFreq = flood2022.add(flood2023).add(flood2024).rename('freq');

// Normalise each layer to 0-1
var normalize = function(img, min, max) {{
  return img.subtract(min).divide(max - min).clamp(0, 1);
}};

var freqNorm    = normalize(floodFreq,   0,    3);
var rainfallNorm= normalize(chirps,      400,  {rainfall + 200});
var elevNorm    = ee.Image(1).subtract(normalize(dem, 0, 500));  // invert — low = high risk
var flowNorm    = normalize(flowAcc.log(), 0,  12);

// Weighted composite — mirrors Chronostasis accuracy {accuracy}%
var riskScore = freqNorm.multiply(0.35)
  .add(rainfallNorm.multiply(0.25))
  .add(elevNorm.multiply(0.20))
  .add(flowNorm.multiply(0.20));

// Classify into 3 zones
// High risk (score > 0.65):  {risk_zones['high']:,.0f} km²
// Moderate  (0.35–0.65):     {risk_zones['moderate']:,.0f} km²
// Low risk  (<0.35):         {risk_zones['low']:,.0f} km²
var highRisk  = riskScore.gt(0.65);
var modRisk   = riskScore.gt(0.35).and(riskScore.lte(0.65));
var lowRisk   = riskScore.lte(0.35);

var riskZones = highRisk.multiply(3)
  .add(modRisk.multiply(2))
  .add(lowRisk.multiply(1))
  .rename('risk_class');

print('Risk model computed. Expected zones:');
print('  High risk:  {risk_zones['high']:,.0f} km²');
print('  Moderate:   {risk_zones['moderate']:,.0f} km²');
print('  Low risk:   {risk_zones['low']:,.0f} km²');

// ── 11. VISUALISATION ─────────────────────────────────────────
var sarVis  = {{min: -25, max: 0, palette: ['000000', '404040', 'ffffff']}};
var floodVis= {{min: 0,  max: 1, palette: ['000000', '0000ff']}};
var riskVis = {{min: 1,  max: 3, palette: ['00aa00', 'ffaa00', 'ff0000']}};
var rainVis = {{min: 400,max: {rainfall + 400}, palette: ['ffffff', '0080ff', '000090']}};
var demVis  = {{min: 0,  max: 500, palette: ['006600','ffff00','ffffff']}};

// SAR composite
Map.addLayer(sarYear, sarVis, 'SAR VV — {year} (Monsoon)', true);

// Flood layers
Map.addLayer(flood2022, {{palette: ['0000ff']}}, 'Flood 2022', false);
Map.addLayer(flood2023, {{palette: ['00aaff']}}, 'Flood 2023', false);
Map.addLayer(flood2024, {{palette: ['00ffff']}}, 'Flood 2024', false);
Map.addLayer(floodMask, {{palette: ['0044ff']}}, 'Flood {year} (selected)', true);
Map.addLayer(chronicFlood.selfMask(), {{palette: ['ff0000']}}, 'Chronic Inundation (all 3yr)', true);

// Risk and terrain
Map.addLayer(riskZones.selfMask(), riskVis, 'Multi-Factor Risk Zones', true);
Map.addLayer(chirps,   rainVis, 'CHIRPS Rainfall {year}', false);
Map.addLayer(dem,      demVis,  'SRTM Elevation', false);
Map.addLayer(slope.gt(5).selfMask(), {{palette:['ff8800']}}, 'Slope > 5°', false);
Map.addLayer(flowAcc.log().gt(8).selfMask(), {{palette:['9900ff']}},
             'HydroSHEDS High Flow Accumulation', false);

// Permanent water
Map.addLayer(permWater.selfMask(), {{palette:['00aaff']}}, 'Permanent Water (Landsat NDWI)', false);

// Basin boundary
Map.addLayer(ee.Image().paint(basin, 0, 2),
             {{palette:['ffff00']}}, '{name} Basin Boundary', true);

// ── 12. LEGEND ────────────────────────────────────────────────
var legend = ui.Panel({{style: {{position:'bottom-left', padding:'8px 12px'}}}});
legend.add(ui.Label('CHRONOSTASIS — {name}',
  {{fontWeight:'bold', fontSize:'14px', color:'#1a1a2e'}}));
legend.add(ui.Label('River: {river} | State: {state}',
  {{fontSize:'11px', color:'#555'}}));
legend.add(ui.Label('Year: {year} | Flood: {flood_km2:,.0f} km² | Accuracy: {accuracy}%',
  {{fontSize:'11px', color:'#333'}}));
legend.add(ui.Label(''));

var colors = ['ff0000','ffaa00','00aa00','0044ff','ff3300'];
var labels = [
  'High Risk Zone ({risk_zones['high']:,.0f} km²)',
  'Moderate Risk ({risk_zones['moderate']:,.0f} km²)',
  'Low Risk ({risk_zones['low']:,.0f} km²)',
  'Flood Extent {year}',
  'Chronic Inundation (3yr)',
];
labels.forEach(function(lbl, i) {{
  var row = ui.Panel([
    ui.Label('', {{backgroundColor:'#'+colors[i],
                   padding:'6px',margin:'2px 6px 2px 2px'}}),
    ui.Label(lbl, {{fontSize:'11px',margin:'2px 0'}})
  ], ui.Panel.Layout.flow('horizontal'));
  legend.add(row);
}});

legend.add(ui.Label(''));
legend.add(ui.Label('Chronic districts: {", ".join(districts[:3])}...',
  {{fontSize:'10px', color:'#666'}}));
legend.add(ui.Label('Population at risk: {pop:,}',
  {{fontSize:'10px', color:'#666'}}));
legend.add(ui.Label(''));
legend.add(ui.Label('Generated by Chronostasis OpenEnv v2.0',
  {{fontSize:'9px', color:'#999'}}));
legend.add(ui.Label('huggingface.co/spaces/LunaAmagi/chronostasis',
  {{fontSize:'9px', color:'#999'}}));
Map.add(legend);

// ── 13. EXPORT FLOOD MASK ──────────────────────────────────────
// Uncomment to export the flood extent as a GeoTIFF
/*
Export.image.toDrive({{
  image:        floodMask.toFloat(),
  description:  'chronostasis_{region_id}_flood_{year}',
  folder:       'chronostasis',
  scale:        30,
  region:       basin,
  crs:          'EPSG:4326',
  maxPixels:    1e10,
}});

Export.image.toDrive({{
  image:        riskZones.toFloat(),
  description:  'chronostasis_{region_id}_risk_zones',
  folder:       'chronostasis',
  scale:        30,
  region:       basin,
  crs:          'EPSG:4326',
  maxPixels:    1e10,
}});

Export.image.toDrive({{
  image:        chronicFlood.toFloat(),
  description:  'chronostasis_{region_id}_chronic_{year_start}_{year_end}',
  folder:       'chronostasis',
  scale:        30,
  region:       basin,
  crs:          'EPSG:4326',
  maxPixels:    1e10,
}});
*/

// ── 14. CHART — Year-on-Year Flood Area ───────────────────────
var yearlyStats = ee.List([2022, 2023, 2024]).map(function(yr) {{
  yr = ee.Number(yr);
  var s = sarBase
    .filterDate(ee.Date.fromYMD(yr, 6, 1), ee.Date.fromYMD(yr, 9, 30))
    .median().clip(basin);
  var f = s.lt(sarThreshold).and(permWater.not())
    .focal_min({{radius:1,kernelType:'square',units:'pixels'}})
    .focal_max({{radius:1,kernelType:'square',units:'pixels'}});
  var area = f.multiply(pixelArea).reduceRegion({{
    reducer: ee.Reducer.sum(), geometry: basin, scale: 100, maxPixels:1e9
  }});
  return ee.Feature(null, {{
    'year':      yr,
    'flood_km2': area.get('VV'),
  }});
}});

var chart = ui.Chart.feature.byFeature(
  ee.FeatureCollection(yearlyStats), 'year', ['flood_km2'])
  .setChartType('ColumnChart')
  .setOptions({{
    title:  '{name} — SAR Flood Extent 2022–2024',
    hAxis:  {{title: 'Year'}},
    vAxis:  {{title: 'Flood Extent (km²)'}},
    colors: ['0044ff'],
    legend: {{position: 'none'}},
  }});
print(chart);
"""
    return code.strip()


def generate_multi_basin_comparison_code(regions_data: Dict) -> str:
    """
    Generates GEE code that compares flood extents across
    all 15 basins in a single map.
    """
    # Build basin geometries
    basin_defs = []
    for rid, bounds in BASIN_BOUNDS.items():
        if rid in regions_data:
            r = regions_data[rid]
            basin_defs.append(
                f"  ee.Feature(ee.Geometry.Rectangle("
                f"[{bounds['minLon']},{bounds['minLat']},"
                f"{bounds['maxLon']},{bounds['maxLat']}]), "
                f"{{'name': '{r['name']}', 'river': '{r['river']}', "
                f"'peak_year': {r['peak_year']}, "
                f"'chronic_km2': {r['chronic_km2']}, "
                f"'accuracy_pct': {r['accuracy_pct']}}})"
            )

    basins_js = ",\n".join(basin_defs)

    return f"""// ============================================================
// CHRONOSTASIS — All-India 15 Basin Flood Comparison
// Compares SAR flood extents across all 15 river basins
// Generated by Chronostasis OpenEnv v2.0
// https://huggingface.co/spaces/LunaAmagi/chronostasis
// ============================================================

// ── All 15 basin definitions ──────────────────────────────────
var basins = ee.FeatureCollection([
{basins_js}
]);

Map.setCenter(80.0, 22.0, 5);
Map.setOptions('HYBRID');

// ── SAR collection ────────────────────────────────────────────
var sarBase = ee.ImageCollection('COPERNICUS/S1_GRD')
  .filterBounds(basins.geometry())
  .filter(ee.Filter.eq('instrumentMode', 'IW'))
  .filter(ee.Filter.listContains('transmitterReceiverPolarisation', 'VV'))
  .filter(ee.Filter.eq('orbitProperties_pass', 'DESCENDING'))
  .select('VV');

var sar2022 = sarBase.filterDate('2022-06-01','2022-09-30').median();
var sar2023 = sarBase.filterDate('2023-06-01','2023-09-30').median();
var sar2024 = sarBase.filterDate('2024-06-01','2024-09-30').median();

var sarThreshold = -16;

// ── Flood masks ───────────────────────────────────────────────
var permWater = ee.ImageCollection('LANDSAT/LC08/C02/T1_L2')
  .filterBounds(basins.geometry())
  .filterDate('2020-01-01','2022-05-31')
  .filter(ee.Filter.lt('CLOUD_COVER',20))
  .map(function(img) {{
    var nir   = img.select('SR_B5').multiply(0.0000275).add(-0.2);
    var green = img.select('SR_B3').multiply(0.0000275).add(-0.2);
    return nir.subtract(green).divide(nir.add(green)).rename('NDWI');
  }}).median().gt(0.3);

var makeFlood = function(sar) {{
  return sar.lt(sarThreshold).and(permWater.not())
    .focal_min({{radius:1,kernelType:'square',units:'pixels'}})
    .focal_max({{radius:1,kernelType:'square',units:'pixels'}});
}};

var flood2022 = makeFlood(sar2022);
var flood2023 = makeFlood(sar2023);
var flood2024 = makeFlood(sar2024);

// Chronic = flooded all 3 years
var chronic = flood2022.and(flood2023).and(flood2024);

// Peak year composite = any year
var anyFlood = flood2022.or(flood2023).or(flood2024);

// ── Risk model ────────────────────────────────────────────────
var dem      = ee.Image('USGS/SRTMGL1_003');
var flowAcc  = ee.Image('WWF/HydroSHEDS/15ACC').select('b1');
var chirps   = ee.ImageCollection('UCSB-CHG/CHIRPS/DAILY')
  .filterDate('2022-06-01','2022-09-30').sum();

var floodFreq   = flood2022.add(flood2023).add(flood2024).rename('freq');
var freqNorm    = floodFreq.divide(3);
var rainfallNorm= chirps.subtract(400).divide(1400).clamp(0,1);
var elevNorm    = ee.Image(1).subtract(dem.divide(500).clamp(0,1));
var flowNorm    = flowAcc.log().divide(12).clamp(0,1);

var riskScore = freqNorm.multiply(0.35)
  .add(rainfallNorm.multiply(0.25))
  .add(elevNorm.multiply(0.20))
  .add(flowNorm.multiply(0.20));

var riskZones = riskScore.gt(0.65).multiply(3)
  .add(riskScore.gt(0.35).and(riskScore.lte(0.65)).multiply(2))
  .add(riskScore.lte(0.35).multiply(1))
  .rename('risk_class');

// ── Visualise ─────────────────────────────────────────────────
Map.addLayer(riskZones.selfMask().clip(basins.geometry()),
  {{min:1,max:3,palette:['00aa00','ffaa00','ff0000']}}, 'All-India Risk Zones', true);
Map.addLayer(anyFlood.selfMask().clip(basins.geometry()),
  {{palette:['0044ff']}}, 'Any-Year Flood (2022-2024)', false);
Map.addLayer(chronic.selfMask().clip(basins.geometry()),
  {{palette:['ff0000']}}, 'Chronic Inundation (all 3yr)', true);
Map.addLayer(ee.Image().paint(basins, 0, 2),
  {{palette:['ffff00']}}, '15 Basin Boundaries', true);

print('=== CHRONOSTASIS — All-India 15 Basin Analysis ===');
print('Basins:', basins.aggregate_array('name'));

// Chart: chronic area per basin
var basinStats = basins.map(function(feat) {{
  var geom = feat.geometry();
  var area = chronic.multiply(ee.Image.pixelArea().divide(1e6))
    .reduceRegion({{reducer:ee.Reducer.sum(),geometry:geom,scale:100,maxPixels:1e9}});
  return feat.set('chronic_sar_km2', area.get('VV'));
}});

var chart = ui.Chart.feature.byFeature(basinStats, 'name', ['chronic_sar_km2'])
  .setChartType('ColumnChart')
  .setOptions({{
    title:  'Chronic Inundation Area — All 15 Indian River Basins',
    hAxis:  {{title:'Basin', slantedText:true}},
    vAxis:  {{title:'Chronic Inundation (km²)'}},
    colors: ['ff0000'],
    legend: {{position:'none'}},
  }});
print(chart);
""".strip()