/*/*cppimport
<%
import os
include = os.path.join(os.path.dirname(self.filename),'../..')
cfg['include_dirs'] = [include, f'{include}/hgeom/extern']
-cfg['compiler_args'] = ['-std=c++17', '-w', '-Ofast']
cfg['dependencies'] = []
cfg['parallel'] = True
setup_pybind11(cfg)
%>
*/
/** \file */

#include <pybind11/eigen.h>
#include <pybind11/numpy.h>
#include <pybind11/pybind11.h>
#include <pybind11/stl.h>

#include "hgeom/util/global_rng.hpp"
#include "hgeom/util/numeric.hpp"
#include "hgeom/util/pybind_types.hpp"
#include "hgeom/xbin/xbin.hpp"
#include "parallel_hashmap/phmap.h"

namespace py = pybind11;
using namespace py::literals;
using namespace Eigen;

#include <iostream>
#include <unordered_set>

using std::cout;
using std::endl;

namespace hgeom {
namespace geom {

using namespace util;
using namespace Eigen;

template <typename F>
py::tuple expand_xforms_rand(std::vector<py::array_t<F>> gen_in, int depth,
                             int ntrials, F radius, V3<F> cen,
                             F reset_radius_ratio, bool deterministic = false) {
    std::uniform_int_distribution<> randidx(0, gen_in.size() - 1);
    std::mt19937 deterministic_rng = std::mt19937(0);
    using X = X3<F>;
    std::vector<X> gen;
    for (auto g : gen_in) gen.push_back(xform_py_to_X3(g));
    xbin::XformHash_bt24_BCC6<X, uint64_t> binner(0.1654, 1.74597, 107.745);
    ::phmap::parallel_flat_hash_map<uint64_t, uint64_t> uniq_keys;
    Mx<X> frames(depth, ntrials);

    for (int idepth = 0; idepth < depth; ++idepth) {
        for (int itrial = 0; itrial < ntrials; ++itrial) {
            // compute new frame
            X xdelta;
            if (deterministic)
                xdelta = gen[randidx(deterministic_rng)];
            else
                xdelta = gen[randidx(global_rng())];
            X &x = frames(idepth, itrial);
            x = idepth == 0 ? xdelta : xdelta * frames(idepth - 1, itrial);
            F dist2cen = (x.translation() - cen).squaredNorm();
            if (dist2cen > sqr(reset_radius_ratio) * sqr(radius)) {
                x = frames(idepth / 2,
                           itrial); // reset to prior location in bounds
                continue;
            } else if (dist2cen > sqr(radius)) {
                continue; // out of bounds, don't record
            }
            // record frame
            int id = idepth * ntrials + itrial;
            uniq_keys.insert_or_assign(binner.get_key(x), id);
        }
    }
    std::vector<X> frames_uniq_key;
    for (auto &[k, v] : uniq_keys) {
        int idepth = v / ntrials;
        int itrial = v % ntrials;
        frames_uniq_key.push_back(frames(idepth, itrial));
    }
    std::vector<X> frames_uniq;
    for (int i = 0; i < frames_uniq_key.size(); ++i) {
        bool uniq = true;
        auto x = frames_uniq_key[i];
        for (int j = 0; j < i; ++j) {
            auto y = frames_uniq_key[j];
            auto y2x = x * y.inverse();
            Matrix<F, 3, 3> m = y2x.linear();
            if (y2x.translation().squaredNorm() > 0.0001) continue;
            if (AngleAxis<F>(m).angle() < 0.0001) uniq = false;
        }
        if (uniq) frames_uniq.push_back(x);
    }
    Vx<X> ret(frames_uniq.size());
    for (int i = 0; i < frames_uniq.size(); ++i) ret[i] = frames_uniq[i];
    return py::make_tuple(xform_eigen_to_py(ret), 0);
}

// ZOMG...
template <typename F>
py::tuple expand_xforms_2gen_loop(X3<F> G1, X3<F> G2, int N = 5, F maxrad = 9e9,
                                  F maxrad_intermediate = 9e9) {
    // auto xbuf = Vx<X3<F>>(10000);  // reasonable maximum?
    auto xbuf = std::vector<X3<F>>();

    xbin::XformHash_bt24_BCC6<X3<F>, uint64_t> binner(0.979837589843, 10.0,
                                                      1000.0);
    // V4<F> representative(3.78593478659, 7.72308475, 13.234075748, 1);
    X3<F> I = X3<F>::Identity();

    std::unordered_set<uint64_t> seenit;

    // ZOMG...
    int count = 0;
    for (int i0a = 0; i0a < 2; ++i0a) {
        X3<F> x0a((i0a == 1 ? G1 : I) * I);
        for (int i0b = 0; i0b < 2; ++i0b) {
            X3<F> x0b((i0b == 1 ? G2 : I) * x0a);
            for (int i1a = 0; i1a < 2; ++i1a) {
                X3<F> x1a((i1a == 1 ? G1 : I) * x0b);
                for (int i1b = 0; i1b < 2; ++i1b) {
                    X3<F> x1b((i1b == 1 ? G2 : I) * x1a);
                    for (int i2a = 0; i2a < 2; ++i2a) {
                        X3<F> x2a((i2a == 1 ? G1 : I) * x1b);
                        for (int i2b = 0; i2b < 2; ++i2b) {
                            X3<F> x2b((i2b == 1 ? G2 : I) * x2a);
                            for (int i3a = 0; i3a < 2; ++i3a) {
                                X3<F> x3a((i3a == 1 ? G1 : I) * x2b);
                                for (int i3b = 0; i3b < 2; ++i3b) {
                                    X3<F> x3b((i3b == 1 ? G2 : I) * x3a);
                                    for (int i4a = 0; i4a < 2; ++i4a) {
                                        X3<F> x4a((i4a == 1 ? G1 : I) * x3b);
                                        for (int i4b = 0; i4b < 2; ++i4b) {
                                            X3<F> x4b((i4b == 1 ? G2 : I) *
                                                      x4a);
                                            for (int i5a = 0; i5a < 2; ++i5a) {
                                                X3<F> x5a((i5a == 1 ? G1 : I) *
                                                          x4b);
                                                for (int i5b = 0; i5b < 2;
                                                     ++i5b) {
                                                    X3<F> x5b(
                                                        (i5b == 1 ? G2 : I) *
                                                        x5a);
                                                    for (int i6a = 0; i6a < 2;
                                                         ++i6a) {
                                                        X3<F> x6a((i6a == 1
                                                                       ? G1
                                                                       : I) *
                                                                  x5b);
                                                        for (int i6b = 0;
                                                             i6b < 2; ++i6b) {
                                                            X3<F> x6b(
                                                                (i6b == 1 ? G2
                                                                          : I) *
                                                                x6a);
                                                            for (int i7a = 0;
                                                                 i7a < 2;
                                                                 ++i7a) {
                                                                X3<F> x7a(
                                                                    (i7a == 1
                                                                         ? G1
                                                                         : I) *
                                                                    x6b);
                                                                for (int i7b =
                                                                         0;
                                                                     i7b < 2;
                                                                     ++i7b) {
                                                                    X3<F> x7b(
                                                                        (i7b == 1
                                                                             ? G2
                                                                             : I) *
                                                                        x7a);
                                                                    for (
                                                                        int i8a =
                                                                            0;
                                                                        i8a < 2;
                                                                        ++i8a) {
                                                                        X3<F> x8a(
                                                                            (i8a == 1
                                                                                 ? G1
                                                                                 : I) *
                                                                            x7b);
                                                                        for (
                                                                            int i8b =
                                                                                0;
                                                                            i8b <
                                                                            2;
                                                                            ++i8b) {
                                                                            X3<F> x8b(
                                                                                (i8b == 1
                                                                                     ? G2
                                                                                     : I) *
                                                                                x8a);
                                                                            for (
                                                                                int i9a =
                                                                                    0;
                                                                                i9a <
                                                                                2;
                                                                                ++i9a) {
                                                                                X3<F> x9a(
                                                                                    (i9a == 1
                                                                                         ? G1
                                                                                         : I) *
                                                                                    x8b);
                                                                                for (
                                                                                    int i9b =
                                                                                        0;
                                                                                    i9b <
                                                                                    2;
                                                                                    ++i9b) {
                                                                                    X3<F> x9b(
                                                                                        (i9b == 1
                                                                                             ? G2
                                                                                             : I) *
                                                                                        x9a);
                                                                                    for (
                                                                                        int iAa =
                                                                                            0;
                                                                                        iAa <
                                                                                        2;
                                                                                        ++iAa) {
                                                                                        X3<F> xAa(
                                                                                            (i7a == 1
                                                                                                 ? G1
                                                                                                 : I) *
                                                                                            x9b);
                                                                                        for (
                                                                                            int iAb =
                                                                                                0;
                                                                                            iAb <
                                                                                            2;
                                                                                            ++iAb) {
                                                                                            X3<F> xAb(
                                                                                                (i7b == 1
                                                                                                     ? G2
                                                                                                     : I) *
                                                                                                xAa);
                                                                                            for (
                                                                                                int iBa =
                                                                                                    0;
                                                                                                iBa <
                                                                                                2;
                                                                                                ++iBa) {
                                                                                                X3<F> xBa(
                                                                                                    (i8a == 1
                                                                                                         ? G1
                                                                                                         : I) *
                                                                                                    xAb);
                                                                                                for (
                                                                                                    int iBb =
                                                                                                        0;
                                                                                                    iBb <
                                                                                                    2;
                                                                                                    ++iBb) {
                                                                                                    X3<F> xBb(
                                                                                                        (i8b == 1
                                                                                                             ? G2
                                                                                                             : I) *
                                                                                                        xBa);
                                                                                                    for (
                                                                                                        int iCa =
                                                                                                            0;
                                                                                                        iCa <
                                                                                                        2;
                                                                                                        ++iCa) {
                                                                                                        X3<F> xCa(
                                                                                                            (i9a == 1
                                                                                                                 ? G1
                                                                                                                 : I) *
                                                                                                            xBb);
                                                                                                        for (
                                                                                                            int iCb =
                                                                                                                0;
                                                                                                            iCb <
                                                                                                            2;
                                                                                                            ++iCb) {
                                                                                                            X3<F> xCb(
                                                                                                                (i9b == 1
                                                                                                                     ? G2
                                                                                                                     : I) *
                                                                                                                xCa);
                                                                                                            X3<F> x(
                                                                                                                xCb);
                                                                                                            if (x.translation()
                                                                                                                    .squaredNorm() >
                                                                                                                maxrad_intermediate *
                                                                                                                    maxrad_intermediate)
                                                                                                                continue;
                                                                                                            // V4<F> rep =
                                                                                                            // x * representative;
                                                                                                            count++;

                                                                                                            // offest to aviod cell
                                                                                                            // boundary crap (?)
                                                                                                            auto xtmp =
                                                                                                                x;
                                                                                                            xtmp.translation() += V3<
                                                                                                                F>(
                                                                                                                0.001312,
                                                                                                                0.0097834,
                                                                                                                0.0047847529);
                                                                                                            uint64_t bin =
                                                                                                                binner
                                                                                                                    .get_key(
                                                                                                                        xtmp);
                                                                                                            if (seenit
                                                                                                                    .find(
                                                                                                                        bin) ==
                                                                                                                seenit
                                                                                                                    .end()) {
                                                                                                                // if (true) {
                                                                                                                // print(depth, seenit,
                                                                                                                // "new bin ", 0);
                                                                                                                // cout << bin << endl;
                                                                                                                xbuf.push_back(
                                                                                                                    x);
                                                                                                                seenit
                                                                                                                    .insert(
                                                                                                                        bin);
                                                                                                                // print(depth, seenit,
                                                                                                                // "recurse pos2");b
                                                                                                            }
                                                                                                            // bool
                                                                                                            // redundant(false); for
                                                                                                            // (auto frorigin :
                                                                                                            // seenit) {
                                                                                                            // if ((frorigin - rep)
                                                                                                            // .squaredNorm() <
                                                                                                            // 0.0001)
                                                                                                            // redundant = true;
                                                                                                            // }
                                                                                                            // cout << x(0, 0) << '
                                                                                                            // '
                                                                                                            // << x(0, 1) << ' '
                                                                                                            // << seenit.size()
                                                                                                            // << endl;
                                                                                                            // if (redundant)
                                                                                                            // continue;
                                                                                                        }
                                                                                                    }
                                                                                                    if (1 >=
                                                                                                        N)
                                                                                                        goto DONE;
                                                                                                }
                                                                                            }
                                                                                            if (2 >=
                                                                                                N)
                                                                                                goto DONE;
                                                                                        }
                                                                                    }
                                                                                    if (3 >=
                                                                                        N)
                                                                                        goto DONE;
                                                                                }
                                                                            }
                                                                            if (4 >=
                                                                                N)
                                                                                goto DONE;
                                                                        }
                                                                    }
                                                                    if (5 >= N)
                                                                        goto DONE;
                                                                }
                                                            }
                                                            if (6 >= N)
                                                                goto DONE;
                                                        }
                                                    }
                                                    if (7 >= N) goto DONE;
                                                }
                                            }
                                            if (8 >= N) goto DONE;
                                        }
                                    }
                                    if (9 >= N) goto DONE;
                                }
                            }
                            if (10 >= N) goto DONE;
                        }
                    }
                    if (11 >= N) goto DONE;
                }
            }
            if (12 >= N) goto DONE;
        }
    }
DONE:

    Vx<X3<F>> xuniq = Vx<X3<F>>(10000); // reasonable maximum?
    size_t n = 0;
    for (auto x : xbuf) {
        if (x.translation().squaredNorm() > maxrad * maxrad) continue;
        bool redundant(false);
        for (size_t i = 0; i < n; ++i) {
            auto x2 = xuniq[i];
            auto cartdis = (x2.translation() - x.translation()).squaredNorm();
            auto oridis = (x2.linear() - x.linear()).squaredNorm();
            if (cartdis < 1e-6 && oridis < 1e-6) redundant = true;
        }
        if (!redundant) {
            xuniq[n] = x;
            ++n;
        }
    }
    // if (seenit.size() != n)
    // cout << "CPP  N=" << N << " seenit.size() " << seenit.size() << " uniq "
    // <<
    // n
    // << endl;
    assert(seenit.size() >= n);

    auto ret = xform_eigen_to_py(xuniq, n);

    return py::make_tuple(ret, py::make_tuple((int)seenit.size(), count));
} // namespace hgeom

template <typename F>
py::tuple expand_xforms_2gen_loop_pyarray(py::array_t<F> G1, py::array_t<F> G2,
                                          int N = 5, F maxrad = 9e9,
                                          F maxrad_intermediate = 9e9) {
    auto g1 = xform_py_to_X3(G1);
    auto g2 = xform_py_to_X3(G2);
    return expand_xforms_2gen_loop(g1, g2, N, maxrad, maxrad_intermediate);
}

void print(int depth, std::unordered_set<uint64_t> const &seenit,
           std::string msg, bool newline = 1) {
    for (int i = 0; i < 6 - depth; ++i) cout << "  ";
    cout << seenit.size() << " " << msg;
    if (newline) cout << endl;
}

template <typename F, typename Binner>
void expx(std::vector<X3<F>> generators, Binner &binner, F maxrad_sq, X3<F> pos,
          int depth, Vx<X3<F>> &xbuf, std::unordered_set<uint64_t> &seenit,
          int &count) {
    if (!depth) {
        // print(depth, seenit, "depth==0");
        return;
    }
    ++count;

    // print(depth, seenit, "recurse nochange");
    expx(generators, binner, maxrad_sq, pos, depth - 1, xbuf, seenit, count);

    auto g = generators[depth % generators.size()];
    auto pos2 = pos * g;
    if (pos2.translation().squaredNorm() > maxrad_sq) {
        print(depth, seenit, "out of bounds");
        return;
    }
    uint64_t bin = binner.get_key(pos2);
    if (seenit.find(bin) == seenit.end()) {
        print(depth, seenit, "new bin ", 0);
        cout << bin << endl;
        xbuf[seenit.size()] = pos2;
        seenit.insert(bin);
        // print(depth, seenit, "recurse pos2");b
        expx(generators, binner, maxrad_sq, pos2, depth - 1, xbuf, seenit,
             count);
    } else {
        print(depth, seenit, "old bin");
    }
}

// template <typename F>
// py::array_t<F> expand_xforms(py::array_t<F> G1in, py::array_t<F> G2in,
// int N = 5, F maxrad = 9e9) {
// X3<F> G1 = xform_py_to_X3(G1in);
// X3<F> G2 = xform_py_to_X3(G2in);
// std::vector<X3<F>> G = std::vector<X3<F>>{G1, G2};
//
// Vx<X3<F>> xbuf = Vx<X3<F>>(10000);  // reasonable maximum?
//
// xbin::XformHash_bt24_BCC6<X3<F>, uint64_t> binner(0.979837589843, 10.0,
// 1000.0);
// F maxrad_sq = maxrad * maxrad;
// int count;
// std::unordered_set<uint64_t> seenit;
//
// print(N, seenit, "initial recurse");
// expx(G, binner, maxrad_sq, X3<F>::Identity(), N, xbuf, seenit, count);
//
// cout << "seenit " << seenit.size() << endl;
//
// Vx<M4<F>> xout = Vx<M4<F>>(seenit.size());
// std::vector<V3<F>> seenit2;
// for (int i = 0; i < seenit.size(); ++i) {
// X3<F> x(xbuf[i]);
// bool redundant(false);
// for (auto frorigin : seenit2) {
// if ((frorigin - x.translation()).squaredNorm() < 0.0001) {
// redundant = true;
// break;
// }
// }
// if (redundant) continue;
// xout[seenit2.size()] = x.matrix();
// seenit2.push_back(x.translation());
// }
//
// auto ret = xform_eigen_to_py(xout, seenit.size());
// return ret;
// }

template <typename F>
py::tuple expand_xforms(py::array_t<F> generators, int N = 5, F maxrad = 9e9,
                        F maxrad_intermediate = 9e9) {
    auto gen = xform_py_to_eigen(generators);

    // cout << "GEN " << gen.size() << endl;
    // cout << gen[0].translation() << endl;
    // cout << gen[1].translation() << endl;

    if (gen.size() == 2) {
        return expand_xforms_2gen_loop(gen[0], gen[1], N, maxrad,
                                       maxrad_intermediate);
    } else {
        cout << "only support 2 generators currently" << endl;
    }
    return py::make_tuple(py::none(), py::none());
}

PYBIND11_MODULE(_expand_xforms, m) {
    // m.def("expand_xforms_loop", &expand_xforms_loop<double>);

    m.def("expand_xforms", &expand_xforms<double>, "", "generators"_a,
          "depth"_a, "radius"_a = 9e9, "maxrad_intermediate"_a = 9e9);

    m.def("expand_xforms_2gen_loop_pyarray",
          &expand_xforms_2gen_loop_pyarray<double>, "", "g1"_a, "g2"_a,
          "depth"_a, "radius"_a = 9e9, "maxrad_intermediate"_a = 9e9);

    m.def("expand_xforms_rand", &expand_xforms_rand<double>, "", "generators"_a,
          "depth"_a = 1000, "trials"_a = 1000, "radius"_a = 9e9,
          "cen"_a = Vector3d(0, 0, 0), "reset_radius_ratio"_a = 9e9,
          "deterministic"_a = true);
}

} // namespace geom
} // namespace hgeom
