import gradio as gr
import joblib
import joblib
import numpy as np
from pydantic import BaseModel
import propy
from sklearn.preprocessing import MinMaxScaler

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

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

def extract_features(sequence):
    """Calculate AAC, Dipeptide Composition and normalize features."""
    # Calculate Amino Acid Composition (AAC)
    aac = propy.AAComposition.CalculateAAC(sequence)
    
    # Calculate Dipeptide Composition
    dipeptide_comp = propy.AAComposition.CalculateAADipeptideComposition(sequence)
    
    # Combine both features (AAC and Dipeptide Composition)
    features = np.concatenate((aac, dipeptide_comp))
    
    # Min-Max Normalization
    scaler = MinMaxScaler()
    normalized_features = scaler.fit_transform(features.reshape(-1, 1)).flatten()
    
    return normalized_features

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