from pytessel import PyTessel
from . import PyQInt
import numpy as np

def build_isosurface(filename, mol, basis, coeff, isovalue, dim = 5.0):
    sz = 99
    integrator = PyQInt()
    grid = integrator.build_rectgrid3d(-dim, dim, sz)
    scalarfield = np.reshape(integrator.plot_wavefunction(grid, coeff, basis), (sz, sz, sz))
    unitcell = np.diag(np.ones(3) * 10.0)

    pytessel = PyTessel()
    vertices_p, normals_p, indices_p = pytessel.marching_cubes(scalarfield.flatten(), scalarfield.shape, unitcell.flatten(), isovalue)
    vertices_n, normals_n, indices_n = pytessel.marching_cubes(scalarfield.flatten(), scalarfield.shape, unitcell.flatten(), -isovalue)

    build_abo(filename, mol, vertices_p, normals_p, indices_p, vertices_n, normals_n, indices_n)

def build_abo(filename, mol, vertices_p, normals_p, indices_p, vertices_n, normals_n, indices_n):
    outfile = ('%s.abo' % filename)
    f = open(outfile, 'wb')

    # write number of frames
    f.write(int(1).to_bytes(2, byteorder='little'))

    # write frame index
    f.write(int(1).to_bytes(2, byteorder='little'))

    # write description
    description = 'Orbital generated by PyQInt / PyTessel'
    f.write(len(description).to_bytes(2, byteorder='little'))
    f.write(bytearray(description, encoding='utf8'))

    # write atoms
    f.write(len(mol.charges).to_bytes(2, byteorder='little'))

    # loop over atoms
    for idx,pos in zip(mol.charges, mol.atoms):
        f.write(idx.to_bytes(1, byteorder='little'))
        f.write((np.array(pos[1], dtype=np.float32) * 0.529177).tobytes())

    # set colors
    color_pos = np.array([201, 2, 65, 0xF0]) / 255.0
    color_neg = np.array([2, 182, 201, 0xF0]) / 255.0

    # write number of models
    f.write(int(2).to_bytes(2, byteorder='little'))

    # write model index
    f.write(int(0).to_bytes(2, byteorder='little'))

    # write model color
    f.write(np.array(color_pos, dtype=np.float32).tobytes())

    # write number of vertices
    f.write(len(vertices_p).to_bytes(4, byteorder='little'))

    # loop over vertices and normals
    for v,n in zip(vertices_p, normals_p):
        f.write(np.array(v, dtype=np.float32).tobytes())
        f.write(np.array(n, dtype=np.float32).tobytes())

    # write number of faces2
    f.write((int(len(indices_p) / 3)).to_bytes(4, byteorder='little'))

    # write faces
    f.write(np.array(indices_p, dtype=np.uint32).tobytes())

    # write model index
    f.write(int(1).to_bytes(2, byteorder='little'))

    # write model color
    f.write(np.array(color_neg, dtype=np.float32).tobytes())

    # write number of vertices
    f.write(len(vertices_n).to_bytes(4, byteorder='little'))

    # loop over vertices and normals
    for v,n in zip(vertices_n, normals_n):
        f.write(np.array(v, dtype=np.float32).tobytes())
        f.write(np.array(n, dtype=np.float32).tobytes())

    # write indices
    f.write((int(len(indices_n) / 3)).to_bytes(4, byteorder='little'))

    # write faces
    f.write(np.array(indices_n, dtype=np.uint32).tobytes())

    f.close()
