package AITB_2022.module.complex;

import adf.core.agent.info.*;
import adf.core.component.module.complex.Search;
import adf.core.component.module.algorithm.Clustering;
import adf.core.component.module.algorithm.PathPlanning;
import adf.core.agent.module.ModuleManager;
import adf.core.agent.communication.MessageManager;
import adf.core.agent.develop.DevelopData;

import rescuecore2.standard.entities.*;
import rescuecore2.worldmodel.EntityID;

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

import javax.swing.text.html.parser.Entity;

import static rescuecore2.standard.entities.StandardEntityURN.*;

import adf.core.component.communication.*;
import adf.core.agent.communication.standard.bundle.information.MessageCivilian;

public class AITFBSearch extends Search {
    private EntityID result;
    private Clustering clustering;
    private PathPlanning pathPlanning;
    private int clusterIndex;
    private Set<EntityID> cluster = new HashSet<>();
    private Set<EntityID> searchedEntityIDs = new HashSet<>();
    private Set<Integer> searchedClusterIndex = new HashSet<>();
    private Set<EntityID> priorityBuildings = new HashSet<>();
    private int VoiceRange;

    // 消防隊がすでに通信で情報が送信された市民のIDを保持するリスト
    List<EntityID> sentNotBuriedCivilianIDs = new ArrayList<>();

    public AITFBSearch(AgentInfo ai, WorldInfo wi, ScenarioInfo si, ModuleManager mm, DevelopData dd) {
        super(ai, wi, si, mm, dd);

        String clusteringKey = "SampleSearch.Clustering";
        switch (ai.me().getStandardURN()) {
            case FIRE_BRIGADE:
                clusteringKey += ".Fire";
                break;
            case AMBULANCE_TEAM:
                clusteringKey += ".Ambulance";
                break;
            case POLICE_FORCE:
                clusteringKey += ".Police";
                break;
        }
        switch (si.getMode()) {
            case PRECOMPUTATION_PHASE:
                this.pathPlanning = moduleManager.getModule(
                        "ActionTransport.PathPlanning",
                        "adf.impl.module.algorithm.DijkstraPathPlanning");
                break;
            case PRECOMPUTED:
                this.pathPlanning = moduleManager.getModule(
                        "ActionTransport.PathPlanning",
                        "adf.impl.module.algorithm.DijkstraPathPlanning");
                break;
            case NON_PRECOMPUTE:
                this.pathPlanning = moduleManager.getModule(
                        "ActionTransport.PathPlanning",
                        "adf.impl.module.algorithm.DijkstraPathPlanning");
                break;
        }

        this.clustering = mm.getModule(clusteringKey, "adf.impl.module.algorithm.SampleKMeans");
        this.registerModule(this.clustering);
    }

    @Override
    public EntityID getTarget() {
        return this.result;
    }

    @Override
    public AITFBSearch preparate() {
        super.preparate();
        if (this.getCountPreparate() > 1)
            return this;
        this.pathPlanning.preparate();
        this.clusterIndex = this.clustering.getClusterIndex(this.agentInfo.me());
        this.cluster = this.getNextCluster(this.clusterIndex);
        this.VoiceRange = this.scenarioInfo.getRawConfig().getIntValue("comms.channels.0.range");
        return this;
    }

    @Override
    public Search updateInfo(MessageManager mm) {
        super.updateInfo(mm);

        if (mm.getHeardAgentHelpCount() > 0) {
            Set<EntityID> inRangeBuildingIDs = this.removeAlreadyReachedBuilding(
                    this.worldInfo.getObjectIDsInRange(this.agentInfo.me(), this.VoiceRange)
                            .stream()
                            .filter(objectID -> this.worldInfo.getEntity(objectID) instanceof Building)
                            .filter(objectID -> !(this.worldInfo.getEntity(objectID) instanceof Refuge))
                            .collect(Collectors.toSet()));
            if (this.priorityBuildings.isEmpty()) {
                this.priorityBuildings.addAll(inRangeBuildingIDs);
            } else {
                this.priorityBuildings.retainAll(inRangeBuildingIDs);
            }
        }

        if (this.worldInfo.getEntity(((FireBrigade) (this.agentInfo.me())).getPosition())
                .getStandardURN() == BUILDING) {
            this.worldInfo.getEntitiesOfType(CIVILIAN)
                    .stream()
                    .filter(Civilian.class::isInstance)
                    .map(Civilian.class::cast)
                    .filter(c -> !sentNotBuriedCivilianIDs.contains(c.getID()))
                    .filter(c -> c.getPosition().equals(this.agentInfo.getPosition()))
                    .filter(c -> c.isBuriednessDefined())
                    .filter(c -> c.getBuriedness() <= 0)
                    .forEach(c -> {
                        mm.addMessage(new MessageCivilian(true, c));
                        mm.addMessage(new MessageCivilian(false, c));
                    });
        }

        this.searchedEntityIDs.add(this.agentInfo.getPosition());
        this.priorityBuildings = this.removeNotBrokenBuilding(this.priorityBuildings);
        this.cluster = this.removeNotBrokenBuilding(this.cluster);

        if (this.cluster.isEmpty()) {
            // System.out.println("cluster is empty");
            this.searchedClusterIndex.add(this.clusterIndex);
            // System.out.println("searchedClusterIndex: " + this.searchedClusterIndex);
            this.clusterIndex = this.getNextClusterIndex();
            this.cluster = this
                    .removeAlreadyReachedBuilding(this.removeNotBrokenBuilding(this.getNextCluster(this.clusterIndex)));
        }

        return this;
    }

    @Override
    public Search calc() {
        if (this.result != null && !this.isReached(this.result)) {
            return this;
        }
        this.result = this.gatNextTarget();
        return this;
    }

    private EntityID gatNextTarget() {
        if (!this.priorityBuildings.isEmpty()) {
            EntityID target = this.getNearestTarget(this.priorityBuildings);
            this.priorityBuildings.remove(target);
            return target;
        }
        if (!this.cluster.isEmpty()) {
            EntityID target = this.getNearestTarget(this.cluster);
            this.cluster.remove(target);
            return target;
        }
        // System.out.println("AITFBSearch: No target");
        return null;
    }

    private Boolean isReached(EntityID target) {
        return this.agentInfo.getPosition().equals(target) || this.searchedEntityIDs.contains(target);
    }

    private int getNextClusterIndex() {
        int len = this.clustering.getClusterNumber();
        // nearestClusterIndex
        double nearest = Double.MAX_VALUE;
        int nearestIndex = -1;
        for (int i = 0; i < len; i++) {
            if (this.searchedClusterIndex.contains(i))
                continue;
            Collection<StandardEntity> cluster = this.clustering.getClusterEntities(i);
            double distance = 0;
            for (StandardEntity entity : cluster) {
                distance += this.worldInfo.getDistance(entity.getID(), this.agentInfo.getPosition());
            }
            if (distance < nearest) {
                nearest = distance;
                nearestIndex = i;
            }
        }
        return nearestIndex;
    }

    private Set<EntityID> getNextCluster(int clusterIndex) {
        return this.clustering.getClusterEntityIDs(clusterIndex).stream()
                .map(e -> this.worldInfo.getEntity(e))
                .filter(e -> e instanceof Building) // 建物のみ
                .filter(e -> !(e instanceof Refuge)) // 避難所を除外
                .map(StandardEntity::getID)
                .collect(Collectors.toSet());
    }

    private EntityID getNearestTarget(Set<EntityID> targets) {
        EntityID nearestID = null;
        double min_dis = Double.MAX_VALUE;
        for (EntityID En : targets) {
            double distance = this.getPathDistance(En);
            if (distance < min_dis) {
                nearestID = En;
                min_dis = distance;
            }
        }
        return nearestID;
    }

    private double getPathDistance(EntityID id) {
        EntityID agentPosition = this.agentInfo.getPosition();
        List<EntityID> ids = this.pathPlanning.getResult(agentPosition, id);
        double distance = 0;
        for (int i = 0; i < ids.size() - 1; i++) {
            distance += this.worldInfo.getDistance(ids.get(i), ids.get(i + 1));
        }
        return distance / 30000;
    }

    private Set<EntityID> removeNotBrokenBuilding(Set<EntityID> targets) {
        for (EntityID id : this.worldInfo.getChanged().getChangedEntities()) {
            StandardEntity entity = this.worldInfo.getEntity(id);
            if (entity instanceof Building && !(entity instanceof Refuge)) {
                if (targets.contains(id)) {
                    Building building = (Building) entity;
                    if (building.isBrokennessDefined() && building.getBrokenness() == 0) {
                        targets.remove(id);
                        this.searchedEntityIDs.add(id);
                    }
                }
            }
        }
        return targets;
    }

    private Set<EntityID> removeAlreadyReachedBuilding(Set<EntityID> targets) {
        for (EntityID id : this.searchedEntityIDs) {
            targets.remove(id);
        }
        return targets;
    }
}
