//SquadAction > SquadManager > Squad > Controller

package edu.berkeley.nlp.starcraft.micro;

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 org.bwapi.proxy.model.Position;
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.drawing.DrawingContext;
import edu.berkeley.nlp.starcraft.drawing.GameDrawingContext;
import edu.berkeley.nlp.starcraft.micro.ZerglingTribe.Mood;
import edu.berkeley.nlp.starcraft.model.BaseTracker;
import edu.berkeley.nlp.starcraft.model.Model;
import edu.berkeley.nlp.starcraft.overmind.Underling;
import edu.berkeley.nlp.starcraft.util.Log;
import edu.berkeley.nlp.util.Maxer;

// In charge of all zerglings (was ZerglingSquadManager)
public class CerebrateZasz implements Underling {
	public static final String gCerebrateZaszHUDWindowName = "= CerebrateZasz HUD =";
	public static final Log logger = Log.getLog("zerglings");
	private static final int HUD_X = 400;
	private static final int HUD_Y = 120;
	private static final double NEAR_TOL = 400;
	private static final int MAX_DEFENSE_LEVEL = 5;
	int TRIBE_NAME_MAX_LEN = 8;
	Map<OvermindUnit, ZerglingController> controllerAssignments = new HashMap<OvermindUnit, ZerglingController>();
	Map<ZerglingController, ZerglingTribe> assignments = new HashMap<ZerglingController, ZerglingTribe>();

	// Set of zerglings that should be (re)assigned (new or excommun'd)
	Set<ZerglingController> assignUs = new HashSet<ZerglingController>();

	// Current functioning tribes
	List<ZerglingTribe> tribes = new ArrayList<ZerglingTribe>();

	public CerebrateZasz() {
		/*
		 * // TODO: Comment out to not use windowed drawing.
		 * WindowDrawingContexts.setWindowLocation(gCerebrateZaszHUDWindowName,
		 * new Point(200, 400));
		 * WindowDrawingContexts.setWindowSize(gCerebrateZaszHUDWindowName, new
		 * Dimension(200, 180)); //
		 */
		makeInitialTribes();
	}

	@Override
	public void onFrame(int frameNum) {
		cleanUp();
		reassignUnits();
		if (Model.getInstance().drawFlags.master && Model.getInstance().drawFlags.cerebrateZasz) {
			updateHUD();
		}
		for (ZerglingTribe tribe : new HashSet<ZerglingTribe>(assignments
				.values())) {
			tribe.giveOrders(frameNum);
		}
	}

	// Empty tribes and dead zerglings are annoying, so we purge them
	private void cleanUp() {
		// No more phantom lings
		for (OvermindUnit ling : new ArrayList<OvermindUnit>(
				controllerAssignments.keySet())) {
			if (!ling.isValid() || !ling.exists()) {
				removeController(controllerAssignments.get(ling));
			}
		}

		// Process excommunications
		for (ZerglingController c : assignUs) {
			if (assignments.containsKey(c)) {
				ZerglingTribe tribe = assignments.get(c);
				assignments.remove(c);
				tribe.removeUnit(c);
			}
		}
	}

	private void updateHUD() {
		int i = 0;
		// /*
		// Draw on screen.
		DrawingContext context = GameDrawingContext.INSTANCE;
		int Xoff = HUD_X;
		int Yoff = HUD_Y;
		// */
		/*
		 * DrawingContext context =
		 * WindowDrawingContexts.getDrawingContext(gCerebrateZaszHUDWindowName);
		 * int Xoff = 10; int Yoff = 10; //
		 */
		context.drawTextScreen(Xoff, Yoff + 10 * i++, "= CerebrateZasz HUD =");
		for (ZerglingTribe tribe : tribes) {
			String fmt = "+ %-" + TRIBE_NAME_MAX_LEN + "s  %03d  %s  %s";
			String tribeString = String.format(fmt, tribe.name, tribe.getUs()
					.size(), tribe.getCurrentMood().name(), tribe
					.getDestination().toString());
			context.drawTextScreen(Xoff, Yoff + 10 * i++, tribeString);
		}
	}

	// we create a controller, bind the unit to it, then add it to the
	// unassigned set
	public void addUnit(OvermindUnit u) {
		ZerglingController c = new ZerglingController(u);
		controllerAssignments.put(u, c);
		assignUs.add(c);
	}

	// we add the controller to the list of lings to be reassigned
	public void recycleController(ZerglingController c) {
		assignUs.add(c);
	}

	public void removeUnit(OvermindUnit u) {
		ZerglingController c = controllerAssignments.get(u);
		removeController(c);
	}

	public void removeController(ZerglingController c) {
		if (assignUs.contains(c)) {
			assignUs.remove(c);
		}
		if (assignments.containsKey(c)) {
			ZerglingTribe tribe = assignments.get(c);
			assignments.remove(c);
			tribe.removeUnit(c);
		}
	}

	// The Brain of the Code, we must be extra careful doing these methods.
	List<ZerglingTribe> harassTribes; // Map from their bases
	List<ZerglingTribe> defenseTribes; // Map from our bases
	Set<TilePosition> processed; // Bases that have been processed (temp)
	ZerglingTribe ferry, nomads;
	Map<ZerglingTribe, Integer> defenseLevels; // How many zerglings to bring to

	// a fight

	private void makeInitialTribes() {
		harassTribes = new ArrayList<ZerglingTribe>();
		defenseTribes = new ArrayList<ZerglingTribe>();
		processed = new HashSet<TilePosition>();
		defenseLevels = new HashMap<ZerglingTribe, Integer>();

		TilePosition home = Model.getInstance().getBaseTracker()
				.getStartingMainLocation();
		if (home != null) {
			makeDefenseTribe(home);
		}

		ferry = new ZerglingTribe(Mood.MOVING, this);
		updateFerryTarget();
		tribes.add(ferry);

		TilePosition middle = Model.getInstance().getStaticMapData()
				.getMiddle();
		// TODO Fix this mood one day
		Mood nomadMood = Model.getInstance().baseTracker.getStartingMainLocation() == null ? Mood.BROODING : Mood.EXPLODING;
		nomads = new ZerglingTribe(nomadMood , this);
		nomads.setDestination(new Position(middle));
		tribes.add(nomads);
	}

	private void makeDefenseTribe(TilePosition pos) {
		ZerglingTribe t = new ZerglingTribe(Mood.EXPLODING, this);
		t.setDestination(new Position(pos));
		tribes.add(t);
		defenseTribes.add(t);
		processed.add(pos);
	}

	private void makeHarassTribe(TilePosition pos) {
		ZerglingTribe t = new ZerglingTribe(Mood.BOYISH_ATTITUDE, this);
		t.setDestination(new Position(pos));
		tribes.add(t);
		harassTribes.add(t);
		processed.add(pos);
	}

	private void updateFerryTarget() {

		// Ferry to our base under attack
		for (ZerglingTribe t : defenseTribes) {
			if (badGuysNearTheTribe(t)) {
				ferry.setDestination(t.getDestination());
			}
		}

		// Ferry to an enemy base
		if (harassTribes.size() > 0) {
			Maxer<Position> easiestTarget = new Maxer<Position>();
			for (ZerglingTribe t : harassTribes) {
				int val = getDefenseIndex(t);
				easiestTarget.observe(t.getDestination(), -val);
			}
			if (easiestTarget.max() >= -MAX_DEFENSE_LEVEL) {
				ferry.setDestination(easiestTarget.argMax());
				return;
			}
			Position pos = easiestTarget.argMax();
			if (pos != null) {
				// TODO Set to a place outside of the base
			}
		}
		ferry.setDestination(new Position(Model.getInstance()
				.getStaticMapData().getMiddle()));
	}

	private int getDefenseIndex(ZerglingTribe t) {
		Integer d = defenseLevels.get(t);
		return d == null ? 0 : d;
	}

	private boolean badGuysNearTheTribe(ZerglingTribe t) {
		Position p = t.getDestination();
		for (Unit u : Model.getInstance().opponentTracker
				.getVisibleEnemyUnits()) {
			if (u.getPosition().getDistance(p) < NEAR_TOL) {
				return true;
			}
		}
		return false;
	}

	private void updateTribes() {
		// Make sure all bases have a tribe
		BaseTracker baseTracker = Model.getInstance().getBaseTracker();
		for (TilePosition t : baseTracker.getOrderedSelfBaseLocations()) {
			if (!processed.contains(t)) {
				makeDefenseTribe(t);
			}
		}
		for (TilePosition t : baseTracker.getOrderedEnemyBaseLocations()) {
			if (!processed.contains(t)) {
				makeHarassTribe(t);
			}
		}

		// If an enemy-base tribe is empty, set it to boyish
		// Adjust defense levels for visible bases
		Set<Unit> vis = Model.getInstance().opponentTracker
				.getVisibleEnemyUnits();
		for (ZerglingTribe t : harassTribes) {
			if (!badGuysNearTheTribe(t)) {
				// t.setMood(Mood.BOYISH_ATTITUDE);
			}

			// Defense levels
			if (t.getUs().size() == 0) {
				continue;
			}
			Position p = t.getDestination();
			int level = 0;
			for (Unit u : vis) {
				if (u.getPosition().getDistance(p) < NEAR_TOL) {
					if (u.getType() == UnitType.TERRAN_BUNKER) {
						level += 12;
					}
					if (u.getType().groundWeapon().targetsGround()) {
						level += 2;
					}
				}
				defenseLevels.put(t, level);
			}
		}
	}

	private void reassignUnits() {
		updateTribes(); // Fiddle with tribe parameters
		updateFerryTarget();

		Set<ZerglingController> newAssignUs = new HashSet<ZerglingController>();
		for (ZerglingController c : assignUs) {
			if (!chooseAssignment(c)) {
				newAssignUs.add(c);
			}
		}
		assignUs = newAssignUs;
	}

	private boolean chooseAssignment(ZerglingController c) {
		Position unitPos = c.unit.getPosition();
		for (ZerglingTribe t : harassTribes) {
			if (isClose(unitPos, t.getDestination())) {
				assign(c, t);
				return true;
			}
		}
		for (ZerglingTribe t : defenseTribes) {
			if (isClose(unitPos, t.getDestination())) {
				if (badGuysNearTheTribe(t)) {
					assign(c, t);
					return true;
				}
			}
		}
		if (Model.getInstance().baseTracker.getStartingMainLocation() != null){
			if (!isClose(unitPos, ferry.getDestination())) { 
				assign(c, ferry);
				return true;
			}
		}
		assign(c, nomads);
		return true;
	}

	private void assign(ZerglingController c, ZerglingTribe t) {
		assignments.put(c, t);
		t.addUnit(c);
	}

	private boolean isClose(Position p1, Position p2) {
		return p1.getDistance(p2) < NEAR_TOL;
	}
}
