import os
has_x = ('DISPLAY' in os.environ.keys()) and (os.environ['DISPLAY'] is not None)

import numpy as np
import matplotlib.pyplot as plt

from cellcomplex.property_topomesh.example_topomesh import vtk_ellipsoid_topomesh
from cellcomplex.property_topomesh.extraction import cut_surface_topomesh
from cellcomplex.property_topomesh.analysis import compute_topomesh_property, compute_topomesh_vertex_property_from_faces
from cellcomplex.property_topomesh.utils.matplotlib_tools import mpl_draw_topomesh, mpl_draw_incidence_graph

from cellcomplex.property_topomesh.visualization.vtk_actor_topomesh import VtkActorTopomesh
from cellcomplex.property_topomesh.visualization.vtk_tools import vtk_image_actors

topomesh = vtk_ellipsoid_topomesh(ellipsoid_radius=2,ellipsoid_scales=[1.5,1,1],subdivisions=2)
topomesh = cut_surface_topomesh(topomesh,z_offset=0,below=False)

compute_topomesh_property(topomesh,'normal',2,normal_method='orientation')
compute_topomesh_property(topomesh,'area',2)
compute_topomesh_vertex_property_from_faces(topomesh,'normal',weighting='area',adjacency_sigma=1.2,neighborhood=3)
compute_topomesh_property(topomesh,'mean_curvature',2)

figure = plt.figure(0)
figure.clf()

if has_x:
    figure.add_subplot(1,2,1)
figure.gca().axis('equal')

col = mpl_draw_topomesh(topomesh,figure,2,property_name='mean_curvature',colormap='Reds',intensity_range=(0.25,1))
mpl_draw_topomesh(topomesh,figure,0,property_name='normal',color='g',coef=0.5,linewidth=1)
mpl_draw_topomesh(topomesh,figure,1,color='k',linewidth=0.5)
mpl_draw_topomesh(topomesh,figure,0,color='k',size=10)

cbar = figure.colorbar(col)
cbar.set_label('Face Mean Curvature')

if has_x:
    figure.add_subplot(1,2,2)

    actors = []

    curvature_actor = VtkActorTopomesh(topomesh,2,property_name='mean_curvature')
    curvature_actor.update(colormap='Reds',value_range=(0.25,1))
    actors += [curvature_actor.actor]

    edge_actor = VtkActorTopomesh(topomesh,1,line_glyph='tube',glyph_scale=0.01)
    edge_actor.update(colormap='Greys',value_range=(0,0),linewidth=3,opacity=1)
    actors += [edge_actor.actor]

    vertex_actor = VtkActorTopomesh(topomesh,0,glyph_scale=0.05)
    vertex_actor.update(colormap='Greys',value_range=(0,0),opacity=1)
    actors += [vertex_actor.actor]

    normal_actor = VtkActorTopomesh(topomesh,0,property_name='normal',vector_glyph='arrow',glyph_scale=0.5)
    normal_actor.update(colormap='YlGn')
    actors += [normal_actor.actor]

    figure.gca().imshow(vtk_image_actors(actors,focal_point=(0,0,1),view_up=(0,1,0)))

    figure.gca().axis('off')

    figure.set_size_inches(10,5)
figure.tight_layout()