from clustice.water import molecules_iter, tip4p
from typing import Union, Iterable
import numpy as np


def is_diagonal(matrix):
    # create a dummy matrix
    dummy_matrix = np.ones(matrix.shape, dtype=np.uint8)
    # Fill the diagonal of dummy matrix with 0.
    np.fill_diagonal(dummy_matrix, 0)

    return np.count_nonzero(np.multiply(dummy_matrix, matrix)) == 0


def render(
    dg,
    layout: dict,
    title: str = "Generated by ClustIce",
    cell_matrix: Union[Iterable | None] = None,
    water_model=tip4p,
    shift: Iterable = [0.0, 0.0, 0.0],
):
    s = ""

    labels = []
    positions = []
    for atom_positions, atom_types in molecules_iter(
        dg, layout, water_model=water_model, cell_matrix=cell_matrix
    ):
        labels += atom_types
        positions.append(atom_positions)

    positions = np.array(positions).reshape(-1, 3)
    natom = positions.shape[0]

    if cell_matrix is not None:
        cell_inv = np.linalg.inv(cell_matrix)
        relative = positions @ cell_inv + shift
        relative -= np.floor(relative)
        positions = relative @ cell_matrix

    for order, (label, (x, y, z)) in enumerate(zip(labels, positions)):
        if order > 99998:
            order = 99998
        s += f"{1:>5}DOD  {label:5}{order+1:>4}{x:>8.3f}{y:>8.3f}{z:>8.3f}\n"

    head = title + "\n"
    head += f"{natom}\n"
    if cell_matrix is None:
        cell_matrix = np.diag([100.0, 100.0, 100.0])
    if is_diagonal(cell_matrix):
        cellitems = cell_matrix[0, 0], cell_matrix[1, 1], cell_matrix[2, 2]
    else:
        cellitems = (
            cell_matrix[0, 0],
            cell_matrix[1, 1],
            cell_matrix[2, 2],
            cell_matrix[0, 1],
            cell_matrix[0, 2],
            cell_matrix[1, 0],
            cell_matrix[1, 2],
            cell_matrix[2, 0],
            cell_matrix[2, 1],
        )
    tail = ""
    for i in cellitems:
        tail += f"{i:.6f} "
    tail += "\n"
    return head + s + tail
