#!/usr/bin/env python3
"""Run a few in-memory training steps on DCU without modifying official weights."""

import argparse
import pickle
from pathlib import Path

from _bootstrap import DATA_DIR, WEIGHT_DIR

import numpy as np
import tensorflow as tf
import tf_keras


DEFAULT_DATA_DIR = DATA_DIR / "features" / "cdr_kmer3_ag_kmer1"
DEFAULT_WEIGHT_PATH = WEIGHT_DIR / "cdr_kmer3_ag_kmer1" / "Model99.h5"


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--data-dir", type=Path, default=DEFAULT_DATA_DIR)
    parser.add_argument("--weights", type=Path, default=DEFAULT_WEIGHT_PATH)
    parser.add_argument("--batch-size", type=int, default=8)
    parser.add_argument("--steps", type=int, default=1)
    return parser.parse_args()


def load_training_batch(data_dir, batch_size):
    # Register features.InputFeatures before unpickling the official artifacts.
    import model  # noqa: F401

    with (data_dir / "cdr_features_tr.pickle").open("rb") as reader:
        cdr_features = pickle.load(reader)
    with (data_dir / "ag_features_tr.pickle").open("rb") as reader:
        ag_features = pickle.load(reader)

    if len(cdr_features) != len(ag_features):
        raise ValueError(
            f"Unpaired training features: CDR={len(cdr_features)}, AG={len(ag_features)}"
        )
    if batch_size < 2 or batch_size > len(cdr_features):
        raise ValueError(
            f"--batch-size must be between 2 and {len(cdr_features)}, got {batch_size}"
        )

    selected_cdr = cdr_features[:batch_size]
    selected_ag = ag_features[:batch_size]
    inputs = [
        np.asarray([item.input_ids for item in selected_cdr], dtype=np.int32),
        np.asarray([item.cdr_number_ids for item in selected_cdr], dtype=np.int32),
        np.asarray([item.input_ids for item in selected_ag], dtype=np.int32),
    ]
    labels = np.asarray(
        [[item.label_id] for item in selected_cdr], dtype=np.float32
    )

    expected_shapes = ((batch_size, 24), (batch_size, 24), (batch_size, 2371))
    actual_shapes = tuple(array.shape for array in inputs)
    if actual_shapes != expected_shapes:
        raise ValueError(f"Unexpected input shapes: {actual_shapes}, expected {expected_shapes}")
    if labels.shape != (batch_size, 1):
        raise ValueError(f"Unexpected label shape: {labels.shape}")

    return inputs, labels


def inference_loss(model, inputs, labels):
    predictions = model(inputs, training=False)
    losses = tf_keras.losses.binary_crossentropy(labels, predictions)
    loss = float(tf.reduce_mean(losses).numpy())
    predictions = predictions.numpy()
    if not np.isfinite(loss) or not np.all(np.isfinite(predictions)):
        raise FloatingPointError("Inference loss or predictions contain NaN/Inf")
    return loss


def main():
    args = parse_args()
    if args.steps < 1:
        raise ValueError(f"--steps must be at least 1, got {args.steps}")
    if not args.data_dir.is_dir():
        raise FileNotFoundError(f"Data directory not found: {args.data_dir}")
    if not args.weights.is_file():
        raise FileNotFoundError(f"Weight file not found: {args.weights}")

    gpus = tf.config.list_physical_devices("GPU")
    if not gpus:
        raise RuntimeError(
            "TensorFlow cannot see a DCU. Set HIP_VISIBLE_DEVICES and "
            "CUDA_VISIBLE_DEVICES before starting Python."
        )
    for gpu in gpus:
        tf.config.experimental.set_memory_growth(gpu, True)

    inputs, labels = load_training_batch(args.data_dir, args.batch_size)

    # Import after configuring the accelerator to avoid initializing it early.
    from model import get_model

    with tf.device("/GPU:0"):
        model = get_model()
        model.load_weights(str(args.weights))
        loss_before = inference_loss(model, inputs, labels)
        weights_before = [variable.numpy().copy() for variable in model.trainable_variables]

        training_losses = []
        for _ in range(args.steps):
            loss = model.train_on_batch(inputs, labels)
            training_losses.append(float(np.asarray(loss).reshape(-1)[0]))

        loss_after = inference_loss(model, inputs, labels)

    if not np.all(np.isfinite(training_losses)):
        raise FloatingPointError(f"Training loss contains NaN/Inf: {training_losses}")

    changed_count = 0
    max_weight_change = 0.0
    for before, variable in zip(weights_before, model.trainable_variables):
        after = variable.numpy()
        if not np.all(np.isfinite(after)):
            raise FloatingPointError(f"Trainable variable contains NaN/Inf: {variable.name}")
        change = float(np.max(np.abs(after - before)))
        max_weight_change = max(max_weight_change, change)
        if change > 0.0:
            changed_count += 1

    dcu_variable_count = sum("GPU:0" in variable.device for variable in model.variables)

    print(f"TensorFlow: {tf.__version__}")
    print(f"tf_keras: {tf_keras.__version__}")
    print(f"DCU: {gpus[0]}")
    print(f"Batch size: {args.batch_size}")
    print(f"Steps: {args.steps}")
    print(f"Labels: {labels.reshape(-1).tolist()}")
    print(f"Inference loss before: {loss_before:.10g}")
    print(f"Training losses: {training_losses}")
    print(f"Inference loss after: {loss_after:.10g}")
    print(f"Trainable tensors changed: {changed_count}/{len(weights_before)}")
    print(f"Max absolute weight change: {max_weight_change:.10g}")
    print(f"Variables placed on GPU:0: {dcu_variable_count}/{len(model.variables)}")

    if changed_count == 0 or max_weight_change == 0.0:
        raise AssertionError("No trainable weights changed after train_on_batch")
    if dcu_variable_count == 0:
        raise RuntimeError("No model variables were placed on GPU:0")

    print("Minimal DCU training: PASS")
    print("Official Model99.h5 was loaded read-only and was not overwritten.")


if __name__ == "__main__":
    main()