"""
ODC Import Module with Cognito Authentication
Module ODC tích hợp xác thực Cognito

Usage:
    import new_import_ODC_cognito
    from new_import_ODC_cognito import *
    
    # Setup Cognito authentication
    setup_cognito_auth('train_files/crediential.txt')
    
    # Then use datacube normally
"""

import matplotlib.pyplot as plt

# Common imports and settings
import os, sys
os.environ['USE_PYGEOS'] = '0'
from IPython.display import Markdown
import pandas as pd
pd.set_option("display.max_rows", None)
import xarray as xr

# Datacube
import datacube
from datacube.utils.rio import configure_s3_access
from datacube.utils import masking
from datacube.utils.cog import write_cog

# DEA Tools
from dea_tools.plotting import display_map, rgb
from dea_tools.datahandling import mostcommon_crs

# EASI defaults - Update path to local installation
easi_tools_path = '/media/x79/2A7D-FAA0/remote-sensing'
if easi_tools_path not in sys.path: 
    sys.path.insert(0, easi_tools_path)

# Try to import EASI tools
EASI_AVAILABLE = False
notebook_utils = None
load_s2l2a_with_offset = None

try:
    # Check if dask_gateway is available first
    try:
        import dask_gateway
        dask_gateway_available = True
    except ImportError:
        dask_gateway_available = False
        print("⚠ dask_gateway not installed - will use LocalCluster instead")
    
    # Import EASI tools
    if dask_gateway_available:
        from easi_tools import notebook_utils
        from easi_tools.load_s2l2a import load_s2l2a_with_offset
        EASI_AVAILABLE = True
        print("✅ EASI tools loaded successfully (with Gateway support)")
    else:
        # Import what we can without dask_gateway
        print("⚠ Loading EASI tools without Gateway support...")
        # Don't import notebook_utils if it requires dask_gateway
        from easi_tools.load_s2l2a import load_s2l2a_with_offset
        print("✅ EASI load_s2l2a loaded (without notebook_utils)")
        
except ImportError as e:
    print(f"⚠ EASI tools not available: {e}")
    print("⚠ Using standard datacube functions")
    EASI_AVAILABLE = False
    
# Create fallback notebook_utils if not available
if notebook_utils is None:
    class FallbackNotebookUtils:
        """Fallback implementation when EASI tools not available"""
        
        @staticmethod
        def initialize_dask(use_gateway=False, workers=(1, 10)):
            """Initialize Dask cluster"""
            from dask.distributed import Client, LocalCluster
            
            if use_gateway:
                print("⚠ Dask Gateway not available, using LocalCluster")
            
            n_workers = workers[0] if isinstance(workers, tuple) else workers
            cluster = LocalCluster(
                n_workers=n_workers, 
                threads_per_worker=1,
                memory_limit='4GB'
            )
            client = Client(cluster)
            
            print(f"✅ Dask LocalCluster started")
            print(f"   Workers: {n_workers}")
            print(f"   Dashboard: {client.dashboard_link}")
            
            return cluster, client
        
        @staticmethod
        def mostcommon_crs(dc, query):
            """Get most common CRS - fallback to Vietnam default"""
            try:
                # Try to get CRS from datacube
                datasets = list(dc.find_datasets(**query))
                if datasets and len(datasets) > 0:
                    return str(datasets[0].crs)
            except Exception as e:
                print(f"⚠ Could not determine CRS from datacube: {e}")
            
            # Default CRS for Vietnam
            print("  Using default CRS: EPSG:32648 (Vietnam)")
            return 'EPSG:32648'
    
    notebook_utils = FallbackNotebookUtils()
    print("✅ Fallback notebook_utils created")
    
from dask.distributed import progress

# Data tools
import numpy as np
from datetime import datetime

# ODC algo
from odc.algo import enum_to_bool
from odc.algo import xr_reproject
from datacube.utils.geometry import GeoBox, box

# Holoviews, Datashader and Bokeh
import hvplot.pandas
import hvplot.xarray
import holoviews as hv
import panel as pn
import colorcet as cc
import cartopy.crs as ccrs
from datashader import reductions
from holoviews import opts
hv.extension('bokeh', logo=False)

# ML and Geo tools
from deafrica_tools.bandindices import calculate_indices
from sklearn.ensemble import RandomForestClassifier, RandomForestRegressor
from sklearn.model_selection import train_test_split, GridSearchCV
from sklearn.metrics import accuracy_score, classification_report, mean_squared_error, r2_score
from sklearn.preprocessing import LabelEncoder, StandardScaler, PolynomialFeatures
from sklearn.pipeline import Pipeline
from sklearn.impute import SimpleImputer
from sklearn.linear_model import LinearRegression
from shapely.geometry import Point, Polygon
import geopandas as gpd
from pyproj import CRS
from matplotlib.colors import ListedColormap
from bokeh.models.tickers import FixedTicker
from rioxarray.merge import merge_arrays
import rasterio
import rioxarray
import joblib

# Import utils
try:
    from utils import load_data_geo
except ImportError:
    print("⚠ utils.py not found, load_data_geo() may not work")

# ══════════════════════════════════════════════════════════════════════════════
# COGNITO AUTHENTICATION
# ══════════════════════════════════════════════════════════════════════════════

# Global authenticator instance
_cognito_auth = None

def setup_cognito_auth(credential_file='train_files/crediential.txt', region='ap-southeast-1'):
    """
    Setup Cognito authentication for S3/ODC access
    Thiết lập xác thực Cognito cho truy cập S3/ODC
    
    Args:
        credential_file: Path to credential file containing AWS + Cognito tokens
        region: AWS region
    
    Returns:
        CognitoAuthenticator instance
    """
    global _cognito_auth
    
    try:
        # Import cognito_auth module
        sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
        from cognito_auth import CognitoAuthenticator
        
        print("🔐 Setting up Cognito authentication...")
        
        # Initialize authenticator
        _cognito_auth = CognitoAuthenticator(region=region)
        
        # Load tokens and credentials
        if not _cognito_auth.load_tokens_from_file(credential_file):
            print("✗ Failed to load Cognito tokens")
            return None
        
        # Display token info
        print("\n📋 Token Information:")
        decoded_id, _ = _cognito_auth.print_token_info()
        
        # Get AWS credentials
        if not _cognito_auth.get_credentials_from_cognito():
            print("✗ Failed to get AWS credentials")
            return None
        
        # Set environment credentials
        if not _cognito_auth.set_environment_credentials():
            print("✗ Failed to set environment credentials")
            return None
        
        # Configure S3 access for datacube
        print("\n🌐 Configuring datacube S3 access...")
        configure_s3_access(
            aws_unsigned=False,  # Using credentials
            region_name=region,
            cloud_defaults=True
        )
        
        print("\n✅ Cognito authentication setup complete!")
        print("✅ Ready to use datacube with S3 access\n")
        
        return _cognito_auth
        
    except ImportError as e:
        print(f"✗ Error: cognito_auth module not found: {e}")
        print("  Make sure cognito_auth.py is in the same directory")
        return None
    except Exception as e:
        print(f"✗ Error setting up Cognito auth: {e}")
        import traceback
        traceback.print_exc()
        return None


def get_cognito_auth():
    """
    Get the current Cognito authenticator instance
    Lấy instance Cognito authenticator hiện tại
    """
    return _cognito_auth


# ══════════════════════════════════════════════════════════════════════════════
# DATA LOADING FUNCTIONS
# ══════════════════════════════════════════════════════════════════════════════

def load_data(dc, date_range, longtitude_range, latitude_range, measurements=None):
    """
    Load Sentinel-2 L2A data from datacube
    """
    product = 's2_l2a'
    query = {
        'product': product,
        'x': longtitude_range,
        'y': latitude_range,
        'time': date_range,
    }
    
    if EASI_AVAILABLE:
        native_crs = notebook_utils.mostcommon_crs(dc, query)
    else:
        native_crs = 'EPSG:32648'  # Default for Vietnam
    
    print(f'Most common native CRS: {native_crs}')
    
    if measurements is None:
        measurements = ['red', 'nir', 'scl']

    load_params = {
        'measurements': measurements,
        'output_crs': native_crs,
        'resolution': (-10, 10),
        'group_by': 'solar_day',
        'dask_chunks': {'x': 2048, 'y': 2048},
    }
    
    if EASI_AVAILABLE:
        data = load_s2l2a_with_offset(dc, query | load_params)
    else:
        data = dc.load(**{**query, **load_params})
    
    return data


def load_data_sen1(dc, date_range, longtitude_range, latitude_range):
    """
    Load Sentinel-1 SAR data (VV, VH bands)
    """
    product = 's1_rtc'
    query = {
        'product': product,
        'x': longtitude_range,
        'y': latitude_range,
        'time': date_range,
        'measurements': ['VV', 'VH'],
        'output_crs': 'EPSG:32648',
        'resolution': (-10, 10),
        'group_by': 'solar_day',
        'dask_chunks': {'x': 2048, 'y': 2048},
    }
    
    data = dc.load(**query)
    return data


def mask_clean(data):
    """
    Apply cloud mask to Sentinel-2 data using SCL band
    """
    flag_name = 'scl'
    flag_desc = masking.describe_variable_flags(data[flag_name])
    display(flag_desc)
    display(flag_desc.loc['qa'].values[1])
    
    # Good pixel flags: 2=dark, 4=vegetation, 5=not-vegetated, 6=water
    flags_def = flag_desc.loc['qa'].values[1]
    good_pixel_flags = [flags_def[str(i)] for i in [2, 4, 5, 6]]
    
    good_pixel_mask = enum_to_bool(data[flag_name], good_pixel_flags)
    data_layer_names = [x for x in data.data_vars if x != 'scl']
    
    result = data[data_layer_names].where(good_pixel_mask).persist()
    return result


def fill_nan(ndvi, time_split=None):
    """
    Fill NaN values in NDVI using forward/backward fill
    """
    if time_split is None:
        # Simple fill without time splits
        fill_m = ndvi.bfill(dim='time').ffill(dim='time')
        return fill_m
    
    # Fill with time splits
    rs = []
    for times in time_split:
        tmp = ndvi.sel(time=times)
        fill_ds = tmp.bfill(dim='time').ffill(dim='time')
        rs.append(fill_ds)
    
    merged_ndvi = xr.concat(rs, dim="time")
    fill_m = merged_ndvi.bfill(dim="time").ffill(dim="time")
    return fill_m


def calculate_average(data, variables, resample='1MS'):
    """
    Calculate temporal average and resample
    """
    result = data[variables].resample(time=resample).mean()
    return result


def load_train_data(train_path=None, label_mapping=None):
    """
    Load training data from shapefile or GeoJSON
    """
    if train_path is None:
        print("⚠ No train_path provided")
        return None
    
    train = load_data_geo(train_path)
    
    if label_mapping is not None:
        # Apply label mapping
        train['label_id'] = train['label'].map(label_mapping).astype(int)
    
    return train


def get_data_sen1_and_sen2(train_data, data_sen2, data_sen1):
    """
    Extract Sentinel-1 and Sentinel-2 data for training points
    """
    X_list = []
    y_list = []
    
    for idx, point in train_data.iterrows():
        try:
            lon, lat = point.geometry.x, point.geometry.y
            
            # Extract S2 data
            s2_values = data_sen2.sel(x=lon, y=lat, method='nearest').values.flatten()
            
            # Extract S1 data
            s1_values = data_sen1.sel(x=lon, y=lat, method='nearest').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'])
                
        except Exception as e:
            print(f"⚠ Skip point {idx}: {e}")
            continue
    
    X = np.array(X_list)
    y = np.array(y_list)
    
    print(f"✅ Extracted {len(X)} training samples")
    print(f"   Features: {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
    """
    # 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:")
    print(f"   Train: {len(X_train)} samples")
    print(f"   Val:   {len(X_val)} samples")
    print(f"   Test:  {len(X_test)} samples")
    
    return X_train, X_val, X_test, y_train, y_val, y_test


# ══════════════════════════════════════════════════════════════════════════════
# HELPER FUNCTIONS
# ══════════════════════════════════════════════════════════════════════════════

def load_sen1(name_vh, name_vv):
    """Load Sentinel-1 from local files"""
    dsvv = rioxarray.open_rasterio(name_vv)
    dsvh = rioxarray.open_rasterio(name_vh)
    return dsvh, dsvv


def print_auth_status():
    """Print current authentication status"""
    global _cognito_auth
    
    print("=" * 60)
    print("AUTHENTICATION STATUS")
    print("=" * 60)
    
    if _cognito_auth:
        print("✅ Cognito authentication is active")
        
        # Check credentials
        if _cognito_auth.aws_credentials:
            print("✅ AWS credentials loaded")
            print(f"   Access Key: {_cognito_auth.aws_credentials['AccessKeyId'][:20]}...")
        else:
            print("⚠ No AWS credentials")
        
        # Check tokens
        if _cognito_auth.id_token:
            print("✅ Cognito tokens loaded")
            decoded = _cognito_auth.decode_token(_cognito_auth.id_token)
            if decoded:
                print(f"   User: {decoded.get('cognito:username', 'N/A')}")
                print(f"   Email: {decoded.get('email', 'N/A')}")
        else:
            print("⚠ No Cognito tokens")
    else:
        print("⚠ Cognito authentication not setup")
        print("   Run: setup_cognito_auth('train_files/crediential.txt')")
    
    print("=" * 60)


# ══════════════════════════════════════════════════════════════════════════════
# AUTO SETUP
# ══════════════════════════════════════════════════════════════════════════════

def auto_setup(credential_file='train_files/crediential.txt'):
    """
    Automatically setup Cognito authentication if credential file exists
    """
    if os.path.exists(credential_file):
        print(f"🔍 Found credential file: {credential_file}")
        print("🚀 Auto-setting up Cognito authentication...\n")
        return setup_cognito_auth(credential_file)
    else:
        print(f"⚠ Credential file not found: {credential_file}")
        print("  Using default S3 configuration (unsigned)")
        configure_s3_access(aws_unsigned=True)
        return None


# Print module info on import
print("=" * 70)
print("📦 ODC Module with Cognito Authentication Loaded")
print("=" * 70)
print("\n💡 Quick Start:")
print("   1. setup_cognito_auth('train_files/crediential.txt')")
print("   2. Use datacube normally with authenticated S3 access")
print("\n📚 Functions:")
print("   - setup_cognito_auth()    : Setup Cognito authentication")
print("   - get_cognito_auth()      : Get authenticator instance")
print("   - print_auth_status()     : Show auth status")
print("   - auto_setup()            : Auto-setup if credentials exist")
print("=" * 70)
print()