import os
from pathlib import Path

os.environ.setdefault("KERAS_BACKEND", "jax")

import keras
import numpy as np


keras.utils.set_random_seed(7)
rng = np.random.default_rng(42)

rows = 6000
amount = rng.lognormal(mean=3.2, sigma=0.9, size=rows).astype("float32")
hour = rng.integers(0, 24, size=rows).astype("float32")
country_risk = rng.beta(1.2, 8.0, size=rows).astype("float32")
merchant_risk = rng.beta(1.4, 7.0, size=rows).astype("float32")
device_age_days = rng.exponential(scale=20.0, size=rows).astype("float32")
chargeback_count = rng.poisson(0.12, size=rows).astype("float32")
night_purchase = ((hour < 6) | (hour > 22)).astype("float32")

amount_log = np.log1p(amount)
amount_z = (amount_log - amount_log.mean()) / amount_log.std()
fraud_score = (
    -5.5
    + 1.2 * amount_z
    + 6.4 * country_risk
    + 5.7 * merchant_risk
    + 1.2 * night_purchase
    + 1.6 * (device_age_days < 2)
    + 1.5 * (chargeback_count > 0)
)
fraud_probability = 1.0 / (1.0 + np.exp(-fraud_score))
target = rng.binomial(1, fraud_probability).astype("float32")

features = np.column_stack(
    [
        amount_log,
        hour / 23.0,
        country_risk,
        merchant_risk,
        np.log1p(device_age_days),
        chargeback_count,
        night_purchase,
    ]
).astype("float32")

order = rng.permutation(rows)
train_size = int(rows * 0.8)
train_index = order[:train_size]
validation_index = order[train_size:]
x_train = features[train_index]
y_train = target[train_index]
x_validation = features[validation_index]
y_validation = target[validation_index]

negative = float((y_train == 0).sum())
positive = float((y_train == 1).sum())
total = negative + positive
class_weight = {
    0: total / (2.0 * negative),
    1: total / (2.0 * positive),
}

normalizer = keras.layers.Normalization()
normalizer.adapt(x_train)

model = keras.Sequential(
    [
        keras.layers.Input(shape=(x_train.shape[1],)),
        normalizer,
        keras.layers.Dense(32, activation="relu"),
        keras.layers.Dropout(0.2),
        keras.layers.Dense(16, activation="relu"),
        keras.layers.Dense(1, activation="sigmoid"),
    ]
)

model.compile(
    optimizer=keras.optimizers.Adam(learning_rate=0.005),
    loss=keras.losses.BinaryCrossentropy(),
    metrics=[
        keras.metrics.AUC(curve="PR", name="pr_auc"),
        keras.metrics.Precision(name="precision"),
        keras.metrics.Recall(name="recall"),
    ],
    jit_compile="auto",
)

history = model.fit(
    x_train,
    y_train,
    validation_data=(x_validation, y_validation),
    epochs=8,
    batch_size=128,
    class_weight=class_weight,
    verbose=0,
)

validation_probability = model.predict(
    x_validation,
    batch_size=256,
    verbose=0,
).ravel()
threshold = 0.35
validation_prediction = validation_probability >= threshold
actual_positive = y_validation == 1
actual_negative = y_validation == 0
true_positive = int(np.logical_and(validation_prediction, actual_positive).sum())
false_positive = int(np.logical_and(validation_prediction, actual_negative).sum())
false_negative = int(np.logical_and(~validation_prediction, actual_positive).sum())
true_negative = int(np.logical_and(~validation_prediction, actual_negative).sum())
threshold_precision = true_positive / max(true_positive + false_positive, 1)
threshold_recall = true_positive / max(true_positive + false_negative, 1)

model_path = Path("fraud_classifier.keras")
model.save(model_path, overwrite=True)
loaded_model = keras.saving.load_model(model_path)
sample_transactions = np.array(
    [
        [np.log1p(38.20), 14 / 23.0, 0.04, 0.08, np.log1p(120.0), 0.0, 0.0],
        [np.log1p(620.00), 2 / 23.0, 0.54, 0.61, np.log1p(0.5), 2.0, 1.0],
    ],
    dtype="float32",
)
sample_scores = loaded_model.predict(sample_transactions, verbose=0).ravel()

final_epoch = len(history.history["loss"])
final_val_pr_auc = history.history["val_pr_auc"][-1]
final_val_precision = history.history["val_precision"][-1]
final_val_recall = history.history["val_recall"][-1]

print(f"backend: {keras.config.backend()}")
print(f"train rows: {len(x_train)}")
print(f"validation rows: {len(x_validation)}")
print(f"training fraud rate: {y_train.mean():.3f}")
print(f"class weights: legit={class_weight[0]:.2f}, fraud={class_weight[1]:.2f}")
print(f"epochs completed: {final_epoch}")
print(f"val pr_auc: {final_val_pr_auc:.3f}")
print(f"val precision@0.5: {final_val_precision:.3f}")
print(f"val recall@0.5: {final_val_recall:.3f}")
print(
    f"threshold 0.35 counts: tp={true_positive}, "
    f"fp={false_positive}, fn={false_negative}, tn={true_negative}"
)
print(
    "threshold 0.35 metrics: "
    f"precision={threshold_precision:.3f}, recall={threshold_recall:.3f}"
)
print(f"saved model: {model_path.name}")
print(
    "sample fraud scores: "
    f"routine={sample_scores[0]:.3f}, suspicious={sample_scores[1]:.3f}"
)