#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 <cstring>
#include <vector>

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

  ImGui::SetNextWindowSize(ImVec2(600, 1000), ImGuiCond_FirstUseEver);
  ImGui::SetNextWindowPos(ImVec2(50, 50), ImGuiCond_FirstUseEver);
  
  {
    ImGui::Begin("Woods-Saxon Calculation");
    
    if( ImGui::BeginTabBar("Woods-Saxon", ImGuiTabBarFlags_None)){
      if (ImGui::BeginTabItem("Cal.")){

        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;

        static int e = 0;
        if( ImGui::RadioButton("Arbitary", &e, 0) ){
          A = 1; 
        }
        ImGui::SameLine();
        if( ImGui::RadioButton("Isotope", &e, 1) ){
          rc = 1.25;
          r0 = 1.25;
          rSO = 1.25;
        }

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

        if( e == 1){
          Rc = rc * pow(A, 1./3);
          R0 = r0 * pow(A, 1./3);
          RSO = rSO * pow(A, 1./3);
        }
        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();

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

        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();
        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::EndTabItem();
      }

      if( ImGui::BeginTabItem("Range")){

        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;

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

          static bool saveDataClicked = false;
          if( ImGui::Button("Download data") ){
            // FILE * file_out;
            // file_out = fopen("", "w+");
            saveDataClicked = true;
          }

          if( saveDataClicked ) {
            ImGui::Text("to be impletmented....");
          } 
        }

        ImGui::EndTabItem();
      }

      if( ImGui::BeginTabItem("Fit")){
        ImGui::Text("to be impletment....");
        ImGui::EndTabItem();
      }
      ImGui::EndTabBar();
    }
    ImGui::End();
  }

  /*
  ImGui::SetNextWindowSize(ImVec2(1000, 400), ImGuiCond_FirstUseEver);
  ImGui::SetNextWindowPos(ImVec2(680, 50), ImGuiCond_FirstUseEver);
  {
    ImGui::Begin("Ptolemy");

    const int maxCount = 20;
    static char reactionChar[maxCount][40];
    static char gsSpinChar[maxCount][10];
    static char KEAChar[maxCount][10];
    static char orbitalChar[maxCount][10];
    static char ExChar[maxCount][10];
    static char jpiChar[maxCount][10];
    static char potInChar[maxCount][10];
    static char potOutChar[maxCount][10];
    
    static int count = 1;
    ImGui::InputInt("Number of states", &count, 1);

    const int nCol = 8;
    if( ImGui::BeginTable("Input", nCol, ImGuiTableFlags_None) ){

      {//Header
        ImGui::TableNextColumn(); ImGui::Text("Reaction");
        ImGui::TableNextColumn(); ImGui::Text("G.S. Spin");
        ImGui::TableNextColumn(); ImGui::Text("Energy [MeV/u]");
        ImGui::TableNextColumn(); ImGui::Text("Orbital");
        ImGui::TableNextColumn(); ImGui::Text("Ex [MeV]");
        ImGui::TableNextColumn(); ImGui::Text("J-pi");
        ImGui::TableNextColumn(); ImGui::Text("Pot-in");
        ImGui::TableNextColumn(); ImGui::Text("Pot-out");
        ImGui::TableNextRow();
      }


      for( int i = 0; i < count; i++){
        
        ImGui::TableNextColumn(); ImGui::InputText(" ", reactionChar[i], 40) ;
        ImGui::TableNextColumn(); ImGui::InputText(" ", gsSpinChar[i], 10) ;
        ImGui::TableNextColumn(); ImGui::InputText(" ", KEAChar[i], 10) ;
        ImGui::TableNextColumn(); ImGui::InputText(" ", orbitalChar[i], 10) ;
        ImGui::TableNextColumn(); ImGui::InputText(" ", ExChar[i], 10) ;
        ImGui::TableNextColumn(); ImGui::InputText(" ", jpiChar[i], 10) ;
        ImGui::TableNextColumn(); ImGui::InputText(" ", potInChar[i], 10) ;
        ImGui::TableNextColumn(); ImGui::InputText(" ", potOutChar[i], 10) ;
        ImGui::TableNextRow();
      }
      ImGui::EndTable();
    }

    if( ImGui::Button("Calculate") ){
      
      ///==== create infile

      ///==== run Ptolemy

      ///==== extract d.s.c. from out file

      ///==== plot

    }

    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, "Woods-Saxon Calculation 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;

}