from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
import joblib
import pandas as pd
import numpy as np
import logging

# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = FastAPI()

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

# Load your models
try:
    preprocessor = joblib.load('preprocessor.joblib')
    numerical_model = joblib.load('numerical_model.joblib')
    categorical_model = joblib.load('categorical_model.joblib')
    label_encoders = joblib.load('label_encoders.joblib')
    logger.info("Models loaded successfully")
except Exception as e:
    logger.error(f"Error loading models: {str(e)}")
    raise

# Define target columns (make sure these match your training data exactly)
numerical_targets = ['Nitrogen (N) Need', 'Phosphorus (P2O5) Need', 'Potassium (K2O) Need', 
                     'Organic Matter Need', 'Lime Need', 'Lime Application - Requirement', 
                     'Organic Matter Application - Requirement', '1st Application - Requirement (1)',
                     '1st Application - Requirement (2)', '2nd Application - Requirement (1)',
                     '1st Application - Requirement (1)_2', '1st Application - Requirement (2)_4', 
                     '2nd Application - Requirement (1)_7', '1st Application - Requirement (1)_22', 
                     '1st Application - Requirement (2)_24', '2nd Application - Requirement (1)_27']

categorical_targets = ['Lime Application - Instruction', 'Lime Application', 
                       'Organic Matter Application - Instruction', 'Organic Matter Application', 
                       '1st Application', '1st Application - Type fertilizer (1)', 
                       '1st Application - Type fertilizer (2)', '2nd Application', 
                       '2nd Application - Type fertilizer (1)', '1st Application_1', 
                       '1st Application - Type fertilizer (1)_3', '1st Application - Type fertilizer (2)_5', 
                       '2nd Application_6', '1st Application_21', '1st Application - Type fertilizer (1)_23', 
                       '1st Application - Type fertilizer (2)_25', '2nd Application_26', 
                       '2nd Application - Type fertilizer (1)_28']

class InputData(BaseModel):
    crop_name: str
    target_yield: float
    field_size: float
    ph_water: float
    organic_carbon: float
    total_nitrogen: float
    phosphorus_m3: float
    potassium_exch: float

@app.get("/")
async def root():
    return {"message": "Fertilizer Prediction API is running"}

@app.post("/predict")
async def predict_fertilizer(input_data: InputData):
    try:
        # Convert input data to DataFrame
        sample_data = pd.DataFrame({
            'Crop Name': [input_data.crop_name],
            'Target Yield': [input_data.target_yield],
            'Field Size': [input_data.field_size],
            'pH (water)': [input_data.ph_water],
            'Organic Carbon': [input_data.organic_carbon],
            'Total Nitrogen': [input_data.total_nitrogen],
            'Phosphorus (M3)': [input_data.phosphorus_m3],
            'Potassium (exch.)': [input_data.potassium_exch]
        })

        # Transform input data
        sample_transformed = preprocessor.transform(sample_data)
        logger.info("Input data transformed successfully")

        # Make predictions
        numerical_pred = numerical_model.predict(sample_transformed)
        categorical_pred = categorical_model.predict(sample_transformed)
        logger.info("Predictions made successfully")

        # Decode categorical predictions
        categorical_pred_decoded = {col: label_encoders[col].inverse_transform(categorical_pred[:, i]).tolist()
                                    for i, col in enumerate(categorical_targets)}

        # Combine predictions
        predictions = pd.concat([
            pd.DataFrame(numerical_pred, columns=numerical_targets),
            pd.DataFrame(categorical_pred_decoded)
        ], axis=1)

        # Convert predictions to dictionary
        predictions_dict = predictions.to_dict(orient='records')[0]
        logger.info("Predictions processed successfully")

        return predictions_dict
    except Exception as e:
        logger.error(f"Error in prediction: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

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