#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;
	struct {
		float time;
		float bode;
		float nyquist;
		float rlocus;
	} start;
	struct {
		float time;
		float bode;
		float nyquist;
		float rlocus;
	} len;
	size_t npoles, nzeros;
	ImPlotPoint poles[MAX_ZP], zeros[MAX_ZP];
};

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


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);
}

// Global variable because std::function is trash
int rlocus_curr_row = 0;
ImPlotPoint rlocus_getter(int idx, void *data)
{
	ct::LocusSeries *ls = (ct::LocusSeries *) data;
	return ImPlotPoint(ls->out(rlocus_curr_row, idx).real(),
			ls->out(rlocus_curr_row, idx).imag());
}

ImPlotPoint nyquist_getter(int idx, void *data)
{
	ct::LocusSeries *ls = (ct::LocusSeries *) data;
	return ImPlotPoint(ls->out(0, idx).real(), ls->out(0, idx).imag());
}

ImPlotPoint nyquist_getter_conj(int idx, void *data)
{
	ct::LocusSeries *ls = (ct::LocusSeries *) data;
	return ImPlotPoint(ls->out(0, idx).real(), -ls->out(0, idx).imag());
}

ImPlotPoint bode_getter_ampl(int idx, void *data)
{
	ct::LocusSeries *ls = (ct::LocusSeries *) data;
	return ImPlotPoint(ls->in(idx), std::abs(ls->out(0, idx)));
}

ImPlotPoint bode_getter_phase(int idx, void *data)
{
	ct::LocusSeries *ls = (ct::LocusSeries *) data;
	return ImPlotPoint(ls->in(idx), std::arg(ls->out(0, idx)));
}

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,
		.start = {
			.time = 0.0,
			.bode = 0.001,
			.nyquist = 0.001,
			.rlocus = 0.001,
		},
		.len = {
			.time = 10.,
			.bode = 1000.,
			.nyquist = 1000.,
			.rlocus = 1000.,
		},
		.npoles = 0,
		.nzeros = 0,
	};

	GState prev_gstate = { };

	// Computation state
	CState cstate = {
		.proper = false,
		.time = ct::TimeSeries(gstate.start.time,
				gstate.start.time + gstate.len.time, gstate.n_samples.time),
		.nyquist = ct::LocusSeries(gstate.start.nyquist,
				gstate.start.nyquist + gstate.len.nyquist, gstate.n_samples.nyquist),
		.rlocus = ct::LocusSeries(gstate.start.rlocus,
				gstate.start.rlocus + gstate.len.rlocus, 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::complex gain(gstate.gain_re, gstate.gain_im);
				// Autogain
				if (gstate.autogain)
				{
					gain = 1. / std::abs(tf.dc_gain());
					gstate.gain_re = gain.real();
					gstate.gain_im = gain.imag();
				}

				// tf = ct::cancel_zp(tf);
				ct::TransferFn tf_cl = ct::feedback(tf, -gain);
				cstate.tf_cl = tf_cl;

				if (tf_cl.is_proper())
				{
					// New time series
					cstate.time = ct::TimeSeries(gstate.start.time,
							gstate.start.time + gstate.len.time, gstate.n_samples.time);

					// New root nyquist series
					cstate.nyquist = ct::LocusSeries(gstate.start.nyquist,
							gstate.start.nyquist + gstate.len.nyquist, gstate.n_samples.nyquist);

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

					// Convert to state space
					ct::SSModel ss = ct::ctrb_form(ct::feedback(tf, -gain));

					// Update plots
					ct::step(ss, cstate.time);
					ct::rlocus(tf, cstate.rlocus);
					ct::nyquist(tf, cstate.nyquist);
				}
			}

			// 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, 0.0001, 10000, NULL, ImGuiSliderFlags_Logarithmic);
			ImGui::SliderFloat("Gain (imag)", &gstate.gain_im, 0.0001, 10000, 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.start.time, 0, 100);
			ImGui::SliderFloat("Time Length", &gstate.len.time, 0, 100, NULL, ImGuiSliderFlags_Logarithmic);

			ImGui::InputFloat("Bode Start", &gstate.start.bode, 0, 100);
			ImGui::SliderFloat("Bode Length", &gstate.len.bode, 0, 100, NULL, ImGuiSliderFlags_Logarithmic);

			ImGui::InputFloat("Nyquist Start", &gstate.start.nyquist, 0, 100);
			ImGui::SliderFloat("Nyquist Length", &gstate.len.nyquist, 0.001, 10000, NULL, ImGuiSliderFlags_Logarithmic);

			ImGui::InputFloat("Root Locus Start", &gstate.start.rlocus, 0, 100);
			ImGui::SliderFloat("Root Locus Length", &gstate.len.rlocus, 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::SetupAxisLimits(ImAxis_X1, gstate.start.bode,
				// 		gstate.start.bode + gstate.len.bode, ImPlotCond_Always);
				ImPlot::SetupAxis(ImAxis_X1, "Amplitude |G(iw)| dB", ImPlotAxisFlags_AutoFit);
				ImPlot::SetupAxis(ImAxis_Y1, "Frequency w / (rad / s)", ImPlotAxisFlags_AutoFit);
				ImPlot::PlotLineG("|G(iw)|", bode_getter_ampl, &cstate.nyquist, cstate.nyquist.n_samples);
				ImPlot::EndPlot();
			}

			if (ImPlot::BeginPlot("Bode Phase", ImVec2(-1, 0)))
			{
				ImPlot::SetupAxisScale(ImAxis_X1, ImPlotScale_Log10);
				ImPlot::SetupAxisScale(ImAxis_Y1, ImPlotScale_Linear);
				// ImPlot::SetupAxisLimits(ImAxis_X1, gstate.start.bode,
				// 		gstate.start.bode + gstate.len.bode, ImPlotCond_Always);
				ImPlot::SetupAxis(ImAxis_X1, "Phase arg G(iw) / deg", ImPlotAxisFlags_AutoFit);
				ImPlot::SetupAxis(ImAxis_Y1, "Frequency w / (rad / s)", ImPlotAxisFlags_AutoFit);
				ImPlot::PlotLineG("arg G(iw)", bode_getter_phase, &cstate.nyquist, cstate.nyquist.n_samples);
				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.start.time,
						gstate.start.time + gstate.len.time, 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");
			ImGui::Text("Hold CTRL / SHIFT after dragging to keep it real / imaginary");

			int rlflags = ImPlotFlags_Equal | ImPlotFlags_NoBoxSelect | ImPlotFlags_NoLegend;
			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);
					rlocus_curr_row = i; // global variable!
					ImPlot::PlotLineG("Root Locus", rlocus_getter, &cstate.rlocus, 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 currPoleCol(0,0.5f,0,1);
				const ImVec4 zeroCol(1,0.5f,1,1);
				const ImVec4 currZeroCol(1,0.2f,1,1);
				const float size = 4;
				const int flags = ImPlotDragToolFlags_Delayed;

				auto zeros = cstate.tf_cl.num.roots();
				for (int i = 0; i < zeros.n_elem; i++)
				{
					const double x = zeros(i).real();
					const double y = zeros(i).imag();
					ImGui::PushID(i);
					ImPlot::SetNextMarkerStyle(ImPlotMarker_Diamond, size, currZeroCol, 0);
					ImPlot::PlotScatter("Zeros", &x, &y, 1);
					ImGui::PopID();
				}

				auto roots = cstate.tf_cl.den.roots();
				for (int i = 0; i < roots.n_elem; i++)
				{
					const double x = roots(i).real();
					const double y = roots(i).imag();
					ImGui::PushID(i);
					ImPlot::SetNextMarkerStyle(ImPlotMarker_Diamond, size, currPoleCol, 0);
					ImPlot::PlotScatter("Roots", &x, &y, 1);
					ImGui::PopID();
				}

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

					if (dragging && ImGui::GetIO().KeyCtrl)
						gstate.zeros[i].y = 0;

					if (dragging && ImGui::GetIO().KeyShift)
						gstate.zeros[i].x = 0;
				}

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

					if (dragging && ImGui::GetIO().KeyCtrl)
						gstate.poles[i].y = 0;

					if (dragging && ImGui::GetIO().KeyShift)
						gstate.poles[i].x = 0;
				}

				ImPlot::EndPlot();
			}

			if (ImPlot::BeginPlot("Nyquist Diagram", ImVec2(-1, 0), ImPlotFlags_Equal))
			{
				ImPlot::SetupAxesLimits(-1, 1, -1, 1);
				// ImPlot::SetNextMarkerStyle(ImPlotMarker_Plus);
				ImPlot::PlotLineG("Nyquist Curve", nyquist_getter, &cstate.nyquist, cstate.nyquist.n_samples);
				ImPlot::PlotLineG("Nyquist Curve Conj", nyquist_getter_conj, &cstate.nyquist, cstate.nyquist.n_samples);
				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: