package mas.agent.student;

import java.util.*;

import cz.agents.alite.communication.content.Content;
import mas.agent.MASQueenAgent;
import cz.agents.alite.communication.Message;

public class MyQueenAgent extends MASQueenAgent {

    private int currentAssignment;
    private Map<Integer, Integer> agentView = new HashMap<>();
    private Set<Map<Integer, Integer>> noGoodStore = new HashSet<>();;

    public MyQueenAgent(int agentId, int nAgents) {
        super(agentId, nAgents);
    }

    @Override
    protected void start(int agentId, int nAgents) {
        // this method is called when the agent is initialized.
        currentAssignment = new Random().nextInt(nAgents);
        messageLowerAgents();
    }

    @Override
    protected void processMessages(List<Message> newMessages) {
        for (Message message : newMessages) {
            Content content = message.getContent();

            if(content instanceof HandleOk) {
                processOkMessage(message);
            }
            if(content instanceof HandleNoGood) {
                processNoGoodMessage(message);
            }
            if(content instanceof HandlePossibleSolution) {
                processPossibleSolutionMessage(message);
            }
            if(content instanceof HandleSolutionDoesNotExist) {
                processSolutionDoesNotExist();
            }
            if(content instanceof HandleSolutionFound) {
                processSolutionFound();
            }
        }
    }

    private ArrayList<Integer> getNewAssignments() {
        ArrayList<Integer> list = new ArrayList<>();
        // iterate through all columns
        for (int i = 0; i < nAgents(); i++) {
            // process every column except current column
            if (i != currentAssignment) {
                //first check my agent view
                boolean add = checkAgentView(i);
                //if agent view is ok check my noGoodStore
                if (add) {
                    add = checkNoGoodStore(i);
                }
                // if everything ok I can add this assignment to result list
                if(add) {
                    list.add(i);
                }
            }
        }
        return list;
    }

    private boolean checkAgentView(int i) {
        // iterate through every entry in agent view
        for (Map.Entry<Integer, Integer> mapEntry : agentView.entrySet()) {
            int agentName = mapEntry.getKey();
            int agentAssignment = mapEntry.getValue();
            // check if any constraint is violated
            if(constraintsViolated(getAgentId(), i, agentName, agentAssignment)) {
                return false;
            }
        }
        return true;
    }

    private boolean checkNoGoodStore(int i) {
        boolean add = true;
        // list for storing noGoods which will be removed
        ArrayList<Map<Integer, Integer>> toRemove = new ArrayList<>();
        for (Map<Integer, Integer> noGood : noGoodStore) {
            boolean removeNoGood = false;
            for (Map.Entry<Integer, Integer> mapEntry : noGood.entrySet()) {
                int key = mapEntry.getKey();
                int value = mapEntry.getValue();
                if (key != getAgentId()) {
                    if (agentView.containsKey(key)) {
                        if (agentView.get(key) != value){
                            removeNoGood = true;
                        }
                    } else {
                        removeNoGood = true;
                    }
                } else {
                    if (value == i) {
                        add = false;
                        break;
                    }
                }
            }
            if (removeNoGood) {
                // we cannot remove element directly from set we are iterating
                toRemove.add(noGood);
            }
        }
        for (Map<Integer, Integer> noGood : toRemove) {
            noGoodStore.remove(noGood);
        }
        return add;
    }

    private int sendNoGood() {
        if (!agentView.isEmpty()) {
            int targetAgent = Integer.MIN_VALUE;
            // get agent from my agent view with lowest priority
            for (Integer integer : agentView.keySet()) {
                if (integer > targetAgent)  targetAgent = integer;
            }
            // send him noGood
            sendMessage(String.valueOf(targetAgent), new HandleNoGood(agentView));
            return targetAgent;
        } else {
            // if agent view is empty I can broadcast that no solution exists
            broadcast(new HandleSolutionDoesNotExist());
            notifySolutionDoesNotExist();
            terminated = true;
            return Integer.MIN_VALUE;
        }
    }

    private boolean lowestAgent() {
        return ( nAgents() == agentView.size() + 1 ) && ( nAgents() == getAgentId() + 1 );
    }

    private boolean constraintsViolated(int agent1Name, int agent1Assignment, int agent2Name, int agent2Assignment) {
        // check if agent1 isn't in the same column as agent2 or that agent1 isn't in same diagonal as agent2
        return (agent1Assignment == agent2Assignment) || (Math.abs(agent1Assignment - agent2Assignment) == Math.abs(agent1Name - agent2Name));
    }

    private boolean solutionOkForMe(Map<Integer, Integer> finalAssignments) {
        // iterate through all assignments
        for (Map.Entry<Integer, Integer> mapEntry : finalAssignments.entrySet()) {
            // check if my assignment matches with solution assignment
            if ( (getAgentId() == mapEntry.getKey()) && (currentAssignment != mapEntry.getValue()) ) {
                return false;
            }
        }
        return true;
    }

    private boolean noGoodOkForMe(Map<Integer, Integer> map) {
        // iterate through all assignment pairs
        for (Map.Entry<Integer, Integer> mapEntry : map.entrySet()) {
            int key = mapEntry.getKey();
            int value = mapEntry.getValue();
            // if I encounter key same as my id
            if ( key == getAgentId() ) {
                // check if my current assignment is different from its value
                if ( currentAssignment != value ) {
                    return false;
                }
            }
            // also check if current key is contained in my agent view
            else if ( agentView.containsKey(key) ) {
                // check if this key has different value from key in my agent view
                if ( agentView.get(key) != value ){
                    return false;
                }
            } else {
                return false;
            }
        }
        return true;
    }

    private void backtrack() {
        // get all new available assignments
        ArrayList<Integer> list = getNewAssignments();
        // if there is none, send NoGood
        if (list.isEmpty()) {
            int targetAgent = sendNoGood();
            // if I got some agent
            if (targetAgent != Integer.MIN_VALUE) {
                // I can remove it from my agent view (cause he will change)
                agentView.remove(targetAgent);
                backtrack();
            }
        } else {
            // else choose randomly on of the assignment
            currentAssignment = list.get(new Random().nextInt(list.size()));
            // and send it to lower priority agents
            messageLowerAgents();
            // check if current agent is agent with lowest priority
            if ( lowestAgent() ) {
                sendSolutionAssignments();
            }
        }
    }

    private void messageLowerAgents() {
        // send message to agents with lower priority than me
        Content content = new HandleOk(getAgentId(), currentAssignment);
        for (int i = getAgentId() + 1; i < nAgents(); i++) {
            sendMessage(String.valueOf(i), content);
        }
    }

    private void sendSolutionAssignments() {
        Map<Integer, Integer> finalAssignments = new HashMap<>(agentView);
        // lowest priority agent puts his assignment into assignment map
        finalAssignments.put(getAgentId(), currentAssignment);
        // and send this map to the agent with higher priority
        sendMessage(String.valueOf(getAgentId() - 1), new HandlePossibleSolution(finalAssignments));
    }

    private void processOkMessage(Message message) {
        HandleOk content = (HandleOk) message.getContent();
        int senderName = content.getSenderName();
        int senderAssignment = content.getSenderAssignment();
        agentView.put(senderName, senderAssignment);
        // check if current agent doesn't violate constraints in agent view
        if ( constraintsViolated(senderName, senderAssignment, getAgentId(), currentAssignment) ) {
            backtrack();
        }
        // check if current agent is agent with lowest priority
        if ( lowestAgent() ) {
            sendSolutionAssignments();
        }
    }

    private void processNoGoodMessage(Message message) {
        HandleNoGood handleNoGood = (HandleNoGood) message.getContent();
        // check if sent noGood doesn't violate any constrains
        if( noGoodOkForMe(handleNoGood.getAssignments()) ){
            // if it's ok I can add it to my noGoodStore
            noGoodStore.add(handleNoGood.getAssignments());
            backtrack();
        } else {
            messageLowerAgents();
        }
    }

    private void processPossibleSolutionMessage(Message message) {
        HandlePossibleSolution handlePossibleSolution = (HandlePossibleSolution) message.getContent();
        // check if given solution assignment is ok for me
        if (solutionOkForMe(handlePossibleSolution.getAssignments())) {
            // if I am agent with the highest priority (0)
            if (getAgentId() == 0) {
                // I can broadcast that we found a solution
                broadcast(new HandleSolutionFound());
                // I notify that we found solution and my assignment
                notifySolutionFound(currentAssignment);
            }
            // else I just send the solution to agent with higher priority than me
            else {
                sendMessage(String.valueOf(getAgentId() - 1), handlePossibleSolution);
            }
        } else {
            // if the solution is not ok I send my assignment to agent with lower priority
            messageLowerAgents();
        }
    }

    private void processSolutionDoesNotExist() {
        // notify that no solution exists and terminate
        notifySolutionDoesNotExist();
        terminated = true;
    }

    private void processSolutionFound() {
        // I notify that we found solution and my assignment
        notifySolutionFound(currentAssignment);
    }

}

class HandleNoGood extends Content {

    private Map<Integer, Integer> assignments;

    HandleNoGood(Map<Integer, Integer> assignments) {
        super(assignments);
        this.assignments = new HashMap<>(assignments);
    }

    Map<Integer, Integer> getAssignments() {
        return assignments;
    }

    @Override
    public String toString() {
        return "HandleNoGood [" +
                "assignments=" + assignments +
                ']';
    }
}

class HandleOk  extends Content {

    private int senderName, senderAssignment;

    HandleOk(int senderName, int senderAssignment) {
        super(senderAssignment);
        this.senderName = senderName;
        this.senderAssignment = senderAssignment;
    }

    int getSenderName() {
        return senderName;
    }

    int getSenderAssignment() {
        return senderAssignment;
    }

    @Override
    public String toString() {
        return "HandleOk [" +
                "senderAssignment=" + senderAssignment +
                ']';
    }
}

class HandleSolutionDoesNotExist extends Content {

    HandleSolutionDoesNotExist() {
        super("HandleSolutionDoesNotExist");
    }

}

class HandleSolutionFound extends Content {

     HandleSolutionFound() {
        super("HandleSolutionFound");
    }

}

class HandlePossibleSolution extends Content {

    private final Map<Integer, Integer> assignments;

    HandlePossibleSolution(Map<Integer, Integer> assignments) {
        super(assignments);
        this.assignments = new HashMap<>(assignments);
    }

    Map<Integer, Integer> getAssignments() {
        return assignments;
    }

    @Override
    public String toString() {
        return "HandlePossibleSolution [" +
                "assignments=" + assignments +
                ']';
    }

}