/*
* Tyr is an AI for StarCraft: Broodwar, 
* 
* Please visit https://github.com/SimonPrins/Tyr for further information.
* 
* Copyright 2015 Simon Prins
*
* This file is part of Tyr.
* Tyr is free software: you can redistribute it and/or modify
* it under the terms of the GNU General Public License as published by
* the Free Software Foundation, either version 2 of the License, or
* (at your option) any later version.
* Tyr is distributed in the hope that it will be useful,
* but WITHOUT ANY WARRANTY; without even the implied warranty of
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
* GNU General Public License for more details.
* You should have received a copy of the GNU General Public License
* along with Tyr.  If not, see http://www.gnu.org/licenses/.
*/


package com.tyr.unitgroups;

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 com.tyr.BWTAProxy;
import com.tyr.DebugMessages;
import com.tyr.EnemyManager;
import com.tyr.PositionUtil;
import com.tyr.Tyr;
import com.tyr.agents.Agent;
import com.tyr.buildingplacement.SpaceManager;
import com.tyr.tasks.ObserverSolution;

import bwapi.Color;
import bwapi.Game;
import bwapi.Player;
import bwapi.Position;
import bwapi.Unit;
import bwapi.UnitType;
import bwapi.WeaponType;
import bwta.BaseLocation;
import bwta.Region;


/**
 * A unit group for sending units to attack a certain location.
 * @author Simon
 *
 */
public class DistributeAttackGroup extends IAttackGroup
{
	/**
	 * The maximum range at which we will target enemy units.
	 */
	private double maxTargetRange = 300;
	
	/**
	 * Range at which an enemy is considered a threat.
	 */
	public static int threatRange = 450;
	
	/**
	 * Distance from which to wait before attacking.
	 */
	public static int waitDistance = -400;
	
	private int initialBonusDist = 200;
	
	private Map<Integer, EnemyUnit> enemyMap = new HashMap<Integer, EnemyUnit>();
	
	private Map<Integer, Integer> attackMap = new HashMap<Integer, Integer>();
	
	private List<Unit> enemyTanks = new ArrayList<>();
	
	private int mode = ATTACK;
	
	private static final int ATTACK = 0;

	private static final int RETREAT = 1;
	private static final int SAFE_MOVE_OUT = 2;
	
	private int lastRetreatFrame = -10000;

	public static int requiredForAttack = 15;
	
	public static int requiredForRetreat = 10;
	
	private List<ZealotDropGroup> drops = new ArrayList<>();
	
	/**
	 * A unit group for sending units to attack a certain location.
	 * @param target The target where we want the units to attack.
	 */
	public DistributeAttackGroup(Position target) 
	{
		super(target);
	}
	
	@Override
	public void onFrame(Game game, Player self, Tyr bot)
	{
		Tyr.drawCircle(target, Color.Orange, 64);
		
		for (ZealotDropGroup drop : drops)
			drop.cleanup();
		
		armyLoop : for (int i = units.size() - 1; i >= 0; i--)
		{
			final Agent agent = units.get(i);
			if (agent.unit.getType() == UnitType.Protoss_Shuttle)
			{
				for (ZealotDropGroup drop : drops)
				{
					if (drop.needsShuttle())
					{
						units.remove(i);
						drop.add(agent);
						continue armyLoop;
					}
				}
				ZealotDropGroup drop = new ZealotDropGroup();
				units.remove(i);
				drop.add(agent);
				drops.add(drop);
			}
		}
		
		for (ZealotDropGroup drop : drops)
		{
			if (!drop.needsShuttle())
			{
				for (int i = units.size() - 1; i >= 0; i--)
				{
					if (drop.getZealots() >= 4)
						break;
					final Agent agent = units.get(i);
					if (agent.unit.getType() == UnitType.Protoss_Zealot)
					{
						units.remove(i);
						drop.add(agent);
					}
				}
			}
			
			drop.setTarget(getDropTarget());
			drop.onFrame(game, self, bot);
		}
		
		
		if (mode == SAFE_MOVE_OUT)
		{
			safeMoveOut();
			return;
		}
		
		final Agent observer = ObserverSolution.getArmyObserver();
		final boolean attackingMainOrNatural = Tyr.bot.suspectedEnemy.size() == 1 
				&& target != null
				&& (PositionUtil.distanceSq(Tyr.bot.suspectedEnemy.get(0).getPosition(), target) <= 200 * 200)
					|| (SpaceManager.getEnemyNatural() != null && PositionUtil.distanceSq(SpaceManager.getEnemyNatural().getPosition(), target) <= 200 * 200);
		
		
		for (Agent agent : units)
			if (agent.unit.getType() == UnitType.Protoss_Reaver)
			{
				agent.drawCircle(Color.Purple);
				if (!agent.unit.isTraining())
					agent.unit.train(UnitType.Protoss_Scarab);
			}


		for (int i = enemyTanks.size() - 1; i >= 0; i--)
		{
			final Unit unit = enemyTanks.get(i);
			final Position pos = EnemyManager.getManager().getLastPosition(unit);
			for (Agent agent : units)
				if (pos == null || agent.distanceSquared(pos) <= 200 * 200)
				{
					enemyTanks.remove(i);
					break;
				}
		}
		
		final List<Unit> spiderMines = new ArrayList<>();
		
		for (Unit unit : EnemyManager.getEnemyUnits())
		{
			if (unit.getType() == UnitType.Terran_Vulture_Spider_Mine)
				spiderMines.add(unit);
			
			if (unit.getType() == UnitType.Terran_Siege_Tank_Siege_Mode
					|| unit.getType() == UnitType.Terran_Siege_Tank_Tank_Mode)
			{
				if (!enemyTanks.contains(unit))
					enemyTanks.add(unit);
			}
		}

		boolean lowGroundTanks = false;
		boolean highGroundTanks = false;
		for (Unit tank : enemyTanks)
		{
			if (!lowGroundTanks || !highGroundTanks)
			{
				final Position pos = EnemyManager.getManager().getLastPosition(tank);
				
				if (pos != null)
				{
					final boolean highGround = isOnHighGround(tank);
					if (highGround)
						highGroundTanks = true;
					else
						lowGroundTanks = true;
				}
			}
		}
		
		final Set<Unit> targettableEnemies = new HashSet<>();
		
		for (Unit unit : EnemyManager.getEnemyUnits())
		{
			if (!unit.getType().isWorker() && unit.getType() != UnitType.Terran_Vulture)
				continue;
			for (BaseLocation pos : bot.expands)
			{
				if (PositionUtil.distanceSq(pos.getPosition(), unit) <= 300 * 300)
				{
					targettableEnemies.add(unit);
					break;
				}
			}
			
			for (BaseLocation pos : bot.suspectedEnemy)
			{
				if (PositionUtil.distanceSq(pos.getPosition(), unit) <= 300 * 300)
				{
					targettableEnemies.add(unit);
					break;
				}
			}
		}

		final int correctedWaitDist = waitDistance + initialBonusDist; 
		int waiting = 0;
		int attacking = 0;
		for (Agent agent : units)
		{
			if (agent.distanceSquared(target) < (correctedWaitDist - 200) * (correctedWaitDist - 200))
			{
				agent.drawCircle(Color.Red);
				attacking++;
				waiting++;
				continue;
			}
			else if (agent.distanceSquared(target) < (correctedWaitDist + 400) * (correctedWaitDist + 400))
			{
				if (agent.distanceSquared(target) < correctedWaitDist * correctedWaitDist)
					agent.drawCircle(Color.Blue);
				else
					agent.drawCircle(Color.Teal);
				waiting++;
				continue;
			}
			
			for (Unit tank : enemyTanks)
			{
				final Position pos = EnemyManager.getManager().getLastPosition(tank);
				if (agent.distanceSquared(pos) < (correctedWaitDist + 400) * (correctedWaitDist + 400))
				{
					if (agent.distanceSquared(pos) < correctedWaitDist * correctedWaitDist)
						agent.drawCircle(Color.Blue);
					else
						agent.drawCircle(Color.Teal);
					waiting++;
					break;
				}
			}
		}
		
		List<Integer> removeIDs = new ArrayList<>();
		
		for (Integer id : enemyMap.keySet())
		{
			if (enemyMap.get(id).isDead())
				removeIDs.add(id);
			else
				Agent.clean(enemyMap.get(id).attackingAgents);
		}
		for (Integer id : removeIDs)
			enemyMap.remove(id);
		
		removeIDs = new ArrayList<>();
		for (Integer id : attackMap.keySet())
			if (!enemyMap.containsKey(attackMap.get(id)))
				removeIDs.add(id);
		for (Integer id : removeIDs)
			attackMap.remove(id);
		
		final boolean contain = (EnemyManager.getManager().getAllCount(UnitType.Terran_Command_Center) <= 1 && !lowGroundTanks && highGroundTanks && waiting < 25);
		final boolean commit = !contain && !(waiting <= requiredForRetreat && attacking <= 5);
		
		if (waiting <= 10)
		{
			mode = RETREAT;
			lastRetreatFrame = game.getFrameCount();
		}
		else if (contain && waiting < requiredForAttack)
		{
			mode = RETREAT;
			lastRetreatFrame = game.getFrameCount();
		}
		else if (contain && !lowGroundTanks && EnemyManager.getManager().getAllCount(UnitType.Terran_Command_Center) <= 1 && waiting < 15)
		{
			mode = RETREAT;
			lastRetreatFrame = game.getFrameCount();
		}
		else if (game.getFrameCount() - lastRetreatFrame >= 200
				&& mode == RETREAT 
				&& (!contain) 
				&& waiting >= requiredForAttack)
		{
			initialBonusDist = 0;
			mode = ATTACK;
		}
		
		if (highGroundTanks)
			DebugMessages.addMessage("Highground tanks detected.");
		if (lowGroundTanks)
			DebugMessages.addMessage("Lowground tanks detected.");
		
		if (contain)
			DebugMessages.addMessage("Containing the enemy.");
		else if (commit)
			DebugMessages.addMessage("Attacking the enemy.");
		
		for (Agent agent : units)
		{
			final EnemyUnit currentTarget;
			if (attackMap.containsKey(agent.unit.getID()))
			{
				// Do not retarget too often.
				if ((agent.unit.getID() + game.getFrameCount()) % 15 != 0)
					continue;
				
				boolean removed = false;
				
				if ((agent.unit.getID() + game.getFrameCount()) % 200 != 0)
				{
					final EnemyUnit enemy = enemyMap.get(attackMap.get(agent.unit.getID()));
					if (isOnHighGround(enemy.enemy) 
							|| game.getGroundHeight(agent.unit.getTilePosition()) < game.getGroundHeight(enemy.enemy.getTilePosition())
							|| ((enemy.enemy.getType().isWorker() || enemy.enemy.getType() == UnitType.Terran_Vulture) && !targettableEnemies.contains(enemy.enemy))
							|| (agent.unit.getType() == UnitType.Protoss_Zealot && enemy.enemy.getType() == UnitType.Terran_Siege_Tank_Tank_Mode && agent.distanceSquared(target) >= 1000 * 1000)
							|| (agent.unit.getType().airWeapon() == WeaponType.None && enemy.enemy.isLifted()))
					{
						enemy.attackingAgents.remove(agent);
						attackMap.remove(agent.unit.getID());
						removed = true;
					}
				}
				
				if (!removed)
					currentTarget = enemyMap.get(attackMap.get(agent.unit.getID()));
				else
					currentTarget = null;
			}
			else
				currentTarget = null;
			
			EnemyUnit target = currentTarget;
			
			double enemyDistance = maxTargetRange * maxTargetRange;
			if (target != null)
				enemyDistance = agent.distanceSquared(target.enemy);
			
			if (agent.unit.getType() == UnitType.Protoss_Dragoon)
			{
				boolean mineIsClose = false;
				for (Unit mine : spiderMines)
					if (agent.distanceSquared(mine) <= 300 * 300)
					{
						mineIsClose = true;
						break;
					}
				if (mineIsClose)
					continue;
			}
			
			for (Unit enemy : EnemyManager.getEnemyUnits())
			{
				if (enemy.getType().isBuilding() && !enemy.getType().isResourceDepot())
					continue;
				
				if (enemy.getType() == UnitType.Protoss_Observer || enemy.getType() == UnitType.Terran_Vulture_Spider_Mine)
					continue;
				
				if (!enemy.isVisible())
					continue;
				
				if (enemy.getType() == UnitType.Terran_Vulture && agent.unit.getType() == UnitType.Protoss_Zealot)
					continue;
				
				if (agent.unit.getType().airWeapon() == WeaponType.None && enemy.getType().isFlyer())
					continue;
				
				if (agent.unit.getType().groundWeapon() == WeaponType.None && !enemy.getType().isFlyer())
					continue;
				
				if (Tyr.bot.suspectedEnemy.size() == 1
						&& units.size() < 30 
						&& isOnHighGround(enemy))
					continue;
				
				if (game.getGroundHeight(agent.unit.getTilePosition()) < game.getGroundHeight(enemy.getTilePosition()))
					continue;

				
				final int distanceSq = agent.distanceSquared(enemy);
				final int maxDist;
				if (enemy.getType() == UnitType.Terran_Siege_Tank_Siege_Mode || enemy.getType() == UnitType.Terran_Siege_Tank_Tank_Mode)
					maxDist = 450 * 450;
				else
					maxDist = (int) (maxTargetRange * maxTargetRange);
				
				if (agent.unit.getType() == UnitType.Protoss_Zealot
						&& enemy.getType() == UnitType.Terran_Siege_Tank_Tank_Mode
						&& agent.distanceSquared(this.target) >= 1000 * 1000)
					continue;
					
				if (distanceSq >= maxDist)
					continue;
				
				final EnemyUnit unit;
				if (!enemyMap.containsKey(enemy.getID()))
				{
					unit = new EnemyUnit(enemy);
					enemyMap.put(enemy.getID(), unit);
				}
				else
					unit = enemyMap.get(enemy.getID());
				

				if ((enemy.getType().isWorker() || enemy.getType() == UnitType.Terran_Vulture)
						&& !targettableEnemies.contains(enemy))
					continue;
				
				if (enemy.getType().isWorker()
						|| enemy.getType() == UnitType.Terran_Marine
						|| enemy.getType() == UnitType.Terran_Firebat)
				{
					if (unit.attackingAgents.size() > 1)
						continue;
				}
				
				if (enemy.getType() == UnitType.Terran_Vulture
						|| enemy.getType() == UnitType.Terran_Wraith)
				{
					if (unit.attackingAgents.size() > 2)
						continue;
				}
				
				if (target == null)
				{
					target = unit;
					enemyDistance = distanceSq;
					continue;
				}
				
				if (getPriority(target.enemy.getType()) > getPriority(enemy.getType()))
					continue;
				
				if (getPriority(target.enemy.getType()) < getPriority(enemy.getType()))
				{
					target = unit;
					enemyDistance = distanceSq;
					continue;
				}
				
				if (unit.attackingAgents.size() < target.attackingAgents.size() && !Agent.isTank(target.enemy))
				{
					target = unit;
					enemyDistance = distanceSq;
					continue;
				}
				
				if (unit.attackingAgents.size() > target.attackingAgents.size() && !Agent.isTank(target.enemy))
					continue;
				
				if (distanceSq < enemyDistance)
				{
					target = unit;
					enemyDistance = distanceSq;
				}
			}
			
			if (target != null && target != currentTarget)
			{
				if (currentTarget != null)
					currentTarget.attackingAgents.remove(agent);
				target.attackingAgents.add(agent);
				attackMap.put(agent.unit.getID(), target.enemy.getID());
			}
		}
		
		// Order all units to attack the target.
		for (Agent agent : units)
		{
			if (agent.unit.getType() == UnitType.Protoss_Zealot && !attackMap.containsKey(agent.unit.getID()))
			{
				boolean mineIsClose = false;
				for (Unit mine : spiderMines)
					if (agent.distanceSquared(mine) <= 300 * 300)
					{
						mineIsClose = true;
						break;
					}
				if (mineIsClose)
				{
					agent.unit.move(Tyr.getStartLocation());
					continue;
				}
			}
			final int dist = agent.distanceSquared(target);
			if (mode == RETREAT && dist > (correctedWaitDist - 200) * (correctedWaitDist - 200))
			{
				boolean retreat = dist < correctedWaitDist * correctedWaitDist;
				if (!retreat)
				{
					for (Unit tank : enemyTanks)
					{
						final Position pos = EnemyManager.getManager().getLastPosition(tank);
						if (agent.distanceSquared(pos) < correctedWaitDist * correctedWaitDist)
						{
							retreat = true;
							break;
						}
					}
				}
				if (retreat)
				{
					agent.unit.move(Tyr.tileToPosition(self.getStartLocation()));
					continue;
				}
			}
			
			if (agent.unit.getType() == UnitType.Protoss_Reaver)
			{
				if (agent.unit.isMoving())
				{
					for (Unit enemy : EnemyManager.getEnemyUnits())
						if (!enemy.isLifted() && ! enemy.getType().isFlyer() && agent.distanceSquared(enemy) <= 250*250)
						{
							agent.unit.stop();
							return;
						}
				}
				agent.stutterstep(target, false);
				continue;
			}
			
			if (attackMap.containsKey(agent.unit.getID()))
			{
				game.drawLineMap(agent.unit.getPosition(), enemyMap.get(attackMap.get(agent.unit.getID())).enemy.getPosition(), Color.Red);
				if (enemyMap.get(attackMap.get(agent.unit.getID())) == null || enemyMap.get(attackMap.get(agent.unit.getID())).enemy == null)
					agent.drawCircle(Color.Red, 16);
				else
					agent.attack(enemyMap.get(attackMap.get(agent.unit.getID())).enemy);
			}
			else
			{
				if (attackingMainOrNatural
						&& observer != null 
						&& agent.distanceSquared(Tyr.getStartLocation()) >= 1000 * 1000 
						&& agent.distanceSquared(target) >= 800 * 800)
				{
					final int obsDist = observer.distanceSquared(agent);
					if (obsDist <= 400 * 400 && obsDist >= 200 * 200)
					{
						agent.unit.move(observer.unit.getPosition());
						game.drawLineMap(agent.unit.getPosition(), observer.unit.getPosition(), Color.Yellow);
						continue;
					}
				}
				agent.drawCircle(Color.Orange);
				agent.attack(target);
			}
		}
	}
	
	private Position getDropTarget()
	{
		int dist = 1000000000;
		Position result = Tyr.getStartLocation();
		for (Agent agent : units)
		{
			int newDist = agent.distanceSquared(target); 
			if (newDist < dist)
			{
				result = agent.unit.getPosition();
				dist = newDist;
			}
		}
		return result;
	}

	private void safeMoveOut()
	{
		Agent zealot = null;
		final Set<Agent> safeDragoons = new HashSet<>();
		
		for (Agent agent : units)
		{
			if (zealot == null && agent.unit.getType() == UnitType.Protoss_Zealot)
			{
				zealot = agent;
				continue;
			}
			if (safeDragoons.size() < 5 && agent.unit.getType() == UnitType.Protoss_Dragoon)
			{
				safeDragoons.add(agent);
				agent.drawCircle(Color.Green);
				continue;
			}
			agent.drawCircle(Color.Blue);
		}
		
		if (zealot == null && units.size() <= 5)
			return;
		
		if (zealot == null || zealot.distanceSquared(target) <= (waitDistance + initialBonusDist + 200) * (waitDistance + initialBonusDist + 200))
		{
			mode = RETREAT;
			return;
		}
		
		zealot.drawCircle(Color.Orange);
		
		int dragoonsTooFar = 0;
		for (Agent agent : units)
		{
			if (agent == zealot)
				continue;
			
			if (agent.distanceSquared(Tyr.getStartLocation()) <= 1000 * 1000)
			{
				if (agent.unit.getType() == UnitType.Protoss_Zealot)
					agent.move(target);
				else
					agent.attack(target);
				continue;
			}

			final int dist = agent.distanceSquared(zealot);
			if (safeDragoons.contains(agent))
			{
				if (dist <= 64 * 64)
				{
					Tyr.game.drawLineMap(agent.unit.getPosition(), Tyr.getStartLocation(), Color.Red);
					agent.attack(Tyr.getStartLocation());
					continue;
				}
				else if (dist <= 196 * 196)
				{
					Tyr.game.drawLineMap(agent.unit.getPosition(), zealot.unit.getPosition(), Color.Red);
					agent.attack(zealot.unit.getPosition());
					continue;
				}
				else
				{
					Tyr.game.drawLineMap(agent.unit.getPosition(), zealot.unit.getPosition(), Color.Blue);
					agent.move(zealot.unit.getPosition());
					dragoonsTooFar++;
					continue;
				}
			}
			
			if (agent.unit.getType() == UnitType.Protoss_Zealot)
			{
				if (dist <= 192 * 192)
				{
					Tyr.game.drawLineMap(agent.unit.getPosition(), Tyr.getStartLocation(), Color.Blue);
					agent.move(Tyr.getStartLocation());
					continue;
				}
				else
				{
					Tyr.game.drawLineMap(agent.unit.getPosition(), zealot.unit.getPosition(), Color.Blue);
					agent.move(zealot.unit.getPosition());
					continue;
				}
			}

			if (dist <= 192 * 192)
			{
				Tyr.game.drawLineMap(agent.unit.getPosition(), Tyr.getStartLocation(), Color.Red);
				agent.attack(Tyr.getStartLocation());
				continue;
			}
			else if (dist <= 256 * 256)
			{
				Tyr.game.drawLineMap(agent.unit.getPosition(), zealot.unit.getPosition(), Color.Red);
				agent.attack(zealot.unit.getPosition());
				continue;
			}
			else
			{
				Tyr.game.drawLineMap(agent.unit.getPosition(), zealot.unit.getPosition(), Color.Blue);
				agent.move(zealot.unit.getPosition());
				continue;
			}
		}
		
		if (dragoonsTooFar >= 3 && BWTAProxy.getRegion(zealot.unit.getPosition()) != BWTAProxy.getRegion(Tyr.getStartLocation()))
			zealot.unit.holdPosition();
		else
		{
			Tyr.game.drawLineMap(zealot.unit.getPosition(), target, Color.Orange);
			zealot.move(target);
		}
	}

	private static boolean isOnHighGround(Unit enemy)
	{
		final Region enemyRegion;
		if (Tyr.bot.suspectedEnemy.size() == 0)
			enemyRegion = null;
		else
			enemyRegion = BWTAProxy.getRegion(Tyr.bot.suspectedEnemy.get(0).getPosition());
		
		if (enemyRegion == null)
			Tyr.drawCircle(enemy.getPosition(), Color.Red, 4);
		else if (BWTAProxy.getRegion(enemy.getPosition()) == enemyRegion)
			Tyr.drawCircle(enemy.getPosition(), Color.Yellow, 4);
		
		if (enemyRegion != null && BWTAProxy.getRegion(enemy.getPosition()) == enemyRegion)
			return true;
		else return false;
		
		/*
		if (Tyr.game.getGroundHeight(enemy.getTilePosition()) >= 4)
			Tyr.drawCircle(enemy.getPosition(), Color.Green, 4);
		else
			Tyr.drawCircle(enemy.getPosition(), Color.Blue, 4);
		return Tyr.game.getGroundHeight(enemy.getTilePosition()) >= 4;
		*/
	}
	
	private static int getPriority(UnitType type)
	{
		if (type.isResourceDepot())
			return 2;
		else if (type.isWorker())
			return 1;
		else
			return 3;
	}
	
	private static class EnemyUnit
	{
		public Unit enemy;
		public List<Agent> attackingAgents = new ArrayList<Agent>();
		
		public EnemyUnit(Unit enemy)
		{
			this.enemy = enemy;
		}
		
		public boolean isDead()
		{
			return enemy.getHitPoints() <= 0 || enemy.getRemoveTimer() != 0 || !enemy.exists();
		}
	}
}
