from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, Field
import pickle
import numpy as np
import pandas as pd
from typing import Dict, List, Any
import uvicorn
import os
import traceback
import json

# Function to list files in current directory
def list_files():
    """List all files in the current directory"""
    try:
        files = os.listdir('.')
        print("📁 Files in current directory:")
        for file in files:
            file_path = os.path.join('.', file)
            if os.path.isfile(file_path):
                size = os.path.getsize(file_path)
                print(f"  📄 {file} ({size} bytes)")
            else:
                print(f"  📁 {file}/")
        return files
    except Exception as e:
        print(f"Error listing files: {e}")
        return []

# Enhanced model loading with debugging
def load_fertilizer_model():
    """Load the saved fertilizer recommendation model and components"""
    print("🔄 Starting model loading...")
    
    # List files first
    files = list_files()
    
    # Check if required files exist
    required_files = [
        'fertilizer_recommendation_model.pkl',
        'fertilizer_label_encoder.pkl', 
        'crop_label_encoder.pkl',
        'fertilizer_model_info.pkl'
    ]
    
    print(f"\n🔍 Checking for required files:")
    for file in required_files:
        print(f"  {file}: {'✅ Found' if file in files else '❌ Missing'}")
    
    try:
        # Load model
        with open('fertilizer_recommendation_model.pkl', 'rb') as f:
            model = pickle.load(f)
        print("✅ Model loaded successfully!")
        
        # Load fertilizer encoder
        with open('fertilizer_label_encoder.pkl', 'rb') as f:
            fertilizer_encoder = pickle.load(f)
        print("✅ Fertilizer encoder loaded successfully!")
        
        # Load crop encoder
        with open('crop_label_encoder.pkl', 'rb') as f:
            crop_encoder = pickle.load(f)
        print("✅ Crop encoder loaded successfully!")
        
        # Load model info
        with open('fertilizer_model_info.pkl', 'rb') as f:
            model_info = pickle.load(f)
        print("✅ Model info loaded successfully!")
        
        # Try to load feature statistics (optional)
        feature_stats = None
        try:
            with open('feature_statistics.json', 'r') as f:
                feature_stats = json.load(f)
            print("✅ Feature statistics loaded successfully!")
        except:
            print("⚠️ Feature statistics not found (optional)")
        
        print(f"📊 Model info: {model_info}")
        
        return model, fertilizer_encoder, crop_encoder, model_info, feature_stats
        
    except Exception as e:
        print(f"❌ Error loading model: {e}")
        print(f"📍 Full traceback: {traceback.format_exc()}")
        return None, None, None, None, None

# Input validation function
def validate_input_ranges(nitrogen_percent, phosphorus_ppm, potassium_meq_percent, soil_ph, feature_stats=None):
    """Validate if input values are within reasonable ranges"""
    warnings = []
    
    # Basic range validation
    if nitrogen_percent < 0 or nitrogen_percent > 5:
        warnings.append(f"Nitrogen percentage {nitrogen_percent} should be between 0-5%")
    
    if phosphorus_ppm < 0 or phosphorus_ppm > 200:
        warnings.append(f"Phosphorus {phosphorus_ppm}ppm should be between 0-200ppm")
    
    if potassium_meq_percent < 0 or potassium_meq_percent > 10:
        warnings.append(f"Potassium {potassium_meq_percent}meq% should be between 0-10meq%")
    
    if soil_ph < 3 or soil_ph > 10:
        warnings.append(f"Soil pH {soil_ph} should be between 3-10")
    
    # Advanced validation using training data statistics if available
    if feature_stats:
        try:
            numerical_stats = feature_stats.get('numerical_features', feature_stats)
            
            params = {
                'total_Nitrogen_percent': nitrogen_percent,
                'phosphorus_Olsen_ppm': phosphorus_ppm,
                'potassium_meq_percent': potassium_meq_percent,
                'soil_pH': soil_ph
            }
            
            for param_name, value in params.items():
                if param_name in numerical_stats:
                    param_stats = numerical_stats[param_name]
                    min_val = param_stats['min']
                    max_val = param_stats['max']
                    mean_val = param_stats['mean']
                    
                    if value < min_val or value > max_val:
                        warnings.append(f"{param_name}: {value} is outside training range [{min_val:.2f}, {max_val:.2f}]")
        except Exception as e:
            warnings.append(f"Could not perform advanced validation: {e}")
    
    return warnings

# Initialize FastAPI app
app = FastAPI(
    title="🌾 Smart Fertilizer Recommender API",
    description="Predict fertilizer recommendations based on soil nutrient parameters and crop type using Random Forest ML model",
    version="1.0.0",
    docs_url="/docs",
    redoc_url="/redoc"
)

# Load model at startup
print("🚀 Starting application...")
model, fertilizer_encoder, crop_encoder, model_info, feature_stats = load_fertilizer_model()

# Pydantic models for request/response
class SoilParameters(BaseModel):
    nitrogen_percent: float = Field(..., ge=0, le=5, description="Total Nitrogen percentage (0-5%)")
    phosphorus_ppm: float = Field(..., ge=0, le=200, description="Phosphorus Olsen in ppm (0-200)")
    soil_ph: float = Field(..., ge=3, le=10, description="Soil pH level (3-10)")
    potassium_meq_percent: float = Field(..., ge=0, le=10, description="Potassium in meq% (0-10%)")
    crop: str = Field(..., description="Crop name (e.g., 'maize', 'beans', 'wheat')")
    
    class Config:
        schema_extra = {
            "example": {
                "nitrogen_percent": 0.08,
                "phosphorus_ppm": 23.0,
                "soil_ph": 5.76,
                "potassium_meq_percent": 1.36,
                "crop": "maize"
            }
        }

class FertilizerRecommendation(BaseModel):
    fertilizer: str
    confidence: float

class PredictionResponse(BaseModel):
    primary_recommendation: str = Field(..., description="Primary fertilizer recommendation")
    confidence: float = Field(..., description="Confidence percentage for primary recommendation")
    all_recommendations: List[FertilizerRecommendation] = Field(..., description="Top 3 fertilizer recommendations with confidence")
    input_parameters: Dict[str, Any] = Field(..., description="Input parameters used for prediction")
    validation_warnings: List[str] = Field(default=[], description="Input validation warnings")

class ModelInfoResponse(BaseModel):
    model_name: str
    model_type: str
    train_accuracy: float
    test_accuracy: float
    cv_mean: float
    cv_std: float
    n_classes: int
    fertilizer_classes: List[str]
    crop_classes: List[str]

# Add a debug endpoint to check files
@app.get("/debug/files")
async def debug_files():
    """Debug endpoint to list all files"""
    files = list_files()
    return {
        "current_directory": os.getcwd(),
        "files": files,
        "model_loaded": model is not None,
        "fertilizer_encoder_loaded": fertilizer_encoder is not None,
        "crop_encoder_loaded": crop_encoder is not None,
        "model_info_loaded": model_info is not None,
        "feature_stats_loaded": feature_stats is not None
    }

# API Endpoints
@app.get("/")
async def root():
    """Welcome message and API information"""
    return {
        "message": "🌾 Smart Fertilizer Recommender API",
        "description": "Use /predict endpoint to get fertilizer recommendations based on soil parameters and crop type",
        "model": model_info['model_name'] if model_info else "Model not loaded",
        "model_loaded": model is not None,
        "available_crops": crop_encoder.classes_.tolist() if crop_encoder else [],
        "available_fertilizers": len(fertilizer_encoder.classes_) if fertilizer_encoder else 0,
        "docs": "Visit /docs for interactive API documentation",
        "debug": "Visit /debug/files to see available files"
    }

@app.get("/health")
async def health_check():
    """Enhanced health check endpoint"""
    return {
        "status": "healthy",
        "model_loaded": model is not None,
        "encoders_loaded": fertilizer_encoder is not None and crop_encoder is not None,
        "model_info_loaded": model_info is not None,
        "current_directory": os.getcwd(),
        "files_count": len(os.listdir('.')) if os.path.exists('.') else 0
    }

@app.get("/model-info", response_model=ModelInfoResponse)
async def get_model_info():
    """Get information about the trained model"""
    if model_info is None:
        raise HTTPException(status_code=500, detail="Model information not available")
    
    return ModelInfoResponse(**model_info)

@app.get("/crops")
async def get_available_crops():
    """Get list of available crops"""
    if crop_encoder is None:
        raise HTTPException(status_code=500, detail="Crop encoder not loaded")
    
    return {
        "available_crops": crop_encoder.classes_.tolist(),
        "total_crops": len(crop_encoder.classes_)
    }

@app.get("/fertilizers")
async def get_available_fertilizers():
    """Get list of available fertilizers"""
    if fertilizer_encoder is None:
        raise HTTPException(status_code=500, detail="Fertilizer encoder not loaded")
    
    return {
        "available_fertilizers": fertilizer_encoder.classes_.tolist(),
        "total_fertilizers": len(fertilizer_encoder.classes_)
    }

@app.post("/predict", response_model=PredictionResponse)
async def predict_fertilizer(soil_params: SoilParameters):
    """
    Predict fertilizer recommendation based on soil parameters and crop type
    
    - **nitrogen_percent**: Total Nitrogen percentage (0-5%)
    - **phosphorus_ppm**: Phosphorus Olsen in ppm (0-200)
    - **soil_ph**: Soil pH level (3-10)
    - **potassium_meq_percent**: Potassium in meq% (0-10%)
    - **crop**: Crop name (use /crops endpoint to see available options)
    
    Returns fertilizer recommendation with confidence scores
    """
    if model is None or fertilizer_encoder is None or crop_encoder is None:
        raise HTTPException(
            status_code=500, 
            detail={
                "error": "Model components not loaded",
                "debug_info": {
                    "model_loaded": model is not None,
                    "fertilizer_encoder_loaded": fertilizer_encoder is not None,
                    "crop_encoder_loaded": crop_encoder is not None,
                    "files_in_directory": os.listdir('.') if os.path.exists('.') else []
                }
            }
        )
    
    try:
        # Validate input ranges
        validation_warnings = validate_input_ranges(
            soil_params.nitrogen_percent, 
            soil_params.phosphorus_ppm,
            soil_params.potassium_meq_percent, 
            soil_params.soil_ph,
            feature_stats
        )
        
        # Validate crop name
        if soil_params.crop not in crop_encoder.classes_:
            available_crops = crop_encoder.classes_.tolist()
            raise HTTPException(
                status_code=400, 
                detail=f"Unknown crop '{soil_params.crop}'. Available crops: {available_crops}"
            )
        
        # Encode crop
        crop_encoded = crop_encoder.transform([soil_params.crop])[0]
        
        # Prepare input data (5 features: N, P, K, pH, crop_encoded)
        input_data = np.array([[
            soil_params.nitrogen_percent,
            soil_params.phosphorus_ppm,
            soil_params.potassium_meq_percent,
            soil_params.soil_ph,
            crop_encoded
        ]])
        
        # Make prediction
        prediction_encoded = model.predict(input_data)[0]
        prediction_proba = model.predict_proba(input_data)[0]
        
        # Decode primary prediction
        predicted_fertilizer = fertilizer_encoder.inverse_transform([prediction_encoded])[0]
        
        # Get top 3 recommendations with probabilities
        top_indices = np.argsort(prediction_proba)[::-1][:3]
        recommendations = []
        
        for idx in top_indices:
            fertilizer_name = fertilizer_encoder.inverse_transform([idx])[0]
            confidence = prediction_proba[idx] * 100
            recommendations.append(FertilizerRecommendation(
                fertilizer=fertilizer_name,
                confidence=round(confidence, 2)
            ))
        
        # Prepare response
        response = PredictionResponse(
            primary_recommendation=predicted_fertilizer,
            confidence=round(max(prediction_proba) * 100, 2),
            all_recommendations=recommendations,
            input_parameters={
                "nitrogen_percent": soil_params.nitrogen_percent,
                "phosphorus_ppm": soil_params.phosphorus_ppm,
                "soil_ph": soil_params.soil_ph,
                "potassium_meq_percent": soil_params.potassium_meq_percent,
                "crop": soil_params.crop
            },
            validation_warnings=validation_warnings
        )
        
        return response
        
    except HTTPException:
        raise
    except Exception as e:
        print(f"Prediction error: {e}")
        print(f"Traceback: {traceback.format_exc()}")
        raise HTTPException(status_code=400, detail=f"Prediction error: {str(e)}")

@app.post("/batch-predict")
async def batch_predict(soil_samples: List[SoilParameters]):
    """
    Predict fertilizer recommendations for multiple soil samples
    """
    if model is None or fertilizer_encoder is None or crop_encoder is None:
        raise HTTPException(status_code=500, detail="Model components not loaded")
    
    if len(soil_samples) > 50:
        raise HTTPException(status_code=400, detail="Maximum 50 samples allowed per batch")
    
    try:
        predictions = []
        
        for i, sample in enumerate(soil_samples):
            try:
                # Validate crop
                if sample.crop not in crop_encoder.classes_:
                    predictions.append({
                        "sample_id": i + 1,
                        "error": f"Unknown crop '{sample.crop}'"
                    })
                    continue
                
                # Encode crop
                crop_encoded = crop_encoder.transform([sample.crop])[0]
                
                # Prepare input data
                input_data = np.array([[
                    sample.nitrogen_percent,
                    sample.phosphorus_ppm, 
                    sample.potassium_meq_percent,
                    sample.soil_ph,
                    crop_encoded
                ]])
                
                # Make prediction
                prediction_encoded = model.predict(input_data)[0]
                prediction_proba = model.predict_proba(input_data)[0]
                predicted_fertilizer = fertilizer_encoder.inverse_transform([prediction_encoded])[0]
                
                predictions.append({
                    "sample_id": i + 1,
                    "primary_recommendation": predicted_fertilizer,
                    "confidence": round(max(prediction_proba) * 100, 2),
                    "input_parameters": {
                        "nitrogen_percent": sample.nitrogen_percent,
                        "phosphorus_ppm": sample.phosphorus_ppm,
                        "soil_ph": sample.soil_ph,
                        "potassium_meq_percent": sample.potassium_meq_percent,
                        "crop": sample.crop
                    }
                })
                
            except Exception as e:
                predictions.append({
                    "sample_id": i + 1,
                    "error": str(e)
                })
        
        return {
            "total_samples": len(soil_samples),
            "predictions": predictions
        }
        
    except Exception as e:
        print(f"Batch prediction error: {e}")
        print(f"Traceback: {traceback.format_exc()}")
        raise HTTPException(status_code=400, detail=f"Batch prediction error: {str(e)}")

# Run the app
if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)