#!/usr/bin/env python3
"""
Simplified FastAPI Crop Yield Prediction API
"""

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, Field
from typing import Optional
import pandas as pd
import numpy as np
import joblib
import warnings
import os
from datetime import datetime

# Import model classes - joblib setup is handled in start.py
try:
    from models import DataPreprocessor
except ImportError as e:
    print(f"Warning: Could not import DataPreprocessor: {e}")
    DataPreprocessor = None

warnings.filterwarnings('ignore')

# Simple predictor class to avoid import conflicts
class CropYieldPredictor:
    """Simplified prediction class that loads Random Forest model."""
    
    def __init__(self, models_dir='models', quiet=False):
        self.models_dir = models_dir
        self.model = None
        self.preprocessor = None
        self.quiet = quiet
        self.fallback_mode = False
        
        if not quiet:
            print(f"🚀 Initializing Random Forest Crop Yield Predictor...")
        
        self.load_models()
        
    def load_models(self):
        """Load Random Forest model and preprocessor."""
        if not self.quiet:
            print("📥 Loading trained model...")
        
        try:
            # Try using the model_loader
            from model_loader import load_models_safely
            model, preprocessor = load_models_safely(self.models_dir)
            
            if model is not None and preprocessor is not None:
                self.model = model
                self.preprocessor = preprocessor
                if not self.quiet:
                    print("  ✅ Models loaded successfully using model_loader")
                return
            else:
                raise Exception("Model loader failed")
                
        except Exception as e:
            if not self.quiet:
                print(f"⚠️ Primary model loading failed: {e}")
                print("⚠️ Switching to fallback mode - limited functionality")
            
            # Fallback: create a simple mock predictor
            self.fallback_mode = True
            self.model = None
            self.preprocessor = self._create_fallback_preprocessor()
            
            if not self.quiet:
                print("✅ Fallback mode initialized")
    
    def _create_fallback_preprocessor(self):
        """Create a simple fallback preprocessor for basic functionality"""
        class FallbackPreprocessor:
            def __init__(self):
                self.label_encoders = {
                    'State': type('MockEncoder', (), {'classes_': ['Punjab', 'Maharashtra', 'Karnataka', 'Gujarat', 'Rajasthan']}),
                    'Crop': type('MockEncoder', (), {'classes_': ['Rice', 'Wheat', 'Cotton', 'Sugarcane', 'Maize']}),
                    'District': type('MockEncoder', (), {'classes_': ['Default District']})
                }
        
        return FallbackPreprocessor()
    
    def predict_yield(self, input_data):
        """Make yield prediction using Random Forest model or fallback."""
        
        if self.fallback_mode:
            # Simple fallback prediction based on basic rules
            try:
                if isinstance(input_data, dict):
                    area = input_data.get('Area', 10)
                    production = input_data.get('Production', 25)
                    rainfall = input_data.get('Annual_Rainfall', 1000)
                    crop = input_data.get('Crop', 'Rice')
                    
                    # Simple formula based on typical crop yields
                    base_yield = {
                        'Rice': 2500, 'Wheat': 3000, 'Cotton': 1200, 
                        'Sugarcane': 60000, 'Maize': 2800
                    }.get(crop.title(), 2000)
                    
                    # Adjust for rainfall
                    rainfall_factor = min(1.2, max(0.8, rainfall / 1000))
                    
                    # Simple prediction
                    prediction = base_yield * rainfall_factor
                    return prediction, None
                else:
                    return "Error: Invalid input format", None
            except Exception as e:
                return f"Fallback Error: {str(e)}", None
        
        try:
            # Convert input to DataFrame
            if isinstance(input_data, dict):
                df = pd.DataFrame([input_data])
            else:
                df = input_data.copy()
            
            # Prepare features
            X, processed_data = self.preprocessor.prepare_features(df)
            
            # Transform data
            X_processed = self.preprocessor.transform(X)
            
            # Make prediction with Random Forest
            try:
                prediction = self.model.predict(X_processed)[0]
                prediction = max(0, prediction)  # Ensure non-negative yield
                return prediction, processed_data
                
            except Exception as e:
                return f"Error: {str(e)}", None
            
        except Exception as e:
            return f"Error: {str(e)}", None
    
    def get_crop_options(self):
        """Get available crop options from the preprocessor."""
        if hasattr(self.preprocessor, 'label_encoders') and 'Crop' in self.preprocessor.label_encoders:
            return list(self.preprocessor.label_encoders['Crop'].classes_)
        return []
    
    def get_state_options(self):
        """Get available state options from the preprocessor."""
        if hasattr(self.preprocessor, 'label_encoders') and 'State' in self.preprocessor.label_encoders:
            return list(self.preprocessor.label_encoders['State'].classes_)
        return []
    
    def get_season_options(self):
        """Get available season options."""
        return ['Kharif', 'Rabi', 'Summer', 'Whole Year', 'Autumn', 'Winter', 'Total']

# Initialize FastAPI app
app = FastAPI(
    title="🌾 SIH Crop Yield Prediction API",
    description="""**Smart India Hackathon Project**
    
    Predict crop yields using advanced Machine Learning models including Random Forest, XGBoost, and PyTorch neural networks.
    
    **Features:**
    - 🤖 Multiple ML models (Random Forest, XGBoost, PyTorch)
    - 🌱 Supports major Indian crops (Rice, Wheat, Cotton, etc.)
    - 🏛️ State-wise predictions across India
    - 🔄 Intelligent fallback system for reliability
    - ⚡ Fast predictions (~100-500ms)
    
    **Perfect for:**
    - Farmers planning crop yields
    - Agricultural consultants
    - Government agricultural departments
    - Research and academic studies
    """,
    version="2.0.0",
    docs_url="/docs",
    redoc_url="/redoc",
    contact={
        "name": "SIH Team - Crop Yield Prediction",
        "url": "https://github.com/AshrafGalibShaik/SIH-2",
    },
    license_info={
        "name": "MIT License",
        "url": "https://opensource.org/licenses/MIT",
    },
)

# Pydantic models for request and response
class CropPredictionRequest(BaseModel):
    year: int = Field(..., description="Crop year (e.g., 2024)", example=2024)
    state: str = Field(..., description="State name", example="Punjab")
    crop: str = Field(..., description="Crop name", example="Rice")
    season: str = Field(..., description="Season", example="Kharif")
    area: float = Field(..., description="Area in hectares", example=10.0)
    production: float = Field(..., description="Production in tons", example=25.0)
    rainfall: Optional[float] = Field(1000.0, description="Annual rainfall in mm", example=1200)
    fertilizer: Optional[float] = Field(50.0, description="Fertilizer usage in kg", example=75)
    pesticide: Optional[float] = Field(5.0, description="Pesticide usage in kg", example=8)

class CropPredictionResponse(BaseModel):
    model: str = Field(..., description="Model used for prediction", example="Random Forest")
    predicted_yield: str = Field(..., description="Predicted yield with units", example="2017.7 kg/hectare")
    total_expected_production: str = Field(..., description="Total expected production with units", example="20.18 tons")
    assessment: str = Field(..., description="Yield assessment", example="Good yield expected")

class ErrorResponse(BaseModel):
    error: str = Field(..., description="Error message")

# Global predictor instance - initialize on startup
predictor = None

@app.on_event("startup")
async def startup_event():
    """Initialize ML models on startup for better performance"""
    global predictor
    print("🚀 Initializing ML models on startup...")
    try:
        predictor = CropYieldPredictor(quiet=False)
        print("✅ Startup initialization completed successfully")
    except Exception as e:
        print(f"⚠️  Startup initialization failed: {e}")
        print("📋 API will use fallback mode for predictions")
        predictor = "failed"

def get_predictor():
    """Get the pre-initialized predictor"""
    global predictor
    if predictor is None:
        # This shouldn't happen with startup event, but fallback just in case
        try:
            predictor = CropYieldPredictor(quiet=True)
            print("✅ Crop Yield Predictor initialized (fallback)")
        except Exception as e:
            print(f"❌ Failed to initialize predictor: {e}")
            predictor = "failed"
    return predictor if predictor != "failed" else None


@app.get("/")
async def root():
    """Root endpoint"""
    return {
        "message": "Crop Yield Prediction API. Use /docs for interactive API documentation.",
        "version": "1.0.0",
        "endpoints": {
            "predict": "/predict",
            "health": "/health",
            "docs": "/docs",
            "available_options": "/available-options"
        }
    }

@app.get("/health")
async def health_check():
    """Health check endpoint"""
    predictor_instance = get_predictor()
    return {
        "status": "healthy",
        "timestamp": datetime.now().isoformat(),
        "model_loaded": predictor_instance is not None
    }

@app.post("/predict", response_model=CropPredictionResponse, responses={400: {"model": ErrorResponse}})
async def predict_yield(request: CropPredictionRequest):
    """
    Predict crop yield based on input parameters
    
    Input format:
    {
      "year": 2024,
      "state": "Punjab",
      "crop": "Rice",
      "season": "Kharif",
      "area": 10.0,
      "production": 25.0,
      "rainfall": 1200,
      "fertilizer": 75,
      "pesticide": 8
    }
    
    Returns prediction with model type, predicted yield, total production, and assessment.
    """
    try:
        predictor_instance = get_predictor()
        if predictor_instance is None:
            raise HTTPException(status_code=500, detail="Predictor not initialized. Please check if trained models are available.")
        
        # Convert request to internal format
        input_data = {
            'Crop_Year': request.year,
            'State': request.state,
            'District': "Unknown",  # Default district
            'Crop': request.crop,
            'Season': request.season,
            'Area': request.area,
            'Production': request.production,
            'Annual_Rainfall': request.rainfall,
            'Fertilizer': request.fertilizer,
            'Pesticide': request.pesticide
        }
        
        # Make prediction
        prediction, _ = predictor_instance.predict_yield(input_data)
        
        if isinstance(prediction, str) and 'Error' in prediction:
            raise HTTPException(status_code=400, detail=prediction)
        
        # Calculate total production
        total_production = (prediction * request.area) / 1000.0  # Convert to tons
        
        # Determine assessment
        if prediction > 3000:
            assessment = "Excellent yield expected"
        elif prediction > 2000:
            assessment = "Good yield expected"
        elif prediction > 1000:
            assessment = "Moderate yield expected"
        else:
            assessment = "Low yield expected"
        
        # Format response
        model_type = "Random Forest" if not predictor_instance.fallback_mode else "Fallback Model (Rule-based)"
        response = CropPredictionResponse(
            model=model_type,
            predicted_yield=f"{round(prediction, 2)} kg/hectare",
            total_expected_production=f"{round(total_production, 2)} tons",
            assessment=assessment
        )
        
        return response
        
    except HTTPException:
        raise  # Re-raise HTTP exceptions
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Prediction failed: {str(e)}")

@app.get("/available-options")
async def get_available_options():
    """Get available crops, states, and seasons"""
    try:
        predictor_instance = get_predictor()
        if predictor_instance is None:
            raise HTTPException(status_code=500, detail="Predictor not initialized")
        
        return {
            "states": predictor_instance.get_state_options()[:10],  # Limit to first 10 for readability
            "crops": predictor_instance.get_crop_options()[:10],    # Limit to first 10 for readability
            "seasons": predictor_instance.get_season_options(),
            "note": "This shows first 10 states and crops. All are supported in predictions."
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

# Note: Server startup is handled by start.py for Railway deployment
# This prevents conflicts between different startup methods