local AH = wesnoth.require "ai/lua/ai_helper.lua"
local BC = wesnoth.require "ai/lua/battle_calcs.lua"
local M = wesnoth.map

local function get_wolves(cfg)
    local wolves = AH.get_units_with_moves {
        side = wesnoth.current.side,
        { "and", wml.get_child(cfg, "filter") }
    }
    return wolves
end

local function get_prey(cfg)
    -- Note: we cannot pass wml.get_child() directly to AH.get_attackable_enemies()
    -- as the former returns two values and the latter takes optional arguments
    local filter_second = wml.get_child(cfg, "filter_second")
    local prey = AH.get_attackable_enemies(filter_second)
    return prey
end

local ca_wolves_move = {}

function ca_wolves_move:evaluation(cfg)
    if (not get_wolves(cfg)[1]) then return 0 end
    if (not get_prey(cfg)[1]) then return 0 end
    return cfg.ca_score
end

function ca_wolves_move:execution(cfg)
    local wolves = get_wolves(cfg)
    local prey = get_prey(cfg)

    local avoid_units = AH.get_attackable_enemies({ type = cfg.avoid_type })
    local avoid_map = BC.get_attack_map(avoid_units).units

    -- Find prey that is closest to the wolves
    local min_dist, target = 9e99
    for _,prey_unit in ipairs(prey) do
        local dist = 0
        for _,wolf in ipairs(wolves) do
            dist = dist + M.distance_between(wolf.x, wolf.y, prey_unit.x, prey_unit.y)
        end
        if (dist < min_dist) then
            min_dist, target = dist, prey_unit
        end
    end

    -- Now sort wolf from farthest to closest
    table.sort(wolves, function(a, b)
        return M.distance_between(a.x, a.y, target.x, target.y) > M.distance_between(b.x, b.y, target.x, target.y)
    end)

    -- First wolf moves toward target, but tries to stay away from map edges
    local width, height = wesnoth.get_map_size()
    local wolf1 = AH.find_best_move(wolves[1], function(x, y)
        local dist_1t = M.distance_between(x, y, target.x, target.y)
        local rating = - dist_1t
        if (x <= 5) then rating = rating - (6 - x) / 1.4 end
        if (y <= 5) then rating = rating - (6 - y) / 1.4 end
        if (width - x <= 5) then rating = rating - (6 - (width - x)) / 1.4 end
        if (height - y <= 5) then rating = rating - (6 - (height - y)) / 1.4 end

       -- Hexes that avoid_type units can reach get a massive penalty
       if avoid_map:get(x, y) then rating = rating - 1000 end

       return rating
    end)

    local move_result = AH.movefull_stopunit(ai, wolves[1], wolf1)
    -- If the wolf was ambushed, return and reconsider; also if an event removed a wolf
    if (AH.is_incomplete_move(move_result)) then return end
    for _,check_wolf in ipairs(wolves) do
        if (not check_wolf) or (not check_wolf.valid) then return end
    end

    for i = 2,#wolves do
        move = AH.find_best_move(wolves[i], function(x,y)
            local rating = 0

            -- We ideally want wolves to be 2-3 hexes from each other
            -- but this requirement gets weaker and weaker with increasing wolf number
            for j = 1,i-1 do
                local dst = M.distance_between(x, y, wolves[j].x, wolves[j].y)
                rating = rating - (dst - 2.7 * j)^2 / j
            end

            -- Same distance from Wolf 1 and target for all the wolves
            local dist_t = M.distance_between(x, y, target.x, target.y)
            local dist_1t = M.distance_between(wolf1[1], wolf1[2], target.x, target.y)
            rating = rating - (dist_t - dist_1t)^2

            -- Hexes that avoid_type units can reach get a massive penalty
            if avoid_map:get(x, y) then rating = rating - 1000 end

            return rating
        end)

        local move_result = AH.movefull_stopunit(ai, wolves[i], move)
        -- If the wolf was ambushed, return and reconsider; also if an event removed a wolf
        if (AH.is_incomplete_move(move_result)) then return end
        for _,check_wolf in ipairs(wolves) do
            if (not check_wolf) or (not check_wolf.valid) then return end
        end
    end
end

return ca_wolves_move
