<filename>env/Assets/Scripts/CubeSatAgent.cs﻿using System.Collections;
using System.Collections.Generic;
using System.Linq;
using UnityEngine;
using Unity.MLAgents;
using Unity.MLAgents.Actuators;
using Unity.MLAgents.Sensors;

public class CubeSatAgent : Agent
{
    EnvironmentParameters m_ResetParams;
    BufferSensorComponent m_BufferSensor;

    int MINSIZE = 20;
    int MAXSIZE = 100;

    int iTaskSize;
    int[] vTaskMemory;
    int[] vTaskPower;
    bool[] vTaskAssinged;

    public override void Initialize()
    {
        m_ResetParams = Academy.Instance.EnvironmentParameters;
        m_BufferSensor = GetComponent<BufferSensorComponent>();
    }

    public override void OnEpisodeBegin()
    {
        InitEpisode();
    }

    void InitEpisode()
    {
        iTaskSize = Random.Range(MINSIZE, MAXSIZE);
        vTaskMemory = new int[iTaskSize];
        vTaskPower = new int[iTaskSize];
        vTaskAssinged = new bool[MAXSIZE];
        
        for (int i = 0; i < iTaskSize; ++i)
        {
            vTaskMemory[i] = Random.Range(10, 90); // MB
            vTaskPower[i] = Random.Range(15, 70); // Gigacycle/sec
        }
    }

    public override void CollectObservations(VectorSensor sensor)
    {
        for (int i = 0; i < iTaskSize; ++i)
        {
            float[] buf = new float[MAXSIZE + 2];
            buf[i] = 1.0f;
            buf[MAXSIZE] = vTaskMemory[i] * 1e-2f;
            buf[MAXSIZE + 1] = vTaskPower[i] * 1e-2f;
            m_BufferSensor.AppendObservation(buf);
        }
    }

    public override void OnActionReceived(ActionBuffers actions)
    {
        vTaskAssinged[actions.DiscreteActions[0]] = true;
        var afterReward = GetReward();

        //Debug.Log("----------------------------");
        //Debug.Log(afterReward);
        //for (int i = 0; i < iTaskSize; ++i)
        //{
        //    if(vTaskAssinged[i])
        //    {
        //        Debug.Log(vTaskMemory[i].ToString() + "/" + vTaskPower[i].ToString());
        //        break;
        //    }
        //}
        if (actions.DiscreteActions[0] >= iTaskSize)
        {
            SetReward(-1f);
        }
        else
        {
            SetReward(afterReward);
        }
    }

    float GetReward()
    {
        var Xtran_c = 0.8f;    // $/MB
        var Xtran_l = 1.2f;    // $/MB
        var Xcomp_c = 0.8f;    // $/Gigacycles
        var Xcomp_l = 3.0f;     // $/Gigacycles


        float result = 0;
        for (int i = 0; i < iTaskSize; ++i)
        {
            float Ttran = 0f;
            float Ptran = 0f;
            float Tcomp = 0f;
            float Pcomp = 0f;

            if (vTaskAssinged[i]) // CubeSats
            {
                var bandwidth = 2e+3f * (1f / 20f);    // MHz
                var power = 1e+2f * (1f / 20f);         // Gigacycles/sec
                var SNR = 1f;
                Ttran = vTaskMemory[i] / (bandwidth * Mathf.Log(1 + SNR, 2));
                Ptran = Xtran_c * bandwidth;
                Tcomp = vTaskPower[i] / power;
                Pcomp = Xcomp_c * power;
            }
            else // LMS
            {
                var bandwidth = 2e+3f * (1f / 100f);    // MHz
                var power = 8e+2f * (1f / 100f);         // Gigacycles/sec
                var SNR = 1f;
                Ttran = vTaskMemory[i] / (bandwidth * Mathf.Log(1 + SNR, 2));
                Ptran = Xtran_l * bandwidth;
                Tcomp = vTaskPower[i] / power;
                Pcomp = Xcomp_l * power;
            }

            float Tser = Ttran + Tcomp;
            float Pser = Ptran + Pcomp;
            float eta = 0.7f;
            result += (eta * Tser) + ((1 - eta) * Pser);
        }
        return 1e+3f / result;
    }
}
