ReCo.jl/src/RL/Envs/OriginCompassEnv.jl

"""
This environment corresponds to the strategy of the origin of the coordinates system with a compass.
"""

using ..ReCo: ReCo

struct OriginCompassEnv <: Env
    shared::EnvSharedProps

    distance_state_space::Vector{Interval}
    direction_angle_state_space::Vector{Interval}
    position_angle_state_space::Vector{Interval}

    function OriginCompassEnv(;
        n_distance_states::Int64=3,
        n_direction_angle_states::Int64=3,
        n_position_angle_states::Int64=8,
        args,
    )
        @assert n_distance_states > 1
        @assert n_direction_angle_states > 1
        @assert n_position_angle_states > 1

        direction_angle_state_space = gen_angle_state_space(n_direction_angle_states)
        position_angle_state_space = gen_angle_state_space(n_position_angle_states)

        min_distance = 0.0
        max_distance = sqrt(2) * args.half_box_len

        distance_state_space = gen_distance_state_space(
            min_distance, max_distance, n_distance_states
        )

        n_states = n_distance_states * n_direction_angle_states * n_position_angle_states

        state_spaces_labels = gen_state_spaces_labels(
            ("d", "\\theta", "\\alpha"),
            (distance_state_space, direction_angle_state_space, position_angle_state_space),
        )

        shared = EnvSharedProps(
            n_states,
            (n_distance_states, n_direction_angle_states, n_position_angle_states),
            state_spaces_labels,
        )

        return new(
            shared,
            distance_state_space,
            direction_angle_state_space,
            position_angle_state_space,
        )
    end
end

mutable struct OriginCompassEnvHelper <: EnvHelper
    shared::EnvHelperSharedProps

    center_of_mass::SVector{2,Float64}
    gyration_tensor_eigvec_to_smaller_eigval::SVector{2,Float64}
    gyration_tensor_eigvec_to_bigger_eigval::SVector{2,Float64}

    half_box_len::Float64
    max_elliptical_distance::Float64

    function OriginCompassEnvHelper(shared::EnvHelperSharedProps, half_box_len::Float64)
        max_elliptical_distance = sqrt(
            half_box_len^2 + (half_box_len / shared.elliptical_b_a_ratio)^2
        )

        return new(
            shared,
            SVector(0.0, 0.0),
            SVector(0.0, 0.0),
            SVector(0.0, 0.0),
            half_box_len,
            max_elliptical_distance,
        )
    end
end

function gen_env_helper(::OriginCompassEnv, env_helper_shared::EnvHelperSharedProps; args)
    return OriginCompassEnvHelper(env_helper_shared, args.half_box_len)
end

function pre_integration_hook!(::OriginCompassEnvHelper)
    return nothing
end

function state_update_helper_hook!(
    ::OriginCompassEnvHelper,
    id1::Int64,
    id2::Int64,
    r⃗₁₂::SVector{2,Float64},
    distance²::Float64,
)
    return nothing
end

function state_update_hook!(
    env_helper::OriginCompassEnvHelper, particles::Vector{ReCo.Particle}
)
    n_particles = env_helper.shared.n_particles

    env = env_helper.shared.env

    env_helper.center_of_mass = ReCo.center_of_mass(particles, env_helper.half_box_len)

    for particle_id in 1:n_particles
        vec_to_origin = -particles[particle_id].c
        distance_to_center_of_mass = ReCo.norm2d(vec_to_origin)
        distance_state_ind = find_state_ind(
            distance_to_center_of_mass, env.distance_state_space
        )

        si, co = sincos(particles[particle_id].φ)
        direction_angle = ReCo.angle2(SVector(co, si), vec_to_origin)
        direction_state_ind = find_state_ind(
            direction_angle, env.direction_angle_state_space
        )

        position_angle = atan(si, co)
        position_state_ind = find_state_ind(position_angle, env.position_angle_state_space)

        state_id = env.shared.state_id_tensor[
            distance_state_ind, direction_state_ind, position_state_ind
        ]

        env_helper.shared.states_id[particle_id] = state_id
    end

    v1, v2 = ReCo.gyration_tensor_eigvecs(
        particles, env_helper.half_box_len, env_helper.center_of_mass
    )

    env_helper.gyration_tensor_eigvec_to_smaller_eigval = v1
    env_helper.gyration_tensor_eigvec_to_bigger_eigval = v2

    return nothing
end

function update_reward!(
    env::OriginCompassEnv, env_helper::OriginCompassEnvHelper, particle::ReCo.Particle
)
    elliptical_distance = ReCo.elliptical_distance(
        particle.c,
        env_helper.center_of_mass,
        env_helper.gyration_tensor_eigvec_to_smaller_eigval,
        env_helper.gyration_tensor_eigvec_to_bigger_eigval,
        env_helper.shared.elliptical_b_a_ratio,
        env_helper.half_box_len,
    )

    reward = minimizing_reward(elliptical_distance, env_helper.max_elliptical_distance)

    set_normalized_reward!(env, reward, env_helper)

    return nothing
end