"""
Load Sentinel data without ODC Database
Sử dụng Microsoft Planetary Computer STAC API

Usage:
    from load_data_no_odc import load_sentinel2_stac, load_sentinel1_stac
"""

import xarray as xr
import numpy as np
import pandas as pd
import geopandas as gpd
from pystac_client import Client
import planetary_computer
import odc.stac
from datetime import datetime
from sklearn.model_selection import train_test_split


def load_sentinel2_stac(
    bbox,  # (lon_min, lat_min, lon_max, lat_max)
    date_range,  # ("2022-09-01", "2023-10-01")
    bands=None,
    resolution=10,
    chunks={'time': 1, 'x': 2048, 'y': 2048}
):
    """
    Load Sentinel-2 L2A data from Microsoft Planetary Computer
    
    Args:
        bbox: Bounding box (lon_min, lat_min, lon_max, lat_max)
        date_range: Tuple of start and end dates
        bands: List of bands to load (default: ['red', 'nir', 'scl'])
        resolution: Spatial resolution in meters
        chunks: Dask chunk sizes
    
    Returns:
        xarray.Dataset with Sentinel-2 data
    """
    if bands is None:
        bands = ['red', 'nir', 'blue', 'green', 'nir08', 'swir16', 'swir22', 'SCL']
    
    print(f"🔍 Searching Sentinel-2 data...")
    print(f"   Bbox: {bbox}")
    print(f"   Date: {date_range[0]} to {date_range[1]}")
    print(f"   Bands: {bands}")
    
    # Connect to Planetary Computer STAC
    catalog = Client.open(
        "https://planetarycomputer.microsoft.com/api/stac/v1",
        modifier=planetary_computer.sign_inplace,
    )
    
    # Search for Sentinel-2 items
    search = catalog.search(
        collections=["sentinel-2-l2a"],
        bbox=bbox,
        datetime=f"{date_range[0]}/{date_range[1]}",
    )
    
    items = list(search.items())
    print(f"✅ Found {len(items)} Sentinel-2 scenes")
    
    if len(items) == 0:
        print("⚠ No data found for the given parameters")
        return None
    
    # Load data using odc-stac
    print("📥 Loading data...")
    try:
        data = odc.stac.load(
            items,
            bands=bands,
            bbox=bbox,
            resolution=resolution,
            chunks=chunks,
            groupby="solar_day",
        )
        
        print(f"✅ Loaded Sentinel-2 data: {data.dims}")
        print(f"   Shape: {dict(data.dims)}")
        print(f"   Bands: {list(data.data_vars)}")
        
        return data
        
    except Exception as e:
        print(f"❌ Error loading data: {e}")
        import traceback
        traceback.print_exc()
        return None


def load_sentinel1_stac(
    bbox,  # (lon_min, lat_min, lon_max, lat_max)
    date_range,  # ("2022-09-01", "2023-10-01")
    bands=None,
    resolution=10,
    chunks={'time': 1, 'x': 2048, 'y': 2048}
):
    """
    Load Sentinel-1 RTC data from Microsoft Planetary Computer
    
    Args:
        bbox: Bounding box (lon_min, lat_min, lon_max, lat_max)
        date_range: Tuple of start and end dates
        bands: List of bands to load (default: ['vv', 'vh'])
        resolution: Spatial resolution in meters
        chunks: Dask chunk sizes
    
    Returns:
        xarray.Dataset with Sentinel-1 data
    """
    if bands is None:
        bands = ['vv', 'vh']
    
    print(f"🔍 Searching Sentinel-1 data...")
    print(f"   Bbox: {bbox}")
    print(f"   Date: {date_range[0]} to {date_range[1]}")
    print(f"   Bands: {bands}")
    
    # Connect to Planetary Computer STAC
    catalog = Client.open(
        "https://planetarycomputer.microsoft.com/api/stac/v1",
        modifier=planetary_computer.sign_inplace,
    )
    
    # Search for Sentinel-1 RTC items
    search = catalog.search(
        collections=["sentinel-1-rtc"],
        bbox=bbox,
        datetime=f"{date_range[0]}/{date_range[1]}",
    )
    
    items = list(search.items())
    print(f"✅ Found {len(items)} Sentinel-1 scenes")
    
    if len(items) == 0:
        print("⚠ No data found for the given parameters")
        return None
    
    # Load data using odc-stac
    print("📥 Loading data...")
    try:
        data = odc.stac.load(
            items,
            bands=bands,
            bbox=bbox,
            resolution=resolution,
            chunks=chunks,
            groupby="solar_day",
        )
        
        # Rename bands to uppercase for consistency
        if 'vv' in data.data_vars:
            data = data.rename({'vv': 'VV', 'vh': 'VH'})
        
        print(f"✅ Loaded Sentinel-1 data: {data.dims}")
        print(f"   Shape: {dict(data.dims)}")
        print(f"   Bands: {list(data.data_vars)}")
        
        return data
        
    except Exception as e:
        print(f"❌ Error loading data: {e}")
        import traceback
        traceback.print_exc()
        return None


def mask_clean_s2(data, scl_band='SCL'):
    """
    Apply cloud mask to Sentinel-2 data using SCL band
    Good pixels: 2=dark, 4=vegetation, 5=not-vegetated, 6=water
    """
    print("🧹 Applying cloud mask...")
    
    if scl_band not in data.data_vars:
        print(f"⚠ Warning: {scl_band} band not found, skipping cloud mask")
        return data
    
    # Good pixel values
    good_pixels = [2, 4, 5, 6]
    
    # Create mask
    mask = xr.zeros_like(data[scl_band], dtype=bool)
    for pixel_val in good_pixels:
        mask = mask | (data[scl_band] == pixel_val)
    
    # Apply mask to all bands except SCL
    data_vars = [v for v in data.data_vars if v != scl_band]
    result = data[data_vars].where(mask)
    
    print(f"✅ Cloud mask applied")
    return result


def calculate_ndvi(data, nir_band='nir08', red_band='red'):
    """
    Calculate NDVI from Sentinel-2 data
    """
    print(f"📊 Calculating NDVI using {nir_band} and {red_band}...")
    
    if nir_band not in data.data_vars or red_band not in data.data_vars:
        print(f"⚠ Warning: Required bands not found")
        print(f"   Available bands: {list(data.data_vars)}")
        return data
    
    nir = data[nir_band]
    red = data[red_band]
    
    ndvi = (nir - red) / (nir + red)
    ndvi = ndvi.rename('NDVI')
    
    # Add NDVI to dataset
    result = xr.merge([data, ndvi])
    
    print(f"✅ NDVI calculated")
    return result


def fill_nan_temporal(data, dim='time'):
    """
    Fill NaN values using forward/backward fill along time dimension
    """
    print(f"🔧 Filling NaN values using temporal interpolation...")
    
    filled = data.bfill(dim=dim).ffill(dim=dim)
    
    print(f"✅ NaN values filled")
    return filled


def resample_monthly(data, method='mean'):
    """
    Resample data to monthly frequency
    """
    print(f"📅 Resampling to monthly using {method}...")
    
    if method == 'mean':
        result = data.resample(time='1MS').mean()
    elif method == 'median':
        result = data.resample(time='1MS').median()
    else:
        raise ValueError(f"Unknown method: {method}")
    
    print(f"✅ Resampled to monthly: {result.dims}")
    return result


# ══════════════════════════════════════════════════════════════════════════════
# CONVENIENCE WRAPPER
# ══════════════════════════════════════════════════════════════════════════════

def load_and_process_s2(bbox, date_range, apply_cloud_mask=True, calculate_indices=True):
    """
    Load and process Sentinel-2 data in one step
    """
    # Load data
    data = load_sentinel2_stac(bbox, date_range)
    
    if data is None:
        return None
    
    # Apply cloud mask
    if apply_cloud_mask:
        data = mask_clean_s2(data)
    
    # Calculate NDVI
    if calculate_indices:
        data = calculate_ndvi(data)
    
    # Fill NaN
    data = fill_nan_temporal(data)
    
    # Resample to monthly
    data_monthly = resample_monthly(data)
    
    return data_monthly


def load_and_process_s1(bbox, date_range):
    """
    Load and process Sentinel-1 data in one step
    """
    # Load data
    data = load_sentinel1_stac(bbox, date_range)
    
    if data is None:
        return None
    
    # Resample to monthly
    data_monthly = resample_monthly(data)
    
    return data_monthly


# ══════════════════════════════════════════════════════════════════════════════
# TRAINING DATA FUNCTIONS
# ══════════════════════════════════════════════════════════════════════════════

def load_train_data(train_path, label_mapping=None):
    """
    Load training data from shapefile or GeoJSON
    
    Args:
        train_path: Path to training data file (.shp, .geojson, etc.)
        label_mapping: Optional dict to map labels to numeric IDs
    
    Returns:
        GeoDataFrame with training points
    """
    print(f"📂 Loading training data from: {train_path}")
    
    try:
        train = gpd.read_file(train_path)
        print(f"✅ Loaded {len(train)} training points")
        print(f"   Columns: {list(train.columns)}")
        
        if label_mapping is not None and 'label' in train.columns:
            train['label_id'] = train['label'].map(label_mapping).astype(int)
            print(f"   Labels mapped: {sorted(train['label_id'].unique())}")
        
        return train
        
    except Exception as e:
        print(f"❌ Error loading training data: {e}")
        return None


def extract_features_at_points(train_data, data_s2, data_s1):
    """
    Extract Sentinel-1 and Sentinel-2 features at training point locations
    
    Args:
        train_data: GeoDataFrame with training points
        data_s2: xarray Dataset with Sentinel-2 data 
        data_s1: xarray Dataset with Sentinel-1 data
    
    Returns:
        X (features), y (labels)
    """
    print(f"🔧 Extracting features from satellite data at {len(train_data)} points...")
    
    X_list = []
    y_list = []
    
    for idx, point in train_data.iterrows():
        try:
            lon, lat = point.geometry.x, point.geometry.y
            
            # Extract S2 data at point location
            s2_point = data_s2.sel(x=lon, y=lat, method='nearest')
            s2_values = s2_point.to_array().values.flatten()
            
            # Extract S1 data at point location
            s1_point = data_s1.sel(x=lon, y=lat, method='nearest')
            s1_values = s1_point.to_array().values.flatten()
            
            # Combine features
            features = np.concatenate([s2_values, s1_values])
            
            # Skip if contains NaN
            if not np.isnan(features).any():
                X_list.append(features)
                y_list.append(point['label_id'])
            else:
                print(f"  ⚠ Skip point {idx}: contains NaN")
                
        except Exception as e:
            print(f"  ⚠ Skip point {idx}: {e}")
            continue
    
    X = np.array(X_list)
    y = np.array(y_list)
    
    print(f"✅ Extracted features from {len(X)} points")
    print(f"   Feature dimension: {X.shape[1]}")
    print(f"   Classes: {sorted(set(y.tolist()))}")
    
    return X, y


def split_train_data(X, y, test_size=0.2, val_size=0.1, random_state=42):
    """
    Split data into train/val/test sets
    
    Args:
        X: Features array
        y: Labels array
        test_size: Proportion for test set
        val_size: Proportion for validation set (from train+val)
        random_state: Random seed
    
    Returns:
        X_train, X_val, X_test, y_train, y_val, y_test
    """
    print(f"📊 Splitting data...")
    
    # First split: train+val vs test
    X_temp, X_test, y_temp, y_test = train_test_split(
        X, y, test_size=test_size, random_state=random_state, stratify=y
    )
    
    # Second split: train vs val
    val_ratio = val_size / (1 - test_size)
    X_train, X_val, y_train, y_val = train_test_split(
        X_temp, y_temp, test_size=val_ratio, random_state=random_state, stratify=y_temp
    )
    
    print(f"✅ Data split complete:")
    print(f"   Train: {len(X_train)} samples ({len(X_train)/len(X)*100:.1f}%)")
    print(f"   Val:   {len(X_val)} samples ({len(X_val)/len(X)*100:.1f}%)")
    print(f"   Test:  {len(X_test)} samples ({len(X_test)/len(X)*100:.1f}%)")
    
    return X_train, X_val, X_test, y_train, y_val, y_test


# Print module info
print("=" * 70)
print("📦 Data Loading Module (No ODC Database Required)")
print("=" * 70)
print("\n💡 Usage:")
print("   from load_data_no_odc import load_and_process_s2, load_and_process_s1")
print("\n   bbox = (lon_min, lat_min, lon_max, lat_max)")
print("   date_range = ('2022-09-01', '2023-10-01')")
print("\n   data_s2 = load_and_process_s2(bbox, date_range)")
print("   data_s1 = load_and_process_s1(bbox, date_range)")
print("\n   # Load training data")
print("   train = load_train_data('train/data.shp', label_mapping)")
print("   X, y = extract_features_at_points(train, data_s2, data_s1)")
print("   X_train, X_val, X_test, y_train, y_val, y_test = split_train_data(X, y)")
print("=" * 70)
print()