﻿using System;
using System.Collections.Generic;
using System.Collections;
using System.Linq;
using System.Text;
using System.Threading;
using System.Threading.Tasks;

using AsyncIO;
using NetMQ;
using NetMQ.Sockets;

namespace PythonPlot
{
  /// <summary>
  /// Class that defines a triple of RGB colors
  /// </summary>
  public class Color
  {
    public ushort r, g, b;
    public Color(ushort ra, ushort ga, ushort ba)
    {
      r = ra;
      g = ga;
      b = ba;
    }

    public static Color Red() { return new Color(255, 0, 0); }
    public static Color Yellow() { return new Color(255, 255, 0); }
    public static Color Blue() { return new Color(0, 0, 255); }
    public static Color Green() { return new Color(0, 255, 0); }
  }

  /// <summary>
  /// Class that encapsulates a Curve to which data can be added.
  /// </summary>
  public class Curve
  {
    string m_name;
    PlotWindow m_plotWindow;

    public string Name
    {
      get
      {
        return m_name;
      }
    }

    /// <summary>
    /// Constructor to create a plot Curve
    /// </summary>
    /// <param name="name">Name of the curve.</param>
    /// <param name="plot">Reference to the plot window in which the curve is going to be plotted.</param>
    public Curve(string name, PlotWindow plotWindow)
    {
      m_name = name;
      m_plotWindow = plotWindow;
    }

    /// <summary>
    /// Add data to the curve
    /// </summary>
    /// <param name="x">X value</param>
    /// <param name="y">Y value</param>
    /// <returns>True if the data is ready to be received. If this method returns false, it could be because: 
    /// a) Something is wrong with the configuration 
    /// b) We are adding data faster than the specified plot frequency, the data will then be ignored.</returns>
    public bool AddData(double x, double y)
    {
      // Is this a valid curve?
      if (m_plotWindow == null)
        return false;

      // This could return false if we are adding data faster than we should
      // It is not really a problem, the data is just discarded
      return m_plotWindow.Client._AddData(m_plotWindow.Title, m_name, x, y);
    }
  }

  /// <summary>
  /// Class that encapsulates a Plot window.
  /// Acts like a simplified interface to the Plot client
  /// </summary>
  public class PlotWindow
  {
    string m_title;
    PlotClient m_client;

    Dictionary<string, Curve> m_curves = new Dictionary<string, Curve>();

    public string Title
    {
      get
      {
        return m_title;
      }
    }

    public PlotClient Client
    {
      get
      {
        return m_client;
      }
    }

    public bool Contains(string curve)
    {
      return m_curves.ContainsKey(curve);
    }

    /// <summary>
    /// Constructor for a plot.
    /// </summary>
    /// <param name="title">Name of this plot window.</param>
    /// <param name="client">Reference to the plot client.</param>
    public PlotWindow(string title, PlotClient client)
    {
      m_title = title;
      m_client = client;
    }

    /// <summary>
    /// Add a new curve into the plot
    /// </summary>
    /// <param name="curve">Name of the curve</param>
    /// <param name="color">Color of the curve</param>
    /// <param name="style">Pen style of the curve</param>
    /// <param name="width">Pen width of the curve</param>
    /// <returns>Return true if curve can be added to an already existing Plot</returns>
    public Curve AddCurve(string curve, Color color, PlotData.Style style = PlotData.Style.SolidLine, ushort width = 1)
    {
      // Test and see if we can add this curve
      bool status = Client._AddCurve(m_title, curve, color, style, width);

      // Something went wrong
      if (!status)
        return null;

      // Create a curve that we can now use to add data.
      Curve curveInstance = new Curve(curve, this);

      m_curves.Add(curve, curveInstance);

      return curveInstance;
    }
  }

  /// <summary>
  /// This class defines the network protocol between the PlotClient and the PlotServer (in Python)
  /// </summary>
  public class PlotData
  {
    public enum Style
    {
      SolidLine,
      DashLine,
      DotLine
    };

    public string command = ""; // reset, addPlot, addCurve, addData, stop
    public string title = "Force"; // Title of the 
    public string curve = "curve";
    public ushort[] color = { 255, 0, 0 };
    public ushort width = 1;
    public Style style = Style.SolidLine;
    public string xLabel = "Time";
    public string yLabel = "Value";
    public string xUnit = "s";
    public string yUnit = "kN";
    public double x_data = 0;
    public double y_data = 0;
  };


  /// <summary>
  /// Interface for remote plotting to a python server
  /// </summary>
  public class PlotClient : RunAbleThread
  {
    private Mutex m_mutex = new Mutex();

    List<PlotData> m_plotData = new List<PlotData>();

    private int m_port = 0;
    private string m_ipaddress;

    private bool m_enabled = true;

    private System.Diagnostics.Stopwatch m_stopWatch = new System.Diagnostics.Stopwatch();
    private bool m_useFrequency = false;
    private double m_frequency = -1;

    private Dictionary<string, double> m_timingTable = new Dictionary<string, double>();

    private Dictionary<string, PlotWindow> m_plotWindows = new Dictionary<string, PlotWindow>();

    /// <summary>
    /// Set/Get the frequency for transmitting with the AddData method
    /// Data sent with higher frequency 
    /// </summary>
    public double Frequency
    {
      get
      {
        return m_frequency;
      }

      set
      {
        m_frequency = value;
        m_useFrequency = value > 0;
      }
    }

    /// <summary>
    /// Returns true if plotting is enabled. If you set it to false, no data will be transmitted to the plot server
    /// </summary>
    public bool Enable
    {
      get
      {
        return m_enabled;
      }

      set
      {
        m_enabled = value;
      }
    }

    /// <summary>
    /// Returns true if a named data field should be sent (via AddData) to the plot server.
    /// </summary>
    /// <param name="dataField">Name of the plot data that should be sent over to the server</param>
    /// <returns>true if a data field should be sent (via AddData) to the plot server. False if we are adding data faster than the specified frequency.</returns>
    private bool TimeToAddData(string dataField)
    {
      if (!m_useFrequency)
        return true;

      double now = m_stopWatch.Elapsed.TotalSeconds;
      double lastTime = 0;

      // When was this data field added last?
      // We do not want to send data faster than the specified plot frequency
      bool hasValue = m_timingTable.TryGetValue(dataField, out lastTime);
      if (!hasValue)
        m_timingTable.Add(dataField, 0);

      // Is it time to send this data?
      double elapsed = now - lastTime;

      bool timeToDoIt = elapsed > 1 / m_frequency;
      if (!hasValue || timeToDoIt)
        m_timingTable[dataField] = now;

      return timeToDoIt;
    }

    /// <summary>
    /// Constructor for PlotClient
    /// </summary>
    /// <param name="port">Specifies the network port used for sending plot data</param>
    /// <param name="ipaddress">Specifies the ip adress of the remote plot server</param>
    public PlotClient(int port = 5555, string ipaddress = "localhost")
    {
      m_port = port;
      m_ipaddress = ipaddress;
      Frequency = 0; // By default we are not filtering based on the frequency of calls to AddData
      ForceDotNet.Force(); // this line is needed to prevent unity freeze after one use
    }

    /// <summary>
    /// Reset the plot server, removes all previous data, ready for new plots.
    /// </summary>
    public void Reset()
    {
      m_timingTable.Clear();
      m_plotWindows.Clear();

      if (!m_enabled)
        return;

      m_mutex.WaitOne();

      PlotData data = new PlotData();
      data.command = "reset";

      m_plotData.Add(data);
      m_stopWatch.Restart();

      m_mutex.ReleaseMutex();
    }

    /// <summary>
    /// Tell server to save current plot as a .png file
    /// </summary>
    /// <param name="title">Filename of saved image</param>
    public void SavePlot(string title)
    {
      if (!m_enabled)
        return;

      m_mutex.WaitOne();

      PlotData plotData = new PlotData();
      plotData.command = "save";
      plotData.title = title;

      m_plotData.Add(plotData);

      m_mutex.ReleaseMutex();
    }

    /// <summary>
    /// Add a new Plot window
    /// </summary>
    /// <param name="title">This title identifies this plot window</param>
    /// <param name="YUnit">Unit of the Y axis for the plot curves in this window</param>
    /// <param name="XUnit">Unit of the X axis (default seconds) for the plot curves in this window</param>
    public PlotWindow AddPlot(string title, string xLabel, string xUnit, string yLabel, string yUnit)
    {
      if (!m_enabled)
        return null;

      if (m_plotWindows.ContainsKey(title)) {
        Log(string.Format("Trying to add an already existing plot window : {0}", title));
        return null;
      }

      var plotWindow = new PlotWindow(title, this);

      m_mutex.WaitOne();

      // Does not matter if the plot already exists, then we will just reuse the previous one
      m_plotWindows.Add(title, plotWindow);

      PlotData plotData = new PlotData();
      plotData.command = "addPlot";
      plotData.xLabel = xLabel;
      plotData.yLabel = yLabel;
      plotData.xUnit = xUnit;
      plotData.yUnit = yUnit;
      plotData.title = title;

      m_plotData.Add(plotData);

      m_mutex.ReleaseMutex();

      return plotWindow;
    }

    /// <summary>
    /// Internal method. Use PlotWindow.AddCurve()
    /// Add a new curve into the named plot window
    /// </summary>
    /// <param name="title">Title of the plot window previously added with AddPlot</param>
    /// <param name="curve">Name of the curve</param>
    /// <param name="color">Color of the curve</param>
    /// <param name="style">Pen style of the curve</param>
    /// <param name="width">Pen width of the curve</param>
    /// <returns>Return true if curve can be added to an already existing Plot</returns>
    public bool _AddCurve(string title, string curve, Color color, PlotData.Style style = PlotData.Style.SolidLine, ushort width = 1)
    {
      if (!m_enabled)
        return false;

      // Does the plot window exist?
      PlotWindow plotWindow = null;
      if (!m_plotWindows.TryGetValue(title, out plotWindow)) {
        Log(string.Format("Trying to add a curve to non existing plot window {0}. Did you forget to call AddPlot()?", title));
        return false;
      }

      // Does the 
      if (plotWindow.Contains(curve)) {
        Log(string.Format("Trying to add an already existing curve: {0}.", curve));
        return false;
      }

      // Prepare the data to be sent.
      PlotData plotData = new PlotData();
      plotData.command = "addCurve";
      plotData.title = title;
      plotData.style = style;
      plotData.curve = curve;
      plotData.width = width;
      plotData.color = new ushort[] { color.r, color.g, color.b };

      // Wait so we get write access and can add data that will be sent over to the server.
      m_mutex.WaitOne();
      m_plotData.Add(plotData);
      m_mutex.ReleaseMutex();

      return true;
    }

    /// <summary>
    /// Internal method. Use Curve.AddData()
    /// Add data (x,y) for a named curve in a named plot window (title)
    /// </summary>
    /// <param name="title">Identifies the Plot previously added using AddPlot</param>
    /// <param name="curve">Identifies the Plot Curve previously added using AddCurve</param>
    /// <param name="x">Value for the x axis</param>
    /// <param name="y">Value for the y axis</param>
    /// <returns>Return true if the data could be added to an existing plot/curve</returns>
    public bool _AddData(string title, string curve, double x, double y)
    {
      if (!m_enabled)
        return false;

      string dataField = title + curve;

      PlotWindow plotWindow = null;
      if (!m_plotWindows.TryGetValue(title, out plotWindow)) 
      {
        Log(string.Format("Trying to add data to non existing plot: {0}. Did you forget to call AddPlot()?", title));
        return false;
      }

      if (!plotWindow.Contains(curve)) {
        Log(string.Format("Trying to add data to non existing curve: {0}. Did you forget to call PlotWindow.AddCurve()?", curve));
        return false;
      }

      // Is it time to send this data? (based on the plot frequency).
      // If not, we will just ignore it and return false.
      if (!TimeToAddData(dataField))
        return false;

      // Yes it is time to send it. Wait for write access.
      m_mutex.WaitOne();

      PlotData data = new PlotData();
      data.command = "addData";
      data.title = title;
      data.curve = curve;
      data.x_data = x;
      data.y_data = y;

      m_plotData.Add(data);
      m_mutex.ReleaseMutex();

      return true;
    }

    /// <summary>
    /// Log a warning to the "host" system.
    /// </summary>
    /// <param name="msg"></param>
    static void Log(string msg)
    {
#if UNITY_2018_1_OR_NEWER
      UnityEngine.Debug.LogWarning(msg);
#else
      Console.WriteLine(msg);
#endif
    }

    /// <summary>
    /// Entry for the thread that runs in the background and send data to the plot server
    /// </summary>
    protected override void Run()
    {
      using (RequestSocket client = new RequestSocket()) {
        // Connect to the remote server
        var address = string.Format("tcp://{0}:{1}", m_ipaddress, m_port);
        try {
          client.Connect(address);
        }
        catch {
          Log("Error connecting to " + address);
          return;
        }

        // Should we continue?
        while (Running) {
          // Wait until there is data in the command queue
          m_mutex.WaitOne();
          bool hasData = m_plotData.Count > 0;

          if (hasData) {
            string jsonString;

            PlotData data = m_plotData.First();

            // Prepare data to send to the client.
            jsonString = UnityEngine.JsonUtility.ToJson(data);
            m_plotData.RemoveAt(0);
            m_mutex.ReleaseMutex();

            try {
              client.SendFrame(jsonString);
            }
            catch {
              Log("Error calling SendFrame");
            }

            string message = null;
            bool gotMessage = false;
            while (Running) {
              TimeSpan requestTimeout = TimeSpan.FromMilliseconds(1000);

              try {
                gotMessage = client.TryReceiveFrameString(requestTimeout, out message); // this returns true if it's successful              
              }
              catch {
                Log(string.Format("Error contacting plot server on address {0}. Make sure it is started!", address));
              }

              if (gotMessage)
                break;
              else {
              }
            }
          }
          else
            m_mutex.ReleaseMutex();
        }
      }

      NetMQConfig.Cleanup(); // this line is needed to prevent unity freeze after one use
    }
  }
}
