import os
import requests
import pandas as pd
import numpy as np
import joblib
import gradio as gr
from datetime import datetime, timedelta
from tensorflow.keras.models import load_model
from tensorflow.keras.preprocessing import image as keras_image
from tensorflow.keras.applications.vgg16 import preprocess_input as vgg_preprocess
from tensorflow.keras.applications.xception import preprocess_input as xce_preprocess
from tensorflow.keras.losses import BinaryFocalCrossentropy
from PIL import Image

# --- CONFIGURATION ---
FOREST_COORDS = {'Pakistan Forest': (34.0, 73.0)}
API_URL = (
    "https://archive-api.open-meteo.com/v1/archive"
    "?latitude={lat}&longitude={lon}"
    "&start_date={start}&end_date={end}"
    "&daily=temperature_2m_max,temperature_2m_min,"
    "precipitation_sum,windspeed_10m_max,"
    "relative_humidity_2m_max,relative_humidity_2m_min"
    "&timezone=UTC"
)

# --- LOAD MODELS ---
def load_models():
    vgg_model = load_model(
        'vgg16_focal_unfreeze_more.keras',
        custom_objects={'BinaryFocalCrossentropy': BinaryFocalCrossentropy}
    )
    def focal_loss_fixed(gamma=2., alpha=.25):
        import tensorflow.keras.backend as K
        def loss_fn(y_true, y_pred):
            eps = K.epsilon()
            y_pred = K.clip(y_pred, eps, 1. - eps)
            ce = -y_true * K.log(y_pred)
            w = alpha * K.pow(1 - y_pred, gamma)
            return K.mean(w * ce, axis=-1)
        return loss_fn
    xce_model = load_model(
        'severity_post_tta.keras',
        custom_objects={'focal_loss_fixed': focal_loss_fixed()}
    )
    rf_model = joblib.load('ensemble_rf_model.pkl')
    xgb_model = joblib.load('ensemble_xgb_model.pkl')
    lr_model = joblib.load('wildfire_logistic_model_synthetic.joblib')
    return vgg_model, xce_model, rf_model, xgb_model, lr_model

vgg_model, xception_model, rf_model, xgb_model, lr_model = load_models()

# --- RULES & TEMPLATES ---
target_map = {0: 'mild', 1: 'moderate', 2: 'severe'}
trend_map = {1: 'increase', 0: 'same', -1: 'decrease'}
task_rules = {
    'mild':    {'decrease': 'mild', 'same': 'mild', 'increase': 'moderate'},
    'moderate': {'decrease': 'mild', 'same': 'moderate', 'increase': 'severe'},
    'severe':  {'decrease': 'moderate', 'same': 'severe', 'increase': 'severe'}
}
templates = {
    'mild': (
        "📌 **Immediate Monitoring:** Although fire intensity is low, assign lookouts to monitor hotspots every 30 minutes. Use handheld IR cameras to detect any hidden flare-ups.\n\n"
        "📌 **Community Alert:** Send SMS alerts to nearby villages reminding them to stay vigilant. Provide clear instructions on how to report any smoke sightings.\n\n"
        "📌 **Fuel Management:** Conduct targeted removal of leaf litter and dry underbrush within a 100 m radius to reduce the chance of flare-ups.\n\n"
        "📌 **Preparedness Drills:** Hold a quick drill with ground crews to review communication protocols and ensure equipment (hoses, pumps) is ready.\n\n"
        "📌 **Public Education:** Distribute flyers on safe fire-watch practices and set up a hotline for rapid reporting."
    ),
    'moderate': (
        "🚒 **Rapid Response:** Dispatch two engine crews and one aerial water-drop helicopter. Coordinate with the regional command center to stage retardant tanks nearby.\n\n"
        "🏃‍♂️ **Evacuation Prep:** Pre-position evacuation buses at community centers. Issue voluntary evacuation notices to residents within 5 km downwind.\n\n"
        "🛠 **Containment Lines:** Construct a 10 m fire break using both hand tools and bulldozers. Apply fire-retardant gel along the anticipated flank.\n\n"
        "🌱 **Fuel Reduction:** Begin mechanical thinning of small trees and brush in high-risk zones adjacent to critical infrastructure.\n\n"
        "📣 **Awareness Campaign:** Launch radio spots explaining what to do if fire approaches, including evacuation routes and shelter locations."
    ),
    'severe': (
        "✈️ **Full Suppression:** Mobilize two air tankers for retardant drops and four ground crews with heavy equipment. Integrate real-time satellite imagery for targeting.\n\n"
        "🚨 **Mandatory Evacuation:** Issue immediate evacuation orders for all residents within a 10 km radius. Open three emergency shelters with medical staff on standby.\n\n"
        "🔥 **Backfire Operations:** Conduct controlled backfires under supervision of senior incident commanders to remove fuel ahead of the main front.\n\n"
        "🌳 **Post-Fire Rehabilitation:** Plan reforestation with fire-resistant native species; stabilize soil to prevent erosion in burn scar areas.\n\n"
        "🗣 **Crisis Communication:** Hold daily press briefings and social media updates. Provide mental-health support hotlines for displaced families."
    )
}

# --- FUNCTIONS ---
def detect_fire(img):
    img_resized = img.resize((224, 224))
    arr = keras_image.img_to_array(img_resized)
    arr = np.expand_dims(arr, axis=0)
    arr = vgg_preprocess(arr)
    pred = vgg_model.predict(arr)[0][0]
    is_fire = pred >= 0.5
    return is_fire, pred

def classify_severity(img):
    img_resized = img.resize((224, 224))
    arr = keras_image.img_to_array(img_resized)
    arr = np.expand_dims(arr, axis=0)
    arr = xce_preprocess(arr)
    feat = np.squeeze(arr)
    feat_flat = feat.flatten().reshape(1, -1)
    
    rf_pred = rf_model.predict_proba(feat_flat)
    xgb_pred = xgb_model.predict_proba(feat_flat)
    avg_pred = (rf_pred + xgb_pred) / 2
    final_class = np.argmax(avg_pred)
    return target_map[final_class]

def fetch_weather_trend(lat, lon):
    today = datetime.utcnow().date()
    start_date = today - timedelta(days=2)
    end_date = today - timedelta(days=1)

    url = API_URL.format(lat=lat, lon=lon, start=start_date, end=end_date)
    response = requests.get(url)
    if response.status_code != 200:
        return 'same'  # fallback if API fails

    data = response.json()
    temp_max = data['daily']['temperature_2m_max']
    wind_max = data['daily']['windspeed_10m_max']
    humidity_min = data['daily']['relative_humidity_2m_min']

    # crude trend logic: hotter, windier = worse
    temp_trend = np.sign(temp_max[-1] - temp_max[0])
    wind_trend = np.sign(wind_max[-1] - wind_max[0])
    humidity_trend = -np.sign(humidity_min[-1] - humidity_min[0])

    overall_trend = temp_trend + wind_trend + humidity_trend
    if overall_trend > 0:
        return 'increase'
    elif overall_trend < 0:
        return 'decrease'
    else:
        return 'same'

def generate_recommendations(original, trend):
    projected = task_rules[original][trend]
    header = (
        f"## 🔥 Wildfire Situation Update\n"
        f"- **Original Severity:** {original.title()}\n"
        f"- **Weather Trend:** {trend.title()}\n"
        f"- **Projected Severity:** {projected.title()}\n\n"
    )
    paras = templates[projected].split("\n\n")
    formatted = "\n\n".join(paras)
    return header + formatted

def pipeline(image):
    img = Image.fromarray(image).convert('RGB')
    fire, prob = detect_fire(img)
    if not fire:
        return (
            f"**No wildfire detected** (probability={prob:.2f})",
            "N/A",
            "N/A",
            "There is currently no sign of wildfire in the image. Continue normal monitoring."
        )
    sev = classify_severity(img)
    trend = fetch_weather_trend(*FOREST_COORDS['Pakistan Forest'])
    recs = generate_recommendations(sev, trend)
    return (
        f"**🔥 Fire Detected** (probability={prob:.2f})",
        sev.title(),
        trend.title(),
        recs
    )

# --- GRADIO APP ---
with gr.Blocks(css=""" 
  .gradio-container {
    background-color: #f5f7fa !important;
  }
  .gradio-textbox textarea {
    background-color: #ffffff !important;
    border: 1px solid #cbd2d9 !important;
    border-radius: 8px !important;
    padding: 12px !important;
    font-size: 1rem !important;
    color: #333333 !important;
    min-height: 3em !important;
  }
  .gradio-accordion {
    background-color: #ffffff !important;
    border: 1px solid #cbd2d9 !important;
    border-radius: 8px !important;
    padding: 8px !important;
  }
  .gradio-button {
    background-color: #0072ce !important;
    color: white !important;
    border-radius: 6px !important;
    padding: 8px 16px !important;
    font-weight: 600 !important;
  }
  .gradio-button:hover {
    background-color: #005bb5 !important;
  }
  .gradio-markdown h1, .gradio-markdown h2 {
    color: #1f2937 !important;
    margin-bottom: 0.5em !important;
  }
""") as demo:
    gr.Markdown("# Wildfire Detection & Management Assistant")
    gr.Markdown("Upload a forest image from Pakistan; the system will detect fire, assess severity, analyze weather trends, and provide in-depth recommendations.")
    
    with gr.Row():
        inp = gr.Image(type="numpy", label="Upload Wildfire Image")
        with gr.Column():
            status = gr.Textbox(label="Fire Status", interactive=False)
            severity = gr.Textbox(label="Severity Level", interactive=False)
            trend = gr.Textbox(label="Weather Trend", interactive=False)
    
    with gr.Accordion("📋 Detailed Recommendations", open=False):
        rec_box = gr.Markdown(label="Recommendations")
    
    btn = gr.Button("Analyze")
    btn.click(fn=pipeline, inputs=inp, outputs=[status, severity, trend, rec_box])
    
    gr.HTML("<p style='font-size:0.8em; color:#666;'>© 2025 ForestAI Labs</p>")

if __name__ == "__main__":
    demo.launch()