function [action_sequence, best_sample_traces, tree, status] = mdp_is_planner(start_world_state, id, time_horizon, P)
% Compute an optimal suquence of actions (policies), each applied for P.action_duration sec.
% also return the sample traces corresponding to this chosen sequence best_sample_traces,
% status containts the average total rewards, tree-search information, etc...
% An exhaustive heuristic best-first tree search over all possible sequences is performed.

% start node
node_id = 1;

if P.importance_sampling
  actions_samples = generate_actions_samples(start_world_state, P);
  number_samples = size(actions_samples, 1);
  for j=1:number_samples
    samples(j) = init_sample(start_world_state, 0, [], [], false);
    actions_sample = actions_samples(j,:);
    samples(j).actions = [actions_sample.action];
    samples(j).likelihood = prod([actions_sample.likelihood]);
  end
else
  % init root with M equivalent samples
  sample = init_sample(start_world_state, 0, [], [], false);
  samples(1:P.number_samples) = sample;  
end
  
disp(['[I] mdp_is_planner: using ' num2str(number_samples) ' samples']);

root = init_node(node_id, samples, [], [], 0);

best_node_id = 0;
max_ave_total_reward = -inf;
pruned_nodes = 0;

% tree stored as a map
tree = containers.Map('KeyType', 'uint32', 'ValueType', 'any');

% list of open (unexpanded) nodes, as a max-priority queue
open_list = PriorityQueue;
open_list.insert(node_priority(root), root);

start_time = start_world_state.car_states(1).t;

violation_count = 0;

tic;

% add and remove new nodes from open list according to priority
while ~open_list.isEmpty()
  % pop new node from end of open list
  [priority, node] = open_list.pop();
  
  if ~node.terminal
    % a cell array, where each cell is indexed by action and contains a list of samples
    action_samples = rollout_samples(node.samples, id, start_time + time_horizon, P);
    
    % iterate over actions
    for i=1:length(action_samples)
      samples = action_samples{i};
      % no valid samples for this action
      if isempty(samples)
        continue
      end
      
      % currently we do not penalize violations -- we simply discard the action      
      if ~isempty([samples.violations])
        violation_count = violation_count + length([samples.violations]);
        continue
      end
      
      parent_samples = [samples.parent_sample];
      
      % pruning is critical
      ave_total_reward = mean([parent_samples.total_reward] + [samples.reward]);
      if ave_total_reward < max_ave_total_reward
        pruned_nodes = pruned_nodes + 1;
        continue
      end
      
      terminal = sum([samples.terminal]);  % OR-terminal nodes
      % all samples should have the same action, so we just grab it from the first one
      action = samples(1).trace(end).car_controls(id).action;
      child = init_node(node_id + 1, samples, node, action, terminal);
      node_id = node_id + 1;
    
      open_list.insert(node_priority(child), child);
    end
    
  else 
    % this is a terminal node
    if node.ave_total_reward > max_ave_total_reward
      best_node_id = node.id;
      max_ave_total_reward = node.ave_total_reward;
    end
    
  end
  % add processed node to tree
  tree(node.id) = node;
end


action_sequence = [];
best_sample_traces = cell(0);

status.max_ave_total_reward = max_ave_total_reward;
status.tree_size = tree.length();
status.pruned_nodes = pruned_nodes;
status.best_node = [];
status.violation_count = violation_count;
status.cpu_time = toc;

% if no solution found
if ~best_node_id
  return
end

best_node = tree(best_node_id)
for j=1:length(best_node.samples)
  best_sample_traces{j} = backup_sample_trace(best_node.samples(j));
end

status.best_node = best_node;

% reconstruct the optimal sequence from the best leaf node up the tree
while best_node_id ~= 0
  best_node = tree(best_node_id);
  action_sequence = [best_node.action, action_sequence];
  best_node_id = best_node.parent_id;
end


function trace = backup_sample_trace(sample)
% backtrace a complete sample trace
trace = [];
while ~isempty(sample) && ~isempty(sample.trace)
  sample.trace
  trace
  trace = [sample.trace, trace];
  sample = sample.parent_sample;
end


function action_samples = rollout_samples(parent_samples, id, final_time, P)
% generate samples for given action policies

action_samples = cell(length(P.action_ids));

for j=1:length(parent_samples)  
  parent_sample = parent_samples(j);
  world_state = parent_sample.trace(end);
  
  if P.importance_sampling
    for k=2:length(P.cars)
      action_policies{k} = @(world_state, P)fixed_action_policy(world_state, parent_sample.actions(k-1), P);
    end
  else
    action_policies = {P.cars.action_policy};
  end
  
  if ~parent_sample.terminal
    % expand children by executing feasible actions
    actions = P.cars(id).feasible_actions(world_state, id, P.cars(id).goal_state, P);
    
    for i=1:length(actions)
      % set this car action policy to the fixed action
      action_policies{id} = @(world_state, P)fixed_action_policy(world_state, actions(i), P);
          
      [action_trace, violations, terminal] = rollout(world_state, action_policies, P.action_duration, P);
      
      r = 0;
      for k=1:length(action_trace)
        r = r + reward(action_trace(k), id, P);
      end
      
      if ~isempty(action_trace)
        if ~terminal && action_trace(end).car_states(1).t > final_time - eps
          terminal = true;
        end
      end
      
      % record all samples including those with violations
      sample = init_sample(action_trace, r, parent_sample, violations, terminal);
      sample.actions = parent_sample.actions;
      sample.likelihood = parent_sample.likelihood;      
      
      % store a list of samples indexed by action type
      action_samples{actions(i).id}(end+1) = sample;
      
    end % iterate over actions
  end
end



function node = init_node(id, samples, parent, action, terminal)

node.id = id;

node.samples = samples;

if ~isempty(parent)
  node.parent_id = parent.id;
  parent_samples = [samples.parent_sample];
  node.ave_total_reward = mean([parent_samples.total_reward] + [samples.reward]);
  node.depth = parent.depth + 1;
else
  node.parent_id = 0;
  node.ave_total_reward = 0;
  node.depth = 0;
end

% which action/action was executed
node.action = action;

% is this a terminal node
node.terminal = terminal;



function sample = init_sample(trace, reward, parent_sample, violations, terminal)

% sequence of high-frequency states from parent to node
sample.trace = trace;

 % reward r(s,a,s') where s' corresponds to this node, i.e. reward for applied action from parent
sample.reward = reward;

sample.parent_sample = parent_sample;

if ~isempty(parent_sample)
  % total reward = "reward to come"
  sample.total_reward = parent_sample.total_reward + reward;
else
  sample.total_reward = 0;
end

% is this a terminal sample
sample.violations = violations;

% is this a terminal sample
sample.terminal = terminal;

% which actions all entities took
sample.actions = [];

sample.likelihood = 1;


function q = node_priority(node)

% heurist function determining the order of node processing
q = node.depth*100000 + node.ave_total_reward;
