#include "BananaBrain.h"

void PathFinder::init()
{
	if (Broodwar->mapHash() == Strategy::kMapHash_Plasma) {
		for (auto& area : bwem_map.Areas()) {
			for (auto& cp : area.ChokePoints()) {
				if (cp->GetAreas().first == &area) {
					Position center = chokepoint_center(cp);
					bool blocked_by_eggs = false;
					std::map<Unit,int> mineral_distances;
					for (auto& unit : Broodwar->getStaticNeutralUnits()) {
						if (!blocked_by_eggs && unit->getType() == UnitTypes::Zerg_Egg && unit->getDistance(center) < 100) {
							blocked_by_eggs = true;
						}
						if (unit->getType().isMineralField() && unit->getResources() == 32) {
							mineral_distances[unit] = unit->getDistance(center);
						}
					}
					if (blocked_by_eggs && mineral_distances.size() >= 2) {
						std::vector<Unit> minerals_sorted = keys_sorted(mineral_distances);
						Unit first_mineral = minerals_sorted[0];
						Unit second_mineral = minerals_sorted[1];
						
						bool forward;
						
						const BWEM::Area* first_mineral_area = bwem_map.GetNearestArea(WalkPosition(first_mineral->getPosition()));
						const BWEM::Area* second_mineral_area = bwem_map.GetNearestArea(WalkPosition(second_mineral->getPosition()));
						if (cp->GetAreas().first == first_mineral_area && cp->GetAreas().second == second_mineral_area) {
							forward = true;
						} else if (cp->GetAreas().first == second_mineral_area && cp->GetAreas().second == first_mineral_area) {
							forward = false;
						} else {
							Position first_area_top(cp->GetAreas().first->Top());
							Position second_area_top(cp->GetAreas().second->Top());
							forward = (second_area_top.getApproxDistance(second_mineral->getPosition()) < first_area_top.getApproxDistance(second_mineral->getPosition()));
						}
						
						MineralWalkData& data = mineral_walk_data_[cp];
						if (forward) {
							data.first_mineral = first_mineral;
							data.second_mineral = second_mineral;
						} else {
							data.first_mineral = second_mineral;
							data.second_mineral = first_mineral;
						}
						data.first_mineral_position = calculate_position_near_egg(data.first_mineral);
						data.second_mineral_position = calculate_position_near_egg(data.second_mineral);
					}
				}
			}
		}
	}
	
	if (Broodwar->mapHash() == Strategy::kMapHash_Hitchhiker) {
		std::vector<const BWEM::ChokePoint*> chokepoints;
		for (auto& area : bwem_map.Areas()) {
			for (auto& cp : area.ChokePoints()) {
				if (!cp->Blocked() && cp->GetAreas().first == &area) {
					if (chokepoint_center(cp) == Position(3876, 1836)) {
						chokepoints.push_back(cp);
					}
				}
			}
		}
		if (!chokepoints.empty()) {
			bwem_map.BlockChokepoints(chokepoints);
		}
	}
	
	if (Broodwar->mapHash() == Strategy::kMapHash_GoldRush) {
		std::vector<const BWEM::ChokePoint*> chokepoints;
		for (auto& area : bwem_map.Areas()) {
			for (auto& cp : area.ChokePoints()) {
				if (!cp->Blocked() && cp->GetAreas().first == &area) {
					Position position = chokepoint_center(cp);
					if (position == Position(428, 3412) ||
						position == Position(1220, 92) ||
						position == Position(3892, 2444) ||
						position == Position(1956, 3900) ||
						position == Position(180, 780) ||
						position == Position(3884, 1108)) {
						chokepoints.push_back(cp);
					}
				}
			}
		}
		if (!chokepoints.empty()) {
			bwem_map.BlockChokepoints(chokepoints);
		}
	}
}

Position PathFinder::calculate_position_near_egg(Unit mineral_unit)
{
	int egg_distance = INT_MAX;
	Unit egg_unit = nullptr;
	for (auto& unit : Broodwar->getStaticNeutralUnits()) {
		if (unit->getType() == UnitTypes::Zerg_Egg) {
			int distance = unit->getDistance(mineral_unit);
			if (distance < egg_distance) {
				egg_unit = unit;
				egg_distance = distance;
			}
		}
	}
	
	Position result;
	if (egg_unit == nullptr) {
		result = mineral_unit->getPosition();
	} else {
		Gap gap(Rect(mineral_unit->getType(), mineral_unit->getPosition()), Rect(egg_unit->getType(), egg_unit->getPosition()));
		result = gap.b();
	}
	return result;
}

Position PathFinder::first_common_path_position(Position position,std::vector<Position> destinations)
{
	if (destinations.empty()) return position;
	
	std::vector<const BWEM::CPPath*> paths;
	for (Position destination : destinations) {
		int distance = -1;
		const BWEM::CPPath& path = bwem_map.GetPath(position, destination, &distance);
		if (distance < 0) return position;
		paths.push_back(&path);
	}
	
	Position result = position;
	const BWEM::CPPath& path0 = *paths[0];
	for (size_t i = 0; i < path0.size(); i++) {
		bool match = true;
		for (size_t j = 1; j < paths.size(); j++) {
			const BWEM::CPPath& path = *paths[j];
			if (i >= path.size() || path0[i] != path[i]) {
				match = false;
				break;
			}
		}
		if (match) {
			const BWEM::ChokePoint* cp = path0[i];
			result = center_position(cp->Center());
		} else {
			break;
		}
	}
	return result;
}

bool PathFinder::execute_path(Unit unit,Position position,std::function<void(void)> command)
{
	if (unit->getType().isWorker() && mineral_walk_continue(unit)) return true;
	
	if (unit->getType().isWorker()) {
		auto mineral_walk_state_it = mineral_walk_state_.find(unit);
		if (mineral_walk_state_it != mineral_walk_state_.end()) {
			MineralWalkState& mineral_walk_state = mineral_walk_state_it->second;
			if (!mineral_walk_state.done) {
				if (Broodwar->getFrameCount() >= mineral_walk_state.expire_frame) {
					mineral_walk_state.done = true;
					mineral_walk_state.expire_frame = Broodwar->getFrameCount() + 90;
				} else {
					if (mineral_walk_to(unit, mineral_walk_state.target, mineral_walk_state.position)) {
						return true;
					} else {
						mineral_walk_state.done = true;
						mineral_walk_state.expire_frame = Broodwar->getFrameCount() + 90;
					}
				}
			} else {
				if (Broodwar->getFrameCount() >= mineral_walk_state.expire_frame) {
					mineral_walk_state_.erase(unit);
				}
			}
		}
	}
	
	if (!unit->isFlying() &&
		has_area(WalkPosition(unit->getPosition())) &&
		has_area(WalkPosition(position))) {
		int distance = -1;
		const BWEM::CPPath& path = bwem_map.GetPath(unit->getPosition(), position, &distance);
		if (distance < 0) return false;
		
		for (auto& choke : path) {
			if (unit->getType().isWorker()) {
				if (mineral_walk_start(unit, choke, center_position(choke->Center()))) return true;
			}
			
			if (distance_to_chokepoint(unit->getPosition(), choke) > 320) {
				unit_move(unit, center_position(choke->Center()));
				return true;
			}
		}
	}
	
	command();
	return true;
}

int PathFinder::distance_to_chokepoint(Position position,const BWEM::ChokePoint* choke)
{
	int center_distance = position.getApproxDistance(center_position(choke->Center()));
	int end1_distance = position.getApproxDistance(center_position(choke->Pos(BWEM::ChokePoint::end1)));
	int end2_distance = position.getApproxDistance(center_position(choke->Pos(BWEM::ChokePoint::end2)));
	return std::min(center_distance, std::min(end1_distance, end2_distance));
}

bool PathFinder::mineral_walk_continue(Unit unit)
{
	bool result = false;
	
	auto mineral_walk_state_it = mineral_walk_state_.find(unit);
	if (mineral_walk_state_it != mineral_walk_state_.end()) {
		MineralWalkState& mineral_walk_state = mineral_walk_state_it->second;
		if (!mineral_walk_state.done) {
			if (Broodwar->getFrameCount() >= mineral_walk_state.expire_frame) {
				mineral_walk_state.done = true;
				mineral_walk_state.expire_frame = Broodwar->getFrameCount() + 90;
			} else {
				if (mineral_walk_to(unit, mineral_walk_state.target, mineral_walk_state.position)) {
					result = true;
				} else {
					mineral_walk_state.done = true;
					mineral_walk_state.expire_frame = Broodwar->getFrameCount() + 90;
				}
			}
		} else {
			if (Broodwar->getFrameCount() >= mineral_walk_state.expire_frame) {
				mineral_walk_state_.erase(unit);
			}
		}
	}
	
	return result;
}

bool PathFinder::mineral_walk_start(Unit unit,const BWEM::ChokePoint* choke,Position waypoint)
{
	bool result = false;
	
	if (mineral_walk_data_.count(choke) > 0 &&
		mineral_walk_state_.count(unit) == 0) {
		const MineralWalkData& data = mineral_walk_data_.at(choke);
		const BWEM::Area* area = bwem_map.GetNearestArea(WalkPosition(unit->getPosition()));
		if (area == choke->GetAreas().first || area == choke->GetAreas().second) {
			auto this_side = (area == choke->GetAreas().first) ? std::make_pair(data.first_mineral, data.first_mineral_position) : std::make_pair(data.second_mineral, data.second_mineral_position);
			auto other_side = (area == choke->GetAreas().second) ? std::make_pair(data.first_mineral, data.first_mineral_position) : std::make_pair(data.second_mineral, data.second_mineral_position);
			if (!other_side.first->exists()) {
				unit_move(unit, this_side.second);
				result = true;
			} else if (mineral_walk_to(unit, other_side.first, this_side.second)) {
				MineralWalkState& mineral_walk_state = mineral_walk_state_[unit];
				mineral_walk_state.expire_frame = Broodwar->getFrameCount() + 240;
				mineral_walk_state.target = other_side.first;
				mineral_walk_state.position = this_side.second;
				result = true;
			}
		}
	}
	
	return result;
}

bool PathFinder::mineral_walk_to(Unit unit,Unit mineral,Position mineral_position)
{
	bool result = true;
	if (mineral->exists()) {
		if (unit->getDistance(mineral) >= 8) {
			if (!unit_has_target(unit, mineral)) unit->gather(mineral);
		} else {
			result = false;
		}
	} else {
		unit_move(unit, mineral_position);
	}
	return result;
}

void PathFinder::close_small_chokepoints_if_needed()
{
	if (!small_chokepoints_closed_ &&
		(training_manager.unit_count_completed(UnitTypes::Protoss_Dragoon) > 0 ||
		 training_manager.unit_count_completed(UnitTypes::Protoss_Reaver) > 0 ||
		 training_manager.unit_count_completed(UnitTypes::Protoss_Archon) > 0 ||
		 training_manager.unit_count_completed(UnitTypes::Protoss_Dark_Archon) > 0 ||
		 training_manager.unit_count_completed(UnitTypes::Terran_Vulture) > 0 ||
		 training_manager.unit_count_completed(UnitTypes::Terran_Siege_Tank_Tank_Mode) > 0 ||
		 training_manager.unit_count_completed(UnitTypes::Terran_Siege_Tank_Siege_Mode) > 0 ||
		 training_manager.unit_count_completed(UnitTypes::Terran_Goliath) > 0 ||
		 training_manager.unit_count_completed(UnitTypes::Zerg_Lurker) > 0 ||
		 training_manager.unit_count_completed(UnitTypes::Zerg_Ultralisk) > 0)) {
			std::vector<const BWEM::ChokePoint*> chokepoints;
			for (auto& area : bwem_map.Areas()) {
				for (auto& cp : area.ChokePoints()) {
					if (!cp->Blocked() && cp->GetAreas().first == &area) {
						Gap gap(cp);
						if (!gap.fits_through(UnitTypes::Protoss_Dragoon) ||
							!gap.fits_through(UnitTypes::Protoss_Reaver)) {
							chokepoints.push_back(cp);
						}
					}
				}
			}
			if (!chokepoints.empty()) {
				bwem_map.BlockChokepoints(chokepoints);
				if (configuration.draw_enabled()) {
					Broodwar << "Closed " << (int)chokepoints.size() << " chokepoints" << std::endl;
				}
			}
			small_chokepoints_closed_ = true;
		}
}
