"""Generate the sample structure and convergence plots used in the CapSci v2 mockups."""
import json
from pathlib import Path

import numpy as np
from ase.build import bulk
from ase.io import write
import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib import font_manager

OUT = Path(__file__).parent / "out"
OUT.mkdir(exist_ok=True)

INK, MUTED, GRID, ACCENT, NEG = "#131416", "#5d5f60", "#e5e7e9", "#1b6c80", "#c4242a"
for name in ("PingFang SC", "Hiragino Sans GB", "Arial Unicode MS"):
    if any(f.name == name for f in font_manager.fontManager.ttflist):
        plt.rcParams["font.family"] = [name]
        break
plt.rcParams.update({"axes.edgecolor": MUTED, "axes.labelcolor": MUTED, "xtick.color": MUTED,
                     "ytick.color": MUTED, "axes.spines.top": False, "axes.spines.right": False,
                     "axes.spines.left": False, "font.size": 11, "svg.fonttype": "none"})

# --- structure: 2x2x2 conventional FCC Ni (32 atoms), 4 Mo placed far apart (stand-in for an SQS candidate)
a0 = 3.52
atoms = bulk("Ni", "fcc", a=a0, cubic=True).repeat((2, 2, 2))
rng = np.random.default_rng(7)
chosen = [int(rng.integers(len(atoms)))]
while len(chosen) < 4:
    d = np.array([min(atoms.get_distance(i, j, mic=True) for j in chosen) for i in range(len(atoms))])
    d[chosen] = -1
    chosen.append(int(np.argmax(d)))
symbols = atoms.get_chemical_symbols()
for i in chosen:
    symbols[i] = "Mo"
atoms.set_chemical_symbols(symbols)
write(OUT / "POSCAR_initial", atoms, format="vasp", direct=True, sort=True)
write(OUT / "initial.cif", atoms)
relaxed = atoms.copy()
relaxed.set_cell(atoms.cell * 1.004, scale_atoms=True)
relaxed.positions += rng.normal(0, 0.015, relaxed.positions.shape)
write(OUT / "CONTCAR", relaxed, format="vasp", direct=True, sort=True)
write(OUT / "relaxed.cif", relaxed)

# --- sample ionic relaxation (illustrative numbers, labelled as sample data in the UI)
E = np.array([-213.80, -214.42, -214.83, -215.07, -215.21, -215.29, -215.34, -215.37, -215.39,
              -215.40, -215.41, -215.415, -215.418, -215.420, -215.421, -215.422, -215.4225, -215.4228])
F = np.array([0.62, 0.41, 0.29, 0.21, 0.16, 0.12, 0.094, 0.075, 0.061, 0.051, 0.044, 0.039, 0.035,
              0.031, 0.028, 0.025, 0.022, 0.019])
steps = np.arange(len(E))
EDIFFG = 0.02


def style(ax):
    ax.grid(axis="y", color=GRID, linewidth=1)
    ax.tick_params(length=0)
    ax.set_axisbelow(True)


def convergence(n, path, w=7.2, h=2.6):
    fig, (a, b) = plt.subplots(1, 2, figsize=(w, h), dpi=200)
    a.plot(steps[:n], E[:n], color=ACCENT, lw=2)
    a.scatter(steps[n - 1], E[n - 1], color=ACCENT, s=18, zorder=3)
    a.set_xlim(0, 17); a.set_ylim(-215.6, -213.6)
    a.set_xlabel("离子步"); a.set_title("总能量 E (eV)", loc="left", color=INK, fontsize=11)
    a.annotate(f"{E[n-1]:.3f}", (steps[n - 1], E[n - 1]), xytext=(-4, 8), textcoords="offset points", ha="right", color=INK, fontsize=10)
    b.semilogy(steps[:n], F[:n], color=ACCENT, lw=2)
    b.scatter(steps[n - 1], F[n - 1], color=ACCENT, s=18, zorder=3)
    b.axhline(EDIFFG, color=INK, lw=1, ls=(0, (4, 3)))
    b.text(0.2, EDIFFG * 1.12, "收敛判据 0.02", color=INK, fontsize=10)
    b.set_xlim(0, 17); b.set_ylim(0.01, 1)
    b.set_xlabel("离子步"); b.set_title("最大原子受力 F_max (eV/Å)", loc="left", color=INK, fontsize=11)
    b.annotate(f"{F[n-1]:.3f}", (steps[n - 1], F[n - 1]), xytext=(-4, -14), textcoords="offset points", ha="right", color=INK, fontsize=10)
    for ax in (a, b):
        style(ax)
    fig.tight_layout()
    fig.savefig(path, transparent=False, facecolor="white")
    plt.close(fig)


convergence(len(E), OUT / "conv_done.png")
convergence(15, OUT / "conv_live.png")


def spark(n, path):
    fig, ax = plt.subplots(figsize=(2.4, 0.9), dpi=200)
    ax.semilogy(steps[:n], F[:n], color=ACCENT, lw=1.6)
    ax.axhline(EDIFFG, color=INK, lw=0.8, ls=(0, (3, 2)))
    ax.set_xlim(0, 17); ax.set_ylim(0.01, 1)
    ax.axis("off")
    fig.subplots_adjust(0, 0, 1, 1)
    fig.savefig(path, facecolor="white")
    plt.close(fig)


spark(len(F), OUT / "spark_done.png")
spark(15, OUT / "spark_live.png")

# --- sample failed SCF for the failure screen: |dE| per electronic step, two harness attempts
k = np.arange(1, 101)
de1 = 10 ** (-1 - 3.2 * (1 - np.exp(-k / 18))) * (1 + 0.6 * np.sin(k / 2.3) ** 2)
k2 = np.arange(1, 201)
de2 = 10 ** (-1 - 3.6 * (1 - np.exp(-k2 / 30))) * (1 + 0.7 * np.sin(k2 / 2.7) ** 2)
fig, ax = plt.subplots(figsize=(7.2, 2.6), dpi=200)
ax.semilogy(k, de1, color=MUTED, lw=1.4, label="第 1 次 · NELM 100")
ax.semilogy(k2, de2, color=NEG, lw=1.6, label="第 2 次 · NELM 200")
ax.axhline(1e-5, color=INK, lw=1, ls=(0, (4, 3)))
ax.text(4, 1.3e-5, "EDIFF 1e-5", color=INK, fontsize=10)
ax.set_xlim(0, 200); ax.set_ylim(1e-6, 1)
ax.set_xlabel("电子步"); ax.set_title("|ΔE| (eV) · 第 4 个离子步", loc="left", color=INK, fontsize=11)
ax.legend(frameon=False, loc="upper right", fontsize=10)
style(ax)
fig.tight_layout()
fig.savefig(OUT / "scf_fail.png", facecolor="white")
plt.close(fig)

json.dump({"mo_indices": chosen, "n_atoms": len(atoms), "E_final": float(E[-1]),
           "E_per_atom": float(E[-1] / len(atoms)), "F_final": float(F[-1]),
           "a_relaxed": float(a0 * 1.004)}, open(OUT / "summary.json", "w"), indent=2)
print(open(OUT / "summary.json").read())