#!/usr/bin/env python3
"""Neural network classifier workflow.

This script provides a practical, repeatable template for:
- loading tabular classification data
- validating inputs and target labels
- splitting data safely
- scaling numeric features without leakage
- building a compact Keras classifier
- training with early stopping
- evaluating the model with classification metrics
- running a lightweight production-readiness check

Expected input:
- A CSV file containing feature columns and one target column.
- The target column may contain strings or numeric labels.

Example usage:
    python nn_classifier_workflow.py --data data.csv --target label

Optional usage for a provided test split:
    python nn_classifier_workflow.py --data train.csv --target label --test-data test.csv

Notes:
- The script avoids hard-coded secrets and external service calls.
- It is designed for structured/tabular data.
- For multiclass tasks, the output layer and loss are adjusted automatically.
"""

from __future__ import annotations

import argparse
import json
import os
import random
from dataclasses import dataclass
from typing import Optional, Tuple

import numpy as np
import pandas as pd
from sklearn.metrics import (
    accuracy_score,
    classification_report,
    confusion_matrix,
    f1_score,
    log_loss,
)
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import LabelEncoder, StandardScaler

from tensorflow import keras
from tensorflow.keras import layers


@dataclass
class DatasetBundle:
    X_train: np.ndarray
    X_val: np.ndarray
    X_test: np.ndarray
    y_train: np.ndarray
    y_val: np.ndarray
    y_test: np.ndarray
    label_encoder: Optional[LabelEncoder]
    scaler: StandardScaler
    feature_names: list


def set_seeds(seed: int) -> None:
    """Set random seeds for repeatable runs."""
    random.seed(seed)
    np.random.seed(seed)
    os.environ["PYTHONHASHSEED"] = str(seed)
    try:
        import tensorflow as tf

        tf.random.set_seed(seed)
    except Exception:
        pass


def load_data(path: str) -> pd.DataFrame:
    if not os.path.exists(path):
        raise FileNotFoundError(f"Data file not found: {path}")
    if not path.lower().endswith(".csv"):
        raise ValueError("Input data must be a CSV file.")
    df = pd.read_csv(path)
    if df.empty:
        raise ValueError("Loaded dataset is empty.")
    return df


def prepare_features(df: pd.DataFrame, target_col: str) -> Tuple[pd.DataFrame, pd.Series]:
    if target_col not in df.columns:
        raise ValueError(f"Target column '{target_col}' not found in dataset.")

    y = df[target_col]
    X = df.drop(columns=[target_col])

    if X.empty:
        raise ValueError("No feature columns remain after removing the target column.")

    # Keep the workflow simple and explicit: only numeric features are allowed here.
    # If categorical features exist, encode them before using this script.
    non_numeric_cols = [c for c in X.columns if not pd.api.types.is_numeric_dtype(X[c])]
    if non_numeric_cols:
        raise ValueError(
            "Non-numeric feature columns detected: " + ", ".join(non_numeric_cols) +
            ". Encode categorical features before training."
        )

    if y.isna().any():
        raise ValueError("Target column contains missing values.")

    if X.isna().any().any():
        raise ValueError("Feature matrix contains missing values. Impute or remove them before training.")

    return X, y


def encode_target(y: pd.Series) -> Tuple[np.ndarray, Optional[LabelEncoder]]:
    if pd.api.types.is_numeric_dtype(y):
        y_values = y.to_numpy()
        unique_values = np.unique(y_values)
        if len(unique_values) < 2:
            raise ValueError("Target must contain at least two classes.")
        return y_values, None

    le = LabelEncoder()
    y_encoded = le.fit_transform(y.astype(str))
    if len(np.unique(y_encoded)) < 2:
        raise ValueError("Target must contain at least two classes.")
    return y_encoded, le


def split_and_scale(
    X: pd.DataFrame,
    y: np.ndarray,
    seed: int,
    test_size: float = 0.30,
    val_size: float = 0.50,
) -> DatasetBundle:
    if not 0 < test_size < 1:
        raise ValueError("test_size must be between 0 and 1.")
    if not 0 < val_size < 1:
        raise ValueError("val_size must be between 0 and 1.")

    stratify = y if len(np.unique(y)) > 1 else None
    X_train, X_temp, y_train, y_temp = train_test_split(
        X,
        y,
        test_size=test_size,
        random_state=seed,
        stratify=stratify,
    )

    stratify_temp = y_temp if len(np.unique(y_temp)) > 1 else None
    X_val, X_test, y_val, y_test = train_test_split(
        X_temp,
        y_temp,
        test_size=val_size,
        random_state=seed,
        stratify=stratify_temp,
    )

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

    return DatasetBundle(
        X_train=X_train_scaled,
        X_val=X_val_scaled,
        X_test=X_test_scaled,
        y_train=y_train,
        y_val=y_val,
        y_test=y_test,
        label_encoder=None,
        scaler=scaler,
        feature_names=list(X.columns),
    )


def build_model(num_features: int, num_classes: int, learning_rate: float) -> keras.Model:
    if num_features < 1:
        raise ValueError("num_features must be at least 1.")
    if num_classes < 2:
        raise ValueError("num_classes must be at least 2.")
    if learning_rate <= 0:
        raise ValueError("learning_rate must be positive.")

    is_binary = num_classes == 2
    output_units = 1 if is_binary else num_classes
    output_activation = "sigmoid" if is_binary else "softmax"
    loss = "binary_crossentropy" if is_binary else "sparse_categorical_crossentropy"

    model = keras.Sequential(
        [
            layers.Input(shape=(num_features,)),
            layers.Dense(64, activation="relu"),
            layers.Dropout(0.2),
            layers.Dense(32, activation="relu"),
            layers.Dense(output_units, activation=output_activation),
        ]
    )

    model.compile(
        optimizer=keras.optimizers.Adam(learning_rate=learning_rate),
        loss=loss,
        metrics=["accuracy"],
    )
    return model


def train_model(
    model: keras.Model,
    X_train: np.ndarray,
    y_train: np.ndarray,
    X_val: np.ndarray,
    y_val: np.ndarray,
    epochs: int,
    batch_size: int,
) -> keras.callbacks.History:
    if epochs < 1:
        raise ValueError("epochs must be at least 1.")
    if batch_size < 1:
        raise ValueError("batch_size must be at least 1.")

    early_stop = keras.callbacks.EarlyStopping(
        monitor="val_loss",
        patience=5,
        restore_best_weights=True,
    )

    history = model.fit(
        X_train,
        y_train,
        validation_data=(X_val, y_val),
        epochs=epochs,
        batch_size=batch_size,
        callbacks=[early_stop],
        verbose=1,
    )
    return history


def predict_scores(model: keras.Model, X: np.ndarray) -> np.ndarray:
    scores = model.predict(X, verbose=0)
    return scores


def evaluate_model(
    model: keras.Model,
    X_test: np.ndarray,
    y_test: np.ndarray,
    num_classes: int,
    label_encoder: Optional[LabelEncoder] = None,
) -> dict:
    is_binary = num_classes == 2
    scores = predict_scores(model, X_test)

    if is_binary:
        scores = scores.reshape(-1)
        y_pred = (scores >= 0.5).astype(int)
        y_prob = np.clip(scores, 1e-7, 1 - 1e-7)
        loss_value = log_loss(y_test, np.column_stack([1 - y_prob, y_prob]))
    else:
        y_pred = np.argmax(scores, axis=1)
        y_prob = np.clip(scores, 1e-7, 1 - 1e-7)
        loss_value = log_loss(y_test, y_prob, labels=list(range(num_classes)))

    accuracy = accuracy_score(y_test, y_pred)
    f1 = f1_score(y_test, y_pred, average="binary" if is_binary else "weighted")
    cm = confusion_matrix(y_test, y_pred).tolist()

    if label_encoder is not None:
        target_names = list(label_encoder.classes_)
        report = classification_report(y_test, y_pred, target_names=target_names, zero_division=0)
    else:
        report = classification_report(y_test, y_pred, zero_division=0)

    return {
        "accuracy": float(accuracy),
        "f1_score": float(f1),
        "log_loss": float(loss_value),
        "confusion_matrix": cm,
        "classification_report": report,
    }


def production_readiness_check(metrics: dict, min_accuracy: float, max_log_loss: float) -> dict:
    checks = {
        "accuracy_threshold": metrics["accuracy"] >= min_accuracy,
        "log_loss_threshold": metrics["log_loss"] <= max_log_loss,
        "non_empty_report": bool(metrics.get("classification_report")),
    }
    checks["ready"] = all(checks.values())
    return checks


def main() -> int:
    parser = argparse.ArgumentParser(
        description="Train and evaluate a neural network classifier on tabular data."
    )
    parser.add_argument("--data", required=True, help="Path to the main CSV dataset.")
    parser.add_argument("--target", required=True, help="Name of the target column.")
    parser.add_argument("--test-data", default=None, help="Optional separate CSV file for testing.")
    parser.add_argument("--seed", type=int, default=42, help="Random seed for repeatability.")
    parser.add_argument("--epochs", type=int, default=50, help="Maximum training epochs.")
    parser.add_argument("--batch-size", type=int, default=32, help="Training batch size.")
    parser.add_argument("--learning-rate", type=float, default=1e-3, help="Adam learning rate.")
    parser.add_argument("--min-accuracy", type=float, default=0.70, help="Readiness threshold for accuracy.")
    parser.add_argument("--max-log-loss", type=float, default=0.80, help="Readiness threshold for log loss.")
    parser.add_argument("--metrics-output", default=None, help="Optional path to save metrics as JSON.")
    args = parser.parse_args()

    if not 0 <= args.min_accuracy <= 1:
        raise ValueError("--min-accuracy must be between 0 and 1.")
    if args.max_log_loss <= 0:
        raise ValueError("--max-log-loss must be positive.")

    set_seeds(args.seed)

    df = load_data(args.data)
    X, y_raw = prepare_features(df, args.target)
    y, label_encoder = encode_target(y_raw)

    if args.test_data:
        test_df = load_data(args.test_data)
        X_test_df, y_test_raw = prepare_features(test_df, args.target)
        y_test, test_label_encoder = encode_target(y_test_raw)
        if label_encoder is None and test_label_encoder is not None:
            raise ValueError("Training labels are numeric but test labels are strings; use a consistent encoding.")
        if label_encoder is not None:
            y_test = label_encoder.transform(y_test_raw.astype(str))

        X_train_df, X_val_df, y_train, y_val = train_test_split(
            X,
            y,
            test_size=0.30,
            random_state=args.seed,
            stratify=y if len(np.unique(y)) > 1 else None,
        )
        scaler = StandardScaler()
        X_train = scaler.fit_transform(X_train_df)
        X_val = scaler.transform(X_val_df)
        X_test = scaler.transform(X_test_df)
    else:
        bundle = split_and_scale(X, y, seed=args.seed)
        X_train, X_val, X_test = bundle.X_train, bundle.X_val, bundle.X_test
        y_train, y_val, y_test = bundle.y_train, bundle.y_val, bundle.y_test
        scaler = bundle.scaler

    num_features = X_train.shape[1]
    num_classes = len(np.unique(y))
    model = build_model(num_features, num_classes, args.learning_rate)
    _history = train_model(model, X_train, y_train, X_val, y_val, args.epochs, args.batch_size)

    metrics = evaluate_model(model, X_test, y_test, num_classes, label_encoder=label_encoder)
    readiness = production_readiness_check(metrics, args.min_accuracy, args.max_log_loss)

    summary = {
        "model_type": "dense_neural_network_classifier",
        "num_features": num_features,
        "num_classes": num_classes,
        "metrics": metrics,
        "readiness": readiness,
        "notes": [
            "Confirm the same preprocessing is used in production.",
            "Review confusion matrix and class-level metrics before deployment.",
            "Re-train and re-validate if feature distributions shift.",
        ],
    }

    print(json.dumps(summary, indent=2))

    if args.metrics_output:
        with open(args.metrics_output, "w", encoding="utf-8") as f:
            json.dump(summary, f, indent=2)

    return 0


if __name__ == "__main__":
    raise SystemExit(main())