import numpy as np
import ifcopenshell

# Edit placements and delete the ones that places spaces
def a2p(o, z, x):
    y = np.cross(z, x)
    r = np.eye(4)
    r[:-1,:-1] = x,y,z
    r[-1,:-1] = o
    return r.T

def get_axis2placement(plc):
    z = np.array(plc.Axis.DirectionRatios if plc.Axis else (0,0,1))
    x = np.array(plc.RefDirection.DirectionRatios if plc.RefDirection else (1,0,0))
    o = plc.Location.Coordinates
    return a2p(o,z,x)

def get_local_placement(plc):
    if plc is None:
        return np.eye(4)
    if plc.PlacementRelTo is None:
        parent = np.eye(4)
    else:
        parent = get_local_placement(plc.PlacementRelTo)

    # OR
    # return np.dot(get_axis2placement(plc.RelativePlacement), parent)
    # OR
    return np.dot(parent, get_axis2placement(plc.RelativePlacement))

def create_axis2placement(f, x, z, o):
    x = f.createIfcDirection(x)
    z = f.createIfcDirection(z)
    o = f.createIfcCartesianPoint(o)
    axes = f.createIfcAxis2Placement3D(o, z, x)

    return axes

f = ifcopenshell.file()

# ----- initial data for testing ------
relPlc_storey = create_axis2placement(
    f,
    x=(1., 0., 0.),
    z=(0., 0., 1.),
    o=(10., 0., 0.)
)
plc_storey = f.createIfcLocalPlacement(None, relPlc_storey)

relPlc_space = create_axis2placement(
    f,
    x=(2**0.5 / 2, 2**0.5 / 2, 0.),
    z=(0., 0., 1.),
    o=(100., 0., 0.)
)
plc_space = f.createIfcLocalPlacement(plc_storey, relPlc_space)

relPlc_ref = create_axis2placement(
    f,
    x=(1., 0., 0.),
    z=(0., 0., 1.),
    o=(0., -50., 0.)
)
plc_ref = f.createIfcLocalPlacement(plc_space, relPlc_ref)

m = get_local_placement(plc_ref)
print(m)
print(m[0][3], m[1][3], m[2][3])

# -------------------

# local_placements = f.by_type('IfcLocalPlacement')
local_placements = [plc_space]
for lp in local_placements:
    # if lp.PlacesObject[0].is_a('IfcSpace'):
    if True: # lp.PlacesObject[0].is_a('IfcSpace'):
        lp_space = lp
        R_space = get_axis2placement(lp_space.RelativePlacement)
        plcRelTo_space = lp_space.PlacementRelTo

        for lp_ref in lp_space.ReferencedByPlacements:
            R_ref = get_axis2placement(lp_ref.RelativePlacement)

            # OR
            # R_new = R_ref.dot(R_space)
            # OR
            R_new = R_space.dot(R_ref)

            print(R_new)

            relPlc_new = create_axis2placement(
                f,
                x=R_new[:-1, 0].tolist(),
                z=R_new[:-1, 2].tolist(),
                o=R_new[:-1, -1].tolist()
            )
            plc_new = f.createIfcLocalPlacement(plcRelTo_space, relPlc_new)
            # assign plc_new to ObjectPlacement

f.write('test.ifc')














#