import os
import joblib
import torch
import torch.nn as nn
import numpy as np
from flask import Flask, request, jsonify
from flask_cors import CORS

# --- 1. Define the same PyTorch Model Architecture ---
# This class must be identical to the one in your training script.
class FraudClassifier(nn.Module):
    def __init__(self):
        super(FraudClassifier, self).__init__()
        self.layer1 = nn.Linear(2, 16) # 2 input features
        self.layer2 = nn.Linear(16, 8)
        self.layer3 = nn.Linear(8, 1)  # 1 output for binary classification
        self.relu = nn.ReLU()
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        x = self.relu(self.layer1(x))
        x = self.relu(self.layer2(x))
        x = self.sigmoid(self.layer3(x))
        return x

# --- 2. Initialize Flask App and Load Assets ---
app = Flask(__name__)
CORS(app) # Enable Cross-Origin Resource Sharing

# The model and scaler are loaded once when the application starts.
model = None
scaler = None

try:
    # Load the PyTorch model
    model = FraudClassifier()
    # Use map_location=torch.device('cpu') for compatibility if the server doesn't have a GPU
    model.load_state_dict(torch.load("realistic_fraud_model.pth", map_location=torch.device('cpu')))
    model.eval() # Set model to evaluation mode
    print("✅ PyTorch model loaded successfully.")

    # Load the scaler
    scaler = joblib.load("scaler.pkl")
    print("✅ Scaler loaded successfully.")

except Exception as e:
    print(f"❌ An error occurred during asset loading: {e}")

# --- 3. API Endpoints ---
@app.route('/', methods=['GET'])
def home():
    """Root endpoint to check API status."""
    return jsonify({
        "status": "API is running",
        "model_loaded": model is not None,
        "scaler_loaded": scaler is not None
    })

@app.route('/predict', methods=['POST'])
def predict():
    """Receives feature data and returns a fraud prediction."""
    if not model or not scaler:
        return jsonify({"error": "Model or scaler not loaded."}), 500

    try:
        data = request.get_json(force=True)
        # Expected input: {"features": [value_ether, fee_ether]}
        features = np.array(data['features']).reshape(1, -1)

        # Apply the same scaling as in training
        scaled_features = scaler.transform(features)

        # Convert to PyTorch tensor
        features_tensor = torch.tensor(scaled_features, dtype=torch.float32)

        # Make prediction
        with torch.no_grad(): # Disable gradient calculation for inference
            probability = model(features_tensor)

        prob_fraud = probability.item() # Get the float value from the tensor
        prediction = 1 if prob_fraud > 0.5 else 0

        # Return the result
        return jsonify({
            'prediction': prediction,
            'is_fraud': bool(prediction == 1),
            'prediction_probability': {
                'not_fraud': 1.0 - prob_fraud,
                'fraud': prob_fraud
            }
        })
    except Exception as e:
        return jsonify({"error": f"An error occurred during prediction: {str(e)}"}), 400

if __name__ == '__main__':
    port = int(os.environ.get("PORT", 8080))
    app.run(host='0.0.0.0', port=port)