package edu.berkeley.nlp.starcraft.overmind;

import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;

import org.bwapi.proxy.model.Game;
import org.bwapi.proxy.model.Order;
import org.bwapi.proxy.model.Player;
import org.bwapi.proxy.model.Position;
import org.bwapi.proxy.model.Race;
import org.bwapi.proxy.model.TilePosition;
import org.bwapi.proxy.model.Unit;
import org.bwapi.proxy.model.UnitType;

import edu.berkeley.nlp.starcraft.OvermindUnit;
import edu.berkeley.nlp.starcraft.resources.ResourceAccount;
import edu.berkeley.nlp.starcraft.resources.ResourceBank;
import edu.berkeley.nlp.starcraft.resources.ResourceInventory;
import edu.berkeley.nlp.starcraft.resources.ResourceType;
import edu.berkeley.nlp.starcraft.resources.ResourceAccount.AllocationType;
import edu.berkeley.nlp.starcraft.resources.ResourceAccount.ReleaseType;
import edu.berkeley.nlp.starcraft.util.Filter;
import edu.berkeley.nlp.starcraft.util.Log;
import edu.berkeley.nlp.util.PriorityQueue;
import edu.berkeley.nlp.starcraft.util.UnitUtils;
import edu.berkeley.nlp.starcraft.util.Utils;
import edu.berkeley.nlp.util.Counter;

public class OvermindBank extends ResourceBank implements UnitEventHandler, Underling {
  static Log logger = Log.getLog("resource.OvermindBank");

  private Map<Unit, AnOvermindUnit> unitMap = new HashMap<Unit, AnOvermindUnit>();
  int clicks = 0;
  int recentClicks = 0;

  void click() {
    clicks++;
    recentClicks++;
  }

  // Stores resources whilst a unit has begun a build move; deducted once the
  // building is built or canceled.
  ResourceAccount buildAccount;
  ResourceAccount unallocatedAccount;
  // Watch these units to see if &when they spend anything
  Map<Unit, UnitType> ongoingBuildAttempts = new HashMap<Unit, UnitType>();
  // In case the builder gets reassigned to another action
  Map<Unit, ResourceAccount> ongoingBuildAccounts = new HashMap<Unit, ResourceAccount>();
  Map<Unit, TilePosition> ongoingBuildPositions = new HashMap<Unit, TilePosition>();

  // guys who might return a resource
  private Set<Unit> previousGasCarriers = new HashSet<Unit>();
  private Set<Unit> previousMinCarriers = new HashSet<Unit>();

  private static final int MINERALS_PER_TRIP = 8;
  private static final int GAS_PER_TRIP = 8;

  // Supply spent in the previous frame, since this seems to be delayed one
  // tick.
  int supplySpentPrevFrame;
  // FIXME: why isn't this next line indicative of a bug -- dlwh
  // After training a unit it deducts the mineral cost TWICE, then adds it back.
  // Minerals can go negative.
  int mineralsSpentPrevFrame;
  int mineralsSpentPrevPrevFrame;
  // TODO: probably the same is true for gas; WARNING: this might be a bug in
  // BWAPI

  Race race;

  public OvermindBank() {
    super("OvermindBank");
    buildAccount = newAccount(null, "BuildAccount");
    unallocatedAccount = newAccount(null, "UnallocatedAccount", getInventory(), getUnits());
    race = Game.getInstance().self().getRace();
  }

  @Override
  public int getGas() {
    // FIXME: need to predict mineral/gas/supply usage and compare with the
    // game's state.
    return Game.getInstance().self().gas();
  }

  @Override
  public int getMinerals() {
    return Game.getInstance().self().minerals() + mineralsSpentPrevPrevFrame;
  }

  @Override
  public int getSupply() {
    Player self = Game.getInstance().self();
    return Math.max(0, (self.supplyTotal() - self.supplyUsed()));
  }

  @Override
  public Set<OvermindUnit> getUnits() {
    return Collections.<OvermindUnit> unmodifiableSet(new HashSet<AnOvermindUnit>(unitMap.values()));
  }

  @Override
  protected void onAccountClosed(ResourceAccount acct) {
    // FIXME implement

  }

  @Override
  protected boolean allowAllocation(ResourceAccount acct, ResourceInventory request, AllocationType type) {
    return true;
  }

  @SuppressWarnings("unchecked")
  @Override
  protected boolean doTransfer(ResourceAccount from, ResourceAccount to, ResourceInventory theRequest,
      List<OvermindUnit> suggestedUnits, ReleaseType fromType, AllocationType toType) {
    Set<AnOvermindUnit> theUnits = new LinkedHashSet<AnOvermindUnit>();

    // make sure we have a clean copy.
    ResourceInventory request = theRequest.clone();

    ResourceInventory inv = request.clone();
    inv.put(ResourceType.MINERAL, 0);
    inv.put(ResourceType.GAS, 0);
    inv.put(ResourceType.SUPPLY, 0);

    for (OvermindUnit unit : suggestedUnits) {
      if (unit instanceof AnOvermindUnit && from.units().contains(unit)) {
        if (inv.getCount(unit.getType()) > 0) {
          boolean changed = theUnits.add((AnOvermindUnit) unit);
          if (changed) {
            inv.incrementCount(ResourceType.unit(unit.getType()), -1);
          }
        }
      }
    }

    if(!inv.isEmpty()) {
      // we assign a simple score to each unit based on its position in each list
      // lowest score wins. Units not listed in a list are assumed to have score
      // -20 or so
      // TODO: make the system return two counters, or something more sensible.
      Counter<OvermindUnit> scores = new Counter<OvermindUnit>();
      List<OvermindUnit> toSuggest = to.getAccountHolder() != null ? to.getAccountHolder()
          .suggestUnitsForAllocation(inv, from.units()) : Collections.<OvermindUnit> emptyList();
      if(toSuggest != null) for(int i = 0; i < toSuggest.size(); ++i) {
        scores.setCount(toSuggest.get(i), toSuggest.size() - i + 20);
      }
      List<OvermindUnit> fromSuggest = from.getAccountHolder() != null ? from.getAccountHolder()
          .suggestUnitsForRelease(inv) : Collections.<OvermindUnit> emptyList();
      if(fromSuggest != null) for (int i = 0; i < fromSuggest.size(); ++i) {
        scores.incrementCount(fromSuggest.get(i), fromSuggest.size() - i + 20);
      }
        
      PriorityQueue<OvermindUnit> s = scores.asPriorityQueue();
      while(s.hasNext()) {
        OvermindUnit ou = s.next();
        if(ou == null) continue;
        AnOvermindUnit unit = (AnOvermindUnit) ou;
        logger.trace("Checking unit " + unit.getType() + " against inv " + inv.getCount(unit.getType()) + " "
            + inv);
        if (inv.getCount(unit.getType()) > 0) {
          boolean changed = theUnits.add(unit);
          if (changed) {
            inv.incrementCount(ResourceType.unit(unit.getType()), -1);
          }
        }
      }
    }
    
    if (!inv.isEmpty()) {
      for (OvermindUnit ou : from.units()) {
        AnOvermindUnit unit = (AnOvermindUnit) ou;
        logger.trace("Checking unit " + unit.getType() + " against inv " + inv.getCount(unit.getType()) + " "
            + inv);
        if (inv.getCount(unit.getType()) > 0) {
          boolean changed = theUnits.add(unit);
          if (changed) {
            inv.incrementCount(ResourceType.unit(unit.getType()), -1);
          }
        }

        logger.trace("inv: " + inv);
      }
    }

    if (!inv.isEmpty()) {
      logger.error("Couldn't satisfy request:" + inv + " from account: " + from.units());
      return false;
    }

    List<OvermindUnit> toList = new ArrayList<OvermindUnit>(theUnits.size());
    for (AnOvermindUnit ou : theUnits) {
      ou.invalidate();
      AnOvermindUnit newOu = new AnOvermindUnit(ou.unit, to, this);
      unitMap.put(ou.unit, newOu);
      toList.add(newOu);
    }
    // yay erasure ugly code. but i'm lazy. -- dlwh
    to.allocate(request, toList, toType);
    from.release(request, (Set) theUnits, fromType);
    return true;
  }

  public void trackBuildOrder(Unit builder, UnitType building, TilePosition position,
      ResourceAccount account, ResourceInventory inv) {
    logger.debug("Tracking build order by " + builder + " of " + building + " at position " + position
        + " with account " + account);

    transfer(account, buildAccount, inv, ResourceAccount.ReleaseType.ESCROW, AllocationType.ESCROW);
    ongoingBuildAttempts.put(builder, building);
    ongoingBuildAccounts.put(builder, account);
    ongoingBuildPositions.put(builder, position);
  }

  public void trackBuildAddonOrder(Unit unit, UnitType type, ResourceAccount account) {
    // TODO Auto-generated method stub

  }

  @Override
  public void notifyUnitCreated(Unit unit) {
  //	logger.info("Creation!");
    ResourceAccount depositAccount = determineLikelyBuilder(unit);

    AnOvermindUnit ou = new AnOvermindUnit(unit, depositAccount, this);
    unitMap.put(unit, ou);
    List<OvermindUnit> single = Collections.<OvermindUnit> singletonList(ou);
    depositAccount.allocate(new ResourceInventory(0, 0, ou.getType().supplyProvided(), single), single,
        AllocationType.UNIT_CREATION);

  }

  private ResourceAccount determineLikelyBuilder(Unit unit) {
    ResourceAccount depositAccount = getUnallocatedAccount();
    final UnitType type = unit.getType();
    // If the new unit is a building under construction, remove resources from
    // the buildBank
    // and remove the building unit from the list of ongoingBuildingAttempts.
    if (type.isAddon()) {
      AnOvermindUnit builder = null;
      for (Map.Entry<Unit, AnOvermindUnit> e : unitMap.entrySet()) {
        if (e.getKey().getAddon() != null && e.getKey().getAddon().equals(unit)) {
          builder = e.getValue();
          break;
        }
      }

      if (builder == null) {
        logger.warn("No addon builder found!" + unit);
      } else {
        ResourceAccount builderAccount = builder.account;
        if (builderAccount != null && !builderAccount.isClosed()) {
          depositAccount = builderAccount;
          logger.debug("New addon " + unit + " going to action " + builderAccount.getAccountHolder());
        }
      }
    } else if (type.isBuilding()) {
      // Try to find who just built this building. This is one of the units who
      // were trying
      // to build at a unit of this type at a location closest to this.
      // FIXME: can this fail?
      List<Unit> candidates = Utils.findAll(ongoingBuildAttempts.keySet(), new Filter<Unit>() {
        @Override
        public boolean test(Unit builder) {
          return ongoingBuildAttempts.get(builder).equals(type);
        }
      });
      double minDistance = Double.POSITIVE_INFINITY;
      Unit builder = null;
      for (Unit candidate : candidates) {
        double distance = unit.getDistance(new Position(ongoingBuildPositions.get(candidate)));
        if (distance < minDistance) {
          minDistance = distance;
          builder = candidate;
        }
      }

      if (builder == null) {
        AnOvermindUnit ou = unitMap.get(unit);
        if (ou != null && ou.account != null) {
          // probably a morph:
          depositAccount = ou.account;
        } else {
          logger.warn("No builder found!  ongoingBuildAttempts=" + ongoingBuildAttempts);
        }
      } else {

        ResourceAccount builderAccount = ongoingBuildAccounts.remove(builder);
        ongoingBuildAttempts.remove(builder);
        ongoingBuildPositions.remove(builder);

        // Everything built is paid for, so must transfer out of buildAccount
        // Transfer first otherwise we overcount the bank free contents...
        buildAccount.release(new ResourceInventory(type.mineralPrice(), type.gasPrice(), type
            .supplyRequired()), ReleaseType.SPENT);
        if (builderAccount != null && !builderAccount.isClosed()) {
          depositAccount = builderAccount;
          
        } 
      } 
    } else if (type.equals(UnitType.ZERG_LARVA)) {
      // do nothing. Larva should go to default account
    } else {
      AnOvermindUnit ou = unitMap.get(unit);
      if (ou != null && ou.account != null) {
        // probably a morph:
        depositAccount = ou.account;
      } else {
        // Assign the unit to the action who owns the unit who probably built
        // this,
        // if we can find any, otherwise to the overmind.
        // TODO: can we detect whether a building has just finished training?
        final UnitType createrType = unit.getType().whatBuilds().getKey();
        List<Unit> candidates = Utils.findAll(Game.getInstance().self().getUnits(), new Filter<Unit>() {
          @Override
          public boolean test(Unit creator) {
            return creator.getType().equals(createrType);
          }
        });

        if (candidates.size() > 0) {
          Unit creator = UnitUtils.getClosest(unit, candidates);
          if (unitMap.containsKey(creator)) {
            ResourceAccount account = unitMap.get(creator).account;
            if (account != null) {
              depositAccount = account;
            }
          }
          
        }
      }
    }
   // logger.info("New unit " + unit + " created for " + depositAccount);
    return depositAccount;
  }

  @Override
  public void notifyUnitDestroyed(Unit unit) {

    // FIXME: figure out where to get the supply from and how to handle having
    // supply deficit
    UnitType type = unit.getPreviousUnitData().getType();
	if (type.supplyProvided() > 0) {
      logger.error("Supply providing unit died, we need to reclaim its supply somehow.");
    }

    logger.warn("lost a " + type.getName() + " " + type.supplyRequired());

    AnOvermindUnit ou = unitMap.get(unit);
    ResourceAccount account = unallocatedAccount;
    if (ou != null && ou.account != null) {
      account = ou.account;
    }
	account.release(new ResourceInventory(type), Collections.<OvermindUnit> singleton(ou),
        ResourceAccount.ReleaseType.UNIT_DESTRUCTION);
    unallocatedAccount.allocate(new ResourceInventory(0, 0, type.supplyRequired()
        - type.supplyProvided()), ResourceAccount.AllocationType.RESOURCE_CREATION);
    unitMap.remove(unit);

    // TODO: is this needed?
    // TODO: does this mean we should refund stuff?
    ongoingBuildAccounts.remove(unit);
    ongoingBuildAttempts.remove(unit);
    ongoingBuildPositions.remove(unit);
  }

  @Override
  public void notifyUnitMorphed(Unit unit) {
    // FIXME vespene geyser code

    // logger.info("Morph! " + unit.getType().getName() + " " +
    // unit.getType().whatBuilds());

    if (unit.getType().equals(UnitType.PROTOSS_ASSIMILATOR) || unit.getType().equals(UnitType.ZERG_EXTRACTOR)
        || unit.getType().equals(UnitType.TERRAN_REFINERY)) {
      logger.warn("Morphed unit " + unit.getType().getName() + " assuming created vespene extractor.");

      notifyUnitCreated(unit);
      //
    } else if (unit.getType().equals(UnitType.RESOURCE_VESPENE_GEYSER)) {
      logger.warn("Morphed unit " + unit.getType().getName() + " assuming destroyed vespene extractor.");
      notifyUnitDestroyed(unit);
    } else if (unit.getPlayer().equals(Game.getInstance().self())) {
      ResourceAccount builder = determineLikelyBuilder(unit);

      ResourceInventory inv = new ResourceInventory(0, 0, unit.getType().supplyProvided(), unit.getType());
      if (unit.getType().equals(UnitType.PROTOSS_ARCHON)
      		|| unit.getType().equals(UnitType.PROTOSS_DARK_ARCHON)
      		|| unit.getType().equals(UnitType.TERRAN_SIEGE_TANK_SIEGE_MODE)
      		|| unit.getType().equals(UnitType.TERRAN_SIEGE_TANK_TANK_MODE)
          || unit.getType().equals(UnitType.ZERG_EGG)
          || (unit.getType().isBuilding() && unit.getType().getRace().equals(Race.ZERG))) {
        inv.put(ResourceType.unit(unit.getType().whatBuilds().getKey()), -unit.getType().whatBuilds()
            .getValue());
      } else if(unit.getType().getRace().equals(Race.ZERG)) {
        inv.put(ResourceType.unit(UnitType.ZERG_EGG), -1);
      } else {
      	logger.info("Spurious morph?" + unit.getType().getName());
      	return;
      }
      AnOvermindUnit au = unitMap.get(unit);
      logger.debug("before: " + builder.inventory() + " " + builder.units());
      builder.allocate(inv, AllocationType.EXCESS);

      if (!unit.getType().equals(UnitType.ZERG_HATCHERY))
        getUnallocatedAccount().allocate(
            new ResourceInventory(0, 0, unit.getType().whatBuilds().getKey().supplyRequired()),
            AllocationType.RESOURCE_CREATION);

      logger.debug("after:" + builder.inventory() + " " + builder.units());
      if (builder.getAccountHolder() != null) {
        builder.getAccountHolder().notifyMorph(au);
      }
      // FIXME: this will not detect an extractor being destroyed, we don't get
      // a morphUnit notification.
    } else {
      // we get this for enemy units too. Don't throw.
      // throw new
      // RuntimeException("Received a morph message on a unit we're not tracking"
      // + unit);
    }

  }

  public void onFrame(int frameNum) {

    if (frameNum % 100 == 0) {
      logger.debug("Game contents: " + getMinerals() + "m " + getGas() + "g " + getSupply() + "s" + " units:"
          + Game.getInstance().self().getUnits());
      logger.debug("Build account contents: " + buildAccount.inventory());
      logger.debug("Avg. APM:" + (clicks * 24.0 / frameNum * 60.0));
      logger.debug("APM:" + (recentClicks * 24.0) / 100 * 60.0);
      recentClicks = 0;
    }

    Collection<Unit> orphans = new HashSet<Unit>(Game.getInstance().self().getUnits());
    for (OvermindUnit u : getUnits()) {
      orphans.remove(((AnOvermindUnit) u).unit);
    }
    if (orphans.size() > 0 && frameNum != 0) {
      logger.warn("Orphan units! " + orphans);
    }

    // debug output
    logger.trace("Game contents: " + getMinerals() + "m " + getGas() + "g " + getSupply() + "s" + " units:"
        + Game.getInstance().self().getUnits());
    logger.trace("Overmind unallocated inventory: " + getUnallocatedInventory());
    logger.trace("Overmind allocated inventory: " + getAllocatedInventory());
    for (ResourceAccount acct : this.getAccounts()) {
      logger.trace("   " + acct.getName() + " has " + acct.inventory());
    }
    logger.trace("min...prevFrame" + this.mineralsSpentPrevFrame + " ...prevprevF "
        + this.mineralsSpentPrevPrevFrame);
    logger.trace("Build account contents: " + buildAccount.inventory());

    // Notice any new units. This should happen only on the first frame, with
    // the initial units.
    // All other units are added by Overmind.onUnitCreate?
    Set<Unit> currentGas = new HashSet<Unit>();
    Set<Unit> currentMinerals = new HashSet<Unit>();
    for (Unit u : Game.getInstance().self().getUnits()) {
      if (!unitMap.containsKey(u)) {
        if (frameNum == 0)
          addInitialUnit(u); // FIXME
        else
          notifyUnitCreated(u);
      }

      // while we're at it, see if anyone is carrying gas/min that didn't used
      // to be, or
      // just returned something.
      if (u.getType().isWorker()) {
        if (u.getOrder().equals(Order.RETURN_GAS)) {
          currentGas.add(u);
        } else if (u.getOrder().equals(Order.RETURN_MINERALS)) {
          currentMinerals.add(u);
        } else if (!u.isCarryingGas() && !u.isCarryingMinerals() && previousGasCarriers.contains(u)) {
          // probably returned gas.
          AnOvermindUnit ou = unitMap.get(u);
          if (ou != null && ou.account != null)
            ou.account.allocate(new ResourceInventory(0, GAS_PER_TRIP, 0), AllocationType.RESOURCE_CREATION);
        } else if (!u.isCarryingGas() && !u.isCarryingMinerals() && previousMinCarriers.contains(u)) {
          // probably returned minerals.
          AnOvermindUnit ou = unitMap.get(u);
          if (ou != null && ou.account != null)
            ou.account.allocate(new ResourceInventory(MINERALS_PER_TRIP, 0, 0),
                AllocationType.RESOURCE_CREATION);
        }
      }

    }

    previousMinCarriers = currentMinerals;
    previousGasCarriers = currentGas;

    if (frameNum % 20 == 0) {
      ResourceInventory gameInv = getInventory();
      logger.debug("Game inv: " + gameInv);
      for (ResourceAccount a : accounts) {
        logger.debug(a.getName() + " " + a.inventory() + " " + a.units());
        gameInv.subtractInventory(a.inventory());
      }
      logger.debug("Missing: " + gameInv);
      if (!gameInv.isEmpty()) {
        logger.error("Mismatch!");
        // TODO: stamp out more mismatches
        // TODO: ensure accounts don't go negative.
        unallocatedAccount.allocate(gameInv, AllocationType.EXCESS);
      }
    }

    // Watch units attempting to build to see if they have canceled their build
    // order. If so, refund them.
    // FIXME: there are some cases this seems to cancel an order too early
    // TODO: track how long ongoingBuildAttempts have been around, throw
    // exception if it's too long
    // suggesting our tracking code is broken.
    HashSet<Unit> toRemove = new HashSet<Unit>();
    for (Unit unit : ongoingBuildAttempts.keySet()) {
      if (!unit.getOrder().equals(Order.BUILD_PROTOSS_1) && !unit.getOrder().equals(Order.BUILD_PROTOSS_2)
          && !unit.getOrder().equals(Order.RESET_COLLISION_2)
          && !unit.getOrder().equals(Order.RESET_COLLISION_1) && !unit.getOrder().equals(Order.BUILD_TERRAN)
          && !unit.getOrder().equals(Order.ZERG_BUILDING_MORPH) && !unit.getOrder().equals(Order.ZERG_UNIT_MORPH)
          && !unit.getOrder().equals(Order.DRONE_BUILD) && !unit.getOrder().equals(Order.DRONE_LAND)
          && !unit.getOrder().equals(Order.DRONE_START_BUILD) && !unit.getOrder().equals(Order.DRONE_LIFTOFF)
          && !unit.getOrder().equals(Order.ZERG_BIRTH) && !unit.getOrder().equals(Order.ZERG_BUILD_SELF)) {
        UnitType type = ongoingBuildAttempts.get(unit);
        ResourceInventory inv = new ResourceInventory(type.mineralPrice(), type.gasPrice(), type
            .supplyRequired());

        ResourceAccount acct = ongoingBuildAccounts.get(unit);
        if (acct == null || acct.isClosed()) {
          acct = getUnallocatedAccount();
        }
        transfer(buildAccount, acct, inv, ReleaseType.MANUAL, AllocationType.REFUND);
        toRemove.add(unit);
        logger.debug("Canceled/failed build attempt by " + unit + " of " + type.getName()
            + " currently executing order " + unit.getOrder().getName() + " transferring refund to "
            + unitMap.get(unit).account);
        
        // FIXME HACK: Just to be sure, tell the builder to stop so it doesn't
        // "accidentally" build anyway
        unit.stop();
      }
    }
    for (Unit u : toRemove) {
      ongoingBuildAttempts.remove(u);
      ongoingBuildAccounts.remove(u);
      ongoingBuildPositions.remove(u);
    }

    // TODO: remove from build account

    supplySpentPrevFrame = 0;
    mineralsSpentPrevPrevFrame = mineralsSpentPrevFrame;
    mineralsSpentPrevFrame = 0;
  }

  private void addInitialUnit(Unit unit) {
    AnOvermindUnit ou = new AnOvermindUnit(unit, unallocatedAccount, this);
    unitMap.put(unit, ou);
    List<OvermindUnit> single = Collections.<OvermindUnit> singletonList(ou);
    unallocatedAccount.allocate(new ResourceInventory(single), single, AllocationType.UNIT_CREATION);
  }

  protected ResourceAccount getUnallocatedAccount() {
    return unallocatedAccount;
  }

  /**
   * Transfers the unallocated account to a new owner. The old owner will be
   * invalidated.
   * 
   * @param holder
   */
  public void setUnallocatedHolder(ResourceAccount.Holder holder) {
    getUnallocatedAccount().transferOwner(holder);
  }

}
