using OneOf;
using RLMatrix; //NuGet v0.2.0

namespace ReinforcementLearning.Trivial;

internal class Program
{
    static void Main(string[] args)
    {
        var options = new PPOAgentOptions(lr: .005f);
        List<IEnvironment<float[]>> environments = [new TrivialEnvironment(), new TrivialEnvironment(), new TrivialEnvironment()];
        var agent = new PPOAgent<float[]>(options, environments, new PPONetProviderBase<float[]>(16, 2));

        Console.WriteLine($"Pre-training: {Evaluate(agent)}");

        for (int i = 0; i < 10000; i++)
        {
            Console.Write(i);
            agent.Step();
            Console.Write("\b\b\b\b\b\b\b\b\b\b\b\b");
        }

        Console.WriteLine($"Post-training: {Evaluate(agent)}");
    }

    private static float Evaluate(PPOAgent<float[]> agent, int count = 1000)
    {
        float[] state = new float[1];
        int correct = 0;
        for (int i = 0; i < count; i++)
        {
            state[0] = TrivialEnvironment.RandomValue();
            (int[] greedyDiscreteActions, float[] meanContinuousActions) = agent.SelectAction(state, false);
            if (state[0] == greedyDiscreteActions[0])
            {
                correct++;
            }
        }
        return (float)correct / count;
    }

    private sealed class TrivialEnvironment : IEnvironment<float[]>
    {
        public const float CorrectAnswerReward = 1;
        public const float WrongAnswerPenalty = -1;

        private readonly float[] state = [RandomValue()];

        public static int RandomValue()
        {
            return Random.Shared.Next(2);
        }

        public int stepCounter { get; set; }
        public int maxSteps { get; set; } = int.MinValue;
        public bool isDone { get; set; }
        public OneOf<int, (int, int)> stateSize { get; set; } = 1;
        public int[] actionSize { get; set; } = [2];

        public float[] GetCurrentState() => state;

        public TrivialEnvironment() => Initialise();

        public void Initialise() => Reset();

        public void Reset()
        {
            stepCounter = 0;
            isDone = false;
        }

        public float Step(int[] actionsIds)
        {
            float input = state[0];
            float output = actionsIds[0];

            state[0] = RandomValue();

            return input == output ? CorrectAnswerReward : WrongAnswerPenalty;
        }
    }
}

private sealed class Env : IContinuousEnvironment<float[]>
{
	public int stepCounter { get; set; }
	public int maxSteps { get; set; }
	public bool isDone { get; set; }
	public OneOf<int, (int, int)> stateSize { get; set; }
	public int[] actionSize { get; set; }
	public (float min, float max)[] continuousActionBounds { get; set; }

	private readonly float[] _inputs;

	public Env()
	{
		_inputs = new float[3];
		Initialise();
	}

	public void Initialise()
	{
		// new physics
		stepCounter = 0;
		maxSteps = 100_000;
		stateSize = 3;
		actionSize = [3];
		continuousActionBounds = [(-1f, 1f), (-1f, 1f), (-1f, 1f)];
		// physics reset
		isDone = false;
	}

	public float[] GetCurrentState()
	{
		_inputs[0] = 1f;
		_inputs[1] = 0.5f;
		_inputs[2] = -1f;
		return _inputs;
	}

	public void Reset()
	{
		// physics reset
		isDone = false;
		stepCounter = 0;
	}

	public float Step(int[] actionsIds)
	{
		throw new NotImplementedException();
	}
	public float Step(int[] discreteActions, float[] continuousActions)
	{
		return continuousActions[0];
	}
}