using System;
using System.Collections.Generic;
using System.Globalization;
using System.IO;
using System.Linq;
using System.Text;
using Newtonsoft.Json;
using Newtonsoft.Json.Linq;
using Newtonsoft.Json.Serialization;
using Unity.Mathematics;
using UnityEngine;
using UnityEngine.Perception.GroundTruth.DataModel;

namespace UnityEngine.Perception.GroundTruth.Consumers
{
    public class PerceptionResolver : DefaultContractResolver
    {
        protected override JsonObjectContract CreateObjectContract(Type objectType)
        {
            var contract = base.CreateObjectContract(objectType);
            if (objectType == typeof(Vector3) ||
                objectType == typeof(Vector2) ||
                objectType == typeof(Color) ||
                objectType == typeof(Quaternion))
            {
                contract.Converter = PerceptionConverter.Instance;
            }

            return contract;
        }
    }

    public class PerceptionConverter : JsonConverter
    {
        public static PerceptionConverter Instance = new PerceptionConverter();

        public override void WriteJson(JsonWriter writer, object value, JsonSerializer serializer)
        {
            switch (value)
            {
                case Vector3 v3:
                {
                    writer.WriteStartArray();
                    writer.WriteValue(v3.x);
                    writer.WriteValue(v3.y);
                    writer.WriteValue(v3.z);
                    writer.WriteEndArray();
                    break;
                }
                case Vector2 v2:
                {
                    writer.WriteStartArray();
                    writer.WriteValue(v2.x);
                    writer.WriteValue(v2.y);
                    writer.WriteEndArray();
                    break;
                }
                case Color32 rgba:
                {
                    writer.WriteStartObject();
                    writer.WritePropertyName("r");
                    writer.WriteValue(rgba.r);
                    writer.WritePropertyName("g");
                    writer.WriteValue(rgba.g);
                    writer.WritePropertyName("b");
                    writer.WriteValue(rgba.b);
                    writer.WritePropertyName("a");
                    writer.WriteValue(rgba.a);
                    writer.WriteEndObject();
                    break;
                }
                case Quaternion quaternion:
                {
                    writer.WriteStartObject();
                    writer.WritePropertyName("x");
                    writer.WriteValue(quaternion.x);
                    writer.WritePropertyName("y");
                    writer.WriteValue(quaternion.y);
                    writer.WritePropertyName("z");
                    writer.WriteValue(quaternion.z);
                    writer.WritePropertyName("w");
                    writer.WriteValue(quaternion.w);
                    writer.WriteEndObject();
                    break;
                }
            }
        }

        public override object ReadJson(JsonReader reader, Type objectType, object existingValue, JsonSerializer serializer)
        {
            return null;
        }

        public override bool CanConvert(Type objectType)
        {
            if (objectType == typeof(Vector3)) return true;
            if (objectType == typeof(Vector2)) return true;
            if (objectType == typeof(Quaternion)) return true;
            return objectType == typeof(Color32);
        }
    }

    public class OldPerceptionConsumer : ConsumerEndpoint
    {
        static readonly string version = "0.1.1";

        const string k_BaseDirectory = "D:/PerceptionOutput/KickinItOldSchool";
        const int k_CapturesPerFile = 150;
        const int k_MetricsPerFile = 150;

        internal JsonSerializer Serializer { get; }= new JsonSerializer { ContractResolver = new PerceptionResolver()};

        string m_CurrentPath;
        string m_DatasetPath;
        string m_RgbPath;

        // ReSharper disable NotAccessedField.Local
        [Serializable]
        struct SensorInfo
        {
            public string id;
            public string modality;
            public string description;
        }
        // ReSharper enable NotAccessedField.Local

        Dictionary<string, SensorInfo> m_SensorMap = new Dictionary<string, SensorInfo>();
        Dictionary<string, AnnotationDefinition> m_RegisteredAnnotations = new Dictionary<string, AnnotationDefinition>();
        Dictionary<string, MetricDefinition> m_RegisteredMetrics = new Dictionary<string, MetricDefinition>();

        List<PerceptionCapture> m_CurrentCaptures = new List<PerceptionCapture>();

        protected override bool IsComplete()
        {
            return true;
        }

        public override object Clone()
        {
            var cloned = ScriptableObject.CreateInstance<OldPerceptionConsumer>();
            return cloned;
        }

        internal string VerifyDirectoryWithGuidExists(string directoryPrefix, bool appendGuid = true)
        {
            var dirs = Directory.GetDirectories(m_CurrentPath);
            var found = string.Empty;

            foreach (var dir in dirs)
            {
                var dirName = new DirectoryInfo(dir).Name;
                if (dirName.StartsWith(directoryPrefix))
                {
                    found = dir;
                    break;
                }
            }

            if (found == string.Empty)
            {
                var dirName = appendGuid ? $"{directoryPrefix}{Guid.NewGuid().ToString()}" : directoryPrefix;
                found = Path.Combine(m_CurrentPath, dirName);
                Directory.CreateDirectory(found);
            }

            return found;
        }

        public override void AnnotationRegistered(AnnotationDefinition annotationDefinition)
        {
            if (m_RegisteredAnnotations.ContainsKey(annotationDefinition.id))
            {
                Debug.LogError("Tried to register an annotation twice");
                return;
            }

            m_RegisteredAnnotations[annotationDefinition.id] = annotationDefinition;
        }

        public override void MetricRegistered(MetricDefinition metricDefinition)
        {
            if (m_RegisteredMetrics.ContainsKey(metricDefinition.id))
            {
                Debug.LogError("Tried to register a metric twice");
                return;
            }

            m_RegisteredMetrics[metricDefinition.id] = metricDefinition;
        }

        public override void SensorRegistered(SensorDefinition sensor)
        {
            if (m_SensorMap.ContainsKey(sensor.id))
            {
                Debug.LogError("Tried to register a sensor twice");
                return;
            }

            m_SensorMap[sensor.id] = new SensorInfo
            {
                id = sensor.id,
                modality = sensor.modality,
                description = sensor.definition
            };
        }

        public string RemoveDatasetPathPrefix(string path)
        {
            return path.Replace(m_CurrentPath + "\\", string.Empty);
        }

        public override void SimulationStarted(SimulationMetadata metadata)
        {
            // Create a directory guid...
            var path = Guid.NewGuid().ToString();

            m_CurrentPath =  Path.Combine(k_BaseDirectory, path);
            Directory.CreateDirectory(m_CurrentPath);

            m_DatasetPath = VerifyDirectoryWithGuidExists("Dataset");
            m_RgbPath = VerifyDirectoryWithGuidExists("RGB");
        }

        public override void FrameGenerated(Frame frame)
        {
            var seqId = frame.sequence.ToString();

            // Only support one image file right now
            var path = "";

            var annotations = new JArray();
            RgbSensor rgbSensor = null;
            if (frame.sensors.Count() == 1)
            {
                var sensor = frame.sensors.First();
                if (sensor is RgbSensor rgb)
                {
                    rgbSensor = rgb;
                    path = WriteOutImageFile(frame.frame, rgb);
                }

                foreach (var annotation in sensor.annotations)
                {
                    m_RegisteredAnnotations.TryGetValue(annotation.Id, out var def);
                    var defId = def?.id ?? string.Empty;
                    var json = OldPerceptionJsonFactory.Convert(this, frame, annotation.Id, defId, annotation);
                    if (json != null) annotations.Add(json);
                }
            }

            foreach (var metric in  frame.metrics)
            {
                AddMetricToReport(metric);
            }

            // Scenarios report frames before captures, so skip it if a capture is not ready
            if (rgbSensor != null)
            {
                var capture = new PerceptionCapture
                {
                    id = $"frame_{frame.frame}",
                    filename = RemoveDatasetPathPrefix(path),
                    format = "PNG",
                    sequence_id = seqId,
                    step = frame.step,
                    timestamp = frame.timestamp,
                    sensor = PerceptionRgbSensor.Convert(rgbSensor),
                    annotations = annotations
                };

                m_CurrentCaptures.Add(capture);
            }

            WriteCaptures();
        }

        void WriteMetrics(bool flush = false)
        {
            if (flush || m_MetricsReady.Count > k_MetricsPerFile)
            {
                WriteMetricsFile(m_MetricOutCount++, m_MetricsReady);
                m_MetricsReady.Clear();
            }
        }

        void WriteCaptures(bool flush = false)
        {
            if (flush || m_CurrentCaptures.Count >= k_CapturesPerFile)
            {
                WriteCaptureFile(m_CurrentCaptureIndex++, m_CurrentCaptures);
                m_CurrentCaptures.Clear();
            }
        }

        public override void SimulationCompleted(CompletionMetadata metadata)
        {
            WriteSensorsFile();
            WriteAnnotationsDefinitionsFile();
            WriteMetricsDefinitionsFile();

            WriteCaptures(true);
            WriteMetrics(true);
        }

        int m_CurrentCaptureIndex;

        string WriteOutImageFile(int frame, RgbSensor rgb)
        {
            var path = Path.Combine(m_RgbPath, $"{rgb.sensorType}_{frame}.png");
            var file = File.Create(path, 4096);
            file.Write(rgb.buffer, 0, rgb.buffer.Length);
            file.Close();
            return path;
        }

        void WriteJTokenToFile(string filePath, PerceptionJson json)
        {
            WriteJTokenToFile(filePath,  JToken.FromObject(json, Serializer));
        }

        void WriteJTokenToFile(string filePath, MetricsJson json)
        {
            WriteJTokenToFile(filePath, JToken.FromObject(json, Serializer));
        }

        static void WriteJTokenToFile(string filePath, JToken json)
        {
            var stringWriter = new StringWriter(new StringBuilder(256), CultureInfo.InvariantCulture);
            using (var jsonTextWriter = new JsonTextWriter(stringWriter))
            {
                jsonTextWriter.Formatting = Formatting.Indented;
                json.WriteTo(jsonTextWriter);
            }

            var contents = stringWriter.ToString();

            File.WriteAllText(filePath, contents);
        }

        void WriteAnnotationsDefinitionsFile()
        {
            var defs = new JArray();

            foreach (var def in m_RegisteredAnnotations.Values)
            {
                defs.Add(OldPerceptionJsonFactory.Convert(this, def.id, def));
            }

            var top = new JObject
            {
                ["version"] = version,
                ["annotation_definitions"] = defs
            };
            var path = Path.Combine(m_DatasetPath, "annotation_definitions.json");
            WriteJTokenToFile(path, top);
        }

        void WriteMetricsDefinitionsFile()
        {
            var defs = new JArray();

            foreach (var def in m_RegisteredMetrics.Values)
            {
                defs.Add(OldPerceptionJsonFactory.Convert(this, def.id, def));
            }

            var top = new JObject
            {
                ["version"] = version,
                ["metric_definitions"] = defs
            };
            var path = Path.Combine(m_DatasetPath, "metric_definitions.json");
            WriteJTokenToFile(path, top);
        }

        void WriteSensorsFile()
        {
            var sub = new JArray();
            foreach (var sensor in m_SensorMap)
            {
                sub.Add(JToken.FromObject(sensor.Value, Serializer));
            }
            var top = new JObject
            {
                ["version"] = version,
                ["sensors"] = sub
            };
            var path = Path.Combine(m_DatasetPath, "sensors.json");
            WriteJTokenToFile(path, top);
        }

        JToken ToJToken(Metric metric)
        {
            string sensorId = null;
            string annotationId = null;
            string defId = null;

            if (!string.IsNullOrEmpty(metric.sensorId))
            {
                sensorId = m_SensorMap[metric.sensorId].id;
            }

            if (!string.IsNullOrEmpty(metric.annotationId))
            {
                annotationId = m_RegisteredAnnotations[metric.annotationId].id;
            }

            if (m_RegisteredMetrics.TryGetValue(metric.Id, out var def))
            {
                defId = def.id;
            }

            return new JObject
            {
                ["capture_id"] =  sensorId,
                ["annotation_id"] = annotationId,
                ["sequence_id"] = metric.sequenceId.ToString(),
                ["step"] = metric.step,
                ["metric_definition"] = defId,
                ["values"] = JToken.FromObject(metric.Values)
            };
        }

        void WriteMetricsFile(int index, IEnumerable<JToken> metrics)
        {
            var top = new MetricsJson
            {
                version = version,
                metrics = metrics
            };

            var path = Path.Combine(m_DatasetPath, $"metrics_{index:000}.json");
            WriteJTokenToFile(path, top);
        }

        int m_MetricOutCount;
        List<JToken> m_MetricsReady = new List<JToken>();
        void AddMetricToReport(Metric metric)
        {
            m_MetricsReady.Add(ToJToken(metric));
            WriteMetrics();
        }


        void WriteCaptureFile(int index, IEnumerable<PerceptionCapture> captures)
        {
            var top = new PerceptionJson
            {
                version = version,
                captures = captures
            };

            var path = Path.Combine(m_DatasetPath, $"captures_{index:000}.json");
            WriteJTokenToFile(path, top);
        }

        // ReSharper disable NotAccessedField.Local
        // ReSharper disable InconsistentNaming
        [Serializable]
        struct PerceptionJson
        {
            // ReSharper disable once MemberHidesStaticFromOuterClass
            public string version;
            public IEnumerable<PerceptionCapture> captures;
        }

        [Serializable]
        struct MetricsJson
        {
            // ReSharper disable once MemberHidesStaticFromOuterClass
            public string version;
            public IEnumerable<JToken> metrics;
        }

        [Serializable]
        struct PerceptionCapture
        {
            public string id;
            public string sequence_id;
            public int step;
            public float timestamp;
            public PerceptionRgbSensor sensor;
            public string filename;
            public string format;
            public JArray annotations;
        }

        static float[][] ToFloatArray(float3x3 inF3)
        {
            return new[]
            {
                new [] { inF3[0][0], inF3[0][1], inF3[0][2] },
                new [] { inF3[1][0], inF3[1][1], inF3[1][2] },
                new [] { inF3[2][0], inF3[2][1], inF3[2][2] }
            };
        }

        [Serializable]
        struct PerceptionRgbSensor
        {
            public string sensor_id;
            public string modality;
            public Vector3 translation;
            public Vector3 rotation;
            public Vector3 velocity;
            public Vector3 acceleration;
            public float[][] camera_intrinsic;
            public string projection;

            public static PerceptionRgbSensor Convert(RgbSensor inRgb)
            {
                return new PerceptionRgbSensor
                {
                    sensor_id = inRgb.Id,
                    modality = inRgb.sensorType,
                    translation = inRgb.position,
                    rotation = inRgb.rotation,
                    velocity = inRgb.velocity,
                    acceleration = inRgb.acceleration,
                    projection = inRgb.projection,
                    camera_intrinsic = ToFloatArray(inRgb.intrinsics)
                };
            }
        }
        // ReSharper enable NotAccessedField.Local
        // ReSharper enable InconsistentNaming
    }
}