import os
import numpy as np
from flask import Flask, request, jsonify
import rasterio
from rasterio.fill import fillnodata
from sklearn.preprocessing import StandardScaler

app = Flask(__name__)

time_steps_image = 4
target_img_size = (224, 224)

def preprocess_dem(dem_path):
    with rasterio.open(dem_path) as src:
        dem_data = src.read(1)
        filled_dem = fillnodata(dem_data, mask=src.read_masks(1))

        processed_dem_path = "processed_dem.tif"
        with rasterio.open(
            processed_dem_path, "w", **src.meta
        ) as dest:
            dest.write(filled_dem, 1)

    return processed_dem_path

def simulate_flood(dem_path, lake_elevation, output_path):
    with rasterio.open(dem_path) as src:
        dem_data = src.read(1)
        meta = src.meta

    flood_map = np.where(dem_data <= lake_elevation, 1, 0)

    with rasterio.open(output_path, "w", **meta) as dest:
        dest.write(flood_map.astype(rasterio.uint8), 1)

    return output_path

@app.route('/predict_flood', methods=['POST'])
def predict_flood():
    try:
        dem_file = request.files['dem']
        lake_elevation = float(request.form['lake_elevation'])

        dem_path = os.path.join("uploads", "input_dem.tif")
        os.makedirs("uploads", exist_ok=True)
        dem_file.save(dem_path)

        processed_dem_path = preprocess_dem(dem_path)

        output_dir = "outputs"
        os.makedirs(output_dir, exist_ok=True)
        flood_map_path = os.path.join(output_dir, "flood_map.tif")
        simulate_flood(processed_dem_path, lake_elevation, flood_map_path)

        return jsonify({"message": "Flood map generated", "flood_map_path": flood_map_path})

    except Exception as e:
        return jsonify({"error": str(e)}), 500

if __name__ == "__main__":
    app.run(debug=True)