import numpy as np
import gradio as gr

from model_inverse import inverse_design


def format_number(value):
    if value is None:
        return "--"
    return f"{value:,.4f}"


def format_percent(value):
    if value is None:
        return "--"
    return f"{value * 100:.2f}%"


def run_inverse_design(ply_number, a1, b1, c1, stress, n_restarts, epochs, use_lbfgs):
    y_target = np.array([a1, b1, c1, stress], dtype=np.float32)
    best = inverse_design(
        ply_number=int(ply_number),
        y_target=y_target,
        n_restarts=int(n_restarts),
        epochs=int(epochs),
        use_lbfgs=bool(use_lbfgs),
    )
    if best["input"] is None or best["output"] is None:
        return "<div class='empty-state'>No valid solution found.</div>"

    input_vals = {
        "Initial Temp": float(best["input"][0]),
        "Punch Velocity": float(best["input"][1]),
        "Cooling Time": float(best["input"][2]),
    }
    output_vals = {
        "A1": float(best["output"][0]),
        "B1": float(best["output"][1]),
        "C1": float(best["output"][2]),
        "Stress": float(best["output"][3]),
    }
    def rel_err(pred, target):
        target = float(target)
        if target == 0.0:
            return None
        return abs(float(pred) - target) / abs(target)

    output_rows = [
        ("A1", output_vals["A1"], rel_err(output_vals["A1"], a1)),
        ("B1", output_vals["B1"], rel_err(output_vals["B1"], b1)),
        ("C1", output_vals["C1"], rel_err(output_vals["C1"], c1)),
        ("Stress", output_vals["Stress"], rel_err(output_vals["Stress"], stress)),
    ]

    input_html = "".join(
        f"<div class='kv-row'><span class='kv-key'>{k}</span><span class='kv-val'>{format_number(v)}</span></div>"
        for k, v in input_vals.items()
    )
    output_html = "".join(
        "<div class='out-row'>"
        f"<div class='out-name'>{name}</div>"
        f"<div class='out-pred'>{format_number(pred)}</div>"
        f"<div class='out-err'>{format_percent(err)}</div>"
        "</div>"
        for name, pred, err in output_rows
    )
    return (
        "<div class='results-grid'>"
        "<div class='panel'>"
        "<div class='panel-title'>Best Input</div>"
        f"{input_html}"
        "</div>"
        "<div class='panel'>"
        "<div class='panel-title'>Predicted Output</div>"
        f"<div class='out-grid'>{output_html}</div>"
        "</div>"
        "</div>"
    )


custom_css = """
:root {
  --ink: #1d1a16;
  --muted: #6b625a;
  --sand: #f6f1e7;
  --clay: #e8dfd2;
  --accent: #c65b2a;
  --card: #fffaf1;
}
.page-wrap {
  background: radial-gradient(1200px 600px at 10% -10%, #fbe9cf 0%, #f6f1e7 45%, #efe7dc 100%);
  border-radius: 18px;
  padding: 18px;
}
.panel {
  background: var(--card);
  border: 1px solid var(--clay);
  border-radius: 16px;
  padding: 16px;
  box-shadow: 0 8px 24px rgba(57, 44, 30, 0.08);
}
.panel-title {
  font-size: 16px;
  letter-spacing: 0.5px;
  text-transform: uppercase;
  color: var(--muted);
  margin-bottom: 10px;
}
.kv-row {
  display: flex;
  justify-content: space-between;
  align-items: center;
  padding: 10px 12px;
  background: #fdf7ec;
  border: 1px dashed #e3d5c4;
  border-radius: 12px;
  margin-bottom: 10px;
}
.kv-key {
  color: var(--ink);
  font-weight: 600;
}
.kv-val {
  color: var(--accent);
  font-weight: 700;
}
.results-grid {
  display: grid;
  grid-template-columns: 1fr 1fr;
  gap: 18px;
}
.out-grid {
  display: grid;
  grid-template-columns: 1fr 1fr 1fr;
  gap: 10px;
}
.out-row {
  background: #fff3e2;
  border: 1px solid #f0d2b7;
  border-radius: 12px;
  padding: 12px;
  display: grid;
  grid-template-columns: 1fr;
  gap: 6px;
}
.out-name {
  font-weight: 700;
  color: var(--ink);
}
.out-pred {
  font-size: 18px;
  color: var(--accent);
  font-weight: 700;
}
.out-err {
  color: var(--muted);
  font-size: 13px;
}
@media (max-width: 1000px) {
  .results-grid {
    grid-template-columns: 1fr;
  }
  .out-grid {
    grid-template-columns: 1fr 1fr;
  }
}
@media (max-width: 600px) {
  .out-grid {
    grid-template-columns: 1fr;
  }
}
.empty-state {
  padding: 16px;
  border-radius: 12px;
  border: 1px dashed #e3d5c4;
  color: var(--muted);
  text-align: center;
  background: #fff7ea;
}
"""


with gr.Blocks(title="Inverse Design Demo", css=custom_css) as demo:
    gr.Markdown("# Inverse Design Demo")
    gr.Markdown("Enter a target output; the model finds processing parameters.")

    with gr.Row():
        with gr.Column():
            gr.Markdown("## Target Output")
            ply_number = gr.Number(label="Ply number", value=2, precision=0)
            a1 = gr.Number(label="A1 (Angle)", value=0.89, precision=4)
            b1 = gr.Number(label="B1 (Angle)", value=0.83, precision=4)
            c1 = gr.Number(label="C1 (Angle)", value=0.12, precision=4)
            stress = gr.Number(label="Max Stress (MPa)", value=180.2, precision=4)
        with gr.Column():
            gr.Markdown("## Optimization Settings")
            n_restarts = gr.Number(label="Restarts", value=5, precision=0)
            epochs = gr.Number(label="Epochs of Optimization", value=100, precision=0)
            use_lbfgs = gr.Checkbox(label="Use LBFGS Optimization", value=True)

    run_btn = gr.Button("Run Inverse Design")

    with gr.Row(elem_classes=["page-wrap"]):
        results_html = gr.HTML()

    run_btn.click(
        run_inverse_design,
        inputs=[ply_number, a1, b1, c1, stress, n_restarts, epochs, use_lbfgs],
        outputs=[results_html],
    )


if __name__ == "__main__":
    demo.launch()