import streamlit as st
import ee
import os
import pandas as pd
import geopandas as gpd
from datetime import datetime
import leafmap.foliumap as leafmap
import re
from shapely.geometry import base
from xml.etree import ElementTree as XET
from concurrent.futures import ThreadPoolExecutor, as_completed
import time
import json
import math

# Set up the page layout
st.set_page_config(layout="wide")

# Custom button styling
m = st.markdown(
    """
    <style>
    div.stButton > button:first-child {
        background-color: #006400;
        color:#ffffff;
    }
    </style>""",
    unsafe_allow_html=True,
)

# Logo and Title
st.write(
    f"""
    <div style="display: flex; justify-content: space-between; align-items: center;">
        <img src="https://huggingface.co/spaces/YashMK89/SATRANG/resolve/main/ISRO_Logo.png" style="width: 20%; margin-right: auto;">
        <img src="https://huggingface.co/spaces/YashMK89/SATRANG/resolve/main/SAC_Logo.png" style="width: 20%; margin-left: auto;">
    </div>
    """,
    unsafe_allow_html=True,
)
st.markdown(
    f"""
    <div style="display: flex; flex-direction: column; align-items: center;">
        <img src="https://huggingface.co/spaces/YashMK89/SATRANG/resolve/main/SATRANG.png" style="width: 30%;">
        <h3 style="text-align: center; margin: 0;">( Spatial and Temporal Aggregation for Remote-sensing Analysis of GEE Data )</h3>
    </div>
    <hr>
    """,
    unsafe_allow_html=True,
)

# Authenticate and initialize Earth Engine
earthengine_credentials = os.environ.get("EE_Authentication")
os.makedirs(os.path.expanduser("~/.config/earthengine/"), exist_ok=True)
with open(os.path.expanduser("~/.config/earthengine/credentials"), "w") as f:
    f.write(earthengine_credentials)
ee.Initialize(project='ee-yashsacisro24')


# Helper function to get reducer
def get_reducer(reducer_name):
    reducers = {
        'mean': ee.Reducer.mean(),
        'sum': ee.Reducer.sum(),
        'median': ee.Reducer.median(),
        'min': ee.Reducer.min(),
        'max': ee.Reducer.max(),
        'count': ee.Reducer.count(),
    }
    return reducers.get(reducer_name.lower(), ee.Reducer.mean())

# Function to convert geometry to Earth Engine format
def convert_to_ee_geometry(geometry):
    if isinstance(geometry, base.BaseGeometry):
        if geometry.is_valid:
            geojson = geometry.__geo_interface__
            return ee.Geometry(geojson)
        else:
            raise ValueError("Invalid geometry: The polygon geometry is not valid.")
    elif isinstance(geometry, dict) or isinstance(geometry, str):
        try:
            if isinstance(geometry, str):
                geometry = json.loads(geometry)
            if 'type' in geometry and 'coordinates' in geometry:
                return ee.Geometry(geometry)
            else:
                raise ValueError("GeoJSON format is invalid.")
        except Exception as e:
            raise ValueError(f"Error parsing GeoJSON: {e}")
    elif isinstance(geometry, str) and geometry.lower().endswith(".kml"):
        try:
            tree = XET.parse(geometry)
            kml_root = tree.getroot()
            kml_namespace = {'kml': 'http://www.opengis.net/kml/2.2'}
            coordinates = kml_root.findall(".//kml:coordinates", kml_namespace)
            if coordinates:
                coords_text = coordinates[0].text.strip()
                coords = coords_text.split()
                coords = [tuple(map(float, coord.split(','))) for coord in coords]
                geojson = {"type": "Polygon", "coordinates": [coords]}
                return ee.Geometry(geojson)
            else:
                raise ValueError("KML does not contain valid coordinates.")
        except Exception as e:
            raise ValueError(f"Error parsing KML: {e}")
    else:
        raise ValueError("Unsupported geometry input type. Supported types are Shapely, GeoJSON, and KML.")

# Function to calculate custom formula
def calculate_custom_formula(image, geometry, selected_bands, custom_formula, reducer_choice, scale=30):
    try:
        # Create a dictionary to hold band values
        band_values = {}
        
        # Get all band names from the image
        band_names = image.bandNames().getInfo()
        
        # Verify all selected bands exist in the image
        for band in selected_bands:
            if band not in band_names:
                raise ValueError(f"Band '{band}' not found in the dataset.")
        
        # Select only the bands we need for calculation
        image = image.select(selected_bands)
        
        # Apply the reducer to get statistics for each band
        reducer = get_reducer(reducer_choice)
        reduced_values = image.reduceRegion(
            reducer=reducer,
            geometry=geometry,
            scale=scale,
            bestEffort=True
        ).getInfo()
        
        # Prepare the formula by replacing band names with their values
        formula = custom_formula
        for band in selected_bands:
            value = reduced_values.get(band, 0)
            if value is None:
                value = 0
            formula = formula.replace(band, str(value))
        
        # Safely evaluate the formula
        result = eval(formula, {"__builtins__": None}, {"math": math})
        
        if not isinstance(result, (int, float)):
            raise ValueError("Formula did not result in a numeric value.")
        
        return ee.Image.constant(result).rename('custom_result')
    
    except ZeroDivisionError:
        st.error("Error: Division by zero in the formula.")
        return ee.Image(0).rename('custom_result').set('error', 'Division by zero')
    except SyntaxError:
        st.error(f"Error: Invalid syntax in formula '{custom_formula}'.")
        return ee.Image(0).rename('custom_result').set('error', 'Invalid syntax')
    except ValueError as e:
        st.error(f"Error: {str(e)}")
        return ee.Image(0).rename('custom_result').set('error', str(e))
    except Exception as e:
        st.error(f"Unexpected error: {e}")
        return ee.Image(0).rename('custom_result').set('error', str(e))

# Cloud Masking Function
def apply_generic_cloud_mask(image, cloud_band, cloud_threshold=None):
    """Apply cloud masking based on common band patterns"""
    if not cloud_band:
        return image
        
    band_names = image.bandNames().getInfo()

    if cloud_band not in band_names:
        return image

    try:
        # Case 1: Bitmask (Landsat-style QA bands)
        if any(keyword in cloud_band.lower() for keyword in ['qa', 'pixel_qa', 'quality']):
            cloud_bit = 1 << 3  # Typically bit 3 for clouds
            shadow_bit = 1 << 4  # Typically bit 4 for shadows
            mask = image.select(cloud_band).bitwiseAnd(cloud_bit).eq(0) \
                .And(image.select(cloud_band).bitwiseAnd(shadow_bit).eq(0))

        # Case 2: Simple cloud mask (0=clear, 1=cloud)
        elif 'cloud' in cloud_band.lower():
            mask = image.select(cloud_band).eq(0)

        # Case 3: CFMask/FMask style (0=clear, 1=water, 2=shadow, 3=snow, 4=cloud)
        elif any(keyword in cloud_band.lower() for keyword in ['cfmask', 'fmask']):
            mask = image.select(cloud_band).lt(2)  # Accept clear and water

        # Default case: assume 0=clear, 1=cloud
        else:
            mask = image.select(cloud_band).eq(0)

        masked_image = image.updateMask(mask)

        if cloud_threshold is not None:
            cloud_area = image.select(cloud_band).gt(0).rename('cloud')
            stats = cloud_area.reduceRegion(
                reducer=ee.Reducer.mean(),
                geometry=image.geometry(),
                scale=30,
                maxPixels=1e9
            )
            cloud_percent = ee.Number(stats.get('cloud')).multiply(100)
            masked_image = masked_image.set('cloud_percent', cloud_percent)

        return masked_image

    except Exception as e:
        st.warning(f"Error applying cloud mask: {str(e)}")
        return image

# Aggregation functions
def aggregate_data_daily(collection):
    def set_day_start(image):
        date = ee.Date(image.get('system:time_start')).format('YYYY-MM-dd')
        return image.set('day_start', date)
    
    collection = collection.map(set_day_start)
    grouped_by_day = collection.aggregate_array('day_start').distinct()
    
    def calculate_daily_mean(day_start):
        daily_collection = collection.filter(ee.Filter.eq('day_start', day_start))
        daily_mean = daily_collection.mean()
        return daily_mean.set('day_start', day_start)
    
    daily_images = ee.List(grouped_by_day.map(calculate_daily_mean))
    return ee.ImageCollection(daily_images)

def aggregate_data_weekly(collection):
    def set_week_start(image):
        date = ee.Date(image.get('system:time_start'))
        days_since_week_start = date.getRelative('day', 'week')
        offset = ee.Number(days_since_week_start).multiply(-1)
        week_start = date.advance(offset, 'day').format('YYYY-MM-dd')
        return image.set('week_start', week_start)
    
    collection = collection.map(set_week_start)
    grouped_by_week = collection.aggregate_array('week_start').distinct()
    
    def calculate_weekly_mean(week_start):
        weekly_collection = collection.filter(ee.Filter.eq('week_start', week_start))
        weekly_mean = weekly_collection.mean()
        return weekly_mean.set('week_start', week_start)
    
    weekly_images = ee.List(grouped_by_week.map(calculate_weekly_mean))
    return ee.ImageCollection(weekly_images)

def aggregate_data_monthly(collection, start_date, end_date):
    collection = collection.filterDate(start_date, end_date)
    collection = collection.map(lambda image: image.set('month', ee.Date(image.get('system:time_start')).format('YYYY-MM')))
    grouped_by_month = collection.aggregate_array('month').distinct()
    
    def calculate_monthly_mean(month):
        monthly_collection = collection.filter(ee.Filter.eq('month', month))
        monthly_mean = monthly_collection.mean()
        return monthly_mean.set('month', month)
    
    monthly_images = ee.List(grouped_by_month.map(calculate_monthly_mean))
    return ee.ImageCollection(monthly_images)

def aggregate_data_yearly(collection):
    collection = collection.map(lambda image: image.set('year', ee.Date(image.get('system:time_start')).format('YYYY')))
    grouped_by_year = collection.aggregate_array('year').distinct()
    
    def calculate_yearly_mean(year):
        yearly_collection = collection.filter(ee.Filter.eq('year', year))
        yearly_mean = yearly_collection.mean()
        return yearly_mean.set('year', year)
    
    yearly_images = ee.List(grouped_by_year.map(calculate_yearly_mean))
    return ee.ImageCollection(yearly_images)

def aggregate_data_custom(collection):
    collection = collection.map(lambda image: image.set('date', ee.Date(image.get('system:time_start')).format('YYYY-MM-dd')))
    return collection

# Worker function for processing a single geometry
def process_single_geometry(row, start_date_str, end_date_str, dataset_id, selected_bands, reducer_choice, shape_type, 
                          aggregation_period, custom_formula, kernel_size=None, include_boundary=None, 
                          default_scale=None, apply_cloud_mask=False, cloud_threshold=None, cloud_band=None):
    try:
        image_count = 0
        if shape_type.lower() == "point":
            latitude = row.get('latitude')
            longitude = row.get('longitude')
            if pd.isna(latitude) or pd.isna(longitude):
                return None, 0
            location_name = row.get('name', f"Location_{row.name}")
            
            if kernel_size == "3x3 Kernel":
                buffer_size = 45  # 90m x 90m
                roi = ee.Geometry.Point([longitude, latitude]).buffer(buffer_size).bounds()
            elif kernel_size == "5x5 Kernel":
                buffer_size = 75  # 150m x 150m
                roi = ee.Geometry.Point([longitude, latitude]).buffer(buffer_size).bounds()
            else:  # Point
                roi = ee.Geometry.Point([longitude, latitude])
                
        elif shape_type.lower() == "polygon":
            polygon_geometry = row.get('geometry')
            location_name = row.get('name', f"Polygon_{row.name}")
            try:
                roi = convert_to_ee_geometry(polygon_geometry)
                if not include_boundary:
                    roi = roi.buffer(-30).bounds()
            except ValueError as e:
                st.warning(f"Skipping invalid geometry for {location_name}: {str(e)}")
                return None, 0

        # Initialize the image collection
        collection = ee.ImageCollection(dataset_id) \
            .filterDate(ee.Date(start_date_str), ee.Date(end_date_str)) \
            .filterBounds(roi) \
            .select(selected_bands)

        # Apply cloud masking if requested
        if apply_cloud_mask and cloud_band:
            # First filter by built-in cloud percentage if available
            if cloud_threshold is not None:
                try:
                    metadata_props = collection.first().propertyNames().getInfo()
                    if 'CLOUDY_PIXEL_PERCENTAGE' in metadata_props:
                        collection = collection.filter(ee.Filter.lte('CLOUDY_PIXEL_PERCENTAGE', cloud_threshold))
                    elif 'cloud_cover' in metadata_props:
                        collection = collection.filter(ee.Filter.lte('cloud_cover', cloud_threshold))
                except Exception as e:
                    pass

            # Apply pixel-level cloud masking
            collection = collection.map(lambda img: apply_generic_cloud_mask(img, cloud_band, cloud_threshold))

            # Filter again using calculated cloud percentage if applicable
            if cloud_threshold is not None:
                try:
                    cloud_prop_exists = 'cloud_percent' in collection.first().propertyNames().getInfo()
                    if cloud_prop_exists:
                        collection = collection.filter(ee.Filter.lte('cloud_percent', cloud_threshold))
                except:
                    pass

        initial_count = collection.size().getInfo()
        image_count += initial_count

        # Apply temporal aggregation
        if aggregation_period.lower() == 'custom (start date to end date)':
            collection = aggregate_data_custom(collection)
        elif aggregation_period.lower() == 'daily':
            collection = aggregate_data_daily(collection)
        elif aggregation_period.lower() == 'weekly':
            collection = aggregate_data_weekly(collection)
        elif aggregation_period.lower() == 'monthly':
            collection = aggregate_data_monthly(collection, start_date_str, end_date_str)
        elif aggregation_period.lower() == 'yearly':
            collection = aggregate_data_yearly(collection)

        image_list = collection.toList(collection.size())
        aggregated_results = []

        for i in range(image_list.size().getInfo()):
            image = ee.Image(image_list.get(i))
            
            # Get the timestamp based on aggregation period
            if aggregation_period.lower() == 'custom (start date to end date)':
                timestamp = image.get('date')
                period_label = 'Date'
                date = ee.String(timestamp).getInfo()
            elif aggregation_period.lower() == 'daily':
                timestamp = image.get('day_start')
                period_label = 'Day'
                date = ee.String(timestamp).getInfo()
            elif aggregation_period.lower() == 'weekly':
                timestamp = image.get('week_start')
                period_label = 'Week'
                date = ee.String(timestamp).getInfo()
            elif aggregation_period.lower() == 'monthly':
                timestamp = image.get('month')
                period_label = 'Month'
                date = ee.String(timestamp).getInfo()
            elif aggregation_period.lower() == 'yearly':
                timestamp = image.get('year')
                period_label = 'Year'
                date = ee.String(timestamp).getInfo()

            # Calculate the custom formula
            index_image = calculate_custom_formula(image, roi, selected_bands, custom_formula, reducer_choice, scale=default_scale)
            
            try:
                index_value = index_image.reduceRegion(
                    reducer=get_reducer(reducer_choice),
                    geometry=roi,
                    scale=default_scale,
                    bestEffort=True
                ).get('custom_result')
                
                calculated_value = index_value.getInfo()
                if calculated_value is None:
                    calculated_value = 0
                    
                if isinstance(calculated_value, (int, float)):
                    result = {
                        'Location Name': location_name,
                        period_label: date,
                        'Start Date': start_date_str,
                        'End Date': end_date_str,
                        'Calculated Value': calculated_value
                    }
                    if shape_type.lower() == 'point':
                        result['Latitude'] = latitude
                        result['Longitude'] = longitude
                    aggregated_results.append(result)
            except Exception as e:
                st.warning(f"Error retrieving value for {location_name}: {e}")

        return aggregated_results, image_count
        
    except Exception as e:
        st.error(f"Error processing {location_name if 'location_name' in locals() else 'unknown location'}: {str(e)}")
        return None, 0

# Main processing function
def process_aggregation(locations_df, start_date_str, end_date_str, dataset_id, selected_bands, reducer_choice, shape_type, 
                      aggregation_period, custom_formula="", kernel_size=None, include_boundary=None, 
                      default_scale=None, apply_cloud_mask=False, cloud_threshold=None, cloud_band=None):
    
    # Auto-detect scale if not provided
    if default_scale is None:
        try:
            collection = ee.ImageCollection(dataset_id)
            default_scale = collection.first().select(0).projection().nominalScale().getInfo()
            if not isinstance(default_scale, (int, float)) or default_scale <= 0:
                default_scale = 30
        except:
            default_scale = 30

    aggregated_results = []
    total_images = 0
    total_steps = len(locations_df)
    progress_bar = st.progress(0)
    progress_text = st.empty()
    start_time = time.time()

    with ThreadPoolExecutor(max_workers=10) as executor:
        futures = []
        for idx, row in locations_df.iterrows():
            future = executor.submit(
                process_single_geometry,
                row,
                start_date_str,
                end_date_str,
                dataset_id,
                selected_bands,
                reducer_choice,
                shape_type,
                aggregation_period,
                custom_formula,
                kernel_size,
                include_boundary,
                default_scale,
                apply_cloud_mask,
                cloud_threshold,
                cloud_band
            )
            futures.append(future)
            
        completed = 0
        for future in as_completed(futures):
            result, image_count = future.result()
            if result:
                aggregated_results.extend(result)
                total_images += image_count
            completed += 1
            progress_percentage = completed / total_steps
            progress_bar.progress(progress_percentage)
            progress_text.markdown(f"Processing: {int(progress_percentage * 100)}% (Total images: {total_images})")

    end_time = time.time()
    processing_time = end_time - start_time

    if aggregated_results:
        result_df = pd.DataFrame(aggregated_results)
        if aggregation_period.lower() == 'custom (start date to end date)':
            agg_dict = {
                'Start Date': 'first',
                'End Date': 'first',
                'Calculated Value': 'mean'
            }
            if shape_type.lower() == 'point':
                agg_dict['Latitude'] = 'first'
                agg_dict['Longitude'] = 'first'
            aggregated_output = result_df.groupby('Location Name').agg(agg_dict).reset_index()
            aggregated_output.rename(columns={'Calculated Value': 'Aggregated Value'}, inplace=True)
            return aggregated_output.to_dict(orient='records'), processing_time, total_images
        else:
            return result_df.to_dict(orient='records'), processing_time, total_images
    return [], processing_time, total_images

# Streamlit UI Starts Here
st.markdown("<h5>Image Collection</h5>", unsafe_allow_html=True)
imagery_base = st.selectbox("Select Imagery Base", ["Sentinel", "Landsat", "MODIS", "VIIRS", "Custom Input"], index=0)
st.markdown("<hr><h5><b>{}</b></h5>".format(imagery_base), unsafe_allow_html=True)

# Initialize data as an empty dictionary
data = {}

if imagery_base == "Sentinel":
    dataset_file = "sentinel_datasets.json"
elif imagery_base == "Landsat":
    dataset_file = "landsat_datasets.json"
elif imagery_base == "MODIS":
    dataset_file = "modis_datasets.json"
elif imagery_base == "VIIRS":
    dataset_file = "viirs_datasets.json"
else:
    dataset_file = ""

# Load dataset JSON files
if imagery_base != "Custom Input":
    try:
        with open(dataset_file) as f:
            data = json.load(f)
    except FileNotFoundError:
        st.error(f"Dataset file '{dataset_file}' not found.")
        data = {}

# Handle Custom Dataset Input
elif imagery_base == "Custom Input":
    custom_dataset_id = st.text_input(
        "Enter Custom Earth Engine Dataset ID (e.g., MODIS/006/MOD13Q1)",
        value="",
        help="Enter the full path of the EE dataset (e.g., 'COPERNICUS/S2_SR')"
    )
    if custom_dataset_id:
        try:
            if custom_dataset_id.startswith("ee.ImageCollection("):
                custom_dataset_id = custom_dataset_id.replace("ee.ImageCollection('", "").replace("')", "").strip()
            collection = ee.ImageCollection(custom_dataset_id)
            first_image = collection.first()
            band_names = first_image.bandNames().getInfo()
            try:
                default_scale = first_image.select(0).projection().nominalScale().getInfo()
                if not isinstance(default_scale, (int, float)) or default_scale <= 0:
                    raise ValueError("Invalid scale from GEE")
            except:
                default_scale = 30
            data = {
                f"Custom Dataset: {custom_dataset_id}": {
                    "sub_options": {custom_dataset_id: f"Custom Dataset ({custom_dataset_id})"},
                    "bands": {custom_dataset_id: band_names}
                }
            }
            st.success(f"✅ Successfully loaded: {custom_dataset_id}")
            st.write(f"Available Bands: {', '.join(band_names)}")
            st.write(f"Native Scale: {default_scale} meters")
        except Exception as e:
            st.error(f"Error loading dataset: {e}")
            data = {}
    else:
        st.warning("Please enter a custom dataset ID to proceed.")
        data = {}

# If no valid dataset available
if not data:
    st.error("No valid dataset available. Please check your inputs.")
    st.stop()

main_selection = st.selectbox(f"Select {imagery_base} Dataset Category", list(data.keys()))
sub_selection = None
dataset_id = None

if main_selection:
    sub_options = data[main_selection]["sub_options"]
    sub_selection = st.selectbox(f"Select Specific {imagery_base} Dataset ID", list(sub_options.keys()))
    if sub_selection:
        st.write(f"You selected: {main_selection} → {sub_options[sub_selection]}")
        st.write(f"Dataset ID: {sub_selection}")
        dataset_id = sub_selection
        try:
            collection = ee.ImageCollection(dataset_id)
            first_image = collection.first()
            default_scale = first_image.select(0).projection().nominalScale().getInfo()
            st.write(f"Default Scale for Selected Dataset: {default_scale} meters")
        except Exception as e:
            st.error(f"Error fetching default scale: {str(e)}")

# Detect Cloud Support
has_cloud_info = False
cloud_band_candidates = []
cloud_metadata_key = None

if main_selection and sub_selection:
    dataset_bands = data[main_selection]["bands"].get(sub_selection, [])
    st.write(f"Available Bands for {sub_options[sub_selection]}: {', '.join(dataset_bands)}")

    try:
        collection = ee.ImageCollection(dataset_id)
        first_image = collection.first()
        metadata_props = first_image.propertyNames().getInfo()

        if 'CLOUDY_PIXEL_PERCENTAGE' in metadata_props:
            cloud_metadata_key = 'CLOUDY_PIXEL_PERCENTAGE'
            has_cloud_info = True
        elif 'cloud_cover' in metadata_props:
            cloud_metadata_key = 'cloud_cover'
            has_cloud_info = True

        all_bands = first_image.bandNames().getInfo()
        cloud_band_candidates = [
            band for band in all_bands 
            if any(keyword in band.lower() for keyword in ['cloud', 'qa', 'quality', 'cfmask', 'fmask', 'pixel_qa', 'qa_pixel'])
        ]
        if cloud_band_candidates:
            has_cloud_info = True
    except Exception as e:
        st.warning(f"Could not detect cloud support: {str(e)}")
        has_cloud_info = False

# Conditional Cloud Masking UI
apply_cloud_mask = False
cloud_threshold = None
cloud_band = None
st.markdown("<hr><h5><b></b></h5>", unsafe_allow_html=True)
if has_cloud_info:
    st.markdown("<hr><h5><b>Cloud Masking</b></h5>", unsafe_allow_html=True)
    apply_cloud_mask = st.checkbox("Apply Cloud Masking", value=False,
                                 help="Enable to filter out cloudy pixels. Only works if the dataset has cloud information")
    
    if apply_cloud_mask:
        cloud_threshold = st.slider(
            "Maximum Cloud Percentage Allowed (0-100%)",
            min_value=0,
            max_value=100,
            value=20,
            help="Images with cloud coverage above this percentage will be excluded"
        )
        
        if cloud_band_candidates:
            cloud_band = st.selectbox(
                "Select Cloud Mask Band",
                options=cloud_band_candidates,
                index=0,
                help="Select the band that contains cloud information"
            )
            st.info("Common cloud mask values: 0=clear, 1=cloud, 2=shadow, 3=snow, 4=water")
        else:
            st.warning("No cloud mask bands detected in this dataset")
            apply_cloud_mask = False
else:
    st.info("This dataset does not support cloud masking.")

# Formula Input
if main_selection and sub_selection:
    dataset_bands = data[main_selection]["bands"].get(sub_selection, [])
    selected_bands = st.multiselect(
        "Select Bands for Calculation",
        options=dataset_bands,
        default=[dataset_bands[0]] if dataset_bands else [],
        help=f"Select bands from: {', '.join(dataset_bands)}"
    )
    if len(selected_bands) < 1:
        st.warning("Please select at least one band.")
        st.stop()
    if selected_bands:
        if len(selected_bands) == 1:
            default_formula = f"{selected_bands[0]}"
            example = f"'{selected_bands[0]} * 2' or '{selected_bands[0]} + 1'"
        else:
            default_formula = f"({selected_bands[0]} - {selected_bands[1]}) / ({selected_bands[0]} + {selected_bands[1]})"
            example = f"'{selected_bands[0]} * {selected_bands[1]} / 2' or '({selected_bands[0]} - {selected_bands[1]}) / ({selected_bands[0]} + {selected_bands[1]})'"
        custom_formula = st.text_input(
            "Enter Custom Formula (e.g (B8 - B4) / (B8 + B4) , B4*B3/2)",
            value=default_formula,
            help=f"Use only these bands: {', '.join(selected_bands)}. Examples: {example}"
        )

        def validate_formula(formula, selected_bands):
            allowed_chars = set(" +-*/()0123456789.")
            terms = re.findall(r'[a-zA-Z][a-zA-Z0-9_]*', formula)
            invalid_terms = [term for term in terms if term not in selected_bands]
            if invalid_terms:
                return False, f"Invalid terms in formula: {', '.join(invalid_terms)}. Use only {', '.join(selected_bands)}."
            if not all(char in allowed_chars or char in ''.join(selected_bands) for char in formula):
                return False, "Formula contains invalid characters."
            return True, ""
        
        is_valid, error_message = validate_formula(custom_formula, selected_bands)
        if not is_valid:
            st.error(error_message)
            st.stop()
        elif not custom_formula:
            st.warning("Please enter a custom formula to proceed.")
            st.stop()
        st.write(f"Custom Formula: {custom_formula}")

reducer_choice = st.selectbox(
    "Select Reducer (e.g, mean , sum , median , min , max , count)",
    ['mean', 'sum', 'median', 'min', 'max', 'count'],
    index=0
)

start_date = st.date_input("Start Date", value=datetime(2024, 11, 1))
end_date = st.date_input("End Date", value=datetime(2024, 12, 1))
start_date_str = start_date.strftime('%Y-%m-%d')
end_date_str = end_date.strftime('%Y-%m-%d')

aggregation_period = st.selectbox(
    "Select Aggregation Period (e.g, Custom(Start Date to End Date) , Daily , Weekly , Monthly , Yearly)",
    ["Custom (Start Date to End Date)", "Daily", "Weekly", "Monthly", "Yearly"],
    index=0
)

shape_type = st.selectbox("Do you want to process 'Point' or 'Polygon' data?", ["Point", "Polygon"])

kernel_size = None
include_boundary = None

if shape_type == "Point":
    kernel_size = st.selectbox(
        "Select Calculation Area(e.g, Point , 3x3 Kernel , 5x5 Kernel)",
        ["Point", "3x3 Kernel", "5x5 Kernel"],
        index=0,
        help="Choose 'Point' for exact point calculation, or a kernel size for area averaging."
    )
elif shape_type == "Polygon":
    include_boundary = st.checkbox(
        "Include Boundary Pixels",
        value=True,
        help="Check to include pixels on the polygon boundary; uncheck to exclude them."
    )

file_upload = st.file_uploader(f"Upload your {shape_type} data (CSV, GeoJSON, KML)", type=["csv", "geojson", "kml"])
locations_df = pd.DataFrame()
original_lat_col = None
original_lon_col = None

if file_upload is not None:
    if shape_type.lower() == "point":
        if file_upload.name.endswith('.csv'):
            locations_df = pd.read_csv(file_upload)
            st.write("Preview of your uploaded data (first 5 rows):")
            st.dataframe(locations_df.head())
            all_columns = locations_df.columns.tolist()
            col1, col2 = st.columns(2)
            with col1:
                original_lat_col = st.selectbox(
                    "Select Latitude Column",
                    options=all_columns,
                    index=all_columns.index('latitude') if 'latitude' in all_columns else 0,
                    help="Select the column containing latitude values"
                )
            with col2:
                original_lon_col = st.selectbox(
                    "Select Longitude Column",
                    options=all_columns,
                    index=all_columns.index('longitude') if 'longitude' in all_columns else 0,
                    help="Select the column containing longitude values"
                )
            if not pd.api.types.is_numeric_dtype(locations_df[original_lat_col]) or not pd.api.types.is_numeric_dtype(locations_df[original_lon_col]):
                st.error("Selected Latitude and Longitude columns must contain numeric values")
                st.stop()
            locations_df = locations_df.rename(columns={
                original_lat_col: 'latitude',
                original_lon_col: 'longitude'
            })
        elif file_upload.name.endswith('.geojson'):
            locations_df = gpd.read_file(file_upload)
            if 'geometry' in locations_df.columns:
                locations_df['latitude'] = locations_df['geometry'].y
                locations_df['longitude'] = locations_df['geometry'].x
                original_lat_col = 'latitude'
                original_lon_col = 'longitude'
            else:
                st.error("GeoJSON file doesn't contain geometry column")
                st.stop()
        elif file_upload.name.endswith('.kml'):
            kml_string = file_upload.read().decode('utf-8')
            try:
                root = XET.fromstring(kml_string)
                ns = {'kml': 'http://www.opengis.net/kml/2.2'}
                points = []
                for placemark in root.findall('.//kml:Placemark', ns):
                    name = placemark.findtext('kml:name', default=f"Point_{len(points)}", namespaces=ns)
                    coords_elem = placemark.find('.//kml:Point/kml:coordinates', ns)
                    if coords_elem is not None:
                        coords_text = coords_elem.text.strip()
                        coords = [c.strip() for c in coords_text.split(',')]
                        if len(coords) >= 2:
                            lon, lat = float(coords[0]), float(coords[1])
                            points.append({'name': name, 'geometry': f"POINT ({lon} {lat})"})
                if not points:
                    st.error("No valid Point data found in the KML file.")
                else:
                    locations_df = gpd.GeoDataFrame(points, geometry=gpd.GeoSeries.from_wkt([p['geometry'] for p in points], crs="EPSG:4326"))
                    locations_df['latitude'] = locations_df['geometry'].y
                    locations_df['longitude'] = locations_df['geometry'].x
                    original_lat_col = 'latitude'
                    original_lon_col = 'longitude'
            except Exception as e:
                st.error(f"Error parsing KML file: {str(e)}")
        if not locations_df.empty and 'latitude' in locations_df.columns and 'longitude' in locations_df.columns:
            m = leafmap.Map(center=[locations_df['latitude'].mean(), locations_df['longitude'].mean()], zoom=10)
            for _, row in locations_df.iterrows():
                latitude = row['latitude']
                longitude = row['longitude']
                if pd.isna(latitude) or pd.isna(longitude):
                    continue
                m.add_marker(location=[latitude, longitude], popup=row.get('name', 'No Name'))

            st.write("Map of Uploaded Points:")
            m.to_streamlit()
    elif shape_type.lower() == "polygon":
        if file_upload.name.endswith('.csv'):
            st.error("CSV upload not supported for polygons. Please upload a GeoJSON or KML file.")
        elif file_upload.name.endswith('.geojson'):
            locations_df = gpd.read_file(file_upload)
            if 'geometry' not in locations_df.columns:
                st.error("GeoJSON file doesn't contain geometry column")
                st.stop()
        elif file_upload.name.endswith('.kml'):
            kml_string = file_upload.read().decode('utf-8')
            try:
                root = XET.fromstring(kml_string)
                ns = {'kml': 'http://www.opengis.net/kml/2.2'}
                polygons = []
                for placemark in root.findall('.//kml:Placemark', ns):
                    name = placemark.findtext('kml:name', default=f"Polygon_{len(polygons)}", namespaces=ns)
                    coords_elem = placemark.find('.//kml:Polygon//kml:coordinates', ns)
                    if coords_elem is not None:
                        coords_text = ' '.join(coords_elem.text.split())
                        coord_pairs = [pair.split(',')[:2] for pair in coords_text.split() if pair]
                        if len(coord_pairs) >= 4:
                            coords_str = " ".join([f"{float(lon)} {float(lat)}" for lon, lat in coord_pairs])
                            polygons.append({'name': name, 'geometry': f"POLYGON (({coords_str}))"})
                if not polygons:
                    st.error("No valid Polygon data found in the KML file.")
                else:
                    locations_df = gpd.GeoDataFrame(polygons, geometry=gpd.GeoSeries.from_wkt([p['geometry'] for p in polygons], crs="EPSG:4326"))
            except Exception as e:
                st.error(f"Error parsing KML file: {str(e)}")
        if not locations_df.empty and 'geometry' in locations_df.columns:
            centroid_lat = locations_df.geometry.centroid.y.mean()
            centroid_lon = locations_df.geometry.centroid.x.mean()
            m = leafmap.Map(center=[centroid_lat, centroid_lon], zoom=10)
            for _, row in locations_df.iterrows():
                polygon = row['geometry']
                if polygon.is_valid:
                    gdf = gpd.GeoDataFrame([row], geometry=[polygon], crs=locations_df.crs)
                    m.add_gdf(gdf=gdf, layer_name=row.get('name', 'Unnamed Polygon'))
            st.write("Map of Uploaded Polygons:")
            m.to_streamlit()

# Add this test code right before your processing button
if st.checkbox("Run Test Query"):
    try:
        test_collection = ee.ImageCollection(dataset_id) \
            .filterDate(start_date_str, end_date_str)
        
        if shape_type == "Point" and not locations_df.empty:
            first_point = ee.Geometry.Point([
                locations_df.iloc[0]['longitude'],
                locations_df.iloc[0]['latitude']
            ])
            test_collection = test_collection.filterBounds(first_point)
        elif shape_type == "Polygon" and not locations_df.empty:
            first_poly = convert_to_ee_geometry(locations_df.iloc[0]['geometry'])
            test_collection = test_collection.filterBounds(first_poly)
        
        image_count = test_collection.size().getInfo()
        st.write(f"Test query found {image_count} images matching your criteria")
        
        if image_count > 0:
            first_image = test_collection.first()
            st.write("First image properties:", first_image.getInfo())
    except Exception as e:
        st.error(f"Test query failed: {str(e)}")

if st.button(f"Calculate {custom_formula}"):
    if not locations_df.empty:
        with st.spinner("Processing Data..."):
            try:
                # TEMPORARY DEBUG - print all parameters
                st.write("DEBUG PARAMETERS:")
                st.write(f"Dataset: {dataset_id}")
                st.write(f"Bands: {selected_bands}")
                st.write(f"Formula: {custom_formula}")
                st.write(f"Cloud Masking: {apply_cloud_mask} (Band: {cloud_band}, Threshold: {cloud_threshold})")
                
                # Run with cloud masking disabled temporarily
                debug_results, debug_time, debug_count = process_aggregation(
                    locations_df.head(1),  # Just process first feature for debugging
                    start_date_str,
                    end_date_str,
                    dataset_id,
                    selected_bands,
                    reducer_choice,
                    shape_type,
                    aggregation_period,
                    custom_formula,
                    kernel_size,
                    include_boundary,
                    default_scale,
                    apply_cloud_mask=True,  # Disable cloud masking for debug
                    cloud_threshold=None,
                    cloud_band=None
                )
                
                # Display debug results in a table
                if debug_results:
                    debug_df = pd.DataFrame(debug_results)
                    st.write("DEBUG RESULTS TABLE:")
                    st.dataframe(debug_df)
                    
                    # Add download button for debug results
                    csv = debug_df.to_csv(index=False).encode('utf-8')
                    st.download_button(
                        label="Download Debug Results as CSV",
                        data=csv,
                        file_name=f"debug_results_{datetime.now().strftime('%Y%m%d_%H%M%S')}.csv",
                        mime='text/csv'
                    )
                
                st.write(f"DEBUG IMAGE COUNT: {debug_count}")
                
                if debug_count > 0:
                    # Now run full processing if debug worked
                    results, processing_time, total_images = process_aggregation(
                        locations_df,
                        start_date_str,
                        end_date_str,
                        dataset_id,
                        selected_bands,
                        reducer_choice,
                        shape_type,
                        aggregation_period,
                        custom_formula,
                        kernel_size,
                        include_boundary,
                        default_scale,
                        apply_cloud_mask,
                        cloud_threshold,
                        cloud_band
                    )
                    
                    if results:
                        result_df = pd.DataFrame(results)
                        st.write("FINAL RESULTS TABLE:")
                        st.dataframe(result_df)
                        
                        # Add download button for final results
                        csv = result_df.to_csv(index=False).encode('utf-8')
                        st.download_button(
                            label="Download Final Results as CSV",
                            data=csv,
                            file_name=f"final_results_{datetime.now().strftime('%Y%m%d_%H%M%S')}.csv",
                            mime='text/csv'
                        )
                        
                        st.success(f"Processed {total_images} images in {processing_time:.2f}s")
                    else:
                        st.warning("Main processing returned no results despite debug success")
                else:
                    st.error("Debug processing failed - check parameters above")
                    
            except Exception as e:
                st.error(f"Processing failed: {str(e)}")