// Implementation of Discrete Conformal Seamless Similarity Mapping
// along the lines of
// [Campen and Zorin 2017]: "Similarity Maps and Field-Guided T-Splines: a Perfect Couple"
// [Campen and Zorin 2017]: "On Discrete Conformal Seamless Similarity Maps"
//
// Author: Marcel Campen
//
// Version 1.0
// 21 July 2017

#include <Eigen/Sparse>
#include <set>
#include <queue>
#include <vector>
#include <igl/Timer.h>

class Mesh {
public:
  std::vector<int> n; // next halfedge of halfedge
  std::vector<int> to; // to vertex of halfedge
  std::vector<int> f; // face of halfedge
  std::vector<int> h; // one halfedge of face
  std::vector<int> out; // one outgoing halfedge of vertex

  std::vector<double> l; // discrete metric (length per edge)

  int n_halfedges() { return n.size(); }
  int n_edges() { return n_halfedges()/2; }
  int n_faces() { return h.size(); }
  int n_vertices() { return out.size(); }

  int e(int h) { return h/2; }
  int opp(int h) { return (h%2 == 0) ? (h+1) : (h-1); }
  int v0(int h) { return to[opp(h)]; }
  int v1(int h) { return to[h]; }
  int h0(int e) { return e*2; }
  int h1(int e) { return e*2+1; }
  double sign(int h) { return (h%2 == 0) ? 1.0 : -1.0; }

  virtual void init() {};

  virtual bool flip_ccw(int _h)
  {
    int ha = _h;
    int hb = opp(_h);
    int f0 = f[ha];
    int f1 = f[hb];
    if(f0 == f1) return false;
    int h2 = n[ha];
    int h3 = n[h2];
    int h4 = n[hb];
    int h5 = n[h4];
    out[to[hb]] = h4;
    out[to[ha]] = h2;
    f[h4] = f0;
    f[h2] = f1;
    h[f0] = h4;
    h[f1] = h2;
    to[ha] = to[h2];
    to[hb] = to[h4];
    n[h5] = h2;
    n[h3] = h4;
    n[h2] = hb;
    n[h4] = ha;
    n[ha] = h3;
    n[hb] = h5;
    return true;
  }

  virtual void get_mesh(std::vector<int>& _n, // next halfedge of halfedge
                        std::vector<int>& _to, // to vertex of halfedge
                        std::vector<int>& _f, // face of halfedge
                        std::vector<int>& _h, // one halfedge of face
                        std::vector<int>& _out) // one outgoing halfedge of vertex
  {
    _n = n;
    _to = to;
    _f = f;
    _h = h;
    _out = out;
  }

  template<typename T>
  std::vector<T> interpolate(const std::vector<T>& u)
  {
    return u;
  }

  bool is_complex()
  {
    int nh = n_halfedges();
    for(int i = 0; i < nh; i++)
    {
      if(to[i] == to[opp(i)]) return true; //contains loop edge
    }
    int nv = n_vertices();
    for(int i = 0; i < nv; i++)
    {
      std::set<int> onering;
      int h = out[i];
      if(h < 0) continue;
      int k = h;
      do {
        int v = to[k];
        if(onering.find(v) != onering.end()) return true; //contains multi-edges
        onering.insert(v);
        k = n[opp(k)];
      } while(k != h);
    }
    return false;
  }
};

class ConformalSeamlessSimilarityMapping {
public:

  Mesh& m;

  std::vector<double> Theta_hat; //target cone angles per vertex
  std::vector<double> kappa_hat; //target holonomy angles per gamma loop
  std::vector< std::vector<int> > gamma; //directed dual loops, represented by halfedges (the ones adjacent to the earlier triangles in the dual loop)

  const double cot_infty = 1e10;

  int n_s;
  int n_e;
  int n_h;
  int n_f;
  int n_v;

  std::vector<double> xi;
  std::vector<double> delta_xi;
  std::vector<double> cot_alpha;
  std::vector<double> alpha;

  Eigen::SparseMatrix<double> A;
  Eigen::VectorXd b;

  std::string log_dir;
  std::string name;
  double eps = 1e-12;

  ConformalSeamlessSimilarityMapping(Mesh& _m, const std::vector<double>& _Theta_hat, const std::vector<double>& _kappa_hat, std::vector< std::vector<int> >& _gamma, std::string _log_dir="", std::string _name="", double _eps=1e-12) : m(_m), Theta_hat(_Theta_hat), kappa_hat(_kappa_hat), gamma(_gamma), log_dir(_log_dir), name(_name), eps(_eps)
  {
    n_s = gamma.size();
    n_e = m.n_edges();
    n_h = m.n_halfedges();
    n_f = m.n_faces();
    n_v = m.n_vertices();

    xi.resize(n_h, 0.0);
    delta_xi.resize(n_h, 0.0);
    cot_alpha.resize(n_h);
    alpha.resize(n_h);
  }

  void log(const char* c)
  {
    std::cout << c << std::endl;
  }

  void compute_angles() // compute alpha and cot_alpha from scaled edge lengths
  {
    #pragma omp parallel for
    for(int f = 0; f < n_f; f++)
    {
      int hi = m.h[f];
      int hj = m.n[hi];
      int hk = m.n[hj];
      // (following "On Discrete Conformal Seamless Similarity Maps")
      double li = m.l[m.e(hi)] * std::exp(1.0/6.0*(xi[hk]-xi[hj]));
      double lj = m.l[m.e(hj)] * std::exp(1.0/6.0*(xi[hi]-xi[hk]));
      double lk = m.l[m.e(hk)] * std::exp(1.0/6.0*(xi[hj]-xi[hi]));
      // (following "A Cotangent Laplacian for Images as Surfaces")
      double s = (li+lj+lk)/2.0;
      double Aijk4 = 4.0*std::sqrt(std::max(0.0, s*(s-li)*(s-lj)*(s-lk)));
      double Ijk = (-li*li+lj*lj+lk*lk);
      double iJk = (li*li-lj*lj+lk*lk);
      double ijK = (li*li+lj*lj-lk*lk);
      cot_alpha[hi] = Aijk4 == 0.0 ? copysign(cot_infty,Ijk) : (Ijk/Aijk4);
      cot_alpha[hj] = Aijk4 == 0.0 ? copysign(cot_infty,iJk) : (iJk/Aijk4);
      cot_alpha[hk] = Aijk4 == 0.0 ? copysign(cot_infty,ijK) : (ijK/Aijk4);

      alpha[hi] = std::acos(std::min(1.0, std::max(-1.0, Ijk/(2.0*lj*lk))));
      alpha[hj] = std::acos(std::min(1.0, std::max(-1.0, iJk/(2.0*lk*li))));
      alpha[hk] = std::acos(std::min(1.0, std::max(-1.0, ijK/(2.0*li*lj))));
    }
  }

  void setup_b() // system right-hand sid
  {
    b.resize(n_v-1 + n_s + n_f-1);
    b.fill(0.0);

    std::vector<double> Theta(n_v, 0.0);
    std::vector<double> kappa(n_s, 0.0);

    for(int h = 0; h < n_h; h++)
    {
      Theta[m.to[m.n[h]]] += alpha[h];
    }
    #pragma omp parallel for
    for(int r = 0; r < n_v-1; r++)
    {
      b[r] = Theta_hat[r] - Theta[r];
    }
    #pragma omp parallel for
    for(int s = 0; s < n_s; s++)
    {
      kappa[s] = 0.0;
      int loop_size = gamma[s].size();
      for(int si = 0; si < loop_size; si++)
      {
        int h = gamma[s][si];
        int hn = m.n[h];
        int hnn = m.n[hn];
        if(m.opp(hn) == gamma[s][(si+1)%loop_size])
          kappa[s] -= alpha[hnn];
        else if(m.opp(hnn) == gamma[s][(si+1)%loop_size])
          kappa[s] += alpha[hn];
        else std::cerr << "ERROR: loop is broken" << std::endl;
      }
      b[n_v-1+s] = kappa_hat[s] - kappa[s];
    }
  }

  void setup_A() // system matrix
  {
    A.resize(n_v-1 + n_s + n_f-1, n_e);
    int loop_trips = 0;
    for(int i = 0; i < n_s; i++)
      loop_trips += gamma[i].size();

    typedef Eigen::Triplet<double> Trip;
    std::vector<Trip> trips;
    trips.clear();
    trips.resize(n_h*2 + loop_trips + (n_f-1)*3);
    #pragma omp parallel for
    for(int h = 0; h < n_h; h++)
    {
      int v0 = m.v0(h);
      int v1 = m.v1(h);
      if(v0 < n_v-1) trips[h*2] = Trip(v0, m.e(h), m.sign(h)*0.5*cot_alpha[h]);
      if(v1 < n_v-1) trips[h*2+1] = Trip(v1, m.e(h), -m.sign(h)*0.5*cot_alpha[h]);
    }

    int base = n_h*2;
    for(int s = 0; s < n_s; s++)
    {
      int loop_size = gamma[s].size();
      #pragma omp parallel for
      for(int si = 0; si < loop_size; si++)
      {
        int h = gamma[s][si];
        trips[base+si] = Trip(n_v-1 + s, m.e(h), m.sign(h)*0.5*(cot_alpha[h]+cot_alpha[m.opp(h)]));
      }
      base += loop_size;
    }

    #pragma omp parallel for
    for(int f = 0; f < n_f-1; f++)
    {
      int hi = m.h[f];
      int hj = m.n[hi];
      int hk = m.n[hj];
      trips[base+f*3] = Trip(n_v-1 + n_s + f, m.e(hi), m.sign(hi));
      trips[base+f*3+1] = Trip(n_v-1 + n_s + f, m.e(hj), m.sign(hj));
      trips[base+f*3+2] = Trip(n_v-1 + n_s + f, m.e(hk), m.sign(hk));
    }

    A.setFromTriplets(trips.begin(), trips.end());
  }

  double I(int i, int j, int k, double lambda = 0.0)
  {
    return m.l[m.e(i)]*std::exp((-xi[j]-delta_xi[j]*lambda)/2) + m.l[m.e(j)]*std::exp((xi[i]+delta_xi[i]*lambda)/2) - m.l[m.e(k)];
  }

  double firstDegeneracy(int& degen, double lambda)
  {
    bool repeat = true;
    while(repeat)
    {
      repeat = false;
      #pragma omp parallel for
      for(int i = 0; i < n_h; i++)
      {
        int j = m.n[i];
        int k = m.n[j];
        double local_lambda = lambda;
        if(I(i,j,k,local_lambda) < 0.0)
        {
          // root finding (from below) by bracketing bisection
          double lo = 0.0;
          double hi = local_lambda;
          for(int r = 0; r < 100; r++)
          {
            double mid = (lo+hi)*0.5;
            if(I(i,j,k,mid) <= 0.0)
              hi = mid;
            else
              lo = mid;
          }

          #pragma omp critical
          {
            if(lo < lambda)
            {
              lambda = lo;
              degen = k;
              repeat = true;
            }
          }
        }
      }
    }
    return lambda;
  }

  double avg_abs(const Eigen::VectorXd& v)
  {
    double res = 0.0;
    int v_size = v.size();
    for(int i = 0; i < v_size; i++)
      res += std::abs(b[i]);
    return res/v_size;
  }

  double max_abs(const Eigen::VectorXd& v)
  {
    double res = 0.0;
    int v_size = v.size();
    for(int i = 0; i < v_size; i++)
      res = std::max(res, std::abs(b[i]));
    return res;
  }

  void compute_metric()
  {

    igl::Timer timer;
    timer.start();

    // double eps = 1e-12; //if max curvature error below eps: consider converged
    int max_iter = 25; //max full Newton steps
    bool converged = false;

    log("computing angles");
    compute_angles();
    log("setup b");
    setup_b();

    std::vector< std::pair<double,double> > errors;
    int n_flips = 0;

    log("starting Newton");
    int degen = -1;
    while(!converged && max_iter > 0)
    {
      double error = max_abs(b);
      if(degen < 0) errors.push_back( std::pair<double,double>(avg_abs(b), error) );
      if(error <= eps) { converged = true; break; }

      double diff = avg_abs(b);

      log("setup A");
      setup_A();

      log("factorize A");
      Eigen::SparseLU< Eigen::SparseMatrix<double> > chol(A);
      log("solve Ax=b");
      Eigen::VectorXd result = chol.solve(b);
      if(chol.info() != Eigen::Success) { log("factorization failed"); return; }
      log("solved");

      #pragma omp parallel for
      for(int i = 0; i < n_e; i++)
      {
        delta_xi[i*2] = result[i];
        delta_xi[i*2+1] = -result[i];
      }

      log("line search");
      double lambda = 1.0;

      int max_linesearch = 25;
      degen = -1;
      while(true) // line search
      {
        log("  checking for degeneration events");
        double first_degen = firstDegeneracy(degen, lambda);
        if(first_degen < lambda)
        {
          lambda = first_degen;
          std::cout << "    degeneracy at lambda = " << lambda << std::endl;
        }

        log("  checking for improvement");
        std::vector<double> xi_old = xi;

        #pragma omp parallel for
        for(int i = 0; i < n_h; i++)
          xi[i] = xi_old[i] + lambda * delta_xi[i];

        compute_angles();
        setup_b();

        if(lambda == 0.0)
        {
          converged = true;
          break;
        }

        double new_diff = avg_abs(b);
        if(new_diff < diff)
        {
          std::cout << "    OK. (" << diff << " -> " << new_diff << ")" << std::endl;
          break;
        }

        lambda *= 0.5;
        if(max_linesearch-- == 0) lambda = 0.0;
        std::cout << "    reduced to    lambda = " << lambda << std::endl;
      }

      if(degen < 0) max_iter--; //no degeneration event

      if(!converged) //flip edge(s) of degeneracy/ies
      {
        std::set<int> degens;
        if(degen >= 0) degens.insert(m.e(degen));
        #pragma omp parallel for
        for(int i = 0; i < n_h; i++) //check for additional (simultaneous) degeneracies
        {
          int j = m.n[i];
          int k = m.n[j];
          if(I(i,j,k) <= 0.0)
          {
            #pragma omp critical
            {
              degens.insert(m.e(k));
            }
          }
        }
        int n_d = degens.size();
        if(n_d == 1) std::cout << "handling a degeneracy by edge flip" << std::endl;
        else if(n_d > 1) std::cout << "handling " << degens.size() << " degeneracies by edge flips" << std::endl;
        for(std::set<int>::iterator it = degens.begin(); it != degens.end(); it++)
        {
          int e = *it;
          int h = m.h0(e);
          int hl = m.n[h];
          int hr = m.n[m.n[m.opp(h)]];

          int hlu = m.n[hl];
          int hru = m.n[m.n[hr]];
          int ho = m.opp(h);
          int hu = h;

          double angle = alpha[m.n[hl]]+alpha[m.n[m.opp(h)]];
          double a = m.l[m.e(hl)] * std::exp(xi[hl]/2);
          double b = m.l[m.e(hr)] * std::exp(-xi[hr]/2);
          m.l[e] = std::sqrt(a*a + b*b - 2.0*a*b*std::cos(angle)) / std::exp((xi[hl]-xi[hr])/2); //intrinsic flip (law of cosines)

          xi[h] = xi[hl]+xi[hr];
          xi[m.opp(h)] = -xi[h];

          if(!m.flip_ccw(h)) { std::cerr << "ERROR: edge could not be flipped." << std::endl; converged = true; break; };
          n_flips++;

          if(m.l[e] <= 0.0)
          {
            m.l[e] = 1e-20;
            std::cerr << "WARNING: numerical issue: flipped edge had zero length.";
          }

          // adjust gamma loops that contain the flipped edge e
          #pragma omp parallel for
          for(int i = 0; i < n_s; ++i)
          {
            std::vector<int>& li = gamma[i];
            int n = li.size();
            for(int j = 0; j < n; ++j)
            {
              int hij = li[j];
              int hij1 = li[(j+1)%n];
              int hij2 = li[(j+2)%n];

              bool change = true;

              if(hij == hru && hij1 == m.opp(hr)) li.insert(li.begin()+j+1,ho);
              else if(hij == hru && hij1 == hu && hij2 == m.opp(hl)) li[(j+1)%n] = ho;
              else if(hij == hru && hij1 == hu && hij2 == m.opp(hlu)) li.erase(li.begin()+((j+1)%n));

              else if(hij == hr && hij1 == m.opp(hru)) li.insert(li.begin()+j+1,hu);
              else if(hij == hr && hij1 == hu && hij2 == m.opp(hlu)) li[(j+1)%n] = hu;
              else if(hij == hr && hij1 == hu && hij2 == m.opp(hl)) li.erase(li.begin()+((j+1)%n));

              else if(hij == hl && hij1 == m.opp(hlu)) li.insert(li.begin()+j+1,hu);
              else if(hij == hl && hij1 == ho && hij2 == m.opp(hru)) li[(j+1)%n] = hu;
              else if(hij == hl && hij1 == ho && hij2 == m.opp(hr)) li.erase(li.begin()+((j+1)%n));

              else if(hij == hlu && hij1 == m.opp(hl)) li.insert(li.begin()+j+1,ho);
              else if(hij == hlu && hij1 == ho && hij2 == m.opp(hr)) li[(j+1)%n] = ho;
              else if(hij == hlu && hij1 == ho && hij2 == m.opp(hru)) li.erase(li.begin()+((j+1)%n));

              else change = false;

              if(change) // cleanup "cusps" in loop
              {
                n = li.size();
                int j0 = j;
                int j1 = (j+1)%n;
                int j2 = (j+2)%n;
                if(li[j0] == m.opp(li[j1]))
                {
                  if(j1 < j0) std::swap(j0,j1);
                  li.erase(li.begin()+j1);
                  li.erase(li.begin()+j0);
                }
                else if(li[j1] == m.opp(li[j2]))
                {
                  if(j2 < j1) std::swap(j1,j2);
                  li.erase(li.begin()+j2);
                  li.erase(li.begin()+j1);
                }
              }
            }
          }
        }
        if(n_d > 0) //recompute angles after flipping
        {
          compute_angles();
          setup_b();

          //sanity check
          for(int i = 0; i < n_h; i++)
          {
            int j = m.n[i];
            int k = m.n[j];
            double indicator = I(i,j,k);
            if(indicator <= 0.0)
            {
              #pragma omp critical
              {
                if(indicator == 0.0) std::cerr << "WARNING: numerical issue: triangle("<<i<<", "<<j<<", "<<k<<") is degenerate after Newton step." << std::endl;
                if(indicator < 0.0) std::cerr << "ERROR: numerical issue: triangle("<<i<<", "<<j<<", "<<k<<") is violating after Newton step." << std::endl;
                degens.insert(m.e(k));
              }
            }
          }
        }
      }
    }

    double error = max_abs(b);

    if(error > eps) std::cerr << "WARNING: the final max error is larger than desired ("<<error<<")." << std::endl;

    std::cout << "\nSTATISTICS:\n";
    std::cout << "Flips: " << n_flips << std::endl;
    std::cout << "Error Decay: (iter, avg, max)" << std::endl;
    for(size_t i = 0; i < errors.size(); i++)
    {
      std::cout << i << ": " << errors[i].first << "  " << errors[i].second << std::endl;
    }
    std::cout << std::endl;
    if(n_flips > 0)
      std::cout << "HINT: The given mesh m has been modified by edge flips. Get the modified mesh by m.get_mesh(...)" << std::endl;
    if(n_flips > 0 && m.is_complex())
      std::cout << "HINT: The modified mesh is non-simple (e.g. contains a loop edge or multiple edges between a pair of vertices). Beware that many mesh data structures and libraries do not support this appropriately." << std::endl;
    std::cout << std::endl;

    // write n_flip and time spent to log
    auto total_time = timer.getElapsedTime();

    std::ofstream mf;
    mf.open(log_dir+"/summary_similarity.csv", std::ios_base::app);
    std::ifstream nf(log_dir+"/summary_similarity.csv");
    if (nf && nf.peek() == std::ifstream::traits_type::eof() ){
      // if stats file is empty then add column names
      nf.close();
      mf << "name, n_flips, max_error_similarity, Th_hat_preset, time\n" ;
    }

    mf << name << ", " << n_flips << ", " << error <<"," << Theta_hat[0] << ","<< total_time << std::endl;
    mf.close();

  }

  void compute_layout(std::vector<double>& u, std::vector<double>& v) //metric -> parametrization
  {
    std::vector<double> phi(n_h);

    u.resize(n_h);
    v.resize(n_h);

    //set starting point
    int h = 0;
    phi[h] = 0.0;
    u[h] = 0.0;
    v[h] = 0.0;
    h = m.n[h];
    phi[h] = xi[h];
    u[h] = m.l[m.e(h)]*std::exp(phi[h]/2);
    v[h] = 0.0;

    // layout the rest of the mesh by BFS
    std::vector<bool> visited(n_f, false);
    std::queue<int> q;
    q.push(h);
    visited[m.f[h]] = true;
    while(!q.empty())
    {
      h = q.front();
      q.pop();

      int hn = m.n[h];
      int hp = m.n[hn];

      phi[hn] = phi[h] + xi[hn];

      double len = m.l[m.e(hn)] * std::exp((phi[h]+phi[hn])/2);

      double ud = u[hp]-u[h];
      double vd = v[hp]-v[h];
      double d = std::sqrt(ud*ud + vd*vd);
      double co = std::cos(alpha[hp]);
      double si = std::sin(alpha[hp]);

      u[hn] = u[h] + (co*ud + si*vd)*len/d;
      v[hn] = v[h] + (co*vd - si*ud)*len/d;

      int hno = m.opp(hn);
      int hpo = m.opp(hp);
      if(!visited[m.f[hno]])
      {
        visited[m.f[hno]] = true;
        phi[hno] = phi[h];
        phi[m.n[m.n[hno]]] = phi[hn];
        u[hno] = u[h];
        v[hno] = v[h];
        u[m.n[m.n[hno]]] = u[hn];
        v[m.n[m.n[hno]]] = v[hn];
        q.push(hno);
      }
      if(!visited[m.f[hpo]])
      {
        visited[m.f[hpo]] = true;
        phi[hpo] = phi[hn];
        phi[m.n[m.n[hpo]]] = phi[hp];
        u[hpo] = u[hn];
        v[hpo] = v[hn];
        u[m.n[m.n[hpo]]] = u[hp];
        v[m.n[m.n[hpo]]] = v[hp];
        q.push(hpo);
      }
    }
  }

  void compute(std::vector<double>& u, std::vector<double>& v) //main method
  {
    compute_metric();
    compute_layout(u, v);
  }

};


/*
Copyright (c) 2012 David Koes and University of Pittsburgh

Permission is hereby granted, free of charge, to any person
obtaining a copy of this software and associated documentation files
(the "Software"), to deal in the Software without restriction,
including without limitation the rights to use, copy, modify, merge,
publish, distribute, sublicense, and/or sell copies of the Software,
and to permit persons to whom the Software is furnished to do so,
subject to the following conditions:  The above copyright notice and
this permission notice shall be included in all copies or substantial
portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS
BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN
ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
*  */

/* 9/18/2012  dkoes
 * This is a simple little program that uses Shake-Rupley and OpenBabel
 * to calculate the solvent accessible surface area of a molecule.
 * I'm trying to make it fast.
 *
 * If no file is given, assume a pdb on stdin.  Always outputs to stdout.
 *
 */

#include <cmath>
#include <cstdio>
#include <sstream>
#include <string>
#include <openbabel/mol.h>
#include <boost/algorithm/string.hpp>
#include <ANN/ANN.h>

#define MAXVDW (3.0) /* assumed maximum vdw */
using namespace std;
using namespace OpenBabel;
using namespace boost::algorithm;

struct Coord
{
	double x;
	double y;
	double z;

	Coord(): x(0), y(0), z(0) {}

	Coord(double a, double b, double c): x(a), y(b), z(c) {}

	double distSq(const Coord& r) const
	{
		double a = x-r.x;
		double b = y-r.y;
		double c = z-r.z;
		return a*a+b*b+c*c;
	}
};

//represents an atom from a pdb file
class Atom
{
	string line; //full line of pdb file
	Coord position; //coordinates
	double r; //vdw radius

public:
	Atom(): r(0) {}

	Atom(const string& l): line(l), r(0) //initialize from pdb file
	{
		//assume this an atom
		stringstream x(line.substr(30,8));
		stringstream y(line.substr(38,8));
		stringstream z(line.substr(46,8));

		x >> position.x;
		y >> position.y;
		z >> position.z;

		string el = line.substr(76,2);
		trim(el);
		//remove any digits or special characters
		trim_if(el, is_any_of("0123456789+-"));
		unsigned anum = etab.GetAtomicNum(el.c_str());
		r = etab.GetVdwRad(anum);
	}

	//return true if pdb line is an atom
	static bool isAtom(const string& l)
	{
		if(l.substr(0,4) == "ATOM" || l.substr(0,6) == "HETATM")
			return true;
		else
			return false;
	}

	const Coord& coord() const { return position; }
	const double radius() const { return r; }

	void print(FILE *out, double bfactor) const
	{
		if(bfactor >= 100)
			bfactor = 99.999;
		fprintf(out, "%s%6.3f%6.3f%s\n",line.substr(0,54).c_str(),bfactor,r,line.substr(66).c_str());
	}
};

void readAtoms(istream& in, vector<Atom>& mol)
{
	mol.clear();
	string line;
	while(getline(in, line))
	{
		if(Atom::isAtom(line))
			mol.push_back(Atom(line));
	}
}

//create uniformish sphere of n points
void generateSpherePoints(unsigned n, vector<Coord>& points)
{
	points.clear(); points.reserve(n);
    double inc = M_PI * (3 - sqrt(5));
    double offset = 2 / double(n);
    for(unsigned k = 0; k < n; k++)
    {
    	double y = k * offset - 1 + (offset / 2);
        double r = sqrt(1 - y*y);
        double phi = k * inc;
        points.push_back(Coord(cos(phi)*r, y, sin(phi)*r));
    }
}

//find all atoms close the the specified atomi,
//it's faster to do this once per an atom and then brute force the
//collisions with the sphere points
void find_neighbor_indices(vector<Atom>& mol, ANNkd_tree& tree,  double probe, unsigned atomi, vector<unsigned>& neighbor_indices)
{
	neighbor_indices.clear();
	ANNidx idxArray[mol.size()];

	const Coord& avec = mol[atomi].coord();
	double radius = mol[atomi].radius() + probe + probe;

	double maxr = radius+MAXVDW;
	ANNcoord querypt[3] = {avec.x, avec.y, avec.z};
	unsigned numneigh = tree.annkFRSearch(querypt, maxr*maxr, mol.size(), idxArray);

	for (unsigned idx = 0; idx < numneigh; idx++)
	{
		unsigned ai = idxArray[idx];
		if(ai != atomi)
		{
			const Atom& a = mol[ai];
			double distSq = avec.distSq(a.coord());
			double r = radius+a.radius();
			if(distSq < r*r)
				neighbor_indices.push_back(ai);
		}
	}
}

//compute asa of each atom of mol and print as it is computed
void printASA(vector<Atom>& mol, const vector<Coord>& points, ANNkd_tree& tree, double probe, FILE * out)
{
    const double Const = 4.0 * M_PI / double(points.size());
    Coord test_point;
    vector<unsigned> neighbor_indices;
    double total = 0;
	for (unsigned ai = 0, na = mol.size(); ai < na; ai++)
	{
		const Atom& a = mol[ai];
		find_neighbor_indices(mol, tree, probe, ai, neighbor_indices);

		double radius = probe + a.radius();
		unsigned n_accessible_point = 0;

		for(unsigned i = 0, n = points.size(); i < n; i++)
		{
			bool is_accessible = true;
			test_point.x = points[i].x*radius + a.coord().x;
			test_point.y = points[i].y*radius + a.coord().y;
			test_point.z = points[i].z*radius + a.coord().z;

			for(unsigned j = 0, m = neighbor_indices.size(); j < m; j++)
			{
				const Atom& atom_j = mol[neighbor_indices[j]];
				double r = atom_j.radius() + probe;
				double diff_sq = test_point.distSq(atom_j.coord());
				if(diff_sq < r*r)
				{
					is_accessible = false;
					break;
				}
			}
			if(is_accessible)
				n_accessible_point++;
		}

		double area = Const*n_accessible_point*radius*radius;
		total += area;
		a.print(out, area);
	}
	fprintf(out,"END\n");
	fprintf(out,"REMARK Total %f\n",total);
}

int main(int argc, char *argv[])
{
	unsigned n_sphere_point = 960;
	double probe = 1.4;

	vector<Atom> mol;
	if(argc > 1)
	{
		ifstream in(argv[1]);
		if(!in)
		{
			cerr << "Error opening " << argv[1] << endl;
			exit(-1);
		}
		readAtoms(in, mol);
	}
	else
	{
		readAtoms(cin, mol);
	}

	//initialize spatial index
	ANNpointArray atompts = annAllocPts(mol.size(),3);
	for(unsigned i = 0, n = mol.size(); i < n; i++)
	{
		atompts[i][0] = mol[i].coord().x;
		atompts[i][1] = mol[i].coord().y;
		atompts[i][2] = mol[i].coord().z;
	}
	ANNkd_tree tree(atompts, mol.size(), 3);

	vector<Coord> sphpoints;
	generateSpherePoints(n_sphere_point, sphpoints);

	printASA(mol, sphpoints, tree, probe, stdout);

	//clean up memory
	annDeallocPts(atompts);
	annClose();
}


#include <Eigen/Core>
#include <Eigen/Geometry>

#include <opencv2/calib3d.hpp>

#include <vector>

#include "random.h"
#include "so3.h"

#include "affine_problem.h"

struct AffineRegion
{
    Eigen::Vector3d X;
    Eigen::Vector3d u1;
    Eigen::Vector3d u2;
    Eigen::Vector3d u3;
    Eigen::Vector3d u4;
};

Eigen::Matrix2d getRelativeAffineMatrix(AffineRegion Au, AffineRegion Av,
                                       Eigen::Vector3d cu, Eigen::Vector3d cv)
{
    // calculate affine matrix from u image to v image
    // solves A*(u-center_u) = (v-center_v)
    // where v is projected point in v image, center_v is projection of PX in v image
    // where u is projected point in u image, center_u is projection of X in u image
    Eigen::Vector3d center_u3 = Au.X-cu;
    Eigen::Vector2d center_u(center_u3[0]/center_u3[2],center_u3[1]/center_u3[2]);

    Eigen::Vector3d u13 = Au.u1 - cu;
    Eigen::Vector3d u23 = Au.u2 - cu;
    Eigen::Vector3d u33 = Au.u3 - cu;
    Eigen::Vector3d u43 = Au.u4 - cu;

    Eigen::Vector2d u1(u13[0]/u13[2],u13[1]/u13[2]);
    Eigen::Vector2d u2(u23[0]/u23[2],u23[1]/u23[2]);
    Eigen::Vector2d u3(u33[0]/u33[2],u33[1]/u33[2]);
    Eigen::Vector2d u4(u43[0]/u43[2],u43[1]/u43[2]);

    Eigen::Vector3d center_v3 = Av.X-cv;
    Eigen::Vector2d center_v(center_v3[0]/center_v3[2],center_v3[1]/center_v3[2]);

    Eigen::Vector3d v13 = Av.u1 - cv;
    Eigen::Vector3d v23 = Av.u2 - cv;
    Eigen::Vector3d v33 = Av.u3 - cv;
    Eigen::Vector3d v43 = Av.u4 - cv;

    Eigen::Vector2d v1(v13[0]/v13[2],v13[1]/v13[2]);
    Eigen::Vector2d v2(v23[0]/v23[2],v23[1]/v23[2]);
    Eigen::Vector2d v3(v33[0]/v33[2],v33[1]/v33[2]);
    Eigen::Vector2d v4(v43[0]/v43[2],v43[1]/v43[2]);

    std::vector<cv::Point2f> srcPoints(4);
    std::vector<cv::Point2f> dstPoints(4);
    srcPoints[0] = cv::Point2f(u1(0),u1(1));
    srcPoints[1] = cv::Point2f(u2(0),u2(1));
    srcPoints[2] = cv::Point2f(u3(0),u3(1));
    srcPoints[3] = cv::Point2f(u4(0),u4(1));
    dstPoints[0] = cv::Point2f(v1(0),v1(1));
    dstPoints[1] = cv::Point2f(v2(0),v2(1));
    dstPoints[2] = cv::Point2f(v3(0),v3(1));
    dstPoints[3] = cv::Point2f(v4(0),v4(1));
    cv::Mat Hmat = cv::findHomography(srcPoints, dstPoints);
    Eigen::MatrixXf H(3,3);
    for ( int i = 0; i < 3; i++ )
        for ( int j = 0; j < 3; j++ )
            H(i,j) = Hmat.at<double>(i,j);

    float s = H(2,0) * center_v(0) + H(2,1) * center_v(1) + H(2,2) ;

    Eigen::Matrix2d affine_matrix;

    affine_matrix << (H(0,0) - H(2,0) * v1(0))/s ,
                     (H(0,1) - H(2,1) * v1(0))/s,
                     (H(1,0) - H(2,0) * v1(1))/s,
                     (H(1,1) - H(2,1) * v1(1))/s;

    return affine_matrix;
}

Eigen::Vector3d add_noise(Eigen::Vector3d p, Eigen::Vector2d noise)
{
    double unorm = p.norm();
    Eigen::Vector3d uproj = p/p(2);
    uproj.head(2) += noise;
    uproj /= uproj.norm();
    uproj *= unorm;
    if ( p.dot(uproj) < 0 ) uproj = -uproj;

    return uproj;
}

Eigen::Vector2d generate_noise(double noise)
{
    return grand2f()*noise;
}

AffineRegion generateAffineRegion(Eigen::Vector3d X)
{
    const double R = -.01;
    AffineRegion a;
    a.X = X;
    a.u1 = offset(X,-R,-R);
    a.u2 = offset(X,R,-R);
    a.u3 = offset(X,-R,R);
    a.u4 = offset(X,R,R);
    return a;
}

void generateAffineProblem( double trans_mag, double angle, double noise, double affine_noise, AffineProblem &prob, int nrays )
{
    Eigen::Vector3d X;    // 3D point
    Eigen::Vector3d cu;   // camera center for first multi-camera rig
    Eigen::Vector3d cv;   // camera center for second multi-camera rig
    Eigen::Vector3d PX;   // 3D point after R,t transformation
    AffineRay u;          // observations in first multi-camera rig (Pluecker line + affine)
    AffineRay v;          // observations in second multi-camera rig

    AffineRegion A1;      // affine region in first camera
    AffineRegion A2;      // affine region in second camera

    Eigen::Vector3d w = rand3f();  // make random axis for rotation
    w = w/w.norm();                // normalize rotation axis
    double theta = angle*M_PI/180; // make random angle for rotation
    prob.R = so3exp(w*theta);                // random rotation
    prob.t = rand3f();                       // random translation
    while ( prob.t.norm() < 1e-10 ) prob.t = rand3f();  // ensure translation is not all zeros

    prob.ray_pairs.clear(); // clear out ray pairs to be safe

    // generate observations
    for ( size_t i = 0; i < nrays; i++ )
    {
        X = rand3f();                          // sample random point
        while ( X.norm() == 0 ) X = rand3f();  // ensure point is not all zeros
        X(2) = 4.;
        X = X/X.norm();                        // normalize point
        double scale_factor = ((double)rand()/RAND_MAX) * 4 + 4;
        X *= scale_factor;                     // stretch point out to distance in range [4,8]

        A1 = generateAffineRegion(X);     // generate affine region

        cu = rand3f();        // sample random camera locations
        cv = rand3f();        // inter-camera
        //cv = cu;            // intra-camera: camera center is same in both rigs

        PX = prob.R*X+prob.t;    // transform point to second rig's frame
        A2.X = PX;               // transform affine region
        A2.u1 = prob.R*A1.u1+prob.t;
        A2.u2 = prob.R*A1.u2+prob.t;
        A2.u3 = prob.R*A1.u3+prob.t;
        A2.u4 = prob.R*A1.u4+prob.t;

        // affine matrix for first rig is identity
        u.x = X-cu;
        u.c = cu;
        u.A << 1, 0,
               0, 1;

        // calculate relative affine matrix for second rig
        v.x = PX-cv;
        v.c = cv;
        v.A = getRelativeAffineMatrix(A1, A2, cu, cv);

        if ( noise > 0. )    // add noise with requested standard deviation
        {
            // add gaussian noise to center points
            u.x = add_noise(u.x, generate_noise(noise));
            v.x = add_noise(v.x, generate_noise(noise));
        }

        if ( affine_noise > 0. )
        {
            // add gaussian noise to affine parameters
            v.A(0,0) += grandf()*affine_noise;
            v.A(0,1) += grandf()*affine_noise;
            v.A(1,0) += grandf()*affine_noise;
            v.A(1,1) += grandf()*affine_noise;
        }

        // store observation
        prob.ray_pairs.push_back( AffineRayPair( u, v ) );
    }
}


#include "teca_vorticity.h"

#include "teca_cartesian_mesh.h"
#include "teca_array_collection.h"
#include "teca_variant_array.h"
#include "teca_metadata.h"

#include <algorithm>
#include <iostream>
#include <string>
#include <cmath>

#if defined(TECA_HAS_BOOST)
#include <boost/program_options.hpp>
#endif

using std::string;
using std::vector;
using std::cerr;
using std::endl;
using std::cos;

//#define TECA_DEBUG

namespace {

template <typename num_t>
constexpr num_t deg_to_rad() { return num_t(M_PI)/num_t(180); }

template <typename num_t>
constexpr num_t earth_radius() { return num_t(6371.0e3); }

// compute vorticicty
// this is based on TECA 1's calculation, and hence
// assumes fixed mesh spacing. here we add periodic bc in lon
// and apply unit stride vector optimization strategy to loops
template <typename num_t, typename pt_t>
void vorticity(num_t *w, const pt_t *lon, const pt_t *lat,
    const num_t *u, const num_t *v, unsigned long n_lon,
    unsigned long n_lat, bool periodic_lon=true)
{
    size_t n_bytes = n_lat*sizeof(num_t);
    num_t *delta_u = static_cast<num_t*>(malloc(n_bytes));

    // delta lon as a function of latitude
    num_t d_lon = (lon[1] - lon[0]) * deg_to_rad<num_t>() * earth_radius<num_t>();
    for (unsigned long j = 0; j < n_lat; ++j)
        delta_u[j] = d_lon * cos(lat[j] * deg_to_rad<num_t>());

    // delta lat
    num_t delta_v = (lat[1] - lat[0]) * deg_to_rad<num_t>() * earth_radius<num_t>();
    num_t dv = num_t(2)*delta_v;

    unsigned long max_i = n_lon - 1;
    unsigned long max_j = n_lat - 1;

    // vorticity
    for (unsigned long j = 1; j < max_j; ++j)
    {
        unsigned long jj = j*n_lon;
        const num_t *uu_2 = u + jj + n_lon;
        const num_t *uu_0 = u + jj - n_lon;
        const num_t *vv_2 = v + jj + 1;
        const num_t *vv_0 = v + jj - 1;
        num_t *ww = w + jj;
        num_t du = num_t(2)*delta_u[j];

        for (unsigned long i = 1; i < max_i; ++i)
        {
            ww[i] = (vv_2[i] - vv_0[i]) / du -
                    (uu_2[i] - uu_0[i]) / dv ;
        }
    }

    if (periodic_lon)
    {
        // periodic in longitude
        for (unsigned long j = 1; j < max_j; ++j)
        {
            unsigned long jj = j*n_lon;
            const num_t *uu_2 = u + jj + n_lon;
            const num_t *uu_0 = u + jj - n_lon;
            const num_t *vv_2 = v + jj + 1;
            const num_t *vv_0 = v + jj + max_i;
            num_t *ww = w + jj;
            num_t du = num_t(2)*delta_u[j];

            ww[0] = (vv_2[0] - vv_0[0]) / du -
                    (uu_2[0] - uu_0[0]) / dv ;
        }

        for (unsigned long j = 1; j < max_j; ++j)
        {
            unsigned long jj = j*n_lon;
            const num_t *uu_2 = u + jj + max_i + n_lon;
            const num_t *uu_0 = u + jj + max_i - n_lon;
            const num_t *vv_2 = v + jj;
            const num_t *vv_0 = v + jj + max_i - 1;
            num_t *ww = w + jj + max_i;
            num_t du = num_t(2)*delta_u[j];

            ww[0] = (vv_2[0] - vv_0[0]) / du -
                    (uu_2[0] - uu_0[0]) / dv ;
        }
    }
    else
    {
        // zero it out
        for (unsigned long j = 1; j < max_j; ++j)
            w[j*n_lon] = num_t();

        for (unsigned long j = 1; j < max_j; ++j)
            w[j*n_lon + max_i] = num_t();
    }

    // extend values into lat boundaries
    num_t *dest = w;
    num_t *src = w + n_lon;
    for (unsigned long i = 0; i < n_lon; ++i)
        dest[i] = src[i+n_lon];

    dest = w + max_j*n_lon;
    src = dest - n_lon;
    for (unsigned long i = 0; i < n_lon; ++i)
        dest[i] = src[i];

    free(delta_u);

    return;
}
};

// --------------------------------------------------------------------------
teca_vorticity::teca_vorticity() :
    component_0_variable(), component_1_variable(),
    vorticity_variable("vorticity")
{
    this->set_number_of_input_connections(1);
    this->set_number_of_output_ports(1);
}

// --------------------------------------------------------------------------
teca_vorticity::~teca_vorticity()
{}

#if defined(TECA_HAS_BOOST)
// --------------------------------------------------------------------------
void teca_vorticity::get_properties_description(
    const string &prefix, options_description &global_opts)
{
    options_description opts("Options for "
        + (prefix.empty()?"teca_vorticity":prefix));

    opts.add_options()
        TECA_POPTS_GET(std::string, prefix, component_0_variable,
            "array containg lon component of the vector")
        TECA_POPTS_GET(std::string, prefix, component_1_variable,
            "array containg lat component of the vector")
        TECA_POPTS_GET(std::string, prefix, vorticity_variable,
            "array to store the computed vorticity in")
        ;

    global_opts.add(opts);
}

// --------------------------------------------------------------------------
void teca_vorticity::set_properties(
    const string &prefix, variables_map &opts)
{
    TECA_POPTS_SET(opts, std::string, prefix, component_0_variable)
    TECA_POPTS_SET(opts, std::string, prefix, component_1_variable)
    TECA_POPTS_SET(opts, std::string, prefix, vorticity_variable)
}
#endif

// --------------------------------------------------------------------------
std::string teca_vorticity::get_component_0_variable(
    const teca_metadata &request)
{
    std::string comp_0_var = this->component_0_variable;

    if (comp_0_var.empty() &&
        request.has("teca_vorticity::component_0_variable"))
            request.get("teca_vorticity::component_0_variable", comp_0_var);

    return comp_0_var;
}

// --------------------------------------------------------------------------
std::string teca_vorticity::get_component_1_variable(
    const teca_metadata &request)
{
    std::string comp_1_var = this->component_1_variable;

    if (comp_1_var.empty() &&
        request.has("teca_vorticity::component_1_variable"))
            request.get("teca_vorticity::component_1_variable", comp_1_var);

    return comp_1_var;
}

// --------------------------------------------------------------------------
std::string teca_vorticity::get_vorticity_variable(
    const teca_metadata &request)
{
    std::string vort_var = this->vorticity_variable;

    if (vort_var.empty())
    {
        if (request.has("teca_vorticity::vorticity_variable"))
            request.get("teca_vorticity::vorticity_variable", vort_var);
        else
            vort_var = "vorticity";
    }

    return vort_var;
}

// --------------------------------------------------------------------------
teca_metadata teca_vorticity::get_output_metadata(
    unsigned int port,
    const std::vector<teca_metadata> &input_md)
{
#ifdef TECA_DEBUG
    cerr << teca_parallel_id()
        << "teca_vorticity::get_output_metadata" << endl;
#endif
    (void)port;

    // add in the array we will generate
    teca_metadata out_md(input_md[0]);
    out_md.append("variables", this->vorticity_variable);

    return out_md;
}

// --------------------------------------------------------------------------
std::vector<teca_metadata> teca_vorticity::get_upstream_request(
    unsigned int port,
    const std::vector<teca_metadata> &input_md,
    const teca_metadata &request)
{
    (void)port;
    (void)input_md;

    vector<teca_metadata> up_reqs;

    // get the name of the arrays we need to request
    std::string comp_0_var = this->get_component_0_variable(request);
    if (comp_0_var.empty())
    {
        TECA_ERROR("component 0 array was not specified")
        return up_reqs;
    }

    std::string comp_1_var = this->get_component_1_variable(request);
    if (comp_1_var.empty())
    {
        TECA_ERROR("component 0 array was not specified")
        return up_reqs;
    }

    // copy the incoming request to preserve the downstream
    // requirements and add the arrays we need
    teca_metadata req(request);

    std::set<std::string> arrays;
    if (req.has("arrays"))
        req.get("arrays", arrays);

    arrays.insert(this->component_0_variable);
    arrays.insert(this->component_1_variable);

    // capture the array we produce
    arrays.erase(this->get_vorticity_variable(request));

    // update the request
    req.set("arrays", arrays);

    // send it up
    up_reqs.push_back(req);
    return up_reqs;
}

// --------------------------------------------------------------------------
const_p_teca_dataset teca_vorticity::execute(
    unsigned int port,
    const std::vector<const_p_teca_dataset> &input_data,
    const teca_metadata &request)
{
#ifdef TECA_DEBUG
    cerr << teca_parallel_id()
        << "teca_vorticity::execute" << endl;
#endif
    (void)port;

    // get the input mesh
    const_p_teca_cartesian_mesh in_mesh
        = std::dynamic_pointer_cast<const teca_cartesian_mesh>(input_data[0]);

    if (!in_mesh)
    {
        TECA_ERROR("teca_cartesian_mesh is required")
        return nullptr;
    }

    // get component 0 array
    std::string comp_0_var = this->get_component_0_variable(request);

    if (comp_0_var.empty())
    {
        TECA_ERROR("component_0_variable was not specified")
        return nullptr;
    }

    const_p_teca_variant_array comp_0
        = in_mesh->get_point_arrays()->get(comp_0_var);

    if (!comp_0)
    {
        TECA_ERROR("requested array \"" << comp_0_var << "\" not present.")
        return nullptr;
    }

    // get component 1 array
    std::string comp_1_var = this->get_component_1_variable(request);

    if (comp_1_var.empty())
    {
        TECA_ERROR("component_1_variable was not specified")
        return nullptr;
    }

    const_p_teca_variant_array comp_1
        = in_mesh->get_point_arrays()->get(comp_1_var);

    if (!comp_1)
    {
        TECA_ERROR("requested array \"" << comp_1_var << "\" not present.")
        return nullptr;
    }

    // get the input coordinate arrays
    const_p_teca_variant_array lon = in_mesh->get_x_coordinates();
    const_p_teca_variant_array lat = in_mesh->get_y_coordinates();

    if (!lon || !lat)
    {
        TECA_ERROR("lat lon mesh cooridinates not present.")
        return nullptr;
    }

    // allocate the output array
    p_teca_variant_array vort = comp_0->new_instance();
    vort->resize(comp_0->size());

    // compute vorticity
    NESTED_TEMPLATE_DISPATCH_FP(
        const teca_variant_array_impl,
        lon.get(), 1,

        const NT1 *p_lon = dynamic_cast<const TT1*>(lon.get())->get();
        const NT1 *p_lat = dynamic_cast<const TT1*>(lat.get())->get();

        NESTED_TEMPLATE_DISPATCH_FP(
            teca_variant_array_impl,
            vort.get(), 2,

            const NT2 *p_comp_0 = dynamic_cast<const TT2*>(comp_0.get())->get();
            const NT2 *p_comp_1 = dynamic_cast<const TT2*>(comp_1.get())->get();
            NT2 *p_vort = dynamic_cast<TT2*>(vort.get())->get();

            ::vorticity(p_vort, p_lon, p_lat,
                p_comp_0, p_comp_1, lon->size(), lat->size());
            )
        )

    // create the output mesh, pass everything through, and
    // add the vorticity array
    p_teca_cartesian_mesh out_mesh = teca_cartesian_mesh::New();

    out_mesh->shallow_copy(
        std::const_pointer_cast<teca_cartesian_mesh>(in_mesh));

    out_mesh->get_point_arrays()->append(
        this->get_vorticity_variable(request), vort);

    return out_mesh;
}


/*!
 *
 * \file affine_transform.hpp
 *
 * \brief File containing n-dimensional affine transformation functions
 *        that can be combined with various interpolation and boundary
 *        functions.
 *
 *
 */
#pragma once
#include <array>
#include <tuple>

#include <Eigen/Dense>
#include <pybind11/eigen.h>
#include <pybind11/numpy.h>
#include <pybind11/pybind11.h>

#include "interpolation.hpp"

/*! \brief Namespace containing affine transform functionality
 *         for transforming n-dimnesional data.
 */
namespace affine_transform
{
/*! \brief Namespace containing implementation details for the affine
 *         transformation functionality.
 */
namespace detail
{
/*! \brief Inner For loop of the transform function.
 *
 *  Loops over the dimensions of the output image and fills it
 *  using the given interpolation and boundary functions to look
 *  up values in the given input image.
 *
 *  @param[in] point            The point from which to start to loop
 *                              in the input image coordinate system
 *  @param[in] input_image      The image from which to extract values
 *  @param[in,out] output_image The image to write to
 *  @param[in] dx               The vectors along which to iterate to
 *                              fill the output image
 *  @param[in,out] chunk        A temporary memory object which is used
 *                              to call the interpolation function
 *  @param[in] background_value A background value to use when accessing
 *                              points outside the input image. (might get
 *                              ignored by some boundary functions)
 *  @param[in] begin            The starting index of the first dimension
 *                              (can be used to parallelise the transform)
 *  @param[in] end              The end index of the first dimension
 *                              (can be used to parallelise the transform)
 *  @param[in] xs               Loop variables used to access the n-dimensional
 *                              data
 *
 *  @tparam Dim                 The dimensions of the images
 *  @tparam T                   The data type of the images
 *  @tparam Func                The interpolation order function
 *  @tparam BoundaryFunc        The boundary function to use
 *  @tparam Xs                  The loop indices, will be int
 */
template <int Dim, typename T, typename Func, typename BoundaryFunc,
          typename... Xs>
constexpr void transform_loop(
    Eigen::Matrix<double, Dim, 1> point,
    const pybind11::detail::unchecked_reference<T, Dim>& input_image,
    pybind11::detail::unchecked_mutable_reference<T, Dim>& output_image,
    const std::array<Eigen::Matrix<double, Dim, 1>, Dim>& dx,
    interpolation::Data<Func, Dim>& chunk, T background_value, int begin,
    int end, Xs... xs)
{
    for (int i = begin; i < end; ++i)
    {
        /* We are in the inner-most loop*/
        if constexpr (Dim == sizeof...(xs) + 1)
        {
            auto x_lower = std::array<int, Dim>{};
            auto x_relative = std::array<double, Dim>{};
            for (size_t l = 0; l < Dim; ++l)
            {
                x_lower[l] = point(l) - (point(l) < 0);
                x_relative[l] = point(l) - x_lower[l];
            }

            interpolation::extract<Func, BoundaryFunc, Dim>(
                chunk, input_image, x_lower, background_value);

            auto interpolate = [&chunk](auto... args) {
                return interpolation::apply_func(chunk, args...);
            };

            output_image(xs..., i) = std::apply(interpolate, x_relative);
        }
        /* Start more loops to iterate over the N-dimensional data*/
        else
        {
            transform_loop<Dim, T, Func, BoundaryFunc>(
                point, input_image, output_image, dx, chunk, background_value,
                0, output_image.shape(sizeof...(xs) + 1), xs..., i);
        }
        point += dx[sizeof...(xs)];
    }
}
}; // namespace detail

/*! \brief Extracts an image from a given one with a given
 *         coordinate system.
 *
 *  @param[in] origin            The origin of the coordinate system for the
 *                               resulting image in the coordinate system of
 *                               the input system.
 *  @param[in] dx                The vectors of the coordinate system of the
 *                               resulting image in the coordinate system of the
 *                               input system.
 *  @param[in] input_image       The image from which to extract data
 *  @param[in,out] output_image  The image in which to store the results
 *  @param[in] background_value  The value to use for points outside the input
 *                               image's domain. Might be ignored by some
 *                               boundary functions.
 *
 *  @tparam Dim                  The dimensionality of the image data
 *  @tparam T                    The datatype of the image data
 *  @tparam Func                 The interpolation order to use
 *  @tparam BoundaryFunc         The boundary function to use
 */
template <int Dim, typename T, template <typename> typename Func,
          typename BoundaryFunc>
void transform(const Eigen::Matrix<double, Dim, 1>& origin,
               const std::array<Eigen::Matrix<double, Dim, 1>, Dim>& dx,
               const pybind11::array_t<T>& input_image,
               pybind11::array_t<T>& output_image, T background_value)
{
    auto input = input_image.template unchecked<Dim>();
    auto output = output_image.template mutable_unchecked<Dim>();

#pragma omp parallel
    {
        typedef Func<T> _Func;

        interpolation::Data<_Func, Dim> chunk;

        int x_len = output.shape(0) / omp_get_num_threads();
        int x_start = x_len * omp_get_thread_num();
        int x_end = x_start + x_len;

        Eigen::Matrix<double, Dim, 1> local_origin = origin + x_start * dx[0];
        if (omp_get_thread_num() == omp_get_num_threads() - 1)
        {
            x_end = output.shape(0);
        }
        detail::transform_loop<Dim, T, _Func, BoundaryFunc>(
            local_origin, input, output, dx, chunk, background_value, x_start,
            x_end);
    }
}

}; // namespace affine_transform


//  (C) Copyright John Maddock 2018.
//  Use, modification and distribution are subject to the
//  Boost Software License, Version 1.0. (See accompanying file
//  LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt)

#include <boost/math/tools/fraction.hpp>
#include <iostream>
#include <complex>

//[golden_ratio_1
template <class T>
struct golden_ratio_fraction
{
   typedef T result_type;

   result_type operator()()
   {
      return 1;
   }
};
//]

//[cf_tan_fraction
template <class T>
struct tan_fraction
{
private:
   T a, b;
public:
   tan_fraction(T v)
      : a(-v * v), b(-1)
   {}

   typedef std::pair<T, T> result_type;

   std::pair<T, T> operator()()
   {
      b += 2;
      return std::make_pair(a, b);
   }
};
//]
//[cf_tan
template <class T>
T tan(T a)
{
   tan_fraction<T> fract(a);
   return a / continued_fraction_b(fract, std::numeric_limits<T>::epsilon());
}
//]
//[cf_expint_fraction
template <class T>
struct expint_fraction
{
   typedef std::pair<T, T> result_type;
   expint_fraction(unsigned n_, T z_) : b(z_ + T(n_)), i(-1), n(n_) {}
   std::pair<T, T> operator()()
   {
      std::pair<T, T> result = std::make_pair(-static_cast<T>((i + 1) * (n + i)), b);
      b += 2;
      ++i;
      return result;
   }
private:
   T b;
   int i;
   unsigned n;
};
//]
//[cf_expint
template <class T>
inline std::complex<T> expint_as_fraction(unsigned n, std::complex<T> const& z)
{
   boost::uintmax_t max_iter = 1000;
   expint_fraction<std::complex<T> > f(n, z);
   std::complex<T> result = boost::math::tools::continued_fraction_b(
      f,
      std::complex<T>(std::numeric_limits<T>::epsilon()),
      max_iter);
   result = exp(-z) / result;
   return result;
}
//]
//[cf_upper_gamma_fraction
template <class T>
struct upper_incomplete_gamma_fract
{
private:
   typedef typename T::value_type scalar_type;
   T z, a;
   int k;
public:
   typedef std::pair<T, T> result_type;

   upper_incomplete_gamma_fract(T a1, T z1)
      : z(z1 - a1 + scalar_type(1)), a(a1), k(0)
   {
   }

   result_type operator()()
   {
      ++k;
      z += scalar_type(2);
      return result_type(scalar_type(k) * (a - scalar_type(k)), z);
   }
};
//]
//[cf_gamma_Q
template <class T>
inline std::complex<T> gamma_Q_as_fraction(const std::complex<T>& a, const std::complex<T>& z)
{
   upper_incomplete_gamma_fract<std::complex<T> > f(a, z);
   std::complex<T> eps(std::numeric_limits<T>::epsilon());
   return pow(z, a) / (exp(z) *(z - a + T(1) + boost::math::tools::continued_fraction_a(f, eps)));
}
//]

int main()
{
   using namespace boost::math::tools;

   //[cf_gr
   golden_ratio_fraction<double> func;
   double gr = continued_fraction_a(
      func,
      std::numeric_limits<double>::epsilon());
   std::cout << "The golden ratio is: " << gr << std::endl;
   //]

   std::cout << tan(0.5) << std::endl;

   std::complex<double> arg(3, 2);
   std::cout << expint_as_fraction(5, arg) << std::endl;

   std::complex<double> a(3, 3), z(3, 2);
   std::cout << gamma_Q_as_fraction(a, z) << std::endl;

   return 0;
}


/* Copyright (C) 2012-2019 IBM Corp.
 * This program is Licensed under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance
 * with the License. You may obtain a copy of the License at
 *   http://www.apache.org/licenses/LICENSE-2.0
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License. See accompanying LICENSE file.
 */
#include <NTL/ZZ.h>
#include <algorithm>
#include <complex>

#include "norms.h"
#include "helib.h"
#include "debugging.h"
#include "ArgMap.h"

NTL_CLIENT
using namespace helib;

bool verbose=false;

// Compute the L-infinity distance between two vectors
double calcMaxDiff(const vector<cx_double>& v1,
                   const vector<cx_double>& v2){

  if(lsize(v1)!=lsize(v2))
    NTL::Error("Vector sizes differ.\nFAILED\n");

  double maxDiff = 0.0;
  for (long i=0; i<lsize(v1); i++) {
    double diffAbs = std::abs(v1[i]-v2[i]);
    if (diffAbs > maxDiff)
      maxDiff = diffAbs;
  }

  return maxDiff;
}
// Compute the max relative difference between two vectors
double calcMaxRelDiff(const vector<cx_double>& v1,
                   const vector<cx_double>& v2)
{
    if(lsize(v1)!=lsize(v2))
        NTL::Error("Vector sizes differ.\nFAILED\n");

    // Compute the largest-magnitude value in the vector
    double maxAbs = 0.0;
    for (auto& x : v1) {
        if (std::abs(x) > maxAbs)
            maxAbs = std::abs(x);
    }
    if (maxAbs<1e-10)
        maxAbs = 1e-10;

    double maxDiff = 0.0;
    for (long i=0; i<lsize(v1); i++) {
        double relDiff = std::abs(v1[i]-v2[i]) / maxAbs;
        if (relDiff > maxDiff)
            maxDiff = relDiff;
    }

    return maxDiff;
}

inline bool cx_equals(const vector<cx_double>& v1,
                      const vector<cx_double>& v2,
                      double epsilon)
{
  return (calcMaxRelDiff(v1,v2) < epsilon);
}

void testBasicArith(const FHEPubKey& publicKey,
                    const FHESecKey& secretKey,
                    const EncryptedArrayCx& ea, double epsilon);
void testComplexArith(const FHEPubKey& publicKey,
                      const FHESecKey& secretKey,
                      const EncryptedArrayCx& ea, double epsilon);
void testRotsNShifts(const FHEPubKey& publicKey,
                     const FHESecKey& secretKey,
                     const EncryptedArrayCx& ea, double epsilon);

void debugCompare(const EncryptedArrayCx& ea, const FHESecKey& sk,
        vector<cx_double>& p, const Ctxt& c, double epsilon)
{
  vector<cx_double> pp;
  ea.decrypt(c, sk, pp);
  std::cout << "    relative-error="<<calcMaxRelDiff(p,pp)
            << ", absolute-error="<<calcMaxRelDiff(p,pp)<<endl;
//  if (!cx_equals(pp, p, epsilon)) {
//    std::cout << "oops:\n"; std::cout << p << "\n";
//    std::cout << pp << "\n";
//    exit(0);
//  }
}

void negateVec(vector<cx_double>& p1)
{
  for (auto& x: p1) x = -x;
}
void add(vector<cx_double>& to, const vector<cx_double>& from)
{
  if (to.size() < from.size())
    to.resize(from.size(), 0);
  for (long i=0; i<from.size(); i++) to[i] += from[i];
}
void sub(vector<cx_double>& to, const vector<cx_double>& from)
{
  if (to.size() < from.size())
    to.resize(from.size(), 0);
  for (long i=0; i<from.size(); i++) to[i] -= from[i];
}
void mul(vector<cx_double>& to, const vector<cx_double>& from)
{
  if (to.size() < from.size())
    to.resize(from.size(), 0);
  for (long i=0; i<from.size(); i++) to[i] *= from[i];
}
void rotate(vector<cx_double>& p, long amt)
{
  long sz = p.size();
  vector<cx_double> tmp(sz);
  for (long i=0; i<sz; i++)
    tmp[((i+amt)%sz +sz)%sz] = p[i];
  p = tmp;
}

/************** Each round consists of the following:
1. c1.multiplyBy(c0)
2. c0 += random constant
3. c2 *= random constant
4. tmp = c1
5. ea.rotate(tmp, random amount in [-nSlots/2, nSlots/2])
6. c2 += tmp
7. ea.rotate(c2, random amount in [1-nSlots, nSlots-1])
8. c1.negate()
9. c3.multiplyBy(c2)
10. c0 -= c3
**************/
void testGeneralOps(const FHEPubKey& publicKey, const FHESecKey& secretKey,
                    const EncryptedArrayCx& ea, double epsilon,
                    long nRounds)
{
  long nslots = ea.size();
  char buffer[32];

  vector<cx_double> p0, p1, p2, p3;
  ea.random(p0);
  ea.random(p1);
  ea.random(p2);
  ea.random(p3);

  Ctxt c0(publicKey), c1(publicKey), c2(publicKey), c3(publicKey);
  ea.encrypt(c0, publicKey, p0, /*size=*/1.0);
  ea.encrypt(c1, publicKey, p1, /*size=*/1.0);
  ea.encrypt(c2, publicKey, p2, /*size=*/1.0);
  ea.encrypt(c3, publicKey, p3, /*size=*/1.0);

  resetAllTimers();
  FHE_NTIMER_START(Circuit);

  for (long i = 0; i < nRounds; i++) {

    if (verbose) std::cout << "*** round " << i << "..."<<endl;

     long shamt = RandomBnd(2*(nslots/2) + 1) - (nslots/2);
                  // random number in [-nslots/2..nslots/2]
     long rotamt = RandomBnd(2*nslots - 1) - (nslots - 1);
                  // random number in [-(nslots-1)..nslots-1]

     // two random constants
     vector<cx_double> const1, const2;
     ea.random(const1);
     ea.random(const2);

     ZZX const1_poly, const2_poly;
     ea.encode(const1_poly, const1, /*size=*/1.0);
     ea.encode(const2_poly, const2, /*size=*/1.0);

     mul(p1, p0);     // c1.multiplyBy(c0)
     c1.multiplyBy(c0);
     if (verbose) {
       CheckCtxt(c1, "c1*=c0");
       debugCompare(ea, secretKey, p1, c1, epsilon);
     }

     add(p0, const1); // c0 += random constant
     c0.addConstant(const1_poly);
     if (verbose) {
       CheckCtxt(c0, "c0+=k1");
       debugCompare(ea, secretKey, p0, c0, epsilon);
     }
     mul(p2, const2); // c2 *= random constant
     c2.multByConstant(const2_poly);
     if (verbose) {
       CheckCtxt(c2, "c2*=k2");
       debugCompare(ea, secretKey, p2, c2, epsilon);
     }
     vector<cx_double> tmp_p(p1); // tmp = c1
     Ctxt tmp(c1);
     sprintf(buffer, "tmp=c1>>=%d", (int)shamt);
     rotate(tmp_p, shamt); // ea.shift(tmp, random amount in [-nSlots/2,nSlots/2])
     ea.rotate(tmp, shamt);
     if (verbose) {
       CheckCtxt(tmp, buffer);
       debugCompare(ea, secretKey, tmp_p, tmp, epsilon);
     }
     add(p2, tmp_p);  // c2 += tmp
     c2 += tmp;
     if (verbose) {
       CheckCtxt(c2, "c2+=tmp");
       debugCompare(ea, secretKey, p2, c2, epsilon);
     }
     sprintf(buffer, "c2>>>=%d", (int)rotamt);
     rotate(p2, rotamt); // ea.rotate(c2, random amount in [1-nSlots, nSlots-1])
     ea.rotate(c2, rotamt);
     if (verbose) {
       CheckCtxt(c2, buffer);
       debugCompare(ea, secretKey, p2, c2, epsilon);
     }
     negateVec(p1); // c1.negate()
     c1.negate();
     if (verbose) {
       CheckCtxt(c1, "c1=-c1");
       debugCompare(ea, secretKey, p1, c1, epsilon);
     }
     mul(p3, p2); // c3.multiplyBy(c2)
     c3.multiplyBy(c2);
     if (verbose) {
       CheckCtxt(c3, "c3*=c2");
       debugCompare(ea, secretKey, p3, c3, epsilon);
     }
     sub(p0, p3); // c0 -= c3
     c0 -= c3;
     if (verbose) {
       CheckCtxt(c0, "c0=-c3");
       debugCompare(ea, secretKey, p0, c0, epsilon);
     }
  }

  c0.cleanUp();
  c1.cleanUp();
  c2.cleanUp();
  c3.cleanUp();

  FHE_NTIMER_STOP(Circuit);

  vector<cx_double> pp0, pp1, pp2, pp3;

  ea.decrypt(c0, secretKey, pp0);
  ea.decrypt(c1, secretKey, pp1);
  ea.decrypt(c2, secretKey, pp2);
  ea.decrypt(c3, secretKey, pp3);

  std::cout << "Test "<<nRounds<<" rounds of mixed operations, ";
  if (cx_equals(pp0, p0,conv<double>(epsilon*c0.getPtxtMag()))
      && cx_equals(pp1, p1,conv<double>(epsilon*c1.getPtxtMag()))
      && cx_equals(pp2, p2,conv<double>(epsilon*c2.getPtxtMag()))
      && cx_equals(pp3, p3,conv<double>(epsilon*c3.getPtxtMag())))
    std::cout << "PASS\n\n";
  else {
    std::cout << "FAIL:\n";
    std::cout << "  max(p0)="<<largestCoeff(p0)
              << ", max(pp0)="<<largestCoeff(pp0)
              << ", maxDiff="<<calcMaxDiff(p0,pp0) << endl;
    std::cout << "  max(p1)="<<largestCoeff(p1)
              << ", max(pp1)="<<largestCoeff(pp1)
              << ", maxDiff="<<calcMaxDiff(p1,pp1) << endl;
    std::cout << "  max(p2)="<<largestCoeff(p2)
              << ", max(pp2)="<<largestCoeff(pp2)
              << ", maxDiff="<<calcMaxDiff(p2,pp2) << endl;
    std::cout << "  max(p3)="<<largestCoeff(p3)
              << ", max(pp3)="<<largestCoeff(pp3)
              << ", maxDiff="<<calcMaxDiff(p3,pp3) << endl<<endl;
  }

  if (verbose) {
    std::cout << endl;
    printAllTimers();
    std::cout << endl;
  }
  resetAllTimers();
   }

int main(int argc, char *argv[])
{

  // Commandline setup

  ArgMap amap;

  long m=16;
  long r=8;
  long L=0;
  double epsilon=0.01; // Accepted accuracy
  long R=1;
  long seed=0;
  bool debug = false;

  amap.arg("m", m, "Cyclotomic index");
  amap.note("e.g., m=1024, m=2047");
  amap.arg("r", r, "Bits of precision");
  amap.arg("R", R, "number of rounds");
  amap.arg("L", L, "Number of bits in modulus", "heuristic");
  amap.arg("ep", epsilon, "Accepted accuracy");
  amap.arg("seed", seed, "PRG seed");
  amap.arg("verbose", verbose, "more printouts");
  amap.arg("debug", debug, "for debugging");

  amap.parse(argc, argv);

  if (seed)
    NTL::SetSeed(ZZ(seed));

  if (R<=0) R=1;
  if (R<=2)
    L = 100*R;
  else
    L = 220*(R-1);

  if (verbose) {
    cout << "** m="<<m<<", #rounds="<<R<<", |q|="<<L
         << ", epsilon="<<epsilon<<endl;
  }
  epsilon /= R;
  try{

    // FHE setup keys, context, SKMs, etc

    FHEcontext context(m, /*p=*/-1, r);
    context.scale=4;
    buildModChain(context, L, /*c=*/2);

    FHESecKey secretKey(context);
    secretKey.GenSecKey(); // A +-1/0 secret key
    addSome1DMatrices(secretKey); // compute key-switching matrices

    const FHEPubKey publicKey = secretKey;
    const EncryptedArrayCx& ea = context.ea->getCx();

    if (verbose) {
      ea.getPAlgebra().printout();
      cout << "r = " << context.alMod.getR() << endl;
      cout << "ctxtPrimes="<<context.ctxtPrimes
           << ", specialPrimes="<<context.specialPrimes<<endl<<endl;
    }
    if (debug) {
        dbgKey = & secretKey;
        dbgEa = (EncryptedArray*) context.ea;
    }
#ifdef DEBUG_PRINTOUT
          dbgKey = & secretKey;
          dbgEa = (EncryptedArray*) context.ea;
#endif //DEBUG_PRINTOUT

    // Run the tests.
    testBasicArith(publicKey, secretKey, ea, epsilon);
    testComplexArith(publicKey, secretKey, ea, epsilon);
    testRotsNShifts(publicKey, secretKey, ea, epsilon);
    testGeneralOps(publicKey, secretKey, ea, epsilon*R, R);
  }
  catch (exception& e) {
    cerr << e.what() << endl;
    cerr << "***Major FAIL***" << endl;
  }

  return 0;
}

void testBasicArith(const FHEPubKey& publicKey,
                    const FHESecKey& secretKey,
                    const EncryptedArrayCx& ea, double epsilon)
{
  if (verbose)  cout << "Test Arithmetic ";
  // Test objects

  Ctxt c1(publicKey), c2(publicKey), c3(publicKey);

  vector<cx_double> vd;
  vector<cx_double> vd1, vd2, vd3;
  ea.random(vd1);
  ea.random(vd2);

  // test encoding of shorter vectors
  vd1.resize(vd1.size()-2);
  ea.encrypt(c1, publicKey, vd1, /*size=*/1.0);
  vd1.resize(vd1.size()+2, 0.0);

  ea.encrypt(c2, publicKey, vd2, /*size=*/1.0);

  // Test - Multiplication
  c1 *= c2;
  for (long i=0; i<lsize(vd1); i++) vd1[i] *= vd2[i];

  ZZX poly;
  ea.random(vd3);
  ea.encode(poly, vd3, /*size=*/1.0);
  c1.addConstant(poly); // vd1*vd2 + vd3
  for (long i=0; i<lsize(vd1); i++) vd1[i] += vd3[i];

  // Test encoding, encryption of a single number
  double xx = NTL::RandomLen_long(16)/double(1L<<16); // random in [0,1]
  ea.encryptOneNum(c2, publicKey, xx);
  c1 += c2;
  for (auto& x : vd1) x += xx;

  // Test - Multiply by a mask
  vector<long> mask(lsize(vd1), 1);
  for (long i=0; i*(i+1)<lsize(mask); i++) {
    mask[i*i] = 0;
    mask[i*(i+1)] = -1;
  }

  ea.encode(poly,mask, /*size=*/1.0);
  c1.multByConstant(poly); // mask*(vd1*vd2 + vd3)
  for (long i=0; i<lsize(vd1); i++) vd1[i] *= mask[i];

  // Test - Addition
  ea.random(vd3);
  ea.encrypt(c3, publicKey, vd3, /*size=*/1.0);
  c1 += c3;
  for (long i=0; i<lsize(vd1); i++) vd1[i] += vd3[i];

  c1.negate();
  c1.addConstant(to_ZZ(1));
  for (long i=0; i<lsize(vd1); i++) vd1[i] = 1.0 - vd1[i];

  // Diff between approxNums HE scheme and plaintext floating
  ea.decrypt(c1, secretKey, vd);
#ifdef DEBUG_PRINTOUT
  printVec(cout<<"res=", vd, 10)<<endl;
  printVec(cout<<"vec=", vd1, 10)<<endl;
#endif
  if (verbose)
    cout << "(max |res-vec|_{infty}="<< calcMaxDiff(vd, vd1) << "): ";

  if (cx_equals(vd, vd1, conv<double>(epsilon*c1.getPtxtMag())))
    cout << "GOOD\n";
  else {
    cout << "BAD:\n";
    std::cout << "  max(vd)="<<largestCoeff(vd)
              << ", max(vd1)="<<largestCoeff(vd1)
              << ", maxDiff="<<calcMaxDiff(vd,vd1) << endl<<endl;
  }
}

void testComplexArith(const FHEPubKey& publicKey,
                      const FHESecKey& secretKey,
                      const EncryptedArrayCx& ea, double epsilon)
{

  // Test complex conjugate
  Ctxt c1(publicKey), c2(publicKey);

  vector<cx_double> vd;
  vector<cx_double> vd1, vd2;
  ea.random(vd1);
  ea.random(vd2);

  ea.encrypt(c1, publicKey, vd1, /*size=*/1.0);
  ea.encrypt(c2, publicKey, vd2, /*size=*/1.0);

  if (verbose)
    cout << "Test Conjugate: ";
  for_each(vd1.begin(), vd1.end(), [](cx_double& d){d=std::conj(d);});
  c1.complexConj();
  ea.decrypt(c1, secretKey, vd);
#ifdef DEBUG_PRINTOUT
  printVec(cout<<"vd1=", vd1, 10)<<endl;
  printVec(cout<<"res=", vd, 10)<<endl;
#endif
  if (cx_equals(vd, vd1, conv<double>(epsilon*c1.getPtxtMag())))
    cout << "GOOD\n";
  else {
    cout << "BAD:\n";
    std::cout << "  max(vd)="<<largestCoeff(vd)
              << ", max(vd1)="<<largestCoeff(vd1)
              << ", maxDiff="<<calcMaxDiff(vd,vd1) << endl<<endl;
  }

  // Test that real and imaginary parts are actually extracted.
  Ctxt realCtxt(c2), imCtxt(c2);
  vector<cx_double> realParts(vd2), real_dec;
  vector<cx_double> imParts(vd2), im_dec;

  if (verbose)
    cout << "Test Real and Im parts: ";
  for_each(realParts.begin(), realParts.end(), [](cx_double& d){d=std::real(d);});
  for_each(imParts.begin(), imParts.end(), [](cx_double& d){d=std::imag(d);});

  ea.extractRealPart(realCtxt);
  ea.decrypt(realCtxt, secretKey, real_dec);

  ea.extractImPart(imCtxt);
  ea.decrypt(imCtxt, secretKey, im_dec);

#ifdef DEBUG_PRINTOUT
  printVec(cout<<"vd2=", vd2, 10)<<endl;
  printVec(cout<<"real=", realParts, 10)<<endl;
  printVec(cout<<"res=", real_dec, 10)<<endl;
  printVec(cout<<"im=", imParts, 10)<<endl;
  printVec(cout<<"res=", im_dec, 10)<<endl;
#endif
  if (cx_equals(realParts,real_dec,conv<double>(epsilon*realCtxt.getPtxtMag()))
      && cx_equals(imParts, im_dec, conv<double>(epsilon*imCtxt.getPtxtMag())))
    cout << "GOOD\n";
  else {
    cout << "BAD:\n";
    std::cout << "  max(re)="<<largestCoeff(realParts)
              << ", max(re1)="<<largestCoeff(real_dec)
              << ", maxDiff="<<calcMaxDiff(realParts,real_dec) << endl;
    std::cout << "  max(im)="<<largestCoeff(imParts)
              << ", max(im1)="<<largestCoeff(im_dec)
              << ", maxDiff="<<calcMaxDiff(imParts,im_dec) << endl<<endl;
  }
}

void testRotsNShifts(const FHEPubKey& publicKey,
                     const FHESecKey& secretKey,
                     const EncryptedArrayCx& ea, double epsilon)
{

  std::srand(std::time(0)); // set seed, current time.
  int nplaces = rand() % static_cast<int>(ea.size()/2.0) + 1;

  if (verbose)
    cout << "Test Rotation of " << nplaces << ": ";

  Ctxt c1(publicKey);
  vector<cx_double> vd1;
  vector<cx_double> vd_dec;
  ea.random(vd1);
  ea.encrypt(c1, publicKey, vd1, /*size=*/1.0);

#ifdef DEBUG_PRINTOUT
  printVec(cout<< "vd1=", vd1, 10)<<endl;
#endif
  std::rotate(vd1.begin(), vd1.end()-nplaces, vd1.end());
  ea.rotate(c1, nplaces);
  c1.reLinearize();
  ea.decrypt(c1, secretKey, vd_dec);
#ifdef DEBUG_PRINTOUT
  printVec(cout<< "vd1(rot)=", vd1, 10)<<endl;
  printVec(cout<<"res: ", vd_dec, 10)<<endl;
#endif

  if (cx_equals(vd1, vd_dec, conv<double>(epsilon*c1.getPtxtMag())))
    cout << "GOOD\n";
  else {
    cout << "BAD:\n";
    std::cout << "  max(vd)="<<largestCoeff(vd_dec)
              << ", max(vd1)="<<largestCoeff(vd1)
              << ", maxDiff="<<calcMaxDiff(vd_dec,vd1) << endl<<endl;
  }
}


#include <pybind11/pybind11.h>
#include <pybind11/eigen.h>
#include<pybind11/numpy.h>
#include<fstream>
#include<iostream>
#include <vector>
#include <Eigen/Dense>
using namespace std;
using namespace Eigen;

ArrayXd find_max(VectorXd mMat)
{
	VectorXd::Index maxRow;
	double max = mMat.maxCoeff(&maxRow);
	Matrix<bool, Dynamic, 1> MM = (mMat.array() == max).matrix();
	ArrayXd Index_max = ArrayXd::Zero(mMat.size());
	for (size_t i = 0; i < MM.rows(); i++)
	{
		if (MM(i))
		{
			Index_max[i] = 1.0;
		}

	}
	return Index_max;
}

MatrixXd select_M_PN(MatrixXd C, VectorXd P)
{
	MatrixXd res(C.rows(), int(P.sum()));
	Eigen::Index j = 0;
	for (Eigen::Index i = 0; i < C.cols(); ++i)
	{
		if (P(i) > 0) res.col(j++) = C.col(i);
	}
	return res;

}

MatrixXd select_V_PN(VectorXd V, VectorXd P)
{
	VectorXd res(int(P.sum()));
	Eigen::Index j = 0;
	for (Eigen::Index i = 0; i < V.size(); ++i)
	{
		if (P(i) > 0) res(j++) = V(i);
	}
	return res;

}

VectorXd lsqnonneg(MatrixXd C, VectorXd d,double tol)
/*
min_{x>0}||Cx-d||_2^2
Reference: Lawson and Hanson,
"Solving Least Squares Problems",
Prentice-Hall, 1974.
*/
{
	d = C.transpose() * d;
	C = C.transpose() * C;
	int n = C.rows();
	// Initialize vector of n zeros and Infs(to be used later)
	VectorXd nZeros = VectorXd::Zero(n);
	VectorXd wz = nZeros;
	//Initialize set of non - active columns to null
	VectorXd P = nZeros;
	// Initialize set of active columns to all
	// and the initial point to zeros
	VectorXd Z = VectorXd::Ones(n);
	VectorXd x = nZeros;
	VectorXd w = d - C * x;
	VectorXd resid = w;
	w = C.transpose() * w;
	//Set up iteration criterion
	int outeriter = 0;
	int iter = 0;
	int itmax = 3 * n;
	VectorXd z = nZeros;
	VectorXd x_new = VectorXd::Zero(n+1);
	// Outer loop to put variables into set to hold positive coefficients
	while (Z.sum() > 0 && (w.array()*Z.array() > tol).any())
	{
		outeriter += 1;
		z = nZeros;
		wz = (wz.array()*(1 - P.array())).matrix() - P * DBL_MAX;
		wz = (wz.array()*(1 - Z.array()) + w.array()*Z.array()).matrix();
		ArrayXd t = find_max(wz);
		P = (P.array() * (1 - t) + t).matrix();
		Z = (Z.array()*(1 - t)).matrix();
		if (P.sum() == 1)
		{
			VectorXd::Index maxRow;
			P.maxCoeff(&maxRow);
			double bb = (C.col(maxRow)).transpose() * (C.col(maxRow));
			double aa = ((C.col(maxRow)).transpose() * d);
			z(maxRow) = aa / bb;
		}
		if (P.sum() >= 2)
		{
			MatrixXd C_ing = select_M_PN(C, P);
			VectorXd z_ing = (C_ing.transpose() * C_ing).ldlt().solve(C_ing.transpose() * d);
			Eigen::Index j = 0;
			for (size_t i = 0; i < P.size(); i++)
			{
				if (P(i) > 0) z(i) = z_ing(j++);
			}
		}
		while (((z.array()*P.array() + 1-P.array()) <= 0).any())
		{
			iter = iter + 1;
			if (iter > itmax)
			{
				// cout << "optimfun:lsqnonneg: IterationCountExceeded" << endl;
				x = z;
				double error = (resid.array()*resid.array()).sum();
				x_new << error,	x;
				return x_new;
			}

			// Find indices where intermediate solution z is approximately negative
			vector<double> alpha_v;
			for (Eigen::Index i = 0; i < P.size(); i++)
			{
				if (z(i) <= 0 && P(i) > 0) alpha_v.push_back(x(i) / (x(i) - z(i)));
			}
			x = x + (*min_element(alpha_v.begin(), alpha_v.end())) * (z - x);
			for (size_t i = 0; i < Z.size(); i++)
			{
				if ((abs(x(i)) < tol && P(i) > 0)||Z(i)>0)
				{
					Z(i) = 1;
				}
				else
				{
					Z(i) = 0;
				}

			}
			P = (1 - Z.array()).matrix();
			z = nZeros;
			if (P.sum() == 1)
			{
				VectorXd::Index maxRow;
				P.maxCoeff(&maxRow);
				double bb = (C.col(maxRow)).transpose() * (C.col(maxRow));
				double aa = ((C.col(maxRow)).transpose() * d);
				z(maxRow) = aa / bb;
			}
			if (P.sum() >= 2)
			{
				MatrixXd C_ing = select_M_PN(C, P);
				VectorXd z_ing = (C_ing.transpose() * C_ing).ldlt().solve(C_ing.transpose() * d);
				Eigen::Index j = 0;
				for (size_t i = 0; i < P.size(); i++)
				{
					if (P(i) > 0) z(i) = z_ing(j++);
				}
			}
		}
		x = z;
		resid = (d - C * x);
		w = C.transpose() * resid;
	}
	double error = (resid.array()*resid.array()).sum();
	x_new <<error, x;
	return x_new;
}

namespace py = pybind11;
PYBIND11_MODULE(libNNLS, m)
{
m.doc() = "Non-negative least squares";
m.def("lsqnonneg", &lsqnonneg);
}


#include <iostream>
#include <Eigen/Dense>
#include <Eigen/Core>
#include <vector>
#include <chrono>
#include <unistd.h>
using namespace std;
using namespace Eigen;
using namespace std::chrono;

#define DIM 1000
#define NUM_CORES 4
#define ENABLE_PARALLEL

float get_time()
{
    return (float)duration_cast< milliseconds >(system_clock::now().time_since_epoch()).count();
}

void init_A(MatrixXd &A)
{
    for(int i = 0; i < DIM; i++)
    {
        for(int j = 0; j < DIM; j++)
        {
            A(i,j) = static_cast <float> (rand());
        }
    }
}

void init_b(VectorXd &b)
{
    for(int i = 0; i < DIM; i++)
    {
        b[i] = static_cast <float> (rand());
    }
}

void test_QRpseudoInv(MatrixXd &A, VectorXd &b)
{
    A.colPivHouseholderQr().solve(b);
}

void test_QR(MatrixXd &A, VectorXd &b)
{
    A.householderQr().solve(b);
}

void test_CompleteOrtho(MatrixXd &A, VectorXd &b)
{
    A.completeOrthogonalDecomposition().solve(b);
}

void test_LDLT(MatrixXd &A, VectorXd &b)
{
    A.ldlt().solve(b);
}

int main()
{
    /**********Declaration & Initialization**********/
    MatrixXd A(DIM,DIM);
    VectorXd b(DIM);
    init_A(A);
    init_b(b);
    /**********Declaration & Initialization**********/

#ifdef ENABLE_PARALLEL
    Eigen::initParallel();
    Eigen::setNbThreads(NUM_CORES);
#endif

    auto start_ = chrono::steady_clock::now();
    test_QRpseudoInv(A,b);
    auto end_ = chrono::steady_clock::now();
    cout << "colPivHouseholderQr : "
        << chrono::duration_cast<chrono::nanoseconds>(end_ - start_).count()
        << " ns" << endl;

    start_ = chrono::steady_clock::now();
    test_QR(A,b);
    end_ = chrono::steady_clock::now();
    cout << "householderQr : "
        << chrono::duration_cast<chrono::nanoseconds>(end_ - start_).count()
        << " ns" << endl;

    start_ = chrono::steady_clock::now();
    test_CompleteOrtho(A,b);
    end_ = chrono::steady_clock::now();
    cout << "completeOrthogonalDecomposition : "
        << chrono::duration_cast<chrono::nanoseconds>(end_ - start_).count()
        << " ns" << endl;

    start_ = chrono::steady_clock::now();
    test_LDLT(A,b);
    end_ = chrono::steady_clock::now();
    cout << "LDLT: "
        << chrono::duration_cast<chrono::nanoseconds>(end_ - start_).count()
        << " ns" << endl;

}


#include "cxxopts.hpp"

// see https://www.gitmemory.com/issue/ddemidov/amgcl/162/668662041 for inspiration

#include <amgcl/io/mm.hpp>
#include <amgcl/adapter/crs_tuple.hpp>
#include <amgcl/amg.hpp>

#if defined(SOLVER_BACKEND_VEXCL)
#  include <amgcl/value_type/static_matrix.hpp>
#  include <amgcl/adapter/block_matrix.hpp>
#  include <amgcl/backend/vexcl.hpp>
#  include <amgcl/backend/vexcl_static_matrix.hpp>
   typedef amgcl::backend::vexcl<double> Backend;
#elif defined(SOLVER_BACKEND_CUDA)
#  include <amgcl/adapter/eigen.hpp>
#  include <amgcl/backend/cuda.hpp>
#  include <amgcl/relaxation/cusparse_ilu0.hpp>
   using Backend = amgcl::backend::cuda<double>;
#else
   using Backend = amgcl::backend::builtin<double>;
#endif

#include <amgcl/make_solver.hpp>
#include <amgcl/make_block_solver.hpp>
#include <amgcl/solver/bicgstab.hpp>
#include <amgcl/solver/gmres.hpp>
#include <amgcl/solver/runtime.hpp>
#include <amgcl/preconditioner/runtime.hpp>
#include <amgcl/amg.hpp>
#include <amgcl/coarsening/smoothed_aggregation.hpp>
#include <amgcl/coarsening/plain_aggregates.hpp>
#include <amgcl/coarsening/ruge_stuben.hpp>
#include <amgcl/coarsening/aggregation.hpp>
#include <amgcl/relaxation/spai0.hpp>
#include <amgcl/relaxation/ilu0.hpp>
#include <amgcl/adapter/crs_tuple.hpp>
#include <boost/program_options.hpp>
#include <boost/property_tree/ptree.hpp>
#include <boost/property_tree/json_parser.hpp>

// ./amgcl-block-matrices -A ../data/B_test.mtx  -b 3

#include <amgcl/profiler.hpp>

namespace amgcl { profiler<> prof; }
using amgcl::prof;
using amgcl::precondition;

//---------------------------------------------------------------------------
std::tuple<size_t, double> scalar_solve(
        const boost::property_tree::ptree &prm,
        size_t rows,
        std::vector<ptrdiff_t> const &ptr,
        std::vector<ptrdiff_t> const &col,
        std::vector<double>    const &val,
        std::vector<double>    const &rhs,
        std::vector<double>          &x,
        bool reorder
        )
{
    Backend::params bprm;
    // We use the tuple of CRS arrays to represent the system matrix.
    // Note that std::tie creates a tuple of references, so no data is actually
    // copied here:
    auto A = std::tie(rows, ptr, col, val);

#if defined(SOLVER_BACKEND_VEXCL)
    vex::Context ctx(vex::Filter::Env);
    std::cout << ctx << std::endl;
    bprm.q = ctx;
#elif defined(SOLVER_BACKEND_VIENNACL)
    std::cout
        << viennacl::ocl::current_device().name()
        << " (" << viennacl::ocl::current_device().vendor() << ")\n\n";
#elif defined(SOLVER_BACKEND_CUDA)
    cusparseCreate(&bprm.cusparse_handle);
    {
        int dev;
        cudaGetDevice(&dev);

        cudaDeviceProp prop;
        cudaGetDeviceProperties(&prop, dev);
        std::cout << prop.name << std::endl << std::endl;
    }
#endif

    using Solver1 = amgcl::make_solver<
    amgcl::runtime::preconditioner<Backend>,
    amgcl::runtime::solver::wrapper<Backend>>;

    //using Solver2 = amgcl::make_block_solver<
    using Solver2 = amgcl::make_solver<
    // Use AMG as preconditioner:
    amgcl::amg<
        Backend,
        amgcl::coarsening::smoothed_aggregation,
        //amgcl::relaxation::spai0
        amgcl::relaxation::damped_jacobi
        >,
    // And BiCGStab as iterative solver:
    amgcl::solver::bicgstab<Backend>>;

    using Solver3 = amgcl::make_solver<
    // Use AMG as preconditioner:
    amgcl::amg<
        Backend,
        amgcl::coarsening::ruge_stuben,
        amgcl::relaxation::spai0
        //amgcl::relaxation::ilu0
        >,
    // And GMRES as iterative solver:
    amgcl::solver::gmres<Backend>>;

    try {
        Solver2::params prm;
        Solver2::backend_params bprm;
        // prm.precond.direct_coarse = false;
        prm.precond.direct_coarse = true;
        prm.solver.maxiter = 1200;
        prm.solver.verbose = true;
        prm.solver.tol = 1e-5;
        //prm.precond.coarsening ...

#if defined(SOLVER_BACKEND_VEXCL)
        vex::Context ctx(vex::Filter::Env);
        std::cout << ctx << std::endl;
        bprm.q = ctx;
#elif defined(SOLVER_BACKEND_VIENNACL)
        std::cout
            << viennacl::ocl::current_device().name()
            << " (" << viennacl::ocl::current_device().vendor() << ")\n\n";
#elif defined(SOLVER_BACKEND_CUDA)
        cusparseCreate(&bprm.cusparse_handle);
        {
            int dev;
            cudaGetDevice(&dev);

            cudaDeviceProp prop;
            cudaGetDeviceProperties(&prop, dev);
            std::cout << prop.name << std::endl << std::endl;
        }
#endif

        Solver2 solve( A, prm, bprm );

        auto f_b = Backend::copy_vector(rhs, bprm);
        auto x_b = Backend::copy_vector(x,   bprm);

        //auto [iters, error] = solve(A, rhs, x);
        auto [iters, error] = solve(A, *f_b, *x_b);

#if defined(SOLVER_BACKEND_VEXCL)
        vex::copy(*x_b, x);
#elif defined(SOLVER_BACKEND_VIENNACL)
        viennacl::fast_copy(*x_b, x);
#elif defined(SOLVER_BACKEND_CUDA)
        thrust::copy(x_b->begin(), x_b->end(), x.begin());
#else
        std::copy(&(*x_b)[0], &(*x_b)[0] + rows, &x[0]);
#endif

        // Output the number of iterations, the relative error:
        std::cout << "Iters: " << iters << std::endl
                  << "Error: " << error << std::endl;
        // std::cout << "x = [ ";
        // for (auto el : x) {
        //     std::cout << el << ", ";
        // }
        // std::cout << "]\n";
        return {iters, error};
    } catch(std::runtime_error &e) {
        std::cout << "caught exception: " << e.what() << std::endl;
    }

    return {0, 0.0};
}

int main(int argc, char *argv[])
{
    //namespace po = boost::program_options;
    namespace io = amgcl::io;

    using std::vector;
    using std::string;

    cxxopts::Options options("amgcl_block_matrices", "check AMG for different block sizes");
    options
        .allow_unrecognised_options()
        .add_options()
        ("A,matrix", "System matrix in the MatrixMarket format.",
         cxxopts::value<std::string>())
        ("b,block-size", "The block size of the system matrix. ",
         cxxopts::value<int>()->default_value("1"))
        ("h,help", "Print help")
        ;

    auto vm = options.parse(argc, argv);

    if (vm.count("help")) {
        std::cout << options.help() << std::endl;
        return 0;
    }

    for (int i = 0; i < argc; ++i) {
        if (i) std::cout << " ";
        std::cout << argv[i];
    }
    std::cout << std::endl;

    boost::property_tree::ptree prm;
    if (vm.count("prm-file")) {
        read_json(vm["prm-file"].as<string>(), prm);
    }

    if (vm.count("prm")) {
        for(const string &v : vm["prm"].as<vector<string>>()) {
            amgcl::put(prm, v);
        }
    }

    size_t n, block_size = 1;
    vector<ptrdiff_t> ptr, col;
    vector<double> val;

    if (vm.count("block-size")) {
        block_size = vm["block-size"].as<int>();
    }

    if (vm.count("matrix")) {
        string Afile  = vm["matrix"].as<string>();
        auto [rows, cols] = io::mm_reader(Afile)(ptr, col, val);

        std::cout << "rows: " << rows << std::endl;
        std::cout << "cols: " << cols << std::endl;
        n = rows;
    }

    // We use the tuple of CRS arrays to represent the system matrix.
    // Note that std::tie creates a tuple of references, so no data is actually
    // copied here:
    auto A = std::tie(n, ptr, col, val);
    std::vector<double> rhs (n, 0.0);
    std::vector<double> x   (n, 0.0);

    // Set RHS := Ax where x = 1
    for (size_t i = 0; i < n; ++i) {
        double s = 0;
        for (ptrdiff_t j = ptr[i], e = ptr[i+1]; j < e; ++j) {
            s += val[j];
        }
        rhs[i] = s;
    }

    if (block_size == 1) { // scalar case
        scalar_solve(prm, n, ptr,col, val, rhs, x, false);
    } else {
        // Compose the solver type
        constexpr size_t B = 6; // 3
        using dmat_type = amgcl::static_matrix<double, B, B>;
        using dvec_type = amgcl::static_matrix<double, B, 1>;
        using SBackend = amgcl::backend::builtin<double>;    // the outer iterative solver backend
        using PBackend = amgcl::backend::builtin<float>;     // the PSolver backend
        using UBackend = amgcl::backend::builtin<dmat_type>; // the USolver backend
        using BlockSolver = amgcl::make_block_solver<
                // preconditioner
                amgcl::amg<
                    UBackend,
                    amgcl::coarsening::aggregation,
                    amgcl::relaxation::ilu0
                    >,
                // solver
                amgcl::solver::gmres<UBackend>>;
        using Solver = amgcl::make_solver<
                // preconditioner
                amgcl::amg<
                    UBackend,
                    amgcl::coarsening::aggregation,
                    amgcl::relaxation::ilu0
                    >,
                // solver
                amgcl::solver::gmres<UBackend>>;
        try {
            // see `tutorial/3.CoupCons3D/coupcons3d.cpp`
#if 1
            BlockSolver::params prm;
            prm.solver.maxiter = 1200;
            prm.solver.verbose = true;
            BlockSolver::backend_params bprm;
            prm.precond.direct_coarse = false;
            prm.precond.allow_rebuild = true;
            auto Ab = amgcl::adapter::block_matrix<dmat_type>(A);
            auto rhsb = reinterpret_cast<dvec_type*>(rhs.data());
            auto xb = reinterpret_cast<dvec_type*>(x.data());
            auto RHS = amgcl::make_iterator_range(rhsb, rhsb + n / 3);
            auto X = amgcl::make_iterator_range(xb, xb + n / block_size);
            Solver solve0( Ab, prm, bprm );
            auto [iters0, error0] = solve0(Ab, RHS, X);
            std::cout << "Iters0: " << iters0 << std::endl
                      << "Error0: " << error0 << std::endl;

            BlockSolver solve( A, prm, bprm );
            auto [iters, error] = solve(A, rhs, x);

            std::cout << "Iters: " << iters << std::endl
                      << "Error: " << error << std::endl;
            // std::cout << "x = [ ";
            // for (auto el : x) {
            //     std::cout << el << ", ";
            // }
            // std::cout << "]\n";
#endif
        } catch(std::runtime_error &e) {
            std::cout << "caught exception: " << e.what() << std::endl;
        }
    }
}


#define DEBUG 1
/**
 * File    : F.cpp
 * Author  : Kazune Takahashi
 * Created : 6/16/2020, 5:44:32 PM
 * Powered by Visual Studio Code
 */
#include <algorithm>
#include <bitset>
#include <cassert>
#include <cctype>
#include <chrono>
#include <climits>
#include <cmath>
#include <complex>
#include <cstdint>
#include <cstdio>
#include <cstdlib>
#include <functional>
#include <iomanip>
#include <iostream>
#include <limits>
#include <map>
#include <queue>
#include <random>
#include <set>
#include <stack>
#include <string>
#include <tuple>
#include <unordered_map>
#include <unordered_set>
#include <vector>
// ----- boost -----
#include <boost/integer/common_factor_rt.hpp>
#include <boost/multiprecision/cpp_int.hpp>
#include <boost/rational.hpp>
// ----- using directives and manipulations -----
using namespace std;
using boost::rational;
using boost::integer::gcd; // for C++14 or for cpp_int
using boost::integer::lcm; // for C++14 or for cpp_int
using boost::multiprecision::cpp_int;
using ll = long long;
using ld = long double;
template <typename T>
using max_heap = priority_queue<T>;
template <typename T>
using min_heap = priority_queue<T, vector<T>, greater<T>>;
// ----- constexpr for Mint and Combination -----
constexpr ll MOD{1'000'000'007LL};
// constexpr ll MOD{998'244'353LL}; // be careful
constexpr ll MAX_SIZE{3'000'010LL};
// constexpr ll MAX_SIZE{30'000'010LL}; // if 10^7 is needed
// ----- ch_max and ch_min -----
template <typename T>
bool ch_max(T &left, T right)
{
  if (left < right)
  {
    left = right;
    return true;
  }
  return false;
}
template <typename T>
bool ch_min(T &left, T right)
{
  if (left > right)
  {
    left = right;
    return true;
  }
  return false;
}
// ----- Mint -----
template <ll MOD = MOD>
class Mint
{
public:
  ll x;
  Mint() : x{0LL} {}
  Mint(ll x) : x{(x % MOD + MOD) % MOD} {}
  Mint operator-() const { return x ? MOD - x : 0; }
  Mint &operator+=(Mint const &a)
  {
    if ((x += a.x) >= MOD)
    {
      x -= MOD;
    }
    return *this;
  }
  Mint &operator-=(Mint const &a) { return *this += -a; }
  Mint &operator++() { return *this += 1; }
  Mint operator++(int)
  {
    Mint tmp{*this};
    ++*this;
    return tmp;
  }
  Mint &operator--() { return *this -= 1; }
  Mint operator--(int)
  {
    Mint tmp{*this};
    --*this;
    return tmp;
  }
  Mint &operator*=(Mint const &a)
  {
    (x *= a.x) %= MOD;
    return *this;
  }
  Mint &operator/=(Mint const &a)
  {
    Mint b{a};
    return *this *= b.power(MOD - 2);
  }
  Mint operator+(Mint const &a) const { return Mint(*this) += a; }
  Mint operator-(Mint const &a) const { return Mint(*this) -= a; }
  Mint operator*(Mint const &a) const { return Mint(*this) *= a; }
  Mint operator/(Mint const &a) const { return Mint(*this) /= a; }
  bool operator<(Mint const &a) const { return x < a.x; }
  bool operator<=(Mint const &a) const { return x <= a.x; }
  bool operator>(Mint const &a) const { return x > a.x; }
  bool operator>=(Mint const &a) const { return x >= a.x; }
  bool operator==(Mint const &a) const { return x == a.x; }
  bool operator!=(Mint const &a) const { return !(*this == a); }
  Mint power(ll N) const
  {
    if (N == 0)
    {
      return 1;
    }
    else if (N % 2 == 1)
    {
      return *this * power(N - 1);
    }
    else
    {
      Mint half = power(N / 2);
      return half * half;
    }
  }
};
template <ll MOD>
Mint<MOD> operator+(ll lhs, Mint<MOD> const &rhs) { return rhs + lhs; }
template <ll MOD>
Mint<MOD> operator-(ll lhs, Mint<MOD> const &rhs) { return -rhs + lhs; }
template <ll MOD>
Mint<MOD> operator*(ll lhs, Mint<MOD> const &rhs) { return rhs * lhs; }
template <ll MOD>
Mint<MOD> operator/(ll lhs, Mint<MOD> const &rhs) { return Mint<MOD>{lhs} / rhs; }
template <ll MOD>
istream &operator>>(istream &stream, Mint<MOD> &a) { return stream >> a.x; }
template <ll MOD>
ostream &operator<<(ostream &stream, Mint<MOD> const &a) { return stream << a.x; }
// ----- Combination -----
template <ll MOD = MOD, ll MAX_SIZE = MAX_SIZE>
class Combination
{
public:
  vector<Mint<MOD>> inv, fact, factinv;
  Combination() : inv(MAX_SIZE), fact(MAX_SIZE), factinv(MAX_SIZE)
  {
    inv[1] = 1;
    for (auto i{2LL}; i < MAX_SIZE; i++)
    {
      inv[i] = (-inv[MOD % i]) * (MOD / i);
    }
    fact[0] = factinv[0] = 1;
    for (auto i{1LL}; i < MAX_SIZE; i++)
    {
      fact[i] = Mint<MOD>(i) * fact[i - 1];
      factinv[i] = inv[i] * factinv[i - 1];
    }
  }
  Mint<MOD> operator()(int n, int k)
  {
    if (n >= 0 && k >= 0 && n - k >= 0)
    {
      return fact[n] * factinv[k] * factinv[n - k];
    }
    return 0;
  }
  Mint<MOD> catalan(int x, int y)
  {
    return (*this)(x + y, y) - (*this)(x + y, y - 1);
  }
};
// ----- for C++14 -----
using mint = Mint<MOD>;
using combination = Combination<MOD, MAX_SIZE>;
// ----- for C++17 -----
template <typename T, typename enable_if<is_integral<T>::value>::type * = nullptr>
size_t popcount(T x) { return bitset<64>(x).count(); }
size_t popcount(string const &S) { return bitset<200010>{S}.count(); }
// ----- Infty -----
template <typename T>
constexpr T Infty() { return numeric_limits<T>::max(); }
template <typename T>
constexpr T mInfty() { return numeric_limits<T>::min(); }
// ----- frequently used constexpr -----
// constexpr double epsilon{1e-10};
// constexpr ll infty{1'000'000'000'000'010LL}; // or
// constexpr int infty{1'000'000'010};
// constexpr int dx[4] = {1, 0, -1, 0};
// constexpr int dy[4] = {0, 1, 0, -1};
// ----- Yes() and No() -----
void Yes()
{
  cout << "Yes" << endl;
  exit(0);
}
void No()
{
  cout << "No" << endl;
  exit(0);
}

// ----- ReadGraph -----
// Referring to ymatsux-san's source code: https://atcoder.jp/contests/abc138/submissions/7016619

struct Edge
{
  int src, dst, id;
  // ll cost;
  Edge() {}
  Edge(int src, int dst, int id) : src{src}, dst{dst}, id{id} {}
  // Edge(int src, int dst, ll cost) : src{src}, dst{dst}, cost{cost} {}

  void added_edge(vector<vector<Edge>> &V)
  {
    V[src].push_back(*this);
  }

  void added_rev(vector<vector<Edge>> &V)
  {
    V[dst].push_back(rev());
  }

  Edge rev()
  {
    Edge edge{*this};
    swap(edge.src, edge.dst);
    return edge;
  }
};

tuple<vector<vector<Edge>>, vector<Edge>> ReadGraphWithEdges(int N, int M, bool is_undirected = true, bool is_one_indexed = true)
{
  vector<vector<Edge>> V(N);
  vector<Edge> E(M);
  for (auto i = 0; i < M; ++i)
  {
    int v, w;
    cin >> v >> w;
    if (is_one_indexed)
    {
      --v;
      --w;
    }
    Edge edge{v, w, i};
    edge.added_edge(V);
    if (is_undirected)
    {
      edge.added_rev(V);
    }
    E.push_back(edge);
  }
  return make_tuple(V, E);
}

vector<vector<Edge>> ReadGraph(int N, int M, bool is_undirected = true, bool is_one_indexed = true)
{
  return get<0>(ReadGraphWithEdges(N, M, is_undirected, is_one_indexed));
}

tuple<vector<vector<Edge>>, vector<Edge>> ReadTreeWithEdges(int N)
{
  return ReadGraphWithEdges(N, N - 1);
}

vector<vector<Edge>> ReadTree(int N)
{
  return ReadGraph(N, N - 1);
}

// ----- Solve -----

using Info = tuple<int, int>;

class Solve
{
  int N, M;
  vector<vector<Edge>> V;
  vector<Info> W;
  vector<ll> P;

public:
  Solve(int N, int M, vector<vector<Edge>> const &V) : N{N}, M{M}, V{V}, W(M), P(M)
  {
    for (auto i{0}; i < M; ++i)
    {
      cin >> get<0>(W[i]) >> get<1>(W[i]);
      get<0>(W[i])--;
      get<1>(W[i])--;
    }
  }

  void flush()
  {
    for (auto i{0}; i < M; ++i)
    {
      P[i] = path(get<0>(W[i]), get<1>(W[i]));
#if DEBUG == 1
      cerr << "P[" << i << "] = " << P[i] << endl;
#endif
    }
    ll ans{0};
    for (auto i{0}; i < (1 << M); ++i)
    {
      if (popcount(i) & 1)
      {
        ans -= mask_to_cnt(i);
      }
      else
      {
        ans += mask_to_cnt(i);
      }
    }
    cout << ans << endl;
  }

private:
  ll mask_to_cnt(int mask)
  {
    return 1LL << (N - 1 - popcount(mask_to_path(mask)));
  }

  ll mask_to_path(int mask)
  {
    ll ans{0};
    for (auto i{0}; i < M; ++i)
    {
      if (mask >> i & 1)
      {
        ans |= P[i];
      }
    }
    return ans;
  }

  ll path(int src, int dst)
  {
    vector<int> parents(N, -1), ids(N, -1);
    dfs(src, parents, ids);
    ll ans{0};
    while (dst != src)
    {
      ans |= 1LL << ids[dst];
      dst = parents[dst];
    }
    return ans;
  }

  void dfs(int src, vector<int> &parents, vector<int> &ids)
  {
    for (auto e : V[src])
    {
      if (parents[e.src] == e.dst)
      {
        continue;
      }
      parents[e.dst] = e.src;
      ids[e.dst] = e.id;
      dfs(e.dst, parents, ids);
    }
  }
};

// ----- main() -----

int main()
{
  int N;
  cin >> N;
  auto V{ReadTree(N)};
  int M;
  cin >> M;
  Solve solve(N, M, V);
  solve.flush();
}


/*
 * variable_selection.cpp
 *
 *  Created on: 19 Aug 2020
 *      Author: kovesarki
 */

#include <Eigen/Dense>
#include "include/monopoly.h"
#include <random>
#include <iostream>
#include <vector>
#include <string>

#include "include/tensor_serie.hh"
#include "include/correlation_tensor_serie.hh"

int main(void) {
	std::vector<std::vector<double> > signal(10000, std::vector<double>(3,0.0) );
	std::vector<std::vector<double> > background(10000, std::vector<double>(3,0.0) );

	std::random_device rd;
	std::mt19937 gen(rd());

	std::normal_distribution<> sx(0,1.01);
	std::normal_distribution<> sy(1.0,0.4);
	std::normal_distribution<> sz(0.7,2.0);

	std::normal_distribution<> bx(0.5,0.9);
	std::normal_distribution<> by(1.4,0.7);
	std::normal_distribution<> bz(0.73,2.01);

	std::vector<std::string >  variable_names(3);
	variable_names[0] = "parameter x";
	variable_names[1] = "parameter y";
	variable_names[2] = "parameter z";

	for(int i=0; i!= signal.size(); i++ ){
		signal[i][0] = sx(gen);
		signal[i][1] = sy(gen);
		signal[i][2] = sz(gen);
	}

	for(int i=0; i!= background.size(); i++ ){
		background[i][0] = bx(gen);
		background[i][1] = by(gen);
		background[i][2] = bz(gen);
	}

	tensor_serie<double> m_x(3,2);
	correlation_tensor_serie<double> signal_fourier(3,1); //number of dimensions, number of degree (linear is enough for comparison)
	for(int i=0; i!= signal.size(); i++){
		m_x.create_diad(signal[i]);
		signal_fourier.fill(m_x,1.0,1.0); //training for target y=1.0,weights=1.0
	}
	signal_fourier.normalize();

	correlation_tensor_serie<double> background_fourier(3,1); //number of dimensions, number of degree (linear is enough for comparison)
	for(int i=0; i!= background.size(); i++) {
		m_x.create_diad(background[i]);
		background_fourier.fill(m_x,1.0,1.0);
	}
	background_fourier.normalize();

	signal_fourier.cov_matrix(1);
	background_fourier.cov_matrix(1);

	double a = signal_fourier.sigma_distance(background_fourier);
	std::cout << "significance is " << a << std::endl;

	std::vector<double> sigmas(variable_names.size(), 0.0);
	std::vector<int> variable_ordering(variable_names.size(), 0.0);

	signal_fourier.variable_selection(background_fourier, sigmas, variable_ordering);

	std::cout << "Variable order, from best to worst. " << std::endl
			  << "Significance is meant for the group of best variables" << std::endl
			  << "(the variable, and all the variables above)." << std::endl;
	for(int i=0; i!= variable_ordering.size(); i++)
	{
	  std::cout << sigmas[variable_ordering[i]] << " \t" << variable_names[variable_ordering[i]] << std::endl;
	}

	std::cout << "varnames["<< variable_ordering.size()<< "] = {" << std::endl;
	for(int i=variable_ordering.size()-1; i!= -1; i--)
	{
	  std::cout << "\"" << variable_names[variable_ordering[i]] << "\"";
	  if(i!=0) {std::cout << ","; }
	  else {std::cout << std::endl << "}";}
	  std::cout << std::endl;
	}
}