import xml.etree.ElementTree as ET
from datetime import datetime
from pathlib import Path
from xml.dom import minidom

import numpy as np
from ase import Atom, Atoms
from ase.data import covalent_radii
from ase.data.colors import jmol_colors
from ase.neighborlist import natural_cutoffs, neighbor_list
from IPython.core.display import HTML

BOND_RADIUS = 0.17
MIN_ATOM_RADIUS = BOND_RADIUS + 0.1


def view(atoms: Atoms, show_bonds: bool = False):
    """
    Generate an interactive visualisation of an :class:`ase.Atoms` object.

    Parameters
    ----------
    atoms:
        The atoms object to visualise.
    show_bonds:
        Whether to show bonds between atoms.

    Example
    -------
    Visualise a water molecule in a Jupyter notebook:

    .. code-block:: python
        :emphasize-lines: 5

        from ase.build import molecule
        from load_atoms import view

        ethanol = molecule("CH3CH2OH")
        view(ethanol, show_bonds=True)

    .. raw:: html
        :file: ../../../docs/source/_static/ethanol.html

    Save the visualisation to an :code:`html` file:

    .. code-block:: python
        :emphasize-lines: 3

        from pathlib import Path

        viz = view(ethanol, show_bonds=True)
        Path("ethanol.html").write_text(viz.data)
    """
    scene = x3d_scene(atoms, show_bonds, width="300px", height="300px")

    uid = unique_variable_name()
    js_file = Path(__file__).parent / "view.js"
    js = js_file.read_text().replace("uid", uid)

    x3d_script = (Path(__file__).parent / "x3d.script").read_text()

    return HTML(
        f"""\
<html>
    <body>
    {x3d_script}
    <div id="{uid}">
        {scene}
    </div>
    <script>
        {js}
    </script>
    </body>
</html>
"""
    )


def x3d_scene(
    structure: Atoms,
    show_bonds: bool = False,
    **css_style: str,
) -> str:
    style = {"width": "300px", "height": "300px"}
    style.update(css_style)
    style = " ".join(f'{k}="{v}";' for k, v in style.items())

    scene = x3d_atoms(structure.copy(), show_bonds=show_bonds)

    document = f"""\
<X3D {style}>
<param name="showProgress" value="false" />

<!--Inserting Generated X3D Scene-->
{pretty_print(scene)}
<!--End of Inserted Scene-->

</X3D>
"""
    return document


def x3d_atom(atom: Atom, scale: float = 1.0):
    """Represent an atom as a coloured sphere."""

    radius = max(covalent_radii[atom.number] * scale, MIN_ATOM_RADIUS)
    shape = element(
        "shape",
        children=(appearance_for(atom), element("sphere", radius=f"{radius}")),
    )

    x, y, z = atom.position
    return translate(shape, x, y, z)


def appearance_for(atom: Atom):
    r, g, b = jmol_colors[atom.number]
    return element(
        "appearance",
        child=element("material", diffuseColor=f"{r} {g} {b}"),
    )


def x3d_wireframe_box(box):
    """
    Wireframe representation of a box (3x3 array).

    To draw a box, spanned by vectors a, b and c, it is necessary to
    draw 4 faces, each of which is a parallelogram. The faces are:
    (start from) , (vectors spanning the face)
    1. (0), (a, b)
    2. (c), (a, b) # opposite face to 1.
    3. (0), (a, c)
    4. (b), (a, c) # opposite face to 3.
    """

    a, b, c = box
    faces = [
        wireframe_face(a, b),
        wireframe_face(a, b, origin=c),
        wireframe_face(a, c),
        wireframe_face(a, c, origin=b),
    ]
    return group(faces)


def wireframe_face(vec1, vec2, origin=(0, 0, 0)):
    """Wireframe representation of a face spanned by vec1 and vec2."""

    x1, y1, z1 = vec1
    x2, y2, z2 = vec2

    material = element(
        "material", diffuseColor="0 0 0"
    )  # TODO: make this work in both dark and light mode
    appearance = element("appearance", child=material)

    points = [
        (0, 0, 0),
        (x1, y1, z1),
        (x1 + x2, y1 + y2, z1 + z2),
        (x2, y2, z2),
        (0, 0, 0),
    ]
    points = " ".join(f"{x} {y} {z}" for x, y, z in points)

    coordinates = element("coordinate", point=points)
    lineset = element("lineset", vertexCount="5", child=coordinates)
    shape = element("shape", children=(appearance, lineset))

    x, y, z = origin
    return translate(shape, x, y, z)


def x3d_atoms(atoms: Atoms, show_bonds: bool = False):
    """Convert an atoms object into an x3d representation."""

    # first, ensure that the atoms are within the cell
    if atoms.pbc.any():
        atoms.wrap()
        show_cell = True
    else:
        atoms.center(vacuum=0.1)
        show_cell = False

    scale = 0.6 if show_bonds else 1.0

    atom_spheres = group([x3d_atom(atom, scale) for atom in atoms])  # type: ignore
    bonds = x3d_bonds(atoms, scale) if show_bonds else empty_element()
    wireframe = x3d_wireframe_box(atoms.cell) if show_cell else empty_element()
    cell = group((wireframe, atom_spheres, bonds))

    # we want the cell to be in the middle of the viewport
    # so that we can (a) see the whole cell and (b) rotate around the center
    # therefore we translate so that the center of the cell is at the origin
    cell_center = atoms.cell.sum(axis=0) / 2  # type: ignore
    cell = translate(cell, *(-cell_center))
    # rotate the cell about y (we grab this later to start spinning it)
    cell = rotate(cell, 0, 1, 0, 0.1, id="atoms-rotation")

    # we want the cell, and all atoms, to be visible
    # - sometimes atoms appear outside the cell
    # - sometimes atoms only take up a small part of the cell
    # location of the viewpoint therefore takes both of these into account:
    # the scene is centered on the cell, so we find the furthest point away
    # from the cell center, and use this to determine the
    # distance of the viewpoint
    points = np.vstack((atoms.positions, atoms.cell[:]))  # type: ignore
    max_xyz_extent = get_maximum_extent(points - cell_center)

    # put the camera 2.2x as far away as
    # the largest separation between two points in any of x, y or z
    max_dim = max(max_xyz_extent)
    pos = f"0, 0, {max_dim * 2.2}"

    # NB. viewpoint needs to contain an (empty) child to be valid x3d
    viewpoint = element("viewpoint", position=pos, child=empty_element())

    return element("scene", children=(viewpoint, cell), id="atoms-scene")


def empty_element():
    return element("group")


def x3d_bonds(atoms: Atoms, scale: float = 1.0):
    # set pbc off so that we don't try to draw bonds across periodic boundaries
    atoms = atoms.copy()
    atoms.pbc = False
    i, j = neighbor_list("ij", atoms, cutoff=natural_cutoffs(atoms, mult=1.2))  # type: ignore
    bond_index = list(zip(i.tolist(), j.tolist()))

    return group([x3d_bond(atoms[a], atoms[b], scale) for a, b in bond_index])  # type: ignore


def atom_size(atom: Atom, scale: float = 1.0):
    return max(covalent_radii[atom.number] * scale, MIN_ATOM_RADIUS)


def x3d_bond(a: Atom, b: Atom, scale: float = 1.0):
    # if same element: draw a single cylinder connecting the two
    if a.number == b.number:
        return cyclinder_between(a.position, b.position, appearance_for(a))

    # otherwise, draw two cylinders, one from each atom to a suitable
    # point in the middle, and coloured according the respective atoms
    dist = np.linalg.norm(a.position - b.position)
    r_a = atom_size(a, scale)
    r_b = atom_size(b, scale)
    frac = r_a / dist + (0.5 * (dist - r_a - r_b)) / dist
    halfway = a.position + frac * (b.position - a.position)

    a_to_center = cyclinder_between(a.position, halfway, appearance_for(a))
    center_to_b = cyclinder_between(halfway, b.position, appearance_for(b))

    return group([a_to_center, center_to_b])


def cyclinder_between(a, b, appearance):
    halfway = (a + b) / 2

    z = np.array([0, 1, 0])
    v = b - a
    d = np.linalg.norm(v)
    v = v / d
    angle = np.arccos(z.dot(v))
    axis = np.cross(z, v)

    cylinder = element(
        "shape",
        children=[
            appearance,
            element("cylinder", radius=f"{BOND_RADIUS}", height=f"{d}"),
        ],
    )
    return translate(rotate(cylinder, *axis, angle), *halfway)  # type: ignore


def element(name, child=None, children=None, **attributes) -> ET.Element:
    """Convenience function to make an XML element.

    If child is specified, it is appended to the element.
    If children is specified, they are appended to the element.
    You cannot specify both child and children."""

    # make sure we don't specify both child and children
    if child is not None:
        assert children is None, "Cannot specify both child and children"
        children = [child]
    else:
        children = children or []

    element = ET.Element(name, **attributes)
    for child in children:
        element.append(child)
    return element


def translate(thing, x, y, z):
    """Translate a x3d element by x, y, z."""
    return element("transform", translation=f"{x} {y} {z}", child=thing)


def rotate(thing, x, y, z, angle, **attributes):
    return element(
        "transform",
        rotation=f"{x}, {y}, {z}, {angle}",
        child=thing,
        **attributes,
    )


def group(things):
    """Group a (list of) x3d elements."""
    return element("group", children=things)


def pretty_print(element: ET.Element, indent: int = 2):
    """Pretty print an XML element."""

    byte_string = ET.tostring(element, "utf-8")
    parsed = minidom.parseString(byte_string)
    prettied = parsed.toprettyxml(indent=" " * indent)
    # remove first line - contains an extra, un-needed xml declaration
    lines = prettied.splitlines()[1:]
    return "\n".join(lines)


def get_maximum_extent(xyz):
    """Get the maximum extent of an array of 3d set of points."""

    return np.max(xyz, axis=0) - np.min(xyz, axis=0)


def unique_variable_name():
    now = str(datetime.now().timestamp())
    now = now.replace(".", "")
    return "".join(chr(97 + int(digit)) for digit in now)
