WorldState.cs

using System;
using System.Linq;
using System.Collections.Generic;
using System.Diagnostics;

namespace mapf;

/// <summary>
/// Describes a node in the A* search space.
/// </summary>
public class WorldState : IComparable<IBinaryHeapItem>, IBinaryHeapItem, IHeuristicSearchNode
{
    public int makespan; // Total time steps passed, max(agent makespans)
    public int g { get; set; } // Value depends on Constants.costFunction and Constants.sumOfCostsVariant, Sum of agent makespans until they reach their goal
    public int h { get; set; }
    public int hBonus { get; set; }
    public AgentState[] allAgentsState;
    public WorldState prevStep;
    private int binaryHeapIndex;
    public MDDNode mddNode;
    public int generated;

    public int primaryTieBreaker;
    public int secondaryTieBreaker;
    /// <summary>
    /// Maps from agent num to the number of times the path up to this node collides with that agent
    /// </summary>
    public Dictionary<int, int> conflictCounts;
    /// <summary>
    /// Maps from agent num to a list of the conflict times with it
    /// </summary>
    public Dictionary<int, List<int>> conflictTimes;
    /// <summary>
    /// The min depth (makespan) from which a node may be considered a goal.
    /// TODO: Consider moving out of the node object to a static variable or something.
    ///       It doesn't change between nodes.
    /// </summary>
    public int minGoalTimeStep;
    /// <summary>
    /// The min cost (g) from which a node may be considered a goal.
    /// TODO: Consider moving out of the node object to a static variable or something.
    ///       It doesn't change between nodes.
    /// </summary>
    public int minGoalCost;
    /// <summary>
    /// The last move of all agents that have already moved in this turn.
    /// Used for making sure the next agent move doesn't collide with moves already made.
    /// Used while generating this node, nullified when done.
    /// </summary>
    public Dictionary<TimedMove, int> currentMoves;
    protected static readonly int NOT_SET = -1;
    /// <summary>
    /// For computing expansion delay
    /// </summary>
    public int expandedCountWhenGenerated;
    ///// <summary>
    ///// For lazy heuristics
    ///// </summary>
    //public CBS cbsState;
    /// <summary>
    /// For MStar.
    /// Disjoint sets of agent indices, since only internal agents are considered.
    /// </summary>
    public DisjointSets<int> collisionSets;
    //public ISet<int> currentCollisionSet;
    public ISet<WorldState> backPropagationSet;
    public TimedMove[] plannedMoves;

    /// <summary>
    /// Create a state with the given state for every agent.
    /// </summary>
    /// <param name="allAgentsState"></param>
    /// <param name="minDepth"></param>
    /// <param name="minCost"></param>
    /// <param name="mddNode"></param>
    public WorldState(AgentState[] allAgentsState, int minDepth = -1, int minCost = -1, MDDNode mddNode = null)
    {
        this.allAgentsState = allAgentsState.ToArray();
        this.makespan = allAgentsState.Max(state => state.lastMove.time); // We expect to only find at most two G values within the agent group
        this.CalculateG(); // G not necessarily zero when solving a partially solved problem.
        this.primaryTieBreaker = 0;
        this.secondaryTieBreaker = 0;
        this.conflictCounts = new Dictionary<int, int>();  // Unused if not running under CBS, and we can't tell at this point easily
        this.conflictTimes = new Dictionary<int, List<int>>();  // Unused if not running under CBS, and we can't tell at this point easily
        this.minGoalTimeStep = minDepth;
        this.minGoalCost = minCost;
        if (mddNode == null)
            this.currentMoves = new Dictionary<TimedMove, int>();
        this.goalCost = NOT_SET;
        this.goalSingleCosts = null;
        this.singlePlans = null;
        this.hBonus = 0;
        this.mddNode = mddNode;
    }

    /// <summary>
    /// Copy constructor.
    /// </summary>
    /// <param name="cpy"></param>
    public WorldState(WorldState cpy)
    {
        this.makespan = cpy.makespan;
        this.g = cpy.g;
        this.h = cpy.h;
        // The conflictTimes, conflictCounts and sumConflictCounts are only copied later if necessary.
        this.minGoalTimeStep = cpy.minGoalTimeStep;
        this.minGoalCost = cpy.minGoalCost;
        this.allAgentsState = new AgentState[cpy.allAgentsState.Length];
        for (int i = 0; i < allAgentsState.Length; i++)
        {
            this.allAgentsState[i] = new AgentState(cpy.allAgentsState[i]);
            // Shallow copy - it's still the same lastMove inside the AgentState, until we set a new lastMove there.
        }
        if (cpy.currentMoves != null)
            // cpy is an intermediate node
            this.currentMoves = new Dictionary<TimedMove, int>(dictionary: cpy.currentMoves);
        else
            // cpy is a concrete node
            this.currentMoves = new Dictionary<TimedMove, int>(capacity: cpy.allAgentsState.Length);
        this.goalCost = NOT_SET;
        this.goalSingleCosts = null;
        this.singlePlans = null;
        this.hBonus = 0;
        this.mddNode = cpy.mddNode;
        this.prevStep = cpy;
    }

    /// <summary>
    /// Creates a new state by extracting a subset of the agents from
    /// the original WorldState. We overload the constructor because
    /// while building our pattern database, we rewrite the problem and
    /// therefore need to make a deep copy of the state data structures so
    /// as to not overwrite the original problem. The ultimate solution
    /// would be to rework the code to remove static variables so that we
    /// can instantiate subproblems without affecting the original data
    /// structures.
    /// </summary>
    /// <param name="allAgentsState">A set of agent states in the original problem.</param>
    /// <param name="agentIndicesToCopy">A list of indices referring to the subset of agents we want to extract.</param>
    public WorldState(AgentState[] allAgentsState, List<uint> agentIndicesToCopy)
        // Copy specified agents only
        : this(agentIndicesToCopy.Select(index => new AgentState(allAgentsState[index])).ToArray())
    {}
        
    public bool GoalTest()
    {
        // Check if this is a generalised goal node and its plan is long enough.
        // If we know the optimal solution, it doesn't matter if this is a real goal node or not, we can finish.
        if (this.singlePlans != null)
        {
            // Check if plans are long enough and costly enough
            if (this.singlePlans.All(plan => plan.GetSize() - 1 >= this.minGoalTimeStep))
            {
                if (Constants.costFunction == Constants.CostFunction.SUM_OF_COSTS)
                {
	                if (this.singlePlans.Sum(plan => plan.GetCost()) >= this.minGoalCost)
                        return true;
                }
                else if (Constants.costFunction == Constants.CostFunction.MAKESPAN || Constants.costFunction == Constants.CostFunction.MAKESPAN_THEN_SUM_OF_COSTS)
                {
                    if (this.singlePlans.Max(plan => plan.GetCost()) >= this.minGoalCost)
    	                return true;
                }
                else
                    throw new Exception("Unsupported cost function");
            }
        }

        if (this.g < this.minGoalCost)
            return false;

        if (this.makespan < this.minGoalTimeStep)
            return false;

        return this.h == 0; // This assumes the heuristic is consistent,
                            // or at least has the property of consistent heuristics that only the goal has h==0.
                            // SIC really is a consistent heuristic, so this is fine for now.
                            // TODO: Implement a proper goal test and use it when h==0.
    }

    protected SinglePlan[] singlePlans;

    /// <summary>
    /// Set the optimal solution of this node as a problem instance.
    /// Currently only used by CbsHeuristicForAStar, if a solution was found while running the heuristic.
    /// </summary>
    /// <param name="solution"></param>
    public virtual void SetSolution(SinglePlan[] solution)
    {
        this.singlePlans = SinglePlan.GetSinglePlans(this); // This node may be a partial solution itself, need to start from the real root.
        for (int i = 0; i < solution.Length; ++i)
            this.singlePlans[i].ContinueWith(solution[i]);
    }

    public SinglePlan[] GetSinglePlans()
    {
        if (this.singlePlans != null)
            return this.singlePlans;
        else
            return SinglePlan.GetSinglePlans(this);
    }

    /// <summary>
    /// Returns the optimal plan to the goal through this node, if this is a goal node (of any kind),
    /// else returns the optimal plan to this node.
    /// </summary>
    /// <returns></returns>
    public Plan GetPlan()
    {
        if (this.singlePlans != null)
            return new Plan(this.singlePlans);
        else
            return new Plan(this);
    }

    /// <summary>
    /// For generalized goal nodes.
    /// TODO: Get rid of this and just return the sum/max of the single costs where needed?
    /// </summary>
    protected int goalCost;

    /// <summary>
    /// Returns the optimal cost to the goal from the start through this node.
    /// </summary>
    /// <returns></returns>
    public int GetGoalCost()
    {
        Trace.Assert(this.GoalTest(), "Only call for goal nodes!");

        if (goalCost == NOT_SET) // This is just a proper goal
        {
            if (Constants.costFunction == Constants.CostFunction.SUM_OF_COSTS)
            {
                return this.g;
            }
            else if (Constants.costFunction == Constants.CostFunction.MAKESPAN ||
                Constants.costFunction == Constants.CostFunction.MAKESPAN_THEN_SUM_OF_COSTS)
            {
                return this.makespan;
            }
            return 0; // To quiet the compiler
        }
        else                     // This is a generalised goal node - it stores the optimal path to the goal through it
            return this.goalCost;
    }

    /// <summary>
    /// Set the optimal cost from the start to the goal through this node.
    /// Makes this a generalized goal node.
    /// Currently only used by CbsHeuristicForAStar.
    /// </summary>
    /// <param name="cost"></param>
    public void SetGoalCost(int cost)
    {
        this.goalCost = cost;
    }

    /// <summary>
    /// For generalized goal nodes
    /// </summary>
    protected int[] goalSingleCosts;

    public int[] GetSingleCosts()
    {
        Trace.Assert(this.GoalTest(), "Only call for goal nodes!");

        if (goalSingleCosts == null) // This is just a proper goal
            return allAgentsState.Select(agent => agent.g).ToArray();
        else
            return this.goalSingleCosts;
    }

    /// <summary>
    /// Set the optimal cost from the start to the goal through this node for every agent.
    /// Makes this node a generalized goal node.
    /// Currently only used by CbsHeuristicForAStar.
    /// </summary>
    /// <param name="costs"></param>
    public void SetSingleCosts(int[] costs)
    {
        this.goalSingleCosts = costs;
    }

    /// <summary>
    /// Used when WorldState objects are put in the open list priority queue
    /// </summary>
    /// <param name="other"></param>
    /// <returns></returns>
    public virtual int CompareTo(IBinaryHeapItem other)
    {
        WorldState that = (WorldState)other;
        int thisF = this.f;
        int thatF = that.f;
        if (thisF < thatF)
            return -1;
        if (thisF > thatF)
            return 1;

        return this.TieBreak(that);
    }

    public int TieBreak(WorldState that)
    {
        bool thisIsGoal = this.GoalTest();
        bool thatIsGoal = that.GoalTest();
        if (thisIsGoal == true && thatIsGoal == false) // The elaborate form is necessary to keep the comparison consistent. Otherwise goalA<goalB and goalB<goalA
            return -1;
        if (thatIsGoal == true && thisIsGoal == false)
            return 1;


        // Prefer nodes that contain conflicts with fewer agents - when a conflict is resolved,
        // many times other conflicts with the same agent are resolved automatically thanks to conflict avoidance,
        // especially if the cost increases.
        // For ID, this may instead prefers nodes which conflict with smaller groups.
        // TODO: Ideally, prefer nodes where the minimum vertex cover of the conflict graph is smaller.
        //       Compute an MVC of the conflict graph without the node's agents in the CBS node
        //       before running the low level, and only compare the number of agents the node
        //       conflicts with that aren't in the MVC.
        //       Maybe even compute the MVC of the cardinal conflict graph and of the all conflict
        //       graph separately and tie-break first according to the number of agents we conflict
        //       with that aren't in the MVC of the conflict graph and then the number of agents
        //       we conflict with that aren't in the cardinal conflict graph
        if (this.primaryTieBreaker < that.primaryTieBreaker)
            return -1;
        if (this.primaryTieBreaker > that.primaryTieBreaker)
            return 1;

        // Prefer nodes with fewer conflicts - the probability that some of them are cardinal is lower
        if (this.secondaryTieBreaker < that.secondaryTieBreaker)
            return -1;
        if (this.secondaryTieBreaker > that.secondaryTieBreaker)
            return 1;

        // //M-Star: prefer nodes with smaller collision sets:
        //if (this.collisionSets != null) // than M-Star is running
        //{
        //    // The collision sets change during collision set backpropagation and closed list hits.
        //    // Backpropagation goes from a node's child to the node, so it's tempting to think
        //    // it only happens when the node is already expanded and out of the open list,
        //    // but partial expansion makes that false. 
        //    // Closed list hits can also happen while the node is waiting to be expanded.
        //    // So the max rank can change while the node is in the open list - 
        //    // it can't be used for tie breaking :(.
        //    if (this.collisionSets.maxRank < that.collisionSets.maxRank)
        //        return -1;
        //    if (that.collisionSets.maxRank > this.collisionSets.maxRank)
        //        return 1;
        //}

        // f, collision sets, conflicts and internal conflicts being equal, prefer nodes with a larger g
        // - they're closer to the goal so less nodes would probably be generated by them on the way to it.
        if (this.g < that.g)
            return 1;
        if (this.g > that.g)
            return -1;

        return 0;
    }

    /// <summary>
    /// Calculate and set the g of the state as the sum of the different agent g values.
    /// </summary>
    public virtual void CalculateG()
    {
        if (Constants.costFunction == Constants.CostFunction.SUM_OF_COSTS)
        {
            g = allAgentsState.Sum<AgentState>(agent => agent.g);
        }
        else if (Constants.costFunction == Constants.CostFunction.MAKESPAN ||
            Constants.costFunction == Constants.CostFunction.MAKESPAN_THEN_SUM_OF_COSTS)
        {
            g = makespan;  // Let's hope makespan var is correct
        }
        else
        {
            throw new Exception($"Unsupported cost function {Constants.costFunction}");
        }
    }

    /// <summary>
    /// Prepare for re-insertion into the open list
    /// </summary>
    public virtual void Clear() { }

    public virtual int f
    {
        get
        {
            return this.g + this.h;
        }
    }

    public int GetTargetH(int f) => f - g;

    public override string ToString()
    {
        var builder = new System.Text.StringBuilder($"{generated} f:{f} makespan:{makespan} h:{h} g:{g} ");
        foreach (AgentState temp in allAgentsState)
        {
            builder.Append("|");
            builder.Append(temp.lastMove);
        }
        builder.Append("|");
        return builder.ToString();
    }

    /// <summary>
    /// Returns the last move of all the agents in this state.
    /// </summary>
    /// <returns>A list of Moves</returns>
    public List<Move> GetAgentsMoves()
    {
        return this.allAgentsState.Select<AgentState, Move>(state => state.lastMove).ToList<Move>();
    }

    /// <summary>
    /// Returns the last move of the requested agent.
    /// </summary>
    /// <param name="index"></param>
    /// <returns></returns>
    public Move GetSingleAgentMove(int index)
    {
        return allAgentsState[index].lastMove;
    }

    /// <summary>
    /// BH_Item implementation
    /// </summary>
    /// <returns></returns>
    public int GetIndexInHeap() { return binaryHeapIndex; }

    /// <summary>
    /// BH_Item implementation
    /// </summary>
    /// <returns></returns>
    public void SetIndexInHeap(int index) { binaryHeapIndex = index; }

    /// <summary>
    /// Checks for internal conflicts
    /// </summary>
    /// <returns></returns>
    public bool isValid()
    {
        for (int i = 0; i < this.allAgentsState.Length; i++)
        {
            for (int j = i+1; j < this.allAgentsState.Length; j++)
            {
                // Internal conflict
                if (this.allAgentsState[i].lastMove.IsColliding(this.allAgentsState[j].lastMove))
                    return false;
            }
        }
        return true;
    }

    /// <summary>
    /// Only the agent states are used in the hash.
    /// The g, makespan, h, potentialConflictsCount, sumConflictCounts and others are ignored, as neccesary.
    /// </summary>
    /// <returns></returns>
    public override int GetHashCode()
    {
        int ans = 0;
        unchecked
        {
            for (int i = 0 ; i < allAgentsState.Length; i++)
            {
                ans += allAgentsState[i].GetHashCode() * Constants.PRIMES_FOR_HASHING[i % Constants.PRIMES_FOR_HASHING.Length];
            }
        }
        return ans;
    }

    /// <summary>
    /// Only the AgentStates are compared.
    /// g, makespan, h, potentialConflictsCount, sumConflictCounts and others are ignored, as necessary.
    /// </summary>
    /// <param name="obj"></param>
    /// <returns></returns>
    public override bool Equals(object obj)
    {
        if (obj == null)
            return false;
        WorldState that = (WorldState)obj;
        return this.allAgentsState.SequenceEqual(that.allAgentsState);
    }

    /// <summary>
    /// Counts the number of times this node collides with each agent move in the conflict avoidance table.
    /// Also compute the representative primary and secondary tie-breaking values according to the CAT's avoidance goal
    /// </summary>
    /// <param name="CAT"></param>
    /// <returns></returns>
    public virtual void IncrementConflictCounts(ConflictAvoidanceTable CAT)
    {
        for (int i = 0; i < this.allAgentsState.Length; i++)
        {
            this.allAgentsState[i].lastMove.IncrementConflictCounts(CAT, this.conflictCounts, this.conflictTimes);
        }

        if (CAT.avoidanceGoal == ConflictAvoidanceTable.AvoidanceGoal.MINIMIZE_CONFLICTS)  // For ID, the original rule
            this.primaryTieBreaker = this.conflictCounts.Sum(pair => pair.Value);
        else if (CAT.avoidanceGoal == ConflictAvoidanceTable.AvoidanceGoal.MINIMIZE_CONFLICTING_GROUPS)
            this.primaryTieBreaker = this.conflictCounts.Keys.Count;
        else if (CAT.avoidanceGoal == ConflictAvoidanceTable.AvoidanceGoal.MINIMIZE_CONFLICTING_GROUPS_THEN_CONFLICTS)
            // For CBS, minimizes the number of conflicting groups and then the number of conflicts with them
        {
            this.primaryTieBreaker = this.conflictCounts.Keys.Count;
            this.secondaryTieBreaker = this.conflictCounts.Sum(pair => pair.Value);
        }
        else if (CAT.avoidanceGoal == ConflictAvoidanceTable.AvoidanceGoal.MINIMIZE_LARGEST_CONFLICTING_GROUP_THEN_NUMBER_OF_SUCH_GROUPS)
            // For ID, minimizes the size of the largest group we conflict with and then 
            // the number of conflicting groups with that size. The idea was to minimize conflicts that matter, and conflicts with
            // non-max-size groups don't.
            // Kept mostly for reference.
        {
            if (this.conflictCounts.Count != 0)
            {
                this.primaryTieBreaker = this.conflictCounts.Max(pair => CAT.agentSizes[pair.Key]);
                this.secondaryTieBreaker = this.conflictCounts.Where(pair => CAT.agentSizes[pair.Key] == this.primaryTieBreaker).Count();
            }
            else
            {
                this.primaryTieBreaker = 0;
                this.secondaryTieBreaker = 0;
            }
        }
        else if (CAT.avoidanceGoal == ConflictAvoidanceTable.AvoidanceGoal.MINIMIZE_LARGEST_CONFLICTING_GROUP_THEN_MAXIMIZE_CONFLICT_COUNTS_WITH_OTHERS)
        // For ID, minimizes the size of the largest group we conflict with and then 
        // maximizes the number of conflicts the two groups have with other groups
        {
            if (this.conflictCounts.Count != 0)
            {
                this.primaryTieBreaker = this.conflictCounts.Max(pair => CAT.agentSizes[pair.Key]);
                this.secondaryTieBreaker = -(this.conflictCounts.Sum(pair => pair.Value) +
                    this.conflictCounts.Where(pair => CAT.agentSizes[pair.Key] == this.primaryTieBreaker).Max(pair => CAT.agentConflictCounts[pair.Key]));
            }
            else
            {
                this.primaryTieBreaker = 0;
                this.secondaryTieBreaker = 0;
            }
        }
        else if (CAT.avoidanceGoal == ConflictAvoidanceTable.AvoidanceGoal.MINIMIZE_CONFLICTING_GROUP_SIZE_AND_COUNT)
        {
            this.primaryTieBreaker = this.conflictCounts.Sum(pair => 1 << (CAT.agentSizes[pair.Key] - 1));
        }
    }

    /// <summary>
    /// Currently only used by CbsHeuristicForAStar, where each A* node is converted to a new
    /// MAPF problem to be solved by CBS (or ICTS, theoretically)
    /// </summary>
    /// <param name="initial"></param>
    /// <returns></returns>
    public virtual (ProblemInstance, ISet<CbsConstraint>) ToProblemInstance(ProblemInstance initial)
    {
        // Notice this is not a subproblem in the number of agents but
        // in the steps from the start.
        // It might even be harder if the steps were away from the goal.
        return (initial.Subproblem(this.allAgentsState), new HashSet<CbsConstraint>());
    }

    //public WorldState GetPlanStart(int agentIndex)
    //{
    //    WorldState node = this;
    //    while (node.individualMStarBookmarks[agentIndex] != 0)
    //        node = node.prevStep;
    //    return node;
    //}
}