import os
import torch
import torch.nn as nn
from flask import Flask, request, jsonify

# --- 1. Define the exact same Model Architecture ---
# This is necessary to load the saved state_dict
class FraudClassifier(nn.Module):
    def __init__(self):
        super(FraudClassifier, self).__init__()
        self.layer1 = nn.Linear(2, 16)
        self.layer2 = nn.Linear(16, 8)
        self.layer3 = nn.Linear(8, 1)
        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 App and Load Assets ---
app = Flask(__name__)
model = None
scaler_params = None

try:
    # Load scaler parameters
    scaler_params = torch.load("scaler_params.pt")
    mean = scaler_params['mean']
    std = scaler_params['std']
    print("✅ PyTorch scaler parameters loaded.")

    # Load the trained model
    model = FraudClassifier()
    model.load_state_dict(torch.load("realistic_fraud_model.pth"))
    model.eval() # Set to evaluation mode
    print("✅ PyTorch model loaded successfully.")

except Exception as e:
    print(f"❌ Error loading assets: {e}")

# --- 3. Define 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_params is not None
    })

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

    try:
        data = request.get_json(force=True)
        # Expected input: {"features": [value_ether, fee_ether]}
        features = data['features']
        
        # Convert to a PyTorch tensor
        features_tensor = torch.tensor(features, dtype=torch.float32)

        # Apply the scaling transformation
        scaled_tensor = (features_tensor - mean) / std

        # Get model prediction
        with torch.no_grad():
            probability = model(scaled_tensor)

        prob_fraud = probability.item()
        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)