package mas.agent.student;

import cz.agents.alite.communication.Message;
import cz.agents.alite.communication.content.Content;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Random;
import java.util.Set;
import mas.agent.MASQueenAgent;

/**
 * This example agent illustrates the usage API available in the MASQueenAgent class.
 */
public class MyQueenAgent extends MASQueenAgent {
    private int _selection;
    private List<Integer> _reachAgents;
    private HashMap<Integer, Integer> _view;
    private Set<Map<Integer, Integer>> _noGoods;
    Random rnd;

    public MyQueenAgent(int agentId, int nAgents) {
        // Leave this method as it is...
        super(agentId, nAgents);
        rnd = new Random();
        this._selection = rnd.nextInt(nAgents);
        this._reachAgents = new ArrayList<>();
        this._view = new HashMap<>();
        this._noGoods = new HashSet<>();
        for (int i = agentId + 1; i < nAgents; i++) {
            _reachAgents.add(i);
        }
    }

    @Override
    protected void start(int agentId, int nAgents) {
        sendAllReachable(new HandleOk(getAgentId(), _selection));
    }

    private void sendAllReachable(Content msg){
        for(Integer agent : this._reachAgents){
            sendMessage(agent.toString(), msg);
        }
    }

    @Override
    protected void processMessages(List<Message> newMessages) {
        // This method is called whenever there are any new messages for the robot

        // You can list the messages like this:
        for (Message message : newMessages) {
            Content cont = message.getContent();
            if (cont instanceof HandleOk) {
                HandleOk handle = (HandleOk) cont;
                _view.put(handle.getSender(), handle.getValue());
                this.handle_Ok(handle);
            } else if (cont instanceof HandleNoGood) {
                HandleNoGood handle = (HandleNoGood) cont;
                this.handle_noGood(handle);
            } else if (cont instanceof HandleNoSolution) {
                notifySolutionDoesNotExist();
            } else if (cont instanceof HandleSolution) {
                HandleSolution handle = (HandleSolution) cont;
                if(handle.getIsFinal()){
                    notifySolutionFound(_selection);
                } else {
                    handle_solutionNotFinal(handle);
                }
            } else {
                System.out.println("---");
            }
        }

    }

    private void handle_Ok(HandleOk handle){
        if (!isConsistent(handle.getSender(), handle.getValue())) {
            // generate new assignment or nogood
            createNewSelection();
        } else if (getAgentId() == nAgents() - 1 && _view.size() == nAgents() - 1) {
            // lowest priority agent will suggest a solution
            HashMap<Integer, Integer> solution = new HashMap<>(_view);
            solution.put(getAgentId(), _selection);
            sendMessage(Integer.toString(getAgentId() - 1), new HandleSolution(solution, false));
        }
    }

    private void handle_noGood(HandleNoGood handle){
        if (!isNoGood(handle, false)) {
            sendAllReachable(new HandleOk(getAgentId(), _selection));
            return;
        }
        _noGoods.add(handle.getView());
        for (Map.Entry<Integer, Integer> entry : handle.getView().entrySet()) {
            if (!_view.containsKey(entry.getKey()) && entry.getKey() != getAgentId()) {
                _view.put(entry.getKey(), entry.getValue());
            }
        }
        createNewSelection();
    }

    private void handle_solutionNotFinal(HandleSolution handle){
        if (validateSolution(handle.getView())) {
            if (getAgentId() == 0) {
                broadcast(new HandleSolution(_view,true));
                notifySolutionFound(_selection);
            } else {
                sendMessage(Integer.toString(getAgentId() - 1), handle);
            }
        } else {
            sendAllReachable(new HandleOk(getAgentId(), _selection));
        }
    }

    private void createNewSelection(){
        Integer validSelection = -1;
        ArrayList<Integer> validList = new ArrayList<>();
        for(int i=0; i<nAgents();i++) {
            if (i != this._selection) {
                boolean isValid = true;
                for (Map.Entry<Integer, Integer> entry : _view.entrySet()) {
                    if (!isConsistent(entry.getKey(), entry.getValue(), getAgentId(), i)) {
                        isValid = false;
                        break;
                    }
                }
                Iterator<Map<Integer, Integer>> it = _noGoods.iterator();
                while (it.hasNext()) {
                    Map<Integer, Integer> noGood = it.next();
                    for (Map.Entry<Integer, Integer> entry : noGood.entrySet()) {
                        if (entry.getKey() == getAgentId() && entry.getValue() == i) {
                            isValid = false;
                            break;
                        }
                    }
                }
                if (isValid) {
                    validList.add(i);
                }
            }
        }
        if (validList.isEmpty()) {
            validSelection =  -1;
        } else {
            Random r = new Random();
            validSelection = validList.get(rnd.nextInt(validList.size()));
        }
        if (validSelection == -1){
            int index = createNoGood();
            if (index != -1){
                this._view.remove(index);
                createNewSelection();
            }
        } else {
            this._selection = validSelection;
            if (getAgentId() == nAgents() - 1 && _view.size() == nAgents() - 1) {
                HashMap<Integer, Integer> solu = new HashMap<>(_view);
                solu.put(getAgentId(), _selection);
                sendMessage(getAgentId() - 1 + "", new HandleSolution(solu,false));
            } else {
                sendAllReachable(new HandleOk(getAgentId(), this._selection));
            }
        }
    }

    private int createNoGood() {
        if (_view.isEmpty()) {
            broadcast(new HandleNoSolution());
            notifySolutionDoesNotExist();
            return -1;
        } else {
            int agent = -1;
                for (Integer i : _view.keySet()) {
                    if (i > agent) {
                        agent = i;
                    }
                }
            sendMessage("" + agent, new HandleNoGood(_view));
            return agent;
        }
    }

    private boolean isConsistent(int sender, int assignment) {
        return isConsistent(sender, assignment, getAgentId(), _selection);
    }

    private boolean isConsistent(int senderId, int senderAssignment, int myId, int myAssignment) {
        //column check
        if (myAssignment == senderAssignment) {
            return false;
            // diagonal check
        } else if (myId - senderId == Math.abs(myAssignment - senderAssignment)) {
            return false;
        }
        return true;
    }

    private boolean isNoGood(HandleNoGood noGood, boolean skipMe) {
        return isNoGood(noGood.getView().entrySet(), skipMe);
    }

    private boolean isNoGood(Set<Map.Entry<Integer, Integer>> noGood, boolean skipMe) {
        boolean valid = true;
        for (Map.Entry<Integer, Integer> entry : noGood) {
            if (skipMe) {
                if (entry.getKey() == getAgentId()) {
                    continue;
                }
            } else {
                if (entry.getKey() == getAgentId()) {
                    if (entry.getValue() != _selection) {
                        valid = false;
                        break;
                    } else {
                        continue;
                    }
                }
            }
            if (_view.containsKey(entry.getKey())) {
                if (_view.get(entry.getKey()).intValue() != entry.getValue()){
                    valid = false;
                    break;
                }
            } else {
                valid = false;
                break;
            }
        }
        return valid;
    }

    private boolean validateSolution(Map<Integer, Integer> assignments) {
        for (Map.Entry<Integer, Integer> entry : assignments.entrySet()) {
            if (entry.getKey() == getAgentId()) {
                if (entry.getValue() == _selection) {
                    continue;
                } else {
                    return false;
                }
            }
            if (entry.getKey() < getAgentId()) {
                if (_view.containsKey(entry.getKey())) {
                    if (entry.getValue().intValue() != _view.get(entry.getKey()).intValue()) {
                        return false;
                    }
                } else {
                    return false;
                }
            }
        }
        return true;
    }
}