import csv
import os
from pathlib import Path

os.environ["KERAS_BACKEND"] = "jax"

import keras
import numpy as np


keras.utils.set_random_seed(13)

model_path = Path("batch-score-model.keras")
input_path = Path("input_batch.csv")
output_path = Path("predictions.csv")

x_train = np.array(
    [
        [0.10, 0.20, 0.30],
        [0.90, 0.70, 0.60],
        [0.30, 0.80, 0.50],
        [0.75, 0.25, 0.90],
        [0.15, 0.35, 0.40],
        [0.88, 0.66, 0.55],
    ],
    dtype="float32",
)
y_train = (x_train.sum(axis=1) > 1.4).astype("float32")

model = keras.Sequential(
    [
        keras.Input(shape=(3,), name="features"),
        keras.layers.Dense(8, activation="relu"),
        keras.layers.Dense(1, activation="sigmoid", name="probability"),
    ]
)
model.compile(
    optimizer=keras.optimizers.Adam(learning_rate=0.04),
    loss=keras.losses.BinaryCrossentropy(),
)
model.fit(x_train, y_train, epochs=20, batch_size=3, verbose=0)
model.save(model_path)

with input_path.open("w", newline="") as handle:
    writer = csv.writer(handle)
    writer.writerow(["account_id", "feature_a", "feature_b", "feature_c"])
    writer.writerow(["A1001", "0.12", "0.44", "0.20"])
    writer.writerow(["A1002", "0.80", "0.70", "0.65"])
    writer.writerow(["A1003", "0.35", "0.50", "0.55"])

loaded_model = keras.models.load_model(model_path)

account_ids = []
features = []
with input_path.open(newline="") as handle:
    reader = csv.DictReader(handle)
    for row in reader:
        account_ids.append(row["account_id"])
        features.append([float(row["feature_a"]), float(row["feature_b"]), float(row["feature_c"])])

feature_array = np.asarray(features, dtype="float32")
predictions = loaded_model.predict(feature_array, batch_size=2, verbose=0)

with output_path.open("w", newline="") as handle:
    writer = csv.writer(handle)
    writer.writerow(["account_id", "score"])
    for account_id, score in zip(account_ids, predictions[:, 0]):
        writer.writerow([account_id, f"{score:.6f}"])

print(f"backend: {keras.backend.backend()}")
print(f"loaded model: {model_path}")
print(f"input rows: {len(account_ids)}")
print(f"prediction shape: {predictions.shape}")
print(f"output file: {output_path}")
print("preview:")
with output_path.open(newline="") as handle:
    for line in handle.read().strip().splitlines()[:4]:
        print(line)