import os
import numpy as np
import tensorflow as tf

os.environ.setdefault("TF_FORCE_GPU_ALLOW_GROWTH", "true")

def build_model():
    model = tf.keras.Sequential([
        tf.keras.layers.Conv2D(64, (3, 3), activation="relu", input_shape=(36, 22, 1), padding="same"),
        tf.keras.layers.Dropout(0.5),
        tf.keras.layers.MaxPooling2D((2, 2), padding="same"),
        tf.keras.layers.Conv2D(32, (3, 4), activation="relu", padding="same"),
        tf.keras.layers.Dropout(0.4),
        tf.keras.layers.MaxPooling2D((2, 2), padding="same"),
        tf.keras.layers.Conv2D(32, (4, 4), activation="relu", padding="same"),
        tf.keras.layers.Dropout(0.3),
        tf.keras.layers.Flatten(),
        tf.keras.layers.Dense(64, activation="relu"),
        tf.keras.layers.Dropout(0.4),
        tf.keras.layers.Dense(2, activation="softmax"),
    ])
    model.compile(
        loss="categorical_crossentropy",
        optimizer="adam",
        metrics=["accuracy"],
    )
    return model

def mcc_score(y_pred, y_real):
    y_pred = y_pred.astype(int)
    y_real = y_real.astype(int)
    tp = float(np.sum((y_pred == 1) & (y_real == 1)))
    tn = float(np.sum((y_pred == 0) & (y_real == 0)))
    fp = float(np.sum((y_pred == 1) & (y_real == 0)))
    fn = float(np.sum((y_pred == 0) & (y_real == 1)))
    denom = np.sqrt((tp + fp) * (tp + fn) * (tn + fp) * (tn + fn))
    return np.nan if denom == 0 else (tp * tn - fp * fn) / denom

def train_one(name, data_file, out_dir):
    print(f"\n==== training {name} ====")
    data = np.load(data_file)
    x_train = data["x_train"].astype("float32")
    y_train = data["y_train"].astype("float32")
    x_test = data["x_test"].astype("float32")
    y_test = data["y_test"].astype("float32")

    print("x_train", x_train.shape, "y_train", y_train.shape)
    print("x_test ", x_test.shape, "y_test ", y_test.shape)
    print("GPUs:", tf.config.list_physical_devices("GPU"))

    with tf.device("/GPU:0"):
        model = build_model()
        model.fit(
            x_train,
            y_train,
            epochs=30,
            batch_size=50,
            validation_split=0.2,
            verbose=2,
        )
        prob = model.predict(x_test, batch_size=128, verbose=0)

    y_real = np.argmax(y_test, axis=1)
    y_pred = np.argmax(prob, axis=1)
    acc = float(np.mean(y_real == y_pred))
    mcc = float(mcc_score(y_pred, y_real))

    print(f"\n{name} accuracy: {acc:.4f}")
    print(f"{name} mcc: {mcc:.4f}")
    print("confusion matrix rows=real cols=pred")
    cm = np.zeros((2, 2), dtype=int)
    for r, p in zip(y_real, y_pred):
        cm[r, p] += 1
    print(cm)

    model.save(out_dir)
    #model.save(out_dir + ".keras")
    np.savez_compressed(out_dir + "_eval.npz", prob=prob, y_real=y_real, y_pred=y_pred, cm=cm, acc=acc, mcc=mcc)
    print("saved", out_dir)

train_one("CTLA-4", "model/CNN/c1_data.npz", "weight/CNN/model_c1_dcu")
train_one("PD-1", "model/CNN/p1_data.npz", "weight/CNN/model_p1_dcu")
print("\nCNN DCU training OK")