import gradio as gr
import requests
import folium
from datetime import datetime, timedelta
import tempfile
import os
from PIL import Image
import io

# Montana Mountain Peaks coordinates
MONTANA_PEAKS = {
    "Lone Peak (Big Sky)": (45.27806, -111.45028),  # 45°16′41″N 111°27′01″W
    "Sacajawea Peak": (45.89583, -110.96861),  # 45°53′45″N 110°58′7″W
    "Pioneer Mountain": (45.231835, -111.450505)  # 45°13′55″N 111°27′2″W
}

def download_image(url):
    """Download image from URL and return as PIL Image."""
    try:
        response = requests.get(url, timeout=10)
        response.raise_for_status()
        image = Image.open(io.BytesIO(response.content))
        return image
    except Exception as e:
        print(f"Error downloading image from {url}: {str(e)}")
        return None

def get_snow_forecast(lat, lon):
    """Get snow forecast data from NOAA's snow forecast API."""
    try:
        # Get the forecast URL for the coordinates
        points_url = f"https://api.weather.gov/points/{lat},{lon}"
        response = requests.get(points_url, timeout=10)
        response.raise_for_status()
        forecast_url = response.json()['properties']['forecast']
        
        # Get the actual forecast with snow data
        forecast_response = requests.get(forecast_url, timeout=10)
        forecast_response.raise_for_status()
        forecast_data = forecast_response.json()
        
        # Process forecast periods to extract snow data
        snow_data = []
        
        for period in forecast_data['properties']['periods']:
            period_data = {
                'time': period['startTime'],
                'temperature': period['temperature'],
                'snowfall': 0  # Default to no snow
            }
            
            # Check if snow is mentioned in the forecast
            forecast_text = period['detailedForecast'].lower()
            
            if any(word in forecast_text for word in ['snow', 'flurries', 'wintry mix']):
                import re
                # Look for snow amounts in the text
                amount_patterns = [
                    r'(\d+(?:\.\d+)?)\s*(?:to\s*\d+(?:\.\d+)?)?\s*inches? of snow',
                    r'snow accumulation of (\d+(?:\.\d+)?)\s*(?:to\s*\d+(?:\.\d+)?)?\s*inches',
                    r'accumulation of (\d+(?:\.\d+)?)\s*(?:to\s*\d+(?:\.\d+)?)?\s*inches',
                ]
                
                for pattern in amount_patterns:
                    match = re.search(pattern, forecast_text)
                    if match:
                        try:
                            snow_amount = float(match.group(1))
                            period_data['snowfall'] = snow_amount / 39.37  # Convert inches to meters
                            break
                        except (ValueError, IndexError):
                            continue
            
            snow_data.append(period_data)
        
        return snow_data
    except Exception as e:
        print(f"Error in snow forecast: {str(e)}")
        return []
    except Exception as e:
        print(f"Error fetching snow forecast: {str(e)}")
        return None

def calculate_snow_density(temp_c, is_new_snow=True):
    """Calculate approximate snow density based on temperature.
    Returns density in kg/m³."""
    if is_new_snow:
        if temp_c <= -20:
            return 50  # Very cold, light powder
        elif temp_c <= -10:
            return 70  # Cold powder
        elif temp_c <= -5:
            return 100  # Moderate temperature powder
        elif temp_c <= 0:
            return 150  # Near freezing powder
        else:
            return 200  # Wet snow
    else:
        # For existing snowpack, density would be higher
        return 350  # Average settled snow density

def get_noaa_forecast(lat, lon):
    """Get NOAA forecast data for given coordinates."""
    points_url = f"https://api.weather.gov/points/{lat},{lon}"
    
    try:
        # Get the forecast URL for the coordinates
        response = requests.get(points_url, timeout=10)
        response.raise_for_status()
        forecast_url = response.json()['properties']['forecast']
        
        # Get the actual forecast
        forecast_response = requests.get(forecast_url, timeout=10)
        forecast_response.raise_for_status()
        forecast_data = forecast_response.json()
        
        # Get snow forecast data
        snow_data = get_snow_forecast(lat, lon)
        
        # Format the forecast text
        periods = forecast_data['properties']['periods']
        forecast_text = ""
        
        for i, period in enumerate(periods[:6]):  # Show next 3 days (day and night)
            forecast_text += f"\n{period['name']}:\n"
            forecast_text += f"Temperature: {period['temperature']}°{period['temperatureUnit']}\n"
            forecast_text += f"Wind: {period['windSpeed']} {period['windDirection']}\n"
            
            # Convert temperature to Celsius for snow density calculation
            temp_c = (period['temperature'] - 32) * 5/9 if period['temperatureUnit'] == 'F' else period['temperature']
            
            # Add snow-specific forecasts if available
            if i < len(snow_data) and snow_data[i]['snowfall'] > 0:
                snow_density = calculate_snow_density(temp_c)
                snow_amount = snow_data[i]['snowfall']  # in meters
                swe = snow_amount * (snow_density / 1000)  # Convert density to g/cm³
                
                # Convert to inches for display
                snow_inches = snow_amount * 39.37
                swe_inches = swe * 39.37
                
                forecast_text += f"Predicted Snowfall: {snow_inches:.1f} inches\n"
                forecast_text += f"Snow Water Equivalent: {swe_inches:.2f} inches\n"
                forecast_text += f"Estimated Snow Density: {snow_density} kg/m³ "
                forecast_text += f"({get_snow_quality_description(snow_density)})\n"
            
            forecast_text += f"{period['detailedForecast']}\n"
        
        return forecast_text.strip()
    except requests.exceptions.RequestException as e:
        return f"Error fetching forecast: {str(e)}"

def get_snow_quality_description(density):
    """Return a description of snow quality based on density."""
    if density <= 50:
        return "Ultra-light powder"
    elif density <= 70:
        return "Light powder"
    elif density <= 100:
        return "Moderate powder"
    elif density <= 150:
        return "Packed powder"
    elif density <= 200:
        return "Wet snow"
    else:
        return "Dense/settled snow"

import gradio as gr
import requests
import folium
from datetime import datetime, timedelta
import tempfile
import os
from PIL import Image, ImageDraw, ImageFont
import io

# Montana Mountain Peaks coordinates
MONTANA_PEAKS = {
    "Lone Peak (Big Sky)": (45.27806, -111.45028),  # 45°16′41″N 111°27′01″W
    "Sacajawea Peak": (45.89583, -110.96861),  # 45°53′45″N 110°58′7″W
    "Pioneer Mountain": (45.231835, -111.450505)  # 45°13′55″N 111°27′2″W
}

def get_radar_and_snow_forecasts():
    """Get current radar and forecast loop."""
    try:
        # For the Montana region, get:
        # 1. Current base reflectivity
        # 2. Short-term forecast (0-6 hours)
        # 3. Medium-term forecast (6-12 hours)
        base_url = "https://radar.weather.gov/ridge/Overlays/Foregrounds"
        overlay_url = "https://radar.weather.gov/ridge/Overlays"
        background_url = "https://radar.weather.gov/ridge/Overlays/Backgrounds"
        
        # Build image frames
        frames = []
        
        # Current conditions plus 6 forecast periods (7 total frames)
        for hour in range(7):
            try:
                # Start with the static background
                background = Image.open(requests.get(f"{background_url}/Topo/Short/KMSX_Topo_Short.jpg", timeout=10).content)
                background = background.convert('RGB')
                
                # Add county overlays
                counties = Image.open(requests.get(f"{overlay_url}/Counties/Short/KMSX_Counties_Short.gif", timeout=10).content)
                counties = counties.convert('RGBA')
                background.paste(counties, (0, 0), counties)
                
                # Add highways
                highways = Image.open(requests.get(f"{overlay_url}/Highways/Short/KMSX_Highways_Short.gif", timeout=10).content)
                highways = highways.convert('RGBA')
                background.paste(highways, (0, 0), highways)
                
                # Add current radar or forecast overlay
                if hour == 0:
                    # Current radar
                    radar = Image.open(requests.get(f"{base_url}/Base/KMSX_Base_0.gif", timeout=10).content)
                else:
                    # Forecast radar (using NDFD data)
                    forecast_time = datetime.now() + timedelta(hours=hour)
                    forecast_str = forecast_time.strftime("%Y%m%d_%H")
                    radar = Image.open(requests.get(f"https://forecast.weather.gov/wrf_images/KMSX/Reflectivity/Reflectivity_{forecast_str}.gif", timeout=10).content)
                
                radar = radar.convert('RGBA')
                background.paste(radar, (0, 0), radar)
                
                # Add timestamp
                draw = ImageDraw.Draw(background)
                if hour == 0:
                    time_text = "Current Conditions"
                else:
                    forecast_time = datetime.now() + timedelta(hours=hour)
                    time_text = f"Forecast +{hour}h ({forecast_time.strftime('%I%p')})"
                
                try:
                    font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", 20)
                except:
                    font = ImageFont.load_default()
                
                # Draw text with outline
                x, y = 10, 10
                for dx, dy in [(-1,-1), (-1,1), (1,-1), (1,1)]:
                    draw.text((x+dx, y+dy), time_text, fill='black', font=font)
                draw.text((x, y), time_text, fill='white', font=font)
                
                frames.append(background)
                
            except Exception as e:
                print(f"Error processing frame {hour}: {str(e)}")
                continue
        
        if not frames:
            raise Exception("No frames could be generated")
            
        # Save as animated GIF
        with tempfile.NamedTemporaryFile(suffix='.gif', delete=False) as tmp_file:
            frames[0].save(
                tmp_file.name,
                save_all=True,
                append_images=frames[1:],
                duration=1500,  # 1.5 seconds per frame
                loop=0
            )
            return tmp_file.name
                
    except Exception as e:
        print(f"Error creating forecast display: {str(e)}")
        # Create error message image
        img = Image.new('RGB', (800, 600), color='black')
        draw = ImageDraw.Draw(img)
        message = (
            "Weather radar and forecast temporarily unavailable\n"
            f"{str(e)}\n"
            "Please check weather.gov for current conditions\n"
            f"Last attempt: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}"
        )
        try:
            font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", 20)
        except:
            font = ImageFont.load_default()
            
        try:
            draw.text((400, 300), message,
                     fill='white', anchor="mm", align="center", font=font)
        except:
            # Fallback if font fails
            draw.text((400, 300), message,
                     fill='white', anchor="mm", align="center")
        
        with tempfile.NamedTemporaryFile(suffix='.gif', delete=False) as tmp_file:
            img.save(tmp_file.name, format='GIF')
            return tmp_file.name

def create_map():
    """Create a folium map centered on Montana with optional markers."""
    # Center on Montana
    m = folium.Map(location=[45.5, -111.0], zoom_start=7)
    
    # Add markers for Montana peaks
    for peak_name, coords in MONTANA_PEAKS.items():
        folium.Marker(
            coords,
            popup=f"{peak_name}<br>Lat: {coords[0]:.4f}, Lon: {coords[1]:.4f}",
            tooltip=peak_name
        ).add_to(m)
    
    m.add_child(folium.ClickForLatLng())  # Enable click events
    return m._repr_html_()

def update_weather(lat, lon):
    """Update weather information based on coordinates."""
    try:
        # Validate coordinates
        lat = float(lat)
        lon = float(lon)
        if not (-90 <= lat <= 90 and -180 <= lon <= 180):
            return "Invalid coordinates. Latitude must be between -90 and 90, longitude between -180 and 180.", None, get_map(lat, lon)
        
        # Get forecast
        forecast = get_noaa_forecast(lat, lon)
        
        # Get animated snow forecast
        animated_forecast = get_radar_and_snow_forecasts()
        
        # Create updated map
        map_html = get_map(lat, lon)
        
        return forecast, animated_forecast, map_html
    except ValueError as e:
        return f"Error: Invalid coordinate format - {str(e)}", None, get_map(lat, lon)
    except Exception as e:
        return f"Error: {str(e)}", None, get_map(lat, lon)

def get_map(lat, lon):
    """Create a map centered on the given coordinates with markers."""
    try:
        # Create map with appropriate zoom level
        m = folium.Map(location=[lat, lon], zoom_start=9)
        
        # Add all Montana peaks
        for peak_name, coords in MONTANA_PEAKS.items():
            folium.Marker(
                coords,
                popup=f"{peak_name}<br>Lat: {coords[0]:.4f}, Lon: {coords[1]:.4f}",
                tooltip=peak_name
            ).add_to(m)
        
        # Add current location marker if not a peak
        if (lat, lon) not in MONTANA_PEAKS.values():
            folium.Marker([lat, lon], popup=f"Selected Location<br>Lat: {lat:.4f}, Lon: {lon:.4f}").add_to(m)
        
        m.add_child(folium.ClickForLatLng())  # Enable click events
        return m._repr_html_()
    except Exception as e:
        print(f"Error creating map: {str(e)}")
        return None

def make_peak_click_handler(peak_name):
    """Creates a click handler for a specific peak."""
    def handler():
        coords = MONTANA_PEAKS[peak_name]
        return coords[0], coords[1]
    return handler

# Create Gradio interface
with gr.Blocks(title="Montana Mountain Weather") as demo:
    gr.Markdown("# Montana Mountain Weather")
    
    with gr.Row():
        with gr.Column(scale=1):
            lat_input = gr.Number(
                label="Latitude",
                value=45.5,
                minimum=-90,
                maximum=90
            )
            lon_input = gr.Number(
                label="Longitude",
                value=-111.0,
                minimum=-180,
                maximum=180
            )
            
            # Quick access buttons for Montana peaks
            gr.Markdown("### Quick Access - Montana Peaks")
            peak_buttons = []
            for peak_name in MONTANA_PEAKS.keys():
                peak_buttons.append(gr.Button(f"📍 {peak_name}"))
            
            submit_btn = gr.Button("Get Weather", variant="primary")
        
        with gr.Column(scale=2):
            map_display = gr.HTML(get_map(45.5, -111.0))
    
    with gr.Row():
        with gr.Column(scale=1):
            forecast_output = gr.Textbox(
                label="Forecast",
                lines=12,
                placeholder="Select a location or mountain peak to see the forecast..."
            )
        with gr.Column(scale=1):
            snow_forecast = gr.Image(
                label="12-Hour Snow Forecast Animation",
                show_label=True,
                container=True,
                type="filepath"
            )
    
    # Handle submit button click
    submit_btn.click(
        fn=update_weather,
        inputs=[lat_input, lon_input],
        outputs=[
            forecast_output,
            snow_forecast,
            map_display
        ]
    )
    
    # Handle peak button clicks
    for i, peak_name in enumerate(MONTANA_PEAKS.keys()):
        peak_buttons[i].click(
            fn=make_peak_click_handler(peak_name),
            inputs=[],
            outputs=[lat_input, lon_input]
        ).then(
            fn=update_weather,
            inputs=[lat_input, lon_input],
            outputs=[
                forecast_output,
                snow_forecast,
                map_display
            ]
        )
    
    gr.Markdown("""
    ## Instructions
    1. Use the quick access buttons to check specific Montana peaks
    2. Or enter coordinates manually / click on the map
    3. Click "Get Weather" to see the forecast and radar imagery
    
    **Montana Peaks Included:**
    - Lone Peak (Big Sky): 45°16′41″N 111°27′01″W
    - Sacajawea Peak: 45°53′45″N 110°58′7″W
    - Pioneer Mountain: 45°13′55″N 111°27′2″W
    
    **Weather Display Information:**
    The display rotates through three views:
    1. Local Radar (KMSX - Missoula)
       - Coverage: Western Montana mountain regions
       - Updates: Every ~10 minutes
    2. National Radar Loop
       - Shows broader weather patterns
       - Helps track incoming systems
    3. Short Term Forecast
       - Predicted conditions
       - Weather warnings and advisories
    
    **Radar Color Legend:**
    - Light green: Light precipitation
    - Dark green/Yellow: Moderate precipitation
    - Red: Heavy precipitation
    - Pink/White: Possible snow/mixed precipitation
    - Blue: Light snow
    - Dark blue: Heavy snow
    
    **Snow Quality Guide:**
    - Ultra-light powder: ≤ 50 kg/m³
    - Light powder: 51-70 kg/m³
    - Moderate powder: 71-100 kg/m³
    - Packed powder: 101-150 kg/m³
    - Wet snow: 151-200 kg/m³
    - Dense/settled snow: > 200 kg/m³
    
    **Note**: This app uses the NOAA Weather API and may have occasional delays or service interruptions.
    Mountain weather can change rapidly - always check multiple sources for safety.
    Snow density and quality estimates are approximate and based on temperature only.
    Local conditions, wind, and other factors can significantly affect actual snow conditions.
    """)

# For Hugging Face Spaces, we need to export the demo
demo.queue().launch()