import argparse
from pathlib import Path

import numpy as np


def parse_hidden_dims(s: str) -> list[int]:
    dims = []
    for part in s.split(','):
        part = part.strip()
        if not part:
            continue
        value = int(part)
        if value <= 0:
            raise ValueError(f"hidden dim must be positive, got {value}")
        dims.append(value)
    if not dims:
        raise ValueError("at least one hidden layer size is required")
    return dims


def build_model(input_dim: int, hidden_dims: list[int], dropout: float, l2_reg: float, lr: float):
    import tensorflow as tf
    from tensorflow.keras import Sequential, regularizers
    from tensorflow.keras.layers import Input, Dense, Dropout

    model = Sequential()
    model.add(Input(shape=(input_dim,)))

    for units in hidden_dims:
        model.add(
            Dense(
                units,
                activation='relu',
                kernel_initializer='he_normal',
                kernel_regularizer=regularizers.l2(l2_reg),
            )
        )
        if dropout > 0:
            model.add(Dropout(dropout))

    model.add(Dense(1))

    optimizer = tf.keras.optimizers.Adam(learning_rate=lr)
    model.compile(loss='mae', optimizer=optimizer, metrics=['mae'])
    return model


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--data', required=True, help='prepared npz path')
    parser.add_argument('--epochs', type=int, default=200)
    parser.add_argument('--batch-size', type=int, default=32)
    parser.add_argument('--save-model', default='rent_nn_best.keras')
    parser.add_argument('--hidden-dims', default='256,128,64', help='comma-separated hidden layer sizes')
    parser.add_argument('--dropout', type=float, default=0.2)
    parser.add_argument('--l2', type=float, default=1e-5)
    parser.add_argument('--lr', type=float, default=1e-4)
    parser.add_argument('--patience', type=int, default=15)
    parser.add_argument('--min-delta', type=float, default=1e-4)
    parser.add_argument('--val-split', type=float, default=0.1, help='fraction of X_train used as validation')
    parser.add_argument('--seed', type=int, default=42)
    args = parser.parse_args()

    hidden_dims = parse_hidden_dims(args.hidden_dims)

    d = np.load(args.data)
    X_train = d['X_train'].astype(np.float32)
    y_train = d['y_train'].astype(np.float32)
    X_test = d['X_test'].astype(np.float32)
    y_test = d['y_test'].astype(np.float32)

    import tensorflow as tf
    from tensorflow.keras.callbacks import EarlyStopping, ModelCheckpoint, ReduceLROnPlateau

    tf.keras.utils.set_random_seed(args.seed)

    # Optional GPU memory growth for safer multi-user environments.
    gpus = tf.config.list_physical_devices('GPU')
    for gpu in gpus:
        try:
            tf.config.experimental.set_memory_growth(gpu, True)
        except Exception:
            pass

    print(f'X_train: {X_train.shape}, y_train: {y_train.shape}')
    print(f'X_test : {X_test.shape}, y_test : {y_test.shape}')
    print(f'hidden_dims={hidden_dims}, dropout={args.dropout}, l2={args.l2}, lr={args.lr}')
    print(f'epochs={args.epochs}, batch_size={args.batch_size}, val_split={args.val_split}, seed={args.seed}')

    model = build_model(
        input_dim=X_train.shape[1],
        hidden_dims=hidden_dims,
        dropout=args.dropout,
        l2_reg=args.l2,
        lr=args.lr,
    )

    save_path = Path(args.save_model)
    if save_path.parent != Path('.'):
        save_path.parent.mkdir(parents=True, exist_ok=True)

    callbacks = [
        EarlyStopping(
            monitor='val_mae',
            mode='min',
            patience=args.patience,
            min_delta=args.min_delta,
            restore_best_weights=True,
            verbose=1,
        ),
        ReduceLROnPlateau(
            monitor='val_mae',
            mode='min',
            factor=0.5,
            patience=max(3, args.patience // 3),
            min_lr=1e-6,
            verbose=1,
        ),
        ModelCheckpoint(
            filepath=str(save_path),
            monitor='val_mae',
            mode='min',
            save_best_only=True,
            verbose=1,
        ),
    ]

    history = model.fit(
        X_train,
        y_train,
        batch_size=args.batch_size,
        epochs=args.epochs,
        validation_split=args.val_split,
        shuffle=True,
        callbacks=callbacks,
        verbose=1,
    )

    best_epoch = int(np.argmin(history.history['val_mae'])) + 1
    best_val_mae = float(np.min(history.history['val_mae']))
    print(f'best epoch by val_mae: {best_epoch}')
    print(f'best val mae (万元): {best_val_mae:.4f}')
    print(f'best val mae (元): {best_val_mae * 10000:.2f}')

    # Final test evaluation: test set is only used once here.
    best_model = tf.keras.models.load_model(save_path)
    test_pred = best_model.predict(X_test, verbose=0).reshape(-1)
    test_mae = float(np.mean(np.abs(test_pred - y_test)))
    print(f'final test mae (万元): {test_mae:.4f}')
    print(f'final test mae (元): {test_mae * 10000:.2f}')
    print(f'saved best model: {save_path}')


if __name__ == '__main__':
    main()