/*
 *  Copyright (C) 2004-2021 Edward F. Valeev
 *
 *  This file is part of Libint.
 *
 *  Libint is free software: you can redistribute it and/or modify
 *  it under the terms of the GNU Lesser General Public License as published by
 *  the Free Software Foundation, either version 3 of the License, or
 *  (at your option) any later version.
 *
 *  Libint 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 Lesser General Public License for more details.
 *
 *  You should have received a copy of the GNU Lesser General Public License
 *  along with Libint.  If not, see <http://www.gnu.org/licenses/>.
 *
 */

#ifndef _libint2_src_bin_libint_cgshellinfo_h_
#define _libint2_src_bin_libint_cgshellinfo_h_

#include <cassert>
#include <utility>
#include <algorithm>

namespace libint2 {

  /// Normalization convention for Cartesian Gaussian shells
  enum class CartesianShellNormalization {
    standard,  // same normalization factor for every function in shell (default)
    uniform    // different normalization factors so that each function has same norm
  };

  namespace detail {
    inline int notxyz(int a, int b) {
      assert(a != b);
      int amax = std::max(a,b);
      int amin = std::min(a,b);
      if (amin == 0 && amax == 1)
        return 2;
      if (amin == 0 && amax == 2)
        return 1;
      if (amin == 1 && amax == 2)
        return 0;
      assert(false); // unreachable
      return -1;
    }

   inline std::pair<int,int> notxyz(int a) {
      switch(a) {
      case 0: return std::make_pair(1,2); break;
      case 1: return std::make_pair(0,2); break;
      case 2: return std::make_pair(0,1); break;
      }
      assert(false); // unreachable
      return std::make_pair(-1,-1);
    }
  }


  template <CGShellOrdering Ord, unsigned int lmax> struct CGShellOrderingGenerator {
    static void compute(int (&cartindex)[lmax+1][lmax+1][lmax+1]);
  };
  template <unsigned int lmax> struct CGShellOrderingGenerator<CGShellOrdering_Standard,lmax> {
    static void compute(int (&cartindex)[lmax+1][lmax+1][lmax+1]) {

      // see cgshell_ordering.h
      for (unsigned int am = 0; am <= lmax; ++am) {
        int count = 0;
        for(int i=(am);(i)>=0;(i)--) {
          for(int j=(am)-(i);(j)>=0;(j)--, ++count) {
            cartindex[am][i][j] = count;
          }
        }
      }

    }
  };
  template <unsigned int lmax> struct CGShellOrderingGenerator<CGShellOrdering_GAMESS,lmax> {
    static void compute(int (&cartindex)[lmax+1][lmax+1][lmax+1]) {

      for (unsigned int am = 0; am <= lmax; ++am) {

        if (am == 0) {
          cartindex[0][0][0] = 0;
          continue;
        }
        /// 12/15/2011 GAMESS ordering for g functions different from what I had assumed before (thanks to Luke Roskop)
        if (am == 4) {
          cartindex[4][4][0] = 0;
          cartindex[4][0][4] = 1;
          cartindex[4][0][0] = 2;
          cartindex[4][3][1] = 3;
          cartindex[4][3][0] = 4;
          cartindex[4][1][3] = 5;
          cartindex[4][0][3] = 6;
          cartindex[4][1][0] = 7;
          cartindex[4][0][1] = 8;
          cartindex[4][2][2] = 9;
          cartindex[4][2][0] = 10;
          cartindex[4][0][2] = 11;
          cartindex[4][2][1] = 12;
          cartindex[4][1][2] = 13;
          cartindex[4][1][1] = 14;
          continue;
        }

        unsigned int current_index = 0;
        unsigned int qn[3] = { 0, 0, 0 };

        const int ammin = ((int) am + 2) / 3;
        for (int am1 = am; am1 >= ammin; --am1) {

          for (int xyz1 = 0; xyz1 < 3; ++xyz1) {

            qn[xyz1] = am1;

            // distribute the remaining quanta according to the following rules

            // "nothing to distribute" is a special case
            if (am - am1 == 0) {
              std::pair<int, int> xyz(detail::notxyz( xyz1));
              qn[xyz.first] = 0;
              qn[xyz.second] = 0;
              cartindex[am][qn[0]][qn[1]] = current_index;
              ++current_index;
            } else {
              int am23 = (int) am - qn[xyz1];
              const int maxam23 = std::min((int) qn[xyz1], am23);
              const int minam23 = (am23 + 1) / 2;
              for (int am2 = maxam23; am2 >= minam23; --am2) {
                const int xyz2min = (am2 == qn[xyz1]) ? xyz1 + 1 : 0;
                for (int xyz2 = xyz2min; xyz2 < 3; ++xyz2) {
                  if (xyz1 == xyz2)
                    continue;
                  qn[xyz2] = am2;
                  const int xyz3 = detail::notxyz(xyz1, xyz2);
                  qn[xyz3] = am23 - am2;
                  if ((qn[xyz3] == qn[xyz1] && xyz3 < xyz1) ||
                      (qn[xyz3] == qn[xyz2] && xyz3 < xyz2)
                     )
                    continue;
                  {
                    cartindex[am][qn[0]][qn[1]] = current_index;
                    ++current_index;
                  }
                }
              }
            }

          }
        }
      }
    }
  };
  template <unsigned int lmax> struct CGShellOrderingGenerator<CGShellOrdering_ORCA,lmax> {
    static void compute(int (&cartindex)[lmax+1][lmax+1][lmax+1]) {

      // identical to GAMESS for s through g functions
      // identical to standard for h and higher
      CGShellOrderingGenerator<CGShellOrdering_Standard,lmax>::compute(cartindex);

      for (unsigned int am = 0; am <= std::min(4u,lmax); ++am) {

        if (am == 0) {
          cartindex[0][0][0] = 0;
          continue;
        }
        if (am == 1) {
          cartindex[1][0][0] = 0;
          cartindex[1][1][0] = 1;
          cartindex[1][0][1] = 2;
          continue;
        }
        if (am == 2) {
          cartindex[2][2][0] = 0;
          cartindex[2][0][2] = 1;
          cartindex[2][0][0] = 2;
          cartindex[2][1][1] = 3;
          cartindex[2][1][0] = 4;
          cartindex[2][0][1] = 5;
          continue;
        }
        if (am == 3) {
          cartindex[3][3][0] = 0;
          cartindex[3][0][3] = 1;
          cartindex[3][0][0] = 2;
          cartindex[3][2][1] = 3;
          cartindex[3][2][0] = 4;
          cartindex[3][1][2] = 5;
          cartindex[3][0][2] = 6;
          cartindex[3][1][0] = 7;
          cartindex[3][0][1] = 8;
          cartindex[3][1][1] = 9;
          continue;
        }
        if (am == 4) {
          cartindex[4][4][0] = 0;
          cartindex[4][0][4] = 1;
          cartindex[4][0][0] = 2;
          cartindex[4][3][1] = 3;
          cartindex[4][3][0] = 4;
          cartindex[4][1][3] = 5;
          cartindex[4][0][3] = 6;
          cartindex[4][1][0] = 7;
          cartindex[4][0][1] = 8;
          cartindex[4][2][2] = 9;
          cartindex[4][2][0] = 10;
          cartindex[4][0][2] = 11;
          cartindex[4][2][1] = 12;
          cartindex[4][1][2] = 13;
          cartindex[4][1][1] = 14;
          continue;
        }
        /*
        if (am == 5) {
          cartindex[5][5][0] = 0;
          cartindex[5][4][1] = 1;
          cartindex[5][4][0] = 2;
          cartindex[5][3][2] = 3;
          cartindex[5][3][1] = 4;
          cartindex[5][3][0] = 5;
          cartindex[5][2][3] = 6;
          cartindex[5][2][2] = 7;
          cartindex[5][2][1] = 8;
          cartindex[5][2][0] = 9;
          cartindex[5][1][4] = 10;
          cartindex[5][1][3] = 11;
          cartindex[5][1][2] = 12;
          cartindex[5][1][1] = 13;
          cartindex[5][1][0] = 14;
          cartindex[5][0][5] = 15;
          cartindex[5][0][4] = 16;
          cartindex[5][0][3] = 17;
          cartindex[5][0][2] = 18;
          cartindex[5][0][1] = 19;
          cartindex[5][0][0] = 20;
          continue;
        }
        if (am == 6) {
          cartindex[6][6][0] = 0;
          cartindex[6][5][1] = 1;
          cartindex[6][5][0] = 2;
          cartindex[6][4][2] = 3;
          cartindex[6][4][1] = 4;
          cartindex[6][4][0] = 5;
          cartindex[6][3][3] = 6;
          cartindex[6][3][2] = 7;
          cartindex[6][3][1] = 8;
          cartindex[6][3][0] = 9;
          cartindex[6][2][4] = 10;
          cartindex[6][2][3] = 11;
          cartindex[6][2][2] = 12;
          cartindex[6][2][1] = 13;
          cartindex[6][2][0] = 14;
          cartindex[6][1][5] = 15;
          cartindex[6][1][4] = 16;
          cartindex[6][1][3] = 17;
          cartindex[6][1][2] = 18;
          cartindex[6][1][1] = 19;
          cartindex[6][1][0] = 20;
          cartindex[6][0][6] = 21;
          cartindex[6][0][5] = 22;
          cartindex[6][0][4] = 23;
          cartindex[6][0][3] = 24;
          cartindex[6][0][2] = 25;
          cartindex[6][0][1] = 26;
          cartindex[6][0][0] = 27;
          continue;
        }
        if (am == 7) {
          cartindex[7][7][0] = 0;
          cartindex[7][6][1] = 1;
          cartindex[7][6][0] = 2;
          cartindex[7][5][2] = 3;
          cartindex[7][5][1] = 4;
          cartindex[7][5][0] = 5;
          cartindex[7][4][3] = 6;
          cartindex[7][4][2] = 7;
          cartindex[7][4][1] = 8;
          cartindex[7][4][0] = 9;
          cartindex[7][3][4] = 10;
          cartindex[7][3][3] = 11;
          cartindex[7][3][2] = 12;
          cartindex[7][3][1] = 13;
          cartindex[7][3][0] = 14;
          cartindex[7][2][5] = 15;
          cartindex[7][2][4] = 16;
          cartindex[7][2][3] = 17;
          cartindex[7][2][2] = 18;
          cartindex[7][2][1] = 19;
          cartindex[7][2][0] = 20;
          cartindex[7][1][6] = 21;
          cartindex[7][1][5] = 22;
          cartindex[7][1][4] = 23;
          cartindex[7][1][3] = 24;
          cartindex[7][1][2] = 25;
          cartindex[7][1][1] = 26;
          cartindex[7][1][0] = 27;
          cartindex[7][0][7] = 28;
          cartindex[7][0][6] = 29;
          cartindex[7][0][5] = 30;
          cartindex[7][0][4] = 31;
          cartindex[7][0][3] = 32;
          cartindex[7][0][2] = 33;
          cartindex[7][0][1] = 34;
          cartindex[7][0][0] = 35;
          continue;
        }
        if (am == 8) {
          cartindex[8][8][0] = 0;
          cartindex[8][7][1] = 1;
          cartindex[8][7][0] = 2;
          cartindex[8][6][2] = 3;
          cartindex[8][6][1] = 4;
          cartindex[8][6][0] = 5;
          cartindex[8][5][3] = 6;
          cartindex[8][5][2] = 7;
          cartindex[8][5][1] = 8;
          cartindex[8][5][0] = 9;
          cartindex[8][4][4] = 10;
          cartindex[8][4][3] = 11;
          cartindex[8][4][2] = 12;
          cartindex[8][4][1] = 13;
          cartindex[8][4][0] = 14;
          cartindex[8][3][5] = 15;
          cartindex[8][3][4] = 16;
          cartindex[8][3][3] = 17;
          cartindex[8][3][2] = 18;
          cartindex[8][3][1] = 19;
          cartindex[8][3][0] = 20;
          cartindex[8][2][6] = 21;
          cartindex[8][2][5] = 22;
          cartindex[8][2][4] = 23;
          cartindex[8][2][3] = 24;
          cartindex[8][2][2] = 25;
          cartindex[8][2][1] = 26;
          cartindex[8][2][0] = 27;
          cartindex[8][1][7] = 28;
          cartindex[8][1][6] = 29;
          cartindex[8][1][5] = 30;
          cartindex[8][1][4] = 31;
          cartindex[8][1][3] = 32;
          cartindex[8][1][2] = 33;
          cartindex[8][1][1] = 34;
          cartindex[8][1][0] = 35;
          cartindex[8][0][8] = 36;
          cartindex[8][0][7] = 37;
          cartindex[8][0][6] = 38;
          cartindex[8][0][5] = 39;
          cartindex[8][0][4] = 40;
          cartindex[8][0][3] = 41;
          cartindex[8][0][2] = 42;
          cartindex[8][0][1] = 43;
          cartindex[8][0][0] = 44;
          continue;
        }
        */
      }

#if 0
      for(int l=0; l<=lmax; ++l) {
        for(int i=0; i<=l; ++i) {
          for(int j=0; j<=l-i; ++j) {
            std::cout << "CGShellOrderingGenerator<CGShellOrdering_ORCA>: cartindex[" << l << "]["
                      << i << "][" << j << "] = " << cartindex[l][i][j] << std::endl;
          }
        }
      }
#endif

    }
  };

  // see http://www.cmbi.ru.nl/molden/molden_format.html
  template <unsigned int lmax> struct CGShellOrderingGenerator<CGShellOrdering_MOLDEN,lmax> {
    static void compute(int (&cartindex)[lmax+1][lmax+1][lmax+1]) {

      for (unsigned int am = 0; am <= 4u; ++am) {

        if (am == 0) {
          cartindex[0][0][0] = 0;
          continue;
        }
        if (am == 1) {
          cartindex[1][1][0] = 0;
          cartindex[1][0][1] = 1;
          cartindex[1][0][0] = 2;
          continue;
        }
        if (am == 2) {
          cartindex[2][2][0] = 0;
          cartindex[2][0][2] = 1;
          cartindex[2][0][0] = 2;
          cartindex[2][1][1] = 3;
          cartindex[2][1][0] = 4;
          cartindex[2][0][1] = 5;
          continue;
        }
        if (am == 3) {
          cartindex[3][3][0] = 0;
          cartindex[3][0][3] = 1;
          cartindex[3][0][0] = 2;
          cartindex[3][1][2] = 3;
          cartindex[3][2][1] = 4;
          cartindex[3][2][0] = 5;
          cartindex[3][1][0] = 6;
          cartindex[3][0][1] = 7;
          cartindex[3][0][2] = 8;
          cartindex[3][1][1] = 9;
          continue;
        }
        if (am == 4) {
          cartindex[4][4][0] = 0;
          cartindex[4][0][4] = 1;
          cartindex[4][0][0] = 2;
          cartindex[4][3][1] = 3;
          cartindex[4][3][0] = 4;
          cartindex[4][1][3] = 5;
          cartindex[4][0][3] = 6;
          cartindex[4][1][0] = 7;
          cartindex[4][0][1] = 8;
          cartindex[4][2][2] = 9;
          cartindex[4][2][0] = 10;
          cartindex[4][0][2] = 11;
          cartindex[4][2][1] = 12;
          cartindex[4][1][2] = 13;
          cartindex[4][1][1] = 14;
          continue;
        }
      }
    }
  };

  template <CGShellOrdering Ord, unsigned int lmax = LIBINT_CARTGAUSS_MAX_AM> struct CGShellOrderingData {

    struct Triple {
      Triple() : i(0), j(0), k(0) {}
      Triple(int ii, int jj, int kk) : i(ii), j(jj), k(kk) {}
      int i, j, k;
    };

    CGShellOrderingData() {
      // compute cartindex
      CGShellOrderingGenerator<Ord,lmax>::compute(cartindex);
      // then use it to compute cartindex_to_ijk
      for(unsigned int l=0; l<=lmax; ++l) {
        for(unsigned int i=0; i<=l; ++i) {
          for(unsigned int j=0; j<=l-i; ++j) {
            const int c = cartindex[l][i][j];
            cartindex_to_ijk[l][c] = Triple(i,j,l-i-j);
          }
        }
      }
    }

    int cartindex[lmax+1][lmax+1][lmax+1];
    Triple cartindex_to_ijk[lmax+1][(lmax+1)*(lmax+2)/2];
  };

  /// provides ordering maps for up to angular momentum lmax and ordering specified by CGShellOrderingSpec
  template <typename OrderingData> struct CGShellInfo {
    // computes cartindex xyz from am i j
    static int cartindex(unsigned int am, int i, int j) {
      return data_.cartindex[am][i][j];
    }
    // computes i j k from cartindex xyz
    template <typename I1, typename I2, typename I3>
    static void cartindex_to_ijk(I1 am,
                                 I2 xyz,
                                 I3& i,
                                 I3& j,
                                 I3& k) {
      const typename OrderingData::Triple& ijk = data_.cartindex_to_ijk[am][xyz];
      i = ijk.i;
      j = ijk.j;
      k = ijk.k;
    }

    private:
      static OrderingData data_;
  };

  template <typename OrderingData> OrderingData CGShellInfo<OrderingData>::data_;

};  // namespace libint2

#endif // header guard
