"""
app.py - Gradio interface for PyTorch fraud detection demo

Requirements:
    pip install torch scikit-learn pandas numpy gradio
Run:
    python app.py
Then open the URL that Gradio prints (usually http://127.0.0.1:7860)
"""

import os
import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler

import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader

import gradio as gr

# =========================
# 1. Load or create data
# =========================

CSV_PATH = "synthetic_fraud_transactions_1000.csv"  # adjust if needed


def load_or_create_data(path: str) -> pd.DataFrame:
    if os.path.exists(path):
        print(f"[INFO] Loading existing CSV: {path}")
        return pd.read_csv(path)

    print(f"[INFO] CSV not found at {path}. Generating synthetic data...")

    np.random.seed(42)
    n = 1000

    # Generate synthetic features
    amount = np.random.lognormal(mean=3, sigma=0.7, size=n)  # skewed positive
    amount_log = np.log1p(amount)

    is_foreign_currency = np.random.binomial(1, 0.1, size=n)

    txn_hour_of_day = np.random.randint(0, 24, size=n)

    # Time since last txn in seconds (0 to 2 days, skewed toward small)
    time_since_last_txn_sec = np.random.exponential(scale=2 * 60 * 60, size=n)
    time_since_last_txn_sec = np.clip(time_since_last_txn_sec, 10, 2 * 24 * 60 * 60).astype(int)

    # Distance between transactions in km (0–10,000)
    geo_distance_from_last_txn_km = np.random.exponential(scale=50, size=n)
    geo_distance_from_last_txn_km = np.clip(geo_distance_from_last_txn_km, 0, 10000)

    # Transactions and spend in windows
    txn_count_last_10min = np.random.poisson(lam=0.3, size=n)
    total_spend_last_24h = np.random.gamma(shape=2.0, scale=50.0, size=n)

    is_new_merchant_for_card = np.random.binomial(1, 0.2, size=n)
    merchant_fraud_rate_30d = np.random.beta(a=1.2, b=50, size=n)  # mostly low rates

    is_new_device_for_card = np.random.binomial(1, 0.15, size=n)
    num_cards_on_device_24h = np.random.poisson(lam=1.1, size=n)

    # Simple fraud label: higher probability when risky patterns occur
    base_prob = 0.02
    risk_score = (
        0.03 * (time_since_last_txn_sec < 120)
        + 0.04 * (geo_distance_from_last_txn_km > 500)
        + 0.05 * is_new_device_for_card
        + 0.04 * is_new_merchant_for_card
        + 0.03 * (num_cards_on_device_24h > 3)
        + 0.03 * is_foreign_currency
    )
    fraud_prob = np.clip(base_prob + risk_score, 0, 0.9)
    fraud_label = (np.random.rand(n) < fraud_prob).astype(int)

    df = pd.DataFrame(
        {
            "amount": amount.round(2),
            "amount_log": amount_log.round(4),
            "is_foreign_currency": is_foreign_currency,
            "txn_hour_of_day": txn_hour_of_day,
            "time_since_last_txn_sec": time_since_last_txn_sec,
            "geo_distance_from_last_txn_km": geo_distance_from_last_txn_km.round(2),
            "txn_count_last_10min": txn_count_last_10min,
            "total_spend_last_24h": total_spend_last_24h.round(2),
            "is_new_merchant_for_card": is_new_merchant_for_card,
            "merchant_fraud_rate_30d": merchant_fraud_rate_30d.round(4),
            "is_new_device_for_card": is_new_device_for_card,
            "num_cards_on_device_24h": num_cards_on_device_24h,
            "fraud_label": fraud_label,
        }
    )

    # Save so future runs reuse the same data
    df.to_csv(path, index=False)
    print(f"[INFO] Synthetic dataset saved to: {path}")

    return df


df = load_or_create_data(CSV_PATH)

feature_cols = [
    "amount",
    "amount_log",
    "is_foreign_currency",
    "txn_hour_of_day",
    "time_since_last_txn_sec",
    "geo_distance_from_last_txn_km",
    "txn_count_last_10min",
    "total_spend_last_24h",
    "is_new_merchant_for_card",
    "merchant_fraud_rate_30d",
    "is_new_device_for_card",
    "num_cards_on_device_24h",
]

target_col = "fraud_label"

X = df[feature_cols].values.astype(np.float32)
y = df[target_col].values.astype(np.float32)

X_train, X_val, y_train, y_val = train_test_split(
    X, y, test_size=0.2, random_state=42, stratify=y
)

scaler = StandardScaler()
X_train_scaled = scaler.fit_transform(X_train)
X_val_scaled = scaler.transform(X_val)

# =========================
# 2. PyTorch Dataset & Model
# =========================


class FraudDataset(Dataset):
    def __init__(self, X, y):
        self.X = torch.from_numpy(X).float()
        self.y = torch.from_numpy(y).float()

    def __len__(self):
        return len(self.X)

    def __getitem__(self, idx):
        return self.X[idx], self.y[idx]


train_dataset = FraudDataset(X_train_scaled, y_train)
val_dataset = FraudDataset(X_val_scaled, y_val)

train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=64, shuffle=False)


class FraudNet(nn.Module):
    def __init__(self, input_dim):
        super(FraudNet, self).__init__()
        self.net = nn.Sequential(
            nn.Linear(input_dim, 64),
            nn.BatchNorm1d(64),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(64, 32),
            nn.BatchNorm1d(32),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(32, 1),  # binary logit
        )

    def forward(self, x):
        return self.net(x).squeeze(1)


input_dim = len(feature_cols)
model = FraudNet(input_dim)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

# Handle class imbalance with pos_weight
pos_weight_value = (len(y_train) - y_train.sum()) / (y_train.sum() + 1e-6)
pos_weight = torch.tensor(pos_weight_value, dtype=torch.float32).to(device)
criterion = nn.BCEWithLogitsLoss(pos_weight=pos_weight)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3, weight_decay=1e-4)


def train_one_epoch(model, loader, optimizer, criterion, device):
    model.train()
    running_loss = 0.0
    for X_batch, y_batch in loader:
        X_batch = X_batch.to(device)
        y_batch = y_batch.to(device)

        optimizer.zero_grad()
        logits = model(X_batch)
        loss = criterion(logits, y_batch)
        loss.backward()
        optimizer.step()

        running_loss += loss.item() * X_batch.size(0)
    return running_loss / len(loader.dataset)


def evaluate(model, loader, criterion, device):
    model.eval()
    running_loss = 0.0
    all_logits = []
    all_labels = []
    with torch.no_grad():
        for X_batch, y_batch in loader:
            X_batch = X_batch.to(device)
            y_batch = y_batch.to(device)
            logits = model(X_batch)
            loss = criterion(logits, y_batch)
            running_loss += loss.item() * X_batch.size(0)
            all_logits.append(logits.cpu())
            all_labels.append(y_batch.cpu())
    all_logits = torch.cat(all_logits)
    all_labels = torch.cat(all_labels)
    probs = torch.sigmoid(all_logits)
    preds = (probs >= 0.5).int()
    accuracy = (preds == all_labels.int()).float().mean().item()
    return running_loss / len(loader.dataset), accuracy


# =========================
# 3. Train the model once
# =========================

EPOCHS = 15  # keep small so startup is fast

for epoch in range(1, EPOCHS + 1):
    train_loss = train_one_epoch(model, train_loader, optimizer, criterion, device)
    val_loss, val_acc = evaluate(model, val_loader, criterion, device)
    print(
        f"Epoch {epoch:02d} | "
        f"Train Loss: {train_loss:.4f} | "
        f"Val Loss: {val_loss:.4f} | "
        f"Val Acc: {val_acc:.4f}"
    )

model.eval()  # set evaluation mode

# =========================
# 4. Inference function
# =========================


def predict_fraud(
    amount,
    is_foreign_currency,
    txn_hour_of_day,
    time_since_last_txn_min,
    geo_distance_from_last_txn_km,
    txn_count_last_10min,
    total_spend_last_24h,
    is_new_merchant_for_card,
    merchant_fraud_rate_30d,
    is_new_device_for_card,
    num_cards_on_device_24h,
):
    """
    This function will be called by the Gradio interface.
    It builds a feature vector, scales it, runs the model,
    and returns a probability + label string.
    """

    # Derived features
    amount_log = np.log1p(amount)
    time_since_last_txn_sec = time_since_last_txn_min * 60.0

    # Build feature vector in the SAME ORDER as feature_cols
    x = np.array(
        [
            amount,
            amount_log,
            float(is_foreign_currency),
            float(txn_hour_of_day),
            time_since_last_txn_sec,
            geo_distance_from_last_txn_km,
            float(txn_count_last_10min),
            total_spend_last_24h,
            float(is_new_merchant_for_card),
            merchant_fraud_rate_30d,
            float(is_new_device_for_card),
            float(num_cards_on_device_24h),
        ],
        dtype=np.float32,
    ).reshape(1, -1)

    # Scale
    x_scaled = scaler.transform(x)

    # To torch
    x_tensor = torch.from_numpy(x_scaled).float().to(device)

    with torch.no_grad():
        logit = model(x_tensor)
        prob = torch.sigmoid(logit).item()

    label = "FRAUD" if prob >= 0.5 else "LEGIT"

    explanation = (
        f"Predicted probability of fraud: **{prob:.3f}**\n\n"
        f"Model classification: **{label}** (threshold = 0.5)"
    )

    return label, prob, explanation


# =========================
# 5. Gradio UI
# =========================

with gr.Blocks(title="Fraud Detection Demo (PyTorch + Gradio)") as demo:
    gr.Markdown(
        "# 💳 Fraud Detection Demo\n"
        "Play with transaction features and see how the model classifies them."
    )

    with gr.Row():
        with gr.Column():
            amount = gr.Slider(
                minimum=1,
                maximum=2000,
                value=50,
                step=1,
                label="Transaction amount (in card currency)",
            )
            is_foreign_currency = gr.Checkbox(
                value=False, label="Foreign currency?"
            )
            txn_hour_of_day = gr.Slider(
                minimum=0,
                maximum=23,
                value=14,
                step=1,
                label="Transaction hour of day (0–23)",
            )
            time_since_last_txn_min = gr.Slider(
                minimum=0.1,
                maximum=2880,
                value=60,
                step=1,
                label="Time since last transaction (minutes, up to 2 days)",
            )
            geo_distance_from_last_txn_km = gr.Slider(
                minimum=0,
                maximum=10000,
                value=10,
                step=1,
                label="Distance from last transaction (km)",
            )
            txn_count_last_10min = gr.Slider(
                minimum=0,
                maximum=20,
                value=1,
                step=1,
                label="Number of transactions in last 10 minutes",
            )
        with gr.Column():
            total_spend_last_24h = gr.Slider(
                minimum=0,
                maximum=10000,
                value=200,
                step=10,
                label="Total spend in last 24h",
            )
            is_new_merchant_for_card = gr.Checkbox(
                value=False, label="Is this a new merchant for this card?"
            )
            merchant_fraud_rate_30d = gr.Slider(
                minimum=0.0,
                maximum=0.5,
                value=0.02,
                step=0.01,
                label="Merchant fraud rate (last 30 days)",
            )
            is_new_device_for_card = gr.Checkbox(
                value=False, label="Is this a new device for this card?"
            )
            num_cards_on_device_24h = gr.Slider(
                minimum=1,
                maximum=20,
                value=1,
                step=1,
                label="Number of different cards on this device in last 24h",
            )

    predict_btn = gr.Button("Predict Fraud Risk")

    label_out = gr.Textbox(label="Model Label (FRAUD / LEGIT)", interactive=False)
    prob_out = gr.Number(label="Fraud Probability", precision=4, interactive=False)
    explanation_out = gr.Markdown(label="Explanation")

    predict_btn.click(
        fn=predict_fraud,
        inputs=[
            amount,
            is_foreign_currency,
            txn_hour_of_day,
            time_since_last_txn_min,
            geo_distance_from_last_txn_km,
            txn_count_last_10min,
            total_spend_last_24h,
            is_new_merchant_for_card,
            merchant_fraud_rate_30d,
            is_new_device_for_card,
            num_cards_on_device_24h,
        ],
        outputs=[label_out, prob_out, explanation_out],
    )

if __name__ == "__main__":
    demo.launch()