from change import switch_to_xy,switch_to_location,judge_xy,rotate,judge_rotate
from treenode import Node
import math
def add_node(treenode,location,*args):
    if len(args)==0:
        flag=0 #默认为0
    else:
        flag=args[0] #可能为0，可能为1，分别代表扩展所有和只扩展一步
    type=judge_rotate(treenode,location)
    if type==0: #停止
        sub_node=Node()
        sub_node.location=treenode.location
        sub_node.rotation=treenode.rotation
        treenode.add_child(sub_node)
    elif type==1: #转向
        [x,y]=switch_to_xy(treenode.location)
        [x_flag,y_flag]=judge_xy(treenode.location,location)
        sub_node = Node()
        sub_node.location = treenode.location
        sub_node.rotation = rotate(treenode.rotation, x_flag, y_flag)
        treenode.add_child(sub_node)
        if flag==0:
            sub_sub_node = Node()
            sub_sub_node.location = switch_to_location([x + x_flag, y + y_flag])
            sub_sub_node.rotation = sub_node.rotation
            sub_node.add_child(sub_sub_node)
            return sub_sub_node

    elif type==2: #前进
        [x, y] = switch_to_xy(treenode.location)
        [x_flag, y_flag] = judge_xy(treenode.location, location)
        sub_node = Node()
        sub_node.location = switch_to_location([x + x_flag, y + y_flag])
        sub_node.rotation = treenode.rotation
        treenode.add_child(sub_node)
    else: #后退
        [x, y] = switch_to_xy(treenode.location)
        [x_flag, y_flag] = judge_xy(treenode.location, location)
        epsilon=math.pow(10,-5)
        # 做出后退决定后，需要同时扩展出三次，分别以sub_node、sub_sub_node、sub_sub_sub_node标识
        sub_node = Node()
        sub_node.location = treenode.location
        sub_node.rotation=treenode.rotation+1/2*math.pi
        if abs(sub_node.rotation-2*math.pi)<epsilon:
            sub_node.rotation=0
        treenode.add_child(sub_node)
        if flag==0:
            sub_sub_node = Node()
            sub_sub_node.location = sub_node.location
            sub_sub_node.rotation = sub_node.rotation+1/2*math.pi
            if abs(sub_node.rotation-2*math.pi)<epsilon:
                sub_node.rotation=0
            sub_node.add_child(sub_sub_node)
            sub_sub_sub_node = Node()
            sub_sub_sub_node.location = switch_to_location([x + x_flag, y + y_flag])
            sub_sub_sub_node.rotation = sub_sub_node.rotation
            sub_sub_node.add_child(sub_sub_sub_node)
            return sub_sub_sub_node
    return sub_node