//
// Normalized [0,1] implementation of AIScoreCalculator
//

#include "NormalizedAIScoreCalculator.hpp"

#include <cmath>
#include <unordered_map>

#include "private/AIScoreCalculatorSharedUtilities.hpp"
#include "private/AbstractAIScoreCalculator.hpp"
#include "src/main/cpp/net/eagle0/shardok/ai/AIAttackLocations.hpp"
#include "src/main/cpp/net/eagle0/shardok/ai/AIScoreUtilities.hpp"
#include "src/main/cpp/net/eagle0/shardok/ai/AIStrategy.hpp"
#include "src/main/cpp/net/eagle0/shardok/ai/AIUnitScoreCalculator.hpp"
#include "src/main/cpp/net/eagle0/shardok/ai/score/AIScoreCalculator.hpp"
#include "src/main/cpp/net/eagle0/shardok/library/ShardokEngine.hpp"
#include "src/main/cpp/net/eagle0/shardok/library/action_point_distances/ActionPointDistancesCache.hpp"
#include "src/main/cpp/net/eagle0/shardok/library/settings/GameSettings.hpp"

namespace shardok {

using net::eagle0::shardok::storage::fb::BattalionTypeId;

// Bring shared utilities into scope
using score_calculator_internal::UnitsScoreComponents;

/// Normalized implementation of AIScoreCalculator that produces scores in [0, 1] range.
/// Inherits from AbstractAIScoreCalculator to share common functionality.
class NormalizedAIScoreCalculator : public AbstractAIScoreCalculator {
public:
    NormalizedAIScoreCalculator(
            int maxRounds,
            ActionPoints braveWaterCost,
            int meteorRange,
            double meteorCastVigorCost,
            int minimumFleeOddsThreshold,
            int desperateFleeThreshold,
            std::vector<BattalionTypeSPtr> battalionTypes,
            const APDCache &apdCache,
            const ALCache &alCache)
        : AbstractAIScoreCalculator(
                  maxRounds,
                  braveWaterCost,
                  meteorRange,
                  meteorCastVigorCost,
                  minimumFleeOddsThreshold,
                  desperateFleeThreshold,
                  std::move(battalionTypes),
                  apdCache,
                  alCache) {}

    [[nodiscard]] auto GuessedStateScore(
            bool isDefender,
            const GameStateW &state,
            const AIStrategy &aiStrategy,
            const CoordsSet &allCastleCoords) const -> ScoreValue override;

    // Implement pure virtual methods from AbstractAIScoreCalculator
    [[nodiscard]] auto InterpretDefenderOutcome(GameOutcome outcome) const -> ScoreValue override;
    [[nodiscard]] auto InterpretAttackerOutcome(GameOutcome outcome) const -> ScoreValue override;

    [[nodiscard]] auto AttackerFleeStrategyScoreForState(const GameStateW &gameState) const
            -> ScoreValue override;

    [[nodiscard]] auto CombineAttackerScores(
            const UnitsScoreComponents &components,
            double victoryConditionScore,
            int roundsRemaining) const -> ScoreValue override;

    [[nodiscard]] auto CombineDefenderScatterScores(const UnitsScoreComponents &components) const
            -> ScoreValue override;

    [[nodiscard]] auto CombineDefenderHoldCastlesScores(
            const UnitsScoreComponents &components,
            double victoryConditionScore,
            int roundsRemaining) const -> ScoreValue override;

private:
    [[nodiscard]] auto DefenderFleeStrategyScoreForState(const GameStateW &gameState) const
            -> ScoreValue override;

    /// Applies power transformation to spread out compressed scores for MCTS.
    /// Maps [0,1] → [0,1] but pushes values away from 0.5 toward the extremes.
    /// Terminal states (0.0, 1.0) are unchanged.
    [[nodiscard]] auto TransformForMCTS(ScoreValue score) const -> ScoreValue;
};

// Implementation of NormalizedAIScoreCalculator methods

auto NormalizedAIScoreCalculator::InterpretDefenderOutcome(GameOutcome outcome) const
        -> ScoreValue {
    switch (outcome) {
        case GameOutcome::DEFENDER_VICTORY: return 1.0;
        case GameOutcome::ATTACKER_VICTORY: return 0.0;
        case GameOutcome::DRAW: return 0.5;
        case GameOutcome::FLEE_OUTCOME: return 0.5;
    }
    throw ShardokInternalErrorException("Unknown GameOutcome");
}

auto NormalizedAIScoreCalculator::InterpretAttackerOutcome(GameOutcome outcome) const
        -> ScoreValue {
    switch (outcome) {
        case GameOutcome::ATTACKER_VICTORY: return 1.0;
        case GameOutcome::DEFENDER_VICTORY: return 0.0;
        case GameOutcome::DRAW: return 0.5;
        case GameOutcome::FLEE_OUTCOME: return 0.5;
    }
    throw ShardokInternalErrorException("Unknown GameOutcome");
}

auto NormalizedAIScoreCalculator::CombineDefenderScatterScores(
        const UnitsScoreComponents &components) const -> ScoreValue {
    // For defender, we flip the perspective: defenderValue is (1), attackerValue is (2)
    const double defenderValue = components.defenderUnitsValue;
    const double attackerValue = components.attackerUnitsValue;

    // No victory condition for scatter strategy
    const double denominator = defenderValue + attackerValue;
    if (denominator == 0.0) { return 0.5; }

    return defenderValue / denominator;
}

auto NormalizedAIScoreCalculator::CombineDefenderHoldCastlesScores(
        const UnitsScoreComponents &components,
        const double victoryConditionScore,
        const int /*roundsRemaining*/) const -> ScoreValue {
    // For defender, flip perspective
    const double defenderValue = components.defenderUnitsValue;
    const double attackerValue = components.attackerUnitsValue;

    // Apply normalization
    double numerator;
    double denominator;

    if (victoryConditionScore >= 0) {
        numerator = defenderValue + victoryConditionScore;
        denominator = defenderValue + attackerValue + victoryConditionScore;
    } else {
        numerator = defenderValue;
        denominator = defenderValue + attackerValue - victoryConditionScore;
    }

    if (denominator == 0.0) { return 0.5; }

    return numerator / denominator;
}

auto NormalizedAIScoreCalculator::DefenderFleeStrategyScoreForState(
        const GameStateW & /*gameState*/) const -> ScoreValue {
    // FLEE strategy doesn't fit the [0,1] model well - return 0.5
    return 0.5;
}

auto NormalizedAIScoreCalculator::AttackerFleeStrategyScoreForState(
        const GameStateW & /*gameState*/) const -> ScoreValue {
    // FLEE strategy doesn't fit the [0,1] model well - return 0.5
    return 0.5;
}

auto NormalizedAIScoreCalculator::CombineAttackerScores(
        const UnitsScoreComponents &components,
        const double victoryConditionScore,
        const int /*roundsRemaining*/) const -> ScoreValue {
    const double attackerUnitsValue = components.attackerUnitsValue;
    const double defenderUnitsValue = components.defenderUnitsValue;

    // Apply normalization formula
    double numerator;
    double denominator;

    if (victoryConditionScore >= 0) {
        // Positive victory condition: add to numerator
        numerator = attackerUnitsValue + victoryConditionScore;
        denominator = attackerUnitsValue + defenderUnitsValue + victoryConditionScore;
    } else {
        // Negative victory condition: subtract from denominator (making it larger)
        numerator = attackerUnitsValue;
        denominator = attackerUnitsValue + defenderUnitsValue - victoryConditionScore;
    }

    // Handle edge case of all zeros
    if (denominator == 0.0) { return 0.5; }

    return numerator / denominator;
}

auto NormalizedAIScoreCalculator::TransformForMCTS(ScoreValue score) const -> ScoreValue {
    // Power transformation exponent - lower values spread scores more toward extremes
    // Tuned for MCTS: balances exploration vs exploitation
    // - Too low (e.g., 0.3): over-exploitation like standard scorer
    // - Too high (e.g., 0.9): over-exploration like untransformed normalized
    // - 0.6-0.7: sweet spot for MCTS
    constexpr double EXPONENT = 0.1;

    if (score > 0.5) {
        // Map [0.5, 1.0] → [0.5, 1.0] with power curve
        // (score - 0.5) * 2.0 maps to [0, 1], apply power, then scale back
        return 0.5 + 0.5 * std::pow((score - 0.5) * 2.0, EXPONENT);
    } else {
        // Map [0.0, 0.5] → [0.0, 0.5] with power curve (symmetric)
        return 0.5 - 0.5 * std::pow((0.5 - score) * 2.0, EXPONENT);
    }
}

auto NormalizedAIScoreCalculator::GuessedStateScore(
        const bool isDefender,
        const GameStateW &state,
        const AIStrategy &aiStrategy,
        const CoordsSet &allCastleCoords) const -> ScoreValue {
    const int roundsRemaining = GetMaxRounds() - state->current_round();

    ScoreValue rawScore;
    if (isDefender) {
        rawScore = DefenderScoreForState(state, aiStrategy, allCastleCoords, roundsRemaining);
    } else {
        rawScore = AttackerScoreForState(state, aiStrategy, allCastleCoords, roundsRemaining);
    }

    // Apply power transformation to spread out scores for MCTS
    return TransformForMCTS(rawScore);
}

// Factory function implementation
auto MakeNormalizedAIScoreCalculator(
        const SettingsGetter &settingsGetter,
        const APDCache &apdCache,
        const ALCache &alCache) -> std::unique_ptr<AIScoreCalculator> {
    // Extract all battalion types into a vector indexed by BattalionTypeId
    std::vector<BattalionTypeSPtr> battalionTypes(BattalionTypeId::BattalionTypeId_MAX + 1);
    for (int typeId = BattalionTypeId::BattalionTypeId_MIN;
         typeId <= BattalionTypeId::BattalionTypeId_MAX;
         typeId++) {
        auto battalionTypeId = static_cast<BattalionTypeId>(typeId);
        battalionTypes[battalionTypeId] = settingsGetter.GetBattalionType(battalionTypeId);
    }

    return std::make_unique<NormalizedAIScoreCalculator>(
            settingsGetter.Backing().max_rounds(),
            settingsGetter.Backing().brave_water_action_point_cost(),
            settingsGetter.Backing().meteor_range(),
            settingsGetter.Backing().meteor_cast_vigor_cost(),
            settingsGetter.Backing().ai_minimum_flee_odds_threshold(),
            settingsGetter.Backing().ai_desperate_flee_threshold(),
            std::move(battalionTypes),
            apdCache,
            alCache);
}

}  // namespace shardok
