"""
Library of SPAM functions for plotting a tetrahedral mesh
Copyright (C) 2020 SPAM Contributors

This program is free software: you can redistribute it and/or modify it
under the terms of the GNU General Public License as published by the Free
Software Foundation, either version 3 of the License, or (at your option)
any later version.

This program is distributed in the hope that it will be useful, but WITHOUT
ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License for
more details.

You should have received a copy of the GNU General Public License along with
this program.  If not, see <http://www.gnu.org/licenses/>.
"""

from __future__ import print_function

import numpy
import matplotlib.pyplot as plt
import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
from progressbar import ProgressBar
pbar = ProgressBar()

def plotDelaunay3D(points,connectivity,center_fraction,tet_strains,strain_range,axstr=1):
    """ 
    Plots Delaunay triangulation with strains.

    Parameters
    ----------
        points: Particles' x,y,z coordinates at original step in Np x 3 array. (Just used for plotting)

        connectivity: Delaunay triangulation generated by mesh_sandbox.voro_to_connectivity_3d

        center_fraction: Fraction of center of specimen to plot.

        tet_strains: List of strain tensors in each tetrahedron, from mesh_sandbox.voro_connectivity_to_strain_3d.

        strain_range: Upper and lower bounds for plotting strain.

        axstr: flag for the strain component to plot
            1 = 'XX' component
            2 = 'YY' component
            3 = 'ZZ' component
            4 = volumetric strain, XX+YY+ZZ
            5 = equivalent strain, sqrt(2/3*e_{ij}^{dev}e_{ij}^{dev})

    Returns
    -------
        Nothing, just creates a plot right now   

    Example
    -------
        mesh_sandbox.voro_plot_connectivity_3d(points,connectivity,0.02,tet_sym,0.02,'Z')

    """

    def voro_patch(ax, x, y, z, v, vmin=0, vmax=100, cmap_name='Greens'):
        cmap = mpl.cm.get_cmap(cmap_name)               # Get colormap by name
        c = cmap(mpl.colors.Normalize(vmin, vmax)(v))   # Normalize value and get color
        pc = Poly3DCollection([list(zip(x,y,z))])       # Create PolyCollection from coords
        pc.set_facecolor(c)                             # Set facecolor to mapped value
        pc.set_edgecolor('None')                           # Set edgecolor to black
        ax.add_collection3d(pc)                         # Add PolyCollection to axes
        return pc

    def voro_view(ax, code):
        if code == 2: #view(2) sets the default two-dimensional view, az = 0, el = 90.
            ax.view_init(0, 0)     # (args are reversed from MATLAB)

    # Find centroids within a central x-slice
    xmin = points[:,0].min()
    xmax = points[:,0].max()
    ymin = points[:,1].min()
    ymax = points[:,1].max()
    zmin = points[:,2].min()
    zmax = points[:,2].max()
    xspan = xmax - xmin
    idx = np.transpose(np.where((points[:,1]>(xmin+(0.5-center_fraction/2)*xspan)) & 
                   (points[:,1]<(xmin+(0.5+center_fraction/2)*xspan))))

    # Now find all tetrahedra containing these centroids
    print("Finding tetrahedra to plot...\n")
    inconnectivity = np.isin(connectivity,idx)
    idxtetbool,_ = np.where(inconnectivity)
    idxtet = np.unique(idxtetbool)

    print("Getting all node positions of tetrahedra to plot...\n")
    combs = [[0,1,2],[1,2,3],[2,3,0],[0,1,3]]

    # All tetrahedra are now unique, but some faces may be plotted twice if
    # they belong to adjacent tetrahedra being plotted. Fix this...
    face_plot_ids = []
    tet_str_plot = []
    for ii in (range(len(idxtet))):
        if (axstr == 1):
            tmpstr = tet_strains[idxtet[ii]][0,0]
        if (axstr == 2):
            tmpstr = tet_strains[idxtet[ii]][1,1]
        if (axstr == 3):
            tmpstr = tet_strains[idxtet[ii]][2,2]
        if (axstr == 4):
            tmpstr = (tet_strains[idxtet[ii]][0,0]+tet_strains[idxtet[ii]][1,1]+
                      tet_strains[idxtet[ii]][2,2])
        if (axstr == 5):
            devstr = tet_strains[idxtet[ii]]
            devstr = devstr - np.identity(3)*np.trace(devstr)
            tmpstr = np.sqrt(0.6666666*np.sum(np.tensordot(devstr,devstr,axes=1)))
        for jj in range(4):
            face_plot_ids.append(np.sort([connectivity[idxtet[ii]][combs[jj][0]],connectivity[idxtet[ii]][combs[jj][1]],connectivity[idxtet[ii]][combs[jj][2]]]))
            tet_str_plot.append(tmpstr)

    # Convert list to np array and remove non-unique rows
    face_plot_ids = np.array(face_plot_ids)
    face_plot_ids = np.sort(face_plot_ids,axis=1)
    y = np.ascontiguousarray(face_plot_ids).view(np.dtype((np.void, face_plot_ids.dtype.itemsize * face_plot_ids.shape[1])))
    _, idx = np.unique(y, return_index=True)
    face_plot_ids_unique = face_plot_ids[idx]
    tet_str_plot = np.array(tet_str_plot)
    tet_str_plot = tet_str_plot[idx]

    # Now construct x,y,z,v to send to patch
    x = []
    y = []
    z = []
    v = []
    for ii in (range((face_plot_ids_unique.shape[0]))):
        tet_coms = points[face_plot_ids_unique[ii],:]
        x.append(tet_coms[:,0])
        y.append(tet_coms[:,1])
        z.append(tet_coms[:,2])
        v.append(tet_str_plot[ii])

    # Initialize figure as 3D projection
    #fig = plt.figure()
    ax = plt.axes(projection='3d')
    ax.set_xlim(xmin,xmax); ax.set_ylim(ymin,ymax); ax.set_zlim(zmin,zmax);

    # Scale v[ii]
    v = np.array(v)/strain_range
    for ii in range(v.shape[0]):
        if (v[ii]<-1):
            v[ii]=-1
        if (v[ii]>1):
            v[ii]=1
    v = (v+1)*50
 
    ztmp = []
    print("Plotting...\n")
    for ii in pbar(range(len(x))):
        ztmp.append(z[ii][0])
        voro_patch(ax,x[ii],y[ii],z[ii],v[ii])

    ax.view_init(elev=0.,azim=-90.)
    plt.axis('off')
    plt.show()
