import numpy as np
import torch
import torch.nn as nn
from flask import Flask, request, jsonify
from sklearn.preprocessing import MinMaxScaler
import joblib

# Define the autoencoder model
class Autoencoder(nn.Module):
    def __init__(self, input_dim):
        super(Autoencoder, self).__init__()
        self.encoder = nn.Sequential(
            nn.Linear(input_dim, 64),
            nn.BatchNorm1d(64),
            nn.ReLU(),
            nn.Dropout(0.2)
        )
        self.decoder = nn.Sequential(
            nn.Linear(64, input_dim),
            nn.BatchNorm1d(input_dim),
            nn.Sigmoid()
        )

    def forward(self, x):
        encoded = self.encoder(x)
        decoded = self.decoder(encoded)
        return decoded

# Load trained model
input_dim = 29  # Assuming 29 features (V1-V28 + Amount)
model = Autoencoder(input_dim)
model.load_state_dict(torch.load("trained_autoencoder.pth"))
model.eval()

# Load the MinMaxScaler
scaler = joblib.load("scaler.pkl")

# Create Flask app
app = Flask(__name__)

@app.route('/')
def home():
    return "Credit Card Fraud Detection API is Running!"

@app.route('/predict', methods=['POST'])
def predict():
    try:
        # Get JSON input
        data = request.get_json()
        X_input = np.array(data['features']).reshape(1, -1)  # Ensure it's in the right shape

        # Scale input data
        X_scaled = scaler.transform(X_input)

        # Convert to PyTorch tensor
        X_tensor = torch.tensor(X_scaled, dtype=torch.float32)

        # Get reconstruction error
        recon = model(X_tensor).detach().numpy()
        recon_error = np.mean((recon - X_scaled) ** 2)

        # Use threshold to classify as fraud (1) or normal (0)
        threshold = 0.01  # Adjust this based on previous experiments
        prediction = 1 if recon_error > threshold else 0

        return jsonify({'fraud_probability': float(recon_error), 'is_fraud': prediction})
    
    except Exception as e:
        return jsonify({'error': str(e)})

if __name__ == '__main__':
    app.run(debug=True)