/* orbital.cc -- OpenGL window for atomic orbitals display.  -*- C++ -*-
   Copyright (C) 2025-2026 Gerardo Ballabio

   This program is free software; you can redistribute it and/or modify
   it under the terms of the GNU General Public License as published by
   the Free Software Foundation; either version 3 of the License, or
   (at your option) any later version.

   This program is distributed in the hope that it will be useful, but
   WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
   GNU General Public License for more details.

   You should have received a copy of the GNU General Public License
   along with this program.  If not, see <http://www.gnu.org/licenses/>. */

#include <cstdlib>
#include <vector>
#include <gtkmm.h>
#include <epoxy/gl.h>
#include "orbital.h"

const int tex_size = 64;
const int tick_duration = 50; // milliseconds
const double pi = 3.141592653589793238;

void gl_orbital::draw_texture(int n, const GLfloat *in,
			      GLfloat scalex, GLfloat scaley)
{
  glActiveTexture(GL_TEXTURE1);
  glBindVertexArray(gl_vertex_array);
  glBindBuffer(GL_ARRAY_BUFFER, gl_buffer);

  GLdouble zoom[] = { 1.0 / scale, 1.0 / scale, 1.0 / scale, 1.0, };
  GLfloat out[4];
  // multiply by current matrix, apply scale
  for (int i=0; i<4; ++i)
    {
      out[i] = 0.0;
      for (int j=0; j<4; ++j)
	out[i] += in[j] * zoom[j] * matrix[i+j*4];
    }
  // normalize
  for (int i=0; i<4; ++i)
    out[i] /= out[3];
  // scale and round to integer
  out[0] = std::floor(out[0] / scalex + 0.5);
  out[1] = std::floor(out[1] / scaley + 0.5);

  // texture coordinates in pixel units
  int ntex = 140 / (n + 25); // start from 5, decrease gradually to 0
  const int size = 6;
  GLfloat rectangle[] =
    {
      -size * 0.5, -size * 0.5, 0.0,
       size * 0.5, -size * 0.5, 0.0,
       size * 0.5,  size * 0.5, 0.0,
      -size * 0.5,  size * 0.5, 0.0,
    };
  for (int i=0; i<4; ++i)
    for (int j=0; j<3; ++j)
      rectangle[i*3+j] += out[j];
  GLfloat origin_x = 0.0,
    origin_y = (GLfloat) size / tex_size * ntex,
    scale_x = (GLfloat) size / tex_size,
    scale_y = (GLfloat) size / tex_size;
  GLfloat trectangle[] =
    {
      origin_x,           origin_y,
      origin_x + scale_x, origin_y,
      origin_x + scale_x, origin_y + scale_y,
      origin_x,           origin_y + scale_y,
    };
  glBufferSubData(GL_ARRAY_BUFFER, 0, sizeof(rectangle), rectangle);
  glBufferSubData(GL_ARRAY_BUFFER, 12 * sizeof(GLfloat), sizeof(trectangle),
		  trectangle);

  glUniform1i(u_sampler, 1);
  glDrawArrays(GL_TRIANGLE_FAN, 0, 4);
}

void gl_orbital::init_rand()
{
  // set up vectors for generating random points

  // radial coordinate
  // build a set of points at increasing distance from each other
  const unsigned nr = 3000;
  const double first = 0.0001, factor = 1.001;
  r.resize(nr);
  fr.resize(nr);
  double step = first;
  r[0] = step;
  for (unsigned i=1; i<nr; ++i)
    r[i] = r[i - 1] + (step *= factor);

  // angular coordinates
  // build a two-dimensional grid of normalized vectors
  const unsigned ntheta = 120, nphi = 240;
  stheta.resize(ntheta + 1);
  ctheta.resize(ntheta + 1);
  sphi.resize(nphi + 1);
  cphi.resize(nphi + 1);
  fa.resize(ntheta * nphi);
  for (unsigned j=0; j<=ntheta; ++j)
    {
      stheta[j] = std::sin(j * pi / ntheta);
      ctheta[j] = std::cos(j * pi / ntheta);
    }
  for (unsigned k=0; k<=nphi; ++k)
    {
      sphi[k] = std::sin(k * 2.0 * pi / nphi);
      cphi[k] = std::cos(k * 2.0 * pi / nphi);
    }
}

void gl_orbital::set_rand()
{
  // populate vectors for generating random points
  // fr, fa are the cumulative density probability function for the radial
  // and angular part, respectively

  rand_scale = default_scale();

  // radial coordinate
  double ff = 0.0, cr = 0.0, dr = 0.0, sr = rand_scale;
  for (unsigned i=0; i<r.size(); ++i)
    {
      // average between start and end point of interval
      double rr = r[i] * sr, f2 = wf.radial(rr) * rr, ff2 = f2 * f2;
      fr[i] = cr += (rr - dr) * (ff + ff2) / 2.0;
      ff = ff2;
      dr = rr;
    }

  // angular coordinates
  double ca = 0.0,
    factor = 2.0 * pi*pi / ((stheta.size() - 1) * (sphi.size() - 1));
  for (unsigned j=0, n=0; j<stheta.size()-1; ++j)
    {
      double f = wf.angular(cphi[0] * stheta[j], sphi[0] * stheta[j],
			    ctheta[j]), ff = f * f;
      for (unsigned k=1; k<sphi.size(); ++k, ++n)
	{
	  double f2 = wf.angular(cphi[k] * stheta[j], sphi[k] * stheta[j],
				 ctheta[j]), ff2 = f2 * f2;
	  fa[n] = ca += factor *
	    (ff + ff2) * (stheta[j] + stheta[j + 1]) / 4.0;
	  ff = ff2;
	}
    }
}

void gl_orbital::set_electron()
{
  // generate random coordinates according to probability density of
  // selected orbital
  // two random numbers would be strictly necessary, but using more makes
  // it easier and they are cheap
  // no need to use a sophisticated generator, std::rand is good enough
  double rand1 = (double) std::rand() / RAND_MAX * fr.back(),
    rand11 = (double) std::rand() / RAND_MAX,
    rand2 = (double) std::rand() / RAND_MAX * fa.back(),
    rand21 = (double) std::rand() / RAND_MAX,
    rand22 = (double) std::rand() / RAND_MAX;
    
  // radial coordinate

  // find interval that contains rand1 by bisection
  int i1 = 0, i2 = fr.size();
  while (i2 - i1 > 1)
    {
      int i = (i1 + i2) / 2;
      if (fr[i - 1] < rand1) // i can't be zero (i2 is at least 2)
	i1 = i;
      else
	i2 = i;
    }
  // i1 is between 0 and fr.size()-1 included
  // rand1 is between fr[i1 - 1] (or 0 if i1==0) and fr[i1]

  // use uniform distribution within selected interval
  // this isn't correct but the error is small
  double r0 = (i1 > 0) ? r[i1 - 1] : 0.0, r1 = r[i1],
    rr = (r0 + (r1 - r0) * rand11) * rand_scale;

  // angular coordinates

  // find interval that contains rand2 by bisection
  unsigned j1 = 0, j2 = fa.size();
  while (j2 - j1 > 1)
    {
      unsigned j = (j1 + j2) / 2;
      if (fa[j - 1] < rand2) // j can't be zero (j2 is at least 2)
	j1 = j;
      else
	j2 = j;
    }
  // j1 is between 0 and fa.size()-1 included
  // rand2 is between fa[j1 - 1] (or 0 if j1==0) and fa[j1]

  // use uniform distribution within selected interval for both coordinates
  // this isn't correct but the error is small
  // interpolating sines and cosines linearly isn't correct either
  // but the error is also small
  unsigned t = j1 / (sphi.size() - 1), p = j1 % (sphi.size() - 1);
  double st = stheta[t] + (stheta[t + 1] - stheta[t]) * rand21,
    ct = ctheta[t] + (ctheta[t + 1] - ctheta[t]) * rand21,
    sp = sphi[p] + (sphi[p + 1] - sphi[p]) * rand22,
    cp = cphi[p] + (cphi[p + 1] - cphi[p]) * rand22;
  double ax = st * cp, ay = st * sp, az = ct;

  // combine radial and angular part
  vertices.push_back(rr * ax);
  vertices.push_back(rr * ay);
  vertices.push_back(rr * az);
  vertices.push_back(1.0);
}

gl_orbital::gl_orbital()
  : begin(0), end(0), draw_electron(true), npoints(600)
{ }

void gl_orbital::init()
{
  // build vertex shader
  const GLchar *vertex_source =
    "#version 330                                     \n"
    "                                                   "
    "uniform mat4 u_matrix;                             "
    "                                                   "
    "in vec3 in_position;                               "
    "in vec2 in_coord;                                  "
    "                                                   "
    "out vec2 tex_coord;                                "
    "                                                   "
    "void main()                                        "
    "{                                                  "
    "  gl_Position = u_matrix * vec4(in_position, 1.0); "
    "  tex_coord = in_coord;                            "
    "}                                                  ";
  gl_vertex_shader = glCreateShader(GL_VERTEX_SHADER);
  glShaderSource(gl_vertex_shader, 1, &vertex_source, NULL);
  glCompileShader(gl_vertex_shader);

  // build fragment shader
  const GLchar *fragment_source =
    "#version 330                                                \n"
    "                                                              "
    "uniform sampler2D u_sampler;                                  "
    "                                                              "
    "in vec2 tex_coord;                                            "
    "                                                              "
    "out vec4 color;                                               "
    "                                                              "
    "void main()                                                   "
    "{                                                             "
    "  /* texture contains alpha channel, color is green */        "
    "  color = vec4(0.0, 1.0, 0.0, texture(u_sampler, tex_coord)); "
    "  /* don't write on depth buffer if alpha is about zero */    "
    "  if (color.a <= 0.2)                                         "
    "    discard;                                                  "
    "}                                                             ";
  gl_fragment_shader = glCreateShader(GL_FRAGMENT_SHADER);
  glShaderSource(gl_fragment_shader, 1, &fragment_source, NULL);
  glCompileShader(gl_fragment_shader);

  // build shader program
  gl_program = glCreateProgram();
  glAttachShader(gl_program, gl_fragment_shader);
  glAttachShader(gl_program, gl_vertex_shader);
  glLinkProgram(gl_program);

  // uniform variables
  u_matrix = glGetUniformLocation(gl_program, "u_matrix");
  u_sampler = glGetUniformLocation(gl_program, "u_sampler");

  // create texture
  const GLubyte texture[] =
    {
      // 1 pixel                           // pad to 8 bytes
      0x00, 0x00, 0x00, 0x00, 0x00, 0x00,  0x00, 0x00,
      0x00, 0x00, 0x00, 0x00, 0x00, 0x00,  0x00, 0x00,
      0x00, 0x00, 0x00, 0xff, 0x00, 0x00,  0x00, 0x00,
      0x00, 0x00, 0x00, 0x00, 0x00, 0x00,  0x00, 0x00,
      0x00, 0x00, 0x00, 0x00, 0x00, 0x00,  0x00, 0x00,
      0x00, 0x00, 0x00, 0x00, 0x00, 0x00,  0x00, 0x00,
      // about 1.5 pixels
      0x00, 0x00, 0x00, 0x00, 0x00, 0x00,  0x00, 0x00,
      0x00, 0x00, 0x00, 0x00, 0x00, 0x00,  0x00, 0x00,
      0x00, 0x00, 0xd0, 0xd0, 0x00, 0x00,  0x00, 0x00,
      0x00, 0x00, 0xd0, 0xd0, 0x00, 0x00,  0x00, 0x00,
      0x00, 0x00, 0x00, 0x00, 0x00, 0x00,  0x00, 0x00,
      0x00, 0x00, 0x00, 0x00, 0x00, 0x00,  0x00, 0x00,
      // 2 pixels
      0x00, 0x00, 0x00, 0x00, 0x00, 0x00,  0x00, 0x00,
      0x00, 0x00, 0x00, 0x00, 0x00, 0x00,  0x00, 0x00,
      0x00, 0x00, 0xff, 0xff, 0x00, 0x00,  0x00, 0x00,
      0x00, 0x00, 0xff, 0xff, 0x00, 0x00,  0x00, 0x00,
      0x00, 0x00, 0x00, 0x00, 0x00, 0x00,  0x00, 0x00,
      0x00, 0x00, 0x00, 0x00, 0x00, 0x00,  0x00, 0x00,
      // 3 pixels
      0x00, 0x00, 0x00, 0x00, 0x00, 0x00,  0x00, 0x00,
      0x00, 0x00, 0xc0, 0xff, 0xc0, 0x00,  0x00, 0x00,
      0x00, 0x00, 0xff, 0xff, 0xff, 0x00,  0x00, 0x00,
      0x00, 0x00, 0xc0, 0xff, 0xc0, 0x00,  0x00, 0x00,
      0x00, 0x00, 0x00, 0x00, 0x00, 0x00,  0x00, 0x00,
      0x00, 0x00, 0x00, 0x00, 0x00, 0x00,  0x00, 0x00,
      // 4 pixels
      0x00, 0x00, 0x00, 0x00, 0x00, 0x00,  0x00, 0x00,
      0x00, 0xc0, 0xff, 0xff, 0xc0, 0x00,  0x00, 0x00,
      0x00, 0xff, 0xff, 0xff, 0xff, 0x00,  0x00, 0x00,
      0x00, 0xff, 0xff, 0xff, 0xff, 0x00,  0x00, 0x00,
      0x00, 0xc0, 0xff, 0xff, 0xc0, 0x00,  0x00, 0x00,
      0x00, 0x00, 0x00, 0x00, 0x00, 0x00,  0x00, 0x00,
      // 6 pixels
      0x00, 0xc0, 0xff, 0xff, 0xc0, 0x00,  0x00, 0x00,
      0xc0, 0xff, 0xff, 0xff, 0xff, 0xc0,  0x00, 0x00,
      0xff, 0xff, 0xff, 0xff, 0xff, 0xff,  0x00, 0x00,
      0xff, 0xff, 0xff, 0xff, 0xff, 0xff,  0x00, 0x00,
      0xc0, 0xff, 0xff, 0xff, 0xff, 0xc0,  0x00, 0x00,
      0x00, 0xc0, 0xff, 0xff, 0xc0, 0x00,  0x00, 0x00,
    };
  glGenTextures(1, &gl_texture);
  glActiveTexture(GL_TEXTURE1);
  glBindTexture(GL_TEXTURE_2D, gl_texture);
  glTexStorage2D(GL_TEXTURE_2D, 1, GL_R8, tex_size, tex_size);
  glTexSubImage2D(GL_TEXTURE_2D, 0, 0, 0, 8, 36, GL_RED, GL_UNSIGNED_BYTE,
		  (void *) texture);

  // build vertex array
  glGenVertexArrays(1, &gl_vertex_array);
  glBindVertexArray(gl_vertex_array);
  glGenBuffers(1, &gl_buffer);
  glBindBuffer(GL_ARRAY_BUFFER, gl_buffer);
  GLfloat quad[] =
    {
      0.0, 0.0,
      1.0, 0.0,
      1.0, 1.0,
      0.0, 1.0,
    };
  glBufferData(GL_ARRAY_BUFFER, 12 * sizeof(GLfloat) + sizeof(quad), NULL,
	       GL_DYNAMIC_DRAW);
  glBufferSubData(GL_ARRAY_BUFFER, 12 * sizeof(GLfloat), sizeof(quad), quad);

  // vertex attributes
  glVertexAttribPointer(0, 3, GL_FLOAT, GL_FALSE, 0, (void *) 0);
  glVertexAttribPointer(1, 2, GL_FLOAT, GL_FALSE, 0,
			(void *) (12 * sizeof(float)));
  glEnableVertexAttribArray(0);
  glEnableVertexAttribArray(1);

  init_rand();
  set(1, 0, 0);
}

void gl_orbital::set_orbital(int width_, int height_)
{
  width = width_;
  height = height_;
}

void gl_orbital::set_matrix(const GLfloat *matrix_)
{
  for (GLuint i=0; i<16; ++i)
    matrix[i] = matrix_[i];
}

bool gl_orbital::set_z(unsigned z_)
{
  queue_rand = true;
  return wf.set_z(z_);
}

bool gl_orbital::set(int n_, int l_, int m_)
{
  queue_rand = true;
  return wf.set(n_, l_, m_);
}

void gl_orbital::set_draw_electron(bool draw_)
{
  draw_electron = draw_;
  if (!draw_electron)
    reset();
}

void gl_orbital::reset()
{
  vertices.resize(0);
  begin = end = 0;
}

void gl_orbital::update_rand()
{
  if (queue_rand)
    {
      queue_rand = false;
      reset();
      set_rand();
    }
}

void gl_orbital::tick()
{
  // add a new point
  // discard one if there are already npoints: first in, first out
  // do the right thing even if npoints has changed
  if (++end > begin + npoints)
    if ((begin = end - npoints) >= npoints)
      {
	// resize vector every npoints points
	vertices.erase(vertices.begin(), vertices.begin() + npoints * 4);
	begin -= npoints;
	end = begin + npoints;
      }
  set_electron();
}

void gl_orbital::draw()
{
  if (!draw_electron)
    return;

  glUseProgram(gl_program);

  // draw textures at their actual pixel size
  // to look right, their coordinates must correspond to whole pixels
  // so the origin of the axes is kept at the center of the window
  // but if width or height are odd, it is shifted by half pixel
  GLfloat scalex = 2.0 / width, scaley = 2.0 / height,
    oddx = (GLfloat) (width - (width / 2) * 2) / width, // half pixel or 0
    oddy = (GLfloat) (height - (height / 2) * 2) / height;
  GLfloat matrix_tex[] = {
    scalex, 0.0,    0.0, 0.0,
    0.0,    scaley, 0.0, 0.0,
    0.0,    0.0,    1.0, 0.0,
    oddx,   oddy,   0.0, 1.0,
  };
  glUniformMatrix4fv(u_matrix, 1, GL_FALSE, matrix_tex);

  for (int i=begin; i<end; ++i)
    draw_texture(end - i, &vertices[i*4], scalex, scaley);

  glUseProgram(0);
}

void OrbitalArea::on_realize()
{
  Area::on_realize();
  orbital.init();
  set_scale(default_scale());
  set_vmax(orbital.max_wave());
  set_vscale(default_vscale() / get_vmax());
  set_f([this]
	(double x, double y, double z){ return orbital.wave(x, y, z); });
}

void OrbitalArea::on_grid_signal_values(double vmax_)
{
  // finalize change of orbital
  set_vmax(orbital.max_wave());
  set_vscale(default_vscale() / get_vmax());
  update_rand();
}

void OrbitalArea::on_grid_signal_scale(double scale_)
{
  scale = scale_;
  frame.set_scale(scale);
  orbital.set_scale(scale);
}

void OrbitalArea::resize(int width_, int height_)
{
  Area::resize(width_, height_);
  orbital.set_orbital(width_, height_);
}

void OrbitalArea::draw()
{
  // multiply matrices
  for (unsigned i=0; i<4; ++i)
    for (unsigned j=0; j<4; ++j)
      {
	matrix[i*4+j] = 0.0;
	for (unsigned k=0; k<4; ++k)
	  matrix[i*4+j] += view_matrix[j+k*4] * model_matrix[i*4+k];
      }

  orbital.set_matrix(matrix);
  frame.set_matrix(matrix);
  grid.set_view_matrix(view_matrix);
  grid.set_model_matrix(model_matrix);

  orbital.draw();
  frame.draw();
  grid.draw();
}

OrbitalArea::OrbitalArea()
{
  set_draw_electron(true);
  set_fix_vmax(true);

  // connect signal
  grid.signal_values().connect
    (sigc::mem_fun(*this, &OrbitalArea::on_grid_signal_values));
}

bool OrbitalArea::set_z(unsigned z_)
{
  bool rvalue = orbital.set_z(z_);
  set_new_scale(default_scale());
  set_new_vscale(default_vscale());
  set_vmax(orbital.max_wave());
  return rvalue;
}

bool OrbitalArea::set(int n_, int l_, int m_)
{
  bool rvalue = orbital.set(n_, l_, m_);
  set_new_scale(default_scale());
  set_new_vscale(default_vscale());
  set_vmax(orbital.max_wave());
  return rvalue;
}

void OrbitalArea::set_draw_electron(bool draw_)
{
  orbital.set_draw_electron(draw_);

  if (orbital.get_draw_electron())
    {
      // activate timeout
      timeout2 = Glib::signal_timeout().connect
	(sigc::mem_fun(*this, &OrbitalArea::on_timeout2), tick_duration);
      // draw orbital with more transparency
      set_ascale(0.4);
    }
  else
    {
      timeout2.disconnect();
      set_ascale(1.0);
    }
}

void OrbitalArea::reset()
{
  init_model_matrix();
  set_scale(default_scale());
  resize(width, height);
  grid.queue_set_values();
}
