#!/usr/bin/env python
# coding: utf-8

# In[115]:


import os

from OCC.Core.IFSelect import IFSelect_RetDone
from OCC.Core.STEPCAFControl import STEPCAFControl_Reader
from OCC.Core.TCollection import TCollection_ExtendedString
from OCC.Core.TDocStd import TDocStd_Document
from OCC.Core.VrmlAPI import VrmlAPI_Writer
from OCC.Core.VrmlData import VrmlData_WorldInfo
from OCC.Core.XCAFDoc import XCAFDoc_ShapeTool, XCAFDoc_DocumentTool, XCAFDoc_DocumentTool_ShapeTool
from OCC.Core.XCAFApp import XCAFApp_Application_GetApplication
from OCC.Core.XCAFDoc import XCAFDoc_DocumentTool_ShapeTool, xcafdoc
from OCC.Core.BRep import BRep_Tool
from OCC.Core.TDF import TDF_LabelSequence, TDF_Label, TDF_ChildIterator, TDF_AttributeIterator, TDF_Tool, TDF_AttributeMap, TDF_Attribute, TDF_ChildIDIterator
from OCC.Core.TCollection import TCollection_AsciiString
from OCC.Core.TDataStd import TDataStd_TreeNode
from OCC.Core.Standard import Standard_GUID #, Standard_Handle
#from OCC.Core.Standard import Standard_Real 
from OCC.Core.TNaming import TNaming_UsedShapes
from OCC.Extend.TopologyUtils import TopologyExplorer
from OCC.Extend.TopologyUtils import is_edge, is_wire, discretize_edge, discretize_wire
from OCC.Extend.DataExchange import read_step_file_with_names_colors, read_step_file
from OCC.Display.WebGl.x3dom_renderer import X3DExporter, X3DomRenderer
from OCC.Core.Tesselator import ShapeTesselator
from OCC.Core.TopAbs import topabs
from OCC.Core.BRepMesh import BRepMesh_IncrementalMesh
from OCC.Core.gp import gp_XYZ, gp_Vec

import xml.etree.ElementTree as ET
from IPython.display import HTML
import x3d.x3d as XX3D
import re


# In[162]:


#stp_filename = 'VentilatorAP203' #'as1_pe_203'
stp_filename = 'as1-oc-214'
#stp_filename = 'as1_pe_203'
#stp_filename = 'screw'
stp_path = os.path.join('assets', stp_filename + '.stp')

# load the STEP file
doc = TDocStd_Document(TCollection_ExtendedString("pythonocc-doc"))

step_reader = STEPCAFControl_Reader()
step_reader.SetColorMode(True)
step_reader.SetLayerMode(True)
step_reader.SetNameMode(True)
step_reader.SetMatMode(True)
step_reader.SetGDTMode(True)

status = step_reader.ReadFile(stp_path)
if status == IFSelect_RetDone:
    step_reader.Transfer(doc)
else:
	raise IOError("STEP file could not be read.")


# In[163]:


doc


# In[176]:


##Copyright 2018 Thomas Paviot (tpaviot@gmail.com)
##(C) 2020 Andreas Plesch
##
##This file is part of pythonOCC.
##
##pythonOCC is free software: you can redistribute it and/or modify
##it under the terms of the GNU Lesser General Public License as published by
##the Free Software Foundation, either version 3 of the License, or
##(at your option) any later version.
##
##pythonOCC 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 Lesser General Public License for more details.
##
##You should have received a copy of the GNU Lesser General Public License
##along with pythonOCC.  If not, see <http://www.gnu.org/licenses/>.

import os

from OCC.Core.TopoDS import TopoDS_Shape
from OCC.Core.BRepMesh import BRepMesh_IncrementalMesh
from OCC.Core.StlAPI import stlapi_Read, StlAPI_Writer
from OCC.Core.BRep import BRep_Builder
from OCC.Core.gp import gp_Pnt, gp_Dir, gp_Pnt2d
from OCC.Core.Bnd import Bnd_Box2d
from OCC.Core.TopoDS import TopoDS_Compound
from OCC.Core.IGESControl import IGESControl_Reader, IGESControl_Writer
from OCC.Core.STEPControl import STEPControl_Reader, STEPControl_Writer, STEPControl_AsIs
from OCC.Core.Interface import Interface_Static_SetCVal
from OCC.Core.IFSelect import IFSelect_RetDone, IFSelect_ItemsByEntity
from OCC.Core.TDocStd import TDocStd_Document
from OCC.Core.XCAFDoc import (XCAFDoc_DocumentTool_ShapeTool,
                              XCAFDoc_DocumentTool_ColorTool)
from OCC.Core.STEPCAFControl import STEPCAFControl_Reader
from OCC.Core.TDF import TDF_LabelSequence, TDF_Label
from OCC.Core.TCollection import TCollection_ExtendedString
from OCC.Core.Quantity import Quantity_Color, Quantity_TOC_RGB
from OCC.Core.TopLoc import TopLoc_Location
from OCC.Core.BRepBuilderAPI import BRepBuilderAPI_Transform

from OCC.Extend.TopologyUtils import (discretize_edge, get_sorted_hlr_edges,
                                      list_of_shapes_to_compound)

try:
    import svgwrite
    HAVE_SVGWRITE = True
except ImportError:
    HAVE_SVGWRITE = False

def read_document_with_names_colors(doc):
    """ Returns list of tuples (topods_shape, label, color)
    Use OCAF.
    """
    output_shapes = {}
    scene = []
    visited = {}
    DEFset = set()
    x3dscene = XX3D.Scene(children=[])
    x3d_doc = XX3D.X3D(Scene = x3dscene)
    
    # Get root assembly
    shape_tool = XCAFDoc_DocumentTool_ShapeTool(doc.Main())
    color_tool = XCAFDoc_DocumentTool_ColorTool(doc.Main())
    
    locs = []

    def _get_sub_shapes(lab, loc, parent, x3dgroup):
        
        name = lab.GetLabelName()
        labelString = lab.EntryDumpToString()
        if labelString in visited:
            return
        visited[labelString] = lab

        if shape_tool.IsAssembly(lab):
            node = {
                    'node' : 'Group',
                    'DEF' : labelString,
                    'name' : name,
                    'children' : []
            }
            x3dnode = XX3D.Group(DEF = node['DEF'], class_ = node['name'], children = [])
            #print('assembly group:', x3dnode.XML())
            l_c = TDF_LabelSequence()
            shape_tool.GetComponents(lab, l_c)
            for i in range(l_c.Length()):
                label = l_c.Value(i + 1)
                print("Group Name DEF :", name, labelString)    
                if shape_tool.IsReference(label):
                    print("########  component label :", label.GetLabelName() )
                    loc = shape_tool.GetLocation(label)
                    print(" Transform  loc DEF          :", loc.HashCode(100))
                    trafo = {
                        'node' : 'Transform',
                        'DEF' : label.EntryDumpToString(),
                        'name' : 'reference location',
                        'transform' : loc,
                        'children': []
                    }
                    x3dtransform = XX3D.Transform(DEF = trafo['DEF'], class_ = trafo['name'], children = [])
                    _x3dapplyLocation(x3dtransform, loc)
                    #print('locator transform:', x3dtransform.XML())
                    
                    locs.append(loc)
                    print("loc chain depth :", len(locs))
                    
                    label_reference = TDF_Label()
                    shape_tool.GetReferredShape(label, label_reference)
                    print("########  Transform USE to DEF ==> referenced label :", label_reference.GetLabelName() )
                    reference_entry = label_reference.EntryDumpToString()
                    if reference_entry not in DEFset:
                        DEFset.add(reference_entry)
                        _get_sub_shapes(label_reference, loc, trafo['children'], x3dtransform.children)
                    else:
                        reference = {
                            'node' : 'Transform',
                            'USE' : reference_entry,
                            'refname' : label_reference.GetLabelName()
                        }
                        trafo['children'].append(reference)
                        x3dtransform.children.append(XX3D.Transform(USE = reference['USE'], class_ = reference['refname']))
                    locs.pop()
                    node['children'].append(trafo)
                    x3dnode.children.append(x3dtransform)
                    #print("added child: ", x3dnode.XML())

        elif shape_tool.IsSimpleShape(lab):
            print("Transform DEF Shape Name :", name, labelString )
            shape = shape_tool.GetShape(lab)
            #print("    all ass locs   :", locs)
#             loc = TopLoc_Location()
#             for l in locs:
#                 loc = loc.Multiplied(l)

            c = _set_color(lab, shape)
            
            n = c.Name(c.Red(), c.Green(), c.Blue())
            print('    instance color Name & RGB: ', n, c.Red(), c.Green(), c.Blue())
            
            labloc = shape_tool.GetLocation(lab)
            print("    Shape Transform: ", labloc.HashCode(100))
     
            #shape_disp = BRepBuilderAPI_Transform(shape, loc.Transformation()).Shape()
            #shape_disp = shape
            #if not shape_disp in output_shapes:
            #    output_shapes[shape_disp] = [lab.GetLabelName(), c]
            
            ##subshapes
            l_subss = TDF_LabelSequence()
            shape_tool.GetSubShapes(lab, l_subss)
            
            #node = {}
            
            if (l_subss.Length() == 0 and labloc.IsIdentity()): # does not need transform
                node = {
                    'node' : 'Shape',
                    'DEF' : labelString,
                    'shape' : lab,
                    'name' : name + '-shape',
                    'color' : f"{c.Red()} {c.Green()} {c.Blue()}"
                }
                x3dnode = XX3D.Shape(
                    DEF = node['DEF'],
                    geometry = _x3dgeofromTShape(node['shape']),
                    appearance = _x3dappfromColor(c),
                    class_ = node['name']
                )
            else: # needs grouping or has transform
                node = {
                    'node' : 'Transform',
                    'DEF' : labelString,
                    'transform' : labloc,
                    'transformhash' : labloc.HashCode(100),
                    'name' : name,
                    'children' : []
                }

                shapenode = {
                    'node' : 'Shape',
                    'shape' : lab,
                    'name' : name + '-shape',
                    'color' : f"{c.Red()} {c.Green()} {c.Blue()}"
                }
                
                node['children'].append(shapenode)
                                
                x3dnode = XX3D.Transform(DEF = node['DEF'], class_ = node['name'], children = [])
                _x3dapplyLocation(x3dnode, labloc)
                
                x3dshape = XX3D.Shape (
                        geometry = _x3dgeofromTShape(shapenode['shape']),
                        appearance = _x3dappfromColor(c),
                        class_ = shapenode['name']
                )
                x3dnode.children.append(x3dshape)
            
            for i in range(l_subss.Length()):
                lab_subs = l_subss.Value(i+1)
                print("########  Transform DEF simpleshape subshape label :", lab.GetLabelName())
                shape_sub = shape_tool.GetShape(lab_subs)

                c = _set_color(lab_subs, shape_sub)
                
                n = c.Name(c.Red(), c.Green(), c.Blue())
                print('    shape color Name & RGB: ', n, c.Red(), c.Green(), c.Blue())
            
                subloc = shape_tool.GetLocation(lab_subs)
                print("    subshape Transform: ", subloc.HashCode(100))
                shapenode = {
                    'node' : 'SubShape',
                    'shape' : lab_subs,
                    'DEF' : lab_subs.EntryDumpToString(),
                    'name' : lab_subs.GetLabelName()+'-subshape',
                    'color' : f"{c.Red()} {c.Green()} {c.Blue()}",
                    'trafo' : subloc                        
                    }
                node['children'].append(shapenode)
                x3dshape = XX3D.Shape (
                    DEF = shapenode['DEF'],
                    class_ = shapenode['name'],
                    geometry = _x3dgeofromTShape(shapenode['shape']),
                    appearance = _x3dappfromColor(c)
                )
                x3dnode.children.append(x3dshape)
                
                #shape_to_disp = BRepBuilderAPI_Transform(shape_sub, loc.Transformation()).Shape()
                #shape_to_disp = shape_sub
                # position the subshape to display
                #if not shape_to_disp in output_shapes:
                #    output_shapes[shape_to_disp] = [lab_subs.GetLabelName(), c]
                    
        parent.append(node)
        #print(x3dgroup.XML())
        x3dgroup.append(x3dnode)
        #print(x3dscene.XML())
        #print(x3dgroup.XML())
    
    def _x3dapplyLocation(x3dtransformnode, location):
        # get translation and rotation from location
        transformation = location.Transformation()
        rot_axis = gp_XYZ()
        #rot_angle = 0.0
        success, rot_angle = transformation.GetRotation(rot_axis)#.GetVectorAndAngle(rot_axis, rot_angle)
        translation = transformation.TranslationPart()
        scale_factor = transformation.ScaleFactor()
        x3dtransformnode.rotation = (rot_axis.X(), rot_axis.Y(), rot_axis.Z(), rot_angle)
        x3dtransformnode.translation = (translation.X(), translation.Y(), translation.Z())
        x3dtransformnode.scale = (scale_factor, scale_factor, scale_factor)

        return
    
    def _x3dgeofromTShape(label):
        shape = shape_tool.GetShape(label)
        tesselator = ShapeTesselator(shape)
        tesselator.Compute(
            compute_edges=False,
            mesh_quality=1,
            parallel=True)
        x3dstring = tesselator.ExportShapeToX3DIndexedFaceSet()#x3dexp._triangle_sets[0] # there should be just one
        element = ET.XML(x3dstring)
        geo = XX3D.Box()
        if (element.tag == 'TriangleSet'):
            coordele = list(element.iter('Coordinate'))[0]
            normalele = list(element.iter('Normal'))[0]
            coord = XX3D.Coordinate(point = _MFVec3ffromString(coordele.attrib['point']))
            normal = XX3D.Normal(vector = _MFVec3ffromString(normalele.attrib['vector']))
            geo = XX3D.TriangleSet(coord = coord, normal=normal, solid=False)
        # get tesselated triangleset or lineset
        return geo
    
    def _MFVec3ffromString(sepString):
        mflist = sepString.split()
        mf = []
        for i in range(len(mflist)):
            if (i % 3 == 2):
                mf.append((float(mflist[i-2]), float(mflist[i-1]), float(mflist[i])))
        return mf
    
    def _x3dappfromColor(c):
        x3dmat = XX3D.Material(diffuseColor = (c.Red(),c.Green(),c.Blue()))
        return XX3D.Appearance(material = x3dmat)

    def _set_color(lab, shape):
        c = Quantity_Color(0.5, 0.5, 0.5, Quantity_TOC_RGB)  # default color
        colorSet = False
        if (color_tool.GetInstanceColor(shape, 0, c) or
            color_tool.GetInstanceColor(shape, 1, c) or
            color_tool.GetInstanceColor(shape, 2, c)):

            colorSet = True

        if not colorSet:
            if (color_tool.GetColor(lab, 0, c) or
                color_tool.GetColor(lab, 1, c) or
                color_tool.GetColor(lab, 2, c)):

                colorSet = True

        if colorSet:
            color_tool.SetInstanceColor(shape, 0, c)
            color_tool.SetInstanceColor(shape, 1, c)
            color_tool.SetInstanceColor(shape, 2, c)
        
        return c

    def _get_shapes():
        labels = TDF_LabelSequence()
        shape_tool.GetFreeShapes(labels)
        #global cnt
        #cnt += 1
       
        print()
        print("Number of shapes at root :", labels.Length())
        print()
        for i in range(labels.Length()):
            root_item = labels.Value(i+1)
            _get_sub_shapes(root_item, None, scene, x3dscene.children)
    _get_shapes()
    print('DONE')
    return { 'scene' : scene, 'x3d_doc' : x3d_doc } 



# In[ ]:


ascene=read_document_with_names_colors(doc)
ascene['scene']


# In[178]:


print(ascene['x3d_doc'].Scene.XML())


# In[179]:


x3domHEAD = '''<script type='text/javascript' src='https://www.x3dom.org/download/dev/x3dom-full.debug.js'> </script> 
<link rel='stylesheet' type='text/css' href='https://www.x3dom.org/download/dev/x3dom.css'></link>'''


# In[180]:


x3dele = list(ET.XML(ascene['x3d_doc'].XML()).iter('X3D'))[0]
x3dHTML = ET.tostring(x3dele, encoding="unicode", short_empty_elements=False)


# In[181]:


HTML(x3domHEAD + x3dHTML)


# In[ ]: