import os
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"

import numpy as np
import torch
import torch.nn as nn
import gradio as gr

MODEL_PATH = "best_10fold_uncertainty_model.pt"
DEVICE = "cpu"   # HF Spaces CPU'da çalıştırmak en stabil seçenek


# ===== Model mimarisi (checkpoint ile birebir) =====
class MLPRegressor(nn.Module):
    def __init__(self, in_dim=5, out_dim=3):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(in_dim, 64),
            nn.ReLU(),
            nn.Dropout(0.1),
            nn.Linear(64, 32),
            nn.ReLU(),
            nn.Dropout(0.1),
            nn.Linear(32, out_dim),
        )

    def forward(self, x):
        return self.net(x)


# ===== Checkpoint yükle =====
if not os.path.exists(MODEL_PATH):
    raise FileNotFoundError(
        f"Model dosyası bulunamadı: {MODEL_PATH}\n"
        "Lütfen best_10fold_uncertainty_model.pt dosyasını app.py ile aynı klasöre koyun."
    )

ckpt = torch.load(MODEL_PATH, map_location=DEVICE)

model = MLPRegressor().to(DEVICE)
model.load_state_dict(ckpt["model_state"])
model.eval()

x_mean = ckpt["x_mean"].cpu().numpy()
x_scale = ckpt["x_scale"].cpu().numpy()
y_mean = ckpt["y_mean"].cpu().numpy()
y_scale = ckpt["y_scale"].cpu().numpy()


def predict(P, v, h, t, E, e_mode):
    """
    Inputs: P, v, h, t, E
    Outputs: bulk, Ms, Hc
    e_mode:
      - "Auto-calc E = P/(v*h*t)"
      - "Manual E"
    """
    # Basit validasyon
    if v <= 0 or h <= 0 or t <= 0:
        return "Hata: v, h, t > 0 olmalı.", None, None, None, None

    if e_mode.startswith("Auto"):
        E = P / (v * h * t)

    if E <= 0:
        return "Hata: E > 0 olmalı.", None, None, None, None

    X = np.array([[P, v, h, t, E]], dtype=np.float32)
    Xs = (X - x_mean) / x_scale

    with torch.no_grad():
        Xs_t = torch.tensor(Xs, dtype=torch.float32, device=DEVICE)
        Ys = model(Xs_t).cpu().numpy()

    Y = Ys * y_scale + y_mean
    bulk, Ms, Hc = Y[0]
    return "OK", float(E), float(bulk), float(Ms), float(Hc)


with gr.Blocks(title="LPBF ANN Simulator (5→3)") as demo:
    gr.Markdown(
        "# LPBF Fe-based ANN Simulator (5 inputs → 3 outputs)\n"
        "**Inputs:** P, v, h, t, E  →  **Outputs:** Bulk density, Ms, Hc\n\n"
        "- **Auto-calc E** seçerseniz fiziksel tutarlılık korunur: `E = P/(v·h·t)`\n"
        "- **Manual E** seçerseniz E’yi elle girersiniz.\n"
    )

    with gr.Row():
        with gr.Column():
            P = gr.Number(value=70, label="Laser power P (W)")
            v = gr.Number(value=900, label="Scan speed v (mm/s)")
            h = gr.Dropdown(choices=[0.02, 0.03], value=0.02, label="Hatch spacing h (mm)")
            t = gr.Number(value=0.05, label="Layer thickness t (mm)")

            e_mode = gr.Radio(
                choices=["Auto-calc E = P/(v*h*t)", "Manual E"],
                value="Auto-calc E = P/(v*h*t)",
                label="Energy density mode"
            )
            E = gr.Number(value=77.7778, label="Energy density E (J/mm³) (Manual modda kullanılır)")

            btn = gr.Button("Predict")

        with gr.Column():
            status = gr.Textbox(label="Status")
            E_used = gr.Number(label="E used (J/mm³)")
            bulk_out = gr.Number(label="Predicted Bulk density (%)")
            Ms_out = gr.Number(label="Predicted Ms (Am²/kg)")
            Hc_out = gr.Number(label="Predicted Hc (kA/m)")

    btn.click(
        fn=predict,
        inputs=[P, v, h, t, E, e_mode],
        outputs=[status, E_used, bulk_out, Ms_out, Hc_out]
    )

demo.launch()