from OCC.Core.TopLoc import TopLoc_Location
from OCC.Core.TDocStd import TDocStd_Document
from OCC.Core.IFSelect import IFSelect_RetDone
from OCC.Core.TDF import TDF_LabelSequence, TDF_Label
from OCC.Core.STEPCAFControl import STEPCAFControl_Reader
from OCC.Core.BRepBuilderAPI import BRepBuilderAPI_Transform
from OCC.Core.Quantity import Quantity_Color, Quantity_TOC_RGB
from OCC.Core.XCAFDoc import XCAFDoc_DocumentTool, XCAFDoc_ColorTool

def readStep(filename):

    output_shapes = {}

    # Create a document handler
    doc = TDocStd_Document("aDoc")

    # Get tools to handle parts
    shape_tool = XCAFDoc_DocumentTool.ShapeTool(doc.Main())
    color_tool = XCAFDoc_DocumentTool.ColorTool(doc.Main())

    step_reader = STEPCAFControl_Reader()
    step_reader.SetColorMode(True)
    step_reader.SetNameMode(True)

    status = step_reader.ReadFile(filename)
    if status == IFSelect_RetDone:
        step_reader.Transfer(doc)

    locs = []

    def getSubShapes(lab, loc):

        l_subss = TDF_LabelSequence()
        shape_tool.GetSubShapes(lab, l_subss)

        l_comps = TDF_LabelSequence()
        shape_tool.GetComponents(lab, l_comps)

        name = lab.GetLabelName()

        # Check if part is Assembly
        if shape_tool.IsAssembly(lab):

            shape = shape_tool.GetShape(lab)

            l_c = TDF_LabelSequence()
            shape_tool.GetComponents(lab, l_c)

            label = l_c.Value(1)
            name = label.GetLabelName()

            if shape_tool.IsReference(label):

                label_reference = TDF_Label()
                shape_tool.GetReferredShape(label, label_reference)
                loc = shape_tool.GetLocation(label)

                print(f"NAME : {name}")
                print("TYPE : ASSEMBLY")

                # Retrieve translation and rotation data
                trans = loc.Transformation()

                rot = trans.GetRotation()
                # angles = rot.GetEulerAngles(gp_EulerSequence.gp_Extrinsic_XYZ)

                print("ROTATION:")
                print(" X  : {}".format(rot.X()))
                print(" Y  : {}".format(rot.Y()))
                print(" Z  : {}".format(rot.Z()))
                print(" W  : {}".format(rot.W()))

                tran = trans.TranslationPart()
                print("TRANSLATION  :")
                print(" X  : {}".format(tran.X()))
                print(" Y  : {}".format(tran.Y()))
                print(" Z  : {}".format(tran.Z()))

            l_c = TDF_LabelSequence()
            shape_tool.GetComponents(lab, l_c)
            for i in range(l_c.Length()):
                label = l_c.Value(i + 1)
                name = label.GetLabelName()
                if shape_tool.IsReference(label):

                    label_reference = TDF_Label()
                    shape_tool.GetReferredShape(label, label_reference)
                    loc = shape_tool.GetLocation(label)

                    locs.append(loc)
                    getSubShapes(label_reference, loc)
                    locs.pop()

        elif shape_tool.IsSimpleShape(lab):

            shape = shape_tool.GetShape(lab)

            # Get Location and Oriention in World
            loc = TopLoc_Location()
            for l in locs:
                loc = loc.Multiplied(l)

            print(f"NAME : {name}")
            print("TYPE : PART")

            # Retrieve translation and rotation data
            trans = loc.Transformation()

            rot = trans.GetRotation()
            print("ROTATION :")
            print(" X  : {}".format(rot.X()))
            print(" Y  : {}".format(rot.Y()))
            print(" Z  : {}".format(rot.Z()))
            print(" W  : {}".format(rot.W()))

            tran = trans.TranslationPart()
            print("TRANSLATION  :")
            print(" X  : {}".format(tran.X()))
            print(" Y  : {}".format(tran.Y()))
            print(" Z  : {}".format(tran.Z()))

            c = Quantity_Color(0.5, 0.5, 0.5, Quantity_TOC_RGB)  # default color
            color_set = False
            if (
                color_tool.GetInstanceColor(shape, 0, c)
                or color_tool.GetInstanceColor(shape, 1, c)
                or color_tool.GetInstanceColor(shape, 2, c)
            ):
                color_tool.SetInstanceColor(shape, 0, c)
                color_tool.SetInstanceColor(shape, 1, c)
                color_tool.SetInstanceColor(shape, 2, c)
                color_set = True
                n = c.Name(c.Red(), c.Green(), c.Blue())
                print(
                    "    instance ONE color Name & RGB: ",
                    c,
                    n,
                    c.Red(),
                    c.Green(),
                    c.Blue(),
                )

            if not color_set:
                if (
                    XCAFDoc_ColorTool.GetColor(lab, 0, c)
                    or XCAFDoc_ColorTool.GetColor(lab, 1, c)
                    or XCAFDoc_ColorTool.GetColor(lab, 2, c)
                ):
                    color_tool.SetInstanceColor(shape, 0, c)
                    color_tool.SetInstanceColor(shape, 1, c)
                    color_tool.SetInstanceColor(shape, 2, c)

                    n = c.Name(c.Red(), c.Green(), c.Blue())
                    print(
                        "    shape ONE color Name & RGB: ",
                        c,
                        n,
                        c.Red(),
                        c.Green(),
                        c.Blue(),
                    )

            shape_disp = BRepBuilderAPI_Transform(shape, loc.Transformation()).Shape()
            if shape_disp not in output_shapes:
                output_shapes[shape_disp] = [lab.GetLabelName(), c]

            if(l_subss.Length() > 0):

                lab_subs = l_subss.Value(1)
                shape_sub = shape_tool.GetShape(lab_subs)

                c = Quantity_Color(0.5, 0.5, 0.5, Quantity_TOC_RGB)  # default color
                color_set = False
                if (
                    color_tool.GetInstanceColor(shape_sub, 0, c)
                    or color_tool.GetInstanceColor(shape_sub, 1, c)
                    or color_tool.GetInstanceColor(shape_sub, 2, c)
                ):
                    color_tool.SetInstanceColor(shape_sub, 0, c)
                    color_tool.SetInstanceColor(shape_sub, 1, c)
                    color_tool.SetInstanceColor(shape_sub, 2, c)
                    color_set = True
                    n = c.Name(c.Red(), c.Green(), c.Blue())
                    print(
                        "    instance TWO color Name & RGB: ",
                        c,
                        n,
                        c.Red(),
                        c.Green(),
                        c.Blue(),
                    )

                if not color_set:
                    if (
                        XCAFDoc_ColorTool.GetColor(lab_subs, 0, c)
                        or XCAFDoc_ColorTool.GetColor(lab_subs, 1, c)
                        or XCAFDoc_ColorTool.GetColor(lab_subs, 2, c)
                    ):
                        color_tool.SetInstanceColor(shape, 0, c)
                        color_tool.SetInstanceColor(shape, 1, c)
                        color_tool.SetInstanceColor(shape, 2, c)

                        print("COLORS :")
                        print(" RED     : {}".format(c.Red()))
                        print(" GREEN   : {}".format(c.Green()))
                        print(" BLUE    : {}\n".format(c.Blue()))

            else:
                print("COLORS:")
                print(" RED     : 0.50000")
                print(" GREEN   : 0.50000")
                print(" BLUE    : 0.50000\n")

    def getShapes():

        labels = TDF_LabelSequence()
        shape_tool.GetFreeShapes(labels)

        for i in range(labels.Length()):
            root_item = labels.Value(i + 1)
            getSubShapes(root_item, None)

    getShapes()

    return output_shapes

def main():

    readStep('Avionics.step')

if __name__ == "__main__":
    main()