Behavior Designer - ProbabilitySelector

using System.Collections.Generic;
using System.Linq;
using UnityEngine;
using BehaviorDesigner.Runtime.Tasks;

[TaskDescription("Runs its children in a random order, weighted by probability. Returns success as soon as a child returns success.")]
[TaskIcon("{SkinColor}RandomSelectorIcon.png")]
public class ProbabilitySelector : Composite {
    [BehaviorDesigner.Runtime.Tasks.Tooltip("The probability of each child task being selected. Values are relative to each other, e.g: 0.5/20 have twice the chance of 0.25/10 of occuring.")]
    public float[] ChildTaskProbabilities;

    /// <summary>Tasks that have not been executed.</summary>
    readonly List<TaskProbability> _availableTasks = new List<TaskProbability>();
    /// <summary>Executed tasks to be reused later.</summary>
    readonly List<TaskProbability> _executedTasks = new List<TaskProbability>();

    int _currentAvailableTasksIndex = 0;
    int _currentChildIndex = 0;
    TaskStatus _lastChildExecutionStatus = TaskStatus.Inactive;

    public override void OnAwake() {
        if (ChildTaskProbabilities.Length != Children.Count) { Debug.LogErrorFormat("Incorrect amount of ChildTaskProbabilities in ProbabilitySelector. (ChildTaskProbabilities: {0}, Child Tasks: {1})", ChildTaskProbabilities.Length, Children.Count); }

        for (int i = 0; i < ChildTaskProbabilities.Length && i < Children.Count; i++) { _availableTasks.Add(new TaskProbability(i)); }
    }

    public override void OnStart() {
        // Transfer any previously completed tasks to _availableTasks
        _availableTasks.AddRange(_executedTasks);
        _executedTasks.Clear();

        UpdateTaskProbabilities();

        UpdateCurrentTaskIndex();
    }

    void UpdateTaskProbabilities() {
        float total = 0;
        float currentValue = 0;

        foreach (TaskProbability task in _availableTasks) { total += ChildTaskProbabilities[task.Index]; }
        foreach (TaskProbability task in _availableTasks) {
            currentValue += ChildTaskProbabilities[task.Index] / total;
            task.ProbabilityValue = currentValue;
        }
    }

    void UpdateCurrentTaskIndex() {
        float value = Random.value;

        for (int i = 0; i < _availableTasks.Count; i++) {
            TaskProbability task = _availableTasks[i];

            // If this isn't the last available task, test to see if the value is in its probability range
            if (i < _availableTasks.Count - 1 && value > task.ProbabilityValue) { continue; }

            _currentAvailableTasksIndex = i;
            _currentChildIndex = task.Index;

            break;
        }
    }

    public override int CurrentChildIndex() {
        return _currentChildIndex;
    }

    public override bool CanExecute() {
        // Continue execution if there are still available tasks and no task has returned success
        return _availableTasks.Count > 0 && _lastChildExecutionStatus != TaskStatus.Success;
    }

    public override void OnChildExecuted(TaskStatus childStatus) {
        if (_availableTasks.Count > 0) {
            // Transfer executed task to _executedTasks
            _executedTasks.Add(_availableTasks[_currentAvailableTasksIndex]);
            _availableTasks.RemoveAt(_currentAvailableTasksIndex);

            // If there are still available tasks, select the next task
            if (_availableTasks.Count > 0) {
                UpdateTaskProbabilities();

                UpdateCurrentTaskIndex();
            }
        }

        _lastChildExecutionStatus = childStatus;
    }

    public override void OnConditionalAbort(int childIndex) {
        _lastChildExecutionStatus = TaskStatus.Inactive;
    }

    public override void OnEnd() {
        _lastChildExecutionStatus = TaskStatus.Inactive;
    }

    class TaskProbability {
        /// <summary>The index of this task in <see cref="ProbabilitySelector.ChildTaskProbabilities"/></summary>
        public readonly int Index;
        /// <summary>The maximum normalised value at which this task will be executed. (Provided the random value is also greater than the <see cref="ProbabilityValue"/>s of any preceding tasks in <see cref="_availableTasks"/>.)</summary>
        public float ProbabilityValue;

        public TaskProbability(int index) {
            Index = index;
        }
    }
}