"""
Google Earth Engine Sentinel Data Retrieval Module

This module provides functionality to retrieve and process Sentinel satellite data
from Google Earth Engine for water-related analysis and monitoring.

Features:
- User-configurable authentication and project setup
- Sentinel-2 surface reflectance data retrieval with custom masking
- Water body analysis using spectral indices
- Time series extraction for specific locations
- Anomaly detection capabilities
- Interactive mapping and visualization

Requirements:
- Google Earth Engine Python API (earthengine-api)
- geemap for visualization
- Valid GEE account and authentication
"""

import ee
import pandas as pd
import numpy as np
import geemap
import warnings
from datetime import datetime, timedelta
from typing import Dict, List, Tuple, Optional, Union
from tqdm import tqdm
import matplotlib.pyplot as plt
import seaborn as sns

warnings.filterwarnings('ignore')


class SentinelConfig:
    """
    Configuration class for Sentinel data retrieval.
    Users can customize these settings for their specific needs.
    """
    
    def __init__(self):
        # Authentication settings (user must configure)
        self.gee_project = None  # User must set their GEE project ID
        
        # Default masks to apply
        self.masks_to_apply = [
            "water",                      # Keep only SCL = 6 (water)
            "no_cloud_shadow",            # Remove cloud shadow (SCL = 3)
            "no_clouds",                  # Remove clouds (SCL = 8, 9, 10)
            "no_snow_ice",                # Remove snow/ice (SCL = 11)
            "no_saturated",               # Remove saturated/defective pixels (SCL = 1)
        ]
        
        # Processing parameters
        self.cloudy_pixel_percentage = 20    # Max allowed cloud percentage
        self.buffer_meters = 20              # Buffer size around point, in meters
        self.scale = 20                      # Processing scale in meters
        
        # Sentinel-2 bands to retrieve
        self.bands = ['B2', 'B3', 'B4', 'B5', 'B6', 'B7', 'B8', 'B8A', 'B9', 'B11', 'B12', 'SCL']
        
        # Visualization parameters
        self.ndwi_vis = {
            'min': 0,
            'max': 1,
            'palette': ['white', 'cyan', 'blue']
        }
    
    def set_project(self, project_id: str):
        """Set the Google Earth Engine project ID."""
        self.gee_project = project_id
    
    def set_masks(self, masks: List[str]):
        """Set the masks to apply during processing."""
        self.masks_to_apply = masks
    
    def set_buffer(self, buffer_meters: int):
        """Set the buffer size around points."""
        self.buffer_meters = buffer_meters
    
    def set_cloud_threshold(self, threshold: int):
        """Set the maximum cloud percentage."""
        self.cloudy_pixel_percentage = threshold


class SentinelDataRetriever:
    """
    A class to retrieve and process Sentinel satellite data from Google Earth Engine.
    """
    
    def __init__(self, config: Optional[SentinelConfig] = None):
        """
        Initialize the Sentinel Data Retriever.
        
        Parameters:
        -----------
        config : SentinelConfig, optional
            Configuration object with user settings
        """
        self.config = config if config else SentinelConfig()
        self.initialized = False
    
    def authenticate_and_initialize(self, project_id: str, force_auth: bool = False):
        """
        Authenticate with Google Earth Engine and initialize.
        
        Parameters:
        -----------
        project_id : str
            Your Google Earth Engine project ID
        force_auth : bool, default False
            Whether to force re-authentication
        
        Returns:
        --------
        bool : True if successful, False otherwise
        """
        try:
            if force_auth:
                print("🔐 Starting authentication process...")
                ee.Authenticate()
            
            print(f"🚀 Initializing Google Earth Engine with project: {project_id}")
            ee.Initialize(project=project_id)
            
            self.config.set_project(project_id)
            self.initialized = True
            
            print("✅ Google Earth Engine initialized successfully!")
            return True
            
        except Exception as e:
            print(f"❌ Failed to initialize Google Earth Engine: {e}")
            print("💡 Make sure you have:")
            print("   1. Valid GEE account access")
            print("   2. Correct project ID")
            print("   3. Run: earthengine authenticate (if needed)")
            return False
    
    def create_point_geometry(self, longitude: float, latitude: float) -> ee.Geometry.Point:
        """
        Create a point geometry for data extraction.
        
        Parameters:
        -----------
        longitude : float
            Longitude coordinate
        latitude : float
            Latitude coordinate
        
        Returns:
        --------
        ee.Geometry.Point : Earth Engine point geometry
        """
        if not self.initialized:
            raise RuntimeError("GEE not initialized. Call authenticate_and_initialize() first.")
        
        return ee.Geometry.Point([longitude, latitude])
    
    def get_available_dates(self, point: ee.Geometry.Point, start_date: str, end_date: str) -> List[str]:
        """
        Get available Sentinel-2 dates for a location.
        
        Parameters:
        -----------
        point : ee.Geometry.Point
            Location point
        start_date : str
            Start date in 'YYYY-MM-DD' format
        end_date : str
            End date in 'YYYY-MM-DD' format
        
        Returns:
        --------
        List[str] : List of available dates
        """
        buffered_point = point.buffer(self.config.buffer_meters)
        
        s2 = (ee.ImageCollection('COPERNICUS/S2_SR_HARMONIZED')
              .filterBounds(buffered_point)
              .filterDate(start_date, end_date))
        
        def get_date(image):
            return ee.Feature(None, {'date': image.date().format('YYYY-MM-dd HH:mm:ss')})
        
        dates_fc = ee.FeatureCollection(s2.map(get_date))
        date_list = dates_fc.aggregate_array('date').getInfo()
        
        print(f'Available Sentinel-2 dates at this site:')
        for d in date_list:
            print(d)
        print(f"Total: {len(date_list)} images from {start_date} to {end_date}")
        
        return date_list
    
    def dynamic_scl_mask(self, image):
        """
        Apply dynamic Scene Classification Layer (SCL) masks based on user configuration.
        
        Parameters:
        -----------
        image : ee.Image
            Sentinel-2 image
        
        Returns:
        --------
        ee.Image : Masked image
        """
        scl = image.select('SCL')
        mask = ee.Image.constant(1)
        
        if "water" in self.config.masks_to_apply:
            mask = mask.And(scl.eq(6))
        if "no_cloud_shadow" in self.config.masks_to_apply:
            mask = mask.And(scl.neq(3))
        if "no_clouds" in self.config.masks_to_apply:
            cloud_mask = scl.neq(8).And(scl.neq(9)).And(scl.neq(10))
            mask = mask.And(cloud_mask)
        if "no_snow_ice" in self.config.masks_to_apply:
            mask = mask.And(scl.neq(11))
        if "no_saturated" in self.config.masks_to_apply:
            mask = mask.And(scl.neq(1))
        
        return image.updateMask(mask)
    
    def extract_time_series(self, 
                           longitude: float, 
                           latitude: float, 
                           start_date: str, 
                           end_date: str) -> pd.DataFrame:
        """
        Extract time series data for a specific location.
        
        Parameters:
        -----------
        longitude : float
            Longitude coordinate
        latitude : float
            Latitude coordinate
        start_date : str
            Start date in 'YYYY-MM-DD' format
        end_date : str
            End date in 'YYYY-MM-DD' format
        
        Returns:
        --------
        pd.DataFrame : Time series data with spectral bands
        """
        if not self.initialized:
            raise RuntimeError("GEE not initialized. Call authenticate_and_initialize() first.")
        
        print(f"🛰️  Extracting Sentinel-2 time series...")
        print(f"📍 Location: {latitude:.6f}°N, {longitude:.6f}°W")
        print(f"📅 Period: {start_date} to {end_date}")
        print(f"🎭 Masks: {self.config.masks_to_apply}")
        
        point = self.create_point_geometry(longitude, latitude)
        buffered_point = point.buffer(self.config.buffer_meters)
        
        # Load and filter image collection
        s2 = (ee.ImageCollection('COPERNICUS/S2_SR_HARMONIZED')
              .filterBounds(buffered_point)
              .filterDate(start_date, end_date)
              .filter(ee.Filter.lt('CLOUDY_PIXEL_PERCENTAGE', self.config.cloudy_pixel_percentage))
              .select(self.config.bands)
              .map(self.dynamic_scl_mask))
        
        def extract_features(image):
            date = image.date().format('YYYY-MM-dd HH:mm:ss')
            values = image.reduceRegion(
                reducer=ee.Reducer.mean(),
                geometry=buffered_point,
                scale=self.config.scale,
                maxPixels=1e8
            )
            return ee.Feature(None, values.set('date', date))
        
        print("🔄 Processing images and extracting features...")
        features = s2.map(extract_features)
        fc = ee.FeatureCollection(features).filter(ee.Filter.notNull(['B2']))
        
        # Convert to pandas DataFrame
        data = fc.getInfo()
        print(f"Number of features: {len(data['features'])}")
        
        rows = [f['properties'] for f in data['features']]
        df = pd.DataFrame(rows)
        
        if len(df) > 0:
            df['date'] = pd.to_datetime(df['date'])
            df = df.sort_values('date').reset_index(drop=True)
            
            print(f"✅ Retrieved {len(df)} records")
            print(f"📊 Date range: {df['date'].min()} to {df['date'].max()}")
        else:
            print("⚠️  No data retrieved. Try adjusting date range or masks.")
        
        return df
    
    def create_interactive_map(self, 
                              longitude: float, 
                              latitude: float, 
                              start_date: str, 
                              end_date: str,
                              zoom: int = 15) -> geemap.Map:
        """
        Create an interactive map showing the location and NDWI visualization.
        
        Parameters:
        -----------
        longitude : float
            Longitude coordinate
        latitude : float
            Latitude coordinate
        start_date : str
            Start date for image composite
        end_date : str
            End date for image composite
        zoom : int, default 15
            Map zoom level
        
        Returns:
        --------
        geemap.Map : Interactive map object
        """
        if not self.initialized:
            raise RuntimeError("GEE not initialized. Call authenticate_and_initialize() first.")
        
        print("🗺️  Creating interactive map...")
        
        point = self.create_point_geometry(longitude, latitude)
        buffer = point.buffer(self.config.buffer_meters)
        
        # Create median composite
        image = (ee.ImageCollection("COPERNICUS/S2_SR_HARMONIZED")
                .filterBounds(buffer)
                .filterDate(start_date, end_date)
                .filter(ee.Filter.lt("CLOUDY_PIXEL_PERCENTAGE", self.config.cloudy_pixel_percentage))
                .median())
        
        # Calculate NDWI
        ndwi = image.normalizedDifference(["B3", "B8"]).rename("NDWI")
        
        # Create map
        Map = geemap.Map()
        Map.centerObject(buffer, zoom=zoom)
        Map.addLayer(ndwi, self.config.ndwi_vis, "NDWI (Water)")
        Map.addLayer(point, {'color': 'yellow'}, 'Point')
        Map.addLayer(buffer, {'color': 'red'}, f'Buffer ({self.config.buffer_meters}m)')
        
        print(f"✅ Map created with NDWI visualization")
        print(f"🎯 Centered on: {latitude:.6f}°N, {longitude:.6f}°W")
        
        return Map
    
    def calculate_water_indices(self, df: pd.DataFrame) -> pd.DataFrame:
        """
        Calculate water-related spectral indices from the DataFrame.
        
        Parameters:
        -----------
        df : pd.DataFrame
            DataFrame with Sentinel-2 band data
        
        Returns:
        --------
        pd.DataFrame : DataFrame with added indices
        """
        if len(df) == 0:
            return df
        
        df_copy = df.copy()
        
        # NDWI (Normalized Difference Water Index)
        if 'B3' in df_copy.columns and 'B8' in df_copy.columns:
            df_copy['NDWI'] = (df_copy['B3'] - df_copy['B8']) / (df_copy['B3'] + df_copy['B8'])
        
        # MNDWI (Modified NDWI)
        if 'B3' in df_copy.columns and 'B11' in df_copy.columns:
            df_copy['MNDWI'] = (df_copy['B3'] - df_copy['B11']) / (df_copy['B3'] + df_copy['B11'])
        
        # NDVI (Normalized Difference Vegetation Index)
        if 'B4' in df_copy.columns and 'B8' in df_copy.columns:
            df_copy['NDVI'] = (df_copy['B8'] - df_copy['B4']) / (df_copy['B8'] + df_copy['B4'])
        
        # Water turbidity proxy
        if 'B4' in df_copy.columns and 'B2' in df_copy.columns:
            df_copy['TURBIDITY'] = df_copy['B4'] / df_copy['B2']
        
        print("✅ Water indices calculated: NDWI, MNDWI, NDVI, TURBIDITY")
        return df_copy
    
    def plot_time_series(self, df: pd.DataFrame, indices: List[str] = None, save_path: str = None):
        """
        Plot time series of spectral indices.
        
        Parameters:
        -----------
        df : pd.DataFrame
            DataFrame with time series data
        indices : List[str], optional
            List of indices to plot (default: ['NDWI', 'MNDWI', 'NDVI'])
        save_path : str, optional
            Path to save the plot
        """
        if len(df) == 0:
            print("⚠️  No data to plot")
            return
        
        if indices is None:
            indices = ['NDWI', 'MNDWI', 'NDVI']
        
        # Filter available indices
        available_indices = [idx for idx in indices if idx in df.columns]
        
        if not available_indices:
            print("⚠️  No indices available to plot")
            return
        
        fig, axes = plt.subplots(len(available_indices), 1, figsize=(12, 4*len(available_indices)))
        if len(available_indices) == 1:
            axes = [axes]
        
        for i, idx in enumerate(available_indices):
            axes[i].plot(df['date'], df[idx], 'o-', markersize=4, linewidth=1.5)
            axes[i].set_ylabel(idx)
            axes[i].set_title(f'{idx} Time Series')
            axes[i].grid(True, alpha=0.3)
            axes[i].tick_params(axis='x', rotation=45)
        
        plt.tight_layout()
        
        if save_path:
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
            print(f"📊 Plot saved: {save_path}")
        
        plt.show()


# Convenience functions for easy access
def setup_gee_authentication(project_id: str, force_auth: bool = False) -> SentinelDataRetriever:
    """
    Set up Google Earth Engine authentication and return a configured retriever.
    
    Parameters:
    -----------
    project_id : str
        Your Google Earth Engine project ID
    force_auth : bool, default False
        Whether to force re-authentication
    
    Returns:
    --------
    SentinelDataRetriever : Configured and initialized retriever
    
    Example:
    --------
    >>> retriever = setup_gee_authentication('your-project-id')
    """
    config = SentinelConfig()
    retriever = SentinelDataRetriever(config)
    
    if retriever.authenticate_and_initialize(project_id, force_auth):
        return retriever
    else:
        raise RuntimeError("Failed to initialize Google Earth Engine")


def get_sentinel_time_series(project_id: str,
                           longitude: float,
                           latitude: float,
                           start_date: str,
                           end_date: str,
                           masks: List[str] = None,
                           buffer_meters: int = 20) -> pd.DataFrame:
    """
    Quick function to get Sentinel-2 time series data.
    
    Parameters:
    -----------
    project_id : str
        Your Google Earth Engine project ID
    longitude : float
        Longitude coordinate
    latitude : float
        Latitude coordinate
    start_date : str
        Start date in 'YYYY-MM-DD' format
    end_date : str
        End date in 'YYYY-MM-DD' format
    masks : List[str], optional
        Custom masks to apply
    buffer_meters : int, default 20
        Buffer size around point
    
    Returns:
    --------
    pd.DataFrame : Time series data
    
    Example:
    --------
    >>> df = get_sentinel_time_series(
    ...     'your-project-id',
    ...     -95.0644278, 29.7785861,
    ...     '2023-01-01', '2023-12-31'
    ... )
    """
    retriever = setup_gee_authentication(project_id)
    
    if masks:
        retriever.config.set_masks(masks)
    if buffer_meters != 20:
        retriever.config.set_buffer(buffer_meters)
    
    df = retriever.extract_time_series(longitude, latitude, start_date, end_date)
    return retriever.calculate_water_indices(df)
