package AIT_2022.module.complex.fb;

import java.util.*;
import java.util.stream.Collectors;

import static java.util.Comparator.*;

import adf.core.agent.communication.MessageManager;
import adf.core.agent.communication.standard.bundle.centralized.CommandAmbulance;
import adf.core.agent.communication.standard.bundle.centralized.CommandFire;
import adf.core.agent.communication.standard.bundle.centralized.CommandPolice;
import adf.core.agent.communication.standard.bundle.information.MessageAmbulanceTeam;
import adf.core.agent.communication.standard.bundle.information.MessageCivilian;
import adf.core.agent.communication.standard.bundle.information.MessageFireBrigade;
import adf.core.agent.communication.standard.bundle.information.MessagePoliceForce;
import adf.core.agent.develop.DevelopData;
import adf.core.agent.info.AgentInfo;
import adf.core.agent.info.ScenarioInfo;
import adf.core.agent.info.WorldInfo;
import adf.core.agent.module.ModuleManager;
import adf.core.component.communication.CommunicationMessage;
import adf.core.component.module.AbstractModule;
import rescuecore2.misc.Pair;
import rescuecore2.standard.entities.*;
import rescuecore2.worldmodel.EntityID;
import static rescuecore2.standard.entities.StandardEntityURN.*;

import AIT_2022.module.complex.common.EntityFilter;
import AIT_2022.module.complex.common.LifeXpectancy;

public class TargetSelector extends AbstractModule {
    private boolean needRescueForMe = false;
    private EntityID result = null;
    private EntityID lastPos = null;
    private Set<EntityID> rescueCivTask = new HashSet<>();
    private Set<EntityID> rescueAgentTask = new HashSet<>();
    private Set<EntityID> loadCivTask = new HashSet<>();
    private Set<EntityID> avoidTasks = new HashSet<>();
    private Set<EntityID> clearTasks = new HashSet<>();
    private Map<EntityID, Integer> avoidTasksCounter = new HashMap<>();
    private Set<EntityID> nonTarget = new HashSet<>();
    private Set<EntityID> sendTarget = new HashSet<>();
    private int[] f = { 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 };
    private Pair<Integer, Integer> centor = null;
    private Set<EntityID> allocatedEintites = new HashSet<>();

    final private int FB_COUNTER = f.length;
    final private int RESCUE_MAX_FB = 3;
    private final boolean VOICE = false;
    private final boolean RADIO = true;

    public TargetSelector(AgentInfo ai, WorldInfo wi, 
            ScenarioInfo si, ModuleManager mm, DevelopData dd) {
        super(ai, wi, si, mm, dd);
        this.worldInfo = wi;
        this.agentInfo = ai;
    }

    @Override
    public AbstractModule updateInfo(MessageManager mm) {
        this.receiveMessage(mm);
        this.updateTask();
        // this.sendMessage(mm);
        this.lastPos = this.agentInfo.getPosition();
        return this;
    }

    @Override
    public AbstractModule calc() {
        final EntityID me = this.agentInfo.getID();
        this.needRescueForMe = this.rescueTarget(me);
        if (needRescueForMe) return this;

        Set<EntityID> changes = this.worldInfo.getChanged().getChangedEntities();
        changes.add(me);

        if (!rescueTarget(this.result) || changes.contains(this.result))
            this.result = null;

        Map<EntityID, Set<EntityID>> civPos = new HashMap<>();
        Map<EntityID, Set<EntityID>> ambPos = new HashMap<>();
        Map<EntityID, Set<EntityID>> fbPos = new HashMap<>();
        Set<EntityID> buriedAMB = new HashSet<>();
        Set<EntityID> buriedFB = new HashSet<>();
        Set<EntityID> buriedPF = new HashSet<>();

        for (EntityID e : changes) {
            final StandardEntity se = this.worldInfo.getEntity(e);
            if (!(se instanceof Human)) continue;
            final StandardEntity targetPosition = this.worldInfo.getPosition(e);
            if (targetPosition.getStandardURN().equals(REFUGE)) continue;
            final boolean resueTarget = this.rescueTarget(e);
            final boolean loadTarget = this.loadTarget(e);
            final StandardEntityURN targetURN = se.getStandardURN();
            final EntityID targetPositionID = targetPosition.getID();
            if (targetURN.equals(AMBULANCE_TEAM)) {
                if (resueTarget) {
                    buriedAMB.add(e);
                } else {
                    Set<EntityID> AMBs = ambPos.get(targetPositionID);
                    if (AMBs == null)
                        AMBs = new HashSet<>();
                    AMBs.add(e);
                    ambPos.put(targetPositionID, AMBs);
                }
            } else if (targetURN.equals(FIRE_BRIGADE)) {
                if (resueTarget) {
                    buriedFB.add(e);
                } else {
                    Set<EntityID> FBs = fbPos.get(targetPositionID);
                    if (FBs == null)
                        FBs = new HashSet<>();
                    FBs.add(e);
                    fbPos.put(targetPositionID, FBs);
                }
            } else if (targetURN.equals(POLICE_FORCE)) {
                if (resueTarget)
                    buriedPF.add(e);
            } else if (targetURN.equals(CIVILIAN)) {
                if (resueTarget) {
                    Set<EntityID> CIVs = civPos.get(targetPositionID);
                    if (CIVs == null)
                        CIVs = new HashSet<>();
                    CIVs.add(e);
                    civPos.put(targetPositionID, CIVs);
                } else if (loadTarget) {
                    this.loadCivTask.add(e);
                }
            }
        }

        //埋没してるエージェントがいる
        this.result = this.getLowerAgentBury(buriedFB);
        if (this.result != null) return this;
        this.result = this.getLowerAgentBury(buriedPF);
        if (this.result != null) return this;
        this.result = this.getLowerAgentBury(buriedAMB);
        if (this.result != null) return this;

        int n = 0;
        for (Set<EntityID> es : fbPos.values())
            n += es.size();
        final int idx = this.agentInfo.getTime() % FB_COUNTER;
        this.f[idx] = n;

        if (civPos.isEmpty()) {
            if (this.result == null)
                this.result = this.selectReceiveTask();
            return this;
        }

        final StandardEntity position = this.worldInfo.getPosition(me);
        final EntityID myPositionID = position.getID();
        Set<EntityID> targets = civPos.get(myPositionID);
        if (targets != null) {
            final Set<EntityID> FBs = fbPos.get(position.getID());
            int cnt = 0;
            for (EntityID e : this.idSort(targets)) {
                if (this.joinRescue(FBs, cnt, e)) {
                    this.result = e;
                    return this;
                } else {
                    if (FBs == null) {
                        if (this.joinRescue(fbPos.get(myPositionID), n, e)) {
                            this.result = e;
                            return this;
                        }
                    } else {
                        this.nonTarget.add(e);
                    }
                }
                cnt++;
            }
            civPos.remove(myPositionID);
        }

        for (Map.Entry<EntityID, Set<EntityID>> es : civPos.entrySet()) {
            final Set<EntityID> FBs = fbPos.get(es.getKey());
            final int FBnum = (FBs == null) ? 1 : FBs.size() + 1;
            int cnt = 0;
            for (EntityID e : this.idSort(es.getValue())) {
                if (this.avoidTasks.contains(e)) continue;
                if (this.canSave(e, Math.max(FBnum, this.getMaxFBcounter()))) {
                    if (this.joinRescue(FBs, cnt)) {
                        this.result = e;
                        return this;
                    } else {
                        this.avoidTasks.add(e);
                    }
                    cnt++;
                }
            }
        }

        if (this.result == null) this.result = this.selectReceiveTask();
        return this;
    }

    public EntityID getResult() {
        return this.result;
    }

    public Set<EntityID> getReceiveTasks() {
        return this.rescueCivTask;
    }

    public void setEntityInCluster(Collection<EntityID> ids) {
        this.allocatedEintites = new HashSet<>(ids);
        this.setClusterCentor();
    }

    private void setClusterCentor() {
        double x = 0.0;
        double y = 0.0;
        for (EntityID e : this.allocatedEintites) {
            final Pair<Integer, Integer> p = this.worldInfo.getLocation(e);
            x += p.first();
            y += p.second();
        }
        final int n = this.allocatedEintites.size();
        this.centor = new Pair<>((int)(x/n),(int)(y/n));
    }

    private boolean rescueTarget(EntityID e) {
        if (e == null) return false;
        final Human h = (Human) this.worldInfo.getEntity(e);
        final int hp = h.getHP();
        final int buried = h.getBuriedness();
        return (hp > 0 && buried > 0);
    }

    private boolean loadTarget(EntityID e) {
        final Human h = (Human) this.worldInfo.getEntity(e);
        if (!(h instanceof Civilian)) return false;
        final int hp = h.getHP();
        final int buried = h.getBuriedness();
        final int damage = h.getDamage();
        return (hp > 0 && buried == 0 && damage > 0);
    }

    private Set<EntityID> idSort(Set<EntityID> list) {
        if (list == null) return null;
        if (list.isEmpty()) return null;
        int[] ids = new int[list.size()];
        int idx = 0;
        for (EntityID e : list) ids[idx++] = e.getValue();
        Arrays.sort(ids);
        Set<EntityID> ret = new HashSet<>();
        for (int i : ids) ret.add(new EntityID(i));
        return ret;
    }

    private boolean joinRescue(Set<EntityID> list, int n) {
        if (list == null) return true;
        list.add(this.agentInfo.getID());
        final int max = Math.min(list.size(), this.getMaxRescueNumber() 
                + this.getMaxRescueNumber() * n);
        int[] ids = new int[list.size()];
        int idx = 0;
        for (EntityID e : list) ids[idx++] = e.getValue();
        Arrays.sort(ids);
        final int myID = this.agentInfo.getID().getValue();
        for (int i = 0; i < max; ++i) 
            if (ids[i] == myID) return true;
        return false;
    }

    private boolean joinRescue(Set<EntityID> list, int n, EntityID id) {
        if (list == null) return true;
        list.add(this.agentInfo.getID());
        final int needNum 
                = this.getMaxRescueNumber(id) + this.getMaxRescueNumber(id) * n;
        final int max = Math.min(list.size(), needNum);
        int[] ids = new int[list.size()];
        int idx = 0;
        for (EntityID e : list) ids[idx++] = e.getValue();
        Arrays.sort(ids);
        final int myID = this.agentInfo.getID().getValue();
        for (int i = 0; i < max; ++i)
            if (ids[i] == myID) return true;
        return false;
    }

    private int getMaxRescueNumber() {
        return RESCUE_MAX_FB + (int) (this.agentInfo.getTime() * 0.01);
    }

    private int getMaxRescueNumber(EntityID id) {
        return this.getNeedFBnum(id) + (int) (this.agentInfo.getTime() * 0.01);
    }

    private EntityID getLowerAgentBury(Set<EntityID> agents) {
        Human h = agents.stream()
                .map(this.worldInfo::getEntity)
                .filter(e -> e instanceof Human)
                .map(Human.class::cast)
                .filter(e -> e.getHP() > 0)
                .min(comparing(a -> a.getBuriedness()))
                .orElse(null);
        return h == null ? null : h.getID();
    }

    private boolean canSave(EntityID id, int n) {
        final StandardEntity s = this.worldInfo.getEntity(id);
        if (((Human) s).getBuriedness() == 0 || ((Human) s).getHP() == 0) return false;

        if (((Civilian) s).getBuriedness() > 0 && ((Civilian) s).getHP() > 0) {
            final LifeXpectancy lx = new LifeXpectancy((Human) s);
            final int ttl = lx.getLifeXpectancy();
            return (ttl - this.rescueTime(id, n) > 0);
        } else {
            return true;
        }
    }

    private boolean canRescueReach(EntityID id) {
        if (this.result == null)
            return false;
        if (this.lastPos.getValue() != this.agentInfo.getPosition().getValue())
            return false;
        if (!(this.agentInfo.getPositionArea() instanceof Road))
            return false;
        return this.worldInfo.getChanged().getChangedEntities().contains(this.result);
    }

    private int rescueTime(EntityID id, int n) {
        return ((Human) this.worldInfo.getEntity(id)).getBuriedness() / n;
    }

    private int getNeedFBnum(EntityID id) {
        LifeXpectancy lx = new LifeXpectancy((Human) this.worldInfo.getEntity(id));
        final int b = ((Human) this.worldInfo.getEntity(id)).getBuriedness();
        final int l = lx.getLifeXpectancy();
        int f = 1;
        while (l - b / f <= 0) ++f;
        return f + 1;
    }

    private int getMaxFBcounter() {
        int max = 0;
        for (int i : this.f)
            if (i > max) max = i;
        return max;
    }

    private void handleMessage(MessageCivilian msg) {
        final EntityID targetID = msg.getAgentID();
        if (this.rescueTarget(targetID)) {
            this.rescueCivTask.add(targetID);
        } else {
            this.nonTarget.add(targetID);
            this.rescueCivTask.remove(targetID);
        }
    }

    private void handleMessage(MessageFireBrigade msg) {
        final EntityID targetID = msg.getAgentID();
        if (this.rescueTarget(targetID)) {
            this.rescueAgentTask.add(targetID);
            this.nonTarget.add(targetID);
        } else {
            this.rescueAgentTask.remove(targetID);
        }
    }

    private void handleMessage(MessageAmbulanceTeam msg) {
        final EntityID targetID = msg.getAgentID();
        if (this.rescueTarget(targetID)) {
            this.rescueAgentTask.add(targetID);
            this.nonTarget.add(targetID);
        } else {
            this.rescueAgentTask.remove(targetID);
        }
    }

    private void handleMessage(MessagePoliceForce msg) {
        final EntityID targetID = msg.getAgentID();
        if (this.rescueTarget(targetID)) {
            this.rescueAgentTask.add(targetID);
            this.nonTarget.add(targetID);
        } else {
            this.rescueAgentTask.remove(targetID);
        }
    }

    private void receiveMessage(MessageManager mm) {
        final List<CommunicationMessage> cmsg 
                = mm.getReceivedMessageList(MessageCivilian.class);
        for (CommunicationMessage tmp : cmsg) {
            final MessageCivilian message = (MessageCivilian) tmp;
            this.handleMessage(message);
        }
        final List<CommunicationMessage> amsg 
                = mm.getReceivedMessageList(MessageAmbulanceTeam.class);
        for (CommunicationMessage tmp : amsg) {
            final MessageAmbulanceTeam message = (MessageAmbulanceTeam) tmp;
            this.handleMessage(message);
        }
        final List<CommunicationMessage> fmsg 
                = mm.getReceivedMessageList(MessageFireBrigade.class);
        for (CommunicationMessage tmp : fmsg) {
            final MessageFireBrigade message = (MessageFireBrigade) tmp;
            this.handleMessage(message);
        }
        final List<CommunicationMessage> pmsg 
                = mm.getReceivedMessageList(MessagePoliceForce.class);
        for (CommunicationMessage tmp : pmsg) {
            final MessagePoliceForce message = (MessagePoliceForce) tmp;
            this.handleMessage(message);
        }
    }

    private void sendMessage(MessageManager mm) {
        EntityFilter filter 
                = new EntityFilter(this.worldInfo, this.agentInfo);
        final EntityID positionID = this.agentInfo.getPosition();
        final EntityID PF = filter.getNearyPF();
        if (this.canRescueReach(this.result)) {
            mm.addMessage(new CommandPolice(
                    RADIO, PF, positionID, CommandPolice.ACTION_CLEAR));
            mm.addMessage(new CommandPolice(
                    VOICE, PF, positionID, CommandPolice.ACTION_CLEAR));
            if (PF == null) {
                this.clearTasks.add(this.result);
                this.avoidTasks.add(this.result);
                this.result = null;
            }
            return;
        }

        final EntityID AT = filter.getNearyAT();
        if (AT != null && !this.loadCivTask.isEmpty()) {
            final Iterator<EntityID> i = this.loadCivTask.iterator();
            final EntityID targetID 
                    = this.worldInfo.getPosition(i.next()).getID();
            mm.addMessage(new CommandAmbulance(
                    VOICE, AT, targetID, CommandAmbulance.ACTION_MOVE));
            i.remove();
            return;
        }

        if (PF != null && !this.clearTasks.isEmpty()) {
            Iterator<EntityID> i = this.clearTasks.iterator();
            final EntityID targetID = i.next();
            mm.addMessage(new CommandPolice(
                    RADIO, PF, targetID, CommandPolice.ACTION_CLEAR));
            mm.addMessage(new CommandPolice(
                    VOICE, PF, targetID, CommandPolice.ACTION_CLEAR));
            this.avoidTasks.remove(targetID);
            i.remove();
            return;
        }

        final EntityID FB = filter.getNearyFB();
        if (FB != null && !this.rescueAgentTask.isEmpty()) {
            final Iterator<EntityID> i = this.rescueAgentTask.iterator();
            final EntityID e = i.next();
            final StandardEntity se = this.worldInfo.getEntity(e);
            if (se instanceof AmbulanceTeam) {
                mm.addMessage(new MessageAmbulanceTeam(
                        VOICE, (AmbulanceTeam) se, CommandAmbulance.ACTION_MOVE, e));
            } else if (se instanceof FireBrigade) {
                mm.addMessage(new MessageFireBrigade(
                        VOICE, (FireBrigade) se, CommandFire.ACTION_MOVE, e));
            } else if (se instanceof PoliceForce) {
                mm.addMessage(new MessagePoliceForce(
                        VOICE, (PoliceForce) se, CommandPolice.ACTION_MOVE, e));
            }
            i.remove();
            return;
        }

        Set<EntityID> targets = new HashSet<>(this.rescueCivTask);
        targets.removeAll(this.sendTarget);
        targets.removeAll(this.nonTarget);
        if (targets.isEmpty()) targets = new HashSet<>(this.sendTarget);
        if (targets.isEmpty()) return;
        final Civilian c = (Civilian) this.worldInfo.getEntity(targets.iterator().next());
        mm.addMessage(new MessageCivilian(VOICE, c));
        this.sendTarget.add(c.getID());
    }

    private void updateTask() {
        Iterator<EntityID> i = this.rescueCivTask.iterator();
        while (i.hasNext()) {
            final EntityID e = i.next();
            if (!this.rescueTarget(e)) {
                this.nonTarget.add(e);
                i.remove();
            }
        }
        i = this.rescueAgentTask.iterator();
        while (i.hasNext()) {
            final EntityID e = i.next();
            if (!this.rescueTarget(e)) {
                this.nonTarget.add(e);
                i.remove();
            }
        }
        i = this.loadCivTask.iterator();
        while (i.hasNext()) {
            final EntityID e = i.next();
            if (!this.loadTarget(e)) {
                this.nonTarget.add(e);
                i.remove();
            }
        }
        this.rescueCivTask.removeAll(this.avoidTasks);
        this.rescueCivTask.removeAll(this.nonTarget);
        this.rescueAgentTask.removeAll(this.nonTarget);
        i = this.avoidTasks.iterator();
        while (i.hasNext()) {
            final EntityID e = i.next();
            if (!this.rescueTarget(e)) {
                this.nonTarget.add(e);
                i.remove();
            }
        }
        i = this.sendTarget.iterator();
        while (i.hasNext()) {
            final EntityID e = i.next();
            if (!this.rescueTarget(e)) {
                this.nonTarget.add(e);
                i.remove();
            }
        }
        if (this.rescueAgentTask.isEmpty()) this.avoidTasks.clear();
    }

    private EntityID getPriorityTask(Set<EntityID> tasks) {
        EntityID ret = null;
        double priority = Double.MAX_VALUE;
        final int MOVE_TIME = 10000;
        for (EntityID e : tasks) {
            final Human h = (Human) this.worldInfo.getEntity(e);
            final Pair<Integer, Integer> p = this.worldInfo.getLocation(e);
            final int outOfArea = (this.allocatedEintites.contains(e) ? 0 : 
                    Math.abs(this.centor.first() - p.first()) 
                    + Math.abs(this.centor.second() - p.second()));
            final double distance = Math.abs(this.agentInfo.getX() - p.first())
                    + Math.abs(this.agentInfo.getY() - p.second());
            final double cost1 = Math.sqrt(distance);
            final LifeXpectancy lx = new LifeXpectancy(h);
            final double cost2 = lx.getLifeXpectancy() - (h.getDamage() * h.getBuriedness());
            final double cost = cost1 + cost2 + outOfArea;
            if (priority > cost) {
                priority = cost;
                ret = e;
            }
        }
        return ret;
    }

    private EntityID selectReceiveTask() {
        if (!this.rescueAgentTask.isEmpty())
            return this.getPriorityTask(this.rescueAgentTask);
        if (!this.rescueCivTask.isEmpty())
            return this.getPriorityTask(this.rescueCivTask);
        return null;
    }
}
