package autumn_2023.module.algorithm;

import adf.core.agent.info.*;
import adf.core.component.module.algorithm.Clustering;
import adf.core.component.module.algorithm.StaticClustering;
import adf.core.agent.module.ModuleManager;
import adf.core.agent.develop.DevelopData;
import adf.core.agent.precompute.PrecomputeData;
import rescuecore2.worldmodel.EntityID;
import rescuecore2.standard.entities.*;
import static rescuecore2.standard.entities.StandardEntityURN.*;
import java.util.*;
import static java.util.stream.Collectors.*;
import static java.util.Comparator.*;

import adf.core.debug.DefaultLogger;
import org.apache.log4j.Logger;

public class KmeansPPClustering extends StaticClustering
{
    // エージェントとクラスタの結びつけを保存
    private Map<EntityID, Integer> assignment = new HashMap<>();
    // グループ分けとグループの保存
    private KmeansPP clusterer;
    // グループ数
    private int n = 0;
    // エージェントの種類
    private StandardEntityURN urn;

    private Logger logger;

    // K-means++ 繰り返し実行回数の定義
    private final static int REP_PRECOMPUTE = 20;

    private final static int REP_PREPARE= 20;

    // 事前計算の結果保存用キーの定義
    // 事前計算の結果はPrecomputeDataクラスを通して保存します．
    // 文字列によるキーと値のペアで保存できます．
    private final static String MODULE_NAME =
    "autumn_2023.module.algorithm.KmeansPPClustering";
    private final static String PD_CLUSTER_N = MODULE_NAME + ".n";
    private final static String PD_CLUSTER_M = MODULE_NAME + ".m";
    private final static String PD_CLUSTER_A = MODULE_NAME + ".a";

    public KmeansPPClustering(
        AgentInfo ai, WorldInfo wi, ScenarioInfo si,
        ModuleManager mm, DevelopData dd)
    {
        super(ai, wi, si, mm, dd);
        this.urn = this.agentInfo.me().getStandardURN();
        this.logger = DefaultLogger.getLogger(agentInfo.me());
    }

    // 保存用キーにエージェントの種類を区別するための接尾辞を足すメソッド
    private String addSuffixToKey(String path)
    {
        return path + "." + this.urn;
    }

    // 保存用キーにエージェントの種類と要素番号を区別するための接尾辞を足すメソッド
    private String addSuffixToKey(String path, int i)
    {
        return this.addSuffixToKey(path) + "." + i;
    }

    // 事前計算 - 計算実行 & 結果保存部
    @Override
    public Clustering precompute(PrecomputeData pd)
    {
        super.precompute(pd);
        // 重複した処理の実行を回避
        if (this.getCountPrecompute() > 1) return this;

        this.initN();
        this.initClusterer();
        // K-means++を実行します．
        // 引数は繰り返し実行回数(K-means++アルゴリズムを参照して下さい)を表します．
        this.clusterer.execute(REP_PRECOMPUTE);
        this.pair();

        // グループ数を保存します．
        // 保存時の引数は (キー, 値)です．
        pd.setInteger(this.addSuffixToKey(PD_CLUSTER_N), this.n);
        for (EntityID agent : this.assignment.keySet())
        {
            int i = this.assignment.get(agent);
            // i番目のクラスタの全ての要素を取得します．
            Collection<EntityID> cluster = this.clusterer.getClusterMembers(i);

            // i番目のクラスタの要素を保存します．
            pd.setEntityIDList(
                this.addSuffixToKey(PD_CLUSTER_M, i),
                new ArrayList<>(cluster));
            // i番目のクラスタと結び付けられたエージェントを保存します．
            pd.setEntityID(this.addSuffixToKey(PD_CLUSTER_A, i), agent);
        }

        //debug
        for(int i = 1; i < this.clusterer.getClusterNumber() + 1; i++) {
            Collection<EntityID> myClusterEntities = this.clusterer.getClusterMembers(i-1);
            logger.debug("cluster="+i+", EntityID="+myClusterEntities.toString());
        }

        return this;
    }

    // 事前計算 - 計算結果読み込み部
    @Override
    public Clustering resume(PrecomputeData pd)
    {
        super.resume(pd);
        // 重複した処理の実行を回避
        if (this.getCountResume() > 1) return this;
        
        // 事前計算の結果の読み込みもPrecomputeDataクラスを通しておこなえます．
        // グループ数を読み込みます．
        this.n = pd.getInteger(this.addSuffixToKey(PD_CLUSTER_N));
        List<Collection<EntityID>> clusters = new ArrayList<>(this.n);
        for (int i=0; i<this.n; ++i)
        {
            // i番目のクラスタの要素を読み込みます．
            List<EntityID> cluster =
                pd.getEntityIDList(this.addSuffixToKey(PD_CLUSTER_M, i));
            // i番目のクラスタと結び付けられたエージェントを読み込みます．
            EntityID agent =
                pd.getEntityID(this.addSuffixToKey(PD_CLUSTER_A, i));

            clusters.add(cluster);
            // エージェントとクラスタの結び付きを復元します．
            this.assignment.put(agent, i);
        }

        // グループ数と各クラスタの要素を表したListでKmeansPPを復元します．
        this.clusterer = new KmeansPP(this.n, clusters);
        return this;
    }

    // 事前計算がおこなわれない場合の処理
    // おそらくprepareの間違い
    @Override
    public Clustering preparate()
    {
        super.preparate();
        // 重複した処理の実行を回避
        if (this.getCountPreparate() > 1) return this;

        this.initN();
        this.initClusterer();
        this.clusterer.execute(REP_PREPARE);
        this.pair();
        return this;
    }

    @Override
    public Clustering calc()
    {
        return this;
    }

    // 他のモジュールがクラスタ数を取得する際に使います
    @Override
    public int getClusterNumber()
    {
        return this.n;
    }

    // 他のモジュールがentityに関連したクラスタの番号を取得する際に使います
    @Override
    public int getClusterIndex(StandardEntity entity)
    {
        return this.getClusterIndex(entity.getID());
    }

    // 他のモジュールがidに関連したクラスタの番号を取得する際に使います
    @Override
    public int getClusterIndex(EntityID id)
    {
        if (!this.assignment.containsKey(id)) return -1;
        return this.assignment.get(id);
    }

    // 他のモジュールがi番目のクラスタ要素をStandardEntityで取得する際に使います
    @Override
    public Collection<StandardEntity> getClusterEntities(int i)
    {
        if (i < 0 || i >= this.n) return null;

        Collection<EntityID> ids = this.getClusterEntityIDs(i);
        Collection<StandardEntity> ret = new ArrayList<>(ids.size());
        for (EntityID id : ids)
        {
            ret.add(this.worldInfo.getEntity(id));
        }
        return ret;
    }

    // 他のモジュールがi番目のクラスタ要素をEntityIDで取得する際に使います
    @Override
    public Collection<EntityID> getClusterEntityIDs(int i)
    {
        if (i < 0 || i >= this.n) return null;
        return this.clusterer.getClusterMembers(i);
    }

    private void initN()
    {
        switch (this.urn)
        {
            // 消防隊
            case FIRE_BRIGADE:
                // グループ数 = シミュレーション全体の消防隊数
                this.n = this.scenarioInfo.getScenarioAgentsFb();
                break;
            // 土木隊
            case POLICE_FORCE:
                // グループ数 = シミュレーション全体の土木隊数
                this.n = this.scenarioInfo.getScenarioAgentsPf();
                break;
            // 救急隊
            case AMBULANCE_TEAM:
                // グループ数 = シミュレーション全体の救急隊数
                this.n = this.scenarioInfo.getScenarioAgentsAt();
                break;
            default:
                this.n = 0;
        }
    }
    private void initClusterer()
    {
        // 次のオブジェクトを全て取得します．
        // 道路/消火栓
        // 建物/ガソリンスタンド
        // 避難所
        // 土木隊司令所/消防隊司令所/救急隊司令所
        List<StandardEntity> entities = new ArrayList<>(
            this.worldInfo.getEntitiesOfType(
                ROAD, HYDRANT,
                BUILDING, GAS_STATION,
                REFUGE,
                POLICE_OFFICE, FIRE_STATION, AMBULANCE_CENTRE));
        // リストをIDでソートします．
        entities.sort(comparing(e -> e.getID().getValue()));

        // データをID/X座標/Y座標の配列に整形します．
        int size = entities.size();
        EntityID[] is = new EntityID[size];
        double[] xs = new double[size];
        double[] ys = new double[size];

        for (int i=0; i<size; ++i)
        {
            // StandardEntityクラスでは座標を得ることができません．
            // 地図上のオブジェクトを表すAreaクラスにダウンキャストします．
            // (エージェントの場合はHumanクラス)
            Area area = (Area)entities.get(i);
            is[i] = area.getID();
            xs[i] = area.getX();
            ys[i] = area.getY();
        }

        // KmeansPPの初期化をおこないます．
        this.clusterer = new KmeansPP(is, xs, ys, this.n);
    }

    private void pair()
    {
        // 同種類のエージェントを全て取得します．
        List<StandardEntity> agents = new ArrayList<>(
            this.worldInfo.getEntitiesOfType(this.urn));
        //リストをIDでソートします．
        agents.sort(comparing(e -> e.getID().getValue()));

        // エージェントとクラスタの間の距離を二次元配列に整形します．
        int[][] costs = new int[this.n][this.n];
        for (int row=0; row<this.n; ++row)
        {
            Human agent = (Human)agents.get(row);
            double x = agent.getX();
            double y = agent.getY();

            for (int col=0; col<this.n; ++col)
            {
                // col番目のクラスタのX/Y座標を取得します．
                double cx = this.clusterer.getClusterX(col);
                double cy = this.clusterer.getClusterY(col);
                // エージェントの座標とクラスタの座標間の距離を計算します．
                costs[row][col] = (int)Math.hypot(cx-x, cy-y);
            }
        }

        // Hungarianアルゴリズムに距離の二次元配列を渡して結果を得る．
        // 結果 row番目: row番目のエージェントが割り当てられたクラスタの番号
        int[] result = Hungarian.execute(costs);
        for (int row=0; row<n; ++row)
        {
            EntityID id = agents.get(row).getID();
            // エージェントのIDとクラスタの番号を結びつけて保存する．
            this.assignment.put(id, result[row]);
        }
    }
}
