#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;


static vector<string> orbitalStr;
static vector<double> AList;
static vector<vector<double>> enLevelList;

GLFWwindow* g_window;
ImVec4 clear_color = ImVec4(0.427, 0.702f, 0.243f, 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();
  }

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

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

    static int e = 0;
    ImGui::RadioButton("Arbitary", &e, 0); ImGui::SameLine();
    ImGui::RadioButton("Isotope", &e, 1);
    ImGui::Separator();

    ImGui::PushItemWidth(100);
    static int nucleon = 0;
    ImGui::RadioButton("Neutron", &nucleon, 0); ImGui::SameLine();
    ImGui::RadioButton("Proton", &nucleon, 1);
    
    ImGui::BeginDisabled(e==0); ImGui::DragInt("A ", &A, 1, 1, 200); ImGui::EndDisabled();
    ImGui::SameLine();
    ImGui::DragInt("Z", &Z, 1, 0, 100);

    ImGui::BeginDisabled(e==1); ImGui::DragFloat("Rc [fm] ", &Rc, 0.01, 1, 10); ImGui::EndDisabled();
    ImGui::SameLine();
    ImGui::BeginDisabled(e==0); 
    ImGui::DragFloat("rc [fm] ", &rc, 0.01, 1, 10); 
    ImGui::EndDisabled();

    if( e == 0 ) {A = 1; rc = Rc / pow(A, 1./3);}
    if( e == 1 ) {
      rc = 1.25;
      r0 = 1.25;
      rSO = 1.25;
      Rc = rc * pow(A, 1./3);
    }

    ImGui::Text("Eff. mass : %.3f ", mass);

    ImGui::Separator();

    ImGui::DragFloat("V0 [MeV]", &V0, 0.1, -100, 0); 

    ImGui::BeginDisabled(e==1); ImGui::DragFloat("R0 [fm] ", &R0, 0.01, 1, 10); ImGui::EndDisabled();
    ImGui::SameLine();
    ImGui::BeginDisabled(e==0); ImGui::DragFloat("r0 [fm] ", &r0, 0.01, 1, 10); ImGui::EndDisabled();
    if( e == 0 ) {A = 1; r0 = R0 / pow(A, 1./3);}
    if( e == 1 ) R0 = r0 * pow(A, 1./3);

    ImGui::DragFloat("a0 [fm]", &a0, 0.01, 0.1, 2); 
    ImGui::DragFloat("VSO [MeV]", &VSO, 0.1, 0, 40);
    ImGui::BeginDisabled(e==1); ImGui::DragFloat("RSO [fm] ", &RSO, 0.01, 1, 10); ImGui::EndDisabled();
    ImGui::SameLine();
    ImGui::BeginDisabled(e==0); ImGui::DragFloat("rSO [fm] ", &rSO, 0.01, 1, 10); ImGui::EndDisabled();
    if( e == 0 ) {A = 1; rSO = RSO / pow(A, 1./3);}
    if( e == 1 ) RSO = rSO * pow(A, 1./3);
    ImGui::DragFloat("aSO [fm]", &aSO, 0.01, 0.1, 2);

    ImGui::Separator();
    ImGui::DragInt("nStep", &nStep, 50, 100, 400); 
    ImGui::SameLine();
    ImGui::DragFloat("dr [fm]", &dr, 0.001, 0.01, 0.1);

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

      if( nucleon == 0) {
        ws.IsNeutron();
      }else{
        ws.IsProton();
      }

      if( e == 0) {
        ws.SetNucleus(1,Z);
      }else{
        ws.SetNucleus(A,Z);
        mass = 931.5 * (ws.GetA() )/(1.008664 + ws.GetA());
        ws.SetMass(mass);
      }

      if( Z > 0 ) 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();

    ImGui::SetNextItemOpen(true, ImGuiCond_Always);
    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::TreePop();
    }

    if( ImPlot::BeginPlot("Plot", ImVec2(ImGui::GetWindowWidth()-20, 400)) ){
      ImPlot::SetupLegend(ImPlotLocation_SouthEast);
      ImPlot::SetupAxes("r [fm]"," Energy [MeV]");
      ImPlot::SetupAxesLimits(0, 10, floor(V0*1.1), 10);
      ImPlot::SetupAxisLimitsConstraints(ImAxis_X1, 0, 30);
      ImPlot::SetupAxisLimitsConstraints(ImAxis_Y1, floor(V0*1.1), abs(floor(V0*1.1)));
      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();
  }

  //*============= WS Trend
  {
    ImGui::SetNextWindowSize(ImVec2(600, 800), ImGuiCond_FirstUseEver);
    ImGui::SetNextWindowPos(ImVec2(100, 200), ImGuiCond_FirstUseEver);
    static float V0 = -45, r0 = 1.25, a0 = 0.6;
    static float VSO = 28, rSO = 1.25,  aSO = 0.6;
    static int Z = 0, N = 1, nStep = 300;
    static float rc = 1.25, dr = 0.1;
    static int A[2] = {10, 20};
    static int NRange[2] = {1, 20};
    static int ZRange[2] = {1, 20};
    static float mass = 939.565;
    static float kappa = -1;

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

    static int e = 0;
    ImGui::RadioButton("Stablility", &e, 0); ImGui::SameLine();
    ImGui::RadioButton("var. A", &e, 1); ImGui::SameLine();
    ImGui::RadioButton("var. N", &e, 2); ImGui::SameLine();
    ImGui::RadioButton("var. Z", &e, 3);
    ImGui::Separator();

    ImGui::PushItemWidth(150);
    static int nucleon = 0; //neutron or proton
    if( e != 1 ){
      ImGui::RadioButton("Neutron", &nucleon, 0); ImGui::SameLine();
      ImGui::RadioButton("Proton", &nucleon, 1);
    }else{
      ImGui::Text("Neutron only");
    }

    if( e == 0 || e == 1 ){
      ImGui::DragInt2("A range", A, 1, 1, 300);
    }

    if( e == 2){

      ImGui::DragInt("Z", &Z, 0, 0, 100);
      ImGui::DragInt2("N range", NRange, 1, 1, 300);
    }

    if( e == 3){
      ImGui::DragInt("N", &N, 1, 1, 200);
      ImGui::DragInt2("Z range", ZRange, 1, 0, 300);
    }

    if( e != 1){
      ImGui::DragFloat("kappa", &kappa, 0.1, -10, 10); 
    }

    ImGui::Separator();

    ImGui::DragFloat("V0 [MeV]", &V0, 0.1, -100, 0); 
    ImGui::DragFloat("r0 [fm] ", &r0, 0.01, 1, 10); 
    ImGui::DragFloat("a0 [fm]", &a0, 0.01, 0.1, 2); 
    ImGui::DragFloat("VSO [MeV]", &VSO, 0.1, 0, 40);
    ImGui::DragFloat("rSO [fm] ", &rSO, 0.01, 1, 10); 
    ImGui::DragFloat("aSO [fm]", &aSO, 0.01, 0.1, 2);

    ImGui::Separator();
    ImGui::DragInt("nStep", &nStep, 50, 100, 400); 
    ImGui::SameLine();
    ImGui::DragFloat("dr [fm]", &dr, 0.001, 0.01, 0.1);


    int ARange[2] ;
    int fixedNZ;
    if( e == 0 || e == 1 ) {ARange[0] = A[0]; ARange[1] = A[1]; fixedNZ = 0;}
    if( e == 2 ) {ARange[0] = NRange[0] + Z; ARange[1] = NRange[1] + Z; fixedNZ = Z;}
    if( e == 3 ) {ARange[0] = ZRange[0] + N; ARange[1] = ZRange[1] + N; fixedNZ = N;}
    int maxA = ARange[1];
    if( maxA < A[0]) {
      maxA = ARange[0];
      ARange[0] = ARange[1];
      ARange[1] = maxA;
    }

    orbitalStr.clear();
    orbitalStr.push_back("0s1/2");   ///2
    orbitalStr.push_back("0p3/2");
    orbitalStr.push_back("0p1/2");   ///8
    if( maxA > 8 ) {
      orbitalStr.push_back("0d5/2");
      orbitalStr.push_back("1s1/2");
      orbitalStr.push_back("0d3/2");   ///20
    }
    if( maxA > 20 ){
      orbitalStr.push_back("0f7/2");   ///28
      orbitalStr.push_back("1p3/2");
      orbitalStr.push_back("0f5/2");
    }
    if( maxA > 40 ){
      orbitalStr.push_back("1p1/2");   ///40
      orbitalStr.push_back("0g9/2");   ///50
      orbitalStr.push_back("0g7/2");   
      orbitalStr.push_back("1d5/2");   
      orbitalStr.push_back("0h11/2"); 
      orbitalStr.push_back("1d3/2");   
    }
    if( maxA > 82 ){
      orbitalStr.push_back("2s1/2");   ///82
      orbitalStr.push_back("0h9/2"); 
      orbitalStr.push_back("1f7/2"); 
      orbitalStr.push_back("0i13/2"); 
      orbitalStr.push_back("2p3/2"); 
      orbitalStr.push_back("1f5/2"); 
    }
    if( maxA > 126 ){
      orbitalStr.push_back("2p1/2");    ///126 
      orbitalStr.push_back("1g9/2"); 
      orbitalStr.push_back("0i11/2"); 
      orbitalStr.push_back("0j15/2"); 
      orbitalStr.push_back("0j15/2"); 
      orbitalStr.push_back("2d5/2"); 
      orbitalStr.push_back("3s1/2"); 
      orbitalStr.push_back("2d3/2"); 
      orbitalStr.push_back("1g7/2");    ///184 
    }

    static float progress = 0;
    static bool cal = false;
    static bool finsihed = false;
    static int ANZ = ARange[0];

    if( ImGui::Button("Cal WS Trend") ){
      cal = true;
      finsihed = false;
      AList.clear();
      enLevelList.clear();
      ANZ = ARange[0];
    }

    if(cal) {
      WoodsSaxon ws;

      ws.Setr0(r0);
      ws.SetrSO(rSO);
      ws.Setrc(rc);
      ws.SetVSO(VSO);
      ws.SetRange2(0.001, dr, nStep);
      if( nucleon == 0 ) {
        ws.IsNeutron();
      }else{
        ws.IsProton();
      }
      
      if( ANZ <= ARange[1] ){
         
        AList.push_back(ANZ);
        if( e == 0){
          ws.SetNucleus ( ANZ, 0.008 + 0.495 * ANZ - 0.00076 * ANZ *ANZ + 1.4e-6 * ANZ * ANZ * ANZ); // aproximate stable vallage
          ws.CalRadius();
          ws.SetV0 ( V0 * ( 1 +  kappa * ( 2.0* ws.GetN() - ws.GetA() ) / ws.GetA() ) ); 
        }

        if( e == 1){
          ws.SetNucleus ( ANZ, 0);
          ws.CalRadius();
          ws.SetV0( V0 );
        }

        if( e == 2 ){ /// change in N, ANZ = N
          ws.SetNucleus( ANZ , fixedNZ );
          ws.CalRadius();

          ws.SetV0 ( V0 * ( 1 +  kappa * ( 2.0* ws.GetN() - ws.GetA() ) / ws.GetA() ) ); 
        }
    
        if( e == 3 ){ /// change in Z, ANZ = Z
          ws.SetNucleus( ANZ , ANZ - fixedNZ );
          ws.CalRadius();
          ws.SetV0 ( V0 * ( 1 +  kappa * ( 2.0* ws.GetN() - ws.GetA() ) / ws.GetA() ) ); 
        }

        double reducedMass = 931.5 * (ws.GetA() )/(1.008664 + ws.GetA());
        ws.SetMass(reducedMass);

        ws.ClearVector();
        ws.CalWSEnergies(false, 7, 500, 1e-7, 400, 0.2, false);

        //ws.PrintWSParas();
        // ws.PrintEnergyLevels();

        vector<double> exList;
        for( int i = 0; i < orbitalStr.size(); i++){
          bool isMatched = false;
          for( int j = 0; j < (ws.energy).size(); j++){
            if( orbitalStr[i] == ws.orbString[j] ) {
              exList.push_back(ws.energy[j]);
              isMatched = true;
            }
          }
          if( isMatched == false ) exList.push_back(0);
        }
        enLevelList.push_back(exList);

        progress = (ANZ - ARange[0]) * 1.0 / (ARange[1] - ARange[0]);

        ANZ ++;
      }else{
        cal = false;
        finsihed = true;
      }
    }
    ImGui::SameLine(); ImGui::ProgressBar(progress);
    
    if( finsihed ){
      if( ImPlot::BeginPlot("Plot", ImVec2(ImGui::GetWindowWidth()-20, 400)) ){
        ImPlot::SetupLegend(ImPlotLocation_East, ImPlotLegendFlags_Outside);
        ImPlot::SetupAxes("A"," Energy [MeV]");
        ImPlot::SetupAxesLimits(ARange[0], ARange[1], -50, 0);
        // ImPlot::SetupAxisLimitsConstraints(ImAxis_X1, ARange[0], ARange[1]);
        // ImPlot::SetupAxisLimitsConstraints(ImAxis_Y1, floor(V0)*2, 0);
        const double * kaka = AList.data();
        for( int i = 0; i < orbitalStr.size(); i++){
          vector<double> jaja;
          for( int j = 0; j < AList.size(); j++) jaja.push_back(enLevelList[j][i]);
          const double * haha = jaja.data();
          ImPlot::PlotLine(orbitalStr[i].c_str(), kaka, haha, AList.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;

}