import type { SimulationState } from './types';
import { hasLineOfSight } from './sensors';

/**
 * Fitness calculation for NEAT Arena.
 * 
 * Fitness rewards:
 * - +10 per hit on opponent
 * - -10 per being hit
 * - -0.002 per tick (time penalty to encourage aggression)
 * - -0.2 per shot fired (ammo management)
 * - +0.01 per tick when aiming well at visible opponent
 */

export const FITNESS_CONFIG = {
    HIT_REWARD: 10.0,      // Kill
    DAMAGE_REWARD: 4.0,    // Per hit dealt (High reward for hitting)
    HIT_PENALTY: 1.0,      // Per hit taken (Reduced to 1.0 to encourage aggression/trading)
    TIME_PENALTY: 0.002,   // Per tick
    SHOT_PENALTY: 0.0,     // REMOVED: Free shooting encourages exploration
    AIM_REWARD: 0.01,      // Increased: Stronger guide signal
    MOVE_REWARD: 0.001,    // Per tick moving
};

export interface FitnessTracker {
    agentId: number;
    fitness: number;
    
    // For incremental calculation
    lastKills: number;
    lastHitsTaken: number;
    lastHitsDealt: number;
    shotsFired: number;
}


/**
 * Create a new fitness tracker
 */
export function createFitnessTracker(agentId: number): FitnessTracker {
    return {
        agentId,
        fitness: 0,
        lastKills: 0,
        lastHitsTaken: 0,
        lastHitsDealt: 0,
        shotsFired: 0,
    };
}

/**
 * Update fitness based on current simulation state
 */
export function updateFitness(tracker: FitnessTracker, state: SimulationState): FitnessTracker {
    const agent = state.agents.find(a => a.id === tracker.agentId)!;
    const opponent = state.agents.find(a => a.id !== tracker.agentId)!;
    
    const newTracker = { ...tracker };
    
    // Reward for new kills
    const newKills = agent.kills - tracker.lastKills;
    newTracker.fitness += newKills * FITNESS_CONFIG.HIT_REWARD;
    newTracker.lastKills = agent.kills;
    
    // Reward for HITS DEALT (Direct Damage)
    // We infer hits dealt by checking opponent's hit counter increase
    const currentHitsDealt = opponent.hits; // Assuming opponent.hits tracks times they were hit
    const newHitsDealt = currentHitsDealt - tracker.lastHitsDealt;
    
    if (newHitsDealt > 0) {
        // +2.0 per hit. 5 hits = 10 pts (Kill equivalent).
        // Makes shooting visibly rewarding immediately.
        newTracker.fitness += newHitsDealt * FITNESS_CONFIG.DAMAGE_REWARD; 
    }
    newTracker.lastHitsDealt = currentHitsDealt;

    // Penalty for being hit (Hits Taken)
    const newHitsTaken = agent.hits - tracker.lastHitsTaken;
    newTracker.fitness -= newHitsTaken * FITNESS_CONFIG.HIT_PENALTY; 
    newTracker.lastHitsTaken = agent.hits;
    
    // Time penalty (encourages finishing quickly)
    newTracker.fitness -= FITNESS_CONFIG.TIME_PENALTY;
    
    // Check if agent fired this tick
    if (agent.fireCooldown === 10) {
        newTracker.shotsFired++;
        newTracker.fitness -= FITNESS_CONFIG.SHOT_PENALTY; // Tiny penalty just to prevent spamming empty space
    }
    
    // Reward for aiming at visible opponent (Guide Signal ONLY)
    if (hasLineOfSight(agent, opponent, state.map.walls)) {
        const dx = opponent.position.x - agent.position.x;
        const dy = opponent.position.y - agent.position.y;
        const angleToOpponent = Math.atan2(dy, dx);
        
        let angleDiff = angleToOpponent - agent.aimAngle;
        while (angleDiff > Math.PI) angleDiff -= 2 * Math.PI;
        while (angleDiff < -Math.PI) angleDiff += 2 * Math.PI;
        
        const cosAngleDiff = Math.cos(angleDiff);
        
        // Reduced from 0.05 to 0.005. 
        // Max total aim points = 1.5. 
        // One bullet hit (2.0) is worth more than perfect aiming all match.
        newTracker.fitness += ((cosAngleDiff + 1) * 0.5) * FITNESS_CONFIG.AIM_REWARD;
    }

    // Small reward for movement
    const speed = Math.sqrt(agent.velocity.x**2 + agent.velocity.y**2);
    if (speed > 0.1) {
        newTracker.fitness += FITNESS_CONFIG.MOVE_REWARD;
    }
    
    return newTracker;
}