#stp_filename = 'VentilatorAP203' #'as1_pe_203'
stp_filename = 'as1-oc-214'
#stp_filename = 'as1_pe_203'
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.")

##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()

    # 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):
        #global cnt, lvl
        #cnt += 1
        #print("\n[%d] level %d, handling LABEL %s\n" % (cnt, lvl, _get_label_name(lab)))
        #print()
        #print(lab.DumpToString())
        #print()
        #print("Is Assembly    :", shape_tool.IsAssembly(lab))
        #print("Is Free        :", shape_tool.IsFree(lab))
        #print("Is Shape       :", shape_tool.IsShape(lab))
        #print("Is Compound    :", shape_tool.IsCompound(lab))
        #print("Is Component   :", shape_tool.IsComponent(lab))
        #print("Is SimpleShape :", shape_tool.IsSimpleShape(lab))
        #print("Is Reference   :", shape_tool.IsReference(lab))

        #users = TDF_LabelSequence()
        #users_cnt = shape_tool.GetUsers(lab, users)
        #print("Nr Users       :", users_cnt)

        #print("Nb subshapes   :", l_subss.Length())
        #l_comps = TDF_LabelSequence()
        #shape_tool.GetComponents(lab, l_comps)
        #print("Nb components  :", l_comps.Length())
        #print()
        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' : []
            }
            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': []
                    }
                    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'])
                    else:
                        trafo['children'].append({
                        'node' : 'Transform',
                        'USE' : reference_entry,
                        'refname' : label_reference.GetLabelName()
                        })
                    locs.pop()
                    node['children'].append(trafo)

        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))
            
            node = {
                'node' : 'Transform',
                'DEF' : labelString,
                'name' : name,
            }
            
            #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)
            
            if (l_subss.Length() == 0 and labloc.IsIdentity()): # does not need transform
                node['node'] = 'Shape'
                node['shape'] = lab
                node['name'] = name+'-shape'
                node['color'] = f"{c.Red()} {c.Green()} {c.Blue()}"
            else: # needs grouping or has transform
                node['transform'] = labloc
                node['children'] = []
                node['transformhash'] = labloc.HashCode(100)
                node['children'].append({
                    'node' : 'Shape',
                    'shape' : lab,
                    'name' : name+'-shape',
                    'color' : f"{c.Red()} {c.Green()} {c.Blue()}"
                })
            
            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))
                node['children'].append({
                    'node' : 'SubShape',
                    'shape' : lab_subs,
                    'DEF' : lab_subs.EntryDumpToString(),
                    'name' : lab_subs.GetLabelName()+'-subshape',
                    'color' : f"{c.Red()} {c.Green()} {c.Blue()}",
                    'trafo' : subloc                        
                    })
                #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)    

    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)
    _get_shapes()
    print('DONE')
    return scene