using System;
using System.Collections.Generic;
using Unity.InferenceEngine;
using Unity.MLAgents.Inference;
using UnityEngine;
using DeviceType = Unity.InferenceEngine.DeviceType;

namespace Unity.MLAgents.Sensors
{
    /// <summary>
    /// Allows sensors to write to both TensorProxy and float arrays/lists.
    /// </summary>
    public class ObservationWriter
    {
        IList<float> m_Data;
        int m_Offset;

        TensorProxy m_Proxy;
        int m_Batch;
        int m_Capacity;

        TensorShape m_TensorShape;

        /// <summary>
        /// Initializes a new instance of the <see cref="ObservationWriter"/> class.
        /// </summary>
        public ObservationWriter() { }

        /// <summary>
        /// Set the writer to write to an IList at the given channelOffset.
        /// </summary>
        /// <param name="data">Float array or list that will be written to.</param>
        /// <param name="observationSpec">ObservationSpec of the observation to be written</param>
        /// <param name="offset">Offset from the start of the float data to write to.</param>
        internal void SetTarget(IList<float> data, ObservationSpec observationSpec, int offset)
        {
            SetTarget(data, observationSpec.Shape, offset);
        }

        /// <summary>
        /// Set the writer to write to an IList at the given channelOffset.
        /// </summary>
        /// <param name="data">Float array or list that will be written to.</param>
        /// <param name="shape">Shape of the observations to be written.</param>
        /// <param name="offset">Offset from the start of the float data to write to.</param>
        internal void SetTarget(IList<float> data, InplaceArray<int> shape, int offset)
        {
            m_Data = data;
            m_Offset = offset;
            m_Proxy = null;
            m_Batch = 0;

            if (shape.Length == 1)
            {
                m_TensorShape = new TensorShape(m_Batch, shape[0]);
            }
            else if (shape.Length == 2)
            {
                m_TensorShape = new TensorShape(new[] { m_Batch, 1, shape[0], shape[1] });
            }
            else
            {
                m_TensorShape = new TensorShape(m_Batch, shape[0], shape[1], shape[2]);
            }
        }

        /// <summary>
        /// Set the writer to write to a TensorProxy at the given batch and channel offset.
        /// </summary>
        /// <param name="tensorProxy">Tensor proxy that will be written to.</param>
        /// <param name="batchIndex">Batch index in the tensor proxy (i.e. the index of the Agent).</param>
        /// <param name="channelOffset">Offset from the start of the channel to write to.</param>
        internal void SetTarget(TensorProxy tensorProxy, int batchIndex, int channelOffset)
        {
            m_Proxy = tensorProxy;
            m_Batch = batchIndex;
            m_Offset = channelOffset;
            m_Data = null;
            m_TensorShape = m_Proxy.data.shape;
            m_Capacity = m_TensorShape.rank >= 2 ? m_TensorShape[1] : 0;
        }

        /// <summary>
        /// 1D write access at a specified index. Use AddList if possible instead.
        /// </summary>
        /// <param name="index">Index to write to.</param>
        public float this[int index]
        {
            set
            {
                if (m_Data != null)
                {
                    m_Data[index + m_Offset] = value;
                }
                else
                {
                    if (index + m_Offset < 0 || index + m_Offset >= m_Capacity)
                        return;

                    m_Proxy.data.CompleteAllPendingOperations();
                    ((Tensor<float>)m_Proxy.data)[m_Batch, index + m_Offset] = value;
                }
            }
        }

        /// <summary>
        /// Write access at the specified channel and width.
        /// </summary>
        /// <param name="ch">Channels</param>
        /// <param name="w">Width</param>
        public float this[int ch, int w]
        {
            set
            {
                if (m_Data != null)
                {
                    m_Data[ch * m_TensorShape[m_TensorShape.length - 1] + w] = value;
                }
                else
                {
                    m_Proxy.data.CompleteAllPendingOperations();

                    ((Tensor<float>)m_Proxy.data)[m_Batch, ch, w] = value;
                }
            }
        }

        /// <summary>
        /// 3D write access at the specified height, width, and channel.
        /// </summary>
        /// <param name="h">Height</param>
        /// <param name="w">Width</param>
        /// <param name="ch">Channels</param>
        public float this[int ch, int h, int w]
        {
            set
            {
                if (m_Data != null)
                {
                    if (h < 0 || h >= m_TensorShape.Height())
                    {
                        throw new IndexOutOfRangeException($"height value {h} must be in range [0, {m_TensorShape.Height() - 1}]");
                    }

                    if (w < 0 || w >= m_TensorShape.Width())
                    {
                        throw new IndexOutOfRangeException($"width value {w} must be in range [0, {m_TensorShape.Width() - 1}]");
                    }

                    if (ch < 0 || ch >= m_TensorShape.Channels())
                    {
                        throw new IndexOutOfRangeException($"channel value {ch} must be in range [0, {m_TensorShape.Channels() - 1}]");
                    }

                    var index = m_TensorShape.Index(m_Batch, ch + m_Offset, h, w);
                    m_Data[index] = value;
                }
                else
                {
                    if (ch + m_Offset < 0 || ch + m_Offset >= m_TensorShape.Channels() ||
                        h < 0 || h >= m_TensorShape.Height() ||
                        w < 0 || w >= m_TensorShape.Width())
                        return;

                    m_Proxy.data.CompleteAllPendingOperations();
                    ((Tensor<float>)m_Proxy.data)[m_Batch, ch + m_Offset, h, w] = value;
                }
            }
        }

        /// <summary>
        /// Write the list of floats.
        /// </summary>
        /// <param name="data">The actual list of floats to write.</param>
        /// <param name="writeOffset">Optional write offset to start writing from.</param>
        public void AddList(IList<float> data, int writeOffset = 0)
        {
            if (m_Data != null)
            {
                for (var index = 0; index < data.Count; index++)
                {
                    var val = data[index];
                    m_Data[index + m_Offset + writeOffset] = val;
                }
            }
            else
            {
                m_Proxy.data.CompleteAllPendingOperations();

                var maxCount = Math.Min(data.Count, Math.Max(0, m_Capacity - m_Offset - writeOffset));
                for (var index = 0; index < maxCount; index++)
                {
                    ((Tensor<float>)m_Proxy.data)[m_Batch, index + m_Offset + writeOffset] = data[index];
                }
            }
        }

        /// <summary>
        /// Write the Vector3 components.
        /// </summary>
        /// <param name="vec">The Vector3 to be written.</param>
        /// <param name="writeOffset">Optional write offset.</param>
        public void Add(Vector3 vec, int writeOffset = 0)
        {
            if (m_Data != null)
            {
                m_Data[m_Offset + writeOffset + 0] = vec.x;
                m_Data[m_Offset + writeOffset + 1] = vec.y;
                m_Data[m_Offset + writeOffset + 2] = vec.z;
            }
            else
            {
                var start = m_Offset + writeOffset;
                var remaining = m_Capacity - start;
                if (remaining <= 0) return;

                m_Proxy.data.CompleteAllPendingOperations();
                ((Tensor<float>)m_Proxy.data)[m_Batch, start + 0] = vec.x;
                if (remaining <= 1) return;
                ((Tensor<float>)m_Proxy.data)[m_Batch, start + 1] = vec.y;
                if (remaining <= 2) return;
                ((Tensor<float>)m_Proxy.data)[m_Batch, start + 2] = vec.z;
            }
        }

        /// <summary>
        /// Write the Vector4 components.
        /// </summary>
        /// <param name="vec">The Vector4 to be written.</param>
        /// <param name="writeOffset">Optional write offset.</param>
        public void Add(Vector4 vec, int writeOffset = 0)
        {
            if (m_Data != null)
            {
                m_Data[m_Offset + writeOffset + 0] = vec.x;
                m_Data[m_Offset + writeOffset + 1] = vec.y;
                m_Data[m_Offset + writeOffset + 2] = vec.z;
                m_Data[m_Offset + writeOffset + 3] = vec.w;
            }
            else
            {
                var start = m_Offset + writeOffset;
                var remaining = m_Capacity - start;
                if (remaining <= 0) return;

                m_Proxy.data.CompleteAllPendingOperations();
                ((Tensor<float>)m_Proxy.data)[m_Batch, start + 0] = vec.x;
                if (remaining <= 1) return;
                ((Tensor<float>)m_Proxy.data)[m_Batch, start + 1] = vec.y;
                if (remaining <= 2) return;
                ((Tensor<float>)m_Proxy.data)[m_Batch, start + 2] = vec.z;
                if (remaining <= 3) return;
                ((Tensor<float>)m_Proxy.data)[m_Batch, start + 3] = vec.w;
            }
        }

        /// <summary>
        /// Write the Quaternion components.
        /// </summary>
        /// <param name="quat">The Quaternion to be written.</param>
        /// <param name="writeOffset">Optional write offset.</param>
        public void Add(Quaternion quat, int writeOffset = 0)
        {
            if (m_Data != null)
            {
                m_Data[m_Offset + writeOffset + 0] = quat.x;
                m_Data[m_Offset + writeOffset + 1] = quat.y;
                m_Data[m_Offset + writeOffset + 2] = quat.z;
                m_Data[m_Offset + writeOffset + 3] = quat.w;
            }
            else
            {
                var start = m_Offset + writeOffset;
                var remaining = m_Capacity - start;
                if (remaining <= 0) return;

                m_Proxy.data.CompleteAllPendingOperations();
                ((Tensor<float>)m_Proxy.data)[m_Batch, start + 0] = quat.x;
                if (remaining <= 1) return;
                ((Tensor<float>)m_Proxy.data)[m_Batch, start + 1] = quat.y;
                if (remaining <= 2) return;
                ((Tensor<float>)m_Proxy.data)[m_Batch, start + 2] = quat.z;
                if (remaining <= 3) return;
                ((Tensor<float>)m_Proxy.data)[m_Batch, start + 3] = quat.w;
            }
        }
    }

    /// <summary>
    /// Provides extension methods for the ObservationWriter.
    /// </summary>
    public static class ObservationWriterExtension
    {
        /// <summary>
        /// Writes a Texture2D into a ObservationWriter.
        /// </summary>
        /// <param name="obsWriter">
        /// Writer to fill with Texture data.
        /// </param>
        /// <param name="texture">
        /// The texture to be put into the tensor.
        /// </param>
        /// <param name="grayScale">
        /// If set to <c>true</c> the textures will be converted to grayscale before
        /// being stored in the tensor.
        /// </param>
        /// <returns>The number of floats written</returns>
        public static int WriteTexture(
            this ObservationWriter obsWriter,
            Texture2D texture,
            bool grayScale)
        {
            if (texture.format == TextureFormat.RGB24)
            {
                return obsWriter.WriteTextureRGB24(texture, grayScale);
            }

            var width = texture.width;
            var height = texture.height;

            var texturePixels = texture.GetPixels32();

            // During training, we convert from Texture to PNG before sending to the trainer, which has the
            // effect of flipping the image. We need another flip here at inference time to match this.
            for (var h = height - 1; h >= 0; h--)
            {
                for (var w = 0; w < width; w++)
                {
                    var currentPixel = texturePixels[(height - h - 1) * width + w];

                    if (grayScale)
                    {
                        obsWriter[0, h, w] =
                            (currentPixel.r + currentPixel.g + currentPixel.b) / 3f / 255.0f;
                    }
                    else
                    {
                        // For Color32, the r, g and b values are between 0 and 255.
                        obsWriter[0, h, w] = currentPixel.r / 255.0f;
                        obsWriter[1, h, w] = currentPixel.g / 255.0f;
                        obsWriter[2, h, w] = currentPixel.b / 255.0f;
                    }
                }
            }

            return height * width * (grayScale ? 1 : 3);
        }

        internal static int WriteTextureRGB24(
            this ObservationWriter obsWriter,
            Texture2D texture,
            bool grayScale
        )
        {
            var width = texture.width;
            var height = texture.height;

            var rawBytes = texture.GetRawTextureData<byte>();

            // During training, we convert from Texture to PNG before sending to the trainer, which has the
            // effect of flipping the image. We need another flip here at inference time to match this.
            for (var h = height - 1; h >= 0; h--)
            {
                for (var w = 0; w < width; w++)
                {
                    var offset = (height - h - 1) * width + w;
                    var r = rawBytes[3 * offset];
                    var g = rawBytes[3 * offset + 1];
                    var b = rawBytes[3 * offset + 2];

                    if (grayScale)
                    {
                        obsWriter[0, h, w] = (r + g + b) / 3f / 255.0f;
                    }
                    else
                    {
                        // For Color32, the r, g and b values are between 0 and 255.
                        obsWriter[0, h, w] = r / 255.0f;
                        obsWriter[1, h, w] = g / 255.0f;
                        obsWriter[2, h, w] = b / 255.0f;
                    }
                }
            }

            return height * width * (grayScale ? 1 : 3);
        }
    }
}