#pragma once
#include "PlayerData.h"
#include "../CommonValues.h"

namespace RLGSC {
	struct ScoreLine {
		int teamGoals[2] = { 0,0 };

		int operator[](size_t index) const {
			return teamGoals[index];
		}

		int& operator[](size_t index) {
			return teamGoals[index];
		}
	};

	// https://github.com/AechPro/rocket-league-gym-sim/blob/main/rlgym_sim/utils/gamestates/game_state.py
	struct GameState {
		float deltaTime = 0; // Time that has passed since last update

		ScoreLine scoreLine;
		int lastTouchCarID = -1;
		std::vector<PlayerData> players;

		BallState ballState;
		PhysObj ball, ballInv;

		std::array<bool, CommonValues::BOOST_LOCATIONS_AMOUNT> boostPads;
		std::array<bool, CommonValues::BOOST_LOCATIONS_AMOUNT> boostPadsInv;

		std::array<float, CommonValues::BOOST_LOCATIONS_AMOUNT> boostPadTimers;
		std::array<float, CommonValues::BOOST_LOCATIONS_AMOUNT> boostPadTimersInv;

		// Last arena we updated with
		// Can be used to determine current arena from within reward function, for example
		// NOTE: Could be null
		Arena* lastArena = NULL;

		// Last tick count when updated
		uint64_t lastTickCount = 0;

		GameState() = default;
		explicit GameState(Arena* arena) {
			UpdateFromArena(arena);
		}

		const PhysObj& GetBallPhys(bool inverted) const {
			return inverted ? ballInv : ball;
		}

		const auto& GetBoostPads(bool inverted) const {
			return inverted ? boostPadsInv : boostPads;
		}

		void UpdateFromArena(Arena* arena);
	};
} // GlobalBallPred.h
#pragma once

#include <unordered_map>
#include <functional>
#include <vector>
#include <memory>
#include <mutex>

// Include necessary project-specific headers
#include "C:\Epsilon\RLGymPPO_CPP-main\RLGymPPO_CPP\RLGymSim_CPP\RocketSim\src\Sim\Arena\Arena.h"
#include "C:\Epsilon\RLGymPPO_CPP-main\RLGymPPO_CPP\RLGymSim_CPP\src\RLGymSim_CPP\Utils\GlobalBallPred\BallPredTrackerManager.h"
#include "C:\Epsilon\RLGymPPO_CPP-main\RLGymPPO_CPP\RLGymSim_CPP\src\RLGymSim_CPP\Utils\Gamestates\PhysObj.h"



RS_NS_START

struct BallPredEnv {
    BallPredTracker ballPredTracker;

    BallPredEnv(const Arena* arena) :
        ballPredTracker(const_cast<Arena*>(arena), 240) {}

    void Update(const Arena* arena) {
        ballPredTracker.UpdatePredFromArena(const_cast<Arena*>(arena));
    }

    PhysObj GetAtTime(float time, const PhysObj& curBall);
};

class BallPred {
public:
    static BallPredEnv* GetOrMake(const Arena* arena);

    static PhysObj GetAtTime(const Arena* arena, float time, const PhysObj& curBall) {
        return GetOrMake(arena)->GetAtTime(time, curBall);
    }

    static void Update(const Arena* arena) {
        GetOrMake(arena)->Update(arena);
    }
};

RS_NS_END using namespace RLGPC; // RLGymPPO
using namespace RLGSC; // RLGymSim

static std::vector<std::string> g_RewardNames = {};
MetricLoggerSystem g_MetricLoggerSystem;


void OnStep(GameInst* gameInst, const RLGSC::Gym::StepResult& stepResult, Report& gameMetrics) {
	if (auto* combRewards = dynamic_cast<LogCombinedReward*>(gameInst->match->rewardFn)) {
		for (size_t i = 0; i < combRewards->rewardFuncs.size(); i++) {
			float reward = combRewards->lastRewards[i] * combRewards->rewardWeights[i];
			gameMetrics.AccumAvg(combRewards->GetRewardFuncName(i), reward);
		}
	}

	// Fix: Pass stepResult directly instead of stepResult.state
	g_MetricLoggerSystem.LogStep(gameInst, stepResult, gameMetrics);

	// **Update the ball prediction here as well**
	BallPred::Update(gameInst->gym->arena);
}


void OnIteration(Learner* learner, Report& allMetrics) {
	auto allGameMetrics = learner->GetAllGameMetrics();
	g_MetricLoggerSystem.SummarizeIteration(allGameMetrics, allMetrics);

	// Handle rewards
	std::vector<AvgTracker> avgRews(g_RewardNames.size());
	for (size_t i = 0; i < g_RewardNames.size(); i++) {
		for (const auto& gameReport : allGameMetrics) {
			if (!gameReport.data.empty()) {
				avgRews[i] += gameReport.GetAvg(g_RewardNames[i]);
			}
		}
		allMetrics["Rewards/" + g_RewardNames[i]] = avgRews[i].Get();
	}
}




EnvCreateResult EnvCreateFunc() {
	constexpr int TICK_SKIP = 12;
	constexpr float NO_TOUCH_TIMEOUT_SECS = 8.f;
	int team_size = 1;
	std::cout << "Team Size: " << team_size << std::endl;

	const auto PZSUM = [](RewardFunction* rewardFunc, float teamSpirit = 1.0f, float opponentScale = 1.0f) {
		return new ZeroSumReward(rewardFunc, teamSpirit, opponentScale);
		};

	auto obs = new DefaultOBS();
	auto actionParser = new DiscreteAction();
	auto stateSetter = new WeightedSampleSetter(
		{
			{new RandomState(true, true, false), 0.9f},
			{new KickoffState(), 0.1f},
			{new AirDribbleSetup(), 0.0f},
			{new ReplaySetter(team_size), 0.0f},
		}
		);

	std::vector<TerminalCondition*> terminalConditions = {
		new NoTouchCondition(NO_TOUCH_TIMEOUT_SECS * 120 / TICK_SKIP),
		new GoalScoreCondition()
	};

	// Create a temporary reward function
	RewardFunction* tempReward;

	Match* match = new Match(
		tempReward,
		terminalConditions,
		obs,
		actionParser,
		stateSetter,
		team_size,
		true
	);

	Gym* gym = new Gym(match, TICK_SKIP);
	Arena* arena = gym->arena;



	// ======================== REWARDS ============================//
	
	auto rewards = new LogCombinedReward(
		std::vector<std::pair<RewardFunction*, float>>{

			{ PZSUM(new PossessionReward(), 0.0f, 0.25f), 1.0f },
	}
	);

	// Replace the temporary reward function with the actual one
	delete match->rewardFn;
	match->rewardFn = rewards;

	{
		static std::once_flag onceFlag;
		std::call_once(onceFlag, [&] {
			for (int i = 0; i < rewards->rewardFuncs.size(); i++)
				g_RewardNames.push_back(rewards->GetRewardFuncName(i));
			});
	}

	return { match, gym };
}

int main(int argc, char* argv[]) {
	// Initialize RocketSim with collision meshes
	RocketSim::Init("./collision_meshes");


	// Make configuration for the learner
	LearnerConfig cfg = {};
	cfg.renderMode = false;

	for (int i = 1; i < argc; ++i) {
		if (std::strcmp(argv[i], "--render") == 0) {
			cfg.renderMode = true;
			cfg.renderTimeScale = 1.0f;
			cfg.deterministic = true;
		}
	}

	// Play around with these to see what the optimal is for your machine, more isn't always better
	cfg.numThreads = 16;
	cfg.numGamesPerThread = 24;

	// We want a large itr/batch size
	// You'll want to increase this as your bot improves, up to an extent
	int tsPerItr = 100 * 1000;
	cfg.timestepsPerIteration = tsPerItr;
	cfg.ppo.batchSize = 100000;
	cfg.ppo.miniBatchSize = 25000; // Lower this if too much VRAM is being allocated
	cfg.expBufferSize = tsPerItr * 2;

	// This is just set to 1 to match rlgym-ppo example
	// I've found the best value is somewhere between 2 and 4
	// Increasing this will lower SPS, but increase step efficiency
	cfg.ppo.epochs = 3;

	// Reasonable starting entropy
	cfg.ppo.entCoef = 0.01f;

	cfg.skillTrackerConfig.enabled = true;
	cfg.timestepsPerSave = 2'000'000;
	cfg.checkpointsToKeep = 30;
	cfg.skillTrackerConfig.numEnvs = 16;
	cfg.skillTrackerConfig.numThreads = 16;
	cfg.skillTrackerConfig.perModeRatings = true;

	cfg.metricsRunName = "RUN6";
	cfg.metricsProjectName = "MAIN2";
	cfg.metricsGroupName = "RUN";


	// Good learning rate to start, may start to be too high around 100m steps
	cfg.ppo.policyLR = 2e-4; // 1e-4 -> 8e-5 -> 4e-5 -> 2e-5 ->1e-5
	cfg.ppo.criticLR = 2e-4;

	// Default model size
	cfg.ppo.policyLayerSizes = { 512, 512, 512 };
	cfg.ppo.criticLayerSizes = { 512, 512, 512 }; // tomorrow 1548, 1024, 512, 512, 256

	//cfg.ppo.policyLayerSizes = { 1024, 1024, 1024, 1024, 1024 };
	//cfg.ppo.criticLayerSizes = { 1024, 1024, 1024, 1024, 1024 };

	cfg.sendMetrics = true; // Send metrics

	// Make the learner with the environment creation function and the config we just made
	Learner learner = Learner(EnvCreateFunc, cfg);

	learner.stepCallback = OnStep;
	learner.iterationCallback = OnIteration;

	// Start learning!
	learner.Learn();

	return 0;
} #include "Gym.h"

namespace RLGSC {

	template<int PlayerData::* DATA_VAR>
	void IncPlayerCounter(Car* car, void* userInfo) {
		if (!car)
			return;

		Gym* gym = (Gym*)userInfo;
		for (auto& player : gym->prevState.players)
			if (player.carId == car->id)
				(player.*DATA_VAR)++;
	}

	void _ShotEventCallback(Arena* arena, Car* shooter, Car* passer, void* userInfo) {
		IncPlayerCounter<&PlayerData::matchShots>(shooter, userInfo);
		IncPlayerCounter<&PlayerData::matchShotPasses>(passer, userInfo);
	}

	void _GoalEventCallback(Arena* arena, Car* scorer, Car* passer, void* userInfo) {
		IncPlayerCounter<&PlayerData::matchGoals>(scorer, userInfo);
		IncPlayerCounter<&PlayerData::matchAssists>(passer, userInfo);
	}

	void _SaveEventCallback(Arena* arena, Car* saver, void* userInfo) {
		IncPlayerCounter<&PlayerData::matchSaves>(saver, userInfo);
	}

	void _BumpCallback(Arena* arena, Car* bumper, Car* victim, bool isDemo, void* userInfo) {
		if (bumper->team == victim->team)
			return;

		IncPlayerCounter<&PlayerData::matchBumps>(bumper, userInfo);

		if (isDemo)
			IncPlayerCounter<&PlayerData::matchDemos>(bumper, userInfo);
	}

	Gym::Gym(Match* match, int tickSkip, CarConfig carConfig, GameMode gameMode, MutatorConfig mutatorConfig) :
		match(match), tickSkip(tickSkip) {
		arena = Arena::Create(gameMode);
		arena->SetMutatorConfig(mutatorConfig);

		for (int i = 0; i < match->teamSize; i++) {
			arena->AddCar(Team::BLUE, carConfig);
			if (match->spawnOpponents)
				arena->AddCar(Team::ORANGE, carConfig);
		}

		eventTracker.SetShotCallback(_ShotEventCallback, this);
		eventTracker.SetGoalCallback(_GoalEventCallback, this);
		eventTracker.SetSaveCallback(_SaveEventCallback, this);

		arena->SetCarBumpCallback(_BumpCallback, this);
	}

	FList2 Gym::Reset() {
		GameState resetState = match->ResetState(arena);
		match->EpisodeReset(resetState);
		prevState = resetState;
		eventTracker.ResetPersistentInfo();

		FList2 obs = match->BuildObservations(resetState);
		return obs;
	}

	Gym::StepResult Gym::Step(const ActionParser::Input& actionsData) {
		ActionSet actions = match->ParseActions(actionsData, prevState);
		match->prevActions = actions;

		GameState state;

		{ // Step arena with actions
			auto carItr = arena->_cars.begin();
			for (int i = 0; i < actions.size(); i++) {
				(*carItr)->controls = (CarControls)actions[i];
				carItr++;
			}

			arena->Step(1);
			if (arena->gameMode != GameMode::HEATSEEKER)
				eventTracker.Update(arena);
			state = prevState; // All callbacks have been hit
			state.UpdateFromArena(arena);
			arena->Step(tickSkip - 1);
			totalTicks += tickSkip;
			totalSteps++;
		}

		FList2 obs = match->BuildObservations(state);
		bool done = match->IsDone(state);
		FList rewards = match->GetRewards(state, done);
		prevState = state;

		return StepResult{
			obs,
			rewards,
			done,
			state
		};
	}
} #include "Match.h"

namespace RLGSC {
	void Match::EpisodeReset(const GameState& initialState) {
		prevActions = ActionSet(initialState.players.size());
		for (auto cond : terminalConditions)
			cond->Reset(initialState);
		rewardFn->Reset(initialState);
		obsBuilder->Reset(initialState);
	}

	FList2 Match::BuildObservations(const GameState& state) {
		auto result = FList2(state.players.size());

		obsBuilder->PreStep(state);

		for (int i = 0; i < state.players.size(); i++) {
			result[i] =
				obsBuilder->BuildOBS(state.players[i], state, prevActions[i]);
		}

		return result;
	}

	FList Match::GetRewards(const GameState& state, bool done) {
		auto result = FList(state.players.size());

		rewardFn->PreStep(state);
		return rewardFn->GetAllRewards(state, prevActions, done);
	}

	bool Match::IsDone(const GameState& state) {
		for (auto& cond : terminalConditions)
			if (cond->IsTerminal(state))
				return true;

		return false;
	}

	ScoreLine Match::GetScoreLine(const GameState& state) {
		return state.scoreLine;
	}

	ActionSet Match::ParseActions(const ActionParser::Input& actionsData, const GameState& gameState) {
		ActionSet actions = actionParser->ParseActions(actionsData, gameState);

		for (int i = 0; i < gameState.players.size(); i++)
			if (gameState.players[i].carState.isDemoed)
				actions[i] = {};

		return actions;
	}

	GameState Match::ResetState(Arena* arena) {
		GameState newState = stateSetter->ResetState(arena);

		if (newState.players.size() != playerAmount) {
			RG_ERR_CLOSE(
				"Match::ResetState(): New state has a different amount of players, "
				"expected " << playerAmount << " but got " << newState.players.size() << ".\n"
				"Changing number of players at state reset is currently not supported.\n" <<
				"If you want variable player amounts, set a differing player amount per env."
			);
		}

		for (auto& pad : arena->_boostPads)
			pad->SetState({});

		return newState;
	}
}#pragma once
#include "../../BaseInc.h"
#include "../Car/Car.h"
#include "../Ball/Ball.h"
#include "../BoostPad/BoostPad.h"
#include "../CollisionMasks.h"

#include "../../CollisionMeshFile/CollisionMeshFile.h"
#include "../BoostPad/BoostPadGrid/BoostPadGrid.h"
#include "../SuspensionCollisionGrid/SuspensionCollisionGrid.h"
#include "../MutatorConfig/MutatorConfig.h"
#include "ArenaConfig/ArenaConfig.h"

#include "../../../libsrc/bullet3-3.24/BulletCollision/BroadphaseCollision/btDbvtBroadphase.h"
#include "../../../libsrc/bullet3-3.24/BulletCollision/CollisionShapes/btStaticPlaneShape.h"
#include "../../../libsrc/bullet3-3.24/BulletCollision/CollisionShapes/btBvhTriangleMeshShape.h"
#include "../../../libsrc/bullet3-3.24/BulletCollision/CollisionDispatch/btCollisionDispatcher.h"
#include "../../../libsrc/bullet3-3.24/BulletDynamics/ConstraintSolver/btSequentialImpulseConstraintSolver.h"
#include "../../../libsrc/bullet3-3.24/BulletDynamics/Dynamics/btDiscreteDynamicsWorld.h"
#include "../../../libsrc/bullet3-3.24/BulletCollision/CollisionDispatch/btDefaultCollisionConfiguration.h"

RS_NS_START

typedef std::function<void(class Arena* arena, Team scoringTeam, void* userInfo)> GoalScoreEventFn;
typedef std::function<void(class Arena* arena, Car* bumper, Car* victim, bool isDemo, void* userInfo)> CarBumpEventFn;

// The container for all game simulation
// Stores cars, the ball, all arena collisions, and manages the overall game state
class Arena {
public:

	GameMode gameMode;

	uint32_t _lastCarID = 0;
	std::unordered_set<Car*> _cars;
	bool ownsCars = true; // If true, deleting this arena instance deletes all cars

	std::unordered_map<uint32_t, Car*> _carIDMap;
	
	Ball* ball;
	bool ownsBall = true; // If true, deleting this arena instance deletes the ball
	
	std::vector<BoostPad*> _boostPads;
	bool ownsBoostPads = true; // If true, deleing this arena instance deletes all boost pads
	
	BoostPadGrid _boostPadGrid;

	SuspensionCollisionGrid _suspColGrid;

	MutatorConfig _mutatorConfig;

	const MutatorConfig& GetMutatorConfig() { return _mutatorConfig; }
	RSAPI void SetMutatorConfig(const MutatorConfig& mutatorConfig);

	// Time in seconds each tick (1/tickrate)
	float tickTime; 

	// Returns (1 / tickTime)
	float GetTickRate() const {
		return 1 / tickTime;
	}

	// Total ticks this arena instance has been simulated for, never resets
	uint64_t tickCount = 0;

	const std::unordered_set<Car*>& GetCars() { return _cars; }
	const std::vector<BoostPad*>& GetBoostPads() { return _boostPads; }

	// Returns true if added, false if car was already added
	bool _AddCarFromPtr(Car* car);
	RSAPI Car* AddCar(Team team, const CarConfig& config = CAR_CONFIG_OCTANE);

	// Returns false if the car ID was not found in the cars list
	RSAPI bool RemoveCar(uint32_t id);

	// Returns false if the car was not found in the cars list
	// NOTE: If the car was removed, the car will be freed and the pointer will be made invalid
	bool RemoveCar(Car* car) {
		return RemoveCar(car->id);
	}

	RSAPI Car* GetCar(uint32_t id);

	btDiscreteDynamicsWorld _bulletWorld;
	struct {
		btDefaultCollisionConfiguration collisionConfig;
		btCollisionDispatcher collisionDispatcher;
		btOverlappingPairCache* overlappingPairCache;
		btBroadphaseInterface* broadphase;
		btSequentialImpulseConstraintSolver constraintSolver;
	} _bulletWorldParams;

	btRigidBody* _worldCollisionRBs = NULL;
	size_t _worldCollisionRBAmount = 0;
	btBvhTriangleMeshShape* _worldCollisionBvhShapes = NULL;
	btStaticPlaneShape* _worldCollisionPlaneShapes = NULL;

	struct {
		GoalScoreEventFn func = NULL;
		void* userInfo = NULL;
	} _goalScoreCallback;
	RSAPI void SetGoalScoreCallback(GoalScoreEventFn callbackFn, void* userInfo = NULL);

	struct {
		CarBumpEventFn func = NULL;
		void* userInfo = NULL;
	} _carBumpCallback;
	RSAPI void SetCarBumpCallback(CarBumpEventFn callbackFn, void* userInfo = NULL);

	// NOTE: Arena should be destroyed after use
	RSAPI static Arena* Create(GameMode gameMode, const ArenaConfig& arenaConfig = {}, float tickRate = 120);
	
	// Serialize entire arena state including cars, ball, and boostpads
	RSAPI void Serialize(DataStreamOut& out) const;

	// Load new arena from serialized data
	RSAPI static Arena* DeserializeNew(DataStreamIn& in);

	Arena(const Arena& other) = delete; // No copy constructor, use Arena::Clone() instead
	Arena& operator =(const Arena& other) = delete; // No copy operator, use Arena::Clone() instead

	Arena(Arena&& other) = delete; // No move constructor
	Arena& operator =(Arena&& other) = delete; // No move operator

	// Get a deep copy of the arena
	RSAPI Arena* Clone(bool copyCallbacks);

	// NOTE: Car ID will not be restored
	RSAPI Car* DeserializeNewCar(DataStreamIn& in, Team team);

	// Simulate everything in the arena for a given number of ticks
	RSAPI void Step(int ticksToSimulate = 1);

	RSAPI void ResetToRandomKickoff(int seed = -1);

	// Returns true if the ball is probably going in, does not account for wall or ceiling bounces
	// NOTE: Purposefully overestimates, just like the real RL's shot prediction
	// To check which goal it will score in, use the ball's velocity
	// Margin can be manually adjusted with extraMargin (negative to prevent overestimating)
	RSAPI bool IsBallProbablyGoingIn(float maxTime = 2.f, float extraMargin = 0, Team* goalTeamOut = NULL) const;

	// Returns true if the ball is in the net
	// Works for all gamemodes (and does nothing in THE_VOID)
	RSAPI bool IsBallScored() const;

	// Free all associated memory
	RSAPI ~Arena();

	// NOTE: Passed shape pointer will be freed when arena is deconstructed
	template <class T>
	btRigidBody* _AddStaticCollisionShape(
		size_t rbIndex, size_t meshListIndex, T* shape, T* meshList, btVector3 posBT = btVector3(0, 0, 0), 
		bool isHoopsNet = false) {

		static_assert(std::is_base_of<btCollisionShape, T>::value);
		meshList[meshListIndex] = *shape;

		assert(rbIndex < _worldCollisionRBAmount);
		btRigidBody& shapeRB = _worldCollisionRBs[rbIndex];
		shapeRB = btRigidBody(0, NULL, &meshList[meshListIndex]);
		shapeRB.setWorldTransform(btTransform(btMatrix3x3::getIdentity(), posBT));
		shapeRB.setUserPointer(this);
		if (isHoopsNet) {
			_bulletWorld.addRigidBody(&shapeRB, CollisionMasks::HOOPS_NET, CollisionMasks::HOOPS_NET);
		} else {
			_bulletWorld.addRigidBody(&shapeRB);
		}
		return &shapeRB;
	}

	void _SetupArenaCollisionShapes();

	// Static function called by Bullet internally when adding a collision point
	static bool _BulletContactAddedCallback(
		btManifoldPoint& cp,
		const btCollisionObjectWrapper* colObjA, int partID_A, int indexA,
		const btCollisionObjectWrapper* colObjB, int partID_B, int indexB
	);

	void _BtCallback_OnCarBallCollision(Car* car, Ball* ball, btManifoldPoint& manifoldPoint, bool ballIsBodyA);
	void _BtCallback_OnCarCarCollision(Car* car1, Car* car2, btManifoldPoint& manifoldPoint);
	void _BtCallback_OnCarWorldCollision(Car* car, btCollisionObject* worldObject, btManifoldPoint& manifoldPoint);

	const ArenaConfig& GetArenaConfig() const {
		return _config;
	}

	// Backwards compatability
	ArenaMemWeightMode GetMemWeightMode() {
		return _config.memWeightMode;
	}

private:
	
	// Constructor for use by Arena::Create()
	Arena(GameMode gameMode, const ArenaConfig& config, float tickRate = 120);

	// Making this private because horrible memory overflows can happen if you changed it
	ArenaConfig _config;
};

RS_NS_END make this possession reward not have to have a arena as an contructor while still keeping ballpred wprking class PossessionReward : public RewardFunction {
public:
	PossessionReward() = default;

	virtual float GetReward(const PlayerData& player, const GameState& state, const Action& prevAction) override {
		// Access the Arena from the GameState
		Arena* arena = state.lastArena;
		if (arena == nullptr) {
			RG_LOG("PossessionReward: Arena is null. Cannot predict ball state.");
			return 0.0f;
		}

		// Calculate the time for the player to reach the ball
		float timeToReach = TimeToReachCalculator::calculateTimeToReach(player, state.ball.pos);

		// Use the Ball Prediction system to get the predicted ball state
		PhysObj predictedBallState = BallPred::GetAtTime(arena, timeToReach, state.ball);

		// Log the time to reach for debugging purposes
		RG_LOG("Time to reach: " << timeToReach);

		// Return the calculated time as the reward
		return timeToReach;
	}
};