"""
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
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 report generator
from report_generator import generate_training_report, generate_prediction_report

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
    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


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]


@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("/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/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 đủ"""
    cache_dir = Path("dataset_cache")
    
    if not cache_dir.exists():
        return {"exists": False, "files": [], "total_size_mb": 0}
    
    cache_files = []
    total_size = 0
    
    for cache_file in cache_dir.glob("*.joblib"):
        size = cache_file.stat().st_size
        total_size += size
        
        # Try to load metadata from cache
        metadata = {}
        try:
            cached_data = joblib.load(cache_file)
            if 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"Error loading cache metadata: {e}")
        
        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)
    
    return {
        "exists": True,
        "files": cache_files,
        "count": len(cache_files),
        "total_size_mb": round(total_size / 1024 / 1024, 2)
    }



@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"
        
        reports.append({
            "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}"
        })
    
    # 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 - Áp dụng phương pháp từ 02.predict_ODC.ipynb"""
    global prediction_status
    
    try:
        prediction_status["progress"] = "Đang import thư viện..."
        
        # Import required libraries
        import xarray as xr
        import numpy as np
        from datetime import datetime as dt
        import rioxarray
        import dask.array as da
        
        # 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
        model_path = Path("model_train") / config.model_filename
        if not model_path.exists():
            raise FileNotFoundError(f"Model không tồn tại: {config.model_filename}")
        
        model_data = joblib.load(model_path)
        
        # Extract model from dict (models are saved as {'model': xgb_model, 'label_encoder': encoder})
        if isinstance(model_data, dict):
            model = model_data.get('model')
            label_encoder = model_data.get('label_encoder')
        else:
            model = model_data
            label_encoder = None
        
        # Check if it's a CNN model (PyTorch)
        is_cnn_model = hasattr(model, '__class__') and 'CNN' in model.__class__.__name__
        if is_cnn_model:
            prediction_status["progress"] = "Phát hiện PyTorch CNN model..."
            # Import PyTorch if needed
            try:
                import torch
            except ImportError:
                raise ImportError("PyTorch is required for CNN prediction. Install: pip install torch")
        
        prediction_status["progress"] = "Đang kiểm tra cache dữ liệu đầu vào..."
        import hashlib, os
        cache_dir = Path("dataset_cache")
        cache_dir.mkdir(exist_ok=True)
        # Tạo cache key từ bbox, time_range, max_scenes, cloud_cover, resolution
        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"
        
        # 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}"

        if cache_file.exists():
            prediction_status["progress"] = "Đang load dữ liệu từ cache..."
            cached = joblib.load(cache_file)
            s2_data = cached["s2_data"]
            s2_items = cached["s2_items"]
            vh_monthly = cached.get("vh_monthly")
            vv_monthly = cached.get("vv_monthly")
            use_radar = cached.get("use_radar", False)
        else:
            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,
            )
            # ============ BƯỚC 1: TẢI DỮ LIỆU SENTINEL-2 ============
            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..."
            s2_data = load(
                s2_items,
                bbox=bbox,
                chunks={"time": 1, "x": 2048, "y": 2048},
                groupby="solar_day",
                resolution=config.resolution
            )
            # ============ BƯỚC 4: TẢI DỮ LIỆU SENTINEL-1 (Radar)... ============
            prediction_status["progress"] = "Đang tải dữ liệu Sentinel-1 (Radar)..."
            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]
                prediction_status["progress"] = f"Đang xử lý {len(s1_items)} scenes Sentinel-1..."
                s1_data = load(
                    s1_items,
                    bbox=bbox,
                    chunks={"time": 1, "x": 2048, "y": 2048},
                    groupby="sat:absolute_orbit",
                    resolution=config.resolution
                )
                if "vh" in s1_data and "vv" in s1_data:
                    vh = s1_data["vh"].astype('float32')
                    vv = s1_data["vv"].astype('float32')
                    vh_monthly = vh.resample(time="1ME").mean().compute()
                    vv_monthly = vv.resample(time="1ME").mean().compute()
                    use_radar = True
                else:
                    vh_monthly = None
                    vv_monthly = None
                    use_radar = False
            else:
                vh_monthly = None
                vv_monthly = None
                use_radar = False
            # Lưu cache
            joblib.dump({
                "s2_data": s2_data,
                "s2_items": s2_items,
                "vh_monthly": vh_monthly,
                "vv_monthly": vv_monthly,
                "use_radar": use_radar
            }, cache_file)
        
        # ============ BƯỚC 2: TÍNH NDVI VÀ XỬ LÝ MÂY ============
        prediction_status["progress"] = "Đang tính toán NDVI và xử lý mây..."
        
        # Calculate NDVI using Sentinel-2 band names (B08 = NIR, B04 = Red)
        nir = s2_data["B08"].astype('float32')
        red = s2_data["B04"].astype('float32')
        ndvi = (nir - red) / (nir + red + 1e-8)
        
        # Mask clouds using SCL band if available
        if "SCL" in s2_data:
            scl = s2_data["SCL"]
            # SCL values: 4=vegetation, 5=bare soil, 6=water - these are clear
            # 3=cloud shadow, 8=cloud medium, 9=cloud high, 10=cirrus - mask these
            cloud_mask = (scl == 3) | (scl == 8) | (scl == 9) | (scl == 10)
            ndvi = ndvi.where(~cloud_mask)
        
        # ============ BƯỚC 3: ĐIỀN GIÁ TRỊ NAN (FILL NAN) ============
        prediction_status["progress"] = "Đang điền giá trị bị che mây..."
        
        # Fill NaN using forward fill and backward fill
        ndvi_filled = ndvi.ffill(dim='time').bfill(dim='time')
        
        # Resample to monthly average
        prediction_status["progress"] = "Đang tính trung bình NDVI theo tháng..."
        ndvi_monthly = ndvi_filled.resample(time="1ME").mean()
        
        # Compute NDVI (convert from dask to numpy)
        ndvi_monthly = ndvi_monthly.compute()
        
        # ============ BƯỚC 4: TẢI DỮ LIỆU SENTINEL-1 (VH, VV) ============
        # Only load radar if not already in cache
        if not cache_file.exists() or (cache_file.exists() and not use_radar):
            prediction_status["progress"] = "Đang tải dữ liệu Sentinel-1 (Radar)..."
            
            # Initialize catalog if not already done
            if not cache_file.exists():
                # catalog already initialized in the else block above
                pass
            else:
                # Need to initialize catalog for radar search
                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,
                )
            
            # Search Sentinel-1 data
            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]
                prediction_status["progress"] = f"Đang xử lý {len(s1_items)} scenes Sentinel-1..."
                
                # Load Sentinel-1 data (without like= to avoid conflict with bbox/resolution)
                s1_data = load(
                    s1_items,
                    bbox=bbox,
                    chunks={"time": 1, "x": 2048, "y": 2048},
                    groupby="sat:absolute_orbit",
                    resolution=config.resolution
                )
                
                # Extract VH and VV bands
                if "vh" in s1_data and "vv" in s1_data:
                    vh = s1_data["vh"].astype('float32')
                    vv = s1_data["vv"].astype('float32')
                    
                    # Resample to monthly average
                    prediction_status["progress"] = "Đang tính trung bình VH/VV theo tháng..."
                    vh_monthly = vh.resample(time="1ME").mean().compute()
                    vv_monthly = vv.resample(time="1ME").mean().compute()
                    
                    use_radar = True
                else:
                    prediction_status["progress"] = "Không tìm thấy bands VH/VV, tiếp tục với NDVI..."
                    use_radar = False
            else:
                prediction_status["progress"] = "Không có dữ liệu Sentinel-1, tiếp tục với NDVI..."
                use_radar = False
        
        # ============ BƯỚC 5: CHUẨN BỊ FEATURES CHO DỰ ĐOÁN ============
        prediction_status["progress"] = "Đang chuẩn bị features cho dự đoán..."
        
        # Get shape information
        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
        
        # Prepare NDVI features (flatten each time step)
        ndvi_features = []
        for t in range(n_times_ndvi):
            ndvi_t = ndvi_monthly.isel(time=t).values.flatten()
            ndvi_features.append(ndvi_t)
        
        # Stack NDVI features
        features = np.column_stack(ndvi_features)
        
        # Add radar features if available
        if use_radar:
            n_times_vh = len(vh_monthly.time)
            n_times_vv = len(vv_monthly.time)
            
            # Add VH features
            for t in range(min(n_times_vh, n_times_ndvi)):
                vh_t = vh_monthly.isel(time=t).values.flatten()
                # Resize if needed
                if len(vh_t) != n_pixels:
                    vh_t = np.resize(vh_t, n_pixels)
                features = np.column_stack([features, vh_t])
            
            # Add VV features
            for t in range(min(n_times_vv, n_times_ndvi)):
                vv_t = vv_monthly.isel(time=t).values.flatten()
                # Resize if needed
                if len(vv_t) != n_pixels:
                    vv_t = np.resize(vv_t, n_pixels)
                features = np.column_stack([features, vv_t])
        
        # Handle NaN values in features✓ CNN PyTorch: Mạnh nhất với ảnh vệ tinh, tự học features, tương thích GPU tốt, cần pip install torch
        features = np.nan_to_num(features, nan=0.0)
        
        # ============ BƯỚC 6: DỰ ĐOÁN ============
        # Check model's expected feature count and adjust
        try:
            # Get expected number of features from model
            if is_cnn_model:
                # For PyTorch CNN, get n_features from model
                expected_features = model.n_features
            elif hasattr(model, 'n_features_in_'):
                expected_features = model.n_features_in_
            elif hasattr(model, 'feature_names_in_'):
                expected_features = len(model.feature_names_in_)
            else:
                # Try to get from booster for XGBoost
                try:
                    expected_features = model.get_booster().num_features()
                except:
                    expected_features = features.shape[1]
            
            prediction_status["progress"] = f"Model cần {expected_features} features, đang có {features.shape[1]} features..."
            
            # Adjust features to match model
            if features.shape[1] > expected_features:
                # Trim to expected number (use only first N features - NDVI only)
                prediction_status["progress"] = f"Cắt bớt features từ {features.shape[1]} xuống {expected_features}..."
                features = features[:, :expected_features]
            elif features.shape[1] < expected_features:
                # Pad with zeros or repeat last features
                prediction_status["progress"] = f"Thêm features từ {features.shape[1]} lên {expected_features}..."
                n_missing = expected_features - features.shape[1]
                # Repeat last feature column to fill
                padding = np.tile(features[:, -1:], (1, n_missing))
                features = np.column_stack([features, padding])
        except Exception as e:
            prediction_status["progress"] = f"Không thể xác định số features của model, tiếp tục với {features.shape[1]} features..."
        
        prediction_status["progress"] = f"Đang dự đoán với {features.shape[1]} features..."
        
        # Make prediction
        if is_cnn_model:
            # PyTorch CNN prediction
            predictions = model.predict(features)
        else:
            predictions = model.predict(features)
        
        # Decode labels if label_encoder exists
        if label_encoder is not None:
            try:
                predictions = label_encoder.inverse_transform(predictions)
            except:
                pass  # Keep numeric predictions if inverse_transform fails
        
        # Reshape to original shape
        pred_shape = (y_size, x_size)
        predictions_2d = predictions.reshape(pred_shape)
        
        # ============ BƯỚC 7: TẠO OUTPUT VÀ LƯU KẾT QUẢ ============
        prediction_status["progress"] = "Đang tạo bản đồ phân loại..."
        
        # Create output xarray
        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)
        
        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)
            
            # Add colorbar
            cbar = plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
            cbar.set_label('Class', rotation=270, labelpad=15)
            
            # 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],
            "n_times_ndvi": n_times_ndvi,
            "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"):
        predictions.append({
            "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}"
        })
    
    # 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
    
    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["started_at"] = datetime.now().isoformat()
        
        try:
            # Create PredictionConfig from job config
            pred_config = PredictionConfig(**job["config"])
            
            # Run prediction (simplified version)
            # In real implementation, call the actual prediction function
            print(f"[BATCH] Processing job: {job['name']}")
            
            # Simulate prediction (replace with actual prediction call)
            # await run_prediction(pred_config)
            
            # For now, mark as completed
            job["status"] = "completed"
            job["completed_at"] = datetime.now().isoformat()
            job["result"] = {
                "output_file": f"predictions/batch_{job['job_id']}.tif",
                "message": "Prediction completed successfully"
            }
            
        except Exception as e:
            job["error"] = str(e)
            
            # Retry logic
            if job["retries"] < job["max_retries"]:
                job["retries"] += 1
                job["status"] = "queued"  # Retry
                print(f"[BATCH] Job {job['name']} failed, retrying ({job['retries']}/{job['max_retries']})")
                continue
            else:
                job["status"] = "failed"
                job["completed_at"] = datetime.now().isoformat()
                print(f"[BATCH] Job {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:]


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")
