import sys


def insert_transformation(center,text):
    transform_id=text.find("Transform {")+11
    head=text[:transform_id]
    back=text[transform_id:]
    insert_str="\ntranslation %f %f %f\nrotation  0 1 0 0"%(center[0],center[1],center[2])
    return head+insert_str+back

def convert_center(pts_str):
    pts=pts_str.split(',')
    if len(pts)<2:  #该点格式为以 '\n'分割
        pts=pts_str.split('\n')
    sumx=0
    sumy=0
    sumz=0
    size=0
    points=[]
    for pt in pts:
        xyz_str=pt.replace('\n','').rstrip().lstrip()
        xyzstrs=xyz_str.split(' ')
        if len(xyzstrs)<3:
            continue
        [strx,stry,strz]=xyzstrs
        x=float(strx)
        y=float(stry)
        z=float(strz)
        size=size+1
        sumx+=x
        sumy+=y
        sumz+=z
        points.append([x,y,z])
    ax=sumx/size
    ay=sumy/size
    az=sumz/size

    new_pts_str=''
    for i in range(len(points)):
        [x,y,z]=points[i]
        [nx,ny,nz]=[x-ax,y-ay,z-az]
        if i ==len(points)-1:
            new_pts_str=new_pts_str+"%f %f %f"%(nz,nx,ny)
        else:
            new_pts_str=new_pts_str+"%f %f %f,"%(nz,nx,ny)


    return [[ax,ay,az],new_pts_str]


def transform_solid(solid_str,shape_name):
    print("------------------------------------------------------------------------------------------------------:\n")
    if solid_str.find("Transform")<0:
        return [False,""]
    solid_str=solid_str.replace("Transform","DEF Trans_%s Transform"%(shape_name))
    index_beg=solid_str.find("point [")+7
    if index_beg<0:
        return [False,""]
    index_end=solid_str.find("]",index_beg)
    if index_end<index_beg:
        return [False,""]
    pts_str=solid_str[index_beg:index_end]

    [center,new_pts_str]=convert_center(pts_str)
    head=solid_str[:index_beg]
    back=solid_str[index_end:-1]

    new_str=head+new_pts_str+back+'\n'
    #new_trransform_pts_str=insert_transformation(center,new_str)


    solid_head="Solid {\ntranslation %f %f %f\nrotation 0 1 0 0\nchildren [\n"%(center[2],center[0],center[1])
    solid_back="]\nname \"%s_solid\"\nboundingObject USE %s\nphysics Physics {}\n}\n"%(shape_name,shape_name)

    new_solid_str=solid_head+new_str+solid_back

    return [True,new_solid_str]

def transform_file(file,outfile):
    buff=""
    writer=open(outfile, 'w', encoding='utf-8')
    writer.write("#VRML_SIM R2021a utf8\n")
    robot_head="Robot {\ntranslation 0 0.0 0\nrotation 0 1 0,0\nchildren [\n"
    robot_transform_head="Transform {\nchildren [\n"
    writer.write(robot_head)
    writer.write(robot_transform_head)

    shape_count=0
    Shape_name=''
    NoTransformStr=[]
    with open(file, 'r', encoding='utf-8') as f:
        for line in f:
            if line.find("#VRML_SIM R2021a utf8")>=0:
                continue
            if line.find('Shape')>=0:
                Shape_name="Shape_%d"%(shape_count)
                line=line.replace('Shape',"DEF %s Shape"%(Shape_name))
                print("DEF Shape name:%s  line:%s"%(Shape_name,line))
            if line.find("Transform")>=0:    #找到
                if len(buff)>0:
                    [ok,new_solid]=transform_solid(buff,Shape_name)
                    if ok==False:
                        NoTransformStr.append(buff)
                    else:
                        writer.write(new_solid)
                    shape_count+=1
                buff=line
            else:
                buff=buff+line
        #循环结束，处理最后一个solid
        if len(buff)>0:
            [ok,new_solid]=transform_solid(buff,Shape_name)
            if ok==False:
                NoTransformStr.append(buff)
            else:
                writer.write(new_solid)

    robot_transform_back="]\n}\n"
    robot_end="\n]\n}\n"
    writer.write(robot_transform_back)
    writer.write(robot_end)
    for buf in NoTransformStr:
        writer.write(buf)



if __name__=="__main__":
    transform_file("./AGV.wbt","./agv_test.wbt")