from clustice.water import molecules_iter, tip4p
from typing import Union, Iterable
import numpy as np
from solid2 import sphere, union, cylinder
import networkx as nx


def bond(a, b, r, _fn):
    d = b - a
    theta = np.arctan2(np.linalg.norm(d[:2]), d[2])
    phi = np.arctan2(d[1], d[0])
    h = np.linalg.norm(d)
    return (
        cylinder(r=r, h=h, _fn=_fn)
        .rotate(0, np.degrees(theta), 0)
        .rotate(0, 0, np.degrees(phi))
        .translate(a)
    )


def model(
    g: nx.Graph,
    layout: np.ndarray,
    # title: str = "Generated by ClustIce",
    cell_matrix: Union[Iterable | None] = None,
    water_model=tip4p,
    vertex_radius=0.1,
    edge_radius=0.05,
    # shift: Iterable = [0.0, 0.0, 0.0],
):
    objects = union()

    if vertex_radius > 0.0:
        for atom_position in layout:
            obj = sphere(r=vertex_radius, _fs=vertex_radius / 2).translate(
                atom_position
            )
            objects += obj
    if edge_radius > 0.0:
        for a, b in g.edges():
            pa, pb = layout[a], layout[b]
            obj = bond(pa, pb, r=edge_radius, _fn=12)
            objects += obj
    return objects
