import os
import urllib.request
import time
import gradio as gr
import numpy as np
from part1_data import TobaccoAnalyzer
from part2_visualization import VisualizationHandler
from part3 import SAMAnalyzer

# Model setup
MODEL_PATH = 'sam_vit_h_4b8939.pth'

def download_with_progress(url, filename):
    if os.path.exists(filename):
        print(f"Model already exists at {filename}")
        return
    print(f"Downloading {filename}...")
    start_time = time.time()
    urllib.request.urlretrieve(url, filename)
    end_time = time.time()
    print(f"Download completed in {end_time - start_time:.2f} seconds")

# Download SAM model if it doesn't exist
if not os.path.exists(MODEL_PATH):
    try:
        download_with_progress(
            'https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth',
            MODEL_PATH
        )
    except Exception as e:
        print(f"Error downloading SAM model: {e}")
        print("Please ensure you have internet connection and sufficient disk space.")

def analyze_location(location_name):
    """Main analysis function with enhanced historical and forecast data"""
    try:
        analyzer = TobaccoAnalyzer()
        visualizer = VisualizationHandler(analyzer.optimal_conditions)
        
        # Get coordinates from location name
        location_data = analyzer.geocode_location(location_name)
        if not location_data:
            return None, None, "Location not found. Please try a different location name.", None
        
        lat, lon = location_data['lat'], location_data['lon']
        
        # Get weather data
        df = analyzer.get_weather_data(lat, lon, historical_days=90, forecast_days=90)
        if df is None or df.empty:
            return None, None, "Unable to fetch weather data. Please try again.", None
        
        # Separate historical and forecast data
        historical = df[df['type'] == 'historical']
        forecast = df[df['type'] != 'historical']
        
        if historical.empty:
            return None, None, "No historical data available.", None
        
        # Calculate base scores
        temp_score = np.clip((historical['temperature'].mean() - 15) / (30 - 15), 0, 1)
        humidity_score = np.clip((historical['humidity'].mean() - 50) / (80 - 50), 0, 1)
        rainfall_score = np.clip(historical['rainfall'].mean() / 5, 0, 1)
        ndvi_score = np.clip((historical['estimated_ndvi'].mean() + 1) / 2, 0, 1)
        
        # Get trends analysis
        trends = analyzer.analyze_trends(df)
        if trends is None:
            return None, None, "Error calculating trends.", None
        
        # Calculate overall score
        weights = {
            'temperature': 0.3,
            'humidity': 0.2,
            'rainfall': 0.2,
            'ndvi': 0.3
        }
        
        overall_score = (
            temp_score * weights['temperature'] +
            humidity_score * weights['humidity'] +
            rainfall_score * weights['rainfall'] +
            ndvi_score * weights['ndvi']
        )
        
        # Create visualizations
        time_series_plot = visualizer.create_interactive_plots(df)
        gauge_plot = visualizer.create_gauge_chart(overall_score)
        location_map = visualizer.create_enhanced_map(lat, lon, overall_score, historical['estimated_ndvi'].mean())
        
        # Generate analysis text
        analysis_text = f"""
        📍 Location Analysis:
        Location: {location_data['address']}
        Coordinates: {lat:.4f}°N, {lon:.4f}°E
        Region: {location_data['region'] if location_data['region'] else 'Unknown'}
        
        🌡️ Historical Weather Analysis (Past 90 Days):
        Temperature: {historical['temperature'].mean():.1f}°C (±{historical['temperature'].std():.1f}°C)
        Daily Range: {historical['temp_range'].mean():.1f}°C
        Humidity: {historical['humidity'].mean():.1f}% (±{historical['humidity'].std():.1f}%)
        Rainfall: {historical['rainfall'].mean():.1f}mm/day (±{historical['rainfall'].std():.1f}mm)
        
        🌿 Vegetation Analysis:
        Current NDVI: {historical['estimated_ndvi'].mean():.2f}
        Minimum NDVI: {historical['estimated_ndvi'].min():.2f}
        Maximum NDVI: {historical['estimated_ndvi'].max():.2f}
        Vegetation Status: {get_vegetation_status(historical['estimated_ndvi'].mean())}
        
        🔮 Forecast Analysis (Next 90 Days):
        Temperature: {forecast['temperature'].mean():.1f}°C (±{forecast['temperature'].std():.1f}°C)
        Humidity: {forecast['humidity'].mean():.1f}% (±{forecast['humidity'].std():.1f}%)
        Rainfall: {forecast['rainfall'].mean():.1f}mm/day (±{forecast['rainfall'].std():.1f}mm)
        Expected NDVI: {forecast['estimated_ndvi'].mean():.2f}
        
        📊 Growing Condition Scores:
        Temperature Score: {temp_score:.2f}
        Humidity Score: {humidity_score:.2f}
        Rainfall Score: {rainfall_score:.2f}
        Vegetation Score: {ndvi_score:.2f}
        Overall Score: {overall_score:.2f}
        
        🎯 Recommendations:
        {get_recommendations(overall_score, ndvi_score)}
        
        ⚠️ Risk Factors:
        {get_risk_factors(df, trends)}
        
        📝 Additional Notes:
        • Growing Season: {is_growing_season(historical['season'].iloc[-1])}
        • Weather Stability: {get_weather_stability(historical)}
        • Long-term Outlook: {get_long_term_outlook(trends)}
        """
        
        return location_map, analysis_text, time_series_plot, gauge_plot
        
    except Exception as e:
        error_message = f"An error occurred: {str(e)}"
        print(f"Error details: {e}")
        return None, None, error_message, None

def process_and_analyze_image(image):
    """Process and analyze uploaded satellite image"""
    try:
        if image is None:
            return None, "Please upload an image first.", "No image provided"
            
        print("Starting image analysis...")
        analyzer = SAMAnalyzer(model_path=MODEL_PATH)
        
        print("Processing image...")
        veg_index, health_analysis, fig = analyzer.process_image(image)
        
        if veg_index is None:
            return None, "Error processing image.", "Image processing failed"
        
        analysis_text = f"""
        🌿 Vegetation Analysis Results:
        
        📊 Average Vegetation Index: {health_analysis['average_index']:.2f}
        
        🌱 Vegetation Distribution:
        • Low Vegetation: {health_analysis['health_distribution']['low_vegetation']*100:.1f}%
        • Moderate Vegetation: {health_analysis['health_distribution']['moderate_vegetation']*100:.1f}%
        • High Vegetation: {health_analysis['health_distribution']['high_vegetation']*100:.1f}%
        
        📋 Overall Health Status: {health_analysis['overall_health']}
        """
        
        debug_msg = "Analysis completed successfully"
        return fig, analysis_text, debug_msg
        
    except Exception as e:
        error_msg = f"Error during analysis: {str(e)}"
        print(error_msg)
        return None, None, error_msg

def get_vegetation_status(ndvi):
    """Get detailed vegetation status based on NDVI value"""
    if ndvi < 0:
        return "Very Low - Bare soil or water bodies"
    elif ndvi < 0.1:
        return "Low - Very sparse vegetation"
    elif ndvi < 0.2:
        return "Sparse - Stressed vegetation"
    elif ndvi < 0.3:
        return "Moderate - Typical agricultural land"
    elif ndvi < 0.4:
        return "Good - Healthy vegetation"
    elif ndvi < 0.5:
        return "High - Very healthy vegetation"
    elif ndvi < 0.6:
        return "Very High - Dense vegetation"
    else:
        return "Dense - Very dense, healthy vegetation"

def get_recommendations(score, ndvi):
    """Get detailed recommendations based on scores"""
    if score >= 0.8 and ndvi >= 0.6:
        return """
        ✅ Excellent conditions for tobacco growing
        • Proceed with standard planting schedule
        • Regular monitoring recommended
        • Consider expansion opportunities
        """
    elif score >= 0.6:
        return """
        👍 Good conditions with some considerations
        • Implement basic risk mitigation measures
        • Regular monitoring essential
        • Consider crop insurance
        """
    elif score >= 0.4:
        return """
        ⚠️ Marginal conditions - proceed with caution
        • Enhanced monitoring required
        • Strong risk mitigation needed
        • Crop insurance strongly recommended
        • Consider alternative timing
        """
    else:
        return """
        ❌ Poor conditions - high risk
        • Not recommended for planting
        • Consider alternative locations
        • Extensive risk mitigation needed
        • Alternative crops suggested
        """

def get_risk_factors(df, trends):
    """Analyze and return risk factors"""
    risks = []
    
    if df['temp_range'].mean() > 15:
        risks.append("• High daily temperature variations")
    if trends['historical']['temperature']['trend'] < 0:
        risks.append("• Declining temperature trend")
    if df['rainfall'].std() > df['rainfall'].mean():
        risks.append("• Inconsistent rainfall patterns")
    if trends['historical']['rainfall']['trend'] < 0:
        risks.append("• Decreasing rainfall trend")
    if trends['historical']['ndvi']['trend'] < 0:
        risks.append("• Declining vegetation health")
    
    return "\n".join(risks) if risks else "No major risk factors identified"

def is_growing_season(season):
    """Check if current season is suitable for growing"""
    season_suitability = {
        'Main': "Prime growing season - Optimal conditions",
        'Early': "Early growing season - Good potential",
        'Late': "Late growing season - Monitor closely",
        'Dry': "Dry season - Higher risk period"
    }
    return season_suitability.get(season, "Season not identified")

def get_weather_stability(df):
    """Assess weather stability"""
    temp_std = df['temperature'].std()
    if temp_std < 2:
        return "Very stable weather patterns"
    elif temp_std < 4:
        return "Moderately stable weather"
    return "Unstable weather patterns - higher risk"

def get_long_term_outlook(trends):
    """Assess long-term outlook based on trends"""
    temp_trend = trends['historical']['temperature']['trend']
    rain_trend = trends['historical']['rainfall']['trend']
    
    if temp_trend > 0 and rain_trend > 0:
        return "Improving conditions"
    elif temp_trend < 0 and rain_trend < 0:
        return "Deteriorating conditions"
    return "Mixed conditions - monitor closely"

# Create Gradio interface
with gr.Blocks(theme=gr.themes.Base()) as demo:
    gr.Markdown(
        """
        # 🌱 Agricultural Credit Risk Assessment System
        ## Weather, Vegetation, and Credit Scoring Analysis for Tobacco Farming
        """
    )
    
    with gr.Tab("📊 Location Analysis"):
        with gr.Row():
            location_input = gr.Textbox(
                label="Enter Location",
                placeholder="e.g., Tabora, Tanzania",
                scale=4
            )
            analyze_button = gr.Button("Analyze", variant="primary", scale=1)

        with gr.Row():
            with gr.Column(scale=2):
                location_map = gr.HTML(label="NDVI Analysis Map")
            with gr.Column(scale=1):
                analysis_text = gr.Textbox(
                    label="Analysis Results",
                    lines=25,
                    show_label=False
                )
        
        with gr.Row():
            weather_plot = gr.Plot(label="Weather Analysis")
        
        with gr.Row():
            score_gauge = gr.Plot(label="Growing Conditions Score")
        
        gr.Examples(
            examples=[
                ["Tabora, Tanzania"],
                ["Urambo, Tabora, Tanzania"],
                ["Sikonge, Tabora, Tanzania"],
                ["Nzega, Tabora, Tanzania"]
            ],
            inputs=location_input,
            outputs=[location_map, analysis_text, weather_plot, score_gauge],
            fn=analyze_location,
            cache_examples=True
        )
        
        analyze_button.click(
            fn=analyze_location,
            inputs=[location_input],
            outputs=[location_map, analysis_text, weather_plot, score_gauge]
        )
    
    with gr.Tab("🛰️ Satellite Image Analysis"):
        gr.Markdown("""
        ## Satellite Image Analysis with SAM2
        Upload a satellite or aerial image to analyze vegetation health using advanced segmentation.
        """)
        
        with gr.Row():
            image_input = gr.Image(
                label="Upload Satellite/Aerial Image",
                type="numpy"
            )
        
        with gr.Row():
            analyze_image_button = gr.Button("🔍 Analyze Image", variant="primary")
        
        with gr.Row():
            image_plot = gr.Plot(
                label="Vegetation Analysis Results"
            )
            
        with gr.Row():
            image_analysis = gr.Textbox(
                label="Analysis Results",
                lines=10
            )
        
        debug_output = gr.Textbox(
            label="Debug Information",
            lines=3,
            visible=True
        )
        
        analyze_image_button.click(
            fn=process_and_analyze_image,
            inputs=[image_input],
            outputs=[image_plot, image_analysis, debug_output]
        )

# Launch the app
if __name__ == "__main__":
    demo.launch()