text: // Copyright (c) 2013-2016 Anton Kozhevnikov, Thomas Schulthess
  // All rights reserved.
  // 
  // Redistribution and use in source and binary forms, with or without modification, are permitted provided that 
  // the following conditions are met:
  // 
  // 1. Redistributions of source code must retain the above copyright notice, this list of conditions and the 
  //    following disclaimer.
  // 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions 
  //    and the following disclaimer in the documentation and/or other materials provided with the distribution.
  // 
  // THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED 
  // WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A 
  // PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR 
  // ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, 
  // PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER 
  // CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR 
  // OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
  
  #include <sirius.h>
  
  class Free_atom: public sirius::Atom_type
  {
      private:
  
          mdarray<double, 2> free_atom_orbital_density_;
          mdarray<double, 2> free_atom_wave_functions_;
          sirius::Spline<double> free_atom_potential_;
  
      public:
      
          double NIST_LDA_Etot;
          double NIST_ScRLDA_Etot;
  
          Free_atom(Free_atom&& src) = default;
      
          Free_atom(sirius::Simulation_parameters const& param__,
                    const std::string symbol, 
                    const std::string name, 
                    int zn, 
                    double mass, 
                    std::vector<atomic_level_descriptor>& levels_nl) 
              : Atom_type(param__, symbol, name, zn, mass, levels_nl, sirius::lin_exp_grid), 
                NIST_LDA_Etot(0),
                NIST_ScRLDA_Etot(0)
          {
              radial_grid_ = sirius::Radial_grid_exp<double>(2000 + 150 * zn, 1e-7, 20.0 + 0.25 * zn); 
          }
  
          double ground_state(double solver_tol, double energy_tol, double charge_tol, std::vector<double>& enu, bool rel)
          {
              PROFILE("sirius::Free_atom::ground_state");
          
              int np = radial_grid().num_points();
              assert(np > 0);
  
              free_atom_orbital_density_ = mdarray<double, 2>(np, num_atomic_levels()); 
              free_atom_wave_functions_ = mdarray<double, 2>(np, num_atomic_levels());
              
              sirius::XC_functional Ex("XC_LDA_X", 1);
              sirius::XC_functional Ec("XC_LDA_C_VWN", 1);
              Ex.set_relativistic(rel);
          
              std::vector<double> veff(np);
              std::vector<double> vrho(np);
              std::vector<double> vnuc(np);
              for (int i = 0; i < np; i++)
              {
                  vnuc[i] = -zn() * radial_grid().x_inv(i);
                  veff[i] = vnuc[i];
                  vrho[i] = 0;
              }
  
              sirius::Mixer<double>* mixer = new sirius::Broyden1<double>(0, np, 12, 0.8, mpi_comm_self());
              for (int i = 0; i < np; i++) {
                  mixer->input_local(i, vrho[i]);
              }
              mixer->initialize();
          
              sirius::Spline<double> rho(radial_grid());
          
              sirius::Spline<double> f(radial_grid());
          
              std::vector<double> vh(np);
              std::vector<double> vxc(np);
              std::vector<double> exc(np);
              std::vector<double> vx(np);
              std::vector<double> vc(np);
              std::vector<double> ex(np);
              std::vector<double> ec(np);
              std::vector<double> g1;
              std::vector<double> g2;
              std::vector<double> rho_old;
          
              enu.resize(num_atomic_levels());
          
              double energy_tot = 0.0;
              double energy_tot_old;
              double charge_rms;
              double energy_diff;
              double energy_enuc = 0;
              double energy_xc = 0;
              double energy_kin = 0;
              double energy_coul = 0;
          
              bool converged = false;
              
              /* starting values for E_{nu} */
              for (int ist = 0; ist < num_atomic_levels(); ist++) {
                  enu[ist] = -1.0 * zn() / 2 / std::pow(double(atomic_level(ist).n), 2);
              }
              
              for (int iter = 0; iter < 200; iter++) {
                  rho_old = rho.values();
                  
                  std::memset(&rho(0), 0, rho.num_points() * sizeof(double));
                  #pragma omp parallel default(shared)
                  {
                      std::vector<double> rho_t(rho.num_points());
                      std::memset(&rho_t[0], 0, rho.num_points() * sizeof(double));
                  
                      #pragma omp for
                      for (int ist = 0; ist < num_atomic_levels(); ist++) {
                          //relativity_t rt = (rel) ? relativity_t::koelling_harmon : relativity_t::none;
                          relativity_t rt = (rel) ? relativity_t::dirac : relativity_t::none;
                          sirius::Bound_state bound_state(rt, zn(), atomic_level(ist).n, atomic_level(ist).l,
                                                          atomic_level(ist).k, radial_grid(), veff, enu[ist]);
                          enu[ist] = bound_state.enu();
                          auto& bs_rho = bound_state.rho();
                          auto& bs_u = bound_state.u();
                          
                          /* assume a spherical symmetry */
                          for (int i = 0; i < np; i++) {
                              free_atom_orbital_density_(i, ist) = bs_rho(i);
                              free_atom_wave_functions_(i, ist) = bs_u(i);
                              /* sum of squares of spherical harmonics for angular momentm l is (2l+1)/4pi */
                              rho_t[i] += atomic_level(ist).occupancy * free_atom_orbital_density_(i, ist) / fourpi;
                          }
                      }
          
                      #pragma omp critical
                      for (int i = 0; i < rho.num_points(); i++) {
                          rho(i) += rho_t[i];
                      }
                  } 
                  
                  charge_rms = 0.0;
                  for (int i = 0; i < np; i++) {
                      charge_rms += std::pow(rho(i) - rho_old[i], 2);
                  }
                  charge_rms = sqrt(charge_rms / np);
                  
                  rho.interpolate();
          
                  /* compute Hartree potential */
                  rho.integrate(g2, 2);
                  double t1 = rho.integrate(g1, 1);
          
                  for (int i = 0; i < np; i++) {
                      vh[i] = fourpi * (g2[i] / radial_grid(i) + t1 - g1[i]);
                  }
                  
                  /* compute XC potential and energy */
                  Ex.get_lda(rho.num_points(), &rho(0), &vx[0], &ex[0]);
                  Ec.get_lda(rho.num_points(), &rho(0), &vc[0], &ec[0]);
                  for (int ir = 0; ir < rho.num_points(); ir++) {
                     vxc[ir] = (vx[ir] + vc[ir]);
                     exc[ir] = (ex[ir] + ec[ir]);
                  }
                 
                  /* mix old and new effective potential */
                  for (int i = 0; i < np; i++) {
                      mixer->input_local(i, vh[i] + vxc[i]);
                  }
                  mixer->mix(1e-16);
                  for (int i = 0; i < np; i++) {
                      vrho[i] = mixer->output_local(i);
                      veff[i] = vrho[i] + vnuc[i];
                  }
  
                  //= /* mix old and new effective potential */
                  //= for (int i = 0; i < np; i++)
                  //=     veff[i] = (1 - beta) * veff[i] + beta * (vnuc[i] + vh[i] + vxc[i]);
                  
                  /* sum of occupied eigen values */
                  double eval_sum = 0.0;
                  for (int ist = 0; ist < num_atomic_levels(); ist++) {
                      eval_sum += atomic_level(ist).occupancy * enu[ist];
                  }
          
                  for (int i = 0; i < np; i++) {
                      f(i) = (veff[i] - vnuc[i]) * rho(i);
                  }
                  /* kinetic energy */
                  energy_kin = eval_sum - fourpi * (f.interpolate().integrate(2) - zn() * rho.integrate(1));
                  
                  /* XC energy */
                  for (int i = 0; i < np; i++) {
                      f(i) = exc[i] * rho(i);
                  }
                  energy_xc = fourpi * f.interpolate().integrate(2); 
                  
                  /* electron-nuclear energy: \int vnuc(r) * rho(r) r^2 dr */
                  energy_enuc = -fourpi * zn() * rho.integrate(1); 
          
                  /* Coulomb energy */
                  for (int i = 0; i < np; i++) {
                      f(i) = vh[i] * rho(i);
                  }
                  energy_coul = 0.5 * fourpi * f.interpolate().integrate(2);
                  
                  energy_tot_old = energy_tot;
          
                  energy_tot = energy_kin + energy_xc + energy_coul + energy_enuc; 
                  
                  energy_diff = std::abs(energy_tot - energy_tot_old);
                  
                  if (energy_diff < energy_tol && charge_rms < charge_tol) { 
                      converged = true;
                      printf("Converged in %i iterations.\n", iter);
                      break;
                  }
              }
          
              if (!converged) {
                  printf("energy_diff : %18.10f   charge_rms : %18.10f\n", energy_diff, charge_rms);
                  std::stringstream s;
                  s << "atom " << symbol() << " is not converged" << std::endl
                    << "  energy difference : " << energy_diff << std::endl
                    << "  charge difference : " << charge_rms;
                  TERMINATE(s);
              }
              
              free_atom_density_spline_ = sirius::Spline<double>(radial_grid_, rho.values());
              
              free_atom_potential_ = sirius::Spline<double>(radial_grid_, vrho);
  
              double Eref = (rel) ? NIST_ScRLDA_Etot : NIST_LDA_Etot;
  
              printf("\n");
              printf("Radial gird\n");
              printf("-----------\n");
              printf("type             : %s\n", radial_grid().name().c_str());
              printf("number of points : %i\n", np);
              printf("origin           : %20.12f\n", radial_grid(0));
              printf("infinity         : %20.12f\n", radial_grid(np - 1));
              printf("\n");
              printf("Energy\n");
              printf("------\n");
              printf("Ekin  : %20.12f\n", energy_kin);
              printf("Ecoul : %20.12f\n", energy_coul);
              printf("Eenuc : %20.12f\n", energy_enuc);
              printf("Eexc  : %20.12f\n", energy_xc);
              printf("Total : %20.12f\n", energy_tot);
              printf("NIST  : %20.12f\n", Eref);
  
              /* difference between NIST and computed total energy. Comparison is valid only for VWN XC functional. */
              double dE = (Utils::round(energy_tot, 6) - Eref);
              std::cerr << zn() << " " << dE << " # " << symbol() << std::endl;
              
              return energy_tot;
          }
  
          inline double free_atom_orbital_density(int ir, int ist)
          {
              return free_atom_orbital_density_(ir, ist);
          }
  
          inline double free_atom_wave_function(int ir, int ist)
          {
              return free_atom_wave_functions_(ir, ist);
          }
  
          inline double free_atom_potential(double x)
          {
              return free_atom_potential_.at_point(x);
          }
  };
  
  Free_atom init_atom_configuration(const std::string& label, sirius::Simulation_parameters param__)
  {
      json jin;
      std::ifstream("atoms.json") >> jin;
      
      atomic_level_descriptor nlk;
      std::vector<atomic_level_descriptor> levels_nlk;
      
      for (size_t i = 0; i < jin[label]["levels"].size(); i++) {
          nlk.n = jin[label]["levels"][i][0];
          nlk.l = jin[label]["levels"][i][1];
          nlk.k = jin[label]["levels"][i][2];
          nlk.occupancy = jin[label]["levels"][i][3];
          levels_nlk.push_back(nlk);
      }
  
      int zn;
      zn = jin[label]["zn"];
      double mass;
      mass = jin[label]["mass"];
      std::string name;
      name = jin[label]["name"];
      double NIST_LDA_Etot = 0.0;
      NIST_LDA_Etot = jin[label].value("NIST_LDA_Etot", NIST_LDA_Etot);
      double NIST_ScRLDA_Etot = 0.0;
      NIST_ScRLDA_Etot = jin[label].value("NIST_ScRLDA_Etot", NIST_ScRLDA_Etot);
      
      Free_atom a(param__, label, name, zn, mass, levels_nlk);
      a.NIST_LDA_Etot = NIST_LDA_Etot;
      a.NIST_ScRLDA_Etot = NIST_ScRLDA_Etot;
      return std::move(a);
  }
  
  void generate_atom_file(Free_atom& a,
                          double core_cutoff,
                          const std::string& lo_type,
                          int apw_order,
                          double apw_enu,
                          bool auto_enu,
                          bool write_to_xml,
                          bool rel)
  {
      std::vector<double> enu;
      
      printf("\n");
      printf("atom : %s, Z = %i\n", a.symbol().c_str(), a.zn());
      printf("----------------------------------\n");
     
      /* solve a free atom */
      a.ground_state(1e-10, 1e-8, 1e-8, enu, rel);
     
      //JSON_write jw(fname);
      json dict;
      dict["name"] = a.name();
      dict["symbol"] = a.symbol();
      dict["number"] = a.zn();
      dict["mass"] = a.mass();
      dict["rmin"] = a.radial_grid(0);
  
      std::vector<atomic_level_descriptor> core;
      std::vector<atomic_level_descriptor> valence;
      std::string level_symb[] = {"s", "p", "d", "f"};
  
      int nl_c[8][4];
      int nl_v[8][4];
      double e_nl_c[8][4];
      double e_nl_v[8][4];
  
      std::memset(&nl_c[0][0], 0, 32 * sizeof(int));
      std::memset(&nl_v[0][0], 0, 32 * sizeof(int));
      std::memset(&e_nl_c[0][0], 0, 32 * sizeof(double));
      std::memset(&e_nl_v[0][0], 0, 32 * sizeof(double));
  
      printf("\n");
      printf("Core / valence partitioning\n");
      printf("---------------------------\n");
      if (core_cutoff <= 0) {
          printf("core cutoff energy       : %f\n", core_cutoff);
      } else {
          printf("core cutoff radius       : %f\n", core_cutoff);
      }
      sirius::Spline<double> rho_c(a.radial_grid());
      sirius::Spline<double> rho(a.radial_grid());
      sirius::Spline<double> s(a.radial_grid());
      int ncore{0};
      for (int ist = 0; ist < a.num_atomic_levels(); ist++) {
          int n = a.atomic_level(ist).n;
          int l = a.atomic_level(ist).l;
  
          printf("%i%s  occ : %8.4f  energy : %12.6f", a.atomic_level(ist).n, level_symb[a.atomic_level(ist).l].c_str(), 
                                                       a.atomic_level(ist).occupancy, enu[ist]);
          
          /* total density */
          for (int ir = 0; ir < a.radial_grid().num_points(); ir++) {
              s(ir) = a.free_atom_orbital_density(ir, ist);
              rho(ir) += a.atomic_level(ist).occupancy * s(ir) / fourpi;
          }
  
          std::vector<double> g;
          s.interpolate().integrate(g, 2);
          double rc{100};
  
          for (int ir = 0; ir < a.radial_grid().num_points(); ir++) {
              if (1.0 - g[ir] < 1e-5) {
                  rc = a.radial_grid(ir);
                  break;
              }
          }
  
          /* assign this state to core */
          if ((core_cutoff <= 0 && enu[ist] < core_cutoff) ||
              (core_cutoff > 0 && rc < core_cutoff)) {
              core.push_back(a.atomic_level(ist));
              printf("  => core (rc = %f)\n", rc);
  
              for (int ir = 0; ir < a.radial_grid().num_points(); ir++) {
                  rho_c(ir) += a.atomic_level(ist).occupancy * a.free_atom_orbital_density(ir, ist) / fourpi;
              }
  
              nl_c[n][l]++;
              e_nl_c[n][l] += enu[ist];
              ncore += static_cast<int>(a.atomic_level(ist).occupancy + 1e-12);
          } else { /* assign this state to valence */
              valence.push_back(a.atomic_level(ist));
              printf("  => valence (rc = %f)\n", rc);
  
              nl_v[n][l]++;
              e_nl_v[n][l] += enu[ist];
          }
      }
      printf("number of core electrons : %i\n", ncore);
      
      /* average energies for {n,l} level */
      for (int n = 1; n <= 7; n++) {
          for (int l = 0; l < 4; l++) {
              if (nl_v[n][l]) e_nl_v[n][l] /= nl_v[n][l];
              if (nl_c[n][l]) e_nl_c[n][l] /= nl_c[n][l];
          }
      }
      
      /* get the lmax for valence states */
      int lmax{0};
      for (size_t i = 0; i < valence.size(); i++) {
          lmax = std::max(lmax, valence[i].l); 
      }
      printf("lmax: %i\n", lmax);
  
      /* valence principal quantum numbers for each l */
      std::array<std::vector<int>, 4> n_v;
      for (int n = 1; n <= 7; n++) {
          for (int l = 0; l < 4; l++) {
              if (nl_v[n][l]) {
                  n_v[l].push_back(n);
              }
          }
      }
      
      printf("valence n for each l:\n");
      for (int l = 0; l < 4; l++) {
          printf("l: %i, n: ", l);
          for (int n: n_v[l]) {
              printf("%i ", n);
          }
          printf("\n");
      }
  
      //FILE* fout = fopen("rho.dat", "w");
      //for (int ir = 0; ir < a.radial_grid().num_points(); ir++) 
      //{
      //    double x = a.radial_grid(ir);
      //    fprintf(fout, "%12.6f %16.8f\n", x, rho[ir] * x * x);
      //}
      //fclose(fout);
      
      //mdarray<int, 2> nl_st(8, 4);
      //nl_st.zero();
      //FILE* fout = fopen((a.symbol() + "_wfs.dat").c_str(), "w");
      //for (int ist = 0; ist < a.num_atomic_levels(); ist++) {
      //    int n = a.atomic_level(ist).n;
      //    int l = a.atomic_level(ist).l;
      //    if (!nl_st(n, l) && !nl_c[n][l]) {
      //        for (int ir = 0; ir < a.radial_grid().num_points(); ir++) {
      //            double x = a.radial_grid(ir);
      //            fprintf(fout, "%12.6f %16.8f\n", x, a.free_atom_wave_function(ir, ist));
      //        }
      //        fprintf(fout, "\n");
      //        nl_st(n, l) = 1;
      //    }
      //}
      //fclose(fout);
  
      FILE* fout = fopen((a.symbol() + "_wfs.dat").c_str(), "w");
      for (int ir = 0; ir < a.radial_grid().num_points(); ir++) {
          double x = a.radial_grid(ir);
          fprintf(fout, "%12.6f ", x);
  
          for (int ist = 0; ist < a.num_atomic_levels(); ist++) {
              fprintf(fout, "%12.6f ", a.free_atom_wave_function(ir, ist));
          }
          fprintf(fout, "\n");
      }
      fclose(fout);
      
      /* estimate effective infinity */
      std::vector<double> g;
      rho.interpolate().integrate(g, 2);
      double rinf{0};
      for (int ir = a.radial_grid().num_points() - 1; ir >= 0; ir--) {
          if (g[ir] / g.back() < 0.99999999) {
              rinf = a.radial_grid(ir);
              break;
          }
      }
      printf("Effective infinity : %f\n", rinf);
  
      /* estimate core radius */
      double core_radius = (core_cutoff > 0) ? core_cutoff : 0.65;
      if (ncore != 0) {
          std::vector<double> g;
          rho_c.interpolate().integrate(g, 2);
  
          for (int ir = 0; ir < a.radial_grid().num_points(); ir++) {
              if ((ncore - fourpi * g[ir]) / ncore < 1e-5) {
                  core_radius = a.radial_grid(ir);
                  break;
              }
          }
      }
  
      /* good number of MT points */
      int nrmt{1500};
  
      printf("minimum MT radius : %f\n", core_radius);
      dict["rmt"] = core_radius;
      dict["nrmt"] = nrmt;
      dict["rinf"] = rinf;
      
      /* compact representation of core states */
      std::string core_str;
      int nl_core[8][4];
      std::memset(&nl_core[0][0], 0, 32 * sizeof(int));
      for (size_t i = 0; i < core.size(); i++) {
          std::stringstream ss;
          if (!nl_core[core[i].n][core[i].l]) {
              ss << core[i].n;
              core_str += (ss.str() + level_symb[core[i].l]);
              nl_core[core[i].n][core[i].l] = 1;
          }
      }
      dict["core"] = core_str;
  
      dict["valence"] = json::array();
      dict["valence"].push_back(json::object({{"basis", json::array({})}}));
      dict["valence"][0]["basis"].push_back({{"enu", apw_enu},{"dme", 0}, {"auto", 0}});
      if (apw_order == 2) {
          dict["valence"][0]["basis"].push_back({{"enu", apw_enu},{"dme", 1}, {"auto", 0}});
      }
  
      if (auto_enu) {
          for (int l = 0; l <= lmax; l++) {
              /* default value for n */
              int n = l + 1;
              /* next n above the core */
              for (size_t i = 0; i < core.size(); i++) {
                  if (core[i].l == l) {
                      n = std::max(core[i].n + 1, n);
                  }
              }
              /* APW for s,p,d,f is constructed for the highest valence state */ 
              for (size_t i = 0; i < valence.size(); i++) {
                  if (valence[i].l == l) {
                      n = std::max(n, valence[i].n);
                  }
              }
              dict["valence"].push_back(json::object({{"n", n}, {"l", l}, {"basis", json::array({})}}));
              dict["valence"].back()["basis"].push_back({{"enu", apw_enu},{"dme", 0}, {"auto", 1}});
              if (apw_order == 2) {
                  dict["valence"].back()["basis"].push_back({{"enu", apw_enu},{"dme", 1}, {"auto", 1}});
              }
          }
      }
  
      int idxlo{0};
      for (int n = 1; n <= 7; n++) {
          for (int l = 0; l < 4; l++) {
              if (nl_v[n][l]) {
                  if (lo_type.find("lo1") != std::string::npos) {
                      a.add_lo_descriptor(idxlo, n, l, e_nl_v[n][l], 0, 1);
                      a.add_lo_descriptor(idxlo, n, l, e_nl_v[n][l], 1, 1);
                      idxlo++;
                  }
  
                  if (lo_type.find("lo2") != std::string::npos) {
                      a.add_lo_descriptor(idxlo, n, l, e_nl_v[n][l], 1, 1);
                      a.add_lo_descriptor(idxlo, n, l, e_nl_v[n][l], 2, 1);
                      idxlo++;
                  }
  
                  if (lo_type.find("LO1") != std::string::npos) {
                      a.add_lo_descriptor(idxlo, 0, l, 0.15, 0, 0);
                      a.add_lo_descriptor(idxlo, 0, l, 0.15, 1, 0);
                      a.add_lo_descriptor(idxlo, n, l, e_nl_v[n][l], 0, 1);
                      idxlo++;
                  }
  
                  if (lo_type.find("LO2") != std::string::npos) {
                      a.add_lo_descriptor(idxlo, 0, l, 1.15, 0, 0);
                      a.add_lo_descriptor(idxlo, 0, l, 1.15, 1, 0);
                      a.add_lo_descriptor(idxlo, n + 1, l, e_nl_v[n][l] + 1, 0, 1);
                      idxlo++;
                  }
              }
          }
      }
       
      if (lo_type.find("lo3") != std::string::npos) {
          for (int l = lmax + 1; l < lmax + 4; l++) {
              a.add_lo_descriptor(idxlo, 0, l, 0.15, 0, 0);
              a.add_lo_descriptor(idxlo, 0, l, 0.15, 1, 0);
              idxlo++;
   
              a.add_lo_descriptor(idxlo, 0, l, 0.15, 1, 0);
              a.add_lo_descriptor(idxlo, 0, l, 0.15, 2, 0);
              idxlo++;
          }
      }
  
      if (lo_type.find("lo4") != std::string::npos) {
          for (int n = 1; n <= 6; n++) {
              for (int l = 0; l < n; l++) {
                  a.add_lo_descriptor(idxlo, n, l, 0.15, 0, 1);
                  a.add_lo_descriptor(idxlo, n, l, 0.15, 1, 1);
                  idxlo++;
              }
          }
      }
   
      std::vector<double> fa_rho(a.radial_grid().num_points());
      std::vector<double> fa_r(a.radial_grid().num_points());
   
      for (int i = 0; i < a.radial_grid().num_points(); i++) {
          fa_rho[i] = a.free_atom_density(i);
          fa_r[i] = a.radial_grid(i);
      }
  
      sirius::Radial_grid_lin_exp<double> rg(1500, 1e-7, 2.0);
      std::vector<double> x;
      std::vector<double> veff;
      for (int ir = 0; ir < rg.num_points(); ir++) {
          x.push_back(rg[ir]);
          veff.push_back(a.free_atom_potential(rg[ir]) - a.zn() * rg.x_inv(ir));
      }
      a.set_radial_grid(1500, &x[0]);
      
      printf("=== initializing atom ===\n");
      a.init(0);
      sirius::Atom_symmetry_class atom_class(0, a);
      atom_class.set_spherical_potential(veff);
      atom_class.generate_radial_functions(relativity_t::none);
      runtime::pstdout pout(mpi_comm_self());
      atom_class.write_enu(pout);
      pout.flush();
      
      auto lo_to_str = [](local_orbital_descriptor lod) {
          std::stringstream s;
          s << "[";
          for (size_t o = 0; o < lod.rsd_set.size(); o++) {
              s << "{" << "\"n\" : " <<  lod.rsd_set[o].n 
                       << ", \"enu\" : " <<  lod.rsd_set[o].enu 
                       << ", \"dme\" : " <<  lod.rsd_set[o].dme
                       << ", \"auto\" : " <<  lod.rsd_set[o].auto_enu << "}";
              if (o != lod.rsd_set.size() - 1) {
                  s << ", ";
              }
          }
          s << "]";
          return s.str();
      };
  
      //== std::vector<int> inc(a.num_lo_descriptors(), 1);
  
      //== mdarray<double, 2> rad_func(rg.num_points(), a.num_lo_descriptors());
  
      //== for (int i = 0; i < a.num_lo_descriptors(); i++) {
      //==     int idxrf = a.indexr().index_by_idxlo(i);
      //==     int l = a.lo_descriptor(i).l;
  
      //==     sirius::Spline<double> f(rg);
      //==     for (int ir = 0; ir < a.num_mt_points(); ir++) {
      //==         f[ir] = atom_class.radial_function(ir, idxrf);
      //==     }
  
      //==     sirius::Spline<double> s(rg);
  
      //==     for (int j = 0; j < i; j++) {
      //==         int l1 = a.lo_descriptor(j).l;
  
      //==         if (l1 == l && inc[j]) {
      //==             for (int ir = 0; ir < a.num_mt_points(); ir++) {
      //==                 s[ir] = f[ir] * rad_func(ir, j);
      //==             }
      //==             double ovlp = s.interpolate().integrate(2);
      //==             for (int ir = 0; ir < a.num_mt_points(); ir++) {
      //==                 f[ir] -= ovlp * rad_func(ir, j);
      //==             }
      //==         }
      //==     }
      //==     for (int ir = 0; ir < a.num_mt_points(); ir++) {
      //==         s[ir] = std::pow(f[ir], 2);
      //==     }
      //==     double norm = std::sqrt(s.interpolate().integrate(2));
      //==     printf("orbital: %i, norm: %f\n", i, norm);
      //==     if (norm < 0.05) {
      //==         auto s = lo_to_str(a.lo_descriptor(i));
      //==         printf("local orbital %i is linearly dependent\n", i);
      //==         printf("  l: %i, basis: %s\n", l, s.str().c_str());
      //==         inc[i] = 0;
      //==     } else {
      //==         norm = 1 / norm;
      //==         for (int ir = 0; ir < a.num_mt_points(); ir++) {
      //==             rad_func(ir, i) = f[ir] * norm;
      //==             //atom_class.radial_function(ir, idxrf) = rad_func(ir, i);
      //==         }
      //==     }
      //== }
      
      auto inc = atom_class.check_lo_linear_independence(0.0001);
      dict["lo"] = json::array();
  
      for (int j = 0; j < atom_class.num_lo_descriptors(); j++) {
          auto s = lo_to_str(atom_class.lo_descriptor(j));
          if (!inc[j]) {
              printf("X ");
          } else {
              printf("  ");
              dict["lo"].push_back({{"l", atom_class.lo_descriptor(j).l}, {"basis", json::parse(s)}});
          }
          printf("l: %i, basis: %s\n", atom_class.lo_descriptor(j).l, s.c_str());
      }
  
      dict["free_atom"] = json::object();
      dict["free_atom"]["density"] = fa_rho;
      dict["free_atom"]["radial_grid"] = fa_r;
      
      std::ofstream ofs(a.symbol() + std::string(".json"), std::ofstream::out | std::ofstream::trunc);
      ofs << dict.dump(4);
      ofs.close();
  
      //if (write_to_xml)
      //{
      //    std::string fname = a.symbol() + std::string(".xml");
      //    FILE* fout = fopen(fname.c_str(), "w");
      //    fprintf(fout, "<?xml version=\"1.0\" encoding=\"UTF-8\"?>\n");
      //    fprintf(fout, "<spdb>\n");
      //    fprintf(fout, "  <sp chemicalSymbol=\"%s\" name=\"%s\" z=\"%f\" mass=\"%f\">\n", a.symbol().c_str(), a.name().c_str(), -1.0 * a.zn(), a.mass());
      //    fprintf(fout, "    <muffinTin rmin=\"%e\" radius=\"%f\" rinf=\"%f\" radialmeshPoints=\"%i\"/>\n", 1e-6, 2.0, rinf, 1000);
  
      //    for (int ist = 0; ist < a.num_atomic_levels(); ist++)
      //    {
      //        std::string str_core = (enu[ist] < core_cutoff_energy) ? "true" : "false";
  
      //        fprintf(fout, "      <atomicState n=\"%i\" l=\"%i\" kappa=\"%i\" occ=\"%f\" core=\"%s\"/>\n",
      //                a.atomic_level(ist).n,
      //                a.atomic_level(ist).l,
      //                a.atomic_level(ist).k,
      //                a.atomic_level(ist).occupancy,
      //                str_core.c_str());
      //    }
      //    fprintf(fout, "      <basis>\n");
      //    fprintf(fout, "        <default type=\"lapw\" trialEnergy=\"0.15\" searchE=\"false\"/>\n");
      //    for (int l = 0; l < 4; l++)
      //    {
      //        if (n_v[l].size() == 1)
      //        {
      //            int n = n_v[l][0];
  
      //            fprintf(fout, "        <lo l=\"%i\">\n", l);
      //            fprintf(fout, "          <wf matchingOrder=\"0\" trialEnergy=\"%f\" searchE=\"false\"/>\n", e_nl_v[n][l]);
      //            fprintf(fout, "          <wf matchingOrder=\"1\" trialEnergy=\"%f\" searchE=\"false\"/>\n", e_nl_v[n][l]);
      //            fprintf(fout, "        </lo>\n");
  
      //            fprintf(fout, "        <lo l=\"%i\">\n", l);
      //            fprintf(fout, "          <wf matchingOrder=\"1\" trialEnergy=\"%f\" searchE=\"false\"/>\n", e_nl_v[n][l]);
      //            fprintf(fout, "          <wf matchingOrder=\"2\" trialEnergy=\"%f\" searchE=\"false\"/>\n", e_nl_v[n][l]);
      //            fprintf(fout, "        </lo>\n");
      //            
      //            fprintf(fout, "        <lo l=\"%i\">\n", l);
      //            fprintf(fout, "          <wf matchingOrder=\"2\" trialEnergy=\"%f\" searchE=\"false\"/>\n", e_nl_v[n][l]);
      //            fprintf(fout, "          <wf matchingOrder=\"3\" trialEnergy=\"%f\" searchE=\"false\"/>\n", e_nl_v[n][l]);
      //            fprintf(fout, "        </lo>\n");
      //        }
      //        if (n_v[l].size() > 1)
      //        {
      //            for (size_t i = 0; i < n_v[l].size() - 1; i++)
      //            {
      //                int n = n_v[l][i];
      //                int n1 = n_v[l][i + 1];
  
      //                fprintf(fout, "        <lo l=\"%i\">\n", l);
      //                fprintf(fout, "          <wf matchingOrder=\"0\" trialEnergy=\"%f\" searchE=\"false\"/>\n", e_nl_v[n][l]);
      //                fprintf(fout, "          <wf matchingOrder=\"1\" trialEnergy=\"%f\" searchE=\"false\"/>\n", e_nl_v[n][l]);
      //                fprintf(fout, "        </lo>\n");
  
      //                fprintf(fout, "        <lo l=\"%i\">\n", l);
      //                fprintf(fout, "          <wf matchingOrder=\"1\" trialEnergy=\"%f\" searchE=\"false\"/>\n", e_nl_v[n][l]);
      //                fprintf(fout, "          <wf matchingOrder=\"2\" trialEnergy=\"%f\" searchE=\"false\"/>\n", e_nl_v[n][l]);
      //                fprintf(fout, "        </lo>\n");
      //                
      //                fprintf(fout, "        <lo l=\"%i\">\n", l);
      //                fprintf(fout, "          <wf matchingOrder=\"2\" trialEnergy=\"%f\" searchE=\"false\"/>\n", e_nl_v[n][l]);
      //                fprintf(fout, "          <wf matchingOrder=\"0\" trialEnergy=\"%f\" searchE=\"false\"/>\n", e_nl_v[n1][l]);
      //                fprintf(fout, "        </lo>\n");
      //            }
      //            int n = n_v[l].back();
      //            fprintf(fout, "        <lo l=\"%i\">\n", l);
      //            fprintf(fout, "          <wf matchingOrder=\"0\" trialEnergy=\"%f\" searchE=\"false\"/>\n", e_nl_v[n][l]);
      //            fprintf(fout, "          <wf matchingOrder=\"1\" trialEnergy=\"%f\" searchE=\"false\"/>\n", e_nl_v[n][l]);
      //            fprintf(fout, "        </lo>\n");
      //        }
  
  
  
      //        //if (n_v[l].size())
      //        //    fprintf(fout, "        <custom l=\"%i\" type=\"lapw\" trialEnergy=\"%f\" searchE=\"false\"/>\n", l, e_nl_v[n_v[l].front()][l]);
  
      //        //for (int n = 1; n <= 7; n++)
      //        //{
      //        //    if (nl_v[n][l])
      //        //    {
      //        //        fprintf(fout, "        <lo l=\"%i\">\n", l);
      //        //        fprintf(fout, "          <wf matchingOrder=\"0\" trialEnergy=\"%f\" searchE=\"false\"/>\n", e_nl_v[n][l]);
      //        //        fprintf(fout, "          <wf matchingOrder=\"1\" trialEnergy=\"%f\" searchE=\"false\"/>\n", e_nl_v[n][l]);
      //        //        fprintf(fout, "        </lo>\n");
  
      //        //        //if (e_nl_v[n][l] > -5.0)
      //        //        //{
      //        //            fprintf(fout, "        <lo l=\"%i\">\n", l);
      //        //            fprintf(fout, "          <wf matchingOrder=\"1\" trialEnergy=\"%f\" searchE=\"false\"/>\n", e_nl_v[n][l]);
      //        //            fprintf(fout, "          <wf matchingOrder=\"2\" trialEnergy=\"%f\" searchE=\"false\"/>\n", e_nl_v[n][l]);
      //        //            fprintf(fout, "        </lo>\n");
      //        //        //}
      //        //        
      //        //        //if (e_nl_v[n][l] > -1.0)
      //        //        //{
      //        //            fprintf(fout, "        <lo l=\"%i\">\n", l);
      //        //            fprintf(fout, "          <wf matchingOrder=\"2\" trialEnergy=\"%f\" searchE=\"false\"/>\n", e_nl_v[n][l]);
      //        //            fprintf(fout, "          <wf matchingOrder=\"3\" trialEnergy=\"%f\" searchE=\"false\"/>\n", e_nl_v[n][l]);
      //        //            fprintf(fout, "        </lo>\n");
      //        //        //}
  
      //        //        //fprintf(fout, "        <lo l=\"%i\">\n", l);
      //        //        //fprintf(fout, "          <wf matchingOrder=\"0\" trialEnergy=\"%f\" searchE=\"true\"/>\n", e_nl_v[n][l]);
      //        //        //fprintf(fout, "          <wf matchingOrder=\"1\" trialEnergy=\"%f\" searchE=\"true\"/>\n", e_nl_v[n][l]);
      //        //        //fprintf(fout, "        </lo>\n");
      //        //        //
      //        //        //if (e_nl_v[n][l] < -1.0)
      //        //        //{
      //        //        //    fprintf(fout, "        <lo l=\"%i\">\n", l);
      //        //        //    fprintf(fout, "          <wf matchingOrder=\"0\" trialEnergy=\"0.15\" searchE=\"false\"/>\n");
      //        //        //    fprintf(fout, "          <wf matchingOrder=\"1\" trialEnergy=\"0.15\" searchE=\"false\"/>\n");
      //        //        //    fprintf(fout, "          <wf matchingOrder=\"0\" trialEnergy=\"%f\" searchE=\"true\"/>\n", e_nl_v[n][l]);
      //        //        //    fprintf(fout, "        </lo>\n");
      //        //        //}
      //        //    }
      //        //}
      //    }
      //    for (int l = lmax + 1; l <= lmax + 3; l++)
      //    {
      //        fprintf(fout, "        <lo l=\"%i\">\n", l);
      //        fprintf(fout, "          <wf matchingOrder=\"0\" trialEnergy=\"0.15\" searchE=\"false\"/>\n");
      //        fprintf(fout, "          <wf matchingOrder=\"1\" trialEnergy=\"0.15\" searchE=\"false\"/>\n");
      //        fprintf(fout, "        </lo>\n");
  
      //        fprintf(fout, "        <lo l=\"%i\">\n", l);
      //        fprintf(fout, "          <wf matchingOrder=\"1\" trialEnergy=\"0.15\" searchE=\"false\"/>\n");
      //        fprintf(fout, "          <wf matchingOrder=\"2\" trialEnergy=\"0.15\" searchE=\"false\"/>\n");
      //        fprintf(fout, "        </lo>\n");
      //    }
      //    fprintf(fout, "      </basis>\n");
      //    fprintf(fout, "  </sp>\n");
      //    fprintf(fout, "</spdb>\n");
  
      //    fclose(fout);
      //}
   
  }
  
  int main(int argn, char **argv)
  {
      sirius::initialize(true);
  
      /* handle command line arguments */
      cmd_args args;
      args.register_key("--symbol=", "{string} symbol of a chemical element");
      args.register_key("--type=", "{lo1, lo2, lo3, LO1, LO2} type of local orbital basis");
      args.register_key("--core=", "{double} cutoff for core states: energy (in Ha, if <0), radius (in a.u. if >0)");
      args.register_key("--order=", "{int} order of augmentation");
      args.register_key("--apw_enu=", "{double} default value for APW linearization energies");
      args.register_key("--auto_enu", "allow search of APW linearization energies");
      args.register_key("--xml", "xml output for Exciting code");
      args.register_key("--rel", "use scalar-relativistic solver");
      args.parse_args(argn, argv);
      
      if (argn == 1 || args.exist("help")) {
          printf("\n");
          printf("Atom (L)APW+lo basis generation.\n");
          printf("\n");
          printf("Usage: %s [options]\n", argv[0]);
          args.print_help();
          printf("\n");
          printf("Definition of the local orbital types:\n");
          printf("  lo  : 2nd order local orbitals composed of u(E) and udot(E),\n");
          printf("        where E is the energy of the bound-state level {n,l}\n");
          printf("  LO  : 3rd order local orbitals composed of u(E), udot(E) and u(E1),\n");
          printf("        where E and E1 are the energies of the bound-state levels {n,l} and {n+1,l}\n");
          printf("\n");
          printf("Examples:\n");
          printf("\n");
          printf("  generate default basis for lithium:\n");
          printf("    ./atom --symbol=Li\n"); 
          printf("\n");
          printf("  generate high precision basis for titanium:\n");
          printf("    ./atom --type=lo+LO --symbol=Ti\n"); 
          printf("\n");
          printf("  make all states of iron to be valence:\n");
          printf("    ./atom --core=-1000 --symbol=Fe\n"); 
          printf("\n");
          return 0;
      }
  
      auto symbol = args.value<std::string>("symbol");
  
      double core_cutoff = args.value<double>("core", -10.0);
  
      std::string lo_type = args.value<std::string>("type", "lo1");
  
      int apw_order = args.value<int>("order", 2);
  
      double apw_enu = args.value<double>("apw_enu", 0.15);
  
      bool auto_enu = args.exist("auto_enu");
  
      bool write_to_xml = args.exist("xml");
  
      bool rel = args.exist("rel");
     
      sirius::Simulation_parameters param;
      param.set_lmax_apw(-1);
      Free_atom a = init_atom_configuration(symbol, param);
      
      generate_atom_file(a, core_cutoff, lo_type, apw_order, apw_enu, auto_enu, write_to_xml, rel);
  
      sirius::finalize();
  }
  