# Include files
include "string.bzz"
include "vec2.bzz"
include "utilities.bzz"
include "barrier.bzz"

############################################

# Global variables

RANGE = 200 # rab range in cm, get from argos file

N_SONS = 10 # Maximum number of sons

TRIGGER_VSTIG = 101 # ID of trigger virtual stigmergy
BARRIER_VSTIG = 102 # ID of barrier virtual stigmergy

################################################################
################################################################

# Tree utility functions

function parent_select() {

   
    # Selects potential parents
    var candidates = {}
    candidates = neighbors.filter(function(rid, data) {  
      return ((knowledge.level[rid] > 0) and (data.distance < (RANGE - 10.0)) and (knowledge.free[rid] == 1))})
    # and (data.distance > 10.0)

    # Selects closest parent candidate as parent
    var temp = {}
    temp = candidates.reduce(function(rid, data, accum) {
      accum.min = math.min(accum.min, data.distance)
      return accum
    }, {.min = RANGE})
    var min = temp.min


    var flag = 0
    foreach(knowledge.distance, function(key, value) {
      if ((flag == 0) and (candidates.data[key] != nil)) {
        if (value == min) {
          tree.parent.id = key
          tree.parent.distance = value
          tree.parent.azimuth = knowledge.azimuth[key]
          flag = 1
				}
      }
      #break (when implemented)
    })

}

####################################

function count() {

    if (nb_sons == 0) {
      eq_level = level
    }

    else if (nb_sons >= 1) {
      var temp = {}
      temp = sons.reduce(function(rid, data, accum) {
        accum.sum = accum.sum + tree.sons[rid].eq_level
        return accum
      }, {.sum = 0}) 
      eq_level = temp.sum - (nb_sons - 1) * level
    }

}

####################################

function change_frame(p01, p1, theta) {

    var p0 = {
      .x = math.cos(theta) * p1.x + math.sin(theta) * p1.y + p01.x,
      .y = -math.sin(theta) * p1.x + math.cos(theta) * p1.y + p01.y
    }
    return p0

}

####################################

transform_accum = function(rid, data) {

    # Son contribution in frame F1
    var p1 = {
      .x = tree.sons[rid].accum_x, 
      .y = tree.sons[rid].accum_y 
    }
 
    # Rotation angle between F0 and F1
    var theta = tree.sons[rid].angle_parent - data.azimuth - math.pi

    # Translation vector from F0 to F1
    var p01 = {
      .x = 0,
      .y = 0
    }
    
    var p0 = {}

    if (tree.sons[rid].angle_parent != nil) {
      var rot_angle = radians_interval(theta)
      p0 = change_frame(p01, p1, rot_angle)
    }
    else {
      p0.x = p1.x
      p0.y = p1.y
    }

    return p0
}

####################################

function centroid() {

    # If the robot has a parent
    if ((tree.parent != nil) or (root == 1)) {
      var sum_F1 = { .x = 0, .y = 0}

      # If the robot has at least one son
      if (nb_sons > 0) {
        var temp = {}
        # Expresses son contrib (in F1) in its own reference frame (F0)
        temp = sons.map(transform_accum)
        # Sums son contributions expressed in F0 frame
        sum_F1 = temp.reduce(function(rid, data, accum) {
          accum.x = accum.x + data.x
          accum.y = accum.y + data.y
          #log("id ", rid, " sum_x ", accum.x, " sum_y ", accum.y)
          return accum
        }, {.x = 0, .y = 0})  
      }

      # If the robot is not the root
      if ((root == 0) and (tree.parent.id != nil)) {
        #var nb_descendants = eq_level - level
        var p0 = knowledge.distance[tree.parent.id]#tree.parent.distance
        var theta = knowledge.azimuth[tree.parent.id]#tree.parent.azimuth
        # Adds current robot contribution to centroid sum
        accum_x = sum_F1.x - (nb_descendants + 1) * p0 * math.cos(theta)
        accum_y = sum_F1.y - (nb_descendants + 1) * p0 * math.sin(theta)
      }
      # If the robot is the root
      else if ((root == 1) and (robot_count != 0)) {
        # Centroid coordinates in root ref frame
        accum_x = sum_F1.x / robot_count
        accum_y = sum_F1.y / robot_count
      }
    }
}


################################################################
################################################################

# Tree reconfiguration functions

function tree_config() {
    statef()
}

function end_fun() {
    if (root == 1) {
      log("centroid X: ", accum_x, " Y ", accum_y)
    }
}

####################################

function root_select() {

    log(id," root_select")
    if (tree.parent.id != nil) {
      if(neighbors.data[tree.parent.id] != nil) {
        angle_parent = neighbors.data[tree.parent.id].azimuth
      }
    }

    if (root == 1) {

      # Listens for new root acknowledgment 
      
      foreach(knowledge.ackn, function(key, value) {
        if (value == better_root) {
          #log(id, " got it")
          root = 0
          level = 0
          setleds(0,0,0)
        }
      })
      
      if (ack == 1) {
        # Triggers transition to new state
        trigger.put("a", 1)
      }
      else if ((root == 1) and (better_root != id) and (trigger.get("a") != 1)) {
        setleds(255,0,0)
        better_root = id

			  # Centroid position in root reference frame (F0)
			  var c0 = math.vec2.new(accum_x, accum_y)

        # Distance from current root to centroid
        var dist_rc = math.vec2.length(c0)
        #log("root ", id, " dist_centr ", dist_rc)

        # Distances from neighbors to centroid 
        var dist_centroid = {}
        dist_centroid = neighbors.map(function(rid, data) {
				  # Neighbour position in frame F0
          var p0 = math.vec2.newp(data.distance, data.azimuth)
          # Difference vector between p0 and c0
          var v = math.vec2.sub(p0, c0)
          # Distance between robot and centroid
          return math.vec2.length(v)
        })

        # Minimal distance to centroid
        var temp = {}
        temp = dist_centroid.reduce(function(rid, data, accum) {
          accum.min = math.min(accum.min, data)
          return accum
        }, {.min = dist_rc})
        var min = temp.min
        #log("min ", min)

        # If the current root is the closest to the centroid
        if(dist_rc == min) {
          ack = 1
        }
        # Otherwise
        else {
			    var flag = 0
          # Selects closest robot to centroid as better root
          foreach(dist_centroid.data, function(key, value) {
            if(flag == 0) {
              if(value == min) {
                better_root = key
                flag = 1
              }
              # break (when implemented)
            }
          })

          
          #log(id, " better root : ", better_root)
          #log("X : ", accum_x, "Y : ", accum_y)
          var angle = neighbors.data[better_root].azimuth
          # Broadcasts
          var message = {
            .better_root = better_root,
            .centroid_x = accum_x,
            .centroid_y = accum_y,
            .angle_old_root = angle
          }
          neighbors.broadcast("msg1", message)
        }
      }
    }

    else if (better_root == nil) {
      # Listen for old root message
      foreach(knowledge.better_root, function(rid, value) {
          if(value == id) {
            
            var theta = neighbors.data[rid].azimuth

            var p1 = {
              .x = knowledge.centroid_x[rid],
              .y = knowledge.centroid_y[rid]
            }

            var p01 = {
              .x = neighbors.data[rid].distance * math.cos(theta),
              .y = neighbors.data[rid].distance * math.sin(theta)
            }

            var p0 = {}

            if (knowledge.angle_old_root[rid] != nil) {
              var rot_angle = radians_interval(knowledge.angle_old_root[rid] - theta - math.pi)
              p0 = change_frame(p01, p1, rot_angle)
            }
            else {
              p0.x = p1.x
              p0.y = p1.y
            }

            accum_x = p0.x
            accum_y = p0.y

            centroid_x = accum_x
            centroid_y = accum_y

            root = 1
            neighbors.broadcast("got_it", id)
          }
       })
    }

    # Transitions to new state when all robots are ready
    if (trigger.get("a") == 1) {
		  barrier_set(ROBOTS, end_fun)
      barrier_ready()
    }

}

####################################

function tree_select() {

    log(id, " tree_select")

    neighbors.map(function(rid, data) {
      knowledge.distance[rid] = data.distance
      knowledge.azimuth[rid] = data.azimuth
      return 1
    })

    if (level == 0) {
      # Finds a parent
      parent_select()
      # Updates robot level
      if (tree.parent.id != nil) {
        #log(" ")
        #log("selected")
        #log("son ", id)
        #log("parent ", tree.parent.id)
        #log(" ")

        level = knowledge.level[tree.parent.id] + 1 
        angle_parent = neighbors.data[tree.parent.id].azimuth
      }
    }
    else {
      # Updates list of sons
      foreach(knowledge.parent, function(key, value) {
        if(value == id) {
          #log(value)
          if(tree.sons[key] == nil) {
            # Updates robot nb_sons
            nb_sons = nb_sons + 1 
            # Updates robot free status
            if (nb_sons >= N_SONS) {
              free = 0
            }
          }
          tree.sons[key] = {
              .distance = knowledge.distance[key],
              .azimuth = knowledge.azimuth[key],
              .eq_level = knowledge.eq_level[key],
              .accum_x = knowledge.accum_x[key],
              .accum_y = knowledge.accum_y[key],
              .angle_parent = knowledge.angle_parent[key]
          }  
        }
      })
    }
    
    # Creates a subset of neighbors to get the sons
    # and parent
    sons = {}
    sons = neighbors.filter(function(rid, data) {
      return (tree.sons[rid] != nil)
    })
    parent = {}
    parent = neighbors.filter(function(rid, data) {
      return (tree.parent.id == rid)
    })

    # Counts robots (updates eq_level)
    count()

    # Updates count of robots in the tree
    if (root == 1) {
      robot_count = eq_level
    }

    nb_descendants = eq_level - level


    # Computes centroid (updates accum_x, accum_y)
    centroid()

    # Broadcast information to other robots
    var message = {
      .level = level,
      .parent = tree.parent.id,
      .eq_level = eq_level,
      .free = free,
      .accum_x = accum_x,
      .accum_y = accum_y,
      .angle_parent = angle_parent
    }
    neighbors.broadcast("msg2", message)
    
    # Triggers transition to new state if root count = ROBOTS
    if (root == 1) {
      if(robot_count == ROBOTS) {
        log("centroid X: ", accum_x, " Y ", accum_y)
        trigger.put("b", 1)
      }
    }

    # Transitions to new state when all robots are ready
    if (trigger.get("b") == 1) {
      barrier_set(ROBOTS, root_select)
      barrier_ready()
     }
}

################################################################
################################################################


# This function is executed every time you press the 'execute' button
function init() {

    trigger = stigmergy.create(TRIGGER_VSTIG)
    barrier = stigmergy.create(BARRIER_VSTIG)

    # Trees
    old_tree = {}
    tree = old_tree
    old_tree.parent = {}
    old_tree.sons = {}

    # Boolean variables
    root = 0 # Root status
    free = 1 # Node status (true if accepts sons)
    ack = 0    

    # Tree variables
    level = 0
    eq_level = 0
    nb_sons = 0
    nb_descendants = 0
    
    # Update root status
    if (id == 0) {
      root = 1 # True if robot is the ro ot
      level = 1
      robot_count = 0
    }
    
    statef = tree_select

    knowledge = {
      .level = {},
      .parent = {},
      .eq_level = {},
      .free = {},
      .accum_x = {},
      .accum_y = {},
      .angle_parent = {},
			.distance = {},	
      .azimuth = {},
      .better_root = {},
      .centroid_x = {},
      .centroid_y = {},
      .angle_old_root = {},
      .ackn = {}
    }

    # Broadcast information to other robots
    var message = {
      .level = level,
      .parent = old_tree.parent.id,
      .eq_level = eq_level,
      .free = free,
      .accum_x = accum_x,
      .accum_y = accum_y,
      .angle_parent = 0.0
    }
    neighbors.broadcast("msg2", message)

    # Listen to information from other robots for root_select
    neighbors.listen("msg1", function(vid, value, rid) {
      knowledge.better_root[rid] = value.better_root
      knowledge.centroid_x[rid] = value.centroid_x
      knowledge.centroid_y[rid] = value.centroid_y
      knowledge.angle_old_root[rid] = value.angle_old_root
      knowledge.angle_parent[rid] = value.angle_parent
    })

    # Listen to information from other robots for tree_select
    neighbors.listen("msg2", function(vid, value, rid) {
      knowledge.level[rid] = value.level
      knowledge.parent[rid] = value.parent
      knowledge.eq_level[rid] = value.eq_level
      knowledge.free[rid] = value.free
      knowledge.accum_x[rid] = value.accum_x
      knowledge.accum_y[rid] = value.accum_y
      knowledge.angle_parent[rid] = value.angle_parent
    })

    neighbors.listen("got_it", function(vid, value, rid) {
      knowledge.ackn[rid] = value
    })

}

############################################

# This function is executed at each time step
function step() {
   
    tree_config()
    goal = math.vec2.new(0.0, 0.0)

}
############################################

# This function is executed every time you press the 'reset'
# button in the GUI. 
function reset() {
     # put your code here
}

############################################

# This function is executed only once, when the robot is removed
# from the simulation
function destroy() {
     # put your code here
}

################