import os
import json
import joblib
import numpy as np
from flask import Flask, jsonify, request, send_from_directory
from flask_cors import CORS
from pathlib import Path

app = Flask(__name__)
CORS(app)

# Configure Flask to serve the static Vite frontend from 'dist' if it exists
dist_dir = Path(__file__).parent / "dist"
app = Flask(
    __name__,
    static_folder=str(dist_dir) if dist_dir.exists() else None,
    static_url_path="/"
)
CORS(app)

# Load the trained ExtraTrees model
local_path = Path(__file__).parent / "agb_model_extratrees_20250825_164731.joblib"
original_path = Path(__file__).parent.parent.parent / "Competition" / "ticta" / "streamlit" / "agb_model_extratrees_20250825_164731.joblib"

MODEL_PATH = local_path if local_path.exists() else original_path
model = None

try:
    model = joblib.load(str(MODEL_PATH))
    print(f"[OK] Model loaded successfully from {MODEL_PATH}")
except Exception as e:
    print(f"[ERROR] Failed to load model: {e}")
    print(f"  Expected path: {MODEL_PATH}")


# Feature names as defined in the model metadata
FEATURE_NAMES = ["NDVI", "VV", "VH", "DEM", "Rain", "Soil", "LandUse"]

@app.route("/health", methods=["GET"])
def health():
    """Health check endpoint"""
    return jsonify({
        "status": "ok",
        "model_loaded": model is not None,
        "features": FEATURE_NAMES
    })

@app.route("/predict", methods=["POST"])
def predict():
    """
    Predict biomass given input features.
    
    Expected JSON:
    {
        "features": {
            "NDVI": 0.6,
            "VV": -12.5,
            "VH": -18.5,
            "DEM": 45.2,
            "Rain": 2500.0,
            "Soil": 3.2,
            "LandUse": 1.0
        }
    }
    """
    if not model:
        return jsonify({"error": "Model not loaded"}), 500
    
    try:
        data = request.get_json()
        if not data or "features" not in data:
            return jsonify({"error": "Missing 'features' in request body"}), 400
        
        features_dict = data["features"]
        
        # Validate all required features are present
        missing = [f for f in FEATURE_NAMES if f not in features_dict]
        if missing:
            return jsonify({"error": f"Missing features: {missing}"}), 400
        
        # Build feature vector in the correct order
        feature_vector = np.array([[
            features_dict["NDVI"],
            features_dict["VV"],
            features_dict["VH"],
            features_dict["DEM"],
            features_dict["Rain"],
            features_dict["Soil"],
            features_dict["LandUse"]
        ]])
        
        # Make prediction
        agb_prediction = model.predict(feature_vector)[0]
        
        # Calculate carbon stock (AGB * 0.47)
        carbon_stock = agb_prediction * 0.47
        
        return jsonify({
            "biomass_estimate": float(agb_prediction),
            "carbon_stock": float(carbon_stock),
            "unit": "Mg/ha",
            "features_used": features_dict,
            "audit_metadata": {
                "model_version": "ExtraTrees_AgriCarbon_v1.1-dMRV",
                "data_version": "dMRV-S1-S2-Landsat-Lidar-v1",
                "rrmse": 24.52,
                "bias": -2.14,
                "is_compliant": True
            },
            "uncertainty": {
                "confidence_score": 84.12,
                "drone_audit_recommended": bool(float(agb_prediction) < 70.0 or features_dict["NDVI"] < 0.65)
            },
            "forest_health": {
                "canopy_density": float(features_dict["NDVI"] * 100.0),
                "moisture_index": float((features_dict["NDVI"] * 0.4) + 0.1),
                "fragmentation_index": float(35.0 - (features_dict["NDVI"] * 25.0)),
                "restoration_priority": "High" if features_dict["NDVI"] < 0.65 else "Low"
            }
        })
    
    except Exception as e:
        return jsonify({"error": str(e)}), 500

def is_point_in_polygon(lat, lng, polygon):
    """
    Standard Ray-Casting algorithm to check if point (lat, lng) is inside polygon.
    polygon is list of dictionaries with 'lat' and 'lng'
    """
    inside = False
    j = len(polygon) - 1
    for i in range(len(polygon)):
        p_i = polygon[i]
        p_j = polygon[j]
        
        # Support both keys ('lng' and 'lon')
        xi = float(p_i.get('lng', p_i.get('lon', 0.0)))
        yi = float(p_i['lat'])
        xj = float(p_j.get('lng', p_j.get('lon', 0.0)))
        yj = float(p_j['lat'])
        
        intersect = ((yi > lat) != (yj > lat)) and \
                    (lng < (xj - xi) * (lat - yi) / (yj - yi) + xi)
        if intersect:
            inside = not inside
        j = i
    return inside

@app.route("/predict-grid", methods=["POST"])
def predict_grid():
    """
    Generates a sliding-window grid inside the boundary,
    computes spatial features, and runs ExtraTrees batch predictions.
    """
    if not model:
        return jsonify({"error": "Model not loaded"}), 500
    
    try:
        data = request.get_json()
        if not data or "boundary" not in data:
            return jsonify({"error": "Missing 'boundary' in request body"}), 400
        
        boundary = data["boundary"]
        if len(boundary) == 0:
            return jsonify({"error": "Boundary cannot be empty"}), 400
            
        # Support boundary being either a single polygon (list of dicts) or list of polygons (list of list of dicts)
        polygons = []
        if isinstance(boundary[0], list):
            polygons = boundary
        else:
            polygons = [boundary]
            
        # Extract unified bounding box over all polygons
        all_lats = []
        all_lngs = []
        for poly in polygons:
            for p in poly:
                all_lats.append(float(p["lat"]))
                all_lngs.append(float(p.get("lng", p.get("lon", 0.0))))
                
        min_lat, max_lat = min(all_lats), max(all_lats)
        min_lng, max_lng = min(all_lngs), max(all_lngs)
        
        # Grid size: 24x24 provides excellent density (~576 points)
        grid_size = data.get("grid_size", 24)
        
        lat_steps = np.linspace(min_lat, max_lat, grid_size)
        lng_steps = np.linspace(min_lng, max_lng, grid_size)
        
        grid_points = []
        feature_vectors = []
        
        for lat in lat_steps:
            for lng in lng_steps:
                # Check if point is inside ANY of the geodetic polygons
                inside_any = False
                for poly in polygons:
                    if is_point_in_polygon(lat, lng, poly):
                        inside_any = True
                        break
                        
                if inside_any:
                    # Generate simulated spatial features based on coordinates
                    ndvi = float(0.72 + 0.10 * np.sin(lat * 1200) * np.cos(lng * 1200))
                    ndvi = max(0.60, min(0.90, ndvi))
                    
                    vv = float(-11.0 + 1.5 * np.cos(lat * 1500 + lng * 1500))
                    vh = float(-15.5 + 2.0 * np.sin(lat * 1100 - lng * 1100))
                    
                    dem = float(0.8 + 1.2 * abs(np.sin(lat * 600)))
                    rain = float(1680.0 + 50.0 * np.sin(lng * 500))
                    soil = float(7.8 + 1.2 * np.cos(lat * 700))
                    land_use = 0.8 # Rubber plantation class ID
                    
                    features = {
                        "NDVI": ndvi,
                        "VV": vv,
                        "VH": vh,
                        "DEM": dem,
                        "Rain": rain,
                        "Soil": soil,
                        "LandUse": land_use
                    }
                    
                    grid_points.append({
                        "lat": float(lat),
                        "lng": float(lng),
                        "features": features
                    })
                    
                    feature_vectors.append([
                        ndvi, vv, vh, dem, rain, soil, land_use
                    ])
        
        if not grid_points:
            return jsonify({
                "grid": [],
                "grid_size": 0,
                "model_version": "ExtraTrees_AgriCarbon_v1.1-dMRV"
            })
        
        # Run batch prediction
        predictions = model.predict(np.array(feature_vectors))
        
        # Format output
        output_grid = []
        for i, p in enumerate(grid_points):
            biomass = float(predictions[i])
            carbon_stock = float(biomass * 0.47)
            
            output_grid.append({
                "lat": p["lat"],
                "lng": p["lng"],
                "biomass": biomass,
                "carbon_stock": carbon_stock,
                "features": p["features"]
            })
        
        return jsonify({
            "grid": output_grid,
            "grid_size": len(output_grid),
            "model_version": "ExtraTrees_AgriCarbon_v1.1-dMRV"
        })
        
    except Exception as e:
        return jsonify({"error": str(e)}), 500

@app.route("/model-info", methods=["GET"])
def model_info():
    """Get model metadata"""
    local_info_path = Path(__file__).parent / "agb_model_extratrees_v1_info.json"
    original_info_path = Path(__file__).parent.parent.parent / "Competition" / "ticta" / "streamlit" / "agb_model_extratrees_v1_info.json"
    
    info_path = local_info_path if local_info_path.exists() else original_info_path
    try:
        with open(info_path, 'r') as f:
            return jsonify(json.load(f))
    except:
        return jsonify({"error": "Model info not found"}), 404

# Serve the static Vite frontend
@app.route("/", defaults={"path": ""})
@app.route("/<path:path>")
def serve(path):
    if path in ["predict", "predict-grid", "health", "model-info"]:
        return jsonify({"error": "Not Found"}), 404
        
    if dist_dir.exists() and (dist_dir / path).exists() and path != "":
        return send_from_directory(str(dist_dir), path)
    elif dist_dir.exists():
        return send_from_directory(str(dist_dir), "index.html")
    else:
        return "Backend is running. Frontend 'dist' folder not found.", 200

if __name__ == "__main__":
    # Hugging Face Spaces passes the port via the PORT environment variable (usually 7860)
    port = int(os.environ.get("PORT", 7860))
    # Turn off debug mode in production to save logs and improve speed
    debug_mode = os.environ.get("FLASK_ENV") == "development"
    app.run(host="0.0.0.0", port=port, debug=debug_mode)