

import os
import io
import uuid
import json
import numpy as np
from base64 import b64encode
from zipfile import ZipFile, ZIP_DEFLATED


def _read_file(filepath):
    '''This function reads file as str

    Parameters
    ----------
    filename : str
        filepath of the file to read as string

    Returns
    -------
    file_content : str
    '''

    with open(filepath, "r") as fp:
        return fp.read()


def plot_shell_mesh(nodes_coordinates,
                    shell_node_indexes,
                    field=None,
                    is_element_field=True):
    ''' Plot a mesh

    Parameters
    ----------
    node_coordinates : np.ndarray
        array of node coordinates for elements
    tria_node_indexes : np.ndarray
        node indexes for the tria shells (3 nodes per shell)
    quad_node_indexes : np.ndarray
        node indexes of quad shells (4 nodes per element)

    Returns
    -------
    html : str
        html code for plotting as string
    '''

    assert(nodes_coordinates.ndim == 2)
    assert(nodes_coordinates.shape[1] == 3)
    assert(shell_node_indexes.ndim == 2)
    assert(shell_node_indexes.shape[1] in [3, 4])
    if isinstance(field, np.ndarray):
        assert(field.ndim == 1)
        if is_element_field:
            assert(field.shape[0] == shell_node_indexes.shape[0])
        else:
            assert(field.shape[0] == nodes_coordinates.shape[0])

    # cast types correctly
    # the types MUST be float32
    nodes_coordinates = nodes_coordinates.astype(np.float32)
    if isinstance(field, np.ndarray):
        field = field.astype(np.float32)

    # distinguish tria and quads
    is_quad = shell_node_indexes[:, 2] != shell_node_indexes[:, 3]
    is_tria = np.logical_not(is_quad)

    # seperate tria and quads ... I know its sad :(
    tria_node_indexes = shell_node_indexes[is_tria][:, :3]
    quad_node_indexes = shell_node_indexes[is_quad]

    # we can only plot tria, therefore we need to split quads
    # into two trias
    quad_node_indexes_tria1 = quad_node_indexes[:, :3]
    # quad_node_indexes_tria2 = quad_node_indexes[:, [True, False, True, True]]
    quad_node_indexes_tria2 = quad_node_indexes[:, [0, 2, 3]]

    # assemble elements for plotting
    # This seems to take a lot of memory and you are right thinking this,
    # the issue is just in order to plot fringe values, we need to output
    # the element values at the 3 corner nodes. Since elements share nodes
    # we can not use the same nodes, thus we need to create multiple nodes
    # at the same position but with different fringe.
    nodes_xyz = np.concatenate([
        nodes_coordinates[tria_node_indexes].reshape((-1, 3)),
        nodes_coordinates[quad_node_indexes_tria1].reshape((-1, 3)),
        nodes_coordinates[quad_node_indexes_tria2].reshape((-1, 3))
    ])

    # fringe value and hover title
    if isinstance(field, np.ndarray):

        if is_element_field:
            n_shells = len(shell_node_indexes)
            n_tria = np.sum(is_tria)
            n_quads = n_shells-n_tria

            # split field according to elements
            field_tria = field[is_tria]
            field_quad = field[is_quad]

            # allocate fringe array
            node_fringe = np.zeros(
                (len(field_tria)+2*len(field_quad), 3), dtype=np.float32)

            # set fringe values
            node_fringe[:n_tria, 0] = field_tria
            node_fringe[:n_tria, 1] = field_tria
            node_fringe[:n_tria, 2] = field_tria

            node_fringe[n_tria:n_tria+n_quads, 0] = field_quad
            node_fringe[n_tria:n_tria+n_quads, 1] = field_quad
            node_fringe[n_tria:n_tria+n_quads, 2] = field_quad

            node_fringe[n_tria+n_quads:n_tria +
                        2*n_quads, 0] = field_quad
            node_fringe[n_tria+n_quads:n_tria +
                        2*n_quads, 1] = field_quad
            node_fringe[n_tria+n_quads:n_tria +
                        2*n_quads, 2] = field_quad

            # flatty paddy
            node_fringe = node_fringe.flatten()
        else:
            # copy & paste ftw
            node_fringe = np.concatenate([
                field[tria_node_indexes].reshape((-1, 3)),
                field[quad_node_indexes_tria1].reshape((-1, 3)),
                field[quad_node_indexes_tria2].reshape((-1, 3))
            ])
            node_fringe = node_fringe.flatten()

        # element text
        node_txt = [str(entry) for entry in node_fringe.flatten()]
    else:
        node_fringe = np.zeros(len(nodes_xyz), dtype=np.float32)
        node_txt = ['']*len(nodes_xyz)

    # zip compression of data for HTML (reduces size)
    zdata = io.BytesIO()
    with ZipFile(zdata, 'w', compression=ZIP_DEFLATED) as zipFile:
        zipFile.writestr('/intensities', node_fringe.tostring())
        zipFile.writestr('/positions',   nodes_xyz.tostring())
        zipFile.writestr('/text', json.dumps(node_txt))
    zdata = b64encode(zdata.getvalue()).decode('utf-8')

    # read html template
    _html_template = _read_file(os.path.join(
        os.path.dirname(__file__), 'resources', 'template.html'))

    # format html template file
    _html_div = _html_template.format(div_id=uuid.uuid4(),
                                      lowIntensity=field.min() if isinstance(field, np.ndarray) else 0,
                                      highIntensity=field.max() if isinstance(field, np.ndarray) else 0,
                                      zdata=zdata)

    # wrap it up with all needed js libraries
    _html_jszip_js = '<script type="text/javascript">%s</script>' % _read_file(
        os.path.join(os.path.dirname(__file__), 'resources', 'jszip.min.js'))
    _html_three_js = '<script type="text/javascript">%s</script>' % _read_file(
        os.path.join(os.path.dirname(__file__), 'resources', 'three.min.js'))
    _html_chroma_js = '<script type="text/javascript">%s</script>' % _read_file(
        os.path.join(os.path.dirname(__file__), 'resources', 'chroma.min.js'))
    _html_jquery_js = '<script type="text/javascript">%s</script>' % _read_file(
        os.path.join(os.path.dirname(__file__), 'resources', 'jquery.min.js'))

    return '''
<!DOCTYPE html>
<html lang="en">
    <head>
    <meta charset="utf-8" />
        {_jquery_js}
        {_jszip_js}
        {_three_js}
        {_chroma_js}
    </head>
    <body>
        {_html_div}
    </body>
</html>'''.format(
        _html_div=_html_div,
        _jszip_js=_html_jszip_js,
        _three_js=_html_three_js,
        _chroma_js=_html_chroma_js,
        _jquery_js=_html_jquery_js)
