/** $Id: EnergyDistribution.cxx 165001 2018-08-27 09:22:30Z jvansanten $
 * @file
 * @author Jakob van Santen <vansanten@wisc.edu>
 *
 * $Revision: 165001 $
 * $Date: 2018-08-27 03:22:30 -0600 (Mon, 27 Aug 2018) $
 */

#include <MuonGun/EnergyDistribution.h>
#include <MuonGun/RadialDistribution.h>
#include <phys-services/I3RandomService.h>
#include <MuonGun/EnsembleSampler.h>

#include <boost/bind.hpp>

namespace I3MuonGun {

EnergyDistribution::~EnergyDistribution() {};

double
EnergyDistribution::operator()(double d, double ct,
    unsigned m, double r, double e) const
{
	return std::exp(GetLog(d, ct, m, r, log_value(std::log(e))));
}

double
EnergyDistribution::Integrate(double d, double ct,
    unsigned m, double r_min, double r_max, double e_min, double e_max) const
{
	// restrict integration to range where density can be nonzero
	double loge_min = std::max(minLog_, std::log(e_min));
	double loge_max = std::min(maxLog_, std::log(e_max));
	r_min = std::max(0., r_min);
	r_max = std::min(GetMaxRadius(), r_max);
	if (m > 1) {
		// Integrate dP/(dlogE dr^2) for numerical stability
		auto integrand = [this,d,ct,m](double r2, double loge)
		{
			double r = std::sqrt(r2);
			return std::exp(GetLog(d, ct, m, r, log_value(loge))
			    - std::log(2*r) + loge);
		};
		boost::array<double, 2> lo = {{r_min*r_min, loge_min}};
		boost::array<double, 2> hi = {{r_max*r_max, loge_max}};
		return I3MuonGun::Integrate(boost::function<double(double,double)>(integrand),
		    lo, hi, 1e-12, 1e-6, 10000);
	} else {
		// For single muons, dP/dr is a delta function at 0
		auto integrand = [this,d,ct,m](double loge)
		{
			return std::exp(GetLog(d, ct, m, 0, log_value(loge)) + loge);
		};
		return I3MuonGun::Integrate(boost::function<double(double)>(integrand), loge_min, loge_max);
	}
}

bool
SplineEnergyDistribution::operator==(const EnergyDistribution &o) const
{
	const SplineEnergyDistribution *other = dynamic_cast<const SplineEnergyDistribution*>(&o);
	if (!other)
		return false;
	else
		return (singles_ == other->singles_ && bundles_ == other->bundles_);
}

SplineEnergyDistribution::SplineEnergyDistribution(const std::string &singles, const std::string &bundles)
    : singles_(singles), bundles_(bundles)
{
	if (singles_.GetNDim() != 3u)
		log_fatal("'%s' does not appear to be a single-muon energy distribution", singles.c_str());
	if (bundles_.GetNDim() != 5u)
		log_fatal("'%s' does not appear to be a muon bundle energy distribution", bundles.c_str());
	SetMin(std::exp(std::max(singles_.GetExtents(2).first, bundles_.GetExtents(2).first)));
	SetMax(std::exp(std::min(singles_.GetExtents(2).second, bundles_.GetExtents(2).second)));
}

double
SplineEnergyDistribution::GetMaxRadius() const
{
	return bundles_.GetExtents(3).second;
}

double
SplineEnergyDistribution::GetLog(double depth, double cos_theta,
    unsigned multiplicity, double radius, log_value log_energy) const
{
	double coords[5] = {cos_theta, depth, static_cast<double>(multiplicity),
	    radius, log_energy};
	double logprob;

	if (radius < 0 || radius > GetMaxRadius() ||
	    log_energy < minLog_ || log_energy > maxLog_) {
		return -std::numeric_limits<double>::infinity();
	} else if (multiplicity < 2) {
		coords[2] = coords[4];
		if (singles_.Eval(coords, &logprob) != 0)
			return -std::numeric_limits<double>::infinity();
	} else if (bundles_.Eval(coords, &logprob) != 0)
		return -std::numeric_limits<double>::infinity();

	// Bundle spline is fit to log(dP/dr^2 dlogE)
	if (multiplicity > 1)
		logprob += std::log(2*radius);

	return logprob;
}

std::vector<std::pair<double,double> >
SplineEnergyDistribution::Generate(I3RandomService &rng, double depth,
    double cos_theta, unsigned multiplicity, unsigned nsamples) const
{
	typedef double (Signature)(double, double);
	typedef EnsembleSampler<Signature> Sampler;

	// the number of walkers must be even, and at least twice the
	// dimensionality of the space
	const unsigned walkers = std::max(4u, nsamples + (nsamples % 2));
	std::vector<Sampler::array_type> initial_ensemble(walkers);
	{
		// Draw starting positions from the MUPAGE parameterization
		BMSSEnergyDistribution proposal;
		proposal.SetMin(GetMin());
		proposal.SetMax(GetMax());
		std::vector<std::pair<double, double> > samples =
		    proposal.Generate(rng, depth, cos_theta, multiplicity, walkers);
		for (unsigned i=0; i < walkers; i++) {
			initial_ensemble[i][0] = samples.at(i).first;
			initial_ensemble[i][1] = samples.at(i).second;
		}
	}

	auto log_posterior = [this,depth,cos_theta,multiplicity](double r, double e)
	{
		return this->GetLog(depth, cos_theta, multiplicity, r,
		    EnergyDistribution::log_value(std::log(e)));
	};
	Sampler sampler(log_posterior, initial_ensemble);

	// Run the sampler for a few cycles to make it independent of the initial
	// ensemble. Fewer than 50 or so burn-in steps is too small to reach the
	// stationary distribution, while more than 100 is a waste of time, as
	// measured with resources/test/test_sampling.py
	for (unsigned i=0; i < 64; i++)
		sampler.Sample(rng);

	// copy the current ensemble into the output
	std::vector<std::pair<double,double> > samples;
	samples.reserve(nsamples);
	const std::vector<Sampler::sample> &ensemble = sampler.Sample(rng);
	unsigned todo = std::min(walkers, nsamples - unsigned(samples.size()));
	for (unsigned j=0; j < todo; j++) {
		samples.push_back(std::make_pair(ensemble[j].point[0], ensemble[j].point[1]));
	}

	// check the acceptance rate for sanity.
	double acceptance_rate = sampler.GetAcceptanceRate();
	if (acceptance_rate < 0.2) {
		log_warn("Ensemble sampler accepted only %.0f%% of samples (too low). It may be stuck in a local maximum.", 100*acceptance_rate);
	} else if (acceptance_rate > 0.9) {
		log_warn("Ensemble sampler accepted %.0f%% of samples (too high). This is an expensive random walk.", 100*acceptance_rate);
	}

	return samples;
}

BMSSEnergyDistribution::BMSSEnergyDistribution() :
    beta_(0.42), g0_(-0.232), g1_(3.961), e0a_(0.0304), e0b_(0.359), e1a_(-0.0077), e1b_(0.659),
    a0_(0.0033), a1_(0.0079), b0a_(0.0407), b0b_(0.0283), b1a_(-0.312), b1b_(6.124),
    q0_(0.0543), q1_(-0.365), c0a_(-0.069), c0b_(0.488), c1_(-0.117),
    d0a_(-0.398), d0b_(3.955), d1a_(0.012), d1b_(-0.35)
{}

bool
BMSSEnergyDistribution::operator==(const EnergyDistribution &o) const
{
	return dynamic_cast<const BMSSEnergyDistribution*>(&o);
}

OffsetPowerLaw
BMSSEnergyDistribution::GetSpectrum(double depth, double cos_theta, unsigned m, double r) const
{
	// Convert to water-equivalent depth
	double h = (200*I3Units::m/I3Units::km)*0.832 + (depth-(200*I3Units::m/I3Units::km))*0.917;
	double bX = beta_*h/cos_theta;
	double g, eps;
	if (m == 1) {
		g = g0_*log(h) + g1_;
		eps = (e0a_*exp(e0b_*h)/cos_theta + e1a_*h + e1b_)*I3Units::TeV;
	} else {
		m = std::min(m, 4u);
		double a = a0_*h + a1_;
		double b = (b0a_*m + b0b_)*h + (b1a_*m + b1b_);
		double q = q0_*h + q1_;
		g = a*r + b*(1 - 0.5*exp(q*r));
		double c = (c0a_*h + c0b_)*exp(c1_*r);
		double d = (d0a_*h + d0b_)*pow(r, d1a_*h + d1b_);
		eps = (c*acos(cos_theta) + d)*I3Units::TeV;
	}

	return OffsetPowerLaw(g, eps*(1-exp(-bX)), GetMin(), GetMax());
}

double
BMSSEnergyDistribution::GetLog(double depth, double cos_theta,
    unsigned multiplicity, double radius, log_value log_energy) const
{

	return BMSSRadialDistribution().GetLog(depth, cos_theta, multiplicity, radius) +
	    GetSpectrum(depth, cos_theta, multiplicity, radius).GetLog(log_energy);
}

double
BMSSEnergyDistribution::GetMaxRadius() const
{
	return 250*I3Units::m;
}

std::vector<std::pair<double,double> >
BMSSEnergyDistribution::Generate(I3RandomService &rng, double depth,
    double cos_theta, unsigned multiplicity, unsigned samples) const
{
	std::vector<std::pair<double,double> > values;
	values.reserve(samples);
	std::pair<double, double> val;
	for (unsigned i=0; i < samples; i++) {
		val.first = BMSSRadialDistribution().Generate(rng, depth, cos_theta, multiplicity);
		val.second = GetSpectrum(depth, cos_theta, multiplicity, val.first).Generate(rng);
		values.push_back(val);
	}

	return values;
}

OffsetPowerLaw::OffsetPowerLaw() : gamma_(NAN), offset_(NAN), emin_(NAN), emax_(NAN)
{}

OffsetPowerLaw::OffsetPowerLaw(double gamma, double offset, double emin, double emax)
    : gamma_(gamma), offset_(offset), emin_(emin), emax_(emax)
{
	if (gamma <= 0)
		log_fatal("Power law index must be > 0");
	else if (gamma == 1) {
		nmin_ = std::log(emin + offset);
		nmax_ = std::log(emax + offset);
		norm_ = 1./(nmax_ - nmin_);
	} else {
		nmin_ = std::pow(emin + offset, 1-gamma);
		nmax_ = std::pow(emax + offset, 1-gamma);
		norm_ = (1-gamma)/(nmax_ - nmin_);
	}
	lognorm_ = std::log(norm_);
}

bool
OffsetPowerLaw::operator==(const OffsetPowerLaw &other) const
{
	return (gamma_ == other.gamma_ && offset_ == other.offset_
	    && emin_ == other.emin_ && emax_ == other.emax_ );
}

double
OffsetPowerLaw::operator()(double energy) const
{
	if (energy <= emax_ && energy >= emin_)
		return norm_*std::pow(energy + offset_, -gamma_);
	else
		return 0.;
}

double
OffsetPowerLaw::GetLog(double energy) const
{
	if (energy <= emax_ && energy >= emin_)
		return lognorm_ - gamma_*std::log(energy + offset_);
	else
		return -std::numeric_limits<double>::infinity();
}

double
OffsetPowerLaw::GetLog(EnergyDistribution::log_value log_energy) const
{
	return GetLog(std::exp(log_energy));
}

double
OffsetPowerLaw::Generate(I3RandomService &rng) const
{
	return InverseSurvivalFunction(rng.Uniform());
}

double
OffsetPowerLaw::InverseSurvivalFunction(double p) const
{
	if (gamma_ == 1)
		return std::exp((1-p)*(nmax_ - nmin_) + nmin_) - offset_;
	else
		return std::pow((1-p)*(nmax_ - nmin_) + nmin_, 1./(1.-gamma_)) - offset_;
}

template <typename Archive>
void
EnergyDistribution::serialize(Archive &ar __attribute__ ((unused)), unsigned version __attribute__ ((unused)))
{}

template <typename Archive>
void
SplineEnergyDistribution::serialize(Archive &ar, unsigned version)
{
	if (version > 0)
		log_fatal_stream("Version "<<version<<" is from the future");

	ar & make_nvp("EnergyDistribution", base_object<EnergyDistribution>(*this));
	ar & make_nvp("SingleEnergy", singles_);
	ar & make_nvp("BundleEnergy", bundles_);
}

template <typename Archive>
void
BMSSEnergyDistribution::serialize(Archive &ar, unsigned version)
{
	if (version > 0)
		log_fatal_stream("Version "<<version<<" is from the future");

	ar & make_nvp("EnergyDistribution", base_object<EnergyDistribution>(*this));
}

template <typename Archive>
void
OffsetPowerLaw::serialize(Archive &ar, unsigned version)
{
	if (version > 0)
		log_fatal_stream("Version "<<version<<" is from the future");

	ar & make_nvp("Gamma", gamma_);
	ar & make_nvp("Offset", offset_);
	ar & make_nvp("MinEnergy", emin_);
	ar & make_nvp("MaxEnergy", emax_);

	*this = OffsetPowerLaw(gamma_, offset_, emin_, emax_);
}

}

I3_SERIALIZABLE(I3MuonGun::EnergyDistribution);
I3_SERIALIZABLE(I3MuonGun::SplineEnergyDistribution);
I3_SERIALIZABLE(I3MuonGun::BMSSEnergyDistribution);
I3_SERIALIZABLE(I3MuonGun::OffsetPowerLaw);


#include <thread>
#include <vector>
#include <armadillo>
#include "RHF.hpp"
#include "esp.hpp"

//
// Electrostatic Potential (ESP)
// : calculates partial atomic charges that fit
//   the quantum mechanical electrostatic potential on the selected grid points
//

using namespace std;
using namespace libint2;

namespace willow { namespace qcmol {

static unsigned int num_threads;

static void esp_esp (const int thread_id,
		     const vector<Atom>& atoms,
		     const BasisSet& bs,
		     const arma::mat& Dm,
		     const vector<arma::vec3>& grids,
		     arma::vec& results);

ESP::ESP (const vector<Atom>& atoms,
	  const BasisSet& bs,
	  const Integrals& ints,
	  const int qm_chg,
	  const bool l_print,
	  const vector<QAtom>& atoms_Q)
  : RHF (atoms, bs, ints, qm_chg, l_print, atoms_Q)
{

  num_threads = std::thread::hardware_concurrency();

  if (num_threads == 0)
    num_threads = 1;

  const arma::mat Dm = densityMatrix();

  esp_grid (atoms);

  // ---- Thread
  vector<arma::vec>   grid_t(num_threads);
  vector<std::thread> t_esp (num_threads);

  for (auto id = 0; id < num_threads; ++id) {
    t_esp[id] = std::thread (esp_esp, id,
			     std::cref(atoms),
			     std::cref(bs),
			     std::cref(Dm),
			     std::cref(m_grids),
			     std::ref(grid_t[id]));
  }

  // Join
  for (auto id = 0; id < num_threads; ++id) {
    t_esp[id].join();
  }

  m_grids_val = grid_t[0];

  for (auto id = 1; id < num_threads; ++id) {
    m_grids_val += grid_t[id];
  }
  //---- (done: Thread) ---

  m_qf = esp_fit (atoms);

  if (l_print) {
    for (auto i = 0; i < atoms.size(); i++)
      cout << "ESP " << i+1 << "  " << m_qf(i) << endl;
  }

}

void esp_esp (const int thread_id,
	      const vector<Atom>& atoms,
	      const BasisSet& bs,
	      const arma::mat& Dm,
	      const vector<arma::vec3>& grids,
	      arma::vec& result)
{

  const auto nsize   = grids.size();

  result = arma::vec(nsize, arma::fill::zeros);

  const auto nshells = bs.size();
  const auto nbf     = bs.nbf();

  libint2::Engine engine (libint2::Operator::nuclear,
			  bs.max_nprim(),
			  bs.max_l(),
			  0);

  const auto& buf = engine.results();

  const auto shell2bf = bs.shell2bf ();

  arma::mat Vm(nbf,nbf);

  for (auto ig = 0; ig < grids.size(); ++ig) {

    if (ig % num_threads != thread_id) continue;

    const auto xg = grids[ig](0);
    const auto yg = grids[ig](1);
    const auto zg = grids[ig](2);

    // Interaction with electron density
    vector<pair<double,array<double,3>>> q;
    q.push_back ( {1.0, {{xg, yg, zg}}} );

    engine.set_params (q);

    // calc Vm
    Vm.zeros();

    for (auto s1 = 0; s1 != nshells; ++s1) {
      auto bf1 = shell2bf[s1];
      auto nbf1 = bs[s1].size();

      for (auto s2 = 0; s2 != nshells; ++s2) {
	auto bf2 = shell2bf[s2];
	auto nbf2 = bs[s2].size();

	engine.compute (bs[s1], bs[s2]);
	const auto* buf0 = buf[0];

	for (auto ib = 0, ij = 0; ib < nbf1; ib++) {
	  for (auto jb = 0; jb < nbf2; jb++, ij++) {
	    const double val = buf0[ij];
	    Vm(bf1+ib,bf2+jb) = val;
	    Vm(bf2+jb,bf1+ib) = val;
	  }
	} // ib

      } // s2
    }  // s1

    //
    // get electrostatic potential on the grid points
    // -- from electron density
    double gval = 0.0;

    for (auto i = 0; i < nbf; i++)
      for (auto j = 0; j < nbf; j++)
	gval += Dm(i,j)*Vm(i,j);

    // get electrostatic potential on the grid points
    // -- from nuclei
    auto zval = 0.0;

    for (auto ia = 0; ia < atoms.size(); ++ia) {
      auto xij = atoms[ia].x - xg;
      auto yij = atoms[ia].y - yg;
      auto zij = atoms[ia].z - zg;
      auto r2  = xij*xij + yij*yij + zij*zij;
      auto r   = sqrt(r2);
      zval += static_cast<double>(atoms[ia].atomic_number)/r;
    }

    result(ig) = gval + zval;
  }

}

arma::vec ESP::esp_fit (const vector<Atom>& atoms)
{

  // set up matrix of linear coefficients
  auto natoms = atoms.size();
  auto ndim   = natoms+1;

  arma::mat am(ndim,ndim);
  arma::vec bv(ndim);

  am.zeros();
  bv.zeros();

  for (auto i = 0; i < atoms.size(); ++i) {
    auto xi = atoms[i].x;
    auto yi = atoms[i].y;
    auto zi = atoms[i].z;

    for (auto j = i; j < atoms.size(); ++j) {
      auto xj = atoms[j].x;
      auto yj = atoms[j].y;
      auto zj = atoms[j].z;

      auto sum = 0.0;

      for (auto k = 0; k < m_grids.size(); ++k) {
	auto xg = m_grids[k](0);
	auto yg = m_grids[k](1);
	auto zg = m_grids[k](2);

	auto rig2 = (xi-xg)*(xi-xg) + (yi-yg)*(yi-yg) + (zi-zg)*(zi-zg);
	auto rjg2 = (xj-xg)*(xj-xg) + (yj-yg)*(yj-yg) + (zj-zg)*(zj-zg);

	sum += 1.0/sqrt(rig2*rjg2);
      }

      am(i,j) = sum;
      am(j,i) = sum;
    }

    am(i,natoms) = 1.0;
    am(natoms,i) = 1.0;
  }

  // construct column vector b

  for (auto i = 0; i < atoms.size(); ++i) {
    auto xi = atoms[i].x;
    auto yi = atoms[i].y;
    auto zi = atoms[i].z;

    auto sum = 0.0;
    for (auto k = 0; k < m_grids.size(); ++k) {
      auto xg = m_grids[k](0);
      auto yg = m_grids[k](1);
      auto zg = m_grids[k](2);
      auto val = m_grids_val(k);

      auto rig2 = (xi-xg)*(xi-xg) + (yi-yg)*(yi-yg) + (zi-zg)*(zi-zg);

      sum += val/sqrt(rig2);
    }
    bv(i) = sum;
  }
  bv(natoms) = 0.0; // b(natoms) = charge;

  arma::mat am_inv = am.i();

  arma::vec qf(natoms);
  qf.zeros();

  for (auto i = 0; i < natoms; ++i) {
    auto sum = 0.0;

    for (auto j = 0; j < ndim; ++j) {
    //for (auto j = 0; j < atoms.size(); ++j) {
      sum = sum + am_inv(i,j)*bv(j);
    }
    qf(i) = sum;
    //cout << " qf " << sum << endl;
  }

  return qf;

}

void ESP::esp_grid (const vector<Atom>& atoms)
{

  vector<double> radius(11);

  // A
  radius[0] = 0.0; // Ghost
  radius[1] = 0.3; // H
  radius[2] = 1.22;// He
  radius[3] = 1.23;// Li
  radius[4] = 0.89; // Be
  radius[5] = 0.88; // B
  radius[6] = 0.77; // C
  radius[7] = 0.70; // N
  radius[8] = 0.66; // O
  radius[9] = 0.58; // F
  radius[10]= 1.60; // Ne

  for (auto i = 1; i <= 10; i++) {
    radius[i] = (radius[i] + 0.7)*ang2bohr;
  }
  double grd_min_x = atoms[0].x;
  double grd_min_y = atoms[0].y;
  double grd_min_z = atoms[0].z;

  double grd_max_x = atoms[0].x;
  double grd_max_y = atoms[0].y;
  double grd_max_z = atoms[0].z;

  for (auto ia = 1; ia < atoms.size(); ++ia) {

    grd_min_x = min(grd_min_x, atoms[ia].x);
    grd_min_y = min(grd_min_y, atoms[ia].y);
    grd_min_z = min(grd_min_z, atoms[ia].z);

    grd_max_x = max(grd_max_x, atoms[ia].x);
    grd_max_y = max(grd_max_y, atoms[ia].y);
    grd_max_z = max(grd_max_z, atoms[ia].z);

  }
  //
  // calculate the grid size
  //
  // A --> au
  const double rcut = 3.0*ang2bohr;
  const double rcut2= rcut*rcut;
  const double spac = 0.5*ang2bohr;

  const int ngrid_x = int ((grd_max_x - grd_min_x + 2.0*rcut)/spac) + 1;
  const int ngrid_y = int ((grd_max_y - grd_min_y + 2.0*rcut)/spac) + 1;
  const int ngrid_z = int ((grd_max_z - grd_min_z + 2.0*rcut)/spac) + 1;

  const auto small = 1.0e-8;
  arma::vec3 vg;

  for(auto iz = 0; iz <  ngrid_z; ++iz)
    for(auto iy = 0; iy <  ngrid_y; ++iy)
      for(auto ix = 0; ix <  ngrid_x; ++ix) {

	vg.zeros();

	vg(0) = grd_min_x - rcut + ix*spac;
	vg(1) = grd_min_y - rcut + iy*spac;
	vg(2) = grd_min_z - rcut + iz*spac;

	double dmin = rcut2;

	bool lupdate = true;

	for (auto ia = 0; ia < atoms.size(); ++ia) {
	  int iz    = atoms[ia].atomic_number;
	  auto rad2 = radius[iz]*radius[iz];

	  auto dx = vg(0) - atoms[ia].x;
	  auto dy = vg(1) - atoms[ia].y;
	  auto dz = vg(2) - atoms[ia].z;

	  auto d2  = dx*dx + dy*dy + dz*dz;

	  if ( (rad2 - d2) > small) {
	    lupdate = false;
	    break;
	  }

	  if ( (dmin-d2) > small) dmin = d2;

	}

	if (lupdate) {
	  if ( (rcut2 - dmin) > small) {
	    m_grids.push_back (vg);
	  }
	}

      }

}

} } // namespace willow::qcmol


//==============================================================================
//         Copyright 2003 - 2012 LASMEA UMR 6602 CNRS/Univ. Clermont II
//         Copyright 2009 - 2012 LRI    UMR 8623 CNRS/Univ Paris Sud XI
//
//          Distributed under the Boost Software License, Version 1.0.
//                 See accompanying file LICENSE.txt or copy at
//                     http://www.boost.org/LICENSE_1_0.txt
//==============================================================================
#ifndef BOOST_SIMD_TOOLBOX_ARITHMETIC_FUNCTIONS_SIMD_VMX_ALTIVEC_REC_HPP_INCLUDED
#define BOOST_SIMD_TOOLBOX_ARITHMETIC_FUNCTIONS_SIMD_VMX_ALTIVEC_REC_HPP_INCLUDED
#ifdef BOOST_SIMD_HAS_VMX_SUPPORT

#include <boost/simd/toolbox/arithmetic/functions/rec.hpp>
#include <boost/dispatch/meta/as_floating.hpp>
#include <boost/simd/include/constants/one.hpp>
#include <boost/simd/include/constants/mzero.hpp>
#include <boost/simd/include/constants/inf.hpp>
#include <boost/simd/include/functions/copysign.hpp>
#include <boost/simd/include/functions/if_else.hpp>
#include <boost/simd/include/functions/is_eqz.hpp>

namespace boost { namespace simd { namespace ext
{
  BOOST_SIMD_FUNCTOR_IMPLEMENTATION( boost::simd::tag::rec_
                                   , boost::simd::tag::altivec_
                                   , (A0)
                                   , ((simd_<floating_<A0>, boost::simd::tag::altivec_>))
                                   )
  {
    typedef typename dispatch::meta::as_floating<A0>::type result_type;

    BOOST_SIMD_FUNCTOR_CALL(1)
    {
      result_type erb   = vec_re(a0());
      result_type vmadd = vec_madd( vec_nmsub(erb(), a0(), One<result_type>()())
                                             , erb()
                                             , erb()
                                             );
      result_type rec_b  = if_else( is_eqz(a0)
                                  , copysign(Inf<result_type>(),a0)
                                  , vmadd
                                  );
      return if_else(is_eqz(erb), copysign(erb,a0), rec_b);
    }
  };
} } }
#endif
#endif


// Copyright (c) 2015-2018 Daniel Cooke
// Use of this source code is governed by the MIT license that can be found in the LICENSE file.

#include "phaser.hpp"

#include <deque>
#include <algorithm>
#include <numeric>
#include <cstddef>
#include <cmath>
#include <utility>
#include <iostream>

#include <boost/functional/hash.hpp>

#include "utils/mappable_algorithms.hpp"
#include "utils/maths.hpp"

#include "timers.hpp"

namespace octopus {

Phaser::Phaser(Phred<double> min_phase_score) : min_phase_score_ {min_phase_score} {}

namespace {

using GenotypeReference    = std::reference_wrapper<const Genotype<Haplotype>>;
using PhaseComplementSet   = std::deque<GenotypeReference>;
using PhaseComplementSets  = std::vector<PhaseComplementSet>;
using PartitionIterator    = std::vector<GenomicRegion>::const_iterator;
using GenotypeChunkVector = std::vector<Genotype<Haplotype>>;

struct GenotypeChunkPartitions
{
    GenotypeReference genotype;
    GenotypeChunkVector chunks;
};

using GenotypeChunkTable = std::vector<GenotypeChunkPartitions>;

template <typename Container>
GenotypeChunkTable
make_chunk_table(const Container& genotypes, PartitionIterator first_parition, PartitionIterator last_partition)
{
    const auto num_partitions = std::distance(first_parition, last_partition);
    GenotypeChunkTable result {};
    result.reserve(genotypes.size());
    for (const auto& genotype : genotypes) {
        result.push_back({std::cref(genotype), GenotypeChunkVector {}});
        result.back().chunks.reserve(num_partitions);
    }
    std::for_each(first_parition, last_partition, [&] (const auto& region) {
        auto chunks = copy_each<Haplotype>(genotypes, region);
        for (std::size_t i {0}; i < genotypes.size(); ++i) {
            result[i].chunks.push_back(std::move(chunks[i]));
        }
    });
    return result;
}

struct GenotypeChunkVectorHash
{
    auto operator()(const GenotypeChunkVector& chunks) const noexcept
    {
        std::size_t result {};
        for (const auto& genotype : chunks) {
            using boost::hash_combine;
            hash_combine(result, std::hash<Genotype<Haplotype>>{}(genotype));
        }
        return result;
    }
};

using PhaseSetMap = std::unordered_map<GenotypeChunkVector, PhaseComplementSet, GenotypeChunkVectorHash>;

PhaseComplementSets make_phase_sets(const GenotypeChunkTable& table)
{
    PhaseSetMap phase_sets {};
    phase_sets.reserve(table.size());
    for (const auto& row : table) {
        phase_sets[row.chunks].push_back(row.genotype);
    }
    PhaseComplementSets result {};
    result.reserve(phase_sets.size());
    for (auto&& p : phase_sets) {
        result.push_back(std::move(p.second));
    }
    return result;
}

template <typename Container>
PhaseComplementSets
generate_phase_complement_sets(const Container& genotypes, PartitionIterator first_parition, PartitionIterator last_partition)
{
    const auto table = make_chunk_table(genotypes, first_parition, last_partition);
    return make_phase_sets(table);
}

template <typename Map>
double marginalise(const PhaseComplementSet& phase_set, const Map& genotype_posteriors)
{
    return std::accumulate(std::cbegin(phase_set), std::cend(phase_set), 0.0,
                           [&] (const auto curr, const auto& genotype) {
                               return curr + genotype_posteriors.at(genotype);
                           });
}

double maximum_entropy(const std::size_t num_elements)
{
    return std::log2(num_elements);
}

template <typename Map>
double calculate_entropy(const PhaseComplementSet& phase_set, const Map& genotype_posteriors)
{
    const auto norm = marginalise(phase_set, genotype_posteriors);
    if (norm <= 0.0) {
        // if norm ~= 0 then every element in the phase must must have probability ~= 0, so it
        // just looks like a uniform distirbution
        return maximum_entropy(phase_set.size());
    }
    return std::max(0.0, -std::accumulate(std::cbegin(phase_set), std::cend(phase_set), 0.0,
                                          [&genotype_posteriors, norm] (const auto curr, const auto& genotype) {
                                              const auto p = genotype_posteriors.at(genotype) / norm;
                                              return curr + p * std::log2(p);
                                          }));
}

template <typename Map>
double calculate_relative_entropy(const PhaseComplementSet& phase_set, const Map& genotype_posteriors)
{
    if (phase_set.size() < 2) return 1.0;
    return 1.0 - calculate_entropy(phase_set, genotype_posteriors) / maximum_entropy(phase_set.size());
}

template <typename Map>
auto calculate_phase_score(const PhaseComplementSet& phase_set, const Map& genotype_posteriors)
{
    return marginalise(phase_set, genotype_posteriors) * calculate_relative_entropy(phase_set, genotype_posteriors);
}

template <typename Map>
Phred<double> calculate_phase_score(const PhaseComplementSets& phase_sets, const Map& genotype_posteriors)
{
    return Phred<double> { Phred<double>::Probability {
        std::max(0.0, 1.0 - std::accumulate(std::cbegin(phase_sets), std::cend(phase_sets), 0.0,
                                           [&] (const auto curr, const auto& phase_set) {
                                               return curr + calculate_phase_score(phase_set, genotype_posteriors);
                                           }))
    }};
}

auto min_phase_score(const double p)
{
    if (maths::almost_one(p)) {
        static const Phred<double> max_possible_score {Phred<double>::Probability {0.0}};
        return max_possible_score;
    } else {
        const auto e = -(p * std::log2(p) + (1.0 - p) * std::log2(1.0 - p)) / maximum_entropy(2);
        return Phred<double> {Phred<double>::Probability {std::max(e, 0.0)}};
    }
}

auto min_phase_score(const Genotype<Haplotype>& called_genotype, const Phaser::SampleGenotypePosteriorMap& genotype_posteriors)
{
    return min_phase_score(genotype_posteriors[called_genotype]);
}

std::vector<GenotypeReference> extract_genotypes(const Phaser::GenotypePosteriorMap& genotype_posteriors)
{
    return extract_key_refs(genotype_posteriors);
}

using GenotypeChunkPosteriorMap = std::unordered_map<Genotype<Haplotype>, double>;

template <typename Container>
auto copy_and_marginalise(const Container& genotypes,
                          const Phaser::SampleGenotypePosteriorMap& genotype_posteriors,
                          const GenomicRegion& region)
{
    auto chunks = copy_each<Haplotype>(genotypes, region);
    GenotypeChunkPosteriorMap chunk_posteriors {genotypes.size()};
    auto chunk_posteriors_itr = std::cbegin(chunks);
    for (const auto& genotype : genotypes) {
        chunk_posteriors[*chunk_posteriors_itr++] += genotype_posteriors[genotype];
    }
    std::sort(std::begin(chunks), std::end(chunks), GenotypeLess {});
    chunks.erase(std::unique(std::begin(chunks), std::end(chunks)), std::end(chunks));
    return std::make_pair(std::move(chunks), std::move(chunk_posteriors));
}

} // namespace

boost::optional<Phaser::PhaseSet>
Phaser::try_phase(const std::vector<Haplotype>& haplotypes,
                  const GenotypePosteriorMap& genotype_posteriors,
                  const std::vector<GenomicRegion>& regions) const
{
    // TODO
    return boost::none;
}

Phaser::PhaseSet::SamplePhaseRegions
force_phase_sample(const GenomicRegion& region,
                   const std::vector<GenomicRegion>& partitions,
                   const std::vector<GenotypeReference>& genotypes,
                   const Phaser::SampleGenotypePosteriorMap& genotype_posteriors,
                   const Phred<double> min_phase_score)
{
    auto first_partition = std::cbegin(partitions);
    auto last_partition  = std::cend(partitions);
    auto phase_set = generate_phase_complement_sets(genotypes, first_partition, last_partition);
    auto phase_score = calculate_phase_score(phase_set, genotype_posteriors);
    if (phase_score >= min_phase_score) {
        return {Phaser::PhaseSet::PhaseRegion {region, phase_score}};
    }
    Phaser::PhaseSet::SamplePhaseRegions result {};
    --last_partition;
    std::vector<Genotype<Haplotype>> chunks;
    GenotypeChunkPosteriorMap chunk_posteriors;
    while (first_partition != std::cend(partitions)) {
        auto curr_region = encompassing_region(first_partition, last_partition);
        std::tie(chunks, chunk_posteriors) = copy_and_marginalise(genotypes, genotype_posteriors, curr_region);
        phase_set = generate_phase_complement_sets(chunks, first_partition, last_partition);
        phase_score = calculate_phase_score(phase_set, chunk_posteriors);
        if (phase_score >= min_phase_score || std::distance(first_partition, last_partition) == 1) {
            result.emplace_back(encompassing_region(first_partition, last_partition), phase_score);
            first_partition = last_partition;
            last_partition  = std::cend(partitions);
        } else {
            --last_partition;
        }
    }
    return result;
}

Phaser::PhaseSet
Phaser::force_phase(const std::vector<Haplotype>& haplotypes,
                    const GenotypePosteriorMap& genotype_posteriors,
                    const std::vector<GenomicRegion>& regions,
                    boost::optional<GenotypeCallMap> genotype_calls) const
{
    assert(!haplotypes.empty());
    assert(!genotype_posteriors.empty1());
    assert(!genotype_posteriors.empty2());
    assert(std::is_sorted(std::cbegin(regions), std::cend(regions)));
    const auto& haplotype_region = haplotypes.front().mapped_region();
    const auto genotypes = extract_genotypes(genotype_posteriors);
    const auto partitions = extract_covered_regions(regions);
    PhaseSet result {haplotype_region};
    result.phase_regions.reserve(genotype_posteriors.size1());
    if (genotypes.front().get().ploidy() == 1 || partitions.size() == 1) {
        for (const auto& p : genotype_posteriors) {
            if (max_phase_score_) {
                result.phase_regions[p.first].emplace_back(haplotype_region, *max_phase_score_);
            } else {
                static const Phred<double> max_possible_score {Phred<double>::Probability {0.0}};
                result.phase_regions[p.first].emplace_back(haplotype_region, max_possible_score);
            }
        }
        return result;
    }
    for (const auto& p : genotype_posteriors) {
        if (genotype_calls && max_phase_score_ && min_phase_score(genotype_calls->at(p.first), p.second) >= *max_phase_score_) {
            result.phase_regions[p.first].emplace_back(haplotype_region, *max_phase_score_);
        } else {
            auto phases = force_phase_sample(haplotype_region, partitions, genotypes, p.second, min_phase_score_);
            if (max_phase_score_) {
                for (auto& phase : phases) phase.score = std::min(phase.score, *max_phase_score_);
            }
            result.phase_regions.emplace(p.first, std::move(phases));
        }
    }
    return result;
}

// non-member methods

bool is_split_phasing(const Phaser::PhaseSet& phase)
{
    return std::any_of(std::cbegin(phase.phase_regions), std::cend(phase.phase_regions),
                       [] (const auto& p) { return p.second.size() > 1; });
}

namespace debug {

void print_phase_sets(const Phaser::PhaseSet& phasings)
{
    print_phase_sets(std::cout, phasings);
}

} // namespace debug

} // namespace Ocotpus


#include <iostream>
#include <string>
#include <array>

#include <Eigen/Core>

#include <franka_typing_aid/examples_common.h>

struct test_struct
{
    int a;
    std::string b;
    std::array<double, 6> c;
};

int main()
{
    test_struct str;
    str.a = 1;
    str.b = "fdsf";
    str.c = {{4, 7.1, 6, 9.9, 2, 34}};

    Eigen::Vector3d test_vec(str.c.data());

    std::cout << test_vec << std::endl;
    std::cout << test_vec.norm() << std::endl;

    return 0;
}