from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import joblib
import numpy as np
from propy import AAComposition
from sklearn.preprocessing import MinMaxScaler

# Initialize FastAPI app
app = FastAPI()

# Load trained SVM model and scaler
model = joblib.load("SVM.joblib")
scaler = MinMaxScaler()

# Define request schema
class SequenceInput(BaseModel):
    sequence: str

def extract_features(sequence: str):
    """Extract Amino Acid Composition (AAC) features and normalize them."""
    try:
        aac = np.array(list(AAComposition.CalculateAADipeptideComposition(sequence)), dtype=float)
        normalized_features = scaler.fit_transform([aac])  # Don't use fit_transform(), only transform()
        return normalized_features
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"Feature extraction failed: {str(e)}")

@app.post("/predict/")
def predict(input_data: SequenceInput):
    """Predict AMP vs Non-AMP from protein sequence."""
    features = extract_features(input_data.sequence)
    prediction = model.predict(features)[0]
    result = "AMP" if prediction == 1 else "Non-AMP"
    return {"prediction": result}

# Run using: uvicorn script_name:app --host 0.0.0.0 --port 8000 --reload