%% @doc Selection strategies for neuroevolution.
%%
%% This module provides different strategies for selecting which individuals
%% survive to the next generation and which are used as parents for breeding.
%%
%% == Selection Strategies ==
%%
%% <ul>
%% <li>`top_n/2' - Select top N% by fitness (truncation selection)</li>
%% <li>`tournament/3' - Tournament selection with configurable size</li>
%% <li>`roulette_wheel/1' - Fitness-proportionate selection</li>
%% <li>`random_select/1' - Uniform random selection</li>
%% </ul>
%%
%% All strategies expect individuals to have fitness values already calculated.
%%
%% @author Macula.io
%% @copyright 2025 Macula.io
-module(neuroevolution_selection).

-include("neuroevolution.hrl").

%% API
-export([
    top_n/2,
    tournament/3,
    roulette_wheel/1,
    random_select/1,
    select_parents/2
]).

%%% ============================================================================
%%% API Functions
%%% ============================================================================

%% @doc Select top N individuals by fitness (truncation selection).
%%
%% Sorts population by fitness descending and returns top N.
%% This is elitist selection - only the best survive.
%%
%% Example:
%% ```
%% %% Select top 10 individuals
%% Survivors = neuroevolution_selection:top_n(Population, 10).
%% '''
-spec top_n(Population, N) -> Survivors when
    Population :: [individual()],
    N :: pos_integer(),
    Survivors :: [individual()].
top_n(Population, N) ->
    Sorted = lists:sort(
        fun(A, B) -> A#individual.fitness >= B#individual.fitness end,
        Population
    ),
    lists:sublist(Sorted, N).

%% @doc Tournament selection.
%%
%% Randomly samples TournamentSize individuals from the population,
%% returns the fittest one. Repeat to select multiple individuals.
%%
%% Tournament selection provides moderate selection pressure while
%% maintaining diversity.
%%
%% Example:
%% ```
%% %% Tournament of size 3
%% Winner = neuroevolution_selection:tournament(Population, 3, 1).
%% '''
-spec tournament(Population, TournamentSize, NumSelections) -> Selected when
    Population :: [individual()],
    TournamentSize :: pos_integer(),
    NumSelections :: pos_integer(),
    Selected :: [individual()].
tournament(Population, TournamentSize, NumSelections) ->
    tournament_loop(Population, TournamentSize, NumSelections, []).

tournament_loop(_Population, _TournamentSize, 0, Acc) ->
    lists:reverse(Acc);
tournament_loop(Population, TournamentSize, Remaining, Acc) ->
    %% Sample TournamentSize individuals
    Contestants = random_sample(Population, TournamentSize),
    %% Select the fittest
    Winner = hd(lists:sort(
        fun(A, B) -> A#individual.fitness >= B#individual.fitness end,
        Contestants
    )),
    tournament_loop(Population, TournamentSize, Remaining - 1, [Winner | Acc]).

%% @doc Roulette wheel (fitness-proportionate) selection.
%%
%% Probability of selection is proportional to fitness.
%% Higher fitness = higher probability of being selected.
%%
%% Note: Handles negative fitness by shifting to positive range.
%%
%% Returns a single selected individual.
-spec roulette_wheel(Population) -> Selected when
    Population :: [individual()],
    Selected :: individual().
roulette_wheel(Population) ->
    %% Shift fitness to positive range if needed
    MinFitness = lists:min([I#individual.fitness || I <- Population]),
    Shift = case MinFitness < 0 of
        true -> abs(MinFitness) + 1.0;
        false -> 0.0
    end,

    %% Calculate total shifted fitness
    TotalFitness = lists:sum([I#individual.fitness + Shift || I <- Population]),

    case TotalFitness =< 0 of
        true ->
            %% All fitness zero or negative - use random selection
            random_select(Population);
        false ->
            %% Spin the wheel
            Spin = rand:uniform() * TotalFitness,
            roulette_select(Population, Shift, Spin, 0.0)
    end.

roulette_select([Individual], _Shift, _Target, _Cumulative) ->
    Individual;
roulette_select([Individual | Rest], Shift, Target, Cumulative) ->
    NewCumulative = Cumulative + Individual#individual.fitness + Shift,
    case NewCumulative >= Target of
        true -> Individual;
        false -> roulette_select(Rest, Shift, Target, NewCumulative)
    end.

%% @doc Uniform random selection.
%%
%% Selects a single individual uniformly at random.
%% All individuals have equal probability regardless of fitness.
-spec random_select(Population) -> Selected when
    Population :: [individual()],
    Selected :: individual().
random_select(Population) ->
    Index = rand:uniform(length(Population)),
    lists:nth(Index, Population).

%% @doc Select two parents for breeding.
%%
%% Uses roulette wheel selection to choose parents, ensuring
%% that two different individuals are selected (if population allows).
-spec select_parents(Population, Config) -> {Parent1, Parent2} when
    Population :: [individual()],
    Config :: neuro_config(),
    Parent1 :: individual(),
    Parent2 :: individual().
select_parents(Population, _Config) when length(Population) < 2 ->
    %% Edge case: only one individual
    Single = hd(Population),
    {Single, Single};
select_parents(Population, _Config) ->
    Parent1 = roulette_wheel(Population),
    %% Select second parent, try to avoid same individual
    Parent2 = select_different(Population, Parent1, 5),
    {Parent1, Parent2}.

%%% ============================================================================
%%% Internal Functions
%%% ============================================================================

%% @private
%% @doc Random sample of N items from a list (with replacement).
-spec random_sample([T], non_neg_integer()) -> [T] when T :: term().
random_sample(_List, 0) ->
    [];
random_sample(List, N) ->
    [random_select_from_list(List) | random_sample(List, N - 1)].

%% @private
random_select_from_list(List) ->
    lists:nth(rand:uniform(length(List)), List).

%% @private
%% @doc Try to select a different individual than the given one.
%% Makes MaxAttempts tries before giving up and returning any individual.
-spec select_different([individual()], individual(), non_neg_integer()) -> individual().
select_different(Population, _Exclude, 0) ->
    random_select(Population);
select_different(Population, Exclude, AttemptsLeft) ->
    Candidate = roulette_wheel(Population),
    case Candidate#individual.id =:= Exclude#individual.id of
        true -> select_different(Population, Exclude, AttemptsLeft - 1);
        false -> Candidate
    end.