"""
Training module for land classification using Sentinel-2 and Sentinel-1 data
from Microsoft Planetary Computer STAC API
"""

import numpy as np
import xarray as xr
import geopandas as gpd
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import LabelEncoder
from sklearn.metrics import classification_report, confusion_matrix
from sklearn.ensemble import RandomForestClassifier
from sklearn.tree import DecisionTreeClassifier
from sklearn.svm import SVC
from xgboost import XGBClassifier
import joblib
from datetime import datetime
import json
import os
import warnings
import hashlib
from pathlib import Path
warnings.filterwarnings('ignore')

# PyTorch for CNN and advanced models
try:
    import torch
    import torch.nn as nn
    import torch.nn.functional as F
    import torch.optim as optim
    from torch.utils.data import TensorDataset, DataLoader
    import torchvision.models as models
    PYTORCH_AVAILABLE = True
except ImportError:
    PYTORCH_AVAILABLE = False
    print("Warning: PyTorch not available. CNN and advanced models will not work.")

# Define CNN model class for PyTorch
class CNNClassifier(nn.Module):
    def __init__(self, n_features, n_classes):
        super(CNNClassifier, self).__init__()
        self.n_features = n_features
        self.n_classes = n_classes
        
        # For small feature sets (like 3 features), use simpler architecture
        if n_features < 8:
            # Simple fully connected network for small features
            self.use_conv = False
            self.fc1 = nn.Linear(n_features, 64)
            self.dropout1 = nn.Dropout(0.3)
            self.fc2 = nn.Linear(64, 128)
            self.dropout2 = nn.Dropout(0.5)
            self.fc3 = nn.Linear(128, n_classes)
        else:
            # CNN architecture for larger feature sets
            self.use_conv = True
            self.conv1 = nn.Conv1d(in_channels=1, out_channels=32, kernel_size=3, padding=1)
            self.pool1 = nn.MaxPool1d(kernel_size=2)
            self.conv2 = nn.Conv1d(in_channels=32, out_channels=64, kernel_size=3, padding=1)
            self.pool2 = nn.MaxPool1d(kernel_size=2)
            
            # Calculate size after convolutions
            conv_output_size = (n_features // 2 // 2) * 64
            
            # Fully connected layers
            self.fc1 = nn.Linear(conv_output_size, 128)
            self.dropout = nn.Dropout(0.5)
            self.fc2 = nn.Linear(128, n_classes)
    
    def forward(self, x):
        # x shape: (batch, n_features) or (batch, 1, n_features)
        if self.use_conv:
            # CNN path for larger feature sets
            if len(x.shape) == 2:
                x = x.unsqueeze(1)  # Add channel dimension
            x = F.relu(self.conv1(x))
            x = self.pool1(x)
            x = F.relu(self.conv2(x))
            x = self.pool2(x)
            x = x.view(x.size(0), -1)  # Flatten
            x = F.relu(self.fc1(x))
            x = self.dropout(x)
            x = self.fc2(x)
        else:
            # Fully connected path for small feature sets
            if len(x.shape) == 3:
                x = x.squeeze(1)  # Remove channel dimension if present
            x = F.relu(self.fc1(x))
            x = self.dropout1(x)
            x = F.relu(self.fc2(x))
            x = self.dropout2(x)
            x = self.fc3(x)
        return x
    
    def predict(self, X):
        """Scikit-learn style predict method"""
        self.eval()
        with torch.no_grad():
            if isinstance(X, np.ndarray):
                X = torch.FloatTensor(X)
            # Handle both 2D and 3D inputs
            if not self.use_conv and len(X.shape) == 3:
                X = X.squeeze(1)
            elif self.use_conv and len(X.shape) == 2:
                X = X.unsqueeze(1)
            outputs = self(X)
            _, predicted = torch.max(outputs, 1)
            return predicted.cpu().numpy()
    
    def score(self, X, y):
        """Scikit-learn style score method"""
        predictions = self.predict(X)
        if isinstance(y, torch.Tensor):
            y = y.cpu().numpy()
        return np.mean(predictions == y)


# Swin-UNet Classifier for feature vectors
class SwinUNetClassifier(nn.Module):
    """
    Swin Transformer U-Net style architecture adapted for feature vector classification.
    Combines hierarchical Swin Transformer blocks with skip connections.
    """
    def __init__(self, n_features, n_classes, embed_dim=128, depths=(2, 2, 6, 2), num_heads=(4, 8, 16, 32)):
        super(SwinUNetClassifier, self).__init__()
        self.n_features = n_features
        self.n_classes = n_classes
        self.embed_dim = embed_dim
        
        # Feature adapter - convert input features to embedding
        self.adapter = nn.Sequential(
            nn.Linear(n_features, embed_dim * 2),
            nn.ReLU(),
            nn.Dropout(0.1),
            nn.Linear(embed_dim * 2, embed_dim)
        )
        
        # Encoder path with hierarchical structure
        # Stage 1 - 1/4 resolution
        self.encoder1 = nn.Sequential(
            nn.Linear(embed_dim, embed_dim),
            nn.LayerNorm(embed_dim),
            nn.GELU(),
            nn.Dropout(0.1)
        )
        self.down1 = nn.Linear(embed_dim, embed_dim * 2)
        
        # Stage 2 - 1/8 resolution
        self.encoder2 = nn.Sequential(
            nn.Linear(embed_dim * 2, embed_dim * 2),
            nn.LayerNorm(embed_dim * 2),
            nn.GELU(),
            nn.Dropout(0.1)
        )
        self.down2 = nn.Linear(embed_dim * 2, embed_dim * 4)
        
        # Stage 3 - 1/16 resolution (bottleneck)
        self.encoder3 = nn.Sequential(
            nn.Linear(embed_dim * 4, embed_dim * 4),
            nn.LayerNorm(embed_dim * 4),
            nn.GELU(),
            nn.Dropout(0.1)
        )
        
        # Decoder path with skip connections
        self.up2 = nn.Linear(embed_dim * 4, embed_dim * 2)
        self.decoder2 = nn.Sequential(
            nn.Linear(embed_dim * 4, embed_dim * 2),  # Concatenated with skip
            nn.LayerNorm(embed_dim * 2),
            nn.GELU(),
            nn.Dropout(0.1)
        )
        
        self.up1 = nn.Linear(embed_dim * 2, embed_dim)
        self.decoder1 = nn.Sequential(
            nn.Linear(embed_dim * 2, embed_dim),  # Concatenated with skip
            nn.LayerNorm(embed_dim),
            nn.GELU(),
            nn.Dropout(0.1)
        )
        
        # Classification head
        self.classifier = nn.Sequential(
            nn.Linear(embed_dim, embed_dim // 2),
            nn.GELU(),
            nn.Dropout(0.3),
            nn.Linear(embed_dim // 2, n_classes)
        )
        
        # Attention mechanism for better feature aggregation
        self.attention = nn.MultiheadAttention(embed_dim, num_heads=4, batch_first=True)
    
    def forward(self, x):
        # x shape: (batch, n_features)
        if len(x.shape) == 3:
            x = x.squeeze(1)
        
        batch_size = x.shape[0]
        
        # Feature adaptation
        x = self.adapter(x)  # (batch, embed_dim)
        
        # Add sequence dimension for attention (treat as sequence of length 1)
        x_seq = x.unsqueeze(1)  # (batch, 1, embed_dim)
        
        # Encoder path
        # Stage 1
        x1 = self.encoder1(x_seq)  # (batch, 1, embed_dim)
        x_down1 = self.down1(x1.squeeze(1))  # (batch, embed_dim*2)
        
        # Stage 2
        x2 = self.encoder2(x_down1.unsqueeze(1))  # (batch, 1, embed_dim*2)
        x_down2 = self.down2(x2.squeeze(1))  # (batch, embed_dim*4)
        
        # Stage 3 (bottleneck)
        x3 = self.encoder3(x_down2.unsqueeze(1))  # (batch, 1, embed_dim*4)
        
        # Decoder path with skip connections
        # Up2
        x_up2 = self.up2(x3.squeeze(1))  # (batch, embed_dim*2)
        x_cat2 = torch.cat([x_up2, x_down1], dim=1)  # (batch, embed_dim*4) - concatenate skip
        # Create proper 3D tensor for decoder
        x_cat2_seq = x_cat2.unsqueeze(1)  # (batch, 1, embed_dim*4)
        x_dec2 = self.decoder2(x_cat2)  # (batch, embed_dim*2)
        
        # Up1
        x_up1 = self.up1(x_dec2)  # (batch, embed_dim)
        x_cat1 = torch.cat([x_up1, x.squeeze(1)], dim=1)  # (batch, embed_dim*2) - concatenate skip
        x_dec1 = self.decoder1(x_cat1)  # (batch, embed_dim)
        
        # Apply attention mechanism for better aggregation
        x_dec1_seq = x_dec1.unsqueeze(1)  # (batch, 1, embed_dim)
        attn_out, _ = self.attention(x_dec1_seq, x_dec1_seq, x_dec1_seq)
        
        # Classification
        output = self.classifier(attn_out.squeeze(1))
        return output
    
    def predict(self, X):
        """Scikit-learn style predict"""
        self.eval()
        with torch.no_grad():
            if isinstance(X, np.ndarray):
                X = torch.FloatTensor(X)
            outputs = self(X)
            _, predicted = torch.max(outputs, 1)
            return predicted.cpu().numpy()
    
    def score(self, X, y):
        """Scikit-learn style score"""
        predictions = self.predict(X)
        if isinstance(y, torch.Tensor):
            y = y.cpu().numpy()
        return np.mean(predictions == y)


# Microsoft Planetary Computer imports
import planetary_computer
from pystac_client import Client
from odc.stac import load as stac_load

# Feature extraction
from feature_extractor import get_feature_extractor


def train_model(
    bbox=[105.6, 9.3, 106.2, 9.8],
    time_range='2023-03-01/2023-05-31',
    max_scenes=12,
    cloud_cover=30,
    resolution=20,
    training_shapefile='train/ST_training data_updated_1130points_new.shp',
    model_type='xgboost',
    n_estimators=100,
    max_depth=20,
    learning_rate=0.1,
    use_gpu=True,
    use_cache=True,
    test_size=0.2,
    feature_mode='simple',  # Changed from 'odc' - simple mode works with B04, B08, SCL only
    output_model_path=None,
    status_callback=None,
    cancel_check=None
):
    """
    Train a land classification model using Sentinel-2 and Sentinel-1 data
    
    Args:
        bbox: [min_lon, min_lat, max_lon, max_lat]
        time_range: "YYYY-MM-DD/YYYY-MM-DD"
        max_scenes: maximum number of scenes to load
        cloud_cover: maximum cloud cover percentage
        resolution: resolution in meters (e.g., 20)
        training_shapefile: path to training shapefile
        n_estimators: number of trees for XGBoost
        max_depth: maximum tree depth
        learning_rate: learning rate for XGBoost
        use_gpu: whether to use GPU for training
        output_model_path: path to save trained model (auto-generated if None)
        status_callback: Optional callback function to report progress
        cancel_check: Optional function that returns True if training should be cancelled
        test_size: Fraction of data to use for test set (0-1)
        feature_mode: 'simple' (3 features), 'temporal' (39 features), 'extended' (15 features), or 'odc' (8 features)
    
    Returns:
        Dictionary containing training results
    """
    
    def update_status(message, progress=None):
        """Helper to update status"""
        if status_callback:
            # Try calling with both arguments, fallback to just message
            try:
                status_callback(message, progress)
            except TypeError:
                status_callback(message)
        print(message)
    
    def check_cancellation():
        """Check if training should be cancelled"""
        if cancel_check and cancel_check():
            raise InterruptedError("Training cancelled by user")
    
    try:
        # Auto-generate output path if not provided
        if output_model_path is None:
            timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
            output_model_path = f'model_train/model_{model_type}_{timestamp}.joblib'
        
        # ============ CACHE SYSTEM ============
        # Create cache directory
        cache_dir = Path("dataset_cache")
        cache_dir.mkdir(exist_ok=True)
        
        # Generate cache key from parameters
        cache_params = f"{bbox}_{time_range}_{max_scenes}_{cloud_cover}_{resolution}"
        cache_key = hashlib.md5(cache_params.encode()).hexdigest()
        cache_file = cache_dir / f"training_data_{cache_key}.joblib"
        
        features = None
        labels = None
        
        # Initialize FeatureExtractor early (will be used for temporal/extended modes)
        update_status(f"Initializing FeatureExtractor (mode={feature_mode})...", 5)
        extractor = get_feature_extractor(mode=feature_mode)
        
        # Try to load from cache
        if use_cache and cache_file.exists():
            update_status(f"📦 Đang load cache: {cache_file.name}...", 5)
            try:
                cached_data = joblib.load(cache_file)
                features = cached_data['features']
                labels = cached_data['labels']
                
                # Validate cached data
                if len(features) == 0:
                    update_status(
                        f"❌ Cache rỗng (0 samples)! Đây là cache từ lần training thất bại trước.\n"
                        f"   Nguyên nhân: Bbox không overlap với shapefile HOẶC tất cả điểm bị NaN.\n"
                        f"   Đang xóa cache lỗi và tải lại dữ liệu...", 10
                    )
                    cache_file.unlink()  # Delete empty cache
                    features = None
                else:
                    update_status(
                        f"✅ Loaded {len(features)} samples từ cache!\n"
                        f"   ⚡ Đã bỏ qua download위성 data (tiết kiệm thời gian)", 50
                    )
                    print(f"[CACHE HIT] Using cached dataset with {len(features)} samples")
            except Exception as e:
                update_status(f"⚠️ Cache bị lỗi: {str(e)}\n   Đang tải lại dữ liệu mới...", 10)
                features = None
        
        # If no cache or cache failed, download data
        if features is None:
            update_status("📡 Cache not found or disabled, downloading satellite data...", 10)
            
            # Connect to Microsoft Planetary Computer
            update_status("Connecting to Microsoft Planetary Computer...", 12)
            catalog = Client.open("https://planetarycomputer.microsoft.com/api/stac/v1")
            check_cancellation()
            
            # Search for Sentinel-2 scenes
            update_status("Searching for Sentinel-2 scenes...", 10)
            query_s2 = catalog.search(
                collections=["sentinel-2-l2a"],
                bbox=bbox,
                datetime=time_range,
                query={"eo:cloud_cover": {"lt": cloud_cover}}
            )
            items_s2 = list(query_s2.item_collection())
            
            check_cancellation()
            
            # Limit scenes
            if len(items_s2) > max_scenes:
                step = len(items_s2) // max_scenes
                items_s2 = items_s2[::step][:max_scenes]
            
            update_status(f"Found {len(items_s2)} Sentinel-2 scenes", 20)
            
            # Sign and load Sentinel-2 data
            update_status("Loading Sentinel-2 data...", 25)
            items_s2 = [planetary_computer.sign(item) for item in items_s2]
            
            # Load different bands based on feature mode
            if feature_mode == 'simple':
                bands_to_load = ["B04", "B08", "SCL"]
            else:  # temporal or extended
                bands_to_load = ["B02", "B03", "B04", "B08", "B11", "SCL"]
            
            ds_s2 = stac_load(
                items_s2,
                bands=bands_to_load,
                crs="EPSG:32648",
                resolution=resolution,
                bbox=bbox,
                patch_url=planetary_computer.sign,
                fail_on_error=False,
            )
            
            # Debug: Print S2 data info
            print(f"[DEBUG S2] Loaded S2 data")
            print(f"[DEBUG S2] Dimensions: {dict(ds_s2.dims)}")
            print(f"[DEBUG S2] Bands: {list(ds_s2.data_vars)}")
            print(f"[DEBUG S2] CRS: {ds_s2.rio.crs if hasattr(ds_s2, 'rio') else 'No CRS'}")
            print(f"[DEBUG S2] Spatial bounds: x=[{float(ds_s2.x.min())}, {float(ds_s2.x.max())}], y=[{float(ds_s2.y.min())}, {float(ds_s2.y.max())}]")
            if 'time' in ds_s2.dims:
                print(f"[DEBUG S2] Time range: {ds_s2.time.min().values} to {ds_s2.time.max().values}")
            
            # Rename for compatibility (simple mode)
            if "B04" in ds_s2 and "red" not in ds_s2:
                ds_s2 = ds_s2.rename({"B04": "red", "B08": "nir", "SCL": "scl"})
            
            check_cancellation()
            
            # Search for Sentinel-1 scenes
            update_status("Searching for Sentinel-1 scenes...", 35)
            query_s1 = catalog.search(
                collections=["sentinel-1-rtc"],
                bbox=bbox,
                datetime=time_range,
            )
            items_s1 = list(query_s1.item_collection())
            
            # Limit scenes
            if len(items_s1) > max_scenes:
                step = len(items_s1) // max_scenes
                items_s1 = items_s1[::step][:max_scenes]
            
            update_status(f"Found {len(items_s1)} Sentinel-1 scenes", 40)
            
            # Sign and load Sentinel-1 data
            update_status("Loading Sentinel-1 data...", 45)
            items_s1 = [planetary_computer.sign(item) for item in items_s1]
            ds_s1 = stac_load(
                items_s1,
                bands=["vv", "vh"],
                crs="EPSG:32648",
                resolution=resolution,
                bbox=bbox,
                patch_url=planetary_computer.sign,
                fail_on_error=False,
            )
            
            # Convert to dB
            ds_s1['vv_db'] = 10 * np.log10(ds_s1['vv'].where(ds_s1['vv'] > 0))
            ds_s1['vh_db'] = 10 * np.log10(ds_s1['vh'].where(ds_s1['vh'] > 0))
            
            # Debug: Print S1 data info
            print(f"[DEBUG S1] Loaded S1 data")
            print(f"[DEBUG S1] Dimensions: {dict(ds_s1.dims)}")
            print(f"[DEBUG S1] Bands: {list(ds_s1.data_vars)}")
            print(f"[DEBUG S1] Spatial bounds: x=[{float(ds_s1.x.min())}, {float(ds_s1.x.max())}], y=[{float(ds_s1.y.min())}, {float(ds_s1.y.max())}]")
            
            check_cancellation()
            
            # Load training data
            update_status("Loading training data...", 55)
            
            # Normalize training shapefile path
            # If path doesn't start with 'train/', add it
            if not training_shapefile.startswith('train/'):
                training_shapefile = f'train/{training_shapefile}'
            
            print(f"[DEBUG] Original training shapefile: {training_shapefile}")
            print(f"[DEBUG] Current working directory: {os.getcwd()}")
            
            # Try to find the file with exact name first
            if not os.path.exists(training_shapefile):
                # File not found, try to find similar files in train directory
                train_dir = Path('train')
                if train_dir.exists():
                    # List all .shp files
                    shp_files = list(train_dir.glob('*.shp'))
                    print(f"[DEBUG] Available shapefile files in train/:")
                    for f in shp_files:
                        print(f"  - {f.name}")
                    
                    # Try to find a matching file (case-insensitive, ignore underscores vs spaces)
                    filename_normalized = os.path.basename(training_shapefile).lower().replace('_', ' ')
                    for shp_file in shp_files:
                        if shp_file.name.lower().replace('_', ' ') == filename_normalized:
                            print(f"[DEBUG] Found matching file: {shp_file}")
                            training_shapefile = str(shp_file)
                            break
                    
                    if not os.path.exists(training_shapefile):
                        raise FileNotFoundError(
                            f"Training shapefile not found: {training_shapefile}\n"
                            f"Available files: {[f.name for f in shp_files]}"
                        )
                else:
                    raise FileNotFoundError(f"Train directory not found: {train_dir}")
            
            print(f"[DEBUG] Final training shapefile path: {training_shapefile}")
            print(f"[DEBUG] File exists: {os.path.exists(training_shapefile)}")
            
            train_gdf = gpd.read_file(training_shapefile)
            
            # Print initial shapefile info
            update_status(f"📍 Loaded {len(train_gdf)} points from shapefile", 56)
            print(f"[DEBUG] Shapefile CRS: {train_gdf.crs}")
            print(f"[DEBUG] Shapefile bounds: {train_gdf.total_bounds}")
            
            # Convert to WGS84 first (if not already) to match bbox coordinates
            original_crs = train_gdf.crs
            if train_gdf.crs and train_gdf.crs.to_epsg() != 4326:
                print(f"📍 Converting training shapefile from {train_gdf.crs} to WGS84")
                train_gdf = train_gdf.to_crs("EPSG:4326")
                print(f"[DEBUG] WGS84 bounds: {train_gdf.total_bounds}")
            
            # Check bbox overlap in WGS84
            shp_bounds = train_gdf.total_bounds  # [minx, miny, maxx, maxy]
            bbox_wgs84 = bbox  # [min_lon, min_lat, max_lon, max_lat]
            
            # Check if there's overlap
            overlap_x = not (shp_bounds[2] < bbox_wgs84[0] or shp_bounds[0] > bbox_wgs84[2])
            overlap_y = not (shp_bounds[3] < bbox_wgs84[1] or shp_bounds[1] > bbox_wgs84[3])
            
            if not (overlap_x and overlap_y):
                update_status(f"⚠️ WARNING: Shapefile and bbox may not overlap!", 57)
                print(f"[WARNING] Shapefile bounds (WGS84): {shp_bounds}")
                print(f"[WARNING] Requested bbox (WGS84): {bbox_wgs84}")
                print(f"[WARNING] This may result in 0 training samples!")
            else:
                # Crop to bbox to see how many points are actually in the region
                train_gdf_cropped = train_gdf.cx[bbox_wgs84[0]:bbox_wgs84[2], bbox_wgs84[1]:bbox_wgs84[3]]
                update_status(f"📍 {len(train_gdf_cropped)} points within bbox", 57)
                if len(train_gdf_cropped) == 0:
                    raise ValueError(
                        f"No training points found within bbox!\n"
                        f"Shapefile bounds: {shp_bounds}\n"
                        f"Requested bbox: {bbox_wgs84}\n"
                        f"Please adjust bbox to cover your training data."
                    )
            
            # Then convert to UTM Zone 48N (EPSG:32648) for extraction
            if train_gdf.crs.to_epsg() != 32648:
                print(f"📍 Converting training shapefile from WGS84 to UTM Zone 48N (EPSG:32648)")
                train_gdf = train_gdf.to_crs('EPSG:32648')
                print(f"[DEBUG] UTM bounds: {train_gdf.total_bounds}")
            
            # Auto-detect label column
            label_column = None
            for col in ['HT_code', 'Ma_LU', 'LU2022', 'Hientrang', 'class', 'Class', 'CLASS']:
                if col in train_gdf.columns:
                    label_column = col
                    break
            
            if label_column is None:
                raise ValueError(f"Cannot find label column in shapefile. Available: {list(train_gdf.columns)}")
            
            # Extract features using FeatureExtractor
            update_status("Extracting features from satellite data...", 60)
            
            print(f"[DEBUG] Starting feature extraction...")
            print(f"[DEBUG] Training GDF has {len(train_gdf)} points")
            print(f"[DEBUG] Training GDF CRS: {train_gdf.crs}")
            print(f"[DEBUG] Training GDF bounds (UTM): {train_gdf.total_bounds}")
            print(f"[DEBUG] Label column: {label_column}")
            
            if feature_mode == 'simple':
                # For simple mode: calculate NDVI first
                ndvi = (ds_s2['nir'] - ds_s2['red']) / (ds_s2['nir'] + ds_s2['red'] + 1e-8)
                # Apply cloud mask
                cloud_mask = ds_s2['scl'].isin([1, 3, 8, 9, 10])
                ndvi_masked = ndvi.where(~cloud_mask)
                
                print(f"[DEBUG] NDVI shape: {ndvi_masked.shape}")
                print(f"[DEBUG] NDVI range: [{float(ndvi_masked.min())}, {float(ndvi_masked.max())}]")
                
                # Extract features at training points
                features = []
                labels = []
                failed_extractions = 0
                
                # Test first point to see what's happening
                first_point = train_gdf.iloc[0]
                print(f"[DEBUG] Testing first point:")
                print(f"  Coords: ({first_point.geometry.x}, {first_point.geometry.y})")
                print(f"  Label: {first_point[label_column]}")
                
                for idx, row in train_gdf.iterrows():
                    point = row.geometry
                    x_coord = point.x
                    y_coord = point.y
                    label = row[label_column]
                    
                    try:
                        ndvi_val = ndvi_masked.sel(x=x_coord, y=y_coord, method='nearest').mean(dim='time').values
                        vh_val = ds_s1['vh_db'].sel(x=x_coord, y=y_coord, method='nearest').mean(dim='time').values
                        vv_val = ds_s1['vv_db'].sel(x=x_coord, y=y_coord, method='nearest').mean(dim='time').values
                        
                        feature_vec = [float(ndvi_val), float(vh_val), float(vv_val)]
                        
                        # Debug first few points
                        if idx < 3:
                            print(f"[DEBUG] Point {idx}: coords=({x_coord:.2f}, {y_coord:.2f}), ndvi={ndvi_val:.3f}, vh={vh_val:.3f}, vv={vv_val:.3f}")
                        
                        if not np.isnan(feature_vec).any():
                            features.append(feature_vec)
                            labels.append(label)
                        else:
                            failed_extractions += 1
                            if idx < 3:
                                print(f"[DEBUG] Point {idx} has NaN: {feature_vec}")
                    except Exception as e:
                        failed_extractions += 1
                        if idx < 3:
                            print(f"[DEBUG] Point {idx} extraction failed: {e}")
                        continue
                
                if failed_extractions > 0:
                    update_status(f"⚠️ {failed_extractions}/{len(train_gdf)} points had NaN/missing data", 65)
                
                features = np.array(features)
                labels = np.array(labels)
            
            else:  # temporal or extended mode
                # Apply cloud mask for temporal/extended modes
                if 'scl' in ds_s2 or 'SCL' in ds_s2:
                    scl_band = ds_s2['scl'] if 'scl' in ds_s2 else ds_s2['SCL']
                    cloud_mask = scl_band.isin([1, 3, 8, 9, 10])
                    for band in ds_s2.data_vars:
                        if band != 'scl' and band != 'SCL':
                            ds_s2[band] = ds_s2[band].where(~cloud_mask)
                
                # Extract features at training points
                features = []
                labels = []
                failed_extractions = 0
                
                for idx, row in train_gdf.iterrows():
                    point = row.geometry
                    x_coord = point.x
                    y_coord = point.y
                    label = row[label_column]
                    
                    try:
                        # Extract point data from S2
                        point_s2 = ds_s2.sel(x=x_coord, y=y_coord, method='nearest')
                        
                        # Extract point data from S1
                        vh_val = ds_s1['vh_db'].sel(x=x_coord, y=y_coord, method='nearest').mean(dim='time').values
                        vv_val = ds_s1['vv_db'].sel(x=x_coord, y=y_coord, method='nearest').mean(dim='time').values
                        
                        # Create minimal dataset for feature extraction
                        point_data = xr.Dataset({
                            'B02': point_s2['B02'],
                            'B03': point_s2['B03'],
                            'B04': point_s2['B04'],
                            'B08': point_s2['B08'],
                            'B11': point_s2['B11']
                        })
                        
                        # Create VH/VV DataArrays (without spatial dims, just time if exists)
                        if 'time' in point_data.dims:
                            vh_da = xr.DataArray([vh_val] * len(point_data.time), dims=['time'])
                            vv_da = xr.DataArray([vv_val] * len(point_data.time), dims=['time'])
                        else:
                            vh_da = xr.DataArray([vh_val])
                            vv_da = xr.DataArray([vv_val])
                        
                        # Extract features using FeatureExtractor
                        # Note: extractor.extract returns (n_pixels, n_features), we take first row
                        feature_vec = extractor.extract(
                            s2_data=point_data,
                            vh_data=vh_da,
                            vv_data=vv_da
                        )
                        
                        # If feature_vec is 2D, take first row
                        if len(feature_vec.shape) > 1:
                            feature_vec = feature_vec[0]
                        
                        if not np.isnan(feature_vec).any():
                            features.append(feature_vec)
                            labels.append(label)
                        else:
                            failed_extractions += 1
                    except Exception as e:
                        failed_extractions += 1
                        continue
                
                if failed_extractions > 0:
                    update_status(f"⚠️ {failed_extractions}/{len(train_gdf)} points had NaN/missing data", 65)
                
                features = np.array(features)
                labels = np.array(labels)
            
            check_cancellation()
            
            update_status(f"Extracted {len(features)} valid training samples", 70)
            
            # ============ VALIDATE SAMPLES ============
            if len(features) == 0:
                error_msg = (
                    f"❌ No valid training samples extracted!\n"
                    f"Possible reasons:\n"
                    f"1. Training shapefile points don't overlap with bbox: {bbox}\n"
                    f"2. All points have NaN values (cloud cover, missing data)\n"
                    f"3. Coordinate system mismatch\n"
                    f"Suggestions:\n"
                    f"- Check if bbox matches your region\n"
                    f"- Try a different time range with less cloud cover\n"
                    f"- Verify training shapefile coordinates are correct"
                )
                raise ValueError(error_msg)
            
            # Warn if very few samples
            if len(features) < 20:
                update_status(f"⚠️ Warning: Only {len(features)} samples extracted. Results may be unreliable.", 70)
            
            # ============ SAVE TO CACHE ============
            if use_cache:
                update_status(f"💾 Saving dataset to cache for future use...", 72)
                try:
                    cache_data = {
                        'features': features,
                        'labels': labels,
                        'bbox': bbox,
                        'time_range': time_range,
                        'resolution': resolution,
                        'feature_mode': feature_mode,
                        'timestamp': datetime.now().isoformat()
                    }
                    joblib.dump(cache_data, cache_file)
                    update_status(f"✅ Cached to {cache_file.name}", 75)
                except Exception as e:
                    update_status(f"⚠️ Cache save failed: {str(e)}", 75)
        
        # Validate samples after cache loading
        if len(features) == 0:
            error_msg = (
                f"❌ No training samples available!\n"
                f"The cached or loaded dataset is empty.\n"
                f"Please try:\n"
                f"1. Clear cache and reload data\n"
                f"2. Check training shapefile and bbox overlap\n"
                f"3. Adjust time range and cloud cover settings"
            )
            raise ValueError(error_msg)
        
        # Encode labels
        label_encoder = LabelEncoder()
        labels_encoded = label_encoder.fit_transform(labels)
        
        # Split data
        X_train, X_test, y_train, y_test = train_test_split(
            features, labels_encoded, test_size=test_size, random_state=42, stratify=labels_encoded
        )
        
        # Train model based on selected type
        update_status(f"Training {model_type.upper()} model...", 75)
        
        device = 'cuda:0' if use_gpu else 'cpu'
        
        if model_type == 'xgboost':
            model = XGBClassifier(
                n_estimators=n_estimators,
                max_depth=max_depth,
                learning_rate=learning_rate,
                device=device if use_gpu else 'cpu',
                tree_method='hist',
                random_state=42,
                eval_metric='mlogloss',
                verbosity=0
            )
        elif model_type == 'random_forest':
            model = RandomForestClassifier(
                n_estimators=n_estimators,
                max_depth=max_depth,
                random_state=42,
                n_jobs=-1,  # Use all cores
                verbose=0
            )
        elif model_type == 'decision_tree':
            model = DecisionTreeClassifier(
                max_depth=max_depth,
                random_state=42
            )
        elif model_type == 'svm':
            model = SVC(
                kernel='rbf',
                random_state=42,
                verbose=False
            )
        elif model_type == 'cnn':
            if not PYTORCH_AVAILABLE:
                raise ImportError("PyTorch is required for CNN. Install: pip install torch")
            
            # CNN requires reshaping data
            n_features = X_train.shape[1]
            n_classes = len(np.unique(y_train))
            
            # Build PyTorch CNN model
            device = torch.device('cuda' if torch.cuda.is_available() and use_gpu else 'cpu')
            update_status(f"Building CNN model on {device}...", 75)
            
            model = CNNClassifier(n_features, n_classes).to(device)
            
            # Convert to PyTorch tensors
            X_train_tensor = torch.FloatTensor(X_train).unsqueeze(1)  # Add channel dim: (N, 1, features)
            y_train_tensor = torch.LongTensor(y_train)
            X_test_tensor = torch.FloatTensor(X_test).unsqueeze(1)
            y_test_tensor = torch.LongTensor(y_test)
            
            # Create data loaders
            train_dataset = TensorDataset(X_train_tensor, y_train_tensor)
            train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
            
            # Loss and optimizer
            criterion = nn.CrossEntropyLoss()
            optimizer = optim.Adam(model.parameters(), lr=0.001)
            
            # Train CNN
            update_status("Training CNN model with PyTorch...", 80)
            epochs = min(50, n_estimators // 2)  # Use n_estimators as epochs
            
            model.train()
            for epoch in range(epochs):
                epoch_loss = 0.0
                for batch_X, batch_y in train_loader:
                    batch_X, batch_y = batch_X.to(device), batch_y.to(device)
                    
                    optimizer.zero_grad()
                    outputs = model(batch_X)
                    loss = criterion(outputs, batch_y)
                    loss.backward()
                    optimizer.step()
                    
                    epoch_loss += loss.item()
                
                if (epoch + 1) % 10 == 0:
                    avg_loss = epoch_loss / len(train_loader)
                    update_status(f"CNN Epoch {epoch+1}/{epochs}, Loss: {avg_loss:.4f}", 80 + (epoch / epochs) * 10)
            
            # Move model to CPU for saving (compatible with non-GPU systems)
            model = model.cpu()
            model.device_used = str(device)
        
        elif model_type == 'swin-unet':
            if not PYTORCH_AVAILABLE:
                raise ImportError("PyTorch is required for Swin-UNet. Install: pip install torch torchvision")
            
            n_features = X_train.shape[1]
            n_classes = len(np.unique(y_train))
            
            device = torch.device('cuda' if torch.cuda.is_available() and use_gpu else 'cpu')
            update_status(f"Building Swin-UNet model on {device}...", 75)
            
            model = SwinUNetClassifier(n_features, n_classes, embed_dim=128).to(device)
            
            # Convert to PyTorch tensors (no unsqueeze needed for Swin-UNet)
            X_train_tensor = torch.FloatTensor(X_train)
            y_train_tensor = torch.LongTensor(y_train)
            X_test_tensor = torch.FloatTensor(X_test)
            y_test_tensor = torch.LongTensor(y_test)
            
            # Create data loaders
            train_dataset = TensorDataset(X_train_tensor, y_train_tensor)
            train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
            
            # Loss and optimizer with weight decay
            criterion = nn.CrossEntropyLoss()
            optimizer = optim.AdamW(model.parameters(), lr=learning_rate, weight_decay=0.01)
            
            # LR scheduler for better convergence
            scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=50)
            
            # Train Swin-UNet
            update_status("Training Swin-UNet model with PyTorch...", 80)
            epochs = min(60, n_estimators // 2)  # Swin-UNet benefits from more epochs
            
            model.train()
            for epoch in range(epochs):
                epoch_loss = 0.0
                for batch_X, batch_y in train_loader:
                    batch_X, batch_y = batch_X.to(device), batch_y.to(device)
                    
                    optimizer.zero_grad()
                    outputs = model(batch_X)
                    loss = criterion(outputs, batch_y)
                    loss.backward()
                    optimizer.step()
                    
                    epoch_loss += loss.item()
                
                scheduler.step()
                
                if (epoch + 1) % 10 == 0:
                    avg_loss = epoch_loss / len(train_loader)
                    lr = optimizer.param_groups[0]['lr']
                    update_status(f"Swin-UNet Epoch {epoch+1}/{epochs}, Loss: {avg_loss:.4f}, LR: {lr:.6f}", 80 + (epoch / epochs) * 10)
            
            model = model.cpu()
            model.device_used = str(device)
        
        else:
            raise ValueError(f"Unknown model type: {model_type}. Choose: xgboost, random_forest, decision_tree, svm, cnn, swin-unet")
        
        # Fit non-neural-network models
        if model_type not in ['cnn', 'swin-unet']:
            model.fit(X_train, y_train)
        
        # Evaluate
        update_status("Evaluating model...", 90)
        if model_type in ['cnn', 'swin-unet']:
            # PyTorch models evaluation
            train_score = model.score(X_train, y_train)
            test_score = model.score(X_test, y_test)
            y_pred = model.predict(X_test)
        else:
            train_score = model.score(X_train, y_train)
            test_score = model.score(X_test, y_test)
            y_pred = model.predict(X_test)
        
        # Generate classification report and confusion matrix
        update_status("Generating classification report...", 92)
        class_names = label_encoder.classes_.tolist()
        
        # Classification report as dict
        from sklearn.metrics import classification_report, confusion_matrix
        cls_report = classification_report(y_test, y_pred, target_names=class_names, output_dict=True, zero_division=0)
        
        # Confusion matrix
        conf_matrix = confusion_matrix(y_test, y_pred).tolist()
        
        # Save model using ModelManager
        update_status("Saving model...", 95)
        os.makedirs(os.path.dirname(output_model_path), exist_ok=True)
        
        # Get feature names from extractor
        if feature_mode == 'temporal':
            # Calculate n_timesteps from data
            n_timesteps = len(features[0]) // 3 - 1  # (NDVI + NDWI + NDBI) * n_timesteps + 3 radar features
            feature_names = extractor.get_feature_names(n_timesteps=n_timesteps)
        else:
            feature_names = extractor.get_feature_names()
        
        # Prepare metadata
        info = {
            "timestamp": datetime.now().isoformat(),
            "data_source": "Microsoft Planetary Computer STAC",
            "collections": ["sentinel-2-l2a", "sentinel-1-rtc"],
            "features": feature_names,
            "feature_mode": feature_mode,
            "training_samples": len(X_train),
            "testing_samples": len(X_test),
            "test_size": test_size,
            "train_accuracy": float(train_score),
            "test_accuracy": float(test_score),
            "model_type": model_type,
            "device": device if model_type == 'xgboost' else 'cpu',
            "n_estimators": n_estimators if model_type in ['xgboost', 'random_forest', 'cnn', 'swin-unet'] else None,
            "max_depth": max_depth if model_type not in ['cnn', 'swin-unet'] else None,
            "learning_rate": learning_rate if model_type in ['xgboost', 'swin-unet'] else None,
            "epochs": min(50, n_estimators // 2) if model_type == 'cnn' else (min(60, n_estimators // 2) if model_type == 'swin-unet' else None),
            "n_features": X_train.shape[1],
            "n_classes": len(np.unique(y_train)),
            "class_names": class_names,
            "classification_report": cls_report,
            "confusion_matrix": conf_matrix,
            "bbox": bbox,
            "time_range": time_range,
            "resolution": resolution
        }
        
        # Use ModelManager to save
        from model_manager import get_model_manager
        model_manager = get_model_manager()
        model_filename = os.path.basename(output_model_path)
        model_manager.save_model(
            model=model,
            metadata=info,
            model_filename=model_filename,
            label_encoder=label_encoder
        )
        
        # Construct info path (model manager saves it in model_train/)
        info_path = os.path.join('model_train', model_filename.replace('.joblib', '_info.json'))
        
        update_status("Training complete!", 100)
        
        return {
            "success": True,
            "model_path": output_model_path,
            "info_path": info_path,
            "train_accuracy": train_score,
            "test_accuracy": test_score,
            "training_samples": len(X_train),
            "testing_samples": len(X_test),
            "test_size": test_size,
            "classes": class_names,
            "classification_report": cls_report,
            "confusion_matrix": conf_matrix,
            "model_type": model_type,
            "bbox": bbox,
            "time_range": time_range,
            "resolution": resolution
        }
        
    except InterruptedError as e:
        update_status(f"Cancelled: {str(e)}", -1)
        return {
            "success": False,
            "error": str(e),
            "cancelled": True
        }
        
    except Exception as e:
        update_status(f"Error: {str(e)}", -1)
        return {
            "success": False,
            "error": str(e)
        }