remote-sensing/api_server.py

"""
API Server for Land Classification Model Training
Cho phép chọn dữ liệu và cấu hình training qua giao diện web
"""

from fastapi import FastAPI, BackgroundTasks, HTTPException, UploadFile, File
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from fastapi.responses import HTMLResponse, FileResponse
from pydantic import BaseModel
from typing import Optional, List
import uvicorn
import joblib
import json
from datetime import datetime
from pathlib import Path
import sys
import numpy as np
import xarray as xr
import rasterio
from rasterio.transform import from_bounds
import asyncio
import hashlib
import traceback

# Import report generator
from report_generator import generate_training_report, generate_prediction_report

# Import Model Manager
from model_manager import ModelManager, get_model_manager

# Import Vietnam provinces data
from vietnam_provinces import get_all_provinces, get_provinces_by_region, get_province_bbox, search_province
from vietnam_provinces_merged import (
    get_all_provinces_32, get_provinces_by_region_32, get_province_bbox_32, 
    search_province_32, get_merged_info, get_provinces_statistics
)

# Import planetary computer libraries (conditional)
try:
    from pystac_client import Client
    import planetary_computer
    import odc.stac
except ImportError:
    Client = None
    planetary_computer = None
    odc = None

app = FastAPI(title="Land Classification Training API", version="1.0.0")

# Enable CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Global training status``
training_status = {
    "is_training": False,
    "progress": "",
    "error": None,
    "result": None,
    "start_time": None,
    "end_time": None,
    "cancel_requested": False
}

# Global prediction status
prediction_status = {
    "is_predicting": False,
    "progress": "",
    "error": None,
    "result": None,
    "output_file": None,
    "start_time": None,
    "end_time": None
}

# Batch prediction queue
batch_queue = []
batch_results = []


class TrainingConfig(BaseModel):
    """Cấu hình training"""
    # Khu vực (bbox)
    min_lon: float = 105.6
    min_lat: float = 9.3
    max_lon: float = 106.2
    max_lat: float = 9.8
    
    # Thời gian
    start_date: str = "2023-03-01"
    end_date: str = "2023-05-31"
    
    # Dữ liệu
    max_scenes: int = 12
    cloud_cover: int = 30
    resolution: int = 20  # 10m hoặc 20m
    
    # Model parameters
    model_type: str = "xgboost"  # xgboost, random_forest, decision_tree, svm, cnn, swin-unet
    n_estimators: int = 100
    max_depth: int = 20
    learning_rate: float = 0.1
    use_gpu: bool = True
    
    # Train/test split
    test_size: float = 0.2  # Tỷ lệ dữ liệu dùng làm test (0-1)
    
    # Cache
    use_cache: bool = True  # Cache dataset để test nhanh hơn
    
    # Training data
    training_shapefile: str = "train/ST_training data_updated_1130points_new.shp"


class PredictionConfig(BaseModel):
    """Cấu hình dự đoán"""
    # Model to use
    model_filename: str
    
    # Khu vực (bbox)
    min_lon: float = 105.6
    min_lat: float = 9.3
    max_lon: float = 106.2
    max_lat: float = 9.8
    
    # Thời gian
    start_date: str = "2023-03-01"
    end_date: str = "2023-05-31"
    
    # Dữ liệu
    max_scenes: int = 12
    cloud_cover: int = 30
    resolution: int = 20
    
    # GPU support for deep learning models
    use_gpu: bool = True


class TrainingStatus(BaseModel):
    """Trạng thái training"""
    is_training: bool
    progress: str
    error: Optional[str]
    result: Optional[dict]
    start_time: Optional[str]
    end_time: Optional[str]


class NDVIConfig(BaseModel):
    """Cấu hình tính NDVI time series"""
    bbox: List[float]  # [min_lon, min_lat, max_lon, max_lat]
    start_date: str
    end_date: str
    max_cloud_cover: int = 30
    resolution: int = 20


class ChangeDetectionWorkflowRequest(BaseModel):
    """Request for change detection workflow"""
    prediction_result: dict
    bbox: List[float]


class ComparePeriodsPredictionConfig(BaseModel):
    """Compare predictions between two time periods"""
    model_filename: str
    min_lon: float
    min_lat: float
    max_lon: float
    max_lat: float
    current_period: dict  # {start_date, end_date}
    prediction_period: dict  # {start_date, end_date}
    max_scenes: int = 12
    cloud_cover: int = 30
    resolution: int = 20
    export_ndvi: bool = True
    export_classification: bool = True


class PredictionWithNDVIConfig(BaseModel):
    """Cấu hình predict kết hợp land classification và NDVI"""
    model_filename: str
    min_lon: float
    min_lat: float
    max_lon: float
    max_lat: float
    start_date: str
    end_date: str
    max_scenes: int = 12
    cloud_cover: int = 30
    resolution: int = 20
    use_gpu: bool = False  # Use GPU for deep learning models
    export_ndvi: bool = True  # Export NDVI raster
    export_classification: bool = True  # Export classification raster


# Serve change detection interface page (moved here after app is defined)
@app.get("/change-detection", response_class=HTMLResponse)
async def change_detection_page():
    html_file = Path(__file__).parent / "change_detection_interface.html"
    if html_file.exists():
        return FileResponse(html_file)
    else:
        return HTMLResponse("<h2>Change Detection Interface not found.</h2>")


@app.get("/", response_class=HTMLResponse)
async def root():
    """Serve main index page with tabs"""
    html_file = Path(__file__).parent / "index.html"
    if html_file.exists():
        return FileResponse(html_file)
    else:
        return HTMLResponse("""
        <html>
            <head><title>Land Classification System</title></head>
            <body>
                <h1>Land Classification System</h1>
                <p>API Documentation: <a href="/docs">/docs</a></p>
                <p>Training: <a href="/training">/training</a></p>
                <p>Prediction: <a href="/prediction">/prediction</a></p>
                <p>Dashboard: <a href="/dashboard">/dashboard</a></p>
            </body>
        </html>
        """)


@app.get("/training", response_class=HTMLResponse)
async def training_page():
    """Serve training interface"""
    html_file = Path(__file__).parent / "training_interface.html"
    if html_file.exists():
        return FileResponse(html_file)
    else:
        raise HTTPException(status_code=404, detail="Training interface không tồn tại")


@app.get("/prediction", response_class=HTMLResponse)
async def prediction_page():
    """Serve prediction interface"""
    html_file = Path(__file__).parent / "prediction_interface.html"
    if html_file.exists():
        return FileResponse(html_file)
    else:
        raise HTTPException(status_code=404, detail="Prediction interface không tồn tại")


@app.get("/dashboard", response_class=HTMLResponse)
async def dashboard():
    """Serve dashboard visualization"""
    html_file = Path(__file__).parent / "dashboard.html"
    if html_file.exists():
        return FileResponse(html_file)
    else:
        raise HTTPException(status_code=404, detail="Dashboard không tồn tại")


@app.get("/batch", response_class=HTMLResponse)
async def batch_page():
    """Serve batch processing interface"""
    html_file = Path(__file__).parent / "batch_interface.html"
    if html_file.exists():
        return FileResponse(html_file)
    else:
        raise HTTPException(status_code=404, detail="Batch interface không tồn tại")


@app.get("/ndvi", response_class=HTMLResponse)
async def ndvi_page():
    """Serve NDVI time series interface"""
    html_file = Path(__file__).parent / "ndvi_interface.html"
    if html_file.exists():
        return FileResponse(html_file)
    else:
        raise HTTPException(status_code=404, detail="NDVI interface không tồn tại")


@app.get("/reports", response_class=HTMLResponse)
async def reports_page():
    """Serve reports management interface"""
    html_file = Path(__file__).parent / "reports_interface.html"
    if html_file.exists():
        return FileResponse(html_file)
    else:
        raise HTTPException(status_code=404, detail="Reports interface không tồn tại")


@app.get("/api/models/list")
async def list_models():
    """Liệt kê tất cả models có sẵn với metadata"""
    try:
        model_manager = get_model_manager()
        models = model_manager.list_models()
        return {
            "success": True,
            "models": models,
            "count": len(models)
        }
    except Exception as e:
        return {
            "success": False,
            "error": str(e),
            "models": []
        }


@app.get("/api/models/{model_filename}/info")
async def get_model_info(model_filename: str):
    """Lấy thông tin chi tiết về model"""
    try:
        model_manager = get_model_manager()
        info = model_manager.get_model_info(model_filename)
        if info is None:
            raise HTTPException(status_code=404, detail=f"Model không tồn tại: {model_filename}")
        return {
            "success": True,
            "model": info
        }
    except HTTPException:
        raise
    except Exception as e:
        return {
            "success": False,
            "error": str(e)
        }


@app.get("/api/models/{model_filename}/validate")
async def validate_model(model_filename: str):
    """Validate model file"""
    try:
        model_manager = get_model_manager()
        validation = model_manager.validate_model(model_filename)
        return {
            "success": True,
            "validation": validation
        }
    except Exception as e:
        return {
            "success": False,
            "error": str(e)
        }


@app.delete("/api/models/{model_filename}")
async def delete_model(model_filename: str):
    """Xóa model"""
    try:
        model_manager = get_model_manager()
        success = model_manager.delete_model(model_filename)
        if not success:
            raise HTTPException(status_code=404, detail=f"Model không tồn tại: {model_filename}")
        return {
            "success": True,
            "message": f"Đã xóa model: {model_filename}"
        }
    except HTTPException:
        raise
    except Exception as e:
        return {
            "success": False,
            "error": str(e)
        }


@app.get("/api/config/presets")
async def get_presets():
    """Lấy các preset cấu hình sẵn"""
    return {
        "presets": [
            {
                "name": "PC - Nhỏ (3 tháng, 20m, 12 scenes)",
                "config": {
                    "min_lon": 105.6, "min_lat": 9.3, "max_lon": 106.2, "max_lat": 9.8,
                    "start_date": "2023-03-01", "end_date": "2023-05-31",
                    "max_scenes": 12, "cloud_cover": 30, "resolution": 20,
                    "test_size": 0.2
                }
            },
            {
                "name": "Server - Trung bình (6 tháng, 10m, 30 scenes)",
                "config": {
                    "min_lon": 105.5, "min_lat": 9.2, "max_lon": 106.4, "max_lat": 10.0,
                    "start_date": "2023-01-01", "end_date": "2023-06-30",
                    "max_scenes": 30, "cloud_cover": 30, "resolution": 10,
                    "test_size": 0.2
                }
            },
            {
                "name": "Full - Lớn (1 năm, 10m, 60 scenes)",
                "config": {
                    "min_lon": 105.5, "min_lat": 9.2, "max_lon": 106.4, "max_lat": 10.0,
                    "start_date": "2022-09-01", "end_date": "2023-10-01",
                    "max_scenes": 60, "cloud_cover": 50, "resolution": 10,
                    "test_size": 0.2
                }
            }
        ]
    }


@app.get("/api/provinces/list")
async def list_provinces():
    """Lấy danh sách tất cả các tỉnh thành Việt Nam"""
    return {
        "provinces": get_all_provinces(),
        "count": len(get_all_provinces())
    }


@app.get("/api/provinces/by-region")
async def list_provinces_by_region():
    """Lấy danh sách tỉnh thành theo vùng miền"""
    return get_provinces_by_region()


@app.get("/api/provinces/{province_name}/bbox")
async def get_province_bbox_api(province_name: str):
    """Lấy bbox của một tỉnh thành"""
    bbox = get_province_bbox(province_name)
    if bbox is None:
        raise HTTPException(status_code=404, detail=f"Không tìm thấy tỉnh: {province_name}")
    return {
        "province": province_name,
        "bbox": bbox,
        "min_lon": bbox[0],
        "min_lat": bbox[1],
        "max_lon": bbox[2],
        "max_lat": bbox[3]
    }


@app.get("/api/provinces/search/{query}")
async def search_provinces(query: str):
    """Tìm kiếm tỉnh thành theo tên"""
    results = search_province(query)
    return {
        "query": query,
        "results": results,
        "count": len(results)
    }


@app.get("/api/provinces-32/list")
async def list_provinces_32():
    """Lấy danh sách 32 tỉnh thành sau sáp nhập"""
    return {
        "provinces": get_all_provinces_32(),
        "count": len(get_all_provinces_32()),
        "note": "32 tỉnh thành sau sáp nhập theo Nghị quyết 1211/2023"
    }


@app.get("/api/provinces-32/by-region")
async def list_provinces_by_region_32():
    """Lấy danh sách 32 tỉnh thành theo vùng miền"""
    return get_provinces_by_region_32()


@app.get("/api/provinces-32/{province_name}/bbox")
async def get_province_bbox_api_32(province_name: str):
    """Lấy bbox của một tỉnh thành (32 tỉnh)"""
    bbox = get_province_bbox_32(province_name)
    if bbox is None:
        raise HTTPException(status_code=404, detail=f"Không tìm thấy tỉnh: {province_name}")
    
    # Get merged info
    info = get_merged_info(province_name)
    
    return {
        "province": province_name,
        "bbox": bbox,
        "min_lon": bbox[0],
        "min_lat": bbox[1],
        "max_lon": bbox[2],
        "max_lat": bbox[3],
        "merged_from": info.get("merged_from"),
        "area_km2": info.get("area_km2"),
        "region": info.get("region")
    }


@app.get("/api/provinces-32/search/{query}")
async def search_provinces_32(query: str):
    """Tìm kiếm tỉnh thành theo tên (32 tỉnh)"""
    results = search_province_32(query)
    return {
        "query": query,
        "results": results,
        "count": len(results)
    }


@app.get("/api/provinces-32/statistics")
async def get_provinces_stats():
    """Thống kê các tỉnh đã sáp nhập"""
    return get_provinces_statistics()


@app.get("/api/training/status", response_model=TrainingStatus)
async def get_training_status():
    """Kiểm tra trạng thái training"""
    return training_status


@app.post("/api/training/start")
async def start_training(config: TrainingConfig, background_tasks: BackgroundTasks):
    """Bắt đầu training với config đã chọn"""
    global training_status
    
    if training_status["is_training"]:
        raise HTTPException(status_code=400, detail="Training đang chạy, vui lòng đợi")
    
    # Reset status
    training_status = {
        "is_training": True,
        "progress": "Đang khởi tạo...",
        "error": None,
        "result": None,
        "start_time": datetime.now().isoformat(),
        "end_time": None
    }
    
    # Run training in background
    background_tasks.add_task(run_training, config)
    
    return {"message": "Training đã bắt đầu", "status": training_status}


@app.post("/api/training/stop")
async def stop_training():
    """Dừng training (nếu đang chạy)"""
    global training_status
    
    if not training_status["is_training"]:
        return {"message": "Không có training nào đang chạy"}
    
    # Set cancel flag - the training will check this and stop
    training_status["cancel_requested"] = True
    training_status["progress"] = "Đang hủy training..."
    
    return {"message": "Đang dừng training..."}


@app.post("/api/cache/clear")
async def clear_cache():
    """Xóa cache dataset"""
    import shutil
    cache_dir = Path("dataset_cache")
    
    if not cache_dir.exists():
        return {"message": "Không có cache để xóa", "deleted": 0}
    
    # Count files
    cache_files = list(cache_dir.glob("*.joblib"))
    count = len(cache_files)
    
    # Delete all cache files
    for cache_file in cache_files:
        try:
            cache_file.unlink()
        except:
            pass
    
    return {"message": f"Đã xóa {count} file cache", "deleted": count}


@app.get("/api/cache/info")
async def get_cache_info():
    """Lấy thông tin về cache với metadata đầy đủ, tự động xóa cache cũ có lazy data"""
    cache_dir = Path("dataset_cache")
    
    if not cache_dir.exists():
        return {"exists": False, "files": [], "total_size_mb": 0}
    
    cache_files = []
    total_size = 0
    deleted_count = 0
    
    for cache_file in cache_dir.glob("*.joblib"):
        # Skip if file doesn't exist (race condition)
        if not cache_file.exists():
            continue
            
        size = cache_file.stat().st_size
        
        # Try to load metadata from cache and check if it's valid
        metadata = {}
        is_valid = True
        try:
            cached_data = joblib.load(cache_file)
            
            # Check if cache contains lazy data (will cause 403 errors)
            if isinstance(cached_data, dict) and "s2_data" in cached_data:
                s2_data_temp = cached_data["s2_data"]
                is_lazy = False
                try:
                    is_lazy = any(hasattr(s2_data_temp[var].data, 'chunks') for var in s2_data_temp.data_vars)
                except:
                    pass
                
                if is_lazy:
                    print(f"[CLEANUP] Deleting cache with lazy data: {cache_file.name}")
                    cache_file.unlink()
                    deleted_count += 1
                    is_valid = False
            
            if is_valid and isinstance(cached_data, dict):
                metadata = {
                    "bbox": cached_data.get("bbox", []),
                    "time_range": cached_data.get("time_range", ""),
                    "resolution": cached_data.get("resolution", 20),
                    "n_samples": len(cached_data.get("features", [])),
                    "created": cached_data.get("timestamp", "")
                }
                # Parse time_range to get start/end dates
                if metadata["time_range"]:
                    time_parts = metadata["time_range"].split("/")
                    if len(time_parts) == 2:
                        metadata["start_date"] = time_parts[0]
                        metadata["end_date"] = time_parts[1]
                # Parse bbox to get min/max lon/lat
                if metadata["bbox"] and len(metadata["bbox"]) == 4:
                    metadata["min_lon"] = metadata["bbox"][0]
                    metadata["min_lat"] = metadata["bbox"][1]
                    metadata["max_lon"] = metadata["bbox"][2]
                    metadata["max_lat"] = metadata["bbox"][3]
        except Exception as e:
            print(f"[CLEANUP] Error loading cache {cache_file.name}: {e}. Deleting...")
            try:
                cache_file.unlink()
                deleted_count += 1
                is_valid = False
            except:
                pass
        
        # Only add valid cache files to the list
        if is_valid:
            total_size += size
            cache_files.append({
                "filename": cache_file.name,
                "size_mb": round(size / 1024 / 1024, 2),
                "modified": datetime.fromtimestamp(cache_file.stat().st_mtime).isoformat(),
                "metadata": metadata
            })
    
    # Sort by modified time (newest first)
    cache_files.sort(key=lambda x: x["modified"], reverse=True)
    
    if deleted_count > 0:
        print(f"[CLEANUP] Deleted {deleted_count} invalid cache files")
    
    return {
        "exists": True,
        "files": cache_files,
        "count": len(cache_files),
        "total_size_mb": round(total_size / 1024 / 1024, 2),
        "deleted_invalid": deleted_count
    }



@app.get("/api/models/list")
async def list_models():
    """Liệt kê các model đã train"""
    model_dir = Path("model_train")
    if not model_dir.exists():
        return {"models": []}

    models = []
    # List all .joblib model files (actual trained models)
    for model_file in model_dir.glob("*.joblib"):
        # Skip any file that contains '_info' in its name
        if '_info' in model_file.stem:
            continue
            
        info = {}
        # Try to find corresponding .json info file
        # Remove .joblib and try with _info.json
        base_name = model_file.stem  # e.g., "model_cnn_20251221_163841"
        info_file = model_dir / f"{base_name}_info.json"
        
        if info_file.exists():
            try:
                with open(info_file) as f:
                    info = json.load(f)
            except Exception as e:
                info = {"error": str(e)}
        
        size_mb = round(model_file.stat().st_size / 1024 / 1024, 2)
        created = datetime.fromtimestamp(model_file.stat().st_mtime).isoformat()
        
        models.append({
            "filename": model_file.name,
            "created": created,
            "size_mb": size_mb,
            "info": info
        })

    # Sort by creation time (newest first)
    models.sort(key=lambda x: x["created"], reverse=True)
    return {"models": models}


# ============ REPORTS API ============

@app.get("/api/reports/list")
async def list_reports():
    """Liệt kê các báo cáo đã tạo"""
    reports_dir = Path("reports")
    reports_dir.mkdir(exist_ok=True)
    
    reports = []
    for report_file in reports_dir.glob("*.html"):
        # Determine report type from filename
        if "training" in report_file.name:
            report_type = "training"
        elif "prediction" in report_file.name:
            report_type = "prediction"
        else:
            report_type = "unknown"
        
        report_info = {
            "filename": report_file.name,
            "type": report_type,
            "created": datetime.fromtimestamp(report_file.stat().st_mtime).isoformat(),
            "size_kb": round(report_file.stat().st_size / 1024, 2),
            "view_url": f"/api/reports/view/{report_file.name}",
            "download_url": f"/api/reports/download/{report_file.name}",
            "is_batch_job": False,
            "batch_metadata": None
        }
        
        # Check if this is a batch job report
        if report_type == "prediction":
            predictions_dir = Path("predictions")
            # Look for batch metadata JSON files that reference this report
            for json_file in predictions_dir.glob("batch_*.json"):
                try:
                    import json
                    with open(json_file, 'r') as f:
                        metadata = json.load(f)
                        if metadata.get("report_filename") == report_file.name or \
                           (metadata.get("batch_job_id") and report_file.name.endswith('.html')):
                            report_info["is_batch_job"] = True
                            report_info["batch_metadata"] = {
                                "batch_job_id": metadata.get("batch_job_id"),
                                "batch_name": metadata.get("batch_name"),
                                "batch_timestamp": metadata.get("batch_timestamp")
                            }
                            break
                except Exception as e:
                    pass
        
        reports.append(report_info)
    
    # Sort by creation time (newest first)
    reports.sort(key=lambda x: x["created"], reverse=True)
    return {"reports": reports, "count": len(reports)}


@app.get("/api/reports/view/{filename}", response_class=HTMLResponse)
async def view_report(filename: str):
    """Xem báo cáo HTML trực tiếp"""
    reports_dir = Path("reports")
    file_path = reports_dir / filename
    
    # Security check
    if ".." in filename or "/" in filename or "\\" in filename:
        raise HTTPException(status_code=400, detail="Invalid filename")
    
    if not file_path.exists():
        raise HTTPException(status_code=404, detail=f"Report không tồn tại: {filename}")
    
    with open(file_path, 'r', encoding='utf-8') as f:
        html_content = f.read()
    
    return HTMLResponse(content=html_content)


@app.get("/api/reports/download/{filename}")
async def download_report(filename: str):
    """Download báo cáo HTML"""
    reports_dir = Path("reports")
    file_path = reports_dir / filename
    
    # Security check
    if ".." in filename or "/" in filename or "\\" in filename:
        raise HTTPException(status_code=400, detail="Invalid filename")
    
    if not file_path.exists():
        raise HTTPException(status_code=404, detail=f"Report không tồn tại: {filename}")
    
    return FileResponse(
        path=str(file_path),
        filename=filename,
        media_type="text/html",
        headers={
            "Content-Disposition": f"attachment; filename={filename}"
        }
    )


@app.delete("/api/reports/delete/{filename}")
async def delete_report(filename: str):
    """Xóa một báo cáo"""
    reports_dir = Path("reports")
    file_path = reports_dir / filename
    
    # Security check
    if ".." in filename or "/" in filename or "\\" in filename:
        raise HTTPException(status_code=400, detail="Invalid filename")
    
    if not file_path.exists():
        raise HTTPException(status_code=404, detail=f"Report không tồn tại: {filename}")
    
    try:
        file_path.unlink()
        return {"message": f"Đã xóa báo cáo: {filename}", "success": True}
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Không thể xóa: {str(e)}")


@app.post("/api/prediction/start")
async def start_prediction(config: PredictionConfig, background_tasks: BackgroundTasks):
    """Bắt đầu dự đoán"""
    global prediction_status
    
    if prediction_status["is_predicting"]:
        raise HTTPException(status_code=400, detail="Đang có dự đoán khác đang chạy")
    
    # Reset status
    prediction_status = {
        "is_predicting": True,
        "progress": "Đang khởi động...",
        "error": None,
        "result": None,
        "output_file": None,
        "start_time": datetime.now().isoformat(),
        "end_time": None
    }
    
    # Run prediction in background
    background_tasks.add_task(run_prediction, config)
    
    return {"message": "Đã bắt đầu dự đoán", "status": prediction_status}


@app.get("/api/prediction/status")
async def get_prediction_status():
    """Kiểm tra trạng thái dự đoán"""
    return prediction_status


async def run_training(config: TrainingConfig):
    """Chạy training process"""
    global training_status
    
    try:
        training_status["cancel_requested"] = False
        training_status["progress"] = "Đang import thư viện..."
        
        # Import training module
        from train_module import train_model
        
        training_status["progress"] = "Đang load dữ liệu Sentinel-2..."
        
        # Function to check if training should be cancelled
        def should_cancel():
            return training_status.get("cancel_requested", False)
        
        # Run training
        result = train_model(
            bbox=[config.min_lon, config.min_lat, config.max_lon, config.max_lat],
            time_range=f"{config.start_date}/{config.end_date}",
            max_scenes=config.max_scenes,
            cloud_cover=config.cloud_cover,
            resolution=config.resolution,
            training_shapefile=config.training_shapefile,
            model_type=config.model_type,
            n_estimators=config.n_estimators,
            max_depth=config.max_depth,
            learning_rate=config.learning_rate,
            use_gpu=config.use_gpu,
            use_cache=config.use_cache,
            test_size=config.test_size,
            status_callback=lambda msg: update_progress(msg),
            cancel_check=should_cancel
        )
        
        if training_status.get("cancel_requested", False):
            training_status["is_training"] = False
            training_status["progress"] = "Đã hủy training"
            training_status["error"] = "Training cancelled by user"
        else:
            training_status["is_training"] = False
            training_status["progress"] = "Hoàn thành! Đang tạo báo cáo..."
            training_status["result"] = result
            
            # Auto generate report
            if result.get("success", False):
                try:
                    report_path, _ = generate_training_report(result)
                    training_status["result"]["report_path"] = report_path
                    training_status["result"]["report_filename"] = Path(report_path).name
                    training_status["progress"] = "Hoàn thành! Báo cáo đã được tạo."
                    print(f"[REPORT] Generated: {report_path}")
                except Exception as e:
                    print(f"[REPORT ERROR] Failed to generate report: {e}")
                    training_status["progress"] = "Hoàn thành! (Không thể tạo báo cáo)"
        
        training_status["end_time"] = datetime.now().isoformat()
        
    except Exception as e:
        training_status["is_training"] = False
        training_status["error"] = str(e)
        training_status["progress"] = f"Lỗi: {str(e)}"
        training_status["end_time"] = datetime.now().isoformat()
        import traceback
        print(traceback.format_exc())


def update_progress(message: str):
    """Cập nhật progress message"""
    global training_status
    training_status["progress"] = message
    print(f"[PROGRESS] {message}")


def update_prediction_progress(message: str):
    """Cập nhật prediction progress message"""
    global prediction_status
    prediction_status["progress"] = message
    print(f"[PREDICTION PROGRESS] {message}")


async def run_prediction(config: PredictionConfig):
    """Chạy prediction process - Sử dụng FeatureExtractor để đồng bộ với training"""
    global prediction_status
    
    try:
        prediction_status["progress"] = "Đang import thư viện..."
        
        # Import required libraries
        import numpy as np
        import xarray as xr
        import rioxarray
        from datetime import datetime as dt
        from feature_extractor import get_feature_extractor
        
        # Validate bbox
        if (config.min_lon < -180 or config.max_lon > 180 or 
            config.min_lat < -90 or config.max_lat > 90):
            raise ValueError(f"Bbox không hợp lệ: ({config.min_lon}, {config.min_lat}, {config.max_lon}, {config.max_lat}). "
                           f"Phải trong phạm vi (-180, -90, 180, 90)")
        
        prediction_status["progress"] = "Đang load model..."
        
        # Load model using ModelManager
        model_manager = get_model_manager()
        model, label_encoder, model_metadata = model_manager.load_model(config.model_filename)
        
        # Get feature_mode and features from metadata (default to 'simple' if not specified)
        feature_mode = model_metadata.get("feature_mode", "simple")
        required_features = model_metadata.get("features", [])
        n_features_expected = model_metadata.get("n_features", len(required_features))

        prediction_status["progress"] = f"Model: {model_metadata.get('model_type', 'unknown')}, mode={feature_mode}, features={n_features_expected}"

        # Initialize FeatureExtractor với đúng mode như lúc training
        extractor = get_feature_extractor(mode=feature_mode)
        
        # Check if it's a PyTorch model (CNN, Swin-UNet, etc.)
        is_pytorch_model = hasattr(model, '__class__') and any(
            name in model.__class__.__name__ for name in ['CNN', 'SwinUNet']
        )
        if is_pytorch_model:
            model_class_name = model.__class__.__name__
            prediction_status["progress"] = f"Phát hiện PyTorch {model_class_name} model..."
            try:
                import torch
            except ImportError:
                raise ImportError(f"PyTorch required for {model_class_name} models. Install: pip install torch")
        
        # Initialize common variables
        bbox = [config.min_lon, config.min_lat, config.max_lon, config.max_lat]
        time_range = f"{config.start_date}/{config.end_date}"
        
        # ============ LOAD SENTINEL-2 DATA ============
        prediction_status["progress"] = "Đang kết nối Microsoft Planetary Computer..."
        import pystac_client
        import planetary_computer
        from odc.stac import load
        
        catalog = pystac_client.Client.open(
            "https://planetarycomputer.microsoft.com/api/stac/v1",
            modifier=planetary_computer.sign_inplace,
        )
        
        prediction_status["progress"] = "Đang tải dữ liệu Sentinel-2..."
        s2_search = catalog.search(
            collections=["sentinel-2-l2a"],
            bbox=bbox,
            datetime=time_range,
            query={"eo:cloud_cover": {"lt": config.cloud_cover}}
        )
        s2_items = list(s2_search.items())
        
        if not s2_items:
            raise ValueError("Không tìm thấy dữ liệu Sentinel-2 cho khu vực và thời gian này")
        
        s2_items = s2_items[:config.max_scenes]
        prediction_status["progress"] = f"Đang xử lý {len(s2_items)} scenes Sentinel-2..."
        
        # 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"]
        
        s2_data = load(
            s2_items,
            bbox=bbox,
            bands=bands_to_load,
            chunks={"time": 1, "x": 2048, "y": 2048},
            groupby="solar_day",
            resolution=config.resolution
        ).compute()
        
        prediction_status["progress"] = "Đã load Sentinel-2 data"
        
        # ============ LOAD SENTINEL-1 DATA (RADAR) ============
        prediction_status["progress"] = "Đang tải dữ liệu Sentinel-1 (Radar)..."
        use_radar = False
        vh_data = None
        vv_data = None
        
        try:
            s1_search = catalog.search(
                collections=["sentinel-1-rtc"],
                bbox=bbox,
                datetime=time_range,
            )
            s1_items = list(s1_search.items())
            
            if s1_items:
                s1_items = s1_items[:config.max_scenes]
                s1_data = load(
                    s1_items,
                    bbox=bbox,
                    bands=["vh", "vv"],
                    chunks={"time": 1, "x": 2048, "y": 2048},
                    groupby="solar_day",
                    resolution=config.resolution
                ).compute()
                
                # Convert to dB
                vh_data = 10 * np.log10(s1_data['vh'].where(s1_data['vh'] > 0))
                vv_data = 10 * np.log10(s1_data['vv'].where(s1_data['vv'] > 0))
                use_radar = True
                prediction_status["progress"] = f"Đã load Sentinel-1 data ({len(s1_items)} scenes)"
            else:
                prediction_status["progress"] = "Không có dữ liệu Sentinel-1, bỏ qua radar features"
        except Exception as e:
            prediction_status["progress"] = f"Lỗi load Sentinel-1: {str(e)}, bỏ qua radar features"
        
        # ============ APPLY CLOUD MASK ============
        prediction_status["progress"] = "Đang xử lý mây..."
        if "SCL" in s2_data:
            scl = s2_data["SCL"]
            # SCL values: 3=cloud shadow, 8=cloud medium, 9=cloud high, 10=cirrus
            cloud_mask = (scl == 3) | (scl == 8) | (scl == 9) | (scl == 10)
            for band in s2_data.data_vars:
                if band != "SCL":
                    s2_data[band] = s2_data[band].where(~cloud_mask)
        
        # ============ EXTRACT FEATURES ============
        prediction_status["progress"] = f"Đang trích xuất features (mode={feature_mode})..."
        

        # Always fill NaN for all bands in s2_data if present
        for band in ["B02", "B03", "B04", "B08", "B11"]:
            if band in s2_data:
                s2_data[band] = s2_data[band].ffill(dim='time').bfill(dim='time')

        # Calculate NDVI if needed (for simple mode)
        ndvi_filled = None
        if feature_mode == 'simple' and 'B08' in s2_data and 'B04' in s2_data:
            nir = s2_data["B08"].astype('float32')
            red = s2_data["B04"].astype('float32')
            ndvi = (nir - red) / (nir + red + 1e-8)
            ndvi_filled = ndvi.ffill(dim='time').bfill(dim='time')

        # Extract features using FeatureExtractor, always pass all possible data
        features = extractor.extract(
            s2_data=s2_data,
            ndvi_data=ndvi_filled,
            vh_data=vh_data,
            vv_data=vv_data
        )

        # Handle NaN values
        features = np.nan_to_num(features, nan=0.0)

        # Ensure features shape matches model expectation
        if features.shape[1] != n_features_expected:
            raise ValueError(f"Số lượng features ({features.shape[1]}) không khớp với model ({n_features_expected}). Hãy kiểm tra lại cấu hình trích xuất đặc trưng và metadata của model.")

        prediction_status["progress"] = f"Đã extract {features.shape[1]} features cho {features.shape[0]} pixels"
        
        # ============ PREDICT ============
        prediction_status["progress"] = "Đang dự đoán..."
        
        # Make prediction (all PyTorch models have the same predict interface)
        predictions = model.predict(features)
        
        # Decode labels if label_encoder exists
        if label_encoder is not None:
            try:
                predictions = label_encoder.inverse_transform(predictions.astype(int))
            except:
                pass
        
        # Reshape to original shape
        if feature_mode == 'simple' and 'B08' in s2_data:
            # Use B08 to get shape
            y_size = len(s2_data.y)
            x_size = len(s2_data.x)
        else:
            y_size = len(s2_data.y)
            x_size = len(s2_data.x)
        
        pred_shape = (y_size, x_size)
        predictions_2d = predictions.reshape(pred_shape)
        
        # ============ CREATE OUTPUT ============
        prediction_status["progress"] = "Đang tạo bản đồ phân loại..."
        
        # Create output xarray
        prediction_da = xr.DataArray(
            predictions_2d,
            coords={
                "y": s2_data.y,
                "x": s2_data.x
            },
            dims=["y", "x"],
            name="classification"
        )
        
        # Save output
        output_dir = Path("predictions")
        output_dir.mkdir(exist_ok=True)
        
        timestamp = dt.now().strftime("%Y%m%d_%H%M%S")
        output_file = output_dir / f"prediction_{timestamp}.tif"
        
        prediction_status["progress"] = "Đang lưu kết quả GeoTIFF..."
        
        # Set CRS and save as GeoTIFF
        if hasattr(s2_data, 'rio') and s2_data.rio.crs is not None:
            prediction_da.rio.write_crs(s2_data.rio.crs, inplace=True)
        else:
            prediction_da.rio.write_crs("EPSG:4326", inplace=True)
        
        prediction_da.rio.to_raster(str(output_file), driver="GTiff")
        
        # Generate PNG preview for web display
        prediction_status["progress"] = "Đang tạo PNG preview..."
        png_file = output_dir / f"prediction_{timestamp}.png"
        try:
            import matplotlib
            matplotlib.use('Agg')  # Non-interactive backend
            import matplotlib.pyplot as plt
            
            # Create a figure with prediction result
            fig, ax = plt.subplots(figsize=(12, 10), dpi=150)
            
            # Plot prediction with colormap
            im = ax.imshow(predictions_2d, cmap='tab20', interpolation='nearest')
            ax.set_title(f'Prediction Result - {timestamp}', fontsize=14, fontweight='bold')
            ax.set_xlabel('X (pixels)', fontsize=10)
            ax.set_ylabel('Y (pixels)', fontsize=10)

            # Build mapping from numeric class value -> display label
            class_map = None
            try:
                # 1) Try label_encoder (preferred)
                if label_encoder is not None:
                    try:
                        # label_encoder.classes_ may be strings or numbers
                        le_classes = list(label_encoder.classes_)
                        # If classes are strings like names, we'll map indices -> names
                        if all(isinstance(x, str) for x in le_classes):
                            class_map = {i: name for i, name in enumerate(le_classes)}
                        else:
                            # If classes are numeric labels matching values, map value->str(value)
                            class_map = {int(v): str(v) for v in le_classes}
                    except Exception:
                        class_map = None
            except Exception:
                class_map = None

            # 2) Try model metadata 'class_names' (list ordered by class code)
            if class_map is None and isinstance(model_metadata, dict):
                try:
                    cn = model_metadata.get('class_names')
                    if isinstance(cn, list):
                        class_map = {i: str(name) for i, name in enumerate(cn)}
                except Exception:
                    pass

            # 3) Try invert label_mapping in metadata if exists (name->code)
            if class_map is None and isinstance(model_metadata, dict):
                try:
                    lm = model_metadata.get('label_mapping') or model_metadata.get('labels')
                    if isinstance(lm, dict):
                        # invert mapping: code -> name
                        inv = {}
                        for k, v in lm.items():
                            try:
                                key_int = int(v)
                            except Exception:
                                continue
                            inv[key_int] = str(k)
                        if inv:
                            class_map = inv
                except Exception:
                    pass

            # Create colorbar
            cbar = plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
            cbar.set_label('Class', rotation=270, labelpad=15)

            # If we have a class_map, set ticks and labels
            try:
                if class_map:
                    vals = np.array(sorted(class_map.keys()))
                    cbar.set_ticks(vals)
                    cbar.set_ticklabels([class_map[int(v)] for v in vals])
                else:
                    # fallback: label numeric ticks from min..max
                    if np.issubdtype(predictions_2d.dtype, np.number):
                        minv = int(np.nanmin(predictions_2d))
                        maxv = int(np.nanmax(predictions_2d))
                        ticks = np.arange(minv, maxv + 1)
                        cbar.set_ticks(ticks)
                        cbar.set_ticklabels([str(t) for t in ticks])
            except Exception:
                pass
            
            # Add grid
            ax.grid(True, alpha=0.3, linestyle='--', linewidth=0.5)
            
            # Save PNG
            plt.tight_layout()
            plt.savefig(str(png_file), dpi=150, bbox_inches='tight')
            plt.close(fig)
            
            print(f"[PNG PREVIEW] Created: {png_file}")
        except Exception as e:
            print(f"[PNG PREVIEW ERROR] Failed to create PNG: {e}")
            png_file = None
        
        # Get unique classes for result
        unique_classes = np.unique(predictions_2d)
        unique_classes = unique_classes[~np.isnan(unique_classes)].tolist()
        
        prediction_status["is_predicting"] = False
        prediction_status["progress"] = "Hoàn thành! Đang tạo báo cáo..."
        prediction_status["output_file"] = str(output_file)
        prediction_status["result"] = {
            "output_file": str(output_file),
            "png_file": str(png_file) if png_file else None,
            "shape": list(pred_shape),
            "unique_classes": unique_classes,
            "bbox": bbox,
            "time_range": time_range,
            "n_features": features.shape[1],
            "feature_mode": feature_mode,
            "used_radar": use_radar,
            "model_used": config.model_filename
        }
        
        # Auto generate prediction report
        try:
            report_path, _ = generate_prediction_report(prediction_status["result"])
            prediction_status["result"]["report_path"] = report_path
            prediction_status["result"]["report_filename"] = Path(report_path).name
            prediction_status["progress"] = "Hoàn thành! Báo cáo đã được tạo."
            print(f"[PREDICTION REPORT] Generated: {report_path}")
        except Exception as e:
            print(f"[PREDICTION REPORT ERROR] Failed to generate report: {e}")
            prediction_status["progress"] = "Hoàn thành! (Không thể tạo báo cáo)"
        
        prediction_status["end_time"] = dt.now().isoformat()
        
    except Exception as e:
        prediction_status["is_predicting"] = False
        prediction_status["error"] = str(e)
        prediction_status["progress"] = f"Lỗi: {str(e)}"
        prediction_status["end_time"] = dt.now().isoformat()
        import traceback
        print(traceback.format_exc())


@app.get("/api/predictions/list")
async def list_predictions():
    """Lấy danh sách các file prediction đã tạo"""
    predictions_dir = Path("predictions")
    predictions_dir.mkdir(exist_ok=True)
    
    predictions = []
    for pred_file in predictions_dir.glob("*.tif"):
        pred_info = {
            "filename": pred_file.name,
            "created": datetime.fromtimestamp(pred_file.stat().st_mtime).isoformat(),
            "size_mb": round(pred_file.stat().st_size / 1024 / 1024, 2),
            "download_url": f"/api/predictions/download/{pred_file.name}",
            "is_batch_job": pred_file.name.startswith("batch_"),
            "batch_metadata": None
        }
        
        # Try to load batch metadata from JSON sidecar if exists
        json_file = pred_file.with_suffix('.json')
        if json_file.exists():
            try:
                import json
                with open(json_file, 'r') as f:
                    metadata = json.load(f)
                    pred_info["batch_metadata"] = {
                        "batch_job_id": metadata.get("batch_job_id"),
                        "batch_name": metadata.get("batch_name"),
                        "batch_timestamp": metadata.get("batch_timestamp")
                    }
            except Exception as e:
                print(f"[METADATA ERROR] Failed to load {json_file}: {e}")
        
        # Check PNG preview
        png_file = pred_file.with_suffix('.png')
        pred_info["has_preview"] = png_file.exists()
        if png_file.exists():
            pred_info["preview_url"] = f"/api/predictions/preview/{png_file.name}"
        
        predictions.append(pred_info)
    
    # Sort by creation time (newest first)
    predictions.sort(key=lambda x: x["created"], reverse=True)
    return {"predictions": predictions}


@app.get("/api/predictions/download/{filename}")
async def download_prediction(filename: str):
    """Download file prediction GeoTIFF"""
    predictions_dir = Path("predictions")
    file_path = predictions_dir / filename
    
    # Security check: ensure filename doesn't contain path traversal
    if ".." in filename or "/" in filename or "\\" in filename:
        raise HTTPException(status_code=400, detail="Invalid filename")
    
    if not file_path.exists():
        raise HTTPException(status_code=404, detail=f"File không tồn tại: {filename}")
    
    return FileResponse(
        path=str(file_path),
        filename=filename,
        media_type="image/tiff",
        headers={
            "Content-Disposition": f"attachment; filename={filename}"
        }
    )


@app.get("/api/predictions/preview/{filename}")
async def preview_prediction_png(filename: str):
    """Preview PNG image of prediction"""
    predictions_dir = Path("predictions")
    file_path = predictions_dir / filename
    
    # Security check
    if ".." in filename or "/" in filename or "\\" in filename:
        raise HTTPException(status_code=400, detail="Invalid filename")
    
    if not file_path.exists():
        raise HTTPException(status_code=404, detail=f"PNG preview không tồn tại: {filename}")
    
    return FileResponse(
        path=str(file_path),
        media_type="image/png"
    )


@app.get("/api/predictions/preview/{filename}")
async def preview_prediction_png(filename: str):
    """Preview PNG image of prediction"""
    predictions_dir = Path("predictions")
    file_path = predictions_dir / filename
    
    # Security check
    if ".." in filename or "/" in filename or "\\" in filename:
        raise HTTPException(status_code=400, detail="Invalid filename")
    
    if not file_path.exists():
        raise HTTPException(status_code=404, detail=f"PNG preview không tồn tại: {filename}")
    
    return FileResponse(
        path=str(file_path),
        media_type="image/png"
    )


# ============ DASHBOARD & VISUALIZATION API ============

@app.get("/api/dashboard/accuracy-trends")
async def get_accuracy_trends():
    """Lấy dữ liệu accuracy trends của các models theo thời gian"""
    model_dir = Path("model_train")
    if not model_dir.exists():
        return {"trends": [], "models": []}
    
    trends_data = []
    for info_file in sorted(model_dir.glob("*.json")):
        try:
            with open(info_file) as f:
                info = json.load(f)
            
            # Extract relevant data
            if "training_date" in info and "metrics" in info:
                trends_data.append({
                    "date": info["training_date"],
                    "model_name": info.get("model_type", "unknown"),
                    "accuracy": info["metrics"].get("accuracy", 0),
                    "f1_score": info["metrics"].get("macro avg", {}).get("f1-score", 0),
                    "precision": info["metrics"].get("macro avg", {}).get("precision", 0),
                    "recall": info["metrics"].get("macro avg", {}).get("recall", 0),
                    "filename": info_file.stem + ".joblib"
                })
        except Exception as e:
            print(f"Error loading {info_file}: {e}")
            continue
    
    # Sort by date
    trends_data.sort(key=lambda x: x["date"])
    
    return {
        "trends": trends_data,
        "models": list(set(d["model_name"] for d in trends_data))
    }


@app.get("/api/dashboard/statistics")
async def get_statistics():
    """Lấy thống kê tổng quan: số models, predictions, reports"""
    model_dir = Path("model_train")
    predictions_dir = Path("predictions")
    reports_dir = Path("reports")
    
    # Count items
    n_models = len(list(model_dir.glob("*.joblib"))) if model_dir.exists() else 0
    n_predictions = len(list(predictions_dir.glob("*.tif"))) if predictions_dir.exists() else 0
    n_reports = len(list(reports_dir.glob("*.html"))) if reports_dir.exists() else 0
    
    # Get latest model info
    latest_model = None
    if model_dir.exists():
        model_files = sorted(model_dir.glob("*.json"), key=lambda x: x.stat().st_mtime, reverse=True)
        if model_files:
            try:
                with open(model_files[0]) as f:
                    latest_model = json.load(f)
            except:
                pass
    
    # Get latest prediction
    latest_prediction = None
    if predictions_dir.exists():
        pred_files = sorted(predictions_dir.glob("*.tif"), key=lambda x: x.stat().st_mtime, reverse=True)
        if pred_files:
            latest_prediction = {
                "filename": pred_files[0].name,
                "created": datetime.fromtimestamp(pred_files[0].stat().st_mtime).isoformat(),
                "size_mb": round(pred_files[0].stat().st_size / 1024 / 1024, 2)
            }
    
    return {
        "models": {
            "total": n_models,
            "latest": latest_model
        },
        "predictions": {
            "total": n_predictions,
            "latest": latest_prediction
        },
        "reports": {
            "total": n_reports
        },
        "training_status": training_status,
        "prediction_status": prediction_status
    }


@app.get("/api/dashboard/class-distribution/{model_filename}")
async def get_class_distribution(model_filename: str):
    """Lấy phân bố các lớp từ model info"""
    # Convert model filename to info filename
    # e.g., model_cnn_20251221_163841.joblib -> model_cnn_20251221_163841_info.json
    base_name = model_filename.replace(".joblib", "")
    info_file = Path("model_train") / f"{base_name}_info.json"
    
    if not info_file.exists():
        raise HTTPException(status_code=404, detail="Model info không tồn tại")
    
    with open(info_file) as f:
        info = json.load(f)
    
    # Extract class distribution from classification report
    class_dist = {}
    if "classification_report" in info:
        for class_name, metrics in info["classification_report"].items():
            if isinstance(metrics, dict) and "support" in metrics:
                class_dist[class_name] = int(metrics["support"])
    
    return {
        "model": model_filename,
        "class_distribution": class_dist,
        "total_samples": sum(class_dist.values()) if class_dist else 0
    }


# ============ BATCH PROCESSING API ============

class BatchPredictionItem(BaseModel):
    """Một item trong batch prediction"""
    name: str
    min_lon: float
    min_lat: float
    max_lon: float
    max_lat: float
    start_date: str = "2023-03-01"
    end_date: str = "2023-05-31"
    max_scenes: int = 12
    cloud_cover: int = 30
    resolution: int = 20


class BatchPredictionConfig(BaseModel):
    """Cấu hình cho batch prediction"""
    model_filename: str
    items: List[BatchPredictionItem]
    auto_retry: bool = True
    max_retries: int = 3


@app.post("/api/batch/start")
async def start_batch_prediction(config: BatchPredictionConfig, background_tasks: BackgroundTasks):
    """Bắt đầu batch prediction"""
    global batch_queue, batch_results
    
    # Create batch jobs
    batch_id = datetime.now().strftime("%Y%m%d_%H%M%S")
    
    for idx, item in enumerate(config.items):
        job = {
            "batch_id": batch_id,
            "job_id": f"{batch_id}_{idx}",
            "name": item.name,
            "status": "queued",
            "progress": 0,
            "error": None,
            "result": None,
            "retries": 0,
            "max_retries": config.max_retries if config.auto_retry else 0,
            "config": {
                "model_filename": config.model_filename,
                "min_lon": item.min_lon,
                "min_lat": item.min_lat,
                "max_lon": item.max_lon,
                "max_lat": item.max_lat,
                "start_date": item.start_date,
                "end_date": item.end_date,
                "max_scenes": item.max_scenes,
                "cloud_cover": item.cloud_cover,
                "resolution": item.resolution
            },
            "created_at": datetime.now().isoformat()
        }
        batch_queue.append(job)
    
    # Start processing in background
    background_tasks.add_task(process_batch_queue)
    
    return {
        "message": f"Đã tạo {len(config.items)} batch jobs",
        "batch_id": batch_id,
        "total_jobs": len(config.items)
    }


@app.get("/api/batch/status")
async def get_batch_status():
    """Lấy trạng thái của batch queue"""
    global batch_queue, batch_results
    
    queued = [j for j in batch_queue if j["status"] == "queued"]
    running = [j for j in batch_queue if j["status"] == "running"]
    completed = [j for j in batch_results if j["status"] == "completed"]
    failed = [j for j in batch_results if j["status"] == "failed"]
    
    return {
        "queue": {
            "queued": len(queued),
            "running": len(running),
            "completed": len(completed),
            "failed": len(failed),
            "total": len(batch_queue) + len(batch_results)
        },
        "jobs": {
            "queued": queued[:5],  # Show first 5
            "running": running,
            "recent_completed": completed[:10],  # Show last 10
            "recent_failed": failed[:10]
        }
    }


@app.get("/api/batch/results/{batch_id}")
async def get_batch_results(batch_id: str):
    """Lấy kết quả của một batch"""
    global batch_results
    
    results = [j for j in batch_results if j["batch_id"] == batch_id]
    
    if not results:
        # Check if still in queue
        queued = [j for j in batch_queue if j["batch_id"] == batch_id]
        if queued:
            return {
                "batch_id": batch_id,
                "status": "processing",
                "jobs": queued
            }
        else:
            raise HTTPException(status_code=404, detail="Batch không tồn tại")
    
    return {
        "batch_id": batch_id,
        "status": "completed",
        "jobs": results,
        "summary": {
            "total": len(results),
            "successful": len([j for j in results if j["status"] == "completed"]),
            "failed": len([j for j in results if j["status"] == "failed"])
        }
    }


@app.post("/api/batch/cancel/{batch_id}")
async def cancel_batch(batch_id: str):
    """Hủy một batch đang chạy"""
    global batch_queue
    
    # Remove from queue
    removed = 0
    batch_queue_copy = batch_queue.copy()
    for job in batch_queue_copy:
        if job["batch_id"] == batch_id and job["status"] == "queued":
            batch_queue.remove(job)
            removed += 1
    
    return {
        "message": f"Đã hủy {removed} jobs",
        "batch_id": batch_id
    }


async def process_batch_queue():
    """Process batch prediction queue"""
    global batch_queue, batch_results
    
    import asyncio
    
    while batch_queue:
        # Get next job
        job = None
        for j in batch_queue:
            if j["status"] == "queued":
                job = j
                break
        
        if not job:
            break
        
        # Mark as running
        job["status"] = "running"
        job["progress"] = 0
        job["started_at"] = datetime.now().isoformat()
        
        try:
            # Create PredictionConfig from job config
            pred_config = PredictionConfig(**job["config"])
            
            print(f"[BATCH] Processing job {job['job_id']}: {job['name']}")
            job["progress"] = 5
            
            # Run prediction synchronously (in the same thread to avoid conflicts)
            await asyncio.to_thread(run_batch_prediction, job, pred_config)
            
            # Check if prediction was successful
            if job.get("result") and not job.get("error"):
                job["status"] = "completed"
                job["progress"] = 100
                job["completed_at"] = datetime.now().isoformat()
                print(f"[BATCH] Job {job['job_id']} completed successfully")
            else:
                raise Exception(job.get("error", "Unknown error during prediction"))
            
        except Exception as e:
            job["error"] = str(e)
            
            # Retry logic
            if job["retries"] < job["max_retries"]:
                job["retries"] += 1
                job["status"] = "queued"  # Retry
                job["progress"] = 0
                print(f"[BATCH] Job {job['job_id']} ({job['name']}) failed, retrying ({job['retries']}/{job['max_retries']}): {e}")
                continue
            else:
                job["status"] = "failed"
                job["progress"] = 0
                job["completed_at"] = datetime.now().isoformat()
                print(f"[BATCH] Job {job['job_id']} ({job['name']}) failed permanently: {e}")
        
        # Move to results
        batch_queue.remove(job)
        batch_results.append(job)
        
        # Keep only last 100 results
        if len(batch_results) > 100:
            batch_results = batch_results[-100:]


def run_batch_prediction(job: dict, config: PredictionConfig):
    """Run prediction for a single batch job"""
    try:
        job["progress"] = 10
        
        # Import required libraries
        import xarray as xr
        import numpy as np
        from datetime import datetime as dt
        import rioxarray
        import dask.array as da
        
        job["progress"] = 15
        
        # Load model using ModelManager
        model_manager = get_model_manager()
        model, label_encoder, model_metadata = model_manager.load_model(config.model_filename)
        
        # Check if it's a PyTorch model (CNN, Swin-UNet, etc.)
        is_pytorch_model = hasattr(model, '__class__') and any(
            name in model.__class__.__name__ for name in ['CNN', 'SwinUNet']
        )
        
        job["progress"] = 20
        
        # Load data from Microsoft Planetary Computer
        import pystac_client
        import planetary_computer
        from odc.stac import load
        
        catalog = pystac_client.Client.open(
            "https://planetarycomputer.microsoft.com/api/stac/v1",
            modifier=planetary_computer.sign_inplace,
        )
        
        bbox = [config.min_lon, config.min_lat, config.max_lon, config.max_lat]
        time_range = f"{config.start_date}/{config.end_date}"
        
        job["progress"] = 25
        
        # Search Sentinel-2
        s2_search = catalog.search(
            collections=["sentinel-2-l2a"],
            bbox=bbox,
            datetime=time_range,
            query={"eo:cloud_cover": {"lt": config.cloud_cover}}
        )
        
        s2_items = list(s2_search.items())
        if not s2_items:
            raise ValueError("Không tìm thấy dữ liệu Sentinel-2")
        
        s2_items = s2_items[:config.max_scenes]
        
        job["progress"] = 35
        
        # Load Sentinel-2 data
        s2_data = load(
            s2_items,
            bbox=bbox,
            chunks={"time": 1, "x": 2048, "y": 2048},
            groupby="solar_day",
            resolution=config.resolution
        )
        
        job["progress"] = 50
        
        # Calculate NDVI
        nir = s2_data["B08"].astype('float32')
        red = s2_data["B04"].astype('float32')
        ndvi = (nir - red) / (nir + red + 1e-8)
        
        # Mask clouds if SCL available
        if "SCL" in s2_data:
            scl = s2_data["SCL"]
            cloud_mask = (scl == 3) | (scl == 8) | (scl == 9) | (scl == 10)
            ndvi = ndvi.where(~cloud_mask)
        
        # Fill NaN and resample
        ndvi_filled = ndvi.ffill(dim='time').bfill(dim='time')
        ndvi_monthly = ndvi_filled.resample(time="1ME").mean().compute()
        
        job["progress"] = 70
        
        # Prepare features
        n_times_ndvi = len(ndvi_monthly.time)
        y_size = len(ndvi_monthly.y)
        x_size = len(ndvi_monthly.x)
        n_pixels = y_size * x_size
        
        ndvi_features = []
        for t in range(n_times_ndvi):
            ndvi_t = ndvi_monthly.isel(time=t).values.flatten()
            ndvi_features.append(ndvi_t)
        
        features = np.column_stack(ndvi_features)
        features = np.nan_to_num(features, nan=0.0)
        
        job["progress"] = 80
        
        # Adjust features to match model expectations
        try:
            if is_pytorch_model:
                expected_features = model.n_features
            elif hasattr(model, 'n_features_in_'):
                expected_features = model.n_features_in_
            else:
                try:
                    expected_features = model.get_booster().num_features()
                except:
                    expected_features = features.shape[1]
            
            if features.shape[1] > expected_features:
                features = features[:, :expected_features]
            elif features.shape[1] < expected_features:
                n_missing = expected_features - features.shape[1]
                padding = np.tile(features[:, -1:], (1, n_missing))
                features = np.column_stack([features, padding])
        except:
            pass
        
        # Predict (all models have same predict interface)
        predictions = model.predict(features)
        
        # Decode labels
        if label_encoder is not None:
            try:
                predictions = label_encoder.inverse_transform(predictions)
            except:
                pass
        
        job["progress"] = 90
        
        # Reshape and create output
        pred_shape = (y_size, x_size)
        predictions_2d = predictions.reshape(pred_shape)
        
        prediction_da = xr.DataArray(
            predictions_2d,
            coords={"y": ndvi_monthly.y, "x": ndvi_monthly.x},
            dims=["y", "x"],
            name="classification"
        )
        
        # Save output
        output_dir = Path("predictions")
        output_dir.mkdir(exist_ok=True)
        
        output_file = output_dir / f"batch_{job['job_id']}_{job['name'].replace(' ', '_')}.tif"
        
        if hasattr(s2_data, 'rio') and s2_data.rio.crs is not None:
            prediction_da.rio.write_crs(s2_data.rio.crs, inplace=True)
        else:
            prediction_da.rio.write_crs("EPSG:4326", inplace=True)
        
        prediction_da.rio.to_raster(str(output_file), driver="GTiff")
        
        # Generate PNG preview
        png_file = output_dir / f"batch_{job['job_id']}_{job['name'].replace(' ', '_')}.png"
        try:
            import matplotlib
            matplotlib.use('Agg')
            import matplotlib.pyplot as plt
            
            fig, ax = plt.subplots(figsize=(12, 10), dpi=150)
            im = ax.imshow(predictions_2d, cmap='tab20', interpolation='nearest')
            ax.set_title(f"{job['name']} - Batch {job['job_id']}", fontsize=14, fontweight='bold')
            ax.set_xlabel('X (pixels)', fontsize=10)
            ax.set_ylabel('Y (pixels)', fontsize=10)
            
            cbar = plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
            cbar.set_label('Class', rotation=270, labelpad=15)
            try:
                # Build mapping from numeric class value -> label (reuse logic from above)
                class_map = None
                if label_encoder is not None:
                    try:
                        le_classes = list(label_encoder.classes_)
                        if all(isinstance(x, str) for x in le_classes):
                            class_map = {i: name for i, name in enumerate(le_classes)}
                        else:
                            class_map = {int(v): str(v) for v in le_classes}
                    except Exception:
                        class_map = None

                if class_map is None and isinstance(model_metadata, dict):
                    cn = model_metadata.get('class_names')
                    if isinstance(cn, list):
                        class_map = {i: str(name) for i, name in enumerate(cn)}

                if class_map is None and isinstance(model_metadata, dict):
                    lm = model_metadata.get('label_mapping') or model_metadata.get('labels')
                    if isinstance(lm, dict):
                        inv = {}
                        for k, v in lm.items():
                            try:
                                key_int = int(v)
                            except Exception:
                                continue
                            inv[key_int] = str(k)
                        if inv:
                            class_map = inv

                if class_map:
                    vals = np.array(sorted(class_map.keys()))
                    cbar.set_ticks(vals)
                    cbar.set_ticklabels([class_map[int(v)] for v in vals])
                else:
                    if np.issubdtype(predictions_2d.dtype, np.number):
                        minv = int(np.nanmin(predictions_2d))
                        maxv = int(np.nanmax(predictions_2d))
                        ticks = np.arange(minv, maxv + 1)
                        cbar.set_ticks(ticks)
                        cbar.set_ticklabels([str(t) for t in ticks])
            except Exception:
                pass
            ax.grid(True, alpha=0.3, linestyle='--', linewidth=0.5)
            
            plt.tight_layout()
            plt.savefig(str(png_file), dpi=150, bbox_inches='tight')
            plt.close(fig)
        except Exception as e:
            print(f"[BATCH PNG ERROR] {e}")
            png_file = None
        
        # Get unique classes
        unique_classes = np.unique(predictions_2d)
        unique_classes = unique_classes[~np.isnan(unique_classes)].tolist()
        
        # Store result in job with batch metadata
        job["result"] = {
            "output_file": str(output_file),
            "png_file": str(png_file) if png_file else None,
            "shape": list(pred_shape),
            "unique_classes": unique_classes,
            "bbox": bbox,
            "time_range": time_range,
            "n_features": features.shape[1],
            "n_times_ndvi": n_times_ndvi,
            "model_used": config.model_filename,
            "batch_job_id": job["job_id"],
            "batch_name": job["name"],
            "batch_timestamp": datetime.now().isoformat()
        }
        
        # Save batch metadata to JSON sidecar file for persistence
        metadata_file = output_file.with_suffix('.json')
        try:
            import json
            with open(metadata_file, 'w') as f:
                json.dump(job["result"], f, indent=2, default=str)
            print(f"[BATCH METADATA] Saved to {metadata_file}")
        except Exception as e:
            print(f"[BATCH METADATA ERROR] Failed to save metadata: {e}")
        
        # Auto generate prediction report for batch job
        try:
            from report_generator import generate_prediction_report
            report_path, _ = generate_prediction_report(job["result"])
            job["result"]["report_path"] = report_path
            job["result"]["report_filename"] = Path(report_path).name
            print(f"[BATCH REPORT] Generated prediction report: {report_path}")
        except Exception as e:
            print(f"[BATCH REPORT ERROR] Failed to generate report: {e}")
        
        job["progress"] = 100
        
    except Exception as e:
        job["error"] = str(e)
        import traceback
        print(f"[BATCH ERROR] Job {job['job_id']}: {traceback.format_exc()}")


# ============ CHANGE DETECTION API ============

def rasterize_ground_truth(shapefile_path, out_shape, bbox, class_column="class"):
    """Rasterize ground truth shapefile to match prediction raster shape."""
    try:
        import geopandas as gpd
        from rasterio import features as rio_features
        
        gdf = gpd.read_file(shapefile_path)
        minx, miny, maxx, maxy = bbox
        
        # Crop to bbox
        gdf = gdf.cx[minx:maxx, miny:maxy]
        
        if class_column not in gdf.columns:
            raise ValueError(f"Shapefile missing '{class_column}' column. Available: {list(gdf.columns)}")
        
        # Create transform for rasterization
        transform = from_bounds(minx, miny, maxx, maxy, out_shape[1], out_shape[0])
        
        # Prepare geometries and values for rasterization
        shapes = zip(gdf.geometry, gdf[class_column])
        
        # Rasterize
        gt_raster = rio_features.rasterize(
            shapes,
            out_shape=out_shape,
            fill=-1,
            transform=transform,
            dtype="int16"
        )
        
        return gt_raster
    except Exception as e:
        print(f"[RASTERIZE ERROR] {e}")
        raise


@app.post("/api/change-detection/predict")
async def change_detection_predict_workflow(
    model_filename: str,
    min_lon: float,
    min_lat: float,
    max_lon: float,
    max_lat: float,
    start_date: str,
    end_date: str,
    max_scenes: int = 12,
    cloud_cover: int = 30,
    resolution: int = 20
):
    """
    Complete workflow: Predict + Compare with Ground Truth
    1. Load Sentinel-2 data for bbox and date range
    2. Run prediction using trained model
    3. Rasterize ground truth from training shapefile
    4. Compare and generate change detection results
    """
    try:
        # --- STEP 1: LOAD MODEL ---
        model_manager = get_model_manager()
        model, label_encoder, model_metadata = model_manager.load_model(model_filename)
        
        print(f"[CHANGE DETECTION] Loaded model: {model_filename}")
        print(f"  - Type: {model_metadata.get('model_type', 'unknown')}")
        print(f"  - Features: {model_metadata.get('features', [])}")
        
        # --- STEP 2: LOAD SENTINEL-2 DATA ---
        bbox = [min_lon, min_lat, max_lon, max_lat]
        time_range = f"{start_date}/{end_date}"
        
        catalog = Client.open(
            "https://planetarycomputer.microsoft.com/api/stac/v1",
            modifier=planetary_computer.sign_inplace
        )
        
        search = catalog.search(
            collections=["sentinel-2-l2a"],
            bbox=bbox,
            datetime=time_range,
            query={"eo:cloud_cover": {"lt": cloud_cover}}
        )
        
        items = list(search.items())[:max_scenes]
        print(f"[CHANGE DETECTION] Found {len(items)} Sentinel-2 scenes")
        
        if len(items) == 0:
            raise HTTPException(status_code=404, detail="No Sentinel-2 data found for the given area and date range")
        
        # Load data
        signed_items = [planetary_computer.sign(item) for item in items]
        data = odc.stac.load(
            signed_items,
            bbox=bbox,
            bands=["B02", "B03", "B04", "B08"],
            resolution=resolution,
            chunks={"x": 2048, "y": 2048}
        ).compute()
        
        # --- STEP 3: CALCULATE NDVI ---
        print("[CHANGE DETECTION] Calculating NDVI...")
        nir = data["B08"].astype('float32')
        red = data["B04"].astype('float32')
        ndvi = (nir - red) / (nir + red + 1e-8)
        
        # Handle clouds if SCL available
        if "SCL" in data:
            scl = data["SCL"]
            cloud_mask = (scl == 3) | (scl == 8) | (scl == 9) | (scl == 10)
            ndvi = ndvi.where(~cloud_mask)
        
        # --- STEP 4: PREPARE FEATURES ---
        ndvi_filled = ndvi.ffill(dim='time').bfill(dim='time')
        ndvi_mean = np.nanmean(ndvi_filled.values, axis=0)
        
        height, width = ndvi_mean.shape
        n_pixels = height * width
        
        # Prepare features
        features = ndvi_mean.flatten().reshape(-1, 1)
        valid_mask = ~np.isnan(features[:, 0])
        features_clean = features[valid_mask]
        
        # --- STEP 5: PREDICT ---
        print("[CHANGE DETECTION] Running prediction...")
        predictions = model.predict(features_clean)
        
        # Decode labels if needed
        if label_encoder is not None:
            try:
                predictions = label_encoder.inverse_transform(predictions)
            except:
                pass
        
        # Reshape to raster
        prediction_raster = np.full(n_pixels, -1, dtype=np.int16)
        prediction_raster[valid_mask] = predictions.astype(np.int16)
        prediction_raster = prediction_raster.reshape(height, width)
        
        # --- STEP 6: COMPARE WITH GROUND TRUTH ---
        print("[CHANGE DETECTION] Comparing with ground truth...")
        gt_shapefile = "train/ST_training data_updated_1130points_new.shp"
        gt_raster = rasterize_ground_truth(gt_shapefile, (height, width), bbox, class_column="class")
        
        # Calculate changes
        mask_valid = (gt_raster >= 0) & (prediction_raster >= 0)
        changes = gt_raster[mask_valid] != prediction_raster[mask_valid]
        n_total = np.count_nonzero(mask_valid)
        n_changed = np.count_nonzero(changes)
        
        # Create change matrix
        from collections import Counter
        change_pairs = list(zip(gt_raster[mask_valid][changes], prediction_raster[mask_valid][changes]))
        change_counter = Counter(change_pairs)
        change_matrix = {f"{int(gt)}->{int(pred)}": int(cnt) for (gt, pred), cnt in change_counter.items()}
        
        # Create change map
        change_map = np.full((height, width), -1, dtype=np.int8)
        change_map[mask_valid] = changes.astype(np.int8)
        
        # --- STEP 7: SAVE RESULTS ---
        output_dir = Path("predictions")
        output_dir.mkdir(exist_ok=True)
        
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        change_file = output_dir / f"change_map_{timestamp}.tif"
        
        transform = from_bounds(min_lon, min_lat, max_lon, max_lat, width, height)
        
        with rasterio.open(
            change_file, 'w',
            driver='GTiff',
            height=height,
            width=width,
            count=1,
            dtype=change_map.dtype,
            crs='EPSG:4326',
            transform=transform
        ) as dst:
            dst.write(change_map, 1)
        
        print(f"[CHANGE DETECTION] Saved change map to {change_file}")
        
        # --- RETURN RESULTS ---
        return {
            "success": True,
            "n_scenes": len(items),
            "ndvi_stats": {
                "mean": float(np.nanmean(ndvi_mean)),
                "min": float(np.nanmin(ndvi_mean)),
                "max": float(np.nanmax(ndvi_mean)),
                "std": float(np.nanstd(ndvi_mean))
            },
            "class_distribution": {
                int(cls): int(count) 
                for cls, count in zip(*np.unique(predictions, return_counts=True))
            },
            "change_detection": {
                "n_total_pixels": int(n_total),
                "n_changed_pixels": int(n_changed),
                "change_rate": float(n_changed) / n_total if n_total > 0 else 0.0,
                "change_matrix": change_matrix,
                "message": f"Detected {n_changed} changes out of {n_total} valid pixels ({(n_changed/n_total*100):.1f}%)" if n_total > 0 else "No valid pixels for comparison"
            },
            "change_map_file": str(change_file),
            "timestamp": timestamp
        }
        
    except HTTPException:
        raise
    except Exception as e:
        print(f"[CHANGE DETECTION ERROR] {e}")
        import traceback
        traceback.print_exc()
        raise HTTPException(status_code=500, detail=f"Change detection workflow failed: {str(e)}")


@app.post("/api/change-detection/workflow")
async def change_detection_workflow(request: ChangeDetectionWorkflowRequest):
    """Workflow: compare prediction with ground truth training data."""
    from collections import Counter
    
    try:
        prediction_result = request.prediction_result
        bbox = request.bbox
        
        if not prediction_result or "output_files" not in prediction_result:
            raise ValueError("Invalid prediction result")
        
        # Get classification raster from prediction
        class_file = None
        for f in prediction_result.get("output_files", []):
            if f.get("type") == "classification":
                class_file = f.get("path")
                break
        
        if not class_file:
            raise ValueError("No classification raster in prediction result")
        
        # Load prediction raster
        with rasterio.open(class_file) as pred_ds:
            pred_arr = pred_ds.read(1)
            pred_crs = pred_ds.crs
            pred_transform = pred_ds.transform
        
        # Rasterize ground truth training data
        gt_shapefile = "train/ST_training data_updated_1130points_new.shp"
        gt_raster = rasterize_ground_truth(gt_shapefile, pred_arr.shape, bbox, class_column="class")
        
        # Calculate change detection
        mask_valid = (gt_raster >= 0) & (pred_arr >= 0) & ~np.isnan(gt_raster) & ~np.isnan(pred_arr)
        changes = gt_raster[mask_valid] != pred_arr[mask_valid]
        n_total = np.count_nonzero(mask_valid)
        n_changed = np.count_nonzero(changes)
        
        # Create change pairs matrix
        change_pairs = list(zip(gt_raster[mask_valid][changes], pred_arr[mask_valid][changes]))
        change_counter = Counter(change_pairs)
        change_matrix = {f"{int(gt)}->{int(pred)}": int(cnt) for (gt, pred), cnt in change_counter.items()}
        
        # Create change map
        change_map = np.full(pred_arr.shape, -1, dtype=np.int8)
        change_map[mask_valid] = changes.astype(np.int8)
        
        # Change rate
        change_rate = float(n_changed) / n_total if n_total > 0 else 0.0
        
        # Save change map
        change_dir = Path("predictions")
        change_dir.mkdir(exist_ok=True)
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        change_file = change_dir / f"change_map_{timestamp}.tif"
        
        with rasterio.open(
            change_file, 'w',
            driver='GTiff',
            height=change_map.shape[0],
            width=change_map.shape[1],
            count=1,
            dtype=change_map.dtype,
            crs=pred_crs,
            transform=pred_transform
        ) as dst:
            dst.write(change_map, 1)
        
        return {
            "success": True,
            "change_detection": {
                "n_total_pixels": int(n_total),
                "n_changed_pixels": int(n_changed),
                "change_rate": change_rate,
                "change_matrix": change_matrix,
                "message": f"Detected {n_changed} changes out of {n_total} valid pixels ({change_rate*100:.1f}%)"
            },
            "change_map_file": str(change_file),
            "timestamp": timestamp
        }
    
    except Exception as e:
        print(f"[CHANGE DETECTION WORKFLOW ERROR] {str(e)}")
        import traceback
        traceback.print_exc()
        raise HTTPException(status_code=500, detail=f"Change detection workflow failed: {str(e)}")


@app.post("/api/change-detection/compare-periods")
async def compare_periods(request: ComparePeriodsPredictionConfig, background_tasks: BackgroundTasks):
    """Compare land use classification between two time periods."""
    from collections import Counter
    
    try:
        # Extract parameters
        model_filename = request.model_filename
        bbox = [request.min_lon, request.min_lat, request.max_lon, request.max_lat]
        
        current_start = request.current_period["start_date"]
        current_end = request.current_period["end_date"]
        pred_start = request.prediction_period["start_date"]
        pred_end = request.prediction_period["end_date"]
        
        max_scenes = request.max_scenes
        cloud_cover = request.cloud_cover
        resolution = request.resolution
        
        print(f"[COMPARE PERIODS] Current: {current_start} to {current_end} | Prediction: {pred_start} to {pred_end}")
        
        # Step 1: Predict on current period
        print(f"[COMPARE PERIODS] Step 1: Predicting current period...")
        current_result = await predict_with_ndvi(PredictionWithNDVIConfig(
            model_filename=model_filename,
            min_lon=request.min_lon,
            min_lat=request.min_lat,
            max_lon=request.max_lon,
            max_lat=request.max_lat,
            start_date=current_start,
            end_date=current_end,
            max_scenes=max_scenes,
            cloud_cover=cloud_cover,
            resolution=resolution,
            export_classification=True,
            export_ndvi=False
        ), background_tasks)
        
        # Get current classification raster
        current_class_file = None
        for f in current_result.get("output_files", []):
            if f.get("type") == "classification":
                current_class_file = f.get("path")
                break
        
        if not current_class_file:
            raise ValueError("No classification raster for current period")
        
        # Step 2: Predict on prediction period
        print(f"[COMPARE PERIODS] Step 2: Predicting future period...")
        pred_result = await predict_with_ndvi(PredictionWithNDVIConfig(
            model_filename=model_filename,
            min_lon=request.min_lon,
            min_lat=request.min_lat,
            max_lon=request.max_lon,
            max_lat=request.max_lat,
            start_date=pred_start,
            end_date=pred_end,
            max_scenes=max_scenes,
            cloud_cover=cloud_cover,
            resolution=resolution,
            export_classification=True,
            export_ndvi=False
        ), background_tasks)
        
        # Get prediction classification raster
        pred_class_file = None
        for f in pred_result.get("output_files", []):
            if f.get("type") == "classification":
                pred_class_file = f.get("path")
                break
        
        if not pred_class_file:
            raise ValueError("No classification raster for prediction period")
        
        # Step 3: Load both rasters
        print(f"[COMPARE PERIODS] Step 3: Comparing classifications...")
        with rasterio.open(current_class_file) as src:
            current_arr = src.read(1)
            crs = src.crs
            transform = src.transform
        
        with rasterio.open(pred_class_file) as src:
            pred_arr = src.read(1)
        
        # Ensure same shape
        if current_arr.shape != pred_arr.shape:
            raise ValueError(f"Shape mismatch: current {current_arr.shape} vs prediction {pred_arr.shape}")
        
        # Calculate changes
        mask_valid = ~np.isnan(current_arr) & ~np.isnan(pred_arr)
        changes = current_arr[mask_valid] != pred_arr[mask_valid]
        n_total = np.count_nonzero(mask_valid)
        n_changed = np.count_nonzero(changes)
        
        # Create change pairs
        change_pairs = list(zip(current_arr[mask_valid][changes], pred_arr[mask_valid][changes]))
        change_counter = Counter(change_pairs)
        change_matrix = {f"{int(curr)}->{int(pred)}": int(cnt) 
                        for (curr, pred), cnt in change_counter.items()}
        
        # Change rate
        change_rate = float(n_changed) / n_total if n_total > 0 else 0.0
        
        # Save change map
        change_dir = Path("predictions")
        change_dir.mkdir(exist_ok=True)
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        change_file = change_dir / f"change_map_{timestamp}.tif"
        
        change_map = np.zeros(current_arr.shape, dtype=np.int8)
        change_map[mask_valid] = changes.astype(np.int8)
        
        with rasterio.open(
            change_file, 'w',
            driver='GTiff',
            height=change_map.shape[0],
            width=change_map.shape[1],
            count=1,
            dtype=change_map.dtype,
            crs=crs,
            transform=transform
        ) as dst:
            dst.write(change_map, 1)
        
        # Extract class distributions
        current_classes = np.unique(current_arr[~np.isnan(current_arr)]).astype(int)
        current_dist = {int(c): int(np.count_nonzero(current_arr == c)) for c in current_classes}
        
        pred_classes = np.unique(pred_arr[~np.isnan(pred_arr)]).astype(int)
        pred_dist = {int(c): int(np.count_nonzero(pred_arr == c)) for c in pred_classes}
        
        return {
            "success": True,
            "current_classification": {
                "n_scenes": current_result.get("n_scenes"),
                "resolution": current_result.get("resolution"),
                "class_distribution": current_dist
            },
            "prediction_classification": {
                "n_scenes": pred_result.get("n_scenes"),
                "resolution": pred_result.get("resolution"),
                "class_distribution": pred_dist
            },
            "change_detection": {
                "n_total_pixels": int(n_total),
                "n_changed_pixels": int(n_changed),
                "change_rate": change_rate,
                "change_matrix": change_matrix,
                "message": f"Detected {n_changed} changes out of {n_total} valid pixels ({change_rate*100:.1f}%)"
            },
            "change_map_file": str(change_file),
            "timestamp": timestamp
        }
    
    except Exception as e:
        print(f"[COMPARE PERIODS ERROR] {str(e)}")
        import traceback
        traceback.print_exc()
        raise HTTPException(status_code=500, detail=f"Period comparison failed: {str(e)}")


@app.post("/api/change-detection")
async def change_detection_api(
    prediction_file: UploadFile = File(...),
    gt_file: UploadFile = File(...)
):
    """Detect changes between two raster files (prediction and ground truth)."""
    import tempfile
    from collections import Counter
    
    try:
        # Save uploaded files temporarily
        pred_tmp = tempfile.NamedTemporaryFile(suffix='.tif', delete=False)
        gt_tmp = tempfile.NamedTemporaryFile(suffix='.tif', delete=False)
        
        try:
            # Write uploaded files to temp
            pred_content = await prediction_file.read()
            gt_content = await gt_file.read()
            
            pred_tmp.write(pred_content)
            gt_tmp.write(gt_content)
            pred_tmp.close()
            gt_tmp.close()
            
            # Read prediction raster
            with rasterio.open(pred_tmp.name) as pred_ds:
                pred_arr = pred_ds.read(1)
                pred_crs = pred_ds.crs
                pred_transform = pred_ds.transform
            
            # Read ground truth raster
            with rasterio.open(gt_tmp.name) as gt_ds:
                gt_arr = gt_ds.read(1)
            
            # Ensure same shape
            if pred_arr.shape != gt_arr.shape:
                raise ValueError(f"Raster shapes don't match: prediction {pred_arr.shape} vs ground truth {gt_arr.shape}")
            
            # Calculate change detection
            mask_valid = (gt_arr >= 0) & (pred_arr >= 0) & ~np.isnan(gt_arr) & ~np.isnan(pred_arr)
            changes = gt_arr[mask_valid] != pred_arr[mask_valid]
            n_total = np.count_nonzero(mask_valid)
            n_changed = np.count_nonzero(changes)
            
            # Create change pairs matrix
            change_pairs = list(zip(gt_arr[mask_valid][changes], pred_arr[mask_valid][changes]))
            change_counter = Counter(change_pairs)
            change_matrix = {f"{int(gt)}->{int(pred)}": int(cnt) for (gt, pred), cnt in change_counter.items()}
            
            # Create change map (0=same, 1=changed, -1=invalid)
            change_map = np.full(pred_arr.shape, -1, dtype=np.int8)
            change_map[mask_valid] = changes.astype(np.int8)
            
            # Change rate
            change_rate = float(n_changed) / n_total if n_total > 0 else 0.0
            
            # Save change map as GeoTIFF
            change_dir = Path("predictions")
            change_dir.mkdir(exist_ok=True)
            
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            change_file = change_dir / f"change_map_{timestamp}.tif"
            
            with rasterio.open(
                change_file, 'w',
                driver='GTiff',
                height=change_map.shape[0],
                width=change_map.shape[1],
                count=1,
                dtype=change_map.dtype,
                crs=pred_crs,
                transform=pred_transform
            ) as dst:
                dst.write(change_map, 1)
            
            # Return results
            return {
                "success": True,
                "change_detection": {
                    "n_total_pixels": int(n_total),
                    "n_changed_pixels": int(n_changed),
                    "change_rate": change_rate,
                    "change_matrix": change_matrix,
                    "message": f"Detected {n_changed} changes out of {n_total} valid pixels ({change_rate*100:.1f}%)"
                },
                "change_map_file": str(change_file),
                "timestamp": timestamp
            }
            
        finally:
            # Cleanup temp files
            try:
                Path(pred_tmp.name).unlink()
                Path(gt_tmp.name).unlink()
            except:
                pass
    
    except Exception as e:
        print(f"[CHANGE DETECTION ERROR] {str(e)}")
        raise HTTPException(status_code=500, detail=f"Change detection failed: {str(e)}")
async def change_detection_api(
    prediction_file: UploadFile = File(...),
    gt_file: UploadFile = File(...)
):
    """Detect changes between prediction raster and ground truth raster."""
    import tempfile
    from collections import Counter
    
    try:
        # Save uploaded files temporarily
        pred_tmp = tempfile.NamedTemporaryFile(suffix='.tif', delete=False)
        gt_tmp = tempfile.NamedTemporaryFile(suffix='.tif', delete=False)
        
        try:
            # Write uploaded files to temp
            pred_content = await prediction_file.read()
            gt_content = await gt_file.read()
            
            pred_tmp.write(pred_content)
            gt_tmp.write(gt_content)
            pred_tmp.close()
            gt_tmp.close()
            
            # Read prediction raster
            import rasterio
            with rasterio.open(pred_tmp.name) as pred_ds:
                pred_arr = pred_ds.read(1)
                pred_crs = pred_ds.crs
                pred_transform = pred_ds.transform
            
            # Read ground truth raster
            with rasterio.open(gt_tmp.name) as gt_ds:
                gt_arr = gt_ds.read(1)
            
            # Ensure same shape
            if pred_arr.shape != gt_arr.shape:
                raise ValueError(f"Raster shapes don't match: prediction {pred_arr.shape} vs ground truth {gt_arr.shape}")
            
            # Calculate change detection
            mask_valid = (gt_arr >= 0) & (pred_arr >= 0) & ~np.isnan(gt_arr) & ~np.isnan(pred_arr)
            changes = gt_arr[mask_valid] != pred_arr[mask_valid]
            n_total = np.count_nonzero(mask_valid)
            n_changed = np.count_nonzero(changes)
            
            # Create change pairs matrix
            change_pairs = list(zip(gt_arr[mask_valid][changes], pred_arr[mask_valid][changes]))
            change_counter = Counter(change_pairs)
            change_matrix = {f"{int(gt)}->{int(pred)}": int(cnt) for (gt, pred), cnt in change_counter.items()}
            
            # Create change map (0=same, 1=changed, -1=invalid)
            change_map = np.full(pred_arr.shape, -1, dtype=np.int8)
            change_map[mask_valid] = changes.astype(np.int8)
            
            # Change rate
            change_rate = float(n_changed) / n_total if n_total > 0 else 0.0
            
            # Save change map as GeoTIFF
            change_dir = Path("predictions")
            change_dir.mkdir(exist_ok=True)
            
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            change_file = change_dir / f"change_map_{timestamp}.tif"
            
            with rasterio.open(
                change_file, 'w',
                driver='GTiff',
                height=change_map.shape[0],
                width=change_map.shape[1],
                count=1,
                dtype=change_map.dtype,
                crs=pred_crs,
                transform=pred_transform
            ) as dst:
                dst.write(change_map, 1)
            
            # Return results
            return {
                "success": True,
                "change_detection": {
                    "n_total_pixels": int(n_total),
                    "n_changed_pixels": int(n_changed),
                    "change_rate": change_rate,
                    "change_matrix": change_matrix,
                    "message": f"Detected {n_changed} changes out of {n_total} valid pixels ({change_rate*100:.1f}%)"
                },
                "change_map_file": str(change_file),
                "timestamp": timestamp
            }
            
        finally:
            # Cleanup temp files
            try:
                Path(pred_tmp.name).unlink()
                Path(gt_tmp.name).unlink()
            except:
                pass
    
    except Exception as e:
        print(f"[CHANGE DETECTION ERROR] {str(e)}")
        raise HTTPException(status_code=500, detail=f"Change detection failed: {str(e)}")


# ============ PREDICTION WITH NDVI API ============

@app.post("/api/predict/with-ndvi")
async def predict_with_ndvi(config: PredictionWithNDVIConfig, background_tasks: BackgroundTasks):
    """Predict land classification và NDVI cho một khu vực"""
    try:
        import numpy as np
        
        # Load model using ModelManager
        model_manager = get_model_manager()
        model, label_encoder, model_metadata = model_manager.load_model(config.model_filename)
        
        print(f"[PREDICT+NDVI] Loaded model: {config.model_filename}")
        print(f"  - Type: {model_metadata.get('model_type', 'unknown')}")
        print(f"  - Features: {model_metadata.get('features', [])}")
        
        # Check cache first
        cache_dir = Path("dataset_cache")
        cache_dir.mkdir(exist_ok=True)
        cache_key = f"pred_{config.min_lon}_{config.min_lat}_{config.max_lon}_{config.max_lat}_{config.start_date}_{config.end_date}_{config.max_scenes}_{config.cloud_cover}_{config.resolution}"
        cache_hash = hashlib.md5(cache_key.encode()).hexdigest()
        cache_file = cache_dir / f"prediction_input_{cache_hash}.joblib"
        
        bbox = [config.min_lon, config.min_lat, config.max_lon, config.max_lat]
        
        # Load from cache or fetch from Microsoft
        if cache_file.exists():
            print(f"[PREDICT+NDVI] Loading from cache: {cache_file.name}")
            cached = joblib.load(cache_file)
            # Extract s2_data from cache (already computed in cache)
            s2_data = cached["s2_data"]
            # Check if needed bands are available in cache
            available_bands = list(s2_data.data_vars.keys())
            needed_bands = ["B02", "B03", "B04", "B08"]
            if all(band in available_bands for band in needed_bands):
                # Use cached data directly (no need to compute again)
                data = s2_data[needed_bands]
                print(f"[PREDICT+NDVI] Using cached bands: {needed_bands}")
                # Set items to match the number of time slices in the cached data
                items = [None] * s2_data.sizes.get("time", 1)
            else:
                raise HTTPException(status_code=400, 
                    detail=f"Cache thiếu bands cần thiết. Có: {available_bands}, Cần: {needed_bands}")
        else:
            print(f"[PREDICT+NDVI] No cache found, fetching from Microsoft Planetary Computer")
            # Connect to Microsoft Planetary Computer
            catalog = Client.open(
                "https://planetarycomputer.microsoft.com/api/stac/v1",
                modifier=planetary_computer.sign_inplace
            )
            
            time_range = f"{config.start_date}/{config.end_date}"
            
            # Search for Sentinel-2 data
            search = catalog.search(
                collections=["sentinel-2-l2a"],
                bbox=bbox,
                datetime=time_range,
                query={"eo:cloud_cover": {"lt": config.cloud_cover}}
            )
            
            items = list(search.items())[:config.max_scenes]
            print(f"[PREDICT+NDVI] Found {len(items)} Sentinel-2 scenes")
            
            if len(items) == 0:
                raise HTTPException(status_code=404, detail="Không tìm thấy dữ liệu vệ tinh")
            
            # Sign items to refresh SAS tokens (keep as pystac.Item, not dict)
            signed_items = [planetary_computer.sign(item) for item in items]
            
            # Load all bands needed for features
            data = odc.stac.load(
                signed_items,
                bbox=bbox,
                bands=["B02", "B03", "B04", "B08"],  # Blue, Green, Red, NIR
                resolution=config.resolution,
                chunks={"x": 2048, "y": 2048}
            ).compute()
        
        print(f"[PREDICT+NDVI] Loaded data shape: {data.dims}")
        
        # Get expected number of features from model metadata
        expected_n_features = model_metadata.get("n_features", 3)
        print(f"[PREDICT+NDVI] Model expects {expected_n_features} features")
        
        # Calculate NDVI and other indices
        blue = data["B02"].values
        green = data["B03"].values
        red = data["B04"].values
        nir = data["B08"].values
        
        # Calculate indices for each time step
        # NDVI = (NIR - Red) / (NIR + Red)
        ndvi = (nir - red) / (nir + red + 1e-8)
        
        # NDWI = (Green - NIR) / (Green + NIR)
        ndwi = (green - nir) / (green + nir + 1e-8)
        
        # NDBI = (SWIR - NIR) / (SWIR + NIR) - we use Red as proxy
        ndbi = (red - nir) / (red + nir + 1e-8)
        
        # Prepare features for prediction
        height, width = ndvi.shape[1:3]  # Skip time dimension
        n_pixels = height * width
        n_times = ndvi.shape[0]
        
        print(f"[PREDICT+NDVI] Data has {n_times} time steps, spatial size: {height}x{width}")
        
        # Build features based on what model expects
        # Model metadata should tell us what features were used
        model_features = model_metadata.get("features", ["NDVI_mean", "VH_dB_mean", "VV_dB_mean"])
        
        # If model was trained with temporal features (multiple time steps)
        if expected_n_features > 10:  # Likely temporal features
            print(f"[PREDICT+NDVI] Building temporal features (all time steps)")
            # Use all time steps for each index
            feature_list = []
            
            # Add NDVI for each time step
            for t in range(n_times):
                feature_list.append(ndvi[t].flatten())
            
            # If model has more features, add NDWI and NDBI time series
            if expected_n_features >= n_times * 2:
                for t in range(n_times):
                    feature_list.append(ndwi[t].flatten())
            
            if expected_n_features >= n_times * 3:
                for t in range(n_times):
                    feature_list.append(ndbi[t].flatten())
            
            features = np.stack(feature_list, axis=1)
            
            # Adjust to match expected features
            if features.shape[1] < expected_n_features:
                # Pad with mean values
                n_missing = expected_n_features - features.shape[1]
                padding = np.tile(features[:, -1:], (1, n_missing))
                features = np.column_stack([features, padding])
            elif features.shape[1] > expected_n_features:
                # Trim to expected
                features = features[:, :expected_n_features]
        else:
            # Use mean values (aggregate features)
            print(f"[PREDICT+NDVI] Building aggregate features (mean values)")
            # Average over time dimension
            ndvi_mean = np.nanmean(ndvi, axis=0)
            ndwi_mean = np.nanmean(ndwi, axis=0)
            ndbi_mean = np.nanmean(ndbi, axis=0)
            
            # Reshape for prediction
            features = np.stack([ndvi_mean.flatten(), ndwi_mean.flatten(), ndbi_mean.flatten()], axis=1)
            
            # Adjust to match expected features if needed
            if features.shape[1] < expected_n_features:
                n_missing = expected_n_features - features.shape[1]
                padding = np.tile(features[:, -1:], (1, n_missing))
                features = np.column_stack([features, padding])
            elif features.shape[1] > expected_n_features:
                features = features[:, :expected_n_features]
        
        print(f"[PREDICT+NDVI] Built features shape: {features.shape}")
        
        # Handle NaN values
        valid_mask = ~np.isnan(features).any(axis=1)
        features_clean = features[valid_mask]
        
        print(f"[PREDICT+NDVI] Predicting {features_clean.shape[0]} valid pixels...")
        
        # Check if model is PyTorch/deep learning model and use GPU if available
        is_pytorch_model = hasattr(model, '__class__') and ('CNN' in model.__class__.__name__ or 'Swin' in model.__class__.__name__ or 'UNet' in model.__class__.__name__)
        
        if is_pytorch_model and config.use_gpu:
            try:
                import torch
                device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
                
                if torch.cuda.is_available():
                    print(f"[PREDICT+NDVI] Using GPU: {torch.cuda.get_device_name(0)}")
                    # Move model to GPU
                    model = model.to(device)
                    
                    # Predict in batches to avoid GPU memory overflow
                    batch_size = 8192  # Adjust based on GPU memory
                    predictions_list = []
                    
                    for i in range(0, len(features_clean), batch_size):
                        batch = features_clean[i:i+batch_size]
                        batch_tensor = torch.from_numpy(batch).float().to(device)
                        
                        with torch.no_grad():
                            batch_pred = model.predict(batch_tensor)
                        
                        # Move back to CPU if needed
                        if isinstance(batch_pred, torch.Tensor):
                            batch_pred = batch_pred.cpu().numpy()
                        
                        predictions_list.append(batch_pred)
                        
                        if (i // batch_size) % 10 == 0:
                            print(f"[PREDICT+NDVI] Processed {i + len(batch)}/{len(features_clean)} pixels on GPU")
                    
                    predictions = np.concatenate(predictions_list)
                    print(f"[PREDICT+NDVI] GPU prediction completed!")
                else:
                    print(f"[PREDICT+NDVI] GPU requested but not available, using CPU")
                    predictions = model.predict(features_clean)
            except Exception as gpu_error:
                print(f"[PREDICT+NDVI] GPU prediction failed: {gpu_error}, falling back to CPU")
                predictions = model.predict(features_clean)
        else:
            # Use CPU for traditional ML models
            if is_pytorch_model and not config.use_gpu:
                print(f"[PREDICT+NDVI] GPU disabled by user, using CPU")
            predictions = model.predict(features_clean)
        
        # Reshape back to raster
        prediction_raster = np.full(n_pixels, -1, dtype=np.int16)
        prediction_raster[valid_mask] = predictions
        prediction_raster = prediction_raster.reshape(height, width)

        # --- CHANGE DETECTION ---
        change_summary = None
        change_map = None
        try:
            # Use training shapefile as ground truth
            gt_shapefile = "train/ST_training data_updated_1130points_new.shp"
            gt_raster = rasterize_ground_truth(gt_shapefile, (height, width), bbox, class_column="class")
            # Compare prediction and ground truth
            mask_valid = (gt_raster >= 0) & (prediction_raster >= 0)
            changes = gt_raster[mask_valid] != prediction_raster[mask_valid]
            n_total = np.count_nonzero(mask_valid)
            n_changed = np.count_nonzero(changes)
            # Per-class change matrix
            from collections import Counter
            change_pairs = list(zip(gt_raster[mask_valid][changes], prediction_raster[mask_valid][changes]))
            change_counter = Counter(change_pairs)
            change_matrix = {f"{int(gt)}->{int(pred)}": int(cnt) for (gt, pred), cnt in change_counter.items()}
            change_summary = {
                "n_total": int(n_total),
                "n_changed": int(n_changed),
                "change_rate": float(n_changed) / n_total if n_total > 0 else 0.0,
                "change_matrix": change_matrix
            }
            # Optionally, create a change map (1=changed, 0=same, -1=invalid)
            change_map = np.full((height, width), -1, dtype=np.int8)
            change_map[mask_valid] = changes.astype(np.int8)
            # Save change map as GeoTIFF
            change_file = output_dir / f"change_map_{timestamp}.tif"
            with rasterio.open(
                change_file, 'w',
                driver='GTiff',
                height=height,
                width=width,
                count=1,
                dtype=change_map.dtype,
                crs='EPSG:4326',
                transform=from_bounds(bbox[0], bbox[1], bbox[2], bbox[3], width, height)
            ) as dst:
                dst.write(change_map, 1)
            output_files.append({"type": "change_map", "path": str(change_file)})
            print(f"[CHANGE DETECTION] Saved change map to {change_file}")
        except Exception as change_exc:
            print(f"[CHANGE DETECTION] Warning: {change_exc}")
        
        # Prepare outputs
        timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
        output_dir = Path("predictions")
        output_dir.mkdir(exist_ok=True)
        
        output_files = []
        
        # Export NDVI if requested
        if config.export_ndvi:
            ndvi_file = output_dir / f"ndvi_{timestamp}.tif"
            transform = from_bounds(bbox[0], bbox[1], bbox[2], bbox[3], width, height)
            
            with rasterio.open(
                ndvi_file, 'w',
                driver='GTiff',
                height=height,
                width=width,
                count=1,
                dtype=ndvi_mean.dtype,
                crs='EPSG:4326',
                transform=transform
            ) as dst:
                dst.write(ndvi_mean, 1)
            
            output_files.append({"type": "ndvi", "path": str(ndvi_file)})
            print(f"[PREDICT+NDVI] Saved NDVI to {ndvi_file}")
        
        # Export classification if requested
        if config.export_classification:
            class_file = output_dir / f"classification_{timestamp}.tif"
            transform = from_bounds(bbox[0], bbox[1], bbox[2], bbox[3], width, height)
            
            with rasterio.open(
                class_file, 'w',
                driver='GTiff',
                height=height,
                width=width,
                count=1,
                dtype=prediction_raster.dtype,
                crs='EPSG:4326',
                transform=transform
            ) as dst:
                dst.write(prediction_raster, 1)
            
            output_files.append({"type": "classification", "path": str(class_file)})
            print(f"[PREDICT+NDVI] Saved classification to {class_file}")
        
        # Calculate statistics
        ndvi_stats = {
            "mean": float(np.nanmean(ndvi_mean)),
            "min": float(np.nanmin(ndvi_mean)),
            "max": float(np.nanmax(ndvi_mean)),
            "std": float(np.nanstd(ndvi_mean))
        }
        
        # Count classes
        unique_classes, counts = np.unique(predictions, return_counts=True)
        class_distribution = {
            int(cls): int(count) for cls, count in zip(unique_classes, counts)
        }
        
        return {
            "success": True,
            "message": "Prediction with NDVI completed",
            "output_files": output_files,
            "ndvi_stats": ndvi_stats,
            "class_distribution": class_distribution,
            "n_scenes": len(items),
            "resolution": config.resolution,
            "bbox": bbox,
            "change_detection": change_summary
        }
        
    except Exception as e:
        print(f"[PREDICT+NDVI ERROR] {str(e)}")
        import traceback
        traceback.print_exc()
        raise HTTPException(status_code=500, detail=str(e))


# ============ NDVI TIME SERIES API ============

@app.post("/api/ndvi/timeseries")
async def calculate_ndvi_timeseries(config: NDVIConfig):
    """Tính NDVI time series cho một khu vực"""
    try:
        import numpy as np
        import xarray as xr
        from pystac_client import Client
        import planetary_computer
        import odc.stac
        
        print(f"[NDVI] Starting calculation for bbox: {config.bbox}, time: {config.start_date} to {config.end_date}")
        
        # Connect to Microsoft Planetary Computer STAC API
        catalog = Client.open(
            "https://planetarycomputer.microsoft.com/api/stac/v1",
            modifier=planetary_computer.sign_inplace
        )
        
        bbox = config.bbox
        time_range = f"{config.start_date}/{config.end_date}"
        
        # Search for Sentinel-2 data
        search = catalog.search(
            collections=["sentinel-2-l2a"],
            bbox=bbox,
            datetime=time_range,
            query={"eo:cloud_cover": {"lt": config.max_cloud_cover}}
        )
        
        items = list(search.items())
        print(f"[NDVI] Found {len(items)} Sentinel-2 scenes")
        
        if len(items) == 0:
            raise HTTPException(status_code=404, detail="Không tìm thấy dữ liệu Sentinel-2 cho khu vực và thời gian này")
        
        # Load data for each time step
        ndvi_timeseries = []
        
        for item in items:
            try:
                # Load NIR (B08) and Red (B04) bands
                data = odc.stac.load(
                    [item],
                    bbox=bbox,
                    bands=["B04", "B08"],  # Red and NIR
                    resolution=config.resolution,
                    chunks={"x": 2048, "y": 2048}
                ).compute()
                
                if data is None or len(data.keys()) == 0:
                    continue
                
                # Calculate NDVI = (NIR - Red) / (NIR + Red)
                nir = data["B08"].values
                red = data["B04"].values
                
                # Avoid division by zero
                denominator = nir + red
                denominator = np.where(denominator == 0, np.nan, denominator)
                
                ndvi = (nir - red) / denominator
                
                # Calculate mean NDVI (ignore NaN values)
                mean_ndvi = float(np.nanmean(ndvi))
                
                # Get date from item
                date_str = item.datetime.strftime("%Y-%m-%d")
                
                ndvi_timeseries.append({
                    "date": date_str,
                    "ndvi": mean_ndvi
                })
                
                print(f"[NDVI] {date_str}: NDVI = {mean_ndvi:.3f}")
                
            except Exception as e:
                print(f"[NDVI WARNING] Failed to process item {item.id}: {e}")
                continue
        
        if len(ndvi_timeseries) == 0:
            raise HTTPException(status_code=500, detail="Không thể tính NDVI cho bất kỳ ảnh nào")
        
        # Sort by date
        ndvi_timeseries.sort(key=lambda x: x["date"])
        
        # Calculate statistics
        ndvi_values = [item["ndvi"] for item in ndvi_timeseries]
        mean_ndvi = float(np.mean(ndvi_values))
        min_ndvi = float(np.min(ndvi_values))
        max_ndvi = float(np.max(ndvi_values))
        
        result = {
            "timeseries": ndvi_timeseries,
            "n_images": len(ndvi_timeseries),
            "mean_ndvi": mean_ndvi,
            "min_ndvi": min_ndvi,
            "max_ndvi": max_ndvi,
            "bbox": bbox,
            "time_range": time_range
        }
        
        print(f"[NDVI] Calculation complete. Mean NDVI: {mean_ndvi:.3f}, Images: {len(ndvi_timeseries)}")
        
        return result
        
    except HTTPException:
        raise
    except Exception as e:
        print(f"[NDVI ERROR] {e}")
        import traceback
        traceback.print_exc()
        raise HTTPException(status_code=500, detail=f"Lỗi khi tính NDVI: {str(e)}")


if __name__ == "__main__":
    print("=" * 70)
    print("🚀 LAND CLASSIFICATION TRAINING API SERVER")
    print("=" * 70)
    print("\n📍 Endpoints:")
    print("   - Web Interface: http://localhost:8000")
    print("   - API Docs: http://localhost:8000/docs")
    print("   - Start Training: POST http://localhost:8000/api/training/start")
    print("   - Check Status: GET http://localhost:8000/api/training/status")
    print("\n" + "=" * 70)
    
    uvicorn.run(app, host="0.0.0.0", port=8000, log_level="info")