//
// Abstract MCTS Node structure for game-agnostic implementation
//

#ifndef EAGLE0_ABSTRACT_MCTSNODE_HPP
#define EAGLE0_ABSTRACT_MCTSNODE_HPP

#include <cmath>
#include <iomanip>
#include <iostream>
#include <limits>
#include <memory>
#include <sstream>
#include <vector>

#include "MCTSAction.hpp"
#include "MCTSGameState.hpp"
#include "MCTSTypes.hpp"

namespace shardok {
namespace mcts {

// Node type for MCTS tree
enum class NodeType {
    DECISION,  // Player chooses an action (standard MCTS node)
    CHANCE     // Nature determines outcome (for probabilistic actions)
};

// Abstract MCTS Node structure
struct MCTSNode {
    // Node type
    NodeType nodeType = NodeType::DECISION;
    // Action information
    std::unique_ptr<MCTSAction> action;  // The action that led to this node (null for root)
    size_t actionIndex = SIZE_MAX;       // Index in the original actions array (SIZE_MAX for root)

    // Score information
    double immediateScore = 0.0;
    double lookaheadScore = 0.0;

    // Game state after this action
    std::unique_ptr<MCTSGameState> gameState;

    // MCTS statistics
    int visitCount = 0;
    double totalReward = 0.0;
    double averageReward = 0.0;
    mutable double ucb1Value = 0.0;
    double actionWeight = 1.0;  // Prior probability/weight for this action (from heuristics)

    // Tree structure
    std::vector<std::unique_ptr<MCTSNode>> children;
    size_t nextUntriedActionIndex = 0;  // Next action to expand
    size_t totalActions = 0;            // Total number of available actions
    MCTSNode* parent = nullptr;

    // Chance node specific fields (only used when nodeType == CHANCE)
    std::vector<double> outcomeProbabilities;  // Probability of each outcome
    std::vector<double> outcomeRolls;          // Representative roll for each outcome
    double unvisitedOutcomeScore = 0.0;        // Default value for outcomes not expanded yet

    // Game context
    MCTSPlayerId playerId;
    int depth = 0;
    bool isTerminal = false;
    int playerFlips = 0;  // Number of times the active player has changed from root player
    bool isMaximizingPlayer = true;  // True if this node is maximizing for root player

    // Transposition detection
    uint64_t stateHash = 0;
    bool isRedundant = false;  // True if this node represents a duplicate state

    // Constructor for root node
    MCTSNode(std::unique_ptr<MCTSGameState> state, MCTSPlayerId pid, int d)
        : gameState(std::move(state)),
          playerId(pid),
          depth(d),
          playerFlips(0),
          isMaximizingPlayer(true) {
        if (gameState) {
            stateHash = gameState->hash();
            isTerminal = gameState->isTerminal();
        }
    }

    // Constructor for child node
    MCTSNode(
            std::unique_ptr<MCTSAction> act,
            std::unique_ptr<MCTSGameState> state,
            MCTSPlayerId pid,
            int d,
            size_t actIdx = SIZE_MAX,
            int flips = 0,
            bool isMaximizing = true,
            double weight = 1.0)
        : action(std::move(act)),
          actionIndex(actIdx),
          gameState(std::move(state)),
          actionWeight(weight),
          playerId(pid),
          depth(d),
          playerFlips(flips),
          isMaximizingPlayer(isMaximizing) {
        if (gameState) {
            stateHash = gameState->hash();
            isTerminal = gameState->isTerminal();
        }
    }

    // Iterative destructor to avoid stack overflow with deep trees
    ~MCTSNode() {
        std::vector<std::unique_ptr<MCTSNode>> nodesToDestroy;
        nodesToDestroy.swap(children);

        while (!nodesToDestroy.empty()) {
            std::vector<std::unique_ptr<MCTSNode>> currentBatch;
            currentBatch.swap(nodesToDestroy);

            for (const auto& node : currentBatch) {
                if (node && !node->children.empty()) {
                    for (auto& child : node->children) {
                        nodesToDestroy.push_back(std::move(child));
                    }
                    node->children.clear();
                }
            }
        }
    }

    // Calculate UCB1 value for this node from parent's perspective
    // Uses prior-weighted formula similar to AlphaGo:
    //   UCB = Q + c * P * sqrt(N_parent) / (1 + N_child)
    // Where P is the action weight (prior probability from heuristics)
    [[nodiscard]] double CalculateUCB1(
            const double explorationConstant,
            const int parentVisitCount,
            const bool parentIsMaximizing) const {
        // Exploitation: use lookahead score (minimax value)
        // For minimizing nodes, negate the score to prefer low child values
        const double exploitationValue = parentIsMaximizing ? lookaheadScore : -lookaheadScore;

        // Exploration: prior-weighted formula (AlphaGo-style)
        // Actions with weight 0.0 (like FLEE_COMMAND) get no exploration bonus
        // Unvisited nodes get: c * weight * sqrt(N_parent)
        // This prevents bad actions from dominating exploration due to infinite UCB
        const double explorationValue = explorationConstant * actionWeight *
                                        std::sqrt(parentVisitCount) / (1.0 + visitCount);

        return exploitationValue + explorationValue;
    }

    // Check if this node can be expanded
    [[nodiscard]] bool CanExpand() const { return nextUntriedActionIndex < totalActions; }

    // Check if this is a chance node
    [[nodiscard]] bool IsChanceNode() const { return nodeType == NodeType::CHANCE; }

    // Check if this is a decision node
    [[nodiscard]] bool IsDecisionNode() const { return nodeType == NodeType::DECISION; }

    // Get best child from chance node (probability-weighted selection)
    // For chance nodes, we want to explore outcomes proportionally to their probability
    [[nodiscard]] MCTSNode* GetBestChanceChild() const {
        if (children.empty() || !IsChanceNode()) return nullptr;

        // Find the outcome that is most under-explored relative to its probability
        // Expected visits for outcome i: total_visits * probability[i]
        // Actual visits: child[i]->visitCount
        // Deficit: expected - actual
        size_t bestIndex = 0;
        double bestDeficit = -std::numeric_limits<double>::max();

        for (size_t i = 0; i < children.size(); i++) {
            if (!children[i] || children[i]->isRedundant) continue;

            const double expectedVisits = visitCount * outcomeProbabilities[i];
            const double actualVisits = static_cast<double>(children[i]->visitCount);
            const double deficit = expectedVisits - actualVisits;

            if (deficit > bestDeficit) {
                bestDeficit = deficit;
                bestIndex = i;
            }
        }

        return children[bestIndex].get();
    }

    // Get best child based on UCB1
    [[nodiscard]] MCTSNode* GetBestChild(const double explorationConstant) const {
        if (children.empty()) return nullptr;

        MCTSNode* bestChild = nullptr;
        double bestValue = -std::numeric_limits<double>::max();

        for (auto& child : children) {
            // Skip redundant nodes
            if (child->isRedundant) continue;

            // Calculate UCB1 value using the helper function
            const double value =
                    child->CalculateUCB1(explorationConstant, visitCount, isMaximizingPlayer);

            // Debug logging for UCB selection
            static bool enableUCBDebug = false;
            if (enableUCBDebug && child->visitCount > 0) {
                const double exploitationValue =
                        isMaximizingPlayer ? child->lookaheadScore : -child->lookaheadScore;
                const double explorationValue =
                        explorationConstant * std::sqrt(std::log(visitCount) / child->visitCount);
                std::ostringstream line;
                line << std::fixed << std::setprecision(2)
                     << "    UCB: " << (isMaximizingPlayer ? "MAX" : "MIN")
                     << " lookahead=" << child->lookaheadScore << " expl=" << exploitationValue
                     << " (+" << explorationValue << ") = " << value << " ["
                     << (child->action ? child->action->getDescription() : "root") << "]";
                std::cout << line.str() << '\n';
            }

            if (value > bestValue) {
                bestValue = value;
                bestChild = child.get();
            }
        }

        return bestChild;
    }

    // Get best child based on visit count (for final selection)
    [[nodiscard]] MCTSNode* GetBestFinalChild() const {
        if (children.empty()) return nullptr;

        MCTSNode* bestChild = nullptr;
        int bestVisits = 0;
        double bestScore = isMaximizingPlayer ? -std::numeric_limits<double>::max()
                                              : std::numeric_limits<double>::max();

        for (const auto& child : children) {
            // Skip redundant nodes
            if (child->isRedundant) continue;

            // Prefer most-visited node (robust child selection)
            if (child->visitCount > bestVisits) {
                bestVisits = child->visitCount;
                bestScore = child->lookaheadScore;
                bestChild = child.get();
            } else if (child->visitCount == bestVisits) {
                // Tie-break on lookahead score (minimax value, not poisoned average)
                // Maximizing: prefer higher score (better for root player)
                // Minimizing: prefer lower score (worse for root player)
                const bool shouldReplace = isMaximizingPlayer ? (child->lookaheadScore > bestScore)
                                                              : (child->lookaheadScore < bestScore);
                if (shouldReplace) {
                    bestScore = child->lookaheadScore;
                    bestChild = child.get();
                }
            }
        }

        // If no child was visited, fall back to lookahead score
        if (!bestChild && !children.empty()) {
            for (const auto& child : children) {
                if (child->isRedundant) continue;

                const bool shouldReplace = isMaximizingPlayer ? (child->lookaheadScore > bestScore)
                                                              : (child->lookaheadScore < bestScore);
                if (shouldReplace) {
                    bestScore = child->lookaheadScore;
                    bestChild = child.get();
                }
            }
        }

        return bestChild;
    }
};

}  // namespace mcts
}  // namespace shardok

#endif  // EAGLE0_ABSTRACT_MCTSNODE_HPP
