import streamlit as st
import ee
import geemap.foliumap as geemap
import folium
import pandas as pd
import numpy as np
import plotly.express as px
import plotly.graph_objects as go
from datetime import datetime, date
import json

# Configure Streamlit page
st.set_page_config(
    page_title="Interactive Landsat 9 Analysis",
    page_icon="🛰️",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Initialize Earth Engine function
@st.cache_resource
def init_ee_with_project(project_id=None):
    try:
        if project_id:
            ee.Initialize(project=project_id)
        else:
            ee.Initialize()
        return True, None
    except Exception as e:
        return False, str(e)

# Main title and description
st.title("🛰️ Interactive Landsat 9 OLI/TIRS Analysis")
st.markdown("""
This interactive application allows you to analyze Landsat 9 satellite imagery with various spectral indices 
and visualizations. Customize your analysis parameters using the sidebar controls.
""")

# Earth Engine initialization with project handling
ee_initialized = False

# Try to initialize without project first
success, error = init_ee_with_project()

if not success:
    st.warning("Earth Engine initialization failed. Please provide your Google Cloud Project ID.")
    st.info("""
    **To get your Project ID:**
    1. Go to https://console.cloud.google.com/
    2. Create a new project or select an existing one
    3. Enable the Earth Engine API for your project
    4. Copy the Project ID from the project selector
    """)
    
    project_id = st.text_input(
        "Enter your Google Cloud Project ID:",
        help="Find this in your Google Cloud Console dashboard"
    )
    
    if project_id:
        success, error = init_ee_with_project(project_id)
        if success:
            st.success(f"Successfully initialized Earth Engine with project: {project_id}")
            ee_initialized = True
        else:
            st.error(f"Failed to initialize with project {project_id}: {error}")
    
    if not ee_initialized:
        st.stop()
else:
    ee_initialized = True
    st.success("Earth Engine initialized successfully!")

# Sidebar controls
st.sidebar.header("🎛️ Analysis Parameters")

# Location settings
st.sidebar.subheader("📍 Location Settings")
center_lat = st.sidebar.number_input("Center Latitude", value=34.741, format="%.3f")
center_lon = st.sidebar.number_input("Center Longitude", value=71.878, format="%.3f")
buffer_size = st.sidebar.slider("Buffer Size (km)", min_value=10, max_value=100, value=50)

# Date range settings
st.sidebar.subheader("📅 Date Range")
start_date = st.sidebar.date_input("Start Date", value=date(2022, 1, 1))
end_date = st.sidebar.date_input("End Date", value=date(2022, 12, 31))

# Cloud cover filter
cloud_cover = st.sidebar.slider("Maximum Cloud Cover (%)", min_value=0, max_value=100, value=20)

# Visualization options
st.sidebar.subheader("🎨 Visualization Options")
vis_options = {
    "True Color (432)": {"bands": ["B4", "B3", "B2"], "min": 8000, "max": 18000},
    "False Color (543)": {"bands": ["B5", "B4", "B3"], "min": 8000, "max": 20000},
    "Agriculture (654)": {"bands": ["B6", "B5", "B4"], "min": 8000, "max": 20000},
    "Geology (764)": {"bands": ["B7", "B6", "B4"], "min": 8000, "max": 20000},
    "Bathymetric (431)": {"bands": ["B4", "B3", "B1"], "min": 8000, "max": 18000}
}

selected_vis = st.sidebar.selectbox("Select Band Combination", list(vis_options.keys()))

# Index calculations
st.sidebar.subheader("📊 Spectral Indices")
show_indices = st.sidebar.multiselect(
    "Select Indices to Display",
    ["NDVI", "EVI", "SAVI", "NDWI", "MNDWI", "NDBI", "NBR", "NDSI"],
    default=["NDVI", "NDWI"]
)

# Analysis function
@st.cache_data
def perform_analysis(lat, lon, buffer_km, start_dt, end_dt, max_cloud):
    # Define area of interest
    aoi = ee.Geometry.Point([lon, lat]).buffer(buffer_km * 1000)
    
    # Load Landsat 9 dataset
    dataset = ee.ImageCollection('LANDSAT/LC09/C02/T1') \
                .filterDate(start_dt.strftime('%Y-%m-%d'), end_dt.strftime('%Y-%m-%d')) \
                .filterBounds(aoi) \
                .filter(ee.Filter.lt('CLOUD_COVER', max_cloud)) \
                .sort('CLOUD_COVER')
    
    # Get dataset info
    size = dataset.size().getInfo()
    
    if size == 0:
        return None, None, "No images found for the specified criteria."
    
    # Create composite
    composite = dataset.median().clip(aoi)
    
    # Calculate indices
    indices = {}
    
    # Vegetation indices
    indices['NDVI'] = composite.normalizedDifference(['B5', 'B4']).rename('NDVI')
    indices['EVI'] = composite.expression(
        '2.5 * ((NIR - RED) / (NIR + 6 * RED - 7.5 * BLUE + 1))', {
            'NIR': composite.select('B5'),
            'RED': composite.select('B4'),
            'BLUE': composite.select('B2')
        }).rename('EVI')
    indices['SAVI'] = composite.expression(
        '((NIR - RED) / (NIR + RED + 0.5)) * (1 + 0.5)', {
            'NIR': composite.select('B5'),
            'RED': composite.select('B4')
        }).rename('SAVI')
    
    # Water indices
    indices['NDWI'] = composite.normalizedDifference(['B3', 'B5']).rename('NDWI')
    indices['MNDWI'] = composite.normalizedDifference(['B3', 'B6']).rename('MNDWI')
    
    # Urban indices
    indices['NDBI'] = composite.normalizedDifference(['B6', 'B5']).rename('NDBI')
    
    # Burn index
    indices['NBR'] = composite.normalizedDifference(['B5', 'B7']).rename('NBR')
    
    # Snow index
    indices['NDSI'] = composite.normalizedDifference(['B3', 'B6']).rename('NDSI')
    
    # Calculate statistics
    stats = {}
    for name, index in indices.items():
        stat = index.reduceRegion(
            reducer=ee.Reducer.mean().combine(ee.Reducer.stdDev(), sharedInputs=True),
            geometry=aoi,
            scale=30,
            maxPixels=1e9
        ).getInfo()
        stats[name] = stat
    
    return composite, indices, stats, aoi, size

# Run analysis button
if st.sidebar.button("🚀 Run Analysis", type="primary"):
    with st.spinner("Analyzing Landsat 9 imagery..."):
        try:
            result = perform_analysis(
                center_lat, center_lon, buffer_size, 
                start_date, end_date, cloud_cover
            )
            
            if result[0] is None:
                st.error(result[2])
            else:
                composite, indices, stats, aoi, image_count = result
                st.success(f"Analysis complete! Found {image_count} images.")
                
                # Store results in session state
                st.session_state.composite = composite
                st.session_state.indices = indices
                st.session_state.stats = stats
                st.session_state.aoi = aoi
                st.session_state.analysis_params = {
                    'lat': center_lat, 'lon': center_lon, 'buffer': buffer_size
                }
                
        except Exception as e:
            st.error(f"Analysis failed: {str(e)}")

# Display results if analysis has been run
if 'composite' in st.session_state:
    # Create tabs for different views
    tab1, tab2, tab3, tab4 = st.tabs(["🗺️ Interactive Map", "📊 Statistics", "📈 Charts", "📋 Data Export"])
    
    with tab1:
        st.subheader("Interactive Satellite Imagery Map")
        
        # Create the map
        Map = geemap.Map(center=[st.session_state.analysis_params['lat'], 
                                st.session_state.analysis_params['lon']], zoom=12)
        
        # Add selected visualization
        vis_params = vis_options[selected_vis]
        Map.addLayer(
            st.session_state.composite.select(vis_params['bands']),
            {
                'min': vis_params['min'],
                'max': vis_params['max'],
                'bands': vis_params['bands']
            },
            selected_vis
        )
        
        # Add selected indices
        index_vis_params = {
            'NDVI': {'min': -0.5, 'max': 0.8, 'palette': ['red', 'yellow', 'green']},
            'EVI': {'min': -0.5, 'max': 0.8, 'palette': ['red', 'yellow', 'green']},
            'SAVI': {'min': -0.5, 'max': 0.8, 'palette': ['red', 'yellow', 'green']},
            'NDWI': {'min': -0.5, 'max': 0.5, 'palette': ['white', 'blue']},
            'MNDWI': {'min': -0.5, 'max': 0.5, 'palette': ['white', 'blue']},
            'NDBI': {'min': -0.5, 'max': 0.5, 'palette': ['blue', 'white', 'red']},
            'NBR': {'min': -0.5, 'max': 0.5, 'palette': ['green', 'yellow', 'red']},
            'NDSI': {'min': 0, 'max': 0.8, 'palette': ['red', 'yellow', 'white']}
        }
        
        for index_name in show_indices:
            if index_name in st.session_state.indices:
                Map.addLayer(
                    st.session_state.indices[index_name].selfMask(),
                    index_vis_params[index_name],
                    index_name,
                    False  # Start with layer hidden
                )
        
        # Add AOI boundary
        Map.addLayer(st.session_state.aoi, {'color': 'yellow'}, 'Area of Interest')
        
        # Display the map
        Map.to_streamlit(height=600)
    
    with tab2:
        st.subheader("📊 Spectral Index Statistics")
        
        # Create statistics dataframe
        stats_data = []
        for index_name, stat_dict in st.session_state.stats.items():
            if index_name in show_indices:
                mean_key = f"{index_name}_mean"
                std_key = f"{index_name}_stdDev"
                
                mean_val = stat_dict.get(mean_key, 0)
                std_val = stat_dict.get(std_key, 0)
                
                stats_data.append({
                    'Index': index_name,
                    'Mean': round(mean_val, 4) if mean_val else 0,
                    'Standard Deviation': round(std_val, 4) if std_val else 0,
                    'Range': f"{round(mean_val - std_val, 4)} to {round(mean_val + std_val, 4)}" if mean_val and std_val else "N/A"
                })
        
        if stats_data:
            df_stats = pd.DataFrame(stats_data)
            st.dataframe(df_stats, use_container_width=True)
            
            # Create bar chart of mean values
            fig_bar = px.bar(
                df_stats, 
                x='Index', 
                y='Mean', 
                title='Mean Values of Spectral Indices',
                color='Mean',
                color_continuous_scale='viridis'
            )
            fig_bar.update_layout(height=400)
            st.plotly_chart(fig_bar, use_container_width=True)
    
    with tab3:
        st.subheader("📈 Data Visualization")
        
        # Index interpretation guide
        with st.expander("📖 Index Interpretation Guide"):
            st.markdown("""
            **Vegetation Indices:**
            - **NDVI**: -1 to 1 (higher = more vegetation)
            - **EVI**: -1 to 1 (enhanced vegetation, reduces atmospheric effects)
            - **SAVI**: -1 to 1 (soil-adjusted vegetation index)
            
            **Water Indices:**
            - **NDWI**: -1 to 1 (higher = more water)
            - **MNDWI**: -1 to 1 (modified NDWI, better for water detection)
            
            **Urban/Built-up:**
            - **NDBI**: -1 to 1 (higher = more built-up areas)
            
            **Environmental:**
            - **NBR**: -1 to 1 (normalized burn ratio for fire detection)
            - **NDSI**: 0 to 1 (normalized difference snow index)
            """)
        
        # Create comparison chart if multiple indices selected
        if len(show_indices) > 1:
            comparison_data = []
            for index_name in show_indices:
                if index_name in st.session_state.stats:
                    stat_dict = st.session_state.stats[index_name]
                    mean_key = f"{index_name}_mean"
                    mean_val = stat_dict.get(mean_key, 0)
                    comparison_data.append({'Index': index_name, 'Mean Value': mean_val})
            
            if comparison_data:
                df_comparison = pd.DataFrame(comparison_data)
                fig_comparison = px.bar(
                    df_comparison,
                    x='Index',
                    y='Mean Value',
                    title='Spectral Index Comparison',
                    color='Mean Value',
                    color_continuous_scale='RdYlGn'
                )
                fig_comparison.update_layout(height=400)
                st.plotly_chart(fig_comparison, use_container_width=True)
    
    with tab4:
        st.subheader("📋 Data Export Options")
        
        col1, col2 = st.columns(2)
        
        with col1:
            st.markdown("**📊 Statistics Export**")
            if st.button("Download Statistics as CSV"):
                if 'stats' in st.session_state:
                    stats_data = []
                    for index_name, stat_dict in st.session_state.stats.items():
                        mean_key = f"{index_name}_mean"
                        std_key = f"{index_name}_stdDev"
                        stats_data.append({
                            'Index': index_name,
                            'Mean': stat_dict.get(mean_key, 0),
                            'Standard_Deviation': stat_dict.get(std_key, 0)
                        })
                    
                    df_export = pd.DataFrame(stats_data)
                    csv = df_export.to_csv(index=False)
                    st.download_button(
                        label="📥 Download CSV",
                        data=csv,
                        file_name=f"landsat9_analysis_{datetime.now().strftime('%Y%m%d_%H%M%S')}.csv",
                        mime="text/csv"
                    )
        
        with col2:
            st.markdown("**🗺️ Map Export**")
            st.info("Use the map's built-in export tools to save visualizations as images.")
        
        # Analysis summary
        st.markdown("**📋 Analysis Summary**")
        summary_info = f"""
        - **Location**: {st.session_state.analysis_params['lat']:.3f}°N, {st.session_state.analysis_params['lon']:.3f}°E
        - **Buffer Size**: {st.session_state.analysis_params['buffer']} km
        - **Date Range**: {start_date} to {end_date}
        - **Cloud Cover**: ≤ {cloud_cover}%
        - **Selected Visualization**: {selected_vis}
        - **Active Indices**: {', '.join(show_indices)}
        """
        st.markdown(summary_info)

else:
    # Instructions when no analysis has been run
    st.info("👈 Configure your analysis parameters in the sidebar and click 'Run Analysis' to get started!")
    
    # Feature overview
    st.markdown("""
    ## 🌟 Features
    
    **🛰️ Satellite Data Analysis**
    - Landsat 9 OLI/TIRS imagery (30m resolution)
    - Customizable date ranges and cloud cover filtering
    - Multiple band combinations for different applications
    
    **📊 Spectral Indices**
    - Vegetation: NDVI, EVI, SAVI
    - Water: NDWI, MNDWI  
    - Urban: NDBI
    - Environmental: NBR (burn), NDSI (snow)
    
    **🎨 Interactive Visualizations**
    - True color, false color, and specialized composites
    - Statistical analysis and charting
    - Export capabilities for further analysis
    
    **🗺️ Interactive Mapping**
    - Zoom, pan, and layer control
    - Real-time visualization switching
    - Area of interest boundary display
    """)

# Footer
st.markdown("---")
st.markdown("Built with ❤️ using Streamlit, Google Earth Engine, and geemap")