TrainingThread.h

#pragma once

#include "main.h"
#include "RLModel.h"
#include "Wall.h"
#include "Tank.h"
#include "AmmoBox.h"
#include "PerformanceMetrics.h"
#include "Globals.h"
#include "A3CModel.h"  // Tämä sisältää A3CWorker ja A3CNetwork määrittelyt
#include "RandomModel.h"
#include "DDPGModel.h"  // Lisätään DDPG-malli

class TrainingThread {
public:
    A3CWorker worker;        // Paikallinen A3C-työntekijä
    
    TrainingThread(int id, A3CNetwork& a3c, std::atomic<int>& episodes, 
                  std::atomic<int>& wins1, std::atomic<int>& wins2, 
                  const Wall& w, float lr)
        : threadId(id)
        , a3cNetwork(a3c)
        , worker(id, a3c)  // Luo A3C-työntekijä
        , totalEpisodes(episodes)
        , tank1Wins(wins1)
        , tank2Wins(wins2)
        , wall(w)
    {
        // Aseta päivitystaajuus korkeammaksi
        worker.setUpdateFrequency(20);
        
        // Luo mallit tankeille
        tank1Model = std::make_unique<A3CModel>(worker);
        tank2Model = std::make_unique<DDPGModel>();
        //tank2Model = std::make_unique<A3CModel>(worker);
        
        // Käynnistä säie kaikille säikeille
        thread = std::thread(&TrainingThread::run, this);
    }

    // Lisää getter-metodi A3C-työntekijälle
    A3CWorker& getWorker() {
        return worker;
    }
    
    // Lisää getter-metodi päivityslaskurille
    int getUpdateCounter() const {
        return worker.getLocalNetwork().updateCounter;
    }

    bool isVisualized() const {
        return visualized;
    }

    void setVisualized(bool isVisualized) {
        std::lock_guard<std::mutex> lock(visualizationMutex);
        visualized = isVisualized;
        if (!isVisualized) {
            // Herätä säie jos se odottaa visualisointia
            visualizationCV.notify_one();
        }
    }

    // Aseta malli tankille 1
    void setTank1Model(std::unique_ptr<RLModel> model) {
        std::lock_guard<std::mutex> lock(visualizationMutex);
        tank1Model = std::move(model);
    }
    
    // Aseta malli tankille 2
    void setTank2Model(std::unique_ptr<RLModel> model) {
        std::lock_guard<std::mutex> lock(visualizationMutex);
        tank2Model = std::move(model);
    }
    
    // Palauta tankin 1 mallin nimi
    std::string getTank1ModelName() const {
        if (tank1Model) {
            return tank1Model->getName();
        }
        return "Ei mallia";
    }
    
    // Palauta tankin 2 mallin nimi
    std::string getTank2ModelName() const {
        if (tank2Model) {
            return tank2Model->getName();
        }
        return "Ei mallia";
    }
    
    // Tallenna tankin 1 malli
    bool saveTank1Model(const std::string& filename) const {
        if (tank1Model) {
            return tank1Model->saveModel(filename);
        }
        return false;
    }
    
    // Tallenna tankin 2 malli
    bool saveTank2Model(const std::string& filename) const {
        if (tank2Model) {
            return tank2Model->saveModel(filename);
        }
        return false;
    }
        // Lataa tankin 1 malli
    bool loadTank1Model(const std::string& filename) {
        if (tank1Model) {
            return tank1Model->loadModel(filename);
        }
        return false;
    }
    
    // Lataa tankin 2 malli
    bool loadTank2Model(const std::string& filename) {
        if (tank2Model) {
            return tank2Model->loadModel(filename);
        }
        return false;
    }
        // Tallenna tankin 1 malli
    void runA3CEpisode(bool isVisualized, Tank& tank1, Tank& tank2, AmmoBox& ammoBox, 
                      float action1[ACTION_SIZE], float action2[ACTION_SIZE],
                      float& value1, float& value2, float& reward1, float& reward2,
                      unsigned long& currentEpisodeStep, bool& episodeRunning) {
        try {
            // Jos episodi ei ole käynnissä, alusta se
            if (!episodeRunning) {
                tank1.reset(GRID_SIZE);
                tank2.reset(GRID_SIZE);
                ammoBox.spawn(tank1.x, tank1.y, tank2.x, tank2.y);
                currentEpisodeStep = 0;
                episodeRunning = true;
                
                // Päivitä tankkien tilat visualisointia varten
                if (isVisualized) {
                    std::lock_guard<std::mutex> lock(visualizationMutex);
                    tanks.first = tank1;
                    tanks.second = tank2;
                    currentAmmoBox = ammoBox;
                    currentEpisodeLength = currentEpisodeStep;
                }
                return;
            }
            
            currentEpisodeStep++;
            
            // Rajoita episodin pituutta
            if (currentEpisodeStep > MAX_EPISODE_LENGTH) {
                episodeRunning = false;
                return;
            }
            
            // Hae tankkien tilat
            float state1[STATE_SIZE], state2[STATE_SIZE];
            float nextState1[STATE_SIZE], nextState2[STATE_SIZE];
            
            tank1.getState(state1, tank2, wall, ammoBox);
            tank2.getState(state2, tank1, wall, ammoBox);
            
            // Suorita toiminnot eri malleilla
            if (tank1Model) {
                tank1Model->forward(state1, action1, value1);
            } else {
                // Oletustoiminta jos mallia ei ole
                for (int i = 0; i < ACTION_SIZE; i++) {
                    action1[i] = 0.0f;
                }
                value1 = 0.0f;
            }
            
            if (tank2Model) {
                tank2Model->forward(state2, action2, value2);
            } else {
                // Oletustoiminta jos mallia ei ole
                for (int i = 0; i < ACTION_SIZE; i++) {
                    action2[i] = 0.0f;
                }
                value2 = 0.0f;
            }
            
            // Tarkista NaN-arvot ja korjaa ne
            for (int i = 0; i < ACTION_SIZE; i++) {
                if (std::isnan(action1[i])) {
                    action1[i] = 0.0f;
                }
                if (std::isnan(action2[i])) {
                    action2[i] = 0.0f;
                }
            }
            
            if (std::isnan(value1)) value1 = 0.0f;
            if (std::isnan(value2)) value2 = 0.0f;
            
            // Normalisoi toiminnot
            for (int i = 0; i < ACTION_SIZE; i++) {
                action1[i] = std::tanh(action1[i]);
                action2[i] = std::tanh(action2[i]);
            }
            
            // Lisää satunnaisuutta toimintoihin, jotta tankit liikkuvat enemmän
            std::random_device rd;
            std::mt19937 gen(rd());
            std::uniform_real_distribution<float> dist(-0.1f, 0.1f);  // Pienennetty satunnaisuutta (-0.15f -> -0.1f)
            
            // Päivitä tankkien tilat - käsittele molempia tankkeja samalla tavalla
            float moveSpeed = 2.0f;  // Pienennetty 2.5f -> 2.0f
            float rotateSpeed = 1.0f;  // Säilytetään rotaationopeus
            float shootThreshold = 0.15f;  // Pienennetään ampumiskynnystä 0.2f -> 0.15f
            
            tank1.move(action1[0] * moveSpeed + dist(gen), wall, &tank2);
            tank2.move(action2[0] * moveSpeed + dist(gen), wall, &tank1);
            tank1.rotateTurret(action1[1] * rotateSpeed, tank2);
            tank2.rotateTurret(action2[1] * rotateSpeed, tank1);
            
            // Käytä samaa ampumiskynnystä molemmille tankeille
            if (action1[2] > shootThreshold) tank1.shoot(currentEpisodeStep);
            if (action2[2] > shootThreshold) tank2.shoot(currentEpisodeStep);
            
            // Päivitä ammukset ja tarkista osumat
            auto hits1 = tank1.updateBullets(tank2, wall);
            auto hits2 = tank2.updateBullets(tank1, wall);
            
            // Tarkista vahingot
            bool damageDealt = false;
            for (int damage : hits1) {
                tank2.health -= damage;
                if (tank2.health <= 0) {
                    tank2.alive = false;
                    break;
                }
                damageDealt = true;
            }
            for (int damage : hits2) {
                tank1.health -= damage;
                if (tank1.health <= 0) {
                    tank1.alive = false;
                    break;
                }
                damageDealt = true;
            }
            
            // Tarkista ammuslaatikot
            tank1.checkAmmoBox(ammoBox);
            tank2.checkAmmoBox(ammoBox);
            
            // Varmista että ammuslaatikko on aina aktiivinen jos jommallakummalla tankilla on ammukset loppu
            if (!ammoBox.active && (tank1.ammo == 0 || tank2.ammo == 0)) {
                ammoBox.spawn(tank1.x, tank1.y, tank2.x, tank2.y);
            }
            
            // Yritä luoda uusi ammuslaatikko jos vanha on kerätty ja kummallakaan ei ole ammukset loppu
            if (!ammoBox.active && tank1.ammo > 0 && tank2.ammo > 0) {
                // Kasvatetaan todennäköisyyttä ammolaatikon ilmestymiselle
                if (rand() % 100 < 30) {  // 30% todennäköisyys joka kierroksella
                    ammoBox.spawn(tank1.x, tank1.y, tank2.x, tank2.y);
                }
            }
            
            // Varmista että ammuslaatikko on aina aktiivinen jos jommallakummalla tankilla on vähän ammuksia
            if (!ammoBox.active && (tank1.ammo < 5 || tank2.ammo < 5)) {
                ammoBox.spawn(tank1.x, tank1.y, tank2.x, tank2.y);
            }
            
            // Laske palkkiot
            reward1 = tank1.calculateReward(tank2, wall, ammoBox, currentEpisodeStep);
            reward2 = tank2.calculateReward(tank1, wall, ammoBox, currentEpisodeStep);
            
            // Tarkista NaN-arvot palkkioissa
            if (std::isnan(reward1)) reward1 = 0.0f;
            if (std::isnan(reward2)) reward2 = 0.0f;
            
            // Päivitä tilat
            tank1.getState(nextState1, tank2, wall, ammoBox);
            tank2.getState(nextState2, tank1, wall, ammoBox);
            
            // Kerää kokemus malleille
            if (tank1Model) {
                tank1Model->collectExperience(state1, action1, reward1, value1);
            }
            
            if (tank2Model) {
                tank2Model->collectExperience(state2, action2, reward2, value2);
            }
            
            // Tarkista pitäisikö päivittää mallit
            if (tank1Model && tank1Model->shouldUpdate()) {
                float nextValue;
                float dummyAction[ACTION_SIZE];
                tank1Model->forward(nextState1, dummyAction, nextValue);
                if (std::isnan(nextValue)) nextValue = 0.0f;
                tank1Model->update(nextValue, !tank1.alive || !tank2.alive);
            }
            
            if (tank2Model && tank2Model->shouldUpdate()) {
                float nextValue;
                float dummyAction[ACTION_SIZE];
                tank2Model->forward(nextState2, dummyAction, nextValue);
                if (std::isnan(nextValue)) nextValue = 0.0f;
                tank2Model->update(nextValue, !tank1.alive || !tank2.alive);
            }
            
            // Pakota päivitys joka 10. kierroksella, vaikka updateFrequency ei olisi täyttynyt
            if (currentEpisodeStep % 10 == 0) {
                if (tank1Model && tank1Model->getRewardsSize() > 0) {
                    float nextValue;
                    float dummyAction[ACTION_SIZE];
                    tank1Model->forward(nextState1, dummyAction, nextValue);
                    if (std::isnan(nextValue)) nextValue = 0.0f;
                    tank1Model->update(nextValue, false);
                }
                
                if (tank2Model && tank2Model->getRewardsSize() > 0) {
                    float nextValue;
                    float dummyAction[ACTION_SIZE];
                    tank2Model->forward(nextState2, dummyAction, nextValue);
                    if (std::isnan(nextValue)) nextValue = 0.0f;
                    tank2Model->update(nextValue, false);
                }
            }
            
            // Tarkista onko episodi päättynyt
            if (!tank1.alive || !tank2.alive || currentEpisodeStep > MAX_EPISODE_LENGTH) {
                // Pakota päivitys episodin päättyessä, vaikka updateFrequency ei olisi täyttynyt
                if (tank1Model && tank1Model->getRewardsSize() > 0) {
                    float nextValue = 0.0f; // Episodin päättyessä seuraavan tilan arvo on 0
                    tank1Model->update(nextValue, true);
                }
                
                if (tank2Model && tank2Model->getRewardsSize() > 0) {
                    float nextValue = 0.0f; // Episodin päättyessä seuraavan tilan arvo on 0
                    tank2Model->update(nextValue, true);
                }
                
                // Päivitä voittotilastot
                if (!tank1.alive && tank2.alive) {
                    tank2Wins++;
                } else if (tank1.alive && !tank2.alive) {
                    tank1Wins++;
                }
                totalEpisodes++;
                
                // Merkitse episodi päättyneeksi
                episodeRunning = false;
            }
            
            // Tarkista ovatko tankit jumissa samassa paikassa
            static int lastTank1X = -1, lastTank1Y = -1;
            static int lastTank2X = -1, lastTank2Y = -1;
            static int stuckCounter = 0;
            
            if (tank1.x == lastTank1X && tank1.y == lastTank1Y && 
                tank2.x == lastTank2X && tank2.y == lastTank2Y) {
                stuckCounter++;
                
                // Jos tankit ovat olleet samassa paikassa 100 kierrosta, resetoi ne
                if (stuckCounter > 100) {
                    tank1.reset(GRID_SIZE);
                    tank2.reset(GRID_SIZE);
                    ammoBox.spawn(tank1.x, tank1.y, tank2.x, tank2.y);
                    stuckCounter = 0;
                }
            } else {
                // Päivitä viimeiset sijainnit
                lastTank1X = tank1.x;
                lastTank1Y = tank1.y;
                lastTank2X = tank2.x;
                lastTank2Y = tank2.y;
                stuckCounter = 0;
            }
            
            // Päivitä hyperparametrit
            updateHyperParameters(reward1);
            
            // Päivitä tankkien tilat visualisointia varten
            if (isVisualized) {
                std::lock_guard<std::mutex> lock(visualizationMutex);
                tanks.first = tank1;
                tanks.second = tank2;
                currentAmmoBox = ammoBox;
                currentEpisodeLength = currentEpisodeStep;
            }
        } catch (const std::exception& e) {
            // Virheenkäsittely
            std::cerr << "Virhe runA3CEpisode-metodissa: " << e.what() << std::endl;
            
            // Resetoi episodi virheen sattuessa
            tank1.reset(GRID_SIZE);
            tank2.reset(GRID_SIZE);
            ammoBox.spawn(tank1.x, tank1.y, tank2.x, tank2.y);
            currentEpisodeStep = 0;
            episodeRunning = true;
            
            // Aseta oletusarvot
            for (int i = 0; i < ACTION_SIZE; i++) {
                action1[i] = 0.0f;
                action2[i] = 0.0f;
            }
            value1 = 0.0f;
            value2 = 0.0f;
            reward1 = 0.0f;
            reward2 = 0.0f;
        }
    }
    
    void visualizationComplete() {
        std::lock_guard<std::mutex> lock(visualizationMutex);
        waitingForVisualization = false;
        visualizationCV.notify_one();
    }

    void run() {
        while (running) {
            // Alusta uusi episodi
            Tank tank1(GRID_SIZE/4, GRID_SIZE/2);
            tank1.reset(GRID_SIZE);
            Tank tank2(3*GRID_SIZE/4, GRID_SIZE/2);
            tank2.reset(GRID_SIZE);
            AmmoBox ammoBox;
            ammoBox.spawn(tank1.x, tank1.y, tank2.x, tank2.y);
            
            float action1[ACTION_SIZE] = {0}, action2[ACTION_SIZE] = {0};
            float value1 = 0, value2 = 0;
            float reward1 = 0, reward2 = 0;
            unsigned long episodeStep = 0;
            bool episodeRunning = true;
            
            // Välitä visualisoinnin tila runEpisode-funktiolle
            bool isVisualized = visualized.load();
            bool wasVisualized = isVisualized;
            
            while (episodeRunning) {
                // Tarkista onko visualisoinnin tila muuttunut
                isVisualized = visualized.load();
                
                // Jos visualisointi kytketään päälle, päivitä tankkien tila visualisointia varten
                if (isVisualized && !wasVisualized) {
                    std::lock_guard<std::mutex> lock(visualizationMutex);
                    tanks.first = tank1;
                    tanks.second = tank2;
                    currentAmmoBox = ammoBox;
                    currentEpisodeLength = episodeStep;
                }
                
                // Suorita yksi askel episodista
                runA3CEpisode(isVisualized, tank1, tank2, ammoBox, 
                             action1, action2, value1, value2, reward1, reward2,
                             episodeStep, episodeRunning);
                
                // Päivitä visualisoinnin edellinen tila
                wasVisualized = isVisualized;
                
                // Säädä viivettä visualisoinnin tilan mukaan
                if (isVisualized) {
                    std::this_thread::sleep_for(std::chrono::milliseconds(1));
                } 
            }
        }
    }

    void stop() {
        running = false;
        if (thread.joinable()) {
            thread.join();
        }
    }

    bool isRunning() const {
        return running;
    }

    unsigned long getCurrentEpisodeLength() const {
        return currentEpisodeLength;
    }

    ~TrainingThread() {
        stop();  // Varmistetaan että säie pysäytetään
    }

    const HyperParameters& getCurrentHyperParameters() const {
        return tuner.getCurrentParams();
    }

    // Lisätään getter-metodit tankkien, ammolaatikon ja visualisointimutexin hakemiseen
    std::mutex& getVisualizationMutex() {
        return visualizationMutex;
    }

    const std::pair<Tank, Tank>& getTanks() const {
        return tanks;
    }

    const AmmoBox& getCurrentAmmoBox() const {
        return currentAmmoBox;
    }

    void resetTanks(Tank& tank1, Tank& tank2, AmmoBox& ammoBox) {
        std::lock_guard<std::mutex> lock(visualizationMutex);
        
        // Varmista että tankit eivät ole seinän sisällä
        if (tank1.x >= GRID_SIZE/2 - 2 && tank1.x < GRID_SIZE/2 + 2 && 
            tank1.y >= GRID_SIZE/3 && tank1.y < GRID_SIZE/3 * 2) {
            tank1.x = GRID_SIZE/4;
            tank1.y = GRID_SIZE/2;
        }
        
        if (tank2.x >= GRID_SIZE/2 - 2 && tank2.x < GRID_SIZE/2 + 2 && 
            tank2.y >= GRID_SIZE/3 && tank2.y < GRID_SIZE/3 * 2) {
            tank2.x = 3*GRID_SIZE/4;
            tank2.y = GRID_SIZE/2;
        }
        
        // Varmista että tankit eivät ole päällekkäin
        if (tank1.x == tank2.x && tank1.y == tank2.y) {
            tank1.x = GRID_SIZE/4;
            tank1.y = GRID_SIZE/2;
            tank2.x = 3*GRID_SIZE/4;
            tank2.y = GRID_SIZE/2;
        }
        
        tanks.first = tank1;
        tanks.second = tank2;
        currentAmmoBox = ammoBox;
    }

private:
    int threadId;
    A3CNetwork& a3cNetwork;  // Viittaus globaaliin A3C-verkkoon
    // A3CWorker worker;        // Paikallinen A3C-työntekijä - siirretty public-osioon
    std::unique_ptr<RLModel> tank1Model; // Malli tankille 1
    std::unique_ptr<RLModel> tank2Model; // Malli tankille 2
    std::atomic<int>& totalEpisodes;
    std::atomic<int>& tank1Wins;
    std::atomic<int>& tank2Wins;
    const Wall& wall;
    std::thread thread;
    std::atomic<bool> running{true};
    unsigned long currentEpisodeLength{0};  // Tämä toimii sekä laskurina että visualisoinnin arvona
    std::mutex visualizationMutex;
    std::pair<Tank, Tank> tanks{Tank(GRID_SIZE/4, GRID_SIZE/2), Tank(3*GRID_SIZE/4, GRID_SIZE/2)};
    AmmoBox currentAmmoBox;
    
    std::atomic<bool> visualized{false};
    bool waitingForVisualization{false};
    std::condition_variable visualizationCV;
    HyperparameterTuner tuner;
    
    void updateHyperParameters(float reward) {
         // Päivitä hyperparametrit molemmille malleille
        if (tank1Model) {
            // Luodaan ei-const kopio hyperparametreista
            HyperParameters params = tank1Model->getHyperParams();
            tuner.updateParameters(reward, params);
            // Asetetaan päivitetyt parametrit takaisin malliin
            tank1Model->setHyperParams(params);
        }
        
        if (tank2Model) {
            // Luodaan ei-const kopio hyperparametreista
            HyperParameters params = tank2Model->getHyperParams();
            tuner.updateParameters(reward, params);
            // Asetetaan päivitetyt parametrit takaisin malliin
            tank2Model->setHyperParams(params);
        }
    }
};