import os
from train_module import train_model
from pathlib import Path

def train_all_models():
    # Ensure output directory exists
    output_dir = Path("land_classification_model")
    output_dir.mkdir(exist_ok=True)
    
    # Define models and their optimized hyperparameters
    models_config = [
        {'type': 'xgboost', 'n_estimators': 200, 'max_depth': 6, 'learning_rate': 0.05, 'use_gpu': True},
        {'type': 'lightgbm', 'n_estimators': 300, 'max_depth': -1, 'learning_rate': 0.05, 'use_gpu': True},
        {'type': 'random_forest', 'n_estimators': 100, 'max_depth': 15, 'use_gpu': False},
        {'type': 'decision_tree', 'max_depth': 12, 'use_gpu': False},
        {'type': 'svm', 'use_gpu': False},
        {'type': 'cnn', 'n_estimators': 30, 'use_gpu': True},  # n_estimators acts as epochs
        {'type': 'swin-unet', 'n_estimators': 80, 'learning_rate': 0.0003, 'use_gpu': True},
        {'type': 'mobilenet-lraspp', 'n_estimators': 50, 'learning_rate': 0.0008, 'use_gpu': True}
    ]
    
    results = []
    
    print("=" * 70)
    print("🚀 BẮT ĐẦU HUẤN LUYỆN TẤT CẢ MÔ HÌNH PHÂN LOẠI LỚP PHỦ")
    print("=" * 70)
    
    for cfg in models_config:
        model_type = cfg['type']
        print(f"\n[{model_type.upper()}] Đang tiến hành huấn luyện...")
        
        # Prepare parameters for train_model
        params = {
            'bbox': [105.5, 9.2, 106.3, 10.0], # Matching the working coordinates
            'time_range': '2023-01-01/2023-04-30',
            'model_type': model_type,
            'feature_mode': 'extended',
            'use_cache': True,
            'output_model_path': f"land_classification_model/model_{model_type}_auto.joblib"
        }
        
        # Merge specific hyperparameters
        for key in ['n_estimators', 'max_depth', 'learning_rate', 'use_gpu']:
            if key in cfg:
                params[key] = cfg[key]
                
        try:
            res = train_model(**params)
            
            # Extract metrics
            if res.get('success'):
                metrics = {
                    'model': model_type.upper(),
                    'accuracy': res.get('test_accuracy', 0.0),
                    'params': f"Estimators:{cfg.get('n_estimators','-')}, Depth:{cfg.get('max_depth','-')}",
                    'status': '✅ Success'
                }
            else:
                metrics = {
                    'model': model_type.upper(),
                    'accuracy': 0.0,
                    'params': '-',
                    'status': f"❌ Failed: {res.get('error', 'Unknown')}"
                }
            results.append(metrics)
            print(f"[{model_type.upper()}] ✅ Xong! Accuracy: {metrics['accuracy']:.4f}")
            
        except Exception as e:
            print(f"[{model_type.upper()}] ❌ LỖI: {e}")
            results.append({
                'model': model_type.upper(),
                'accuracy': 0.0,
                'params': '-',
                'status': f"❌ Error: {str(e)}"
            })
            
    # Print summary table
    print("\n\n" + "=" * 70)
    print("📊 TỔNG HỢP KẾT QUẢ HUẤN LUYỆN")
    print("=" * 70)
    print(f"{'Mô hình':<20} | {'Độ chính xác (Acc)':<20} | {'Tham số':<25} | {'Trạng thái'}")
    print("-" * 70)
    for r in results:
        acc_str = f"{r['accuracy']:.4f}" if isinstance(r['accuracy'], float) else str(r['accuracy'])
        print(f"{r['model']:<20} | {acc_str:<20} | {r['params']:<25} | {r['status']}")
    print("=" * 70)
    
if __name__ == "__main__":
    train_all_models()