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():
    try:
        # Fire detector (VGG16)
        vgg_model = load_model(
            'vgg16_focal_unfreeze_more.keras',
            custom_objects={'BinaryFocalCrossentropy': BinaryFocalCrossentropy}
        )
        # Severity classifier (Xception)
        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()}
        )
        # Ensemble and trend models
        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
    except Exception as e:
        print(f"Error loading models: {e}")
        return None, None, None, None, None

# --- RULES & TEMPLATES ---
# Mapping severity levels and trends
target_map = {0: 'mild', 1: 'moderate', 2: 'severe'}
trend_map = {1: 'increase', 0: 'same', -1: 'decrease'}

# Severity progression rules based on current severity and weather trend
task_rules = {
    'mild':    {'decrease':'mild','same':'mild','increase':'moderate'},
    'moderate':{'decrease':'mild','same':'moderate','increase':'severe'},
    'severe':  {'decrease':'moderate','same':'severe','increase':'severe'}
}

# Detailed recommendations for each severity level
recommendations = {
    'mild': {
        'immediate': "Deploy spot crews for initial attack. Establish command post. Monitor fire behavior with drones or aircraft. Alert local fire stations.",
        'evacuation': "No mass evacuation needed. Notify nearby communities of potential risk. Prepare evacuation routes if conditions change.",
        'containment': "Establish initial fire lines. Use hand crews for direct attack. Position water resources. Clear fuel breaks where feasible.",
        'prevention': "Implement controlled underburning in surrounding areas. Manage vegetation density. Create defensible spaces around structures.",
        'education': "Inform public on fire watch protocols and reporting mechanisms. Train local volunteers in basic firefighting techniques."
    },
    'moderate': {
        'immediate': "Dispatch multiple engines and aerial support. Establish unified command system. Deploy heavy equipment. Request additional resources.",
        'evacuation': "Prepare evacuation zones and staging areas. Advise voluntary evacuation for vulnerable populations. Alert emergency shelters.",
        'containment': "Build substantial fire breaks. Conduct water drops from helicopters. Implement indirect attack strategies. Protect critical infrastructure.",
        'prevention': "Initiate fuel reduction programs in adjacent areas. Create wider buffer zones. Assess watershed protection needs.",
        'education': "Conduct community emergency drills. Launch awareness campaigns on evacuation procedures. Distribute preparedness materials."
    },
    'severe': {
        'immediate': "Implement full suppression with air tankers and multiple resources. Establish incident management team. Request state/federal assistance. Deploy specialized teams.",
        'evacuation': "Issue mandatory evacuation orders. Open multiple emergency shelters. Implement traffic control measures. Assist vulnerable populations.",
        'containment': "Deploy fire retardant lines from aircraft. Consider backfires and burnout operations. Protect critical infrastructure. Establish multiple control lines.",
        'prevention': "Plan for reforestation and erosion control. Harden infrastructure against future fires. Implement watershed protection measures.",
        'education': "Conduct comprehensive emergency response training. Implement risk communication strategies. Develop long-term community resilience programs."
    }
}

# --- PIPELINE FUNCTIONS ---
def detect_fire(img):
    """Detect if a wildfire is present in the image"""
    try:
        if vgg_model is None:
            return True, 0.85  # Fallback if model not loaded
            
        x = keras_image.img_to_array(img.resize((128,128)))[None]
        x = vgg_preprocess(x)
        prob = float(vgg_model.predict(x)[0][0])
        return prob >= 0.5, prob
    except Exception as e:
        print(f"Error in fire detection: {e}")
        return False, 0.0

def classify_severity(img):
    """Classify the severity of the detected wildfire"""
    try:
        if xception_model is None or rf_model is None or xgb_model is None:
            return 'moderate'  # Fallback if models not loaded
            
        x = keras_image.img_to_array(img.resize((224,224)))[None]
        x = xce_preprocess(x)
        preds = xception_model.predict(x)
        rf_p = rf_model.predict(preds)[0]
        xgb_p = xgb_model.predict(preds)[0]
        ensemble = int(round((rf_p + xgb_p)/2))
        return target_map.get(ensemble, 'moderate')
    except Exception as e:
        print(f"Error in severity classification: {e}")
        return 'moderate'  # Default to moderate severity if error occurs

def fetch_weather_trend(lat, lon):
    """Fetch weather data and determine trend"""
    try:
        # Use local weather calculation if API fails
        try:
            end = datetime.utcnow()
            start = end - timedelta(days=1)
            url = API_URL.format(
                lat=lat, lon=lon,
                start=start.strftime('%Y-%m-%d'),
                end=end.strftime('%Y-%m-%d')
            )
            response = requests.get(url, timeout=5)
            if response.status_code != 200:
                raise Exception(f"API returned status code {response.status_code}")
            
            df = pd.DataFrame(response.json().get('daily', {}))
        except Exception as e:
            print(f"API error: {e}. Using synthetic data.")
            # Create synthetic weather data if API fails
            df = pd.DataFrame({
                'date': [(datetime.utcnow() - timedelta(days=i)).strftime('%Y-%m-%d') for i in range(1, -1, -1)],
                'precipitation_sum': [5, 2],
                'temperature_2m_max': [28, 30],
                'temperature_2m_min': [18, 20],
                'relative_humidity_2m_max': [70, 65],
                'relative_humidity_2m_min': [40, 35],
                'windspeed_10m_max': [15, 18]
            })
        
        # Process weather data
        for c in ['precipitation_sum','temperature_2m_max','temperature_2m_min',
                'relative_humidity_2m_max','relative_humidity_2m_min','windspeed_10m_max']:
            df[c] = pd.to_numeric(df.get(c,[]), errors='coerce')
        
        df['precipitation'] = df['precipitation_sum'].fillna(0)
        df['temperature'] = (df['temperature_2m_max'] + df['temperature_2m_min'])/2
        df['humidity'] = (df['relative_humidity_2m_max'] + df['relative_humidity_2m_min'])/2
        df['wind_speed'] = df['windspeed_10m_max']
        
        # Calculate fire risk score based on weather parameters
        df['fire_risk_score'] = (
            0.4*(df['temperature']/55) +
            0.2*(1-df['humidity']/100) +
            0.3*(df['wind_speed']/60) +
            0.1*(1-df['precipitation']/50)
        )
        
        # Prepare features for trend prediction
        feats = df[['temperature','humidity','wind_speed','precipitation','fire_risk_score']]
        feat = feats.fillna(feats.mean()).iloc[-1].values.reshape(1,-1)
        
        # Predict trend using logistic regression model or fallback
        if lr_model is not None:
            trend_cl = lr_model.predict(feat)[0]
            return trend_map.get(trend_cl, 'same')
        else:
            # Fallback logic if model isn't loaded
            if df['fire_risk_score'].iloc[-1] > 0.6:
                return 'increase'
            elif df['fire_risk_score'].iloc[-1] < 0.4:
                return 'decrease'
            return 'same'
            
    except Exception as e:
        print(f"Error in weather trend analysis: {e}")
        return 'same'  # Default to 'same' trend if all else fails

def generate_recommendations(original_severity, weather_trend):
    """Generate comprehensive recommendations based on severity and trend"""
    # Determine projected severity
    projected_severity = task_rules[original_severity][weather_trend]
    
    # Get recommendations for projected severity
    rec = recommendations[projected_severity]
    
    # Create detailed recommendation text
    recommendation_text = f"""**Original Severity:** {original_severity.title()}  
**Weather Trend:** {weather_trend.title()}  
**Projected Severity:** {projected_severity.title()}

### Management Recommendations:

**1. Immediate Actions:**  
{rec['immediate']}

**2. Evacuation Guidelines:**  
{rec['evacuation']}

**3. Short-term Containment:**  
{rec['containment']}

**4. Long-term Prevention & Recovery:**  
{rec['prevention']}

**5. Community Education:**  
{rec['education']}
"""
    return recommendation_text

# --- MAIN PIPELINE ---
def pipeline(image):
    """Main processing pipeline for wildfire detection and analysis"""
    if image is None:
        return "No image provided", "N/A", "N/A", "**Please upload an image to analyze**"
    
    # Convert to PIL Image
    img = Image.fromarray(image).convert('RGB')
    
    # Step 1: Detect fire
    fire, prob = detect_fire(img)
    if not fire:
        return f"No wildfire detected (confidence: {(1-prob)*100:.1f}%)", "N/A", "N/A", "**No wildfire detected. Stay alert and maintain regular monitoring.**"
    
    # Step 2: Classify severity
    severity = classify_severity(img)
    
    # Step 3: Fetch weather trend
    trend = fetch_weather_trend(*FOREST_COORDS['Pakistan Forest'])
    
    # Step 4: Generate recommendations
    recommendations_text = generate_recommendations(severity, trend)
    
    return f"Wildfire detected (confidence: {prob*100:.1f}%)", severity.title(), trend.title(), recommendations_text

# --- LOAD MODELS GLOBALLY ---
vgg_model, xception_model, rf_model, xgb_model, lr_model = load_models()

# --- GRADIO INTERFACE ---
interface = gr.Interface(
    fn=pipeline,
    inputs=gr.Image(type='numpy', label='Upload Wildfire Image'),
    outputs=[
        gr.Textbox(label='Fire Status'),
        gr.Textbox(label='Current Severity Level'),
        gr.Textbox(label='Weather Trend'),
        gr.Markdown(label='Management Recommendations')
    ],
    title='Wildfire Detection & Management Assistant',
    description='Upload an image from a forest region in Pakistan to determine wildfire presence, severity, weather-driven trend, and get expert management recommendations.',
    examples=[],
    theme=gr.themes.Base(),
    allow_flagging='never'
).queue(api_open=True)

if __name__ == '__main__':
    interface.launch(share=False)