// Copyright © 2025 Stephane Capponi and individual contributors. All Rights Reserved.

#pragma once

#include "CoreMinimal.h"
#include "UObject/Interface.h"
#include "UnrealMLAgents/Sensors/ObservationSpec.h"
#include "UnrealMLAgents/Sensors/ObservationType.h"
#include "UnrealMLAgents/DimensionProperty.h"
#include "UnrealMLAgents/Sensors/ObservationWriter.h"
#include "ISensor.generated.h"

/**
 * @interface IISensor
 * @brief Interface for sensors responsible for generating observations in reinforcement learning environments.
 *
 * The `IISensor` interface defines methods for generating, writing, and resetting observations for agents.
 * These observations are used as inputs for decision-making in reinforcement learning models.
 */
UINTERFACE(MinimalAPI)
class UISensor : public UInterface
{
	GENERATED_BODY()
};

/**
 * @class IISensor
 * @brief Interface for sensors that generate agent observations.
 *
 * The `IISensor` interface defines the methods needed to collect and write observations that describe an
 * agent's environment. Sensors implementing this interface are responsible for providing data to the agent's
 * decision-making system.
 */
class UNREALMLAGENTS_API IISensor
{
	GENERATED_BODY()

public:
	/**
	 * @brief Returns a description of the observations that will be generated by the sensor.
	 *
	 * This method returns an `FObservationSpec` object that describes the structure of the observations
	 * the sensor will generate. This includes details such as the shape and type of the observations.
	 *
	 * @return An `FObservationSpec` object describing the sensor's observations.
	 */
	virtual FObservationSpec GetObservationSpec() = 0;

	/**
	 * @brief Writes the sensor's observation data to the provided `ObservationWriter`.
	 *
	 * This method collects the sensor's current observations and writes them into the specified `ObservationWriter`.
	 * It should not modify the internal state of the sensor, as it may be called multiple times per agent step.
	 *
	 * @param Writer The writer to which the observations will be written.
	 * @return The number of elements written.
	 */
	virtual int32 Write(ObservationWriter& Writer) = 0;

	/**
	 * @brief Updates the internal state of the sensor.
	 *
	 * This method is called once per agent step to update the sensor's internal state. Implementations should
	 * handle any necessary updates, such as processing new observations.
	 */
	virtual void Update() = 0;

	/**
	 * @brief Resets the sensor's internal state.
	 *
	 * This method is called at the end of an agent's episode to reset the sensor's state. Most implementations
	 * can leave this empty if no reset behavior is required.
	 */
	virtual void Reset() = 0;

	/**
	 * @brief Returns the name of the sensor.
	 *
	 * The name of the sensor is used to ensure deterministic sorting of sensors for an agent. The naming must be
	 * consistent across all sensors and agents to maintain proper ordering.
	 *
	 * @return The name of the sensor as an `FString`.
	 */
	virtual FString GetName() const = 0;
};

/**
 * @class USensorExtensions
 * @brief Helper class that provides utility methods for sensors.
 *
 * The `USensorExtensions` class contains helper methods that can be used with sensors implementing the `IISensor`
 * interface. These methods are useful for operations such as calculating the total number of elements in an
 * observation.
 */
UCLASS()
class UNREALMLAGENTS_API USensorExtensions : public UObject
{
	GENERATED_BODY()

public:
	/**
	 * @brief Calculates the total number of elements in the sensor's observation.
	 *
	 * This method computes the total number of elements in the sensor's observation by multiplying the dimensions
	 * of the observation's shape.
	 *
	 * @param Sensor The sensor whose observation size is to be calculated.
	 * @return The total number of elements in the observation.
	 */
	UFUNCTION(BlueprintCallable, Category = "Sensor")
	static int32 ObservationSize(TScriptInterface<IISensor> Sensor)
	{
		FObservationSpec ObsSpec = Sensor->GetObservationSpec();
		int32			 Count = 1;
		for (int32 i = 0; i < ObsSpec.GetRank(); ++i)
		{
			Count *= ObsSpec.GetShape()[i];
		}
		return Count;
	}
};

/**
 * @class USensorUtils
 * @brief Utility class for sorting and managing sensors.
 *
 * The `USensorUtils` class contains utility methods for managing and sorting arrays of sensors. These utilities
 * ensure consistent behavior, such as sorting sensors by name.
 */
UCLASS()
class UNREALMLAGENTS_API USensorUtils : public UObject
{
	GENERATED_BODY()

public:
	/**
	 * @brief Sorts an array of sensors by their name.
	 *
	 * This method sorts an array of sensors in place, using the names of the sensors to ensure consistent
	 * sorting behavior. Sorting by name helps maintain deterministic order across sensors.
	 *
	 * @param Sensors The array of sensors to be sorted.
	 */
	UFUNCTION(BlueprintCallable, Category = "Sensor")
	static void SortSensors(TArray<TScriptInterface<IISensor>>& Sensors)
	{
		Sensors.Sort([](const TScriptInterface<IISensor>& A, const TScriptInterface<IISensor>& B) {
			return A->GetName().Compare(B->GetName()) < 0;
		});
	}
};
