package edu.berkeley.nlp.starcraft.macro;

import java.util.ArrayList;
import java.util.Collection;
import java.util.List;

import org.bwapi.proxy.model.Game;
import org.bwapi.proxy.model.TilePosition;
import org.bwapi.proxy.model.UnitType;
import org.bwapi.proxy.model.UpgradeType;

import collections.intCounter.IntCounter;

import edu.berkeley.nlp.starcraft.Action;
import edu.berkeley.nlp.starcraft.ActionProposal;
import edu.berkeley.nlp.starcraft.ActionProposer;
import edu.berkeley.nlp.starcraft.Strategy;
import edu.berkeley.nlp.starcraft.actions.TrainUnitAction;
import edu.berkeley.nlp.starcraft.construction.BuildGraphBuildManager;
import edu.berkeley.nlp.starcraft.model.BaseInfo;
import edu.berkeley.nlp.starcraft.model.BuildState;
import edu.berkeley.nlp.starcraft.model.Model;
import edu.berkeley.nlp.starcraft.model.buildGraph.buildNodes.BuildNode;
import edu.berkeley.nlp.starcraft.model.buildGraph.buildNodes.UnitNode;
import edu.berkeley.nlp.starcraft.model.buildGraph.buildNodes.UpgradeNode;
import edu.berkeley.nlp.starcraft.model.buildState.BuildGraphBuildState;
import edu.berkeley.nlp.starcraft.resources.ResourceInventory;
import edu.berkeley.nlp.starcraft.resources.ResourceSchedule;

/**
 * Does a bunch of reasonable macro-oriented things once our initial build order
 * is complete.
 * 
 * @author dlwh
 * 
 */
public class ZergTechProposer implements ActionProposer {
	Strategy strategy;
	public ZergTechProposer(Strategy strat)
	{
		this.strategy = strat;
	}
  @Override
  public List<ActionProposal> propose(Collection<Action> currentActions, ResourceSchedule resourceSchedule,
      ResourceInventory endingInventory) {
    List<ActionProposal> proposals = new ArrayList<ActionProposal>();
    // expand if we have a lot of gas and nothing to do with it.
    // TODO: Use endingInventory?
    BuildGraphBuildState baseState = BuildGraphBuildState.newBuildState(Game.getInstance().self(),
    		Model.getInstance().getPlannedBuildingsTracker().getPlannedBuildings());
    IntCounter<BuildNode> baseStateCounts = new IntCounter<BuildNode>();
    baseStateCounts.incrementAll(baseState.getCompletedCounts());
    baseStateCounts.incrementAll(baseState.getInProgressCounts());
    BuildGraphBuildManager manager = new BuildGraphBuildManager(strategy);
    final int plannedExtractorCount = baseState.getCount(new UnitNode(UnitType.ZERG_EXTRACTOR));

    int larvaCount = Game.getInstance().self().completedUnitCount(UnitType.ZERG_LARVA);
    int geyserCount = 0;
    int mineralsCount = 0;
    for (TilePosition selfBaseLocation : Model.getInstance().getBaseTracker().getOrderedSelfBaseLocations()) {
      BaseInfo info = Model.getInstance().getBaseTracker().getLocationsToBaseInfos().get(selfBaseLocation);
      geyserCount += info.getVespeneGeyserLocationsToMiningZones().size();
      mineralsCount += info.getMineralFieldsToMiningZones().size();
    }

    if (baseState.getCount(new UnitNode(UnitType.ZERG_MUTALISK)) > 10) {
      IntCounter<BuildNode> state = new IntCounter<BuildNode>(baseStateCounts);
      state.incrementCount(new UpgradeNode(UpgradeType.ZERG_FLYER_ATTACKS, 2), 1);
      state.incrementCount(new UpgradeNode(UpgradeType.ZERG_FLYER_CARAPACE, 2), 1);
      manager.clearQueue();
      manager.addToQueue(state);
      proposals.addAll(manager.propose(currentActions, resourceSchedule, endingInventory));
    }
    
    if( baseState.getCount(new UnitNode(UnitType.ZERG_MUTALISK)) > 12) {
      IntCounter<BuildNode> state = new IntCounter<BuildNode>(baseStateCounts);
      state.incrementCount(new UpgradeNode(UpgradeType.PNEUMATIZED_CARAPACE, 1), 1);
      manager.clearQueue();
      manager.addToQueue(state);
      proposals.addAll(manager.propose(currentActions, resourceSchedule, endingInventory));
    }
    
    if( baseState.getCount(new UnitNode(UnitType.ZERG_ZERGLING)) > 4) {
        IntCounter<BuildNode> state = new IntCounter<BuildNode>(baseStateCounts);
        state.incrementCount(new UpgradeNode(UpgradeType.METABOLIC_BOOST, 1), 1);
        manager.clearQueue();
        manager.addToQueue(state);
        proposals.addAll(manager.propose(currentActions, resourceSchedule, endingInventory));
      }
    
    if( baseState.getCount(new UnitNode(UnitType.ZERG_SPAWNING_POOL)) < 1) {
        IntCounter<BuildNode> state = new IntCounter<BuildNode>(baseStateCounts);
        state.incrementCount(new UnitNode(UnitType.ZERG_SPAWNING_POOL), 1);
        manager.clearQueue();
        manager.addToQueue(state);
        proposals.addAll(manager.propose(currentActions, resourceSchedule, endingInventory));
      }
    
    if( baseState.getCount(new UnitNode(UnitType.ZERG_MUTALISK)) < 1) {
        IntCounter<BuildNode> state = new IntCounter<BuildNode>(baseStateCounts);
        state.incrementCount(new UnitNode(UnitType.ZERG_MUTALISK), 1);
        manager.clearQueue();
        manager.addToQueue(state);
        proposals.addAll(manager.propose(currentActions, resourceSchedule, endingInventory));
      }

    if (plannedExtractorCount < geyserCount) {
        IntCounter<BuildNode> state = new IntCounter<BuildNode>(baseStateCounts);
        state.incrementCount(new UnitNode(UnitType.ZERG_EXTRACTOR), 1);
        manager.addToQueue(state);
        proposals.addAll(manager.propose(currentActions, resourceSchedule, endingInventory));
      }
    
    return proposals;
  }

}
