remote-sensing/train_module.py

"""
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
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
    PYTORCH_AVAILABLE = True
except ImportError:
    PYTORCH_AVAILABLE = False
    print("Warning: PyTorch not available. CNN model 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)

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


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,
    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
    
    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
        
        # Try to load from cache
        if use_cache and cache_file.exists():
            update_status(f"📦 Loading cached dataset from {cache_file.name}...", 5)
            try:
                cached_data = joblib.load(cache_file)
                features = cached_data['features']
                labels = cached_data['labels']
                update_status(f"✅ Loaded {len(features)} samples from cache (skipped satellite download!)", 50)
            except Exception as e:
                update_status(f"⚠️ Cache load failed: {str(e)}, downloading fresh data...", 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]
            ds_s2 = stac_load(
                items_s2,
                bands=["B04", "B08", "SCL"],
                crs="EPSG:32648",
                resolution=resolution,
                bbox=bbox,
                patch_url=planetary_computer.sign,
                fail_on_error=False,
            )
            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))
            
            check_cancellation()
            
            # Calculate NDVI
            update_status("Calculating NDVI...", 50)
            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)
            ndvi_mean = ndvi_masked.mean(dim='time')
            
            # Load training data
            update_status("Loading training data...", 55)
            train_gdf = gpd.read_file(training_shapefile)
            
            if train_gdf.crs != 'EPSG:32648':
                train_gdf = train_gdf.to_crs('EPSG:32648')
            
            # 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
            update_status("Extracting features from training points...", 60)
            features = []
            labels = []
            
            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_mean.sel(x=x_coord, y=y_coord, method='nearest').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 = [ndvi_val, vh_val, vv_val]
                    
                    if not np.isnan(feature_vec).any():
                        features.append(feature_vec)
                        labels.append(label)
                except:
                    continue
            
            features = np.array(features)
            labels = np.array(labels)
            
            check_cancellation()
            
            update_status(f"Extracted {len(features)} valid training samples", 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,
                        '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)
        
        # 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=0.2, 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)
        else:
            raise ValueError(f"Unknown model type: {model_type}. Choose: xgboost, random_forest, decision_tree, svm, cnn")
        
        # Fit non-CNN models
        if model_type != 'cnn':
            model.fit(X_train, y_train)
        
        # Evaluate
        update_status("Evaluating model...", 90)
        if model_type == 'cnn':
            # PyTorch CNN evaluation
            train_score = model.score(X_train, y_train)
            test_score = model.score(X_test, y_test)
        else:
            train_score = model.score(X_train, y_train)
            test_score = model.score(X_test, y_test)
        
        # Save model
        update_status("Saving model...", 95)
        os.makedirs(os.path.dirname(output_model_path), exist_ok=True)
        joblib.dump({'model': model, 'label_encoder': label_encoder}, output_model_path)
        
        # Save model info
        info = {
            "timestamp": datetime.now().isoformat(),
            "data_source": "Microsoft Planetary Computer STAC",
            "collections": ["sentinel-2-l2a", "sentinel-1-rtc"],
            "features": ["NDVI_mean", "VH_dB_mean", "VV_dB_mean"],
            "training_samples": len(X_train),
            "testing_samples": len(X_test),
            "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'] else None,
            "max_depth": max_depth if model_type != 'cnn' else None,
            "learning_rate": learning_rate if model_type == 'xgboost' else None,
            "cnn_epochs": min(50, n_estimators // 2) if model_type == 'cnn' else None,
            "n_features": X_train.shape[1],
            "n_classes": len(np.unique(y_train)),
            "bbox": bbox,
            "time_range": time_range,
            "resolution": resolution
        }
        
        info_path = output_model_path.replace('.joblib', '_info.json')
        with open(info_path, 'w') as f:
            json.dump(info, f, indent=2)
        
        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),
            "classes": label_encoder.classes_.tolist()
        }
        
    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)
        }