package autumn_2023.module.complex;

import autumn_2023.module.algorithm.Hungarian;
import adf.core.agent.communication.MessageManager;
import adf.core.agent.communication.standard.bundle.MessageUtil;
import adf.core.agent.communication.standard.bundle.StandardMessageBundle;
import adf.core.agent.communication.standard.bundle.StandardMessagePriority;
import adf.core.agent.communication.standard.bundle.centralized.MessageReport;
import adf.core.agent.communication.standard.bundle.information.MessageBuilding;
import adf.core.agent.communication.standard.bundle.information.MessageCivilian;
import adf.core.agent.communication.standard.bundle.information.MessageFireBrigade;
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.agent.precompute.PrecomputeData;
import adf.core.component.communication.CommunicationMessage;
import adf.core.component.module.complex.FireTargetAllocator;

import static rescuecore2.standard.entities.StandardEntityURN.AMBULANCE_TEAM;
import static rescuecore2.standard.entities.StandardEntityURN.REFUGE;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.stream.Collector;
import java.util.stream.Collectors;

import rescuecore2.standard.entities.Building;
import rescuecore2.standard.entities.FireBrigade;
import rescuecore2.standard.entities.Human;
import rescuecore2.standard.entities.StandardEntity;
import rescuecore2.standard.entities.StandardEntityURN;
import rescuecore2.worldmodel.EntityID;

public class CentralizedControlFBAllocator extends FireTargetAllocator {

  private Map<EntityID, FireBrigadeInfo> fireBrigadeInfoMap;
  private Set<EntityID> taskSet;
  private Map<EntityID, EntityID> allocateTaskMap;
  private Set<EntityID> finishEntityIDs;

  public CentralizedControlFBAllocator(AgentInfo ai, WorldInfo wi, ScenarioInfo si, ModuleManager moduleManager, DevelopData developData) {
    super(ai, wi, si, moduleManager, developData);
    this.fireBrigadeInfoMap = new HashMap<>();
    this.taskSet = new HashSet<>();
    this.allocateTaskMap = new HashMap<>();
    this.finishEntityIDs = new HashSet<>();
  }


  @Override
  public FireTargetAllocator resume(PrecomputeData precomputeData) {
    super.resume(precomputeData);
    if (this.getCountResume() >= 2) {
      return this;
    }
    for (EntityID id : this.worldInfo
        .getEntityIDsOfType(StandardEntityURN.FIRE_BRIGADE)) {
      this.fireBrigadeInfoMap.put(id, new FireBrigadeInfo(id));
    }
    return this;
  }


  @Override
  public FireTargetAllocator preparate() {
    super.preparate();
    if (this.getCountPrecompute() >= 2) {
      return this;
    }
    for (EntityID id : this.worldInfo
        .getEntityIDsOfType(StandardEntityURN.FIRE_BRIGADE)) {
      this.fireBrigadeInfoMap.put(id, new FireBrigadeInfo(id));
    }
    return this;
  }

  /**
   * calcメソッドで算出した司令内容を返すためのメソッド
   */
  @Override
  public Map<EntityID, EntityID> getResult() {
    return this.convert(this.fireBrigadeInfoMap);
  }

  /**
   * 消防隊情報を
   *  key:命令対象のエージェントのID
   *  value:やるべきタスク(市民)のID
   * のMapに変換するメソッド
   * @param map 消防隊IDと消防隊情報に関連するMap
   * @return 司令情報となるMap
   */
  private Map<EntityID, EntityID> convert(Map<EntityID, FireBrigadeInfo> map) {
    Map<EntityID, EntityID> result = new HashMap<>();
    for (EntityID id : map.keySet()) {
      FireBrigadeInfo info = map.get(id);
      if (info != null && info.target != null) {
        result.put(id, info.target);
      }
    }
    return result;
  }

  /**
   * 分散エージェントへの司令内容を算出するためのメソッド
   * 
   * ハンガリアンアルゴリズムを利用して消防隊に担当させるタスクを決定する
   */
  @Override
  public FireTargetAllocator calc() {
    if(this.agentInfo.getTime() < scenarioInfo.getKernelAgentsIgnoreuntil()) return this;

    if(this.fireBrigadeInfoMap.keySet().size()==0||this.taskSet.size()==0) return this;

    int row = this.fireBrigadeInfoMap.keySet().size();
    int col = this.taskSet.size();
    boolean reverse = false;
    List<EntityID> rowlist = new ArrayList<>(this.fireBrigadeInfoMap.keySet());
    List<EntityID> collist = new ArrayList<>(this.taskSet);

    if(row > col){
      reverse = true;
      int tmp = row;
      row = col;
      col = tmp;
      List<EntityID> tmplist = rowlist;
      rowlist = collist;
      collist = tmplist;
    }

    int[][] costs = new int[row][col];
    for(int i = 0; i<row; i++){
      EntityID rowid = rowlist.get(i);
      StandardEntity rowposi = this.worldInfo.getPosition(rowid);
      for(int j = 0; j<col; j++){
        EntityID colid = collist.get(j);
        int distance = this.worldInfo.getDistance(rowposi, this.worldInfo.getPosition(colid));
        costs[i][j] = distance;
      }
    }

    int[] allocate = Hungarian.execute(costs);
    if(reverse){
      int tmp = row;
      row = col;
      col = tmp;
      List<EntityID> tmplist = rowlist;
      rowlist = collist;
      collist = tmplist;
      int[] reverse_list = new int[row];
      for(int i = 0; i<row; i++) reverse_list[i] = -1;
      for(int i = 0; i<col; i++){
        reverse_list[allocate[i]] = i;
      }
      allocate = reverse_list;
    }

    for(int i = 0; i<row; i++){
      if(allocate[i] == -1) continue;
      this.fireBrigadeInfoMap.get(rowlist.get(i)).target = collist.get(allocate[i]);
    }
    //this.printAllocate(allocate, rowlist, collist, row);
    return this;
  }

  /**
   * 新しく受信したメッセージから情報を取得するためのメソッド
   * 
   * 埋没市民Setを更新
   * 消防隊情報Setを更新
   * 消防隊から司令完了を受けたときタスク完了リストを更新
   */
  @Override
  public FireTargetAllocator updateInfo(MessageManager messageManager) {
    super.updateInfo(messageManager);

    Set<EntityID> recievetasks = messageManager.getReceivedMessageList(MessageCivilian.class).stream()
    .map(e -> {
      MessageCivilian mesciv = (MessageCivilian) e;
      MessageUtil.reflectMessage(this.worldInfo, mesciv);
      return(mesciv);
    })
    .map(MessageCivilian::getAgentID)
    .filter(e -> this.needRescue(this.worldInfo.getEntity(e)))
    .filter(e -> !(finishEntityIDs.contains(e)))
    .collect(Collectors.toSet());
    
    this.taskSet.addAll(recievetasks);

    // 消防隊の情報をリセット
    for(EntityID fbid : this.fireBrigadeInfoMap.keySet()){
      FireBrigadeInfo fbi = this.fireBrigadeInfoMap.get(fbid);
      fbi.resetInfo();
    }
    List<EntityID> fblist = new ArrayList<>();
    List<CommunicationMessage> mesfblist = messageManager.getReceivedMessageList(MessageFireBrigade.class);
    for(CommunicationMessage mes : mesfblist){
      MessageFireBrigade mesfb = (MessageFireBrigade) mes;
      MessageUtil.reflectMessage(this.worldInfo, mesfb);
      EntityID fbID = mesfb.getAgentID();
      fblist.add(fbID);
      FireBrigadeInfo fbi = this.fireBrigadeInfoMap.get(fbID);
      fbi.setInfo(mesfb.getTargetID(), mesfb.getPosition(), mesfb.getBuriedness(),
          (mesfb.getAction() == MessageFireBrigade.ACTION_RESCUE) ? false : true, this.agentInfo.getTime());
    }

    List<EntityID> reportlist = messageManager.getReceivedMessageList(MessageReport.class).stream()
    .map(MessageReport.class::cast)
    .filter(MessageReport::isDone)
    .map(MessageReport::getFromID)
    .collect(Collectors.toList());
    
    this.finishEntityIDs.addAll(reportlist);
    this.taskSet.removeAll(reportlist);

    return this;
  }

  /**
   * ある市民について消防隊による救助が必要かを判定するメソッド
   * 
   * @param entity エンティティ
   * @return true : 救助必要 false : 救助不必要
   */
  private boolean needRescue(StandardEntity entity){
    if (entity == null)
      return false;
    if (!(entity instanceof Human))
      return false;

    Human target = (Human) entity;
    if (!target.isHPDefined() || target.getHP() == 0)
      return false;
    if (!target.isPositionDefined())
      return false;
    if (!target.isBuriednessDefined() || target.getBuriedness() == 0)
      return false;

    StandardEntity position = worldInfo.getPosition(target);
    if (position == null)
      return false;

    StandardEntityURN positionURN = position.getStandardURN();
    if (positionURN == REFUGE || positionURN == AMBULANCE_TEAM)
      return false;

    return true;
  }

  /**
   * ハンガリアンアルゴリズムの割当結果をみるためのメソッド
   * @param allocate 消防隊に対応するタスク割当結果配列
   * @param rowlist 消防隊IDリスト
   * @param collist 埋没市民IDリスト
   * @param row 消防隊の数
   */
  private void printAllocate(int[] allocate, List<EntityID> rowlist, List<EntityID> collist, int row){
    String str = " allocate";
    for(int i = 0; i<row; i++){
      if(allocate[i] == -1) continue;
      str += "\n["+i+"]:"+allocate[i]+" fbID:"+rowlist.get(i)+" tagID:"+collist.get(allocate[i]);
    }
    System.out.println("[FBCenter] time:"+this.agentInfo.getTime()+" id:"+this.agentInfo.getID()+str);
  }

  /**
   * 消防隊情報クラス
   */
  private class FireBrigadeInfo {
    
    EntityID agentID;
    EntityID prevtarget;
    EntityID position;
    int buriedness;
    EntityID target;
    boolean canNewAction;
    int commandTime;

    FireBrigadeInfo(EntityID id) {
      agentID = id;
      prevtarget = null;
      position = null;
      buriedness = -1;
      target = null;
      canNewAction = true;
      commandTime = -1;
    }
    /**
     * 消防隊に関する情報をセットするメソッド
     */
    public void setInfo(EntityID prevtarget, EntityID position, int buriedness, boolean canNewAction, int commandTime){
      this.prevtarget = prevtarget;
      this.position = position;
      this.buriedness = buriedness;
      this.canNewAction = canNewAction;
      this.commandTime = commandTime;
    }
    /**
     * 消防隊に関する情報をリセットするメソッド
     */
    public void resetInfo(){
      prevtarget = null;
      position = null;
      buriedness = -1;
      target = null;
      canNewAction = true;
      commandTime = -1;
    }
  }
}