#!/usr/bin/env python3
"""Evaluate the official test set and weights on a single DCU."""

import argparse
import pickle
import time
from pathlib import Path

from _bootstrap import DATA_DIR, WEIGHT_DIR

import numpy as np
import tensorflow as tf
import tf_keras
from sklearn.metrics import average_precision_score, roc_auc_score


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=64)
    parser.add_argument(
        "--output",
        type=Path,
        help="optional .npz path for labels and predictions",
    )
    return parser.parse_args()


def load_test_set(data_dir):
    # Register features.InputFeatures before loading the official pickles.
    import model  # noqa: F401

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

    if not cdr_features:
        raise ValueError("The test set is empty")
    if len(cdr_features) != len(ag_features):
        raise ValueError(
            f"Unpaired test features: CDR={len(cdr_features)}, AG={len(ag_features)}"
        )

    cdr_labels = np.asarray([item.label_id for item in cdr_features], dtype=np.int32)
    ag_labels = np.asarray([item.label_id for item in ag_features], dtype=np.int32)
    if not np.array_equal(cdr_labels, ag_labels):
        mismatch_count = int(np.count_nonzero(cdr_labels != ag_labels))
        raise ValueError(f"CDR/AG labels differ for {mismatch_count} test samples")
    if set(np.unique(cdr_labels).tolist()) != {0, 1}:
        raise ValueError(f"Expected binary labels 0/1, got {np.unique(cdr_labels)}")

    inputs = [
        np.asarray([item.input_ids for item in cdr_features], dtype=np.int32),
        np.asarray([item.cdr_number_ids for item in cdr_features], dtype=np.int32),
        np.asarray([item.input_ids for item in ag_features], dtype=np.int32),
    ]
    sample_count = len(cdr_features)
    expected_shapes = ((sample_count, 24), (sample_count, 24), (sample_count, 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}")

    return inputs, cdr_labels


def main():
    args = parse_args()
    if args.batch_size < 1:
        raise ValueError(f"--batch-size must be positive, got {args.batch_size}")
    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_test_set(args.data_dir)

    # 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))
        start_time = time.perf_counter()
        predictions = model.predict(
            inputs,
            batch_size=args.batch_size,
            verbose=1,
        ).reshape(-1)
        elapsed_seconds = time.perf_counter() - start_time

    if predictions.shape != labels.shape:
        raise ValueError(
            f"Prediction/label shape mismatch: {predictions.shape} vs {labels.shape}"
        )
    if not np.all(np.isfinite(predictions)):
        raise FloatingPointError("Test predictions contain NaN or Inf")
    if np.any((predictions < 0.0) | (predictions > 1.0)):
        raise ValueError("Test predictions are outside [0, 1]")

    auroc = float(roc_auc_score(labels, predictions))
    aupr = float(average_precision_score(labels, predictions))
    positive_count = int(labels.sum())
    negative_count = int(labels.size - positive_count)
    dcu_variable_count = sum("GPU:0" in variable.device for variable in model.variables)

    print("\nFull official test-set evaluation")
    print(f"TensorFlow: {tf.__version__}")
    print(f"tf_keras: {tf_keras.__version__}")
    print(f"DCU: {gpus[0]}")
    print(f"Samples: {labels.size}")
    print(f"Positive/negative: {positive_count}/{negative_count}")
    print(f"Batch size: {args.batch_size}")
    print(f"Variables placed on GPU:0: {dcu_variable_count}/{len(model.variables)}")
    print(f"Prediction min/max/mean: {predictions.min():.10g} / "
          f"{predictions.max():.10g} / {predictions.mean():.10g}")
    print(f"Elapsed seconds: {elapsed_seconds:.3f}")
    print(f"Throughput samples/s: {labels.size / elapsed_seconds:.3f}")
    print(f"AUROC: {auroc:.10f}")
    print(f"AUPR: {aupr:.10f}")
    print("Paper reference AUROC/AUPR: 0.926 / 0.952")

    if dcu_variable_count == 0:
        raise RuntimeError("No model variables were placed on GPU:0")

    if args.output is not None:
        args.output.parent.mkdir(parents=True, exist_ok=True)
        np.savez_compressed(
            args.output,
            labels=labels,
            predictions=predictions,
            auroc=np.asarray(auroc),
            aupr=np.asarray(aupr),
        )
        print(f"Saved predictions: {args.output}")

    print("Full test-set inference: PASS")


if __name__ == "__main__":
    main()