-- Enhanced missile guidance script
-- - Memory management for destroyed missiles
-- - Energy-efficient updates with variable tick rates
-- - Climb-to-altitude behavior before engaging targets
-- - Performance optimizations for large missile salvos
-- - Exact proportional allocation (largest-remainder)
-- - Per-missile-type guidance based on initial speed (small/medium/large)
-- - Predictive lead applied before dive adjustment
-- - Adaptive smoothing per missile type, with reduced lag when close
-- - Proximity detonation when missile starts moving away from target within 5 meters
-- - Fan-out behavior at cruising altitude to improve effectiveness against countermeasures

-- Persistent tables with cleanup tracking
smoothed_aim_points = smoothed_aim_points or {}
last_target_velocities = last_target_velocities or {}
missile_target_index = missile_target_index or {}
missile_initial_speeds = missile_initial_speeds or {}
missile_last_update = missile_last_update or {}
missile_altitude_phase = missile_altitude_phase or {} -- Tracks if missile is in climb phase
missile_guidance_active = missile_guidance_active or {} -- Tracks if missile is actively guiding
missile_last_distance = missile_last_distance or {} -- Tracks last known distance to target
missile_proximity_armed = missile_proximity_armed or {} -- Tracks if proximity detonation is armed
missile_fan_offset = missile_fan_offset or {} -- Tracks fan-out offset for each missile

-- Missile energy tracking
local missile_max_speeds = {}  -- Tracks max speed achieved by each missile
local missile_energy = {}      -- Normalized energy level (0-1) for each missile
local energy_update_interval = 1.0 -- Check energy every second
local last_energy_update = 0

-- Profiles per missile type
smoothing_profiles = {
    small  = {alpha_min = 0.3,  alpha_max = 0.6,  dive_distance = 1000,  cruise_altitude = 150, update_interval = 0.1},
    medium = {alpha_min = 0.15, alpha_max = 0.5,  dive_distance = 2000, cruise_altitude = 300, update_interval = 0.2},
    large  = {alpha_min = 0.1,  alpha_max = 0.35, dive_distance = 2000, cruise_altitude = 300, update_interval = 0.3}
}

local alpha_vel_scale = 0.05
local EPS = 1e-6

------------------------------------------------------------
-- Helper: Predictive lead aimpoint with performance optimization
function GetLeadAimpoint(targetPos, targetVel, missilePos, missileSpeed, maxLead)
    maxLead = maxLead or 500
    local toTarget = targetPos - missilePos
    local distance = toTarget.magnitude
    
    -- Skip calculation for very slow missiles or very close targets
    if missileSpeed <= 0.1 or distance < 5 then
        return targetPos
    end
    
    -- Cache magnitude calculations
    local timeToImpact = distance / missileSpeed
    local leadOffset = targetVel * timeToImpact
    local leadMag = leadOffset.magnitude
    
    -- Limit lead distance
    if leadMag > maxLead then
        leadOffset = leadOffset * (maxLead / leadMag)
    end
    
    return targetPos + leadOffset
end

------------------------------------------------------------
-- Helper: Exponential smoothing with nil checks
function ExponentialSmooth(prev, current, alpha)
    if not current then return prev end
    if not prev then return current end
    return Vector3(
        alpha * current.x + (1 - alpha) * prev.x,
        alpha * current.y + (1 - alpha) * prev.y,
        alpha * current.z + (1 - alpha) * prev.z
    )
end

------------------------------------------------------------
-- Helper: Clean up missile data for destroyed missiles
function CleanupMissileData(I)
    if not I then return end  -- Safety check
    local current_time = I:GetTime()
    local missile_ids = {}
    
    -- Build set of active missile IDs
    for t = 0, I:GetLuaTransceiverCount() - 1 do
        for m = 0, I:GetLuaControlledMissileCount(t) - 1 do
            local mi = I:GetLuaControlledMissileInfo(t, m)
            if mi and mi.Valid then
                missile_ids[mi.Id] = true
            end
        end
    end
    
    -- Clean up data for missiles that no longer exist
    for id, _ in pairs(missile_initial_speeds) do
        if not missile_ids[id] then
            missile_initial_speeds[id] = nil
            missile_target_index[id] = nil
            missile_last_update[id] = nil
            smoothed_aim_points[id] = nil
            missile_altitude_phase[id] = nil
            missile_guidance_active[id] = nil
        end
    end
end

------------------------------------------------------------
-- Helper: Update missile energy state based on speed
function UpdateMissileEnergy(missile_id, current_speed, current_time, I)
    -- Initialize if needed
    if missile_max_speeds[missile_id] == nil then
        missile_max_speeds[missile_id] = current_speed
        missile_energy[missile_id] = 0.0
    end
    
    -- Only update energy at fixed intervals
    if current_time - (missile_energy[missile_id.."_last_update"] or 0) < energy_update_interval then
        return missile_energy[missile_id]
    end
    
    -- Update max speed if current speed is higher
    if current_speed > missile_max_speeds[missile_id] then
        missile_max_speeds[missile_id] = current_speed
    end
    
    -- Calculate energy level (0-1) based on current speed vs max speed
    if missile_max_speeds[missile_id] > 0 then
        missile_energy[missile_id] = current_speed / missile_max_speeds[missile_id]
    else
        missile_energy[missile_id] = 0
    end
    
    missile_energy[missile_id.."_last_update"] = current_time
    return missile_energy[missile_id]
end

------------------------------------------------------------
-- Helper: Get update interval based on missile energy and type
function GetUpdateInterval(missile_id, profile, current_time, I)
    local base_interval = profile.update_interval or 0.1
    local energy = missile_energy[missile_id] or 1.0
    
    -- Scale update rate based on energy (more frequent when energy is high)
    local scaled_interval = base_interval * (1.5 - (energy * 0.5))
    
    -- Ensure minimum update interval
    return math.max(0.05, scaled_interval)
end

------------------------------------------------------------
-- Helper: Track target acceleration (unused for now)
function GetTargetAcceleration(target_idx, current_vel)
    local last_vel = last_target_velocities[target_idx]
    last_target_velocities[target_idx] = current_vel
    if not last_vel then
        return Vector3(0, 0, 0)
    end
    return (current_vel - last_vel) / 1 -- assuming 1s update
end

------------------------------------------------------------
-- Helper: Adaptive alpha based on target velocity change
function GetAdaptiveAlpha(target_idx, current_vel, profile)
    local last_vel = last_target_velocities[target_idx]
    last_target_velocities[target_idx] = current_vel
    if not last_vel then return profile.alpha_min end
    local vel_change = Vector3.Magnitude(current_vel - last_vel)
    local alpha = profile.alpha_min + math.min(profile.alpha_max - profile.alpha_min, vel_change * alpha_vel_scale)
    return alpha
end

------------------------------------------------------------
-- Helper: Classify missile type by initial speed
function ClassifyMissile(initial_speed)
    if initial_speed >= 35 then
        return "small"
    elseif initial_speed >= 15 then
        return "medium"
    else
        return "large"
    end
end

------------------------------------------------------------
-- Helper: Score target for allocation
function GetTargetScore(ref_pos, tinfo)
    return Vector3.Distance(ref_pos, tinfo.AimPointPosition)
end

------------------------------------------------------------
-- Helper: Exact proportional allocation
function AllocateMissilesExact(targets, missile_count, ref_pos)
    local weights = {}
    local total_weight = 0

    for _, tgt in ipairs(targets) do
        local score = GetTargetScore(ref_pos, tgt.info)
        local weight = 1 / (score + EPS)
        weights[tgt.index] = {weight = weight, tgt = tgt}
        total_weight = total_weight + weight
    end

    if total_weight <= 0 then
        local allocation = {}
        local base = math.floor(missile_count / #targets)
        local rem = missile_count - base * #targets
        for _, tgt in ipairs(targets) do
            allocation[tgt.index] = base
        end
        for i = 1, rem do
            allocation[targets[((i - 1) % #targets) + 1].index] =
            allocation[targets[((i - 1) % #targets) + 1].index] + 1
        end
        return allocation
    end

    local quotas = {}
    local allocation = {}
    local remainders = {}
    local allocated_sum = 0
    for _, tgt in ipairs(targets) do
        local w = weights[tgt.index].weight
        local quota = (w / total_weight) * missile_count
        quotas[tgt.index] = quota
        local flo = math.floor(quota)
        allocation[tgt.index] = flo
        allocated_sum = allocated_sum + flo
        remainders[tgt.index] = quota - flo
    end

    local remaining = missile_count - allocated_sum
    if remaining > 0 then
        local rem_list = {}
        for idx, rem in pairs(remainders) do
            table.insert(rem_list, {index = idx, rem = rem})
        end
        table.sort(rem_list, function(a, b)
            if a.rem == b.rem then return a.index < b.index end
            return a.rem > b.rem
        end)
        local i = 1
        while remaining > 0 do
            local pick = rem_list[((i - 1) % #rem_list) + 1]
            allocation[pick.index] = allocation[pick.index] + 1
            remaining = remaining - 1
            i = i + 1
        end
    end

    return allocation
end

------------------------------------------------------------
-- MAIN UPDATE
function Update(I)
    if not I then return end  -- Safety check
    -- Get current time for update timing
    local current_time = I:GetTime()
    
    -- Clean up data for destroyed missiles
    CleanupMissileData(I)
    
    -- Gather valid targets with caching
    local target_count = I:GetNumberOfTargets(0)
    if target_count == 0 then 
        -- No targets, clean up and return
        CleanupMissileData()
        return 
    end
    
    local targets = {}
    for ti = 0, target_count - 1 do
        local tinfo = I:GetTargetInfo(0, ti)
        if tinfo and tinfo.Valid then
            table.insert(targets, {index = ti, info = tinfo})
        end
    end
    if #targets == 0 then return end

    -- Gather missiles and compute average position
    local missiles = {}
    local total_missiles = 0
    local sum_pos = Vector3(0,0,0)
    for t = 0, I:GetLuaTransceiverCount() - 1 do
        for m = 0, I:GetLuaControlledMissileCount(t) - 1 do
            local mi = I:GetLuaControlledMissileInfo(t, m)
            if mi and mi.Valid then
                total_missiles = total_missiles + 1
                local pos = mi.Position
                sum_pos = sum_pos + pos
                local vel = mi.Velocity
                local speed = Vector3.Magnitude(vel)
                -- Update missile energy state
                local energy = UpdateMissileEnergy(mi.Id, speed, current_time, I)
                table.insert(missiles, {t = t, m = m, id = mi.Id, pos = pos, velocity = vel, speed = speed, energy = energy})
            end
        end
    end
    if total_missiles == 0 then return end
    local avg_pos = sum_pos / total_missiles

    -- Proportional allocation
    local allocation = AllocateMissilesExact(targets, total_missiles, avg_pos)
    local assigned_counts = {}
    for _, tgt in ipairs(targets) do
        assigned_counts[tgt.index] = 0
        if allocation[tgt.index] == nil then allocation[tgt.index] = 0 end
    end

    -- Sort missiles by speed
    table.sort(missiles, function(a,b) return a.speed > b.speed end)

    -- Assign missiles
    local target_index_order = {}
    for _, tgt in ipairs(targets) do table.insert(target_index_order, tgt.index) end
    for _, missile in ipairs(missiles) do
        local missile_id = missile.id
        if not missile_initial_speeds[missile_id] then
            missile_initial_speeds[missile_id] = missile.speed
        end
        local current_tgt = missile_target_index[missile_id]
        if current_tgt and allocation[current_tgt] and assigned_counts[current_tgt] < allocation[current_tgt] then
            assigned_counts[current_tgt] = assigned_counts[current_tgt] + 1
        else
            local found = nil
            for _, idx in ipairs(target_index_order) do
                if assigned_counts[idx] < allocation[idx] then
                    found = idx
                    break
                end
            end
            if not found then found = target_index_order[1] end
            missile_target_index[missile_id] = found
            assigned_counts[found] = assigned_counts[found] + 1
        end
    end

    -- Guidance loop with update rate control and climb behavior
    for _, missile in ipairs(missiles) do
        local t = missile.t
        local m = missile.m
        local missile_id = missile.id
        local missile_pos = missile.pos
        local missile_speed = missile.speed
        
        -- Skip if this missile was updated recently (update rate control)
        local m_initial_speed = missile_initial_speeds[missile_id] or missile_speed
        local missile_type = ClassifyMissile(m_initial_speed)
        local profile = smoothing_profiles[missile_type]
        
        -- Initialize missile state if needed
        if missile_guidance_active[missile_id] == nil then
            missile_guidance_active[missile_id] = false
            missile_altitude_phase[missile_id] = true -- Start in climb phase
        end
        
        -- Check if it's time to update this missile based on energy and type
        local last_update = missile_last_update[missile_id] or 0
        local update_interval = GetUpdateInterval(missile_id, profile, current_time, I)
        if current_time - last_update < update_interval then
            goto continue
        end
        missile_last_update[missile_id] = current_time
        
        -- Skip update if missile is nearly out of energy
        if missile.energy < 0.1 then
            goto continue
        end
        
        -- Get target info first
        local target_idx = missile_target_index[missile_id]
        local target_info = target_idx and I:GetTargetInfo(0, target_idx)
        if not target_info or not target_info.Valid then
            -- No valid target, fly straight
            local straight_ahead = missile_pos + (missile.velocity.normalized * 1000)
            I:SetLuaControlledMissileAimPoint(t, m, 
                straight_ahead.x,
                straight_ahead.y,
                straight_ahead.z
            )
            goto continue
        end
        
        local target_pos = target_info.AimPointPosition
        local target_vel = target_info.Velocity
        local target_distance = (missile_pos - target_pos).magnitude
        if target_distance < 10 then  -- Within 5 meters
            if not missile_proximity_armed[missile_id] then
                missile_proximity_armed[missile_id] = true
                missile_last_distance[missile_id] = target_distance
            elseif missile_last_distance[missile_id] < target_distance then
                -- We're getting further away, detonate!
                I:DetonateLuaControlledMissile(t, m)
                goto continue
            else
                missile_last_distance[missile_id] = target_distance
            end
        else
            missile_proximity_armed[missile_id] = nil
            missile_last_distance[missile_id] = nil
        end
        
        -- Get target info
        local target_idx = missile_target_index[missile_id]
        local target_info = target_idx and I:GetTargetInfo(0, target_idx)
        
        -- Check if target is within 1500m for direct approach
        if target_distance <= 1500 then
            -- Fly directly to target with a small lead
            local lead_point = GetLeadAimpoint(target_pos, target_vel or Vector3(0,0,0), missile_pos, missile_speed or 0, 500)
            local direct_aim = lead_point or target_pos
            
            -- Add small vertical offset for hitbox
            direct_aim = direct_aim + Vector3(0, -2, 0)
            
            I:SetLuaControlledMissileAimPoint(t, m, direct_aim.x, direct_aim.y, direct_aim.z)
            goto continue
        end
        
        -- Handle climb phase with fan-out (only for targets > 1500m away)
        if missile_altitude_phase[missile_id] then
            local cruise_alt = profile.cruise_altitude
            
            -- Generate a consistent fan-out offset based on missile ID
            if not missile_fan_offset[missile_id] then
                -- Create a unique offset based on missile ID
                local id_str = tostring(missile_id)
                local id_hash = 0
                for i = 1, #id_str do
                    id_hash = id_hash + string.byte(id_str, i) * i
                end
                math.randomseed(id_hash)
                -- Calculate angle relative to target direction (within 60 degrees left or right)
                local target_direction = (target_pos - missile_pos).normalized
                local base_angle = math.atan2(target_direction.z, target_direction.x)
                local angle_variation = (math.random() - 0.5) * (2 * math.pi / 3)  -- ±60 degrees
                local angle = base_angle + angle_variation
                
                local distance = 50 + (id_hash % 100)  -- 50-150m offset
                missile_fan_offset[missile_id] = {
                    x = math.cos(angle) * distance,
                    z = math.sin(angle) * distance
                }
            end
            
            -- Calculate climb point with offset
            local offset = missile_fan_offset[missile_id] or {x=0, z=0}
            local climb_point = Vector3(
                missile_pos.x + offset.x,
                cruise_alt,
                missile_pos.z + offset.z
            )
            
            -- Check if we've reached cruise altitude
            if missile_pos.y >= cruise_alt * 0.85 then  -- 85% of cruise alt
                missile_altitude_phase[missile_id] = false
                missile_guidance_active[missile_id] = true
                -- Initialize smoothed point at current position when starting guidance
                smoothed_aim_points[missile_id] = Vector3(climb_point.x, cruise_alt, climb_point.z)
            end
            
            -- Set aim point with offset
            I:SetLuaControlledMissileAimPoint(t, m, 
                climb_point.x,
                cruise_alt,
                climb_point.z
            )
            goto continue
        end
        
        -- Standard guidance for active phase
        if not target_info or not target_info.Valid then
            if missile_guidance_active[missile_id] then
                -- Try to find a new target if current one is lost
                missile_target_index[missile_id] = nil
                missile_guidance_active[missile_id] = false
            end
            goto continue
        end
        
        local target_pos = target_info.AimPointPosition
        local target_vel = target_info.Velocity
        local target_distance = Vector3.Distance(target_pos, missile_pos)
        
        -- Only activate guidance if we're not in climb phase
        if not missile_guidance_active[missile_id] and not missile_altitude_phase[missile_id] then
            missile_guidance_active[missile_id] = true
        end
        
        -- Calculate lead point with altitude management and safety checks
        local lead_point = GetLeadAimpoint(target_pos, target_vel or Vector3(0,0,0), missile_pos, missile_speed or 0, 800)
        local aim_point = lead_point and Vector3(lead_point.x, lead_point.y, lead_point.z) or Vector3(target_pos.x, target_pos.y, target_pos.z)
        
        -- Manage altitude based on distance to target
        local min_dive_altitude = 20  -- Minimum altitude before starting dive
        local start_dive_distance = profile.dive_distance * 1.5  -- Start diving earlier
        local turn_start_distance = profile.dive_distance * 3.5  -- Start turning before diving
        
        -- Debug output
        I:Log("Target distance: " .. target_distance .. ", turn_start_distance: " .. turn_start_distance)
        
        -- Calculate altitude based on distance
        if target_distance > turn_start_distance then
            -- Maintain cruise altitude while far from target
            local base_point = Vector3(target_pos.x, profile.cruise_altitude, target_pos.z)
            
            -- Use the pre-calculated fan-out offset during cruise
            local offset = missile_fan_offset[missile_id] or {x=0, z=0}
            aim_point = Vector3(
                base_point.x + offset.x * 0.5,  -- Reduce offset during cruise
                base_point.y,
                base_point.z + offset.z * 0.5   -- Reduce offset during cruise
            )
            
        elseif target_distance > start_dive_distance then
            -- Start turning towards target while maintaining altitude
            local turn_ratio = (target_distance - start_dive_distance) / (turn_start_distance - start_dive_distance)
            local target_point = Vector3(target_pos.x, profile.cruise_altitude, target_pos.z)
            
            -- Blend between offset point and direct approach
            aim_point = Vector3.Lerp(
                target_point,
                Vector3(target_pos.x, profile.cruise_altitude, target_pos.z),
                turn_ratio
            )
            
        elseif target_distance > profile.dive_distance then
            -- Start descending gradually while turning
            local dive_ratio = (target_distance - profile.dive_distance) / (start_dive_distance - profile.dive_distance)
            local target_alt = min_dive_altitude + (profile.cruise_altitude - min_dive_altitude) * dive_ratio
            aim_point = Vector3(target_pos.x, math.max(target_alt, min_dive_altitude), target_pos.z)
            
        else
            -- In final approach, aim directly at target with small offset
            aim_point = target_pos + Vector3(0, -2, 0)  -- 2m offset to hit center
        end
        
        -- Add small random offset to prevent perfect circular patterns
        if target_distance < 50 then
            aim_point = aim_point + Vector3(
                (math.random() - 0.5) * 2,
                (math.random() - 0.5) * 2,
                (math.random() - 0.5) * 2
            )
        end
        
        -- Adaptive smoothing with distance-based alpha
        local alpha = GetAdaptiveAlpha(target_idx, target_vel, profile)
        if target_distance < 300 then
            alpha = profile.alpha_max -- More responsive at close range
        end
        
        -- Apply smoothing
        smoothed_aim_points[missile_id] =
            ExponentialSmooth(smoothed_aim_points[missile_id], aim_point, alpha)
        
        -- Set final aim point
        I:SetLuaControlledMissileAimPoint(t, m, 
            smoothed_aim_points[missile_id].x,
            smoothed_aim_points[missile_id].y,
            smoothed_aim_points[missile_id].z
        )
        
        ::continue::
    end
end