import json
import os
import time
from dataclasses import dataclass
from typing import Dict, List, Tuple

import numpy as np
from sklearn.metrics import mean_absolute_error, r2_score
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler

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

import tensorflow as tf
from tensorflow.keras.callbacks import EarlyStopping
from tensorflow.keras.layers import Conv1D, Dense, Dropout, Flatten, Input, MaxPooling1D
from tensorflow.keras.models import Sequential
from tensorflow.keras.optimizers import Adam


CURRENT_DIR = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.abspath(os.path.join(CURRENT_DIR, os.pardir))
OUT_DIR = os.path.join(ROOT, "output", "report_dante_v2_v5_v6")
os.makedirs(OUT_DIR, exist_ok=True)

DB_PATH = os.path.join(ROOT, "output", "d1a3o12_250421_forces02_dante_v6_database_live.npz")
DECISION_PATH = os.path.join(ROOT, "output", "d1a3o12_250421_forces02_dante_v6_1_structure_decision.json")


@dataclass
class ModelSpec:
    name: str
    kind: str  # mlp | cnn


def set_tf_device(device_id: int = 1) -> str:
    gpus = tf.config.list_physical_devices("GPU")
    if not gpus:
        return "CPU"
    idx = int(max(0, min(len(gpus) - 1, device_id)))
    try:
        tf.config.set_visible_devices([gpus[idx]], "GPU")
        tf.config.experimental.set_memory_growth(gpus[idx], True)
        return f"GPU:{idx}"
    except RuntimeError:
        return f"GPU:{idx}(runtime_initialized)"


def load_init_dataset() -> Tuple[np.ndarray, np.ndarray, Dict]:
    if not os.path.exists(DB_PATH):
        raise FileNotFoundError(DB_PATH)
    if not os.path.exists(DECISION_PATH):
        raise FileNotFoundError(DECISION_PATH)

    with open(DECISION_PATH, "r", encoding="utf-8") as f:
        decision = json.load(f)

    n_init = int(decision["num_initial"])
    db = np.load(DB_PATH, allow_pickle=True)
    x = np.asarray(db["input_x"], dtype=np.float64)
    y = np.asarray(db["input_y"], dtype=np.float64).reshape(-1)

    if len(x) < n_init:
        raise RuntimeError(f"DB samples {len(x)} < num_initial {n_init}")

    x0 = x[:n_init]
    y0 = y[:n_init]
    return x0, y0, decision


def make_mlp(input_dims: int) -> Sequential:
    model = Sequential(
        [
            Input(shape=(input_dims,)),
            Dense(128, activation="elu"),
            Dropout(0.10),
            Dense(64, activation="elu"),
            Dropout(0.10),
            Dense(32, activation="elu"),
            Dense(1, activation="linear"),
        ]
    )
    model.compile(optimizer=Adam(learning_rate=1e-3), loss="mse", metrics=["mae"])
    return model


def make_cnn(input_dims: int) -> Sequential:
    model = Sequential(
        [
            Input(shape=(input_dims, 1)),
            Conv1D(128, kernel_size=3, padding="same", activation="elu"),
            MaxPooling1D(pool_size=2, strides=1),
            Dropout(0.2),
            Conv1D(64, kernel_size=3, padding="same", activation="elu"),
            MaxPooling1D(pool_size=2, strides=1),
            Dropout(0.2),
            Conv1D(32, kernel_size=3, padding="same", activation="elu"),
            Conv1D(16, kernel_size=3, padding="same", activation="elu"),
            Flatten(),
            Dense(64, activation="elu"),
            Dense(1, activation="linear"),
        ]
    )
    model.compile(optimizer=Adam(learning_rate=1e-3), loss="mse", metrics=["mae"])
    return model


def run_one_split(x: np.ndarray, y: np.ndarray, seed: int, spec: ModelSpec) -> Dict:
    x_tr, x_te, y_tr, y_te = train_test_split(x, y, test_size=0.30, random_state=seed, shuffle=True)

    x_scaler = StandardScaler()
    y_scaler = StandardScaler()
    x_tr_s = x_scaler.fit_transform(x_tr)
    x_te_s = x_scaler.transform(x_te)
    y_tr_s = y_scaler.fit_transform(y_tr.reshape(-1, 1)).reshape(-1)

    if spec.kind == "mlp":
        model = make_mlp(x.shape[1])
        xtr_in = x_tr_s
        xte_in = x_te_s
    elif spec.kind == "cnn":
        model = make_cnn(x.shape[1])
        xtr_in = x_tr_s.reshape(len(x_tr_s), x.shape[1], 1)
        xte_in = x_te_s.reshape(len(x_te_s), x.shape[1], 1)
    else:
        raise ValueError(spec.kind)

    cb = [EarlyStopping(monitor="val_loss", patience=25, restore_best_weights=True)]
    hist = model.fit(
        xtr_in,
        y_tr_s,
        validation_split=0.25,
        batch_size=32,
        epochs=250,
        verbose=0,
        callbacks=cb,
    )

    y_hat_s = model.predict(xte_in, verbose=0).reshape(-1, 1)
    y_hat = y_scaler.inverse_transform(y_hat_s).reshape(-1)

    return {
        "seed": int(seed),
        "r2": float(r2_score(y_te, y_hat)),
        "mae": float(mean_absolute_error(y_te, y_hat)),
        "epochs": int(len(hist.history.get("loss", []))),
        "best_val_loss": float(np.min(hist.history.get("val_loss", [np.nan]))),
    }


def summarize(rows: List[Dict]) -> Dict:
    r2 = np.array([r["r2"] for r in rows], dtype=np.float64)
    mae = np.array([r["mae"] for r in rows], dtype=np.float64)
    return {
        "r2_mean": float(np.mean(r2)),
        "r2_std": float(np.std(r2)),
        "mae_mean": float(np.mean(mae)),
        "mae_std": float(np.std(mae)),
    }


def main() -> None:
    t0 = time.time()
    device = set_tf_device(1)
    x, y, decision = load_init_dataset()

    specs = [
        ModelSpec(name="mlp_128_64_32", kind="mlp"),
        ModelSpec(name="cnn_paper_like", kind="cnn"),
    ]
    seeds = [0, 1, 2, 3, 4]

    results = {}
    for spec in specs:
        rows = [run_one_split(x, y, s, spec) for s in seeds]
        results[spec.name] = {
            "kind": spec.kind,
            "per_split": rows,
            "summary": summarize(rows),
        }

    r2_mlp = results["mlp_128_64_32"]["summary"]["r2_mean"]
    r2_cnn = results["cnn_paper_like"]["summary"]["r2_mean"]

    out = {
        "db_path": DB_PATH,
        "decision_path": DECISION_PATH,
        "device": device,
        "n_init": int(len(x)),
        "dims": int(x.shape[1]),
        "y_stats": {
            "mean": float(np.mean(y)),
            "std": float(np.std(y)),
            "min": float(np.min(y)),
            "max": float(np.max(y)),
        },
        "models": results,
        "delta": {
            "r2_cnn_minus_mlp": float(r2_cnn - r2_mlp),
            "better_model": "cnn_paper_like" if r2_cnn > r2_mlp else "mlp_128_64_32",
        },
        "elapsed_sec": float(time.time() - t0),
    }

    out_json = os.path.join(OUT_DIR, "v6_1_init_surrogate_mlp_vs_cnn.json")
    with open(out_json, "w", encoding="utf-8") as f:
        json.dump(out, f, ensure_ascii=False, indent=2)

    print("saved", out_json)
    print("better_model", out["delta"]["better_model"]) 
    print("delta_r2", out["delta"]["r2_cnn_minus_mlp"])


if __name__ == "__main__":
    main()
