from pathlib import Path
import itertools
import py3Dmol
from ase import Atoms
from ase.io import write
OUT = Path("out")
a = 3.52
pos = [p for p in itertools.product((0, a), repeat=3)]
pos += [(a/2, a/2, 0), (a/2, a/2, a), (a/2, 0, a/2), (a/2, a, a/2), (0, a/2, a/2), (a, a/2, a/2)]
write(OUT / "fcc_cell.xyz", Atoms("Ni14", positions=pos))
COLORS = {"Ni": "#b5bbc1", "Mo": "#1b6c80"}
def page(path, fmt, name, w, h, scale, rot, zoom, box=None):
    v = py3Dmol.view(width=w, height=h)
    v.addModel(open(path).read(), fmt)
    for el, c in COLORS.items():
        v.setStyle({"elem": el}, {"sphere": {"scale": scale, "color": c}})
    if box is not None:
        L = box
        edges = [((0,0,0),(L,0,0)),((0,0,0),(0,L,0)),((0,0,0),(0,0,L)),((L,0,0),(L,L,0)),((L,0,0),(L,0,L)),((0,L,0),(L,L,0)),((0,L,0),(0,L,L)),((0,0,L),(L,0,L)),((0,0,L),(0,L,L)),((L,L,0),(L,L,L)),((L,0,L),(L,L,L)),((0,L,L),(L,L,L))]
        for s, e in edges:
            v.addLine({"start": dict(zip("xyz", s)), "end": dict(zip("xyz", e)), "color": "#8a9096"})
    v.setBackgroundColor("white")
    v.zoomTo()
    for axis, deg in rot:
        v.rotate(deg, axis)
    v.zoom(zoom)
    Path(f"{name}.html").write_text(f"<!doctype html><html><head><meta charset='utf-8'><style>body{{margin:0;background:#fff}}</style></head><body>{v._make_html()}</body></html>")
rot = [("x", -18), ("y", 28)]
page(OUT/"relaxed.cif", "cif", "struct_big", 880, 600, 0.42, rot, 0.92, box=7.04*1.004)
page(OUT/"relaxed.cif", "cif", "struct_thumb", 320, 220, 0.42, rot, 0.9, box=7.04*1.004)
page(OUT/"fcc_cell.xyz", "xyz", "fcc_cell", 220, 200, 0.3, [("x", -15), ("y", 25)], 1.15, box=a)