import sc2
import random

from sc2 import Race, Difficulty, position
from sc2.position import Point2, Point3
from sc2.unit import Unit
from sc2.player import Bot, Computer
from sc2.player import Human
from sc2.ids.unit_typeid import UnitTypeId
from sc2.ids.ability_id import AbilityId
from sc2.ids.upgrade_id import UpgradeId
from sc2.units import Units


class AlphaWannabe(sc2.BotAI):
	newNexusPositions = [];

	async def distribute_workers(self, performanceHeavy=True, onlySaturateGas=False,):
		mineralTags = [x.tag for x in self.mineral_field]
		gas_buildingTags = [x.tag for x in self.gas_buildings]

		workerPool = Units([], self)
		workerPoolTags = set()

	    # Find all gas_buildings that have surplus or deficit
		deficit_gas_buildings = {}
		surplusgas_buildings = {}
		for g in self.gas_buildings.filter(lambda x: x.vespene_contents > 0):
		    # Only loop over gas_buildings that have still gas in them
		    deficit = g.ideal_harvesters - g.assigned_harvesters
		    if deficit > 0:
		        deficit_gas_buildings[g.tag] = {"unit": g, "deficit": deficit}
		    elif deficit < 0:
		        surplusWorkers = self.workers.closer_than(10, g).filter(
		            lambda w: w not in workerPoolTags
		            and len(w.orders) == 1
		            and w.orders[0].ability.id in [AbilityId.HARVEST_GATHER]
		            and w.orders[0].target in gas_buildingTags
		        )
		        for i in range(-deficit):
		            if surplusWorkers.amount > 0:
		                w = surplusWorkers.pop()
		                workerPool.append(w)
		                workerPoolTags.add(w.tag)
		        surplusgas_buildings[g.tag] = {"unit": g, "deficit": deficit}

		# Find all townhalls that have surplus or deficit
		deficitTownhalls = {}
		surplusTownhalls = {}
		if not onlySaturateGas:
		    for th in self.townhalls:
		        deficit = th.ideal_harvesters - th.assigned_harvesters
		        if deficit > 0:
		            deficitTownhalls[th.tag] = {"unit": th, "deficit": deficit}
		        elif deficit < 0:
		            surplusWorkers = self.workers.closer_than(10, th).filter(
		                lambda w: w.tag not in workerPoolTags
		                and len(w.orders) == 1
		                and w.orders[0].ability.id in [AbilityId.HARVEST_GATHER]
		                and w.orders[0].target in mineralTags
		            )
		            # workerPool.extend(surplusWorkers)
		            for i in range(-deficit):
		                if surplusWorkers.amount > 0:
		                    w = surplusWorkers.pop()
		                    workerPool.append(w)
		                    workerPoolTags.add(w.tag)
		            surplusTownhalls[th.tag] = {"unit": th, "deficit": deficit}

		    if all(
		        [
		            len(deficit_gas_buildings) == 0,
		            len(surplusgas_buildings) == 0,
		            len(surplusTownhalls) == 0 or deficitTownhalls == 0,
		        ]
		    ):
		        # Cancel early if there is nothing to balance
		        return

		# Check if deficit in gas less or equal than what we have in surplus, else grab some more workers from surplus bases
		deficitGasCount = sum(
		    gasInfo["deficit"] for gasTag, gasInfo in deficit_gas_buildings.items() if gasInfo["deficit"] > 0
		)
		surplusCount = sum(
		    -gasInfo["deficit"] for gasTag, gasInfo in surplusgas_buildings.items() if gasInfo["deficit"] < 0
		)
		surplusCount += sum(-thInfo["deficit"] for thTag, thInfo in surplusTownhalls.items() if thInfo["deficit"] < 0)

		if deficitGasCount - surplusCount > 0:
		    # Grab workers near the gas who are mining minerals
		    for gTag, gInfo in deficit_gas_buildings.items():
		        if workerPool.amount >= deficitGasCount:
		            break
		        workersNearGas = self.workers.closer_than(10, gInfo["unit"]).filter(
		            lambda w: w.tag not in workerPoolTags
		            and len(w.orders) == 1
		            and w.orders[0].ability.id in [AbilityId.HARVEST_GATHER]
		            and w.orders[0].target in mineralTags
		        )
		        while workersNearGas.amount > 0 and workerPool.amount < deficitGasCount:
		            w = workersNearGas.pop()
		            workerPool.append(w)
		            workerPoolTags.add(w.tag)

		# Now we should have enough workers in the pool to saturate all gases, and if there are workers left over, make them mine at townhalls that have mineral workers deficit
		for gTag, gInfo in deficit_gas_buildings.items():
		    if performanceHeavy:
		        # Sort furthest away to closest (as the pop() function will take the last element)
		        workerPool.sort(key=lambda x: x.distance_to(gInfo["unit"]), reverse=True)
		    for i in range(gInfo["deficit"]):
		        if workerPool.amount > 0:
		            w = workerPool.pop()
		            if len(w.orders) == 1 and w.orders[0].ability.id in [AbilityId.HARVEST_RETURN]:
		                self.do(w.gather(gInfo["unit"], queue=True))
		            else:
		                self.do(w.gather(gInfo["unit"]))

		if not onlySaturateGas:
		    # If we now have left over workers, make them mine at bases with deficit in mineral workers
		    for thTag, thInfo in deficitTownhalls.items():
		        if performanceHeavy:
		            # Sort furthest away to closest (as the pop() function will take the last element)
		            workerPool.sort(key=lambda x: x.distance_to(thInfo["unit"]), reverse=True)
		        for i in range(thInfo["deficit"]):
		            if workerPool.amount > 0:
		                w = workerPool.pop()
		                mf = self.mineral_field.closer_than(10, thInfo["unit"]).closest_to(w)
		                if len(w.orders) == 1 and w.orders[0].ability.id in [AbilityId.HARVEST_RETURN]:
		                    self.do(w.gather(mf, queue=True))
		                else:
		                    self.do(w.gather(mf))

	async def microUnits(self):
		enemies = self.enemy_units | self.enemy_structures
		enemies_can_attack = enemies.filter(lambda unit: unit.can_attack_ground)
		for r in self.units(UnitTypeId.MARINE):

		    # Move to range 15 of closest unit if reaper is below 20 hp and not regenerating
		    enemyThreatsClose = enemies_can_attack.filter(
		        lambda unit: unit.distance_to(r) < 15
		    )  # Threats that can attack the reaper

		    if r.health_percentage < 2 / 5 and enemyThreatsClose:
		        retreatPoints = self.neighbors8(r.position, distance=2) | self.neighbors8(r.position, distance=4)
		        # Filter points that are pathable
		        retreatPoints = {x for x in retreatPoints if self.in_pathing_grid(x)}
		        if retreatPoints:
		            closestEnemy = enemyThreatsClose.closest_to(r)
		            retreatPoint = closestEnemy.position.furthest(retreatPoints)
		            self.do(r.move(retreatPoint))
		            continue  # Continue for loop, dont execute any of the following

		    # Reaper is ready to attack, shoot nearest ground unit
		    enemyGroundUnits = enemies.filter(
		        lambda unit: unit.distance_to(r) < 5 and not unit.is_flying
		    )  # Hardcoded attackrange of 5
		    if r.weapon_cooldown == 0 and enemyGroundUnits:
		        enemyGroundUnits = enemyGroundUnits.sorted(lambda x: x.distance_to(r))
		        closestEnemy = enemyGroundUnits[0]
		        self.do(r.attack(closestEnemy))
		        continue  # Continue for loop, dont execute any of the following

		    # Move to max unit range if enemy is closer than 4
		    enemyThreatsVeryClose = enemies.filter(
		        lambda unit: unit.can_attack_ground and unit.distance_to(r) < 4.5
		    )  # Hardcoded attackrange minus 0.5
		    # Threats that can attack the reaper
		    if r.weapon_cooldown != 0 and enemyThreatsVeryClose:
		        retreatPoints = self.neighbors8(r.position, distance=2) | self.neighbors8(r.position, distance=4)
		        # Filter points that are pathable by a reaper
		        retreatPoints = {x for x in retreatPoints if self.in_pathing_grid(x)}
		        if retreatPoints:
		            closestEnemy = enemyThreatsVeryClose.closest_to(r)
		            retreatPoint = max(retreatPoints, key=lambda x: x.distance_to(closestEnemy) - x.distance_to(r))
		            self.do(r.move(retreatPoint))
		            continue  # Continue for loop, don't execute any of the following

		    # # Move to nearest enemy ground unit/building because no enemy unit is closer than 5
		    # allEnemyGroundUnits = self.enemy_units.not_flying
		    # if allEnemyGroundUnits:
		    #     closestEnemy = allEnemyGroundUnits.closest_to(r)
		    #     self.do(r.move(closestEnemy))
		    #     continue  # Continue for loop, don't execute any of the following

		    # Move to random enemy start location if no enemy buildings have been seen
		    self.do(r.move(random.choice(self.enemy_start_locations)))

	async def buildExpansion(self, buildingId): # UnitTypeId.whatever
		if (self.can_afford(UnitTypeId.NEXUS)):
			workers = self.workers.gathering
			if (workers):
				worker = workers.furthest_to(workers.center)
				# If a placement location was found
				location = await self.get_next_expansion()
				if location:
					self.do(worker.build(UnitTypeId.NEXUS, location), subtract_cost=True)
					self.newNexusPositions.append(location)

	async def trainUnits(self, buildingName, unitName):
		for sg in self.structures(buildingName).ready.idle:
			if self.can_afford(unitName):
				self.do(sg.train(unitName), subtract_cost=True, subtract_supply=True)


	async def warpUnits(self, abilityId, unitTypeId):
		for warpgate in self.structures(UnitTypeId.WARPGATE).ready:
			abilities = await self.get_available_abilities(warpgate)
			if abilityId in abilities:
				closestPylon = self.structures(UnitTypeId.PYLON).ready.closest_to(self.enemy_start_locations[0])
				placement = await self.find_placement(abilityId, closestPylon.position, placement_step=1)
				if placement is None:
					print("can't find valid place to warp in a unit!")
					return
				self.do(warpgate.warp_in(unitTypeId, placement), subtract_cost=True, subtract_supply=True)

	async def trainWorker(self, buildingId, workerId): #UnitTypeId.whatever, UnitTypeId.whatever
		if (
		    self.can_afford(workerId)
		    and self.supply_left > 0
		    and
		    (
		        self.townhalls(buildingId).idle
		        or self.townhalls(buildingId).idle
		    )
		):
		    for th in self.townhalls.idle:
		        self.do(th.train(workerId), subtract_cost=True, subtract_supply=True)







	async def manageIdleWorkers(self):
		if self.townhalls:
		    for w in self.workers.idle:
		        th = self.townhalls.closest_to(w)
		        mfs = self.mineral_field.closer_than(10, th)
		        if mfs:
		            mf = mfs.closest_to(w)
		            self.do(w.gather(mf))


	async def buildStructure(self, structureId): # UnitTypeId
		if (self.can_afford(structureId)):
			workers = self.workers.gathering
			if (workers):
				worker = workers.furthest_to(workers.center)
				# If a placement location was found
				location = await self.find_placement(structureId, worker.position, placement_step=3)
				if location:
	                # Order worker to build exactly on that location
					self.do(worker.build(structureId, location), subtract_cost=True)

	async def buildGeyser(self, geyserId):
		for nexus in self.townhalls.ready:
			vgs = self.vespene_geyser.closer_than(15, nexus)
			for vg in vgs:
			    if (self.can_afford(geyserId)):

				    worker = self.select_build_worker(vg.position)
				    if worker is None:
				        break

				    if not self.gas_buildings or not self.gas_buildings.closer_than(1, vg):
				        self.do(worker.build(geyserId, vg), subtract_cost=True)
				        self.do(worker.stop(queue=True))

	async def upgradeWarpgate(self):
		if (self.can_afford(AbilityId.RESEARCH_WARPGATE)):
			cCore = self.structures(UnitTypeId.CYBERNETICSCORE).ready.first
			self.do(cCore(AbilityId.RESEARCH_WARPGATE), subtract_cost=True)

	async def upgradeBlink(self):
		if (self.can_afford(UpgradeId.BLINKTECH)):
		    tCouncil = self.structures(UnitTypeId.TWILIGHTCOUNCIL).ready.first
		    self.do(tCouncil(AbilityId.RESEARCH_BLINK), subtract_cost=True)

	async def morphGateways(self):
		for gateway in self.structures(UnitTypeId.GATEWAY).ready.idle:
		    if self.already_pending_upgrade(UpgradeId.WARPGATERESEARCH) == 1:
		        self.do(gateway(AbilityId.RESEARCH_WARPGATE))

	async def buildPylon(self):
		if self.can_afford(UnitTypeId.PYLON):
			if (self.newNexusPositions != []):
				wantedBuildLocation = self.newNexusPositions[0] #.position ?
				self.newNexusPositions.pop(0)
			else:
				wantedBuildLocation = self.structures(UnitTypeId.NEXUS).ready.random.position  # can also use self.townhalls)


			pylonPosition = wantedBuildLocation.towards_with_random_angle(self.game_info.map_center, random.randrange(5, 15), 3.1415/3)
			await self.build(UnitTypeId.PYLON, near=pylonPosition)

	async def attackOpponent(self):
		for stalker in self.units(UnitTypeId.STALKER):
			#if I don't have many stalkers, don't do anything
			targets = (self.enemy_units | self.enemy_structures).filter(lambda unit: unit.can_be_attacked)
			if targets:
				target = targets.closest_to(stalker)
				self.do(stalker.attack(target))
			else:
				self.do(stalker.attack(self.enemy_start_locations[0]))
			if (stalker.shield_percentage < 0.5):
				await self.blinkAwayFromEnemy(stalker)

	async def sendUnitsToDefend(self):
		#locationToAttack = self.structures(UnitTypeId.NEXUS).closest_to(self.enemy_start_locations[0]).position
		for stalker in self.units(UnitTypeId.STALKER).ready:
			targets = (self.enemy_units).filter(lambda unit: unit.can_be_attacked)
		    #targets = (self.enemy_units | self.enemy_structures).filter(lambda unit: unit.can_be_attacked)
			if (stalker.shield_percentage < 0.5):
				await self.blinkAwayFromEnemy(stalker)

			if targets:
				target = targets.closest_to(stalker)
				self.do(stalker.attack(target))
			else:
				self.do(stalker.attack(self.structures(UnitTypeId.NEXUS).closest_to(self.enemy_start_locations[0]).position))

	async def blinkAwayFromEnemy(self, stalker):
		blinkAbility = await self.get_available_abilities(stalker)
		enemies = self.enemy_units
		if (enemies == None or enemies == []):
			return
		closestEnemyPosition = enemies.closest_to(stalker).position
		hopefullySafePosition = (-closestEnemyPosition[0], -closestEnemyPosition[1])
		self.do(stalker(AbilityId.EFFECT_BLINK, position.Point2(hopefullySafePosition))) #TODO: fix this!

	# async def blinkOnOpposingArmy(self, stalker): # teleports single stalkers, not the entire army.
	# 	blinkAbility = await self.get_available_abilities(stalker)
	# 	enemies = self.enemy_units
	# 	if (enemies == None or enemies == []):
	# 		return
	# 	closestEnemyPosition = enemies.closest_to(stalker).position
	# 	hopefullySafePosition = (-closestEnemyPosition[0], -closestEnemyPosition[1])
	# 	print(closestEnemyPosition)
	# 	print(hopefullySafePosition)
	# 	self.do(stalker(AbilityId.EFFECT_BLINK, enemies.closest_to(stalker).position * 8))


		# locationToAttack = self.structures(UnitTypeId.NEXUS).random.position
		# for stalker in self.units(UnitTypeId.STALKER):
		# 	self.do(stalker.attack(locationToAttack))





	async def on_step(self, iteration: int):
		if iteration >= 0:

### things that should ALWAYS run - micro, mainly
			#Always do stuff with idle workers/new workers/whatever
			await self.manageIdleWorkers()
			if (iteration % 25 == 0):
				await self.distribute_workers()

			if (self.already_pending_upgrade(UpgradeId.BLINKTECH) == 1 and self.units(UnitTypeId.STALKER).amount >= 20):
				await self.attackOpponent()
			else:
				await self.sendUnitsToDefend()

			#await self.morphGateways()


### things I want to run ASAP - HUGE priority

			# Build a cybernetics core
			if (self.structures(UnitTypeId.GATEWAY).ready and not self.structures(UnitTypeId.CYBERNETICSCORE) and self.already_pending(UnitTypeId.CYBERNETICSCORE) == 0):
				await self.buildStructure(UnitTypeId.CYBERNETICSCORE)
				return

			# Build a twilight council
			if (self.structures(UnitTypeId.CYBERNETICSCORE).ready and not self.structures(UnitTypeId.TWILIGHTCOUNCIL) and self.already_pending(UnitTypeId.TWILIGHTCOUNCIL) == 0):
				await self.buildStructure(UnitTypeId.TWILIGHTCOUNCIL)
				return


			# Research warpgate asap
			if (self.structures(UnitTypeId.CYBERNETICSCORE).ready and self.already_pending_upgrade(UpgradeId.WARPGATERESEARCH) == 0):
				await self.upgradeWarpgate() #Research WarpGate ASAP
				return

			# Upgrade blink asap
			if (self.structures(UnitTypeId.TWILIGHTCOUNCIL).ready and self.already_pending_upgrade(UpgradeId.BLINKTECH) == 0):
				await self.upgradeBlink() # Research blinktech ASAP
				return

### "other" - run whenever if it should/can.
			# Supply stuff starts - deal with supply limits/cap/whatever.
			if (self.supply_cap <= 50):
				if (self.supply_left <= 10 and self.already_pending(UnitTypeId.PYLON) < 1):
					await self.buildPylon()

			if (self.supply_cap <= 100 and self.supply_cap > 50):
				if (self.supply_left <= 20 & self.already_pending(UnitTypeId.PYLON) < 2):
					await self.buildPylon()

			if (self.supply_cap <= 199 and self.supply_cap > 100):
				if (self.supply_left <= 25 & (self.already_pending(UnitTypeId.PYLON) < 3) & (self.supply_cap != 200)):
					await self.buildPylon()
			#Supply stuff ends



			#Build 2 gateways per base expand I have.
			if (self.townhalls.amount * 2 > self.structures(UnitTypeId.GATEWAY).amount + self.structures(UnitTypeId.WARPGATE).amount):
				await self.buildStructure(UnitTypeId.GATEWAY) # Only build 2 gateways per base

			#Train probes, 19 per base (1 refinery per base + 16 mineral patches = 19 probes per expo)
			if (self.townhalls.amount * 19 > self.supply_workers):
				await self.trainWorker(UnitTypeId.NEXUS, UnitTypeId.PROBE)

				# Build 1 assimilator per base - not necessarily at that base.
			if (self.townhalls.amount > self.structures(UnitTypeId.ASSIMILATOR).amount):
				await self.buildGeyser(UnitTypeId.ASSIMILATOR)

			# Build/warp in stalkers whenever possible.
			await self.trainUnits(UnitTypeId.GATEWAY, UnitTypeId.STALKER)
			await self.warpUnits(AbilityId.WARPGATETRAIN_STALKER, UnitTypeId.STALKER)

			# Build 1 nexus at a time. #TODO: Only build 1 nexus at once. (never double/triple expand at the same time)
			if (self.structures(UnitTypeId.NEXUS).amount <= 3): # & self.already_pending(UnitTypeId.NEXUS) < 1
				await self.buildExpansion(UnitTypeId.NEXUS) # Build up to 4 bases, whenever I have enough minerals banked for them.























sc2.run_game(sc2.maps.get("AcropolisLE"), [
    Bot(Race.Protoss, AlphaWannabe()),
    Computer(Race.Protoss, Difficulty.Hard)
], realtime=False)

# force an earlier expo - like 1 or 2 minutes in
# Chrono boost
# I always save 150 mineralss at the start?
# efficiently micro probes at the very start of the game
# maximize pylon coverage ?
# wall off?
# invisibilility detection? (DTs)
# build DTs
# Zealot runbies? STALKER RUNBIES?
# warp prism micro too?
# WARP PRISM DROP ON THEIR THIRD? owo
# teleport in opponent's army if I have more supply of an army? (beat siege tanks faster)
# don't trap my own units in my base
# scout for proxies/any secrets
# if I've basically won for sure, look around hte map and kill ANYTHING else
# build Zealots early game
# scout, if opponent is going to all-in early prepare for it
# bm opponent with nice emotes
# limit number of expansions
# one expo at once, only if fully saturated gas + minerals ( bases & 19)
# blink towards opponents base / one for tank vision, then on top of tanks
# allow saving-up for stuff (if whatever, return - don't run any other code that spends minerals!
# only build an expo if one isn't already building

		# for building pylons, at 84 supply, switch it to be the other way - building more pylons causes mass pylon production lmao.
		# besides that, add more to the sendUnitsToDefend function - if I need to teleport away, teleport away from the enemies.
		# teleport away from the CLOSEST ENEMY!
		# implement stutter-stepping with stalkers
		# if I'm losing the fight, retreat

# Expand/whatever, if I have less than perfect mineral saturation (e.g if a mineral field dissapears)
# Build the assimilator at the new base - not the first/second/third base, etc. --> add that to the newNexusPositions array ? check if enough minerals for both, then do?
# SendUnitsToDefend --> only defend if it's within X distance of a friendly building/nexus/whatever.
# Mass recall!
# chase enemies function
# don't warp in at damaged pylon (probably dying)
# blink into opposing army
# fix "can't find warp-in location"
# get surrounds / concave on the opponent
# fix teleporting to use my position - enemy position or however that works.exe