#include "imgui.h"
#include "imgui_impl_glfw.h"
#include "imgui_impl_opengl3.h"
#include "implot.h"
#include "control.h"

#include <stdio.h>
#include <string.h>
#include <memory.h>

#if defined(IMGUI_IMPL_OPENGL_ES2)
#include <GLES2/gl2.h>
#endif
#undef GLFW_INCLUDE_NONE
#include <GLFW/glfw3.h> // Will drag system OpenGL headers

// [Win32] Our example includes a copy of glfw3.lib pre-compiled with VS2010 to
// maximize ease of testing and compatibility with old VS compilers. To link
// with VS2010-era libraries, VS2015+ requires linking with
// legacy_stdio_definitions.lib, which we do using this pragma. Your own
// project should not be affected, as you are likely to link with a newer
// binary of GLFW that is adequate for your version of Visual Studio.
#if defined(_MSC_VER) && (_MSC_VER >= 1900) && !defined(IMGUI_DISABLE_WIN32_FUNCTIONS)
#pragma comment(lib, "legacy_stdio_definitions")
#endif

#define MAX_ZP 64

struct GState
{
	struct {
		int time;
		int bode;
		int nyquist;
		int rlocus;
	} n_samples;
	float gain_re;
	float gain_im;
	bool autogain;
	float time_start;
	float time_len;
	float bode_start;
	float bode_len;
	float rlocus_start;
	float rlocus_len;
	size_t npoles, nzeros;
	ImPlotPoint poles[MAX_ZP], zeros[MAX_ZP];
};

struct CState
{
	bool proper;
	ct::TimeSeries time;
	ct::LocusSeries rlocus;
};


ImPlotPoint time_real_getter(int idx, void *data)
{
	ct::TimeSeries *ts = (ct::TimeSeries *) data;
	return ImPlotPoint(ts->time(idx), ts->out(idx).real());
}

ImPlotPoint time_imag_getter(int idx, void *data)
{
	ct::TimeSeries *ts = (ct::TimeSeries *) data;
	return ImPlotPoint(ts->time(idx), ts->out(idx).imag());
}

ImPlotPoint step_getter(int idx, void *data)
{
	ct::TimeSeries *ts = (ct::TimeSeries *) data;
	return ImPlotPoint(ts->time(idx), idx != 0);
}

ImPlotPoint rlocus_getter(int idx, void *data)
{
	// return ImPlotPoint(ls->out(row, idx).real(), ls->out(row, idx).imag());
}


static void glfw_error_callback(int error, const char* description)
{
	fprintf(stderr, "Glfw Error %d: %s\n", error, description);
}

int main(int, char**)
{
	// Setup window
	glfwSetErrorCallback(glfw_error_callback);
	if (!glfwInit())
		return 1;

	// Decide GL+GLSL versions
#if defined(IMGUI_IMPL_OPENGL_ES2)
	// GL ES 2.0 + GLSL 100
	const char* glsl_version = "#version 100";
	glfwWindowHint(GLFW_CONTEXT_VERSION_MAJOR, 2);
	glfwWindowHint(GLFW_CONTEXT_VERSION_MINOR, 0);
	glfwWindowHint(GLFW_CLIENT_API, GLFW_OPENGL_ES_API);
#elif defined(__APPLE__)
	// GL 3.2 + GLSL 150
	const char* glsl_version = "#version 150";
	glfwWindowHint(GLFW_CONTEXT_VERSION_MAJOR, 3);
	glfwWindowHint(GLFW_CONTEXT_VERSION_MINOR, 2);
	glfwWindowHint(GLFW_OPENGL_PROFILE, GLFW_OPENGL_CORE_PROFILE);	// 3.2+ only
	glfwWindowHint(GLFW_OPENGL_FORWARD_COMPAT, GL_TRUE);			// Required on Mac
#else
	// GL 3.0 + GLSL 130
	const char* glsl_version = "#version 130";
	glfwWindowHint(GLFW_CONTEXT_VERSION_MAJOR, 3);
	glfwWindowHint(GLFW_CONTEXT_VERSION_MINOR, 0);
	//glfwWindowHint(GLFW_OPENGL_PROFILE, GLFW_OPENGL_CORE_PROFILE);  // 3.2+ only
	//glfwWindowHint(GLFW_OPENGL_FORWARD_COMPAT, GL_TRUE);			  // 3.0+ only
#endif

	// Create window with graphics context
	GLFWwindow* window = glfwCreateWindow(1280, 720, "Fast SISO Tool", NULL, NULL);
	if (window == NULL)
		return 1;
	glfwMakeContextCurrent(window);
	glfwSwapInterval(1); // Enable vsync

	// Setup Dear ImGui context
	IMGUI_CHECKVERSION();
	ImGui::CreateContext();
	ImPlot::CreateContext();
	ImGuiIO& io = ImGui::GetIO(); (void)io;
	//io.ConfigFlags |= ImGuiConfigFlags_NavEnableKeyboard;		// Enable Keyboard Controls
	//io.ConfigFlags |= ImGuiConfigFlags_NavEnableGamepad;		// Enable Gamepad Controls

	// Setup Dear ImGui style
	ImGui::StyleColorsDark();

	// Setup Platform/Renderer backends
	ImGui_ImplGlfw_InitForOpenGL(window, true);
	ImGui_ImplOpenGL3_Init(glsl_version);

	// State variables
	bool show_demo_window = true;
	ImVec4 clear_color = ImVec4(0.45f, 0.55f, 0.60f, 1.00f);

	// Graphics state
	GState gstate = {
		// settings
		.n_samples = {
			.time = 500,
			.bode = 500,
			.nyquist = 500,
			.rlocus = 2000,
		},
		.gain_re = 1.,
		.gain_im = 0.,
		.autogain = false,
		.time_start = 0,
		.time_len = 10,
		.bode_start = 0.01,
		.bode_len = 10000,
		.rlocus_start = 0.001,
		.rlocus_len = 1000,
		.npoles = 0,
		.nzeros = 0,
	};

	GState prev_gstate = { };

	// Computation state
	CState cstate = {
		.proper = false,
		.time = ct::TimeSeries(gstate.time_start,
				gstate.time_start + gstate.time_len, gstate.n_samples.time),
		.rlocus = ct::LocusSeries(gstate.rlocus_start,
				gstate.rlocus_start + gstate.rlocus_len, gstate.n_samples.rlocus),
	};

	// Main loop
	while (!glfwWindowShouldClose(window))
	{
		// Poll and handle events (inputs, window resize, etc.)
		// You can read the io.WantCaptureMouse, io.WantCaptureKeyboard flags
		// to tell if dear imgui wants to use your inputs.
		//
		// - When io.WantCaptureMouse is true, do not dispatch mouse input data
		//	 to your main application.
		//
		// - When io.WantCaptureKeyboard is true, do not dispatch keyboard
		//	 input data to your main application.
		//
		// Generally you may always pass all inputs to dear imgui, and hide
		// them from your application based on those two flags.
		glfwPollEvents();

		// Recompute simulation if necessary
		if (memcmp(&gstate, &prev_gstate, sizeof(GState)))
		{
			if (gstate.npoles || gstate.nzeros)
			{
				// Create new transfer function
				ct::TransferFn tf;

				for (int k = 0; k < gstate.npoles; k++)
					tf.add_pole(ct::complex(gstate.poles[k].x, gstate.poles[k].y));

				for (int k = 0; k < gstate.nzeros; k++)
					tf.add_zero(ct::complex(gstate.zeros[k].x, gstate.zeros[k].y));

				ct::TransferFn tf_cl = ct::feedback(tf);
				// Autogain
				if (gstate.autogain)
				{
					ct::complex gain = 1. / tf_cl.dc_gain();
					tf_cl = gain * tf_cl;
					gstate.gain_re = gain.real();
					gstate.gain_im = gain.imag();
				}
				else
				{
					ct::complex gain(gstate.gain_re, gstate.gain_im);
					tf_cl = gain * tf_cl;
				}

				if (tf_cl.is_proper())
				{
					// Convert to state space
					ct::SSModel ss = ct::ctrb_form(ct::feedback(tf));

					// New time series
					cstate.time = ct::TimeSeries(gstate.time_start,
							gstate.time_start + gstate.time_len, gstate.n_samples.time);

					// Update time domain simulation
					ct::step(ss, cstate.time);

					// New root locus series
					cstate.rlocus = ct::LocusSeries(gstate.rlocus_start,
							gstate.rlocus_start + gstate.rlocus_len, gstate.n_samples.rlocus);

					ct::rlocus(tf, cstate.rlocus);
				}
			}

			// Update graphic state
			memcpy(&prev_gstate, &gstate, sizeof(GState));
		}

		// Start the Dear ImGui frame
		ImGui_ImplOpenGL3_NewFrame();
		ImGui_ImplGlfw_NewFrame();
		ImGui::NewFrame();

		if (show_demo_window)
		{
			ImGui::ShowDemoWindow(&show_demo_window);
			ImPlot::ShowDemoWindow(&show_demo_window);
		}

		if (ImGui::Begin("Settings"))
		{
			static bool zero_index = 0, pole_index = 0;
			static char buf[32];

			ImGui::SeparatorText("Plant G(s)");

			ImGui::Checkbox("Automatic", &gstate.autogain);
			ImGui::SliderFloat("Gain (real)", &gstate.gain_re, -1000, 1000, NULL, ImGuiSliderFlags_Logarithmic);
			ImGui::SliderFloat("Gain (imag)", &gstate.gain_im, -1000, 1000, NULL, ImGuiSliderFlags_Logarithmic);

			ImGui::PushItemWidth(140);
			if (ImGui::BeginListBox("Poles"))
			{
				for (int i = 0; i < gstate.npoles; i++)
				{
					ImGui::PushID(i);
					snprintf(buf, sizeof(buf), "(%.2f, %.2fi)", gstate.poles[i].x, gstate.poles[i].y);
					if (ImGui::Selectable(buf))
					{
						gstate.poles[i].y = 0;
					}

					ImGui::PopID();
				}

				ImGui::EndListBox();
			}
			ImGui::SameLine();
			if (ImGui::BeginListBox("Zeros"))
			{
				for (int i = 0; i < gstate.nzeros; i++)
				{
					ImGui::PushID(i);
					snprintf(buf, sizeof(buf), "(%.2f, %.2fi)", gstate.zeros[i].x, gstate.zeros[i].y);
					if (ImGui::Selectable(buf))
					{
						gstate.zeros[i].y = 0;
					}
					ImGui::PopID();
				}

				ImGui::EndListBox();
			}
			ImGui::PopItemWidth();

			ImGui::SeparatorText("Simulation Ranges");
			ImGui::InputFloat("Time Start", &gstate.time_start, 0, 100);
			ImGui::SliderFloat("Time Length", &gstate.time_len, 0, 100, NULL, ImGuiSliderFlags_Logarithmic);
			ImGui::InputFloat("Bode Start", &gstate.bode_start, 0, 100);
			ImGui::SliderFloat("Bode Length", &gstate.bode_len, 0, 100, NULL, ImGuiSliderFlags_Logarithmic);
			ImGui::InputFloat("Root Locus Start", &gstate.rlocus_start, 0, 100);
			ImGui::SliderFloat("Root Locus Length", &gstate.rlocus_len, 0.001, 10000, NULL, ImGuiSliderFlags_Logarithmic);

			ImGui::SeparatorText("Number of Samples");
			ImGui::SliderInt("Time", &gstate.n_samples.time, 50, 5000);
			ImGui::SliderInt("Bode", &gstate.n_samples.bode, 50, 5000);
			ImGui::SliderInt("Nyquist", &gstate.n_samples.nyquist, 50, 5000);
			ImGui::SliderInt("Root Locus", &gstate.n_samples.rlocus, 50, 5000);
		}
		ImGui::End();

		if (ImGui::Begin("Bode Diagrams"))
		{
			if (ImPlot::BeginPlot("Bode Magnitude", ImVec2(-1, 0)))
			{
				ImPlot::SetupAxisScale(ImAxis_X1, ImPlotScale_Log10);
				ImPlot::SetupAxisScale(ImAxis_Y1, ImPlotScale_Log10);
				ImPlot::SetupAxesLimits(0.01, 10000, 0.001, 10);
				ImPlot::SetupAxis(ImAxis_X1, "Amplitude |G(iw)| dB");
				ImPlot::SetupAxis(ImAxis_Y1, "Frequency w / (rad / s)");
				ImPlot::EndPlot();
			}

			if (ImPlot::BeginPlot("Bode Phase", ImVec2(-1, 0)))
			{
				ImPlot::SetupAxisScale(ImAxis_X1, ImPlotScale_Linear);
				ImPlot::SetupAxesLimits(0.1, 100, -180, 180);
				ImPlot::SetupAxis(ImAxis_X1, "Phase arg G(iw) / deg");
				ImPlot::SetupAxis(ImAxis_Y1, "Frequency w / (rad / s)");
				ImPlot::EndPlot();
			}
		}
		ImGui::End();

		if (ImGui::Begin("Time Domain"))
		{
			if (ImPlot::BeginPlot("Step Response", ImVec2(-1, 0)))
			{
				ImPlot::SetupAxisScale(ImAxis_X1, ImPlotScale_Linear);
				ImPlot::SetupAxis(ImAxis_X1, "Time t / s");
				ImPlot::SetupAxisLimits(ImAxis_X1, gstate.time_start,
						gstate.time_start + gstate.time_len, ImPlotCond_Always);
				ImPlot::SetupAxis(ImAxis_Y1, "Response y(t)", ImPlotAxisFlags_AutoFit);

				ImPlot::PlotLineG("Step", step_getter, &cstate.time, cstate.time.n_samples);
				ImPlot::PlotLineG("G(s), real", time_real_getter, &cstate.time, cstate.time.n_samples);
				ImPlot::PlotLineG("G(s), imag", time_imag_getter, &cstate.time, cstate.time.n_samples);
				ImPlot::EndPlot();
			}
		}
		ImGui::End();

		if (ImGui::Begin("Stability"))
		{
			ImGui::Text("Use CTRL + R / L Click to add a pole / zero");

			int rlflags = ImPlotFlags_Equal | ImPlotFlags_NoBoxSelect;
			if (ImGui::GetIO().KeyCtrl)
				rlflags |= ImPlotFlags_Crosshairs;

			if (ImPlot::BeginPlot("Root Locus", ImVec2(-1, 0), rlflags))
			{
				ImPlot::SetupAxesLimits(-10, 10, -10, 10);
				
				// Plot curren root locus
				for (int i = 0; i < cstate.rlocus.out.n_rows; i++)
				{
					ImGui::PushID(i);
					// ImPlot::PlotLineG("Root Locus", rlocus_getter, &cstate.rlocus.out.row(i), cstate.rlocus.n_samples);
					ImGui::PopID();
				}

				if (ImPlot::IsPlotHovered() && ImGui::GetIO().KeyCtrl)
				{
					ImPlotPoint pt = ImPlot::GetPlotMousePos();
					// Add a pole
					if (ImGui::IsMouseClicked(ImGuiMouseButton_Left))
					{
						if (gstate.npoles < MAX_ZP)
						{
							gstate.poles[gstate.npoles].x = pt.x;
							gstate.poles[gstate.npoles].y = pt.y;
							gstate.npoles += 1;
						}
					}
					// Add a zero
					else if (ImGui::IsMouseClicked(ImGuiMouseButton_Right))
					{
						if (gstate.npoles < MAX_ZP)
						{
							gstate.zeros[gstate.nzeros].x = pt.x;
							gstate.zeros[gstate.nzeros].y = pt.y;
							gstate.nzeros += 1;
						}
					}
				}

				// Plot zeros and poles
				const ImVec4 poleCol(0,0.9f,0,1);
				const ImVec4 zeroCol(1,0.5f,1,1);
				const float size = 4;
				const int flags = ImPlotDragToolFlags_Delayed;

				for (int i = 0; i < gstate.nzeros; i++)
					ImPlot::DragPoint(i, &gstate.zeros[i].x, &gstate.zeros[i].y, zeroCol, size, flags);

				for (int i = 0; i < gstate.npoles; i++)
					ImPlot::DragPoint(gstate.nzeros + i, &gstate.poles[i].x, &gstate.poles[i].y,
							poleCol, size, flags);

				ImPlot::EndPlot();
			}

			if (ImPlot::BeginPlot("Nyquist Diagram", ImVec2(-1, 0), ImPlotFlags_Equal))
			{
				ImPlot::SetupAxesLimits(-1, 1, -1, 1);
				ImPlot::SetNextMarkerStyle(ImPlotMarker_Cross);
				ImPlot::EndPlot();
			}
		}
		ImGui::End();

		// Rendering
		ImGui::Render();
		int display_w, display_h;
		glfwGetFramebufferSize(window, &display_w, &display_h);
		glViewport(0, 0, display_w, display_h);
		glClearColor(clear_color.x * clear_color.w, clear_color.y * clear_color.w,
				clear_color.z * clear_color.w, clear_color.w);
		glClear(GL_COLOR_BUFFER_BIT);
		ImGui_ImplOpenGL3_RenderDrawData(ImGui::GetDrawData());

		glfwSwapBuffers(window);
	}

	// Cleanup
	ImGui_ImplOpenGL3_Shutdown();
	ImGui_ImplGlfw_Shutdown();
	ImPlot::DestroyContext();
	ImGui::DestroyContext();

	glfwDestroyWindow(window);
	glfwTerminate();

	return 0;
}
// vim:ts=2 sw=2 noet: