using System.Collections.Generic;
using System.Linq;
using UnityEngine;

public class PlanNode
{
    public PlanNode parent;
    public float cost;
    public Dictionary<string, int> states;
    public AI_Action action;

    public PlanNode(PlanNode parent, float cost, Dictionary<string, int> allStates, AI_Action action)
    {
        this.parent = parent;
        this.cost = cost;
        states = new Dictionary<string, int>(allStates);
        this.action = action;
    }

    public PlanNode(PlanNode parent, float cost, Dictionary<string, int> allStates, 
                   Dictionary<string, int> agentStates, AI_Action action)
    {
        this.parent = parent;
        this.cost = cost;
        states = new Dictionary<string, int>(allStates);

        foreach(KeyValuePair<string, int> state in agentStates)
        {
            if (!states.ContainsKey(state.Key))
            {
                states.Add(state.Key, state.Value);
            }
        }
        this.action = action;
    }
}

public class AI_Planner
{
    private string m_PlanDescription;

    public Queue<AI_Action> Plan(List<AI_Action> actions, Dictionary<string, int> goal, AIStates agentStates)
    {
        List<AI_Action> usableActions = new List<AI_Action>();
        foreach (AI_Action action in actions)
        {
            if(action.IsAchievable())
                usableActions.Add(action);
        }

        List<PlanNode> leaves = new List<PlanNode>();
        PlanNode start = new PlanNode(null, 0, AI_World.Instance.GetWorld().GetStates(), 
                                     agentStates.GetStates(), null);

        bool success = BuildGraph(start, leaves, usableActions, goal);

        if(!success)
        {
            Debug.Log("No plan could be constructed");
            return null;
        }

        PlanNode cheapest = null;
        foreach (var leaf in leaves)
        {
            if (cheapest == null)
                cheapest = leaf;
            else if (leaf.cost < cheapest.cost)
                cheapest = leaf;
        }

        var result = new List<AI_Action>();
        var n = cheapest;
        while (n != null)
        {
            if(n.action != null)
            {
                result.Insert(0, n.action);
            }
            n = n.parent;
        }

        var queue = new Queue<AI_Action>();
        foreach(var action in result)
        {
            queue.Enqueue(action);
        }

        m_PlanDescription = "";
        int index = 0;
        foreach(AI_Action action in queue)
        {
            m_PlanDescription += index + ": " + action.actionName + "\n";
            index++;
        }

        return queue;
    }

    private bool BuildGraph(PlanNode parent, List<PlanNode> leaves, List<AI_Action> usableActions, 
                          Dictionary<string, int> goal)
    {
        bool foundPath = false;
        foreach(AI_Action action in usableActions)
        {
            if(action.IsAchievableGiven(parent.states))
            {
                Dictionary<string, int> currentState = new Dictionary<string, int>(parent.states);
                foreach(KeyValuePair<string, int> effect in action.effectsDict)
                    if(!currentState.ContainsKey(effect.Key))
                        currentState.Add(effect.Key, effect.Value);

                PlanNode node = new PlanNode(parent, parent.cost + action.cost, currentState, action);

                if(GoalAchieved(goal, currentState))
                {
                    leaves.Add(node);
                    foundPath = true;
                }
                else
                {
                    List<AI_Action> subset = ActionSubset(usableActions, action);
                    bool found = BuildGraph(node, leaves, subset, goal);
                    if(found)
                        foundPath = true;
                }
            }
        }
        return foundPath;
    }

    private bool GoalAchieved(Dictionary<string, int> goals, Dictionary<string, int> state)
    {
        foreach(KeyValuePair<string, int> goal in goals)
        {
            if (!state.ContainsKey(goal.Key))
                return false;
        }
        return true;
    }

    private List<AI_Action> ActionSubset(List<AI_Action> actions, AI_Action removeAction)
    {
        List<AI_Action> subset = new List<AI_Action>();
        foreach(AI_Action action in actions)
        {
            if(!action.Equals(removeAction))
                subset.Add(action);
        }
        return subset;
    }

    public string GetPlanDescription()
    {
        return m_PlanDescription;
    }
}