#include <stdio.h>

#ifdef __EMSCRIPTEN__
#include <emscripten.h>
#endif

#define GLFW_INCLUDE_ES3
#include <GLES3/gl3.h>
#include <GLFW/glfw3.h>

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

#ifndef IMPLOT_DISABLE_OBSOLETE_FUNCTIONS
#define IMPLOT_DISABLE_OBSOLETE_FUNCTIONS
#endif
#include "implot.h"

#include <iostream>
#include <vector>


//======= Woods-Saxon Library
#include "RK4.h"
#include "WS.h"

const int nPt = 300;
static float xValues[nPt]; 
static float WSCValues[nPt]; 
static float WSSOValues[nPt]; 
static int selected = -1;
static float Energy[nPt];
static vector<double> wfr;
static vector<vector<double>> wf;
static vector<double> energies;
static vector<string> orbString;


GLFWwindow* g_window;
ImVec4 clear_color = ImVec4(0.45f, 0.55f, 0.60f, 1.00f);
int g_width;
int g_height;

// Function used by c++ to get the size of the html canvas
EM_JS(int, canvas_get_width, (), {
  return Module.canvas.width;
});

// Function used by c++ to get the size of the html canvas
EM_JS(int, canvas_get_height, (), {
  return Module.canvas.height;
});

// Function called by javascript
EM_JS(void, resizeCanvas, (), {
  js_resizeCanvas();
});

void on_size_changed(){
  glfwSetWindowSize(g_window, g_width, g_height);

  ImGui::SetCurrentContext(ImGui::GetCurrentContext());
}

void loop(){

  int width = canvas_get_width();
  int height = canvas_get_height();

  if (width != g_width || height != g_height){
    g_width = width;
    g_height = height;
    on_size_changed();
  }

  glfwPollEvents();

  ImGui_ImplOpenGL3_NewFrame();
  ImGui_ImplGlfw_NewFrame();
  ImGui::NewFrame();

  // Debug Window
  {
    ImGui::SetNextWindowSize(ImVec2(400, 100), ImGuiCond_FirstUseEver);
    ImGui::SetNextWindowPos(ImVec2(1000, 100), ImGuiCond_FirstUseEver);
    ImGui::Begin("Debug");
    ImGui::Text("Hello, world!");                           // Display some text (you can use a format string too)
    ImGui::ColorEdit3("bg color", (float*)&clear_color); // Edit 3 floats representing a color

    ImGui::Text("Application average %.3f ms/frame (%.1f FPS)", 1000.0f / ImGui::GetIO().Framerate, ImGui::GetIO().Framerate);
    ImGui::End();
  }

  // // // Demo Window
  // if (true){
  //   ImGui::SetNextWindowPos(ImVec2(650, 20), ImGuiCond_FirstUseEver); // Normally user code doesn't need/want to call this because positions are saved in .ini file anyway. Here we just want to make the demo initial state a bit more friendly!
  //   ImGui::ShowDemoWindow();
  // }

  //  Woods-Saxon window
  {
    ImGui::SetNextWindowSize(ImVec2(1000, 1000), ImGuiCond_FirstUseEver);
    static float V0 = -45, R0 = 3.5, a0 = 0.6;
    static float VSO = 28, RSO = 3.5, aSO = 0.6;
    static int Z = 0, nStep = 300;
    static float Rc = 3.5, dr = 0.1;


    ImGui::Begin("Woods-Saxon Calculation");

    ImGui::SliderFloat("V0 [MeV]", &V0, -100, 0);
    ImGui::SliderFloat("R0 [fm]", &R0, 1, 10);
    ImGui::SliderFloat("a0 [fm]", &a0, 0.1, 2);
    ImGui::SliderFloat("VSO [MeV]", &VSO, 0, 40);
    ImGui::SliderFloat("RSO [fm]", &RSO, 1, 10);
    ImGui::SliderFloat("aSO [fm]", &aSO, 0.1, 2);
    ImGui::SliderInt("Z ", &Z, 0, 100);
    ImGui::SliderFloat("Rc [fm] ", &Rc, 1, 10);
    ImGui::SliderInt("nStep", &nStep, 100, 400);
    ImGui::SliderFloat("dr [fm]", &dr, 0.01, 1);


    ImGui::PushStyleVar(ImGuiStyleVar_FramePadding, ImVec2(30, 10));
    if( ImGui::Button("Cal WS Level") ){
      WoodsSaxon ws;

      if( Z == 0 ) {
        ws.SetNucleus(1,1);
        ws.IsNeutron();
      }else{
        ws.SetNucleus(1, Z);
        ws.IsProton();
        ws.SetRc(Rc);
      }

      ws.SetV0( V0 );
      ws.SetR0( R0 );
      ws.Seta0( a0 );
  
      ws.SetVSO( VSO );
      ws.SetRSO( RSO );
      ws.SetaSO( aSO );

      ws.SetRange2(0.0001, dr, nStep);

      ws.CalWSEnergies(false, 7, 100, 1e-7, 50, 0.2, false);

      float dx = nStep * dr / nPt;
      wfr.clear();
      for( int i = 0; i < nPt; i ++){
        float r = i * dx;
        xValues[i] = r;
        WSCValues[i] = V0/(1 + exp(( r - R0)/a0));
        WSSOValues[i] = VSO * exp((r-RSO)/aSO) / pow(1+exp((r-RSO)/aSO), 2) / aSO/ r ;
        wfr.push_back(dr*i);
      }

      wf.clear();
      for( int i = 0; i < ws.orbString.size(); i++){
        wf.push_back(ws.CalWaveFunction(i, abs(V0)/2, ws.energy[i]));
      }

      selected = -1;

      energies.clear();
      orbString.clear();

      energies = ws.energy;
      orbString =  ws.orbString;
      
      //ws.PrintEnergyLevels();
    }
    ImGui::PopStyleVar();

    if( ImGui::TreeNode("WS Energy Levels:") ){
      for( int i = 0; i < energies.size(); i++){

        char buf[32];
        sprintf(buf, "%24.12f MeV %s", energies[i],orbString[i].c_str());
        if( ImGui::Selectable(buf, selected == i) ) selected = i;
        //ImGui::Text("%12.5f MeV %s", ws.energy[i], ws.orbString[i].c_str());
      }
      ImGui::TreePop();
    }

    if( ImPlot::BeginPlot("Plot") ){
      ImPlot::SetupLegend(ImPlotLocation_SouthEast);
      ImPlot::SetupAxes("r [fm]"," Energy [MeV]");
      ImPlot::SetupAxesLimits(0, 10, floor(V0*1.1), 1);
      ImPlot::SetupAxisLimitsConstraints(ImAxis_X1, 0, 30);
      ImPlot::PlotLine("Central", xValues, WSCValues, nPt);
      ImPlot::PlotLine("S-O", xValues, WSSOValues, nPt);

      if( selected >= 0 ){
        for( int i = 0; i < nPt; i ++) Energy[i] = energies[selected];

        ImPlot::PlotLine(orbString[selected].c_str(), xValues, Energy, nPt);
        const double * haha = wf[selected].data();
        const double * kaka = wfr.data();
        ImPlot::PlotLine("WF", kaka, haha, wfr.size());
      }

      ImPlot::EndPlot();
    }

    ImGui::End();
  }

  ImGui::Render();

  int display_w, display_h;
  glfwMakeContextCurrent(g_window);
  glfwGetFramebufferSize(g_window, &display_w, &display_h);
  glViewport(0, 0, display_w, display_h);
  glClearColor(clear_color.x, clear_color.y, clear_color.z, clear_color.w);
  glClear(GL_COLOR_BUFFER_BIT);

  ImGui_ImplOpenGL3_RenderDrawData(ImGui::GetDrawData());
  glfwMakeContextCurrent(g_window);
}

//^ #########################################################
int init_gl(){
  if( !glfwInit() ){
      fprintf( stderr, "Failed to initialize GLFW\n" );
      return 1;
  }

  glfwWindowHint(GLFW_OPENGL_PROFILE, GLFW_OPENGL_CORE_PROFILE); // We don't want the old OpenGL

  // Open a window and create its OpenGL context
  int canvasWidth = g_width;
  int canvasHeight = g_height;
  g_window = glfwCreateWindow(canvasWidth, canvasHeight, "WebGui Demo", NULL, NULL);
  if( g_window == NULL ){
      fprintf( stderr, "Failed to open GLFW window.\n" );
      glfwTerminate();
      return -1;
  }
  glfwMakeContextCurrent(g_window); // Initialize GLEW

  return 0;
}


int init_imgui(){

  // Setup Dear ImGui binding
  IMGUI_CHECKVERSION();
  ImGui::CreateContext();
  ImGui_ImplGlfw_InitForOpenGL(g_window, true);
  ImGui_ImplOpenGL3_Init();

  ImPlot::CreateContext();

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

  ImGuiIO& io = ImGui::GetIO();

  resizeCanvas();

  for( int i = 0; i < nPt; i++){
    xValues[i] = 0; 
    WSCValues[i] = 0; 
  }

  return 0;
}


int init(){
  init_gl();
  init_imgui();
  return 0;
}

void quit(){
  glfwTerminate();
}

extern "C" int main(int argc, char** argv){

  g_width = canvas_get_width();
  g_height = canvas_get_height();
  if (init() != 0) return 1;

  #ifdef __EMSCRIPTEN__
  emscripten_set_main_loop(loop, 0, 1);
  #endif

  quit();

  return 0;

}