"""
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
from pystac_client import Client
import planetary_computer
import odc.stac
from datetime import datetime


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


# 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("=" * 70)
print()