/** $Id: EnergyDependentSurfaceInjector.cxx 136127 2015-08-12 13:54:32Z jvansanten $
 * @file
 * @author Jakob van Santen <vansanten@wisc.edu>
 *
 * $Revision: 136127 $
 * $Date: 2015-08-12 07:54:32 -0600 (Wed, 12 Aug 2015) $
 */

#include <MuonGun/EnergyDependentSurfaceInjector.h>
#include <MuonGun/I3MuonGun.h>
#include <MuonGun/SamplingSurface.h>
#include <MuonGun/Cylinder.h>
#include <MuonGun/Flux.h>
#include <MuonGun/EnergyDistribution.h>
#include <MuonGun/RadialDistribution.h>
#include <dataclasses/I3Constants.h>
#include <dataclasses/physics/I3Particle.h>
#include <dataclasses/physics/I3MCTreeUtils.h>
#include <phys-services/I3RandomService.h>

#include <boost/bind.hpp>
#include <boost/foreach.hpp>

namespace I3MuonGun {

EnergyDependentSurfaceInjector::EnergyDependentSurfaceInjector(CylinderPtr surface, FluxPtr flux,
    boost::shared_ptr<OffsetPowerLaw> energies, RadialDistributionPtr radius, SurfaceScalingFunctionPtr scaling)
    : StaticSurfaceInjector(surface, flux, energies, radius), scalingFunction_(scaling)
{
}

GenerationProbabilityPtr
EnergyDependentSurfaceInjector::Clone() const
{
	return boost::make_shared<EnergyDependentSurfaceInjector>(*this);
}

bool
EnergyDependentSurfaceInjector::IsCompatible(GenerationProbabilityConstPtr o) const
{
	if (!StaticSurfaceInjector::IsCompatible(o))
		return false;
	boost::shared_ptr<const EnergyDependentSurfaceInjector> other
	    = boost::dynamic_pointer_cast<const EnergyDependentSurfaceInjector>(o);
	if (!other)
		return false;
	else
		return *scalingFunction_ == *(other->scalingFunction_);
}

void
EnergyDependentSurfaceInjector::Generate(I3RandomService &rng, I3MCTree &tree,
    BundleConfiguration &bundle) const
{
	I3Direction dir;
	I3Position pos;
	unsigned m;
	double flux;
	double maxflux = (*flux_)(surface_->GetMinDepth(), 1., flux_->GetMinMultiplicity());
	double h, coszen;
	SamplingSurfaceConstPtr surface;
	do {
		// Choose a multiplicity
		bundle.clear();
		m = rng.Integer(flux_->GetMaxMultiplicity() - flux_->GetMinMultiplicity())
		    + flux_->GetMinMultiplicity();
		// Choose an ensemble of energies
		for (unsigned i=0; i < m; i++)
			bundle.push_back(BundleEntry(0., energyGenerator_->Generate(rng)));
		bundle.sort();

		// Choose target surface based on highest-energy muon
		surface = GetTargetSurface(bundle.front().energy);
		// Sample an impact point on the target surface
		surface->SampleImpactRay(pos, dir, rng);
		h = GetDepth(pos.GetZ());
		coszen = cos(dir.GetZenith());

		// Snap the impact point back to the injection surface
		std::pair<double, double> steps = surface_->GetIntersection(pos, dir);
		if (!(steps.first <= 0))
			log_fatal("The target point is outside the injection surface!");
		pos.SetX(pos.GetX() + steps.first*dir.GetX());
		pos.SetY(pos.GetY() + steps.first*dir.GetY());
		pos.SetZ(pos.GetZ() + steps.first*dir.GetZ());

		// Accept or reject the chosen zenith angle and multiplicity
		flux = (*flux_)(surface_->GetMinDepth(), cos(dir.GetZenith()), m);
	} while (rng.Uniform(0., maxflux) > flux);

	I3Particle primary;
	primary.SetPos(pos);
	primary.SetDir(dir);
	primary.SetShape(I3Particle::Primary);
	primary.SetLocationType(I3Particle::Anywhere);
	primary.SetType(I3Particle::unknown);
	primary.SetTime(0.);
	I3MCTreeUtils::AddPrimary(tree, primary);

	// For each muon, draw a radial offset and add an entry
	// to the MCTree
	BOOST_FOREACH(BundleConfiguration::value_type &bspec, bundle) {
		double radius = 0., azimuth = 0.;
		if (m > 1u) {
			radius = radialDistribution_->Generate(rng, h, coszen, m);
			azimuth = rng.Uniform(0., 2*M_PI);
		}

		I3Particle track = CreateParallelTrack(radius, azimuth, *surface, primary);
		track.SetEnergy(bspec.energy);
		bspec.radius = radius;
		I3MCTreeUtils::AppendChild(tree, primary, track);
	}
}

SamplingSurfacePtr
EnergyDependentSurfaceInjector::GetTargetSurface(double energy) const
{
	if (scalingFunction_)
		return scalingFunction_->GetSurface(energy);
	else
		return boost::make_shared<Cylinder>(1600, 800);
}

double
EnergyDependentSurfaceInjector::GetTotalRate(SamplingSurfaceConstPtr surface) const
{
	double total_rate = 0.;
	for (unsigned m = flux_->GetMinMultiplicity(); m <= flux_->GetMaxMultiplicity(); m++)
		total_rate += surface->IntegrateFlux(boost::bind(boost::cref(*flux_), _1, _2, m));
	return total_rate;
}

double
EnergyDependentSurfaceInjector::GetLogGenerationProbability(const I3Particle &axis,
    const BundleConfiguration &bundlespec) const
{
	// Entries are sorted in descending order of energy, so the
	// "minimum" entry has the maximum energy
	SamplingSurfaceConstPtr surface =
	    GetTargetSurface(std::min_element(bundlespec.begin(), bundlespec.end())->energy);
	std::pair<double, double> steps =
	    surface->GetIntersection(axis.GetPos(), axis.GetDir());
	// This shower axis doesn't intersect the target surface. Bail.
	if (!std::isfinite(steps.first))
		return -std::numeric_limits<double>::infinity();

	double h = GetDepth(axis.GetPos().GetZ() + steps.first*axis.GetDir().GetZ());
	double coszen = cos(axis.GetDir().GetZenith());
	unsigned m = static_cast<unsigned>(bundlespec.size());

	// We used the flux to do rejection sampling in zenith and multiplicity. Evaluate
	// the properly-normalized PDF here.
	double logprob = flux_->GetLog(surface_->GetMinDepth(), coszen, m) - GetZenithNorm();
	BOOST_FOREACH(const BundleEntry &track, bundlespec) {
		if (m > 1)
			logprob += radialDistribution_->GetLog(h, coszen, m, track.radius);
		logprob += energyGenerator_->GetLog(track.energy);
	}

	// We only distributed events over the target surface, not the entire injection surface
	return logprob - std::log(surface->GetAcceptance());
}

template <typename Archive>
void
EnergyDependentSurfaceInjector::serialize(Archive &ar, unsigned)
{
	ar & make_nvp("StaticSurfaceInjector", base_object<StaticSurfaceInjector>(*this));
	ar & make_nvp("ScalingFunction", scalingFunction_);
}

SurfaceScalingFunction::~SurfaceScalingFunction() {}

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

ConstantSurfaceScalingFunction::ConstantSurfaceScalingFunction() {}
ConstantSurfaceScalingFunction::ConstantSurfaceScalingFunction(SamplingSurfacePtr surface) : surface_(surface) {}
ConstantSurfaceScalingFunction::~ConstantSurfaceScalingFunction() {}

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

	ar & make_nvp("SurfaceScalingFunction", base_object<SurfaceScalingFunction>(*this));
	ar & make_nvp("Surface", surface_);
}

SamplingSurfacePtr
ConstantSurfaceScalingFunction::GetSurface(double energy __attribute__((unused))) const { return surface_; }

bool
ConstantSurfaceScalingFunction::operator==(const SurfaceScalingFunction &o) const
{
	const ConstantSurfaceScalingFunction *other = dynamic_cast<const ConstantSurfaceScalingFunction*>(&o);
	if (!other)
		return false;
	else
		return (*surface_ == *(other->surface_));
}

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

	ar & make_nvp("SurfaceScalingFunction", base_object<SurfaceScalingFunction>(*this));
	ar & make_nvp("Scale", scale_);
	ar & make_nvp("EnergyScale", energyScale_);
	ar & make_nvp("Offset", offset_);
	ar & make_nvp("Power", power_);
	ar & make_nvp("RadiusBounds", rBounds_);
	ar & make_nvp("ZBounds", zBounds_);
	ar & make_nvp("CenterBounds", centerBounds_);
}

BasicSurfaceScalingFunction::BasicSurfaceScalingFunction() :
    scale_(800., 240957.5), energyScale_(4., 4.), offset_(3.778, 3.622), power_(1.10, 2.23),
    rBounds_(0., 525.), zBounds_(-500., 400.),
    centerBounds_(pair(46.29,-34.88), pair(31.25, 19.64))
{}

BasicSurfaceScalingFunction::~BasicSurfaceScalingFunction() {}

double
BasicSurfaceScalingFunction::GetMargin(double logenergy, double scale, double offset, double power) const
{
	if (logenergy < offset)
		return pow(scale*(offset - logenergy), 1./power);
	else
		return 0.;
}

SamplingSurfacePtr
BasicSurfaceScalingFunction::GetSurface(double energy) const
{
	// Shrink the cylinder down by an energy-dependent amount
	double z = std::max(zBounds_.second -
	    GetMargin(std::log10(energy/energyScale_.first), scale_.first, offset_.first, power_.first), zBounds_.first);

	// Shrink the sides of the cylinder
	double r = std::max(rBounds_.second -
	    GetMargin(std::log10(energy/energyScale_.second), scale_.second, offset_.second, power_.second), rBounds_.first);

	// Move the center smoothly between the configured bounds
	double hscale = r/(rBounds_.second-rBounds_.first);
	I3Position center(centerBounds_.first.first + hscale*(centerBounds_.second.first-centerBounds_.first.first),
	    centerBounds_.first.second + hscale*(centerBounds_.second.second-centerBounds_.first.second),
	    (zBounds_.first + z)/2.);

	return boost::make_shared<Cylinder>(z-zBounds_.first, r, center);
}

void
BasicSurfaceScalingFunction::SetCapScaling(double energyScale, double scale, double offset, double power)
{
	energyScale_.first = energyScale;
	scale_.first = scale;
	offset_.first = offset;
	power_.first = power;
}

void
BasicSurfaceScalingFunction::SetSideScaling(double energyScale, double scale, double offset, double power)
{
	energyScale_.second = energyScale;
	scale_.second = scale;
	offset_.second = offset;
	power_.second = power;
}

void
BasicSurfaceScalingFunction::SetRadiusBounds(double rmin, double rmax)
{
	rBounds_ = std::make_pair(rmin, rmax);
}

void
BasicSurfaceScalingFunction::SetZBounds(double zmin, double zmax)
{
	zBounds_ = std::make_pair(zmin, zmax);
}

void
BasicSurfaceScalingFunction::SetCenterMin(double x, double y)
{
	centerBounds_.first = std::make_pair(x, y);
}

void
BasicSurfaceScalingFunction::SetCenterMax(double x, double y)
{
	centerBounds_.second = std::make_pair(x, y);
}

bool
BasicSurfaceScalingFunction::operator==(const SurfaceScalingFunction &o) const
{
	const BasicSurfaceScalingFunction *other = dynamic_cast<const BasicSurfaceScalingFunction*>(&o);
	if (!other)
		return false;
	else
		return (scale_ == other->scale_ &&
		    energyScale_ == other->energyScale_ &&
		    offset_ == other->offset_ &&
		    power_ == other->power_ &&
		    rBounds_ == other->rBounds_ &&
		    zBounds_ == other->zBounds_ &&
		    centerBounds_ == other->centerBounds_);
}

}

I3_SERIALIZABLE(I3MuonGun::SurfaceScalingFunction);
I3_SERIALIZABLE(I3MuonGun::ConstantSurfaceScalingFunction);
I3_SERIALIZABLE(I3MuonGun::BasicSurfaceScalingFunction);
I3_SERIALIZABLE(I3MuonGun::EnergyDependentSurfaceInjector);


// Distributed under the MIT License.
// See LICENSE.txt for details.

#pragma once

#include <array>
#include <boost/functional/hash.hpp>
#include <cstddef>
#include <limits>
#include <memory>

#include "DataStructures/DataBox/PrefixHelpers.hpp"
#include "DataStructures/DataBox/Prefixes.hpp"
#include "DataStructures/FixedHashMap.hpp"
#include "DataStructures/Tensor/TypeAliases.hpp"
#include "Domain/Structure/MaxNumberOfNeighbors.hpp"
#include "Domain/Tags.hpp"
#include "Evolution/DgSubcell/Tags/Mesh.hpp"
#include "Evolution/DgSubcell/Tags/NeighborData.hpp"
#include "Evolution/Systems/NewtonianEuler/FiniteDifference/Reconstructor.hpp"
#include "Evolution/Systems/NewtonianEuler/Tags.hpp"
#include "PointwiseFunctions/Hydro/EquationsOfState/EquationOfState.hpp"
#include "PointwiseFunctions/Hydro/Tags.hpp"

namespace NewtonianEuler::fd {
/*!
 * \brief Adaptive-order WENO reconstruction hybridizing orders 5 and 3. See
 * ::fd::reconstruction::aoweno_53() for details.
 */
template <size_t Dim>
class AoWeno53Prim : public Reconstructor<Dim> {
 private:
  // Conservative vars tags
  using MassDensityCons = Tags::MassDensityCons;
  using EnergyDensity = Tags::EnergyDensity;
  using MomentumDensity = Tags::MomentumDensity<Dim>;

  // Primitive vars tags
  using MassDensity = Tags::MassDensity<DataVector>;
  using Velocity = Tags::Velocity<DataVector, Dim>;
  using SpecificInternalEnergy = Tags::SpecificInternalEnergy<DataVector>;
  using Pressure = Tags::Pressure<DataVector>;

  using prims_tags =
      tmpl::list<MassDensity, Velocity, SpecificInternalEnergy, Pressure>;
  using cons_tags = tmpl::list<MassDensityCons, MomentumDensity, EnergyDensity>;
  using flux_tags = db::wrap_tags_in<::Tags::Flux, cons_tags, tmpl::size_t<Dim>,
                                     Frame::Inertial>;
  using prim_tags_for_reconstruction =
      tmpl::list<MassDensity, Velocity, Pressure>;

 public:
  struct GammaHi {
    using type = double;
    static constexpr Options::String help = {
        "The linear weight for the 5th-order stencil."};
  };
  struct GammaLo {
    using type = double;
    static constexpr Options::String help = {
        "The linear weight for the central 3rd-order stencil."};
  };
  struct Epsilon {
    using type = double;
    static constexpr Options::String help = {
        "The parameter added to the oscillation indicators to avoid division "
        "by zero"};
  };
  struct NonlinearWeightExponent {
    using type = size_t;
    static constexpr Options::String help = {
        "The exponent q to which the oscillation indicators are raised"};
  };

  using options =
      tmpl::list<GammaHi, GammaLo, Epsilon, NonlinearWeightExponent>;
  static constexpr Options::String help{
      "Monotised central reconstruction scheme using primitive variables."};

  AoWeno53Prim() = default;
  AoWeno53Prim(AoWeno53Prim&&) = default;
  AoWeno53Prim& operator=(AoWeno53Prim&&) = default;
  AoWeno53Prim(const AoWeno53Prim&) = default;
  AoWeno53Prim& operator=(const AoWeno53Prim&) = default;
  ~AoWeno53Prim() override = default;

  AoWeno53Prim(double gamma_hi, double gamma_lo, double epsilon,
               size_t nonlinear_weight_exponent);

  explicit AoWeno53Prim(CkMigrateMessage* msg);

  WRAPPED_PUPable_decl_base_template(Reconstructor<Dim>, AoWeno53Prim);

  auto get_clone() const -> std::unique_ptr<Reconstructor<Dim>> override;

  void pup(PUP::er& p) override;

  size_t ghost_zone_size() const override { return 3; }

  using reconstruction_argument_tags =
      tmpl::list<::Tags::Variables<prims_tags>,
                 hydro::Tags::EquationOfStateBase, domain::Tags::Element<Dim>,
                 evolution::dg::subcell::Tags::
                     NeighborDataForReconstructionAndRdmpTci<Dim>,
                 evolution::dg::subcell::Tags::Mesh<Dim>>;

  template <size_t ThermodynamicDim, typename TagsList>
  void reconstruct(
      gsl::not_null<std::array<Variables<TagsList>, Dim>*> vars_on_lower_face,
      gsl::not_null<std::array<Variables<TagsList>, Dim>*> vars_on_upper_face,
      const Variables<prims_tags>& volume_prims,
      const EquationsOfState::EquationOfState<false, ThermodynamicDim>& eos,
      const Element<Dim>& element,
      const FixedHashMap<
          maximum_number_of_neighbors(Dim) + 1,
          std::pair<Direction<Dim>, ElementId<Dim>>,
          evolution::dg::subcell::NeighborData,
          boost::hash<std::pair<Direction<Dim>, ElementId<Dim>>>>&
          neighbor_data,
      const Mesh<Dim>& subcell_mesh) const;

  /// Called by an element doing DG when the neighbor is doing subcell.
  template <size_t ThermodynamicDim, typename TagsList>
  void reconstruct_fd_neighbor(
      gsl::not_null<Variables<TagsList>*> vars_on_face,
      const Variables<prims_tags>& subcell_volume_prims,
      const EquationsOfState::EquationOfState<false, ThermodynamicDim>& eos,
      const Element<Dim>& element,
      const FixedHashMap<
          maximum_number_of_neighbors(Dim) + 1,
          std::pair<Direction<Dim>, ElementId<Dim>>,
          evolution::dg::subcell::NeighborData,
          boost::hash<std::pair<Direction<Dim>, ElementId<Dim>>>>&
          neighbor_data,
      const Mesh<Dim>& subcell_mesh,
      const Direction<Dim> direction_to_reconstruct) const;

 private:
  template <size_t LocalDim>
  // NOLINTNEXTLINE(readability-redundant-declaration)
  friend bool operator==(const AoWeno53Prim<LocalDim>& lhs,
                         const AoWeno53Prim<LocalDim>& rhs);

  double gamma_hi_ = std::numeric_limits<double>::signaling_NaN();
  double gamma_lo_ = std::numeric_limits<double>::signaling_NaN();
  double epsilon_ = std::numeric_limits<double>::signaling_NaN();
  size_t nonlinear_weight_exponent_ = 0;

  void (*reconstruct_)(gsl::not_null<std::array<gsl::span<double>, Dim>*>,
                       gsl::not_null<std::array<gsl::span<double>, Dim>*>,
                       const gsl::span<const double>&,
                       const DirectionMap<Dim, gsl::span<const double>>&,
                       const Index<Dim>&, size_t, double, double, double);
  void (*reconstruct_lower_neighbor_)(gsl::not_null<DataVector*>,
                                      const DataVector&, const DataVector&,
                                      const Index<Dim>&, const Index<Dim>&,
                                      const Direction<Dim>&, const double&,
                                      const double&, const double&);
  void (*reconstruct_upper_neighbor_)(gsl::not_null<DataVector*>,
                                      const DataVector&, const DataVector&,
                                      const Index<Dim>&, const Index<Dim>&,
                                      const Direction<Dim>&, const double&,
                                      const double&, const double&);
};

template <size_t Dim>
bool operator!=(const AoWeno53Prim<Dim>& lhs, const AoWeno53Prim<Dim>& rhs) {
  return not(lhs == rhs);
}
}  // namespace NewtonianEuler::fd


#define BOOST_TEST_MODULE "test_read_rectangular_box_interaction"

#ifdef BOOST_TEST_DYN_LINK
#include <boost/test/unit_test.hpp>
#else
#include <boost/test/included/unit_test.hpp>
#endif

#include <mjolnir/core/SimulatorTraits.hpp>
#include <mjolnir/core/BoundaryCondition.hpp>
#include <mjolnir/input/read_external_interaction.hpp>

BOOST_AUTO_TEST_CASE(read_rectangular_box_interaction_EXV)
{
    mjolnir::LoggerManager::set_default_logger("test_read_rectangular_box_interaction.log");

    using traits_type = mjolnir::SimulatorTraits<double, mjolnir::UnlimitedBoundary>;
    using real_type   = traits_type::real_type;
    {
        using namespace toml::literals;
        const toml::value v = u8R"(
            interaction = "RectangularBox"
            potential   = "ExcludedVolumeWall"
            box.lower   = [ 0.0,  0.0,  0.0]
            box.upper   = [10.0, 10.0, 10.0]
            box.margin  = 0.4
            epsilon     = 1.0
            parameters  = [
                {index =   0, radius =  1.0},
                {index = 100, radius = 10.0},
            ]
            )"_toml;

        const auto base = mjolnir::read_external_interaction<traits_type>(v);
        BOOST_TEST(static_cast<bool>(base));

        const auto derv = dynamic_cast<mjolnir::RectangularBoxInteraction<
            traits_type, mjolnir::ExcludedVolumeWallPotential<real_type>>*
            >(base.get()); // check the expected type is contained
        BOOST_TEST(static_cast<bool>(derv));

        const auto& interaction = *derv;

        BOOST_TEST(interaction.potential().epsilon() == 1.0);

        BOOST_TEST(interaction.potential().participants().size() ==   2u);
        BOOST_TEST(interaction.potential().participants().at(0)  ==   0u);
        BOOST_TEST(interaction.potential().participants().at(1)  == 100u);

        BOOST_TEST(interaction.potential().parameters().at(  0) == real_type( 1.0), boost::test_tools::tolerance(1e-8));
        BOOST_TEST(interaction.potential().parameters().at(100) == real_type(10.0), boost::test_tools::tolerance(1e-8));

        BOOST_TEST(mjolnir::math::X(interaction.lower()) == real_type(0.0), boost::test_tools::tolerance(1e-8));
        BOOST_TEST(mjolnir::math::Y(interaction.lower()) == real_type(0.0), boost::test_tools::tolerance(1e-8));
        BOOST_TEST(mjolnir::math::Z(interaction.lower()) == real_type(0.0), boost::test_tools::tolerance(1e-8));

        BOOST_TEST(mjolnir::math::X(interaction.upper()) == real_type(10.0), boost::test_tools::tolerance(1e-8));
        BOOST_TEST(mjolnir::math::Y(interaction.upper()) == real_type(10.0), boost::test_tools::tolerance(1e-8));
        BOOST_TEST(mjolnir::math::Z(interaction.upper()) == real_type(10.0), boost::test_tools::tolerance(1e-8));
    }
}


#include <iostream>                    // std::flush
#include <fstream>                     // std::ofstream
#include <string>                      // std::string
#include <tuple>                       // std::tuple
#include <vector>                      // std::vector
// For External Library
#include <torch/torch.h>               // torch
#include <boost/program_options.hpp>   // boost::program_options
// For Original Header
#include "loss.hpp"                    // Loss
#include "networks.hpp"                // UNet_Generator, PatchGAN_Discriminator
#include "dataloader.hpp"              // DataLoader::ImageFolderPairWithPaths
#include "visualizer.hpp"              // visualizer::graph

// Define Namespace
namespace po = boost::program_options;

// -------------------
// Validation Function
// -------------------
void valid(po::variables_map &vm, DataLoader::ImageFolderPairWithPaths &valid_dataloader, torch::Device &device, Loss &criterion_GAN, torch::nn::L1Loss &criterion_L1, UNet_Generator &gen, PatchGAN_Discriminator &dis, const size_t epoch, visualizer::graph &writer, visualizer::graph &writer_gen, visualizer::graph &writer_dis){

    // (0) Initialization and Declaration
    size_t iteration;
    float ave_G_GAN_loss, total_G_GAN_loss;
    float ave_G_L1_loss, total_G_L1_loss;
    float ave_dis_real_loss, total_dis_real_loss;
    float ave_dis_fake_loss, total_dis_fake_loss;
    std::ofstream ofs;
    std::tuple<torch::Tensor, torch::Tensor, std::vector<std::string>, std::vector<std::string>> mini_batch;
    torch::Tensor realI, realO, fakeO, realP, fakeP;
    torch::Tensor dis_real_out, dis_fake_out;
    torch::Tensor G_L1_loss, G_GAN_loss;
    torch::Tensor dis_real_loss, dis_fake_loss;
    torch::Tensor label_real, label_fake;

    // (1) Tensor Forward per Mini Batch
    gen->train();  // Dropout is required to make the generated images diverse
    dis->train();  // Dropout is required to make the generated images diverse
    iteration = 0;
    total_G_GAN_loss = 0.0; total_G_L1_loss = 0.0; total_dis_real_loss = 0.0; total_dis_fake_loss = 0.0;
    while (valid_dataloader(mini_batch)){

        realI = std::get<0>(mini_batch).to(device);
        realO = std::get<1>(mini_batch).to(device);

        // (1.1) Generator and Discriminator Forward
        fakeO = gen->forward(realI);
        fakeP = torch::cat({realI, fakeO}, /*dim=*/1);
        dis_fake_out = dis->forward(fakeP);

        // (1.2) Set Target Label
        label_real = torch::full({dis_fake_out.size(0), dis_fake_out.size(1), dis_fake_out.size(2), dis_fake_out.size(3)}, /*value*/1.0, torch::TensorOptions().dtype(torch::kFloat)).to(device);
        label_fake = torch::full({dis_fake_out.size(0), dis_fake_out.size(1), dis_fake_out.size(2), dis_fake_out.size(3)}, /*value*/0.0, torch::TensorOptions().dtype(torch::kFloat)).to(device);

        // (1.3) Discriminator Loss
        realP = torch::cat({realI, realO}, /*dim=*/1);
        dis_real_out = dis->forward(realP);
        dis_fake_loss = criterion_GAN(dis_fake_out, label_fake);
        dis_real_loss = criterion_GAN(dis_real_out, label_real);

        // (1.4) Generator Loss
        G_GAN_loss = criterion_GAN(dis_fake_out, label_real);
        G_L1_loss = criterion_L1(fakeO, realO) * vm["Lambda"].as<float>();

        // (1.5) Update Loss
        total_G_GAN_loss += G_GAN_loss.item<float>();
        total_G_L1_loss += G_L1_loss.item<float>();
        total_dis_real_loss += dis_real_loss.item<float>();
        total_dis_fake_loss += dis_fake_loss.item<float>();

        iteration++;
    }

    // (2) Calculate Average Loss
    ave_G_GAN_loss = total_G_GAN_loss / (float)iteration;
    ave_G_L1_loss = total_G_L1_loss / (float)iteration;
    ave_dis_real_loss = total_dis_real_loss / (float)iteration;
    ave_dis_fake_loss = total_dis_fake_loss / (float)iteration;

    // (3.1) Record Loss (Log)
    ofs.open("checkpoints/" + vm["dataset"].as<std::string>() + "/log/valid.txt", std::ios::app);
    ofs << "epoch:" << epoch << '/' << vm["epochs"].as<size_t>() << ' ' << std::flush;
    ofs << "G_GAN:" << ave_G_GAN_loss << ' ' << std::flush;
    ofs << "G_L1:" << ave_G_L1_loss << ' ' << std::flush;
    ofs << "D_Real:" << ave_dis_real_loss << ' ' << std::flush;
    ofs << "D_Fake:" << ave_dis_fake_loss << std::endl;
    ofs.close();

    // (3.2) Record Loss (Graph)
    writer.plot(/*base=*/epoch, /*value=*/{ave_G_GAN_loss + ave_G_L1_loss, ave_dis_real_loss + ave_dis_fake_loss});
    writer_gen.plot(/*base=*/epoch, /*value=*/{ave_G_GAN_loss + ave_G_L1_loss, ave_G_GAN_loss, ave_G_L1_loss});
    writer_dis.plot(/*base=*/epoch, /*value=*/{ave_dis_real_loss + ave_dis_fake_loss, ave_dis_real_loss, ave_dis_fake_loss});

    // End Processing
    return;

}


#include <boost/python.hpp>
#include "integrand.h"

using namespace Eigen;

void GetBandHamiltonian(const VectorXd &x, ComplexMatrix &mat, const std::vector<ComplexMatrix> &HR,
    const MatrixXi &R, const double &H, const double &delta)
{
  static Complex I(0,1);
  mat.fill(0);
  for (int i = 0; i < R.rows(); ++i) {
    mat += HR[i]*exp(I*(x(0)*R(i,0) + x(1)*R(i,1) + x(2)*R(i,2)));
  }
  for (int i = 0; i < mat.rows(); ++i)
    mat(i,i) -= H;
  for (int i = 0; i < kNCor*kNLayers; ++i)
    mat(i,i) -= delta;
}

GreenIntegrand::GreenIntegrand(const Eigen::VectorXcd &omega, const double &chemical_potential,
    const std::vector<ComplexMatrix> &hamiltonian_r, const Eigen::MatrixXi &r_vectors,
    const Eigen::MatrixXcd &all_self_energy, const double &magnetic_field, const double &double_counting)
    : md_int::GeneralIntegrand<Complex>(), delta_(double_counting), chemical_potential_(chemical_potential),
    ham_r_(hamiltonian_r), r_vectors_(r_vectors), h_field_(magnetic_field), all_omega_(omega),
    all_self_energy_(all_self_energy), self_energy_(ComplexMatrix::Zero()) {
  set_dim_in(kDim);
  set_dim_out(kMSize*kMSize);
}

// set omega and self energy:
// e.g. tilted 5d band case: 4 blocks for 4 cell, each block is a 5x5 matrix
void GreenIntegrand::set_data(const int &r) {
    w_ = all_omega_(r);

    for (int c = 0; c < kNLayers; ++c)
        for (int i = 0; i < kNCor; ++i)
            for (int j = 0; j < kNCor; ++j)
                self_energy_(kNCor*c+i,kNCor*c+j) = all_self_energy_(r, kNCor*kNCor*c+kNCor*i+j);
}

void GreenIntegrand::CalculateOriginalIntegrand(const VectorXd &x, VectorXcd &y) const
{
    y.resize(dim_out());
    ComplexMatrix tmp, negative_green;
    GetBandHamiltonian(x, tmp, ham_r_, r_vectors_, h_field_, delta_);

    for (int n = 0; n < kMSize; ++n) tmp(n, n) -= w_ + chemical_potential_;
    tmp += self_energy_;

    inv::inverse(tmp, negative_green);
    for (int i = 0; i < kMSize; ++i)
        for (int j = 0; j < kMSize; ++j)
            y(kMSize*i+j) = -negative_green(i, j);
}

HamiltonianIntegrand::HamiltonianIntegrand(const int &num_order, const std::vector<ComplexMatrix> &hamiltonian_r,
    const Eigen::MatrixXi &r_vectors, const double &magnetic_field, const double &double_counting)
    : md_int::GeneralIntegrand<Complex>(), num_order_(num_order), ham_r_(hamiltonian_r),
    r_vectors_(r_vectors), h_field_(magnetic_field), delta_(double_counting) {
  set_dim_in(kDim);
  set_dim_out(num_order_*kMSize*kMSize);
}

void HamiltonianIntegrand::CalculateOriginalIntegrand(const VectorXd &x, VectorXcd &y) const
{
    ComplexMatrix ham, ham_pow(ComplexMatrix::Identity());
    y.resize(dim_out());
    GetBandHamiltonian(x, ham, ham_r_, r_vectors_, h_field_, delta_);
    for (int n = 0; n < num_order_; ++n) {
        ham_pow *= ham;
        for (int i = 0; i < kMSize; ++i)
            for (int j = 0; j < kMSize; ++j)
                y(kMSize*kMSize*n+kMSize*i+j) = ham_pow(i, j);
    }
}


/* Copyright Institute of Sound and Vibration Research - All rights reserved */

#include "visr_renderer.hpp"

// 'error' and 'post' defined as a macro in the Max/MSP API and gets in the way of boost.
#undef error
#undef post

#include "options.hpp"

#include <libefl/basic_matrix.hpp>

#include <libpanning/LoudspeakerArray.h>

#include <libpml/initialise_parameter_library.hpp>

#include <libvisr/signal_flow_context.hpp>

#include <libsignalflows/baseline_renderer.hpp>

#include <boost/filesystem/path.hpp>
#include <boost/filesystem/operations.hpp>

#include <ciso646>
#include <cstdlib>
#include <sstream>

#include <maxmspexternals/libmaxsupport/argument_list.hpp>
#include <maxmspexternals/libmaxsupport/class_registrar.hpp>

/* Super-safe determination of the MAX define for setting the operating system. */
#ifdef __APPLE_CC__
#ifndef MAC_VERSION
#define MAC_VERSION
#undef WIN_VERSION
#endif
#else
#ifdef _MSC_VER
#ifndef WIN_VERSION
#define WIN_VERSION
#endif
#undef MAC_VERSION
#endif
#endif

#define MAXAPI_USE_MSCRT
#include "ext.h"
#include "z_dsp.h"
#include "ext_obex.h"

// Some loony (Max/MSP?) obviously defines max, thus hiding any C++/STL functions of the same name.
#undef max
#undef min // same for min(), although there is no Min/MSP around.

namespace visr
{
namespace maxmsp
{
namespace visr_renderer
{

VisrRenderer::VisrRenderer( t_pxobject & maxProxy, short argc, t_atom *argv )
 : ExternalBase( maxProxy )
 , mConfigInitialised( false )
 , mArrayConfiguration( new panning::LoudspeakerArray )
{
  pml::initialiseParameterLibrary();

  object_post( reinterpret_cast<t_object *>(getMaxProxy()), "VisrRenderer::VisrRenderer() constructor called." );
  try
  {
    ArgumentList args( argc, argv );

    Options cmdLineOptions;
    std::stringstream errMsg;
    switch( cmdLineOptions.parse( static_cast<int>(args.argc()), args.argv(), errMsg ) )
    {
    case Options::ParseResult::Failure:
      object_error( reinterpret_cast<t_object *>(getMaxProxy()), "Error while parsing command line options: ", errMsg.str().c_str() );
      return;
    case Options::ParseResult::Help:
    {
      std::stringstream str;
      cmdLineOptions.printDescription( str );
      object_error( reinterpret_cast<t_object *>(getMaxProxy()), str.str().c_str() );
      return;
    }
    case Options::ParseResult::Version:
      return;
    case Options::ParseResult::Success:
      break; // carry on
    }
    mNumberOfObjects = cmdLineOptions.getOption<std::size_t>( "input-channels" );
    mNumberOfOutputs = cmdLineOptions.getOption<std::size_t>( "output-channels" );
    mNumberOfEqSections = cmdLineOptions.getDefaultedOption( "object-eq-sections", 0 );

    boost::filesystem::path const arrayConfigPath( cmdLineOptions.getOption<std::string>( "array-config" ) );
    if( !exists( arrayConfigPath ) )
    {
      std::stringstream err; err << "The specified loudspeaker array configuration file \""
        << arrayConfigPath.string( ) << "\" does not exist.";
      throw std::invalid_argument( err.str() );
    }
    mArrayConfiguration->loadXmlFile( arrayConfigPath.string() );

    /* Set up the filter matrix for the diffusion filters. */
    std::size_t const diffusionFilterLength = 63; // fixed filter length of the filters in the compiled-in matrix
    std::size_t const diffusionFiltersInFile = 64; // Fixed number of filters in file.
    // First create a filter matrix containing all filters from a initializer list that is compiled into the program.
    efl::BasicMatrix<SampleType> allDiffusionCoeffs( diffusionFiltersInFile,
      diffusionFilterLength,
#include "../../apps/baseline_renderer/files/quasiAllpassFIR_f64_n63_initializer_list.txt"
      , cVectorAlignmentSamples );
    std::size_t const numberOfLoudspeakers = mArrayConfiguration->getNumRegularSpeakers( );
    mDiffusionFilters.reset( new efl::BasicMatrix<SampleType>( numberOfLoudspeakers,
      allDiffusionCoeffs.numberOfColumns(),
      cVectorAlignmentSamples ) );
    for( std::size_t idx( 0 ); idx < numberOfLoudspeakers; ++idx )
    {
      efl::vectorCopy( allDiffusionCoeffs.row( idx ), mDiffusionFilters->row( idx ), diffusionFilterLength, cVectorAlignmentSamples );
    }

    mSceneReceiverPort = cmdLineOptions.getDefaultedOption<std::size_t>( "scene-port", 4242 );
    mTrackingConfiguration = cmdLineOptions.getDefaultedOption<std::string>( "tracking", std::string( ) );

    mReverbConfiguration= cmdLineOptions.getDefaultedOption<std::string>( "reverb-config", std::string() );

    // Creating the inlets
    dsp_setup( getMaxProxy(), (int)mNumberOfObjects );

    // Creating the outlets
    for( std::size_t chIdx = 0; chIdx < mNumberOfOutputs; ++chIdx )
    {
      // Again: Using plainObject->mMaxProxy would require a less nasty cast.
      outlet_new( reinterpret_cast<t_object *>(getMaxProxy()), "signal" );
    }
  }
  catch( std::exception const & ex )
  {
    object_error( reinterpret_cast<t_object *>(getMaxProxy()), "Error in constructor: ", ex.what() );
    return;
  }
  mConfigInitialised = true;
}

VisrRenderer::~VisrRenderer()
{
}

/*virtual*/ void VisrRenderer::getFloat( double f )
{
  // At the moment we do not support any run-time control inputs.
  int inlet = getMaxProxy()->z_in;
#if 1 // whether to show debug messages at all
  {
    std::stringstream stream;
    stream << "VisrRenderer::applyFloat() called with inlet = " << inlet << ", f=" << f << "." << std::endl;
    post( stream.str().c_str() );
  }
#endif
//  switch( inlet )
  {
//  default:
    object_post( reinterpret_cast<t_object *>(getMaxProxy( )), "getFloat(): Received message for unsupported inlet %d.", inlet );
  }
}

/*virtual*/ void VisrRenderer::initDsp( t_object *dsp64, short *count, double samplerate, long maxvectorsize, long flags )
{
  if( not mConfigInitialised )
  {
    error( "VisrRenderer::initDsp(): Configuration in constructor failed, therefore dspInit() cannot be performed." );
    return;
  }
  long const period = maxvectorsize;
  // Request the actual buffer size
  if( maxvectorsize > std::numeric_limits<short>::max() )
  {
    error( "The maximum vector length is larger than the maximum value to be returned by the \"count\" variable." );
  }
  *count = static_cast<short>(period);

  try
  {
    SamplingFrequencyType const samplingFrequency = static_cast<SamplingFrequencyType>(std::round( samplerate ));

    // Assume a fixed length for the interpolation period.
    // Ideally, this roughly matches the update rate of the scene sender.
    // Provide a special case where the period is longer than the default period.
    // In this case the signal flow would produce an exception to avoid a 'staircased' gain trajectory.
    const std::size_t cInterpolationLength = std::max( 2048l, period );

    mContext.reset( new SignalFlowContext{ static_cast<std::size_t>(period), samplingFrequency} );

    /*
    explicit BaselineRenderer( SignalFlowContext const & context,
      char const * name,
      CompositeComponent * parent,
      panning::LoudspeakerArray const & loudspeakerConfiguration,
      std::size_t numberOfInputs,
      std::size_t numberOfOutputs,
      std::size_t interpolationPeriod,
      efl::BasicMatrix<SampleType> const & diffusionFilters,
      std::string const & trackingConfiguration,
      std::size_t sceneReceiverPort,
      std::size_t numberOfObjectEqSections,
      std::string const & reverbConfig,
      bool frequencyDependentPanning );
*/
    mFlow.reset( new signalflows::BaselineRenderer( *mContext, "VisrRenderer", nullptr,
                                                    *mArrayConfiguration,
                                                    mNumberOfObjects, mNumberOfOutputs,
                                                    cInterpolationLength,
                                                    *mDiffusionFilters,
                                                    mTrackingConfiguration,
                                                    mSceneReceiverPort,
                                                    mNumberOfEqSections,
                                                    mReverbConfiguration,
                                                    false /* no frequency-dependent panning */) );
    mFlowWrapper.reset( new maxmsp::SignalFlowWrapper<double>( *mFlow ) );
  }
  catch( std::exception const & e )
  {
    std::stringstream errMsg;
    errMsg << "Error while initialising the signal flow: " << e.what();
    object_error( reinterpret_cast<t_object *>(getMaxProxy( )), errMsg.str( ).c_str( ));
    return;
  }
}

/*virtual*/ void VisrRenderer::perform( t_object *dsp64, double **ins,
                                        long numins, double **outs, long numouts,
                                        long sampleframes, long flags, void *userparam )
{
  if( numins != mNumberOfObjects )
  {
    error( "VisrRenderer::perform(): Parameter \"numins\" does not match the number of configured input channels." );
  }
  if( numouts != mNumberOfOutputs )
  {
    error( "VisrRenderer::perform(): Parameter \"numouts\" does not match the number of configured output channels." );
  }
  if( sampleframes != static_cast<long>(mContext->period()) )
  {
    error( "VisrRenderer::perform( ): The number of requested samples %d  does not match the given block length." );
  }
  try
  {
    mFlowWrapper->processBlock( ins, outs );
  }
  catch( std::exception const & e )
  {
    error( "Error while initialising the signal flow %s", e.what() );
    return;
  }
}

/*virtual*/ void VisrRenderer::assist( void *b, long msg, long arg, char *dst )
{
  if( msg == ASSIST_INLET )
  {
    switch( arg )
    { // if there is more than one inlet or outlet, a switch is necessary
    case 1: sprintf( dst, "(float/signal2 Input) Scalar gain value (linear scale)" ); break;
    default: sprintf( dst, "(signal %ld) Input", arg + 1 ); break;
    }
  }
  else if( msg == ASSIST_OUTLET )
  {
    sprintf( dst, "(signal %ld) Output", arg + 1 );
  }
}

} // namespace visr_renderer
} // namespace maxmsp
} // namespace visr

extern "C"
{
/****************************/
// Initialization routine MAIN
int C74_EXPORT main()
{
  post( "visr::maxmsp::VisrRenderer::main() called." );

  visr::maxmsp::ClassRegistrar<visr::maxmsp::visr_renderer::VisrRenderer>( "visr_renderer~" );

  post("visr::maxmsp::VisrRenderer::main() finished.");

  return 0;
}
} // extern "C"


//
// Created by ryan on 1/5/17.
//
#include <math.h>
#include <algorithm>

#include <ros/ros.h>

#include <Eigen/Eigen>

#include <opencv2/core.hpp>
#include <opencv2/imgproc.hpp>
#include <opencv2/calib3d.hpp>

#include <vision/Process.h>
#include <vision/Target.h>
#include <sstream>

using namespace cv;
using namespace std;

namespace vision {

    /**
     *
     * @param config
     * @return
     */
    Process::Process(Config *config) {
        m_config = config;
    }

    /**
     *
     * @param source
     * @param info
     * @return
     */
    std::vector<TargetPair> Process::ProcessFrame(Mat source, Mat &dest, Mat &processed) {
        auto featureBegin = ros::Time::now();
        vector<Contour> polyContours = featureExtraction(source, dest, processed);
        ROS_DEBUG("Found blobs %lu", polyContours.size());
        auto featureEnd = ros::Time::now();

        auto idBegin = ros::Time::now();
        vector<TargetInfo> foundTargets = targetIdentification(polyContours, dest, processed);
        ROS_DEBUG("Found targets %lu", foundTargets.size());
        auto idEnd = ros::Time::now();

        auto debugBegin = ros::Time::now();
        for(auto target : foundTargets) {
            Point2f vertices[4];
            target.boundingBox.points(vertices);

            drawContours(dest, polyContours, target.polyIndex, Scalar(0,0,255), 2);
            drawContours(processed, polyContours, target.polyIndex, Scalar(0,0,255), 2);
            for (int i = 0; i < 4; i++) {
                line(processed, vertices[i], vertices[(i + 1) % 4], Scalar(255, 0, 0), 2);
            }

            std::stringstream ss;
            ss << target.compositeScore;
            cv::putText(processed, ss.str(), target.boundingBox.center, 1.0, FONT_HERSHEY_PLAIN, cv::Scalar(255,255,255));
        }
        auto debugEnd = ros::Time::now();

        auto pairBegin = ros::Time::now();
        vector<TargetPair> targetPairs = findTargetPairs(foundTargets, dest);
        ROS_DEBUG("Found targets pairs %lu", targetPairs.size());
        auto pairEnd = ros::Time::now();

        for (auto pair : targetPairs) {
            line(processed, pair.left.boundingBox.center, pair.right.boundingBox.center, Scalar(0, 255, 255), 2);
        }

        auto dataBegin = ros::Time::now();
        auto data = dataExtraction(polyContours, targetPairs, dest);
        auto dataEnd = ros::Time::now();

        ROS_INFO_STREAM("Feature (ms): " << (featureEnd - featureBegin).toSec() * 1000.0);
        ROS_INFO_STREAM("Id      (ms): " << (idEnd - idBegin).toSec() * 1000.0);
        ROS_INFO_STREAM("Debug   (ms): " << (debugEnd - debugBegin).toSec() * 1000.0);
        ROS_INFO_STREAM("Extract (ms): " << (dataEnd - dataBegin).toSec() * 1000.0);

        return data;
    }

    /**
     *
     * @param source
     * @param dest
     * @return
     */
    vector<Contour> Process::featureExtraction(cv::Mat source, cv::Mat &dest, Mat &processed) {
        ROS_DEBUG("Feature extraction start");
        /*
         * Image Processing
         * ================
         * Threshold
         * Blur
         * Remove small objects
         * Reconstruction/fix blobs
         * Find Contours of blobs
         */

        auto begin = ros::Time::now();
        Mat frame;
        source.copyTo(frame);
        source.copyTo(dest);
        source.copyTo(processed);
        auto end = ros::Time::now();
        ROS_INFO_STREAM("  Copy (ms): " << (end - begin).toSec() * 1000.0);

        //Convert to HSV color space
        begin = ros::Time::now();
        ROS_DEBUG("Converting to HSV");
        cvtColor(frame, frame, CV_BGR2HSV);
        end = ros::Time::now();
        ROS_INFO_STREAM("  HSV  (ms): " << (end - begin).toSec() * 1000.0);

        //Threshold, find pixels in range
        begin = ros::Time::now();
        ROS_DEBUG_STREAM("Thresholding image to: " << m_config->hsvLow << m_config->hsvHigh);
        inRange(frame, m_config->hsvLow, m_config->hsvHigh, frame);
        end = ros::Time::now();
        ROS_INFO_STREAM("  range (ms): " << (end - begin).toSec() * 1000.0);

        /*
         * Morphology
         * First a closing operation is preformed to help group blobs that are close together
         * into a single blob
         * Then a opening operation is preformed to help remove isolated noisy areas.
         * https://stackoverflow.com/questions/30369031/
         * remove-spurious-small-islands-of-noise-in-an-image-python-opencv
         */
        //Mat closeSE = getStructuringElement(MORPH_RECT, Size(5,5));
        // morphologyEx(frame, frame, MORPH_CLOSE, closeSE);
        //Mat openSE = getStructuringElement(MORPH_RECT, Size(2,2));
        //morphologyEx(frame, frame, MORPH_OPEN, openSE);

        begin = ros::Time::now();
        frame.copyTo(dest);
        cvtColor(dest,dest, CV_GRAY2BGR);
        end = ros::Time::now();
        ROS_INFO_STREAM("  Gray  (ms): " << (end - begin).toSec() * 1000.0);

        //Find contours
        begin = ros::Time::now();
        vector<Contour> contours;
        vector<Vec4i> hierarchy;
        findContours(frame, contours, hierarchy, CV_RETR_TREE, CV_CHAIN_APPROX_SIMPLE);
        ROS_DEBUG_STREAM("Contours: " << contours.size());
        end = ros::Time::now();
        ROS_INFO_STREAM("  Contour(ms): " << (end - begin).toSec() * 1000.0);

        // Commented out due to how episilon interactes with low resolution (small) targets
        //Find Polygon contours
        begin = ros::Time::now();
        vector<Contour> resultContours;
        for(int i = 0; i < contours.size(); i++) {
            Contour contour = contours[i];
//            //NOTE the true at the end to signal that the resulting contour is closed
//            approxPolyDP(contours[i], polyContour, m_config->epsilon, true);
//
//            drawContours(processed, contours, i, Scalar(0,255,255), 2);
//
//
            //If a point is near the border of the image ignore the contour,
            //as there may be issues with a partial target being detected.
            //This could cause bad pose estimation.
            if(nearBorder(contour, source.size(), 10)) {
                ROS_DEBUG("Skipping contour as it is near border");
                continue;
            }

            const double minArea = 200;
            double area = contourArea(contour);
            if (area < minArea) {
                ROS_DEBUG("Skipping contour as it is too small");
                continue;
            }
//
//            //Check to see if the contour has a valid number of sides, if so add it
////            int sides = polyContour.size();
////            if(sides == 4) {
//                //Rectangle
//                polyContours.push_back(polyContour);
////            } else {
//
////            }
            resultContours.push_back(contour);
        }
        end = ros::Time::now();
        ROS_INFO_STREAM("  Loop (ms): " << (end - begin).toSec() * 1000.0);

        //Sort Contours from largest to smallest
        sort(contours.begin(), contours.end(),
             [](Contour a, Contour b) { return contourArea(a) > contourArea(b); }
        );

        ROS_DEBUG("Feature extraction end");

        return resultContours;
    }

    vector<TargetInfo> Process::targetIdentification(std::vector<Contour> polyContours, cv::Mat &dest, Mat &processed) {
        ROS_DEBUG("Target Identification start");
        /*
         * Target Identification
         * ==============
         * Check each blob against each possible target, and
         * score each blob accordingly. And select the best
         * target for data extraction
         */
        vector<TargetInfo> foundTargets;
        for(int i = 0; i < polyContours.size() && i < m_config->blobMax; i++) {
            ROS_DEBUG("Blob %i", i);
            // int size = polyContours[i].size();
            double area = contourArea(polyContours[i]);
            RotatedRect boundingBox = minAreaRect(polyContours[i]);

            double rectScore = scoreRectangularity(area, boundingBox);
            // The left and right targets have the same aspect ratio
            double aspectRatioScore = scoreAspectRatio(boundingBox, Target::kLEFT_TILT_TARGET);

            double angle = targetAngle(boundingBox);

            // This is used to reject vertical targets
            const double minTilt = 2;

            ROS_DEBUG("Blob %i Score: Rect %f AR %f", i, rectScore, aspectRatioScore);
            if(rectScore >= 50 && aspectRatioScore >= 20 && minTilt < fabs(angle)) {
                TargetInfo info;
                info.polyIndex = i;
                info.boundingBox = boundingBox;

                ROS_DEBUG_STREAM("Angle" <<  angle);
                if (angle > 0) {
                    // Tilted counter clock wise, right target
                    info.targetID = TargetID::kRightTilt;
                } else {
                    // Tilted clock wise, left target
                    info.targetID = TargetID::kLeftTilt;
                }
                info.compositeScore = (rectScore * 0.5) + (aspectRatioScore * 0.5);
                foundTargets.push_back(info);
            }
        }
        ROS_DEBUG("Target Identification end");
        return foundTargets;
    }

    vector<TargetPair> Process::findTargetPairs(std::vector<TargetInfo> foundTargets, cv::Mat &dest) {
        ROS_DEBUG("Find target pairs start");
        vector<TargetPair> results;

        for(auto target : foundTargets) {
            switch (target.targetID) {
            case TargetID::kNotATarget:
                ROS_DEBUG("Not a target target");
                break;
            case TargetID::kLeftTilt:
                ROS_DEBUG("Left tilt target");
                break;
            case TargetID::kRightTilt:
                ROS_DEBUG("Right tilt target");
                break;
            default:
                ROS_DEBUG("Unknown target");
                break;

            }
        }

        // Remove unintresting targets
        foundTargets.erase(remove_if(foundTargets.begin(), foundTargets.end(), [](TargetInfo info) {
            bool remove =  !(info.targetID == TargetID::kLeftTilt || info.targetID == TargetID::kRightTilt);
            ROS_DEBUG_STREAM("Remove: " << remove);
            return remove;
        }), foundTargets.end());

        ROS_DEBUG_STREAM("Target count " << foundTargets.size());

        // If there are less than 2 foudn targets there is no point in proceeding
        if (foundTargets.size() < 2) {
            ROS_DEBUG("Less than 2 targets found");
            return results;
        }

        //Sort Targets from left to right x position (sort as increasing x)
        sort(foundTargets.begin(), foundTargets.end(),
             [](TargetInfo a, TargetInfo b) { return a.boundingBox.center.x < b.boundingBox.center.x; }
        );

        for (auto target : foundTargets) {
            ROS_DEBUG_STREAM(target.targetID << ", " <<target.boundingBox.center.x);
        }

        // Try to find a matching pair
        // A matching pair is found when the left side target is a right tilted target and then the next
        // target is a left tilt target.

        // Start with left most target
        // It is assumed that only left and right ta
        auto it = foundTargets.begin();
        while(it != foundTargets.end()) {
            ROS_DEBUG("Finding pair");
            TargetInfo leftSide = *it;
            ROS_DEBUG_STREAM("Left: " << leftSide.targetID << ", contour index: " << leftSide.polyIndex);

            it++; // Advance to next target
            if (leftSide.targetID != TargetID::kRightTilt) {
                ROS_DEBUG_STREAM("Left side target is not tilted right");
                continue; // Try again
            }

            TargetInfo rightSide = *it;
            ROS_DEBUG_STREAM("Right: " << rightSide.targetID << ", contour index: " << rightSide.polyIndex);
            if (rightSide.targetID != TargetID::kLeftTilt) {
                // This is not a right side target, lets try this process again with the left one being this target
                ROS_DEBUG_STREAM("Left side target is not tilted left");
                continue;
            }

            // Its a right target advance iterator for next iteration of this loop
            it++;

            // Now try to determine if this a valid target pair.
            // This can be done by scoring the distance between the 2 targets centers
            double distance = pointToPointDistance(leftSide.boundingBox.center, rightSide.boundingBox.center);

            double leftRatio = longSide(leftSide.boundingBox.size) / distance;
            double rightRatio = longSide(rightSide.boundingBox.size) / distance;

            const double goalRatio = 0.5; // TODO find
            double leftScore = ratioToScore(goalRatio / leftRatio);
            double rightScore = ratioToScore(goalRatio / rightRatio);

            ROS_DEBUG_STREAM("Left ratio: " << leftRatio << ", score: " << leftScore);
            ROS_DEBUG_STREAM("Right ratio: " << rightRatio << ", score: " << rightScore);

            const double minScore = 1; // TOOD find
            if (leftScore > minScore && rightScore > minScore) {
                ROS_DEBUG_STREAM("Found a target pair!");
                // We found a target pair!
                TargetPair pair;
                pair.left = leftSide;
                pair.right = rightSide;
                results.push_back(pair);
            }
        }

        ROS_DEBUG("Find target pairs end");

        return results;
    }

    vector<TargetPair> Process::dataExtraction(vector<Contour> polyContours, vector<TargetPair> targetPairs, Mat &dest) {
        ROS_DEBUG("Data extraction start");
        /**
         * Data extraction
         * ===============
         * Get target position relative to camera.
         */
        vector<TargetPair> results;

        //If no targets are found, no point in proceeding
        if(targetPairs.size() == 0) return results;

        // Use corner points of the bounding boxes are points to supply to solvePnP

        Mat rvec(3,1, DataType<double>::type);
        Mat tvec(3,1, DataType<double>::type);
        for(auto pair : targetPairs) {
            TargetInfo left = pair.left;
            TargetInfo right = pair.right;

            //X - forward   - blue
            //Y - left      - green
            //Z - up        - red

            // TODO: Would more points be better for Solve PNP?

            // Figure out object points
            vector<Point3f> objectPoints;

            // Top left
            Eigen::Vector3f objectTL;
            objectTL << -1, 2.75, 1;
            objectTL = transform(objectTL, -11.313/2, 0, -14.5 * M_PI / 180.0);
            objectPoints.push_back(Point3f(0, objectTL[0], objectTL[1]));

            // Top Right
            Eigen::Vector3f objectTR;
            objectTR << 1, 2.75, 1;
            objectTR = transform(objectTR, 11.313/2, 0, 14.5 * M_PI / 180.0);
            objectPoints.push_back(Point3f(0, objectTR[0], objectTR[1]));

            // Bottom Left
            Eigen::Vector3f objectBL;
            objectBL << -1, -2.75, 1;
            objectBL = transform(objectBL, -11.313/2, 0, -14.5 * M_PI / 180.0);
            objectPoints.push_back(Point3f(0, objectBL[0], objectBL[1]));

            // Bottom Right
            Eigen::Vector3f objectBR;
            objectBR << 1, -2.75, 1;
            objectBR = transform(objectBR, 11.313/2, 0, 14.5 * M_PI / 180.0);
            objectPoints.push_back(Point3f(0, objectBR[0], objectBR[1]));

            ROS_INFO_STREAM("Object Top Left: " << objectTL);
            ROS_INFO_STREAM("Object Top Right: " << objectTR);
            ROS_INFO_STREAM("Object Bottom Left: " << objectBL);
            ROS_INFO_STREAM("Object Bottom Right: " << objectBR);

            // Figure out image points
             vector<Point2f> imagePoints;
            // imagePoints.push_back(extremeBR);
            // imagePoints.push_back(extremeBL);
            // imagePoints.push_back(extremeTL);
            // imagePoints.push_back(extremeTR);

            // TODO: Is it okay to invert the bounding box angle? If we did then the Y axis would also need to inverted
            double leftAngle = left.boundingBox.angle * M_PI / 180.0;
            double leftLength = longSide(left.boundingBox.size);
            double leftShort = shortSide(left.boundingBox.size);

            double rightAngle = right.boundingBox.angle * M_PI / 180.0;
            double rightLength = longSide(right.boundingBox.size);
            double rightShort = shortSide(right.boundingBox.size);

            // Top Left
            Eigen::Vector3f imageTL;
            imageTL << -leftShort / 2.0, leftLength/2.0, 1;
            imageTL = transform(imageTL, left.boundingBox.center.x, left.boundingBox.center.y, leftAngle);
            imagePoints.push_back(Point2f(imageTL[0], imageTL[1]));

            // Top Right
            Eigen::Vector3f imageTR;
            imageTR << rightShort / 2.0, rightLength/2.0, 1;
            imageTR = transform(imageTL, right.boundingBox.center.x, right.boundingBox.center.y, rightAngle);
            imagePoints.push_back(Point2f(imageTR[0], imageTR[1]));

            // Bottom Left
            Eigen::Vector3f imageBL;
            imageBL<< -leftShort / 2.0, -leftLength/2.0, 1;
            imageBL = transform(imageBL, left.boundingBox.center.x, left.boundingBox.center.y, leftAngle);
            imagePoints.push_back(Point2f(imageBL[0], imageBL[1]));

            // Bottom Right
            Eigen::Vector3f imageBR;
            imageBR << rightShort / 2.0, -rightLength/2.0, 1;
            imageBR = transform(imageBL, right.boundingBox.center.x, right.boundingBox.center.y, rightAngle);
            imagePoints.push_back(Point2f(imageBR[0], imageBR[1]));

            //NOTE: SolvePnPRansac makes things worse
            solvePnP(objectPoints, imagePoints, m_config->cameraMatrix, m_config->distCoeffs, rvec, tvec);

            //Need to adjust the translation vector from solvePnP to achieve correct results
            //https://stackoverflow.com/questions/18637494/camera-position-in-world-coordinate-from-cvsolvepnp
            Mat rotationMatrix;
            Rodrigues(rvec, rotationMatrix);
            rotationMatrix = rotationMatrix.t();
            Mat tvecNew = -rotationMatrix * tvec;

            double x, y, z;

            x = tvecNew.at<double>(0,0);
            y = tvecNew.at<double>(1,0);
            z = tvecNew.at<double>(2,0); //In the future this may not be needed
            double distance = sqrt((x*x) + (y*y) + (z*z));
            ROS_DEBUG("Distance %f", distance);

            if(distance > m_config->maxDistance) {
                ROS_WARN("Target too far way, ignoring");
                continue;
            }

            //Not sure if this function is correct
            double yaw, pitch, roll;
            RotationMatrixToYPR(rotationMatrix, yaw, pitch, roll);

            pair.pose.x = x;
            pair.pose.y = y;
            pair.pose.z = z;
            pair.rotation.yaw   = yaw;
            pair.rotation.pitch = pitch;
            pair.rotation.roll  = roll;
            results.push_back(pair);

            //Project target coordinate axis on to image
            vector<Point3d> axis;
            //X, Y, Z
            axis.push_back(Point3d(0, 0, 0));   //Orgin
            axis.push_back(Point3d(12, 0, 0));  //X
            axis.push_back(Point3d(0, 12, 0));  //Y
            axis.push_back(Point3d(0, 0, 12));  //Z

            vector<Point2d> drawp;
            projectPoints(axis, rvec, tvec, m_config->cameraMatrix, m_config->distCoeffs, drawp);

            line(dest, drawp[0], drawp[1], Scalar(255, 0, 0), 4); //X-Blue
            line(dest, drawp[0], drawp[2], Scalar(0, 255, 0), 4); //Y-Green
            line(dest, drawp[0], drawp[3], Scalar(0, 0, 255), 4); //Z-Red
        }
        ROS_DEBUG("Data extraction end");

        return results;
    }

    bool Process::nearBorder(std::vector<cv::Point> points, cv::Size imageSize, int borderSize) {
        for(int i = 0; i < points.size(); i++) {
            Point point = points[i];
            if(point.x <= borderSize) return true;
            if(point.y <= borderSize) return true;
            if((imageSize.width - point.x) <= borderSize) return true;
            if((imageSize.height - point.y) <= borderSize) return true;
        }
        return false;
    }

    void Process::RotationMatrixToYPR(cv::Mat &rotation,
                                      double &yaw, double &pitch, double &roll) {
        yaw = atan(rotation.at<double>(1,0)/rotation.at<double>(0,0));
        pitch = atan(-rotation.at<double>(2,0)/sqrt(pow(rotation.at<double>(2,1),2) + pow(rotation.at<double>(2,2),2)));
        yaw = atan(rotation.at<double>(2,1)/rotation.at<double>(2,2));
    }

    /**
     *
     * @param ratio
     * @return Score between 0-100
     */
    double Process::ratioToScore(double ratio) {
        return max(0.0, min(100.0*(1.0 - fabs(1.0 - ratio)), 100.0));
    }

    /**
     *
     * @param contourArea
     * @param boundingBox
     * @return Score between 0-100
     */
    double Process::scoreRectangularity(double contourArea, RotatedRect &boundingBox) {
        double boundingArea = boundingBox.size.area();
        if(boundingArea != 0) {
            return 100 * (contourArea/ boundingArea);
        } else {
            return 0;
        }
    }

    /**
     *
     * @param boundingBox
     * @param target
     * @return Score between 0-100
     */
    double Process::scoreAspectRatio(RotatedRect &boundingBox, Target target) {
        double rectShort = min(boundingBox.size.height, boundingBox.size.height);
        double rectLong  = max(boundingBox.size.width,  boundingBox.size.width);

        if(boundingBox.size.width > boundingBox.size.height) {
            //Particle is wider than it is tall, divide long by short
            return ratioToScore((rectLong / rectShort) / target.aspectRatio());
        } else {
            //Particle is taller than it is wide, divide short by long
            return ratioToScore((rectShort / rectLong) / target.aspectRatio());
        }
    }

    /**
     * Normalize angle between 2 bounds
     * @param value Angle
     * @param start Start andgle, Ex: 0, -180
     * @param end End angle, Ex: 360, 180
     * @return Normalized angle
     */
    double Process::normalizeAngle(const double value, const double start, const double end) {
        const double width = end - start;
        const double offsetValue = value - start;

        return (offsetValue - (floor(offsetValue/width) * width)) + start;
    }

    /**
     * Calculate the angle (in degrees) that the target is tilted. 0 degress means that the target is facing straight up.
     * @param box Bounding box of target
     * @return Positive angles clockwise, negative angles counter clockwise
     */
    double Process::targetAngle(cv::RotatedRect box) {
        // How Min Area Rect calculates rotated rect: https://namkeenman.wordpress.com/2015/12/18/open-cv-determine-angle-of-rotatedrect-minarearect/
        double angle = box.angle;
        if (box.size.width > box.size.height) {
            angle += 90;
        }
        return normalizeAngle(angle, -180, 180);
    }

    /**
     *
     * @param image
     * @param center
     * @param length
     * @param angle
     * @param color
     */
    void Process::drawBoxLine(cv::Mat image, cv::Point center, double length, double angle, cv::Scalar color) {
        cv::Point endPoint;
        endPoint.x = center.x + 100.0 * cos(angle * M_PI / 180);
        endPoint.y = center.y + 100.0 * sin(angle * M_PI / 180);
        line(image, center, endPoint, color, 2);
    }

    double Process::distanceScore(TargetInfo left, TargetInfo right) {
//        // TODO :This may not be the best way to determine distance between targets.
//
//        double ratioLeft;
//        double ratioRight;
//        double ratio =  (ratioLeft + ratioRight) / 2.0;
//        return ratioToScore()
        return 0;
    }

    double Process::pointToPointDistance(cv::Point a, cv::Point b) {
        double delx = b.x - a.x;
        double dely = b.y - a.y;
        return sqrt(delx * delx + dely * dely);
    }

    cv::Point Process::getMidPoint(cv::Point a, cv::Point b) {
        return cv::Point((a.x + b.x) / 2.0, (a.y + b.y) / 2.0);
    }

    double Process::longSide(cv::Size size) {
        if (size.height > size.width) {
            return size.height;
        }
        return size.width;
    }

    double Process::shortSide(cv::Size size) {
        if (size.height < size.width) {
            return size.height;
        }
        return size.width;
    }

    Eigen::Vector3f Process::transform(Eigen::Vector3f point, double h, double k, double theta) {
        std::cout <<  "Point: " << point << std::endl;
        Eigen::Matrix3f rotate;
        rotate << cos(theta), -sin(theta), 0,
                  sin(theta), cos(theta),  0,
                  0,          0,           1;

        std::cout <<  "Rotate: " << rotate << std::endl;

        Eigen::Matrix3f translate;
        translate << 1, 0, h,
                     0, 1, k,
                     0, 0, 1;
        std::cout <<  "Translate: " << translate << std::endl;

        auto result = (translate * rotate) * point;
        return result;

    }

}


#include "ImageComponent.h"
#include <iostream>
#include <boost/filesystem.hpp>
#include <math.h>
#include "../Log.h"
#include "../Renderer.h"
#include "../Window.h"

Eigen::Vector2i ImageComponent::getTextureSize() const
{
	if(mTexture)
		return mTexture->getSize();
	else
		return Eigen::Vector2i(0, 0);
}

Eigen::Vector2f ImageComponent::getCenter() const
{
	return Eigen::Vector2f(mPosition.x() - (getSize().x() * mOrigin.x()) + getSize().x() / 2,
		mPosition.y() - (getSize().y() * mOrigin.y()) + getSize().y() / 2);
}

ImageComponent::ImageComponent(Window* window, float offsetX, float offsetY, std::string path, float targetWidth, float targetHeight, bool allowUpscale) : GuiComponent(window),
	mTiled(false), mAllowUpscale(allowUpscale), mFlipX(false), mFlipY(false), mOrigin(0.5, 0.5), mTargetSize(targetWidth, targetHeight), mColorShift(0xFFFFFFFF)
{
	setPosition(offsetX, offsetY);

	if(!path.empty())
		setImage(path);
}

ImageComponent::~ImageComponent()
{
}

void ImageComponent::resize()
{
	if(!mTexture)
		return;

	mSize << (float)getTextureSize().x(), (float)getTextureSize().y();

	//(we don't resize tiled images)
	if(!mTiled && (mTargetSize.x() || mTargetSize.y()))
	{
		Eigen::Vector2f resizeScale(Eigen::Vector2f::Zero());

		if(mTargetSize.x() && (mAllowUpscale || mSize.x() > mTargetSize.x()))
		{
			resizeScale[0] = mTargetSize.x() / mSize.x();
		}
		if(mTargetSize.y() && (mAllowUpscale || mSize.y() > mTargetSize.y()))
		{
			resizeScale[1] = mTargetSize.y() / mSize.y();
		}

		if(resizeScale.x() && !resizeScale.y())
			resizeScale[1] = resizeScale.x();
		if(resizeScale[1] && !resizeScale.x())
			resizeScale[0] = resizeScale.y();

		if(resizeScale.x())
			mSize[0] = (mSize.x() * resizeScale.x());
		if(resizeScale.y())
			mSize[1] = (mSize.y() * resizeScale.y());
	}

	if(mTiled)
		mSize = mTargetSize;
}

void ImageComponent::setImage(std::string path)
{
	mPath = path;

	if(mPath.empty() || !mWindow->getResourceManager()->fileExists(mPath))
		mTexture.reset();
	else
		mTexture = TextureResource::get(*mWindow->getResourceManager(), mPath);

	resize();
}

void ImageComponent::setOrigin(float originX, float originY)
{
	mOrigin << originX, originY;
}

void ImageComponent::setTiling(bool tile)
{
	mTiled = tile;

	if(mTiled)
		mAllowUpscale = false;

	resize();
}

void ImageComponent::setResize(float width, float height, bool allowUpscale)
{
	mTargetSize << width, height;
	mAllowUpscale = allowUpscale;
	resize();
}

void ImageComponent::setFlipX(bool flip)
{
	mFlipX = flip;
}

void ImageComponent::setFlipY(bool flip)
{
	mFlipY = flip;
}

void ImageComponent::setColorShift(unsigned int color)
{
	mColorShift = color;
}

void ImageComponent::render(const Eigen::Affine3f& parentTrans)
{
	Eigen::Affine3f trans = parentTrans * getTransform();
	Renderer::setMatrix(trans);

	if(mTexture && getOpacity() > 0)
	{
		GLfloat points[12], texs[12];
		GLubyte colors[6*4];

		if(mTiled)
		{
			float xCount = mSize.x() / getTextureSize().x();
			float yCount = mSize.y() / getTextureSize().y();

			Renderer::buildGLColorArray(colors, (mColorShift >> 8 << 8)| (getOpacity()), 6);
			buildImageArray(0, 0, points, texs, xCount, yCount);
		}else{
			Renderer::buildGLColorArray(colors, (mColorShift >> 8 << 8) | (getOpacity()), 6);
			buildImageArray(0, 0, points, texs);
		}

		drawImageArray(points, texs, colors, 6);
	}

	GuiComponent::renderChildren(trans);
}

void ImageComponent::buildImageArray(int posX, int posY, GLfloat* points, GLfloat* texs, float px, float py)
{
	points[0] = posX - (mSize.x() * mOrigin.x());		points[1] = posY - (mSize.y() * mOrigin.y());
	points[2] = posX - (mSize.x() * mOrigin.x());		points[3] = posY + (mSize.y() * (1 - mOrigin.y()));
	points[4] = posX + (mSize.x() * (1 - mOrigin.x()));		points[5] = posY - (mSize.y() * mOrigin.y());

	points[6] = posX + (mSize.x() * (1 - mOrigin.x()));		points[7] = posY - (mSize.y() * mOrigin.y());
	points[8] = posX - (mSize.x() * mOrigin.x());		points[9] = posY + (mSize.y() * (1 - mOrigin.y()));
	points[10] = posX + (mSize.x() * (1 -mOrigin.x()));		points[11] = posY + (mSize.y() * (1 - mOrigin.y()));

	texs[0] = 0;		texs[1] = py;
	texs[2] = 0;		texs[3] = 0;
	texs[4] = px;		texs[5] = py;

	texs[6] = px;		texs[7] = py;
	texs[8] = 0;		texs[9] = 0;
	texs[10] = px;		texs[11] = 0;

	if(mFlipX)
	{
		for(int i = 0; i < 11; i += 2)
			if(texs[i] == px)
				texs[i] = 0;
			else
				texs[i] = px;
	}
	if(mFlipY)
	{
		for(int i = 1; i < 12; i += 2)
			if(texs[i] == py)
				texs[i] = 0;
			else
				texs[i] = py;
	}
}

void ImageComponent::drawImageArray(GLfloat* points, GLfloat* texs, GLubyte* colors, unsigned int numArrays)
{
	mTexture->bind();

	glEnable(GL_TEXTURE_2D);
	glEnable(GL_BLEND);
	glBlendFunc(GL_SRC_ALPHA, GL_ONE_MINUS_SRC_ALPHA);

	glEnableClientState(GL_VERTEX_ARRAY);
	glEnableClientState(GL_TEXTURE_COORD_ARRAY);

	if(colors != NULL)
	{
		glEnableClientState(GL_COLOR_ARRAY);
		glColorPointer(4, GL_UNSIGNED_BYTE, 0, colors);
	}

	glVertexPointer(2, GL_FLOAT, 0, points);
	glTexCoordPointer(2, GL_FLOAT, 0, texs);

	glDrawArrays(GL_TRIANGLES, 0, numArrays);

	glDisableClientState(GL_VERTEX_ARRAY);
	glDisableClientState(GL_TEXTURE_COORD_ARRAY);

	if(colors != NULL)
		glDisableClientState(GL_COLOR_ARRAY);

	glDisable(GL_TEXTURE_2D);
	glDisable(GL_BLEND);
}

bool ImageComponent::hasImage()
{
	return !mPath.empty();
}

void ImageComponent::copyScreen()
{
	mTexture.reset();

	mTexture = TextureResource::get(*mWindow->getResourceManager(), "");
	mTexture->initFromScreen();

	resize();
}


//
// Created by jnhkm on 2019-03-10.
//

#include "CombatHandler.h"
#include <sstream>
#include <boost/algorithm/string.hpp>

using handler::CombatHandler;

namespace handler {
    CombatHandler::CombatHandler(AccountHandler &accountHandler, WorldHandler &worldHandler) :
        combatInstances{},
        accountHandler(accountHandler),
        worldHandler(worldHandler),
        RNG(std::random_device{}()){}

    void
    CombatHandler::enterCombat(const Character &attacker, const NPC &defender) {
        this->combatInstances.emplace_back(CombatInstance{attacker.getId(), defender.getUniqueId()});
    }

    bool
    CombatHandler::rollMiss() {
        std::bernoulli_distribution missProc(BASE_MISS_CHANCE);

        return missProc(this->RNG);
    }

    bool
    CombatHandler::rollDodge(const float &modifier) {
        std::bernoulli_distribution dodgeProc(BASE_DODGE_CHANCE + modifier);

        return dodgeProc(this->RNG);
    }

    int
    CombatHandler::rollDamage(const int &modifier) {
        std::uniform_int_distribution damageRoll(BASE_MIN_DAMAGE + modifier, BASE_MAX_DAMAGE + modifier);

        return damageRoll(this->RNG);
    }

    bool
    CombatHandler::rollCritical(const float &modifier) {
        std::bernoulli_distribution criticalProc(BASE_CRITICAL_CHANCE + modifier);

        return criticalProc(this->RNG);
    }

    std::string
    CombatHandler::inflictDamage(NPC &npc) {
        std::ostringstream message;
        auto playerId = this->getOpponentId(npc.getUniqueId());
        auto client = this->accountHandler.getClientByPlayerId(playerId);

        if (client.id != AccountHandler::INVALID_ID) {
            auto &player = this->accountHandler.getPlayerByClient(client);

            if (this->rollMiss()) {
                message << "You miss your attack on " << npc.getShortDescription() << "!\n";
                return message.str();
            }

            auto dodgeValue = player.getMutableEquipment().getDodgeValue();
            if (this->rollDodge(dodgeValue)) {
                message << npc.getShortDescription() << " dodges your attack!\n";
                return message.str();
            }

            auto offenceValue = player.getMutableEquipment().getOffenceValue();
            auto damage = this->rollDamage(offenceValue);

            auto criticalValue = player.getMutableEquipment().getCriticalValue();
            if (this->rollCritical(criticalValue)) {
                damage = static_cast<int>(static_cast<float>(damage) * BASE_CRITICAL_DAMAGE_MULTIPLIER);
                message << "You strike a critical hit, inflicting ";
            } else {
                message << "You inflict ";
            }

            int oldHealth = npc.getHealth();
            int newHealth = std::max(npc.getHealth() - damage, 0);
            npc.setHealth(newHealth);

            message << (oldHealth - newHealth) << " HP worth of damage to "
                    << npc.getShortDescription() << " (" << npc.getHealth() << " HP remaining)\n";

            return message.str();

        } else {
            // Case in which the NPC being damaged is a decoy
            auto damage = this->rollDamage();

            int newHealth = std::max(npc.getHealth() - damage, 0);
            npc.setHealth(newHealth);

            return "";
        }
    }

    std::string
    CombatHandler::inflictDamage(Player &player) {
        std::ostringstream message;
        auto client = this->accountHandler.getClientByPlayerId(player.getId());
        auto roomId = this->accountHandler.getRoomIdByClient(client);
        auto npcId = this->getOpponentId(player.getId());
        auto npc = this->worldHandler.findRoom(roomId).getNpcByUniqueId(npcId);
        auto npcName = npc.getShortDescription();

        if (this->rollMiss()) {
            message << npcName << " misses their attack on you!\n";
            return message.str();
        }

        if (this->rollDodge()) {
            message << "You dodge " << npcName << "'s attack!\n";
            return message.str();
        }

        auto damage = this->rollDamage();
        damage -= player.getMutableEquipment().getDefenceValue();
        damage = std::max(damage, 0);

        if (this->rollCritical()) {
            damage = static_cast<int>(static_cast<float>(damage) * BASE_CRITICAL_DAMAGE_MULTIPLIER);
            message << npcName << " strikes a critical hit, inflicting ";
        } else {
            message << npcName << " inflicts ";
        }

        int oldHealth = player.getHealth();
        int newHealth = std::max(player.getHealth() - damage, 0);
        player.setHealth(newHealth);

        message << (oldHealth - newHealth) << " HP worth of damage on you ("
                << player.getHealth() << " HP remaining)\n";

        return message.str();
    }

    std::string
    CombatHandler::winEvent(Player &player, NPC &npc) {
        std::ostringstream message;
        message << "You won the battle!\n";
        this->exitCombat(player.getId(), npc.getUniqueId());

        // Remove NPC from room
        auto &room = this->worldHandler.findRoom(player.getCurrRoomID());
        auto &npcs = room.getMutableNpcs();
        auto npc_it = std::find_if(
            npcs.begin(),
            npcs.end(),
            [&npc](const auto &roomNpc) {
                return (roomNpc.getUniqueId() == npc.getUniqueId());
            });
        npcs.erase(npc_it);

        return message.str();
    }

    std::string
    CombatHandler::loseEvent(Player &player, NPC& npc) {
        std::ostringstream message;
        this->exitCombat(player.getId(), npc.getUniqueId());

        npc.setHealth(Character::STARTING_HEALTH);
        player.setHealth(Character::STARTING_HEALTH);

        auto startRoomId = Player::STARTING_LOCATION;
        // Check if starting room actually exists, pick first room in first area otherwise
        if (!this->worldHandler.roomExists(Player::STARTING_LOCATION)) {
            assert(!this->worldHandler.getWorld().getAreas().empty());
            auto firstArea = this->worldHandler.getWorld().getAreas().at(0);
            assert(!firstArea.getRooms().empty());
            startRoomId = firstArea.getRooms().at(0).getId();
        }

        auto currentRoomId = player.getCurrRoomID();
        auto startRoom = this->worldHandler.findRoom(startRoomId);
        this->worldHandler.movePlayer(player.getId(), currentRoomId, startRoomId);
        player.setCurrRoomID(startRoomId);

        message << "You have fallen in battle and respawned.\n";

        return message.str();
    }

    std::string
    CombatHandler::attack(const Connection &client, const std::string &targetName) {
        std::ostringstream message;

        if (targetName.empty()) {
            message << "You need to specify the name of the person (NPC) to attack.\n";
            return message.str();
        }

        auto roomId = this->accountHandler.getRoomIdByClient(client);
        auto &room = this->worldHandler.findRoom(roomId);

        try {
            // try getting npc with targetName, throws exception if not found.
            auto &npc = room.getNpcByKeyword(targetName);
            auto &player = this->accountHandler.getPlayerByClient(client);

            bool playerInCombat = this->isInCombat(player.getId());
            bool npcInCombat = this->isInCombat(npc.getUniqueId());
            bool npcInCombatWithPlayer = this->areInCombat(player.getId(), npc.getUniqueId());

            if (npcInCombat && !npcInCombatWithPlayer) {
                message << targetName << " is already engaged in combat with someone else!\n";
                return message.str();
            }

            if (playerInCombat && !npcInCombatWithPlayer) {
                message << "You are already engaged in combat with someone else!\n";
                return message.str();
            }

            if (!npcInCombat) {
                this->enterCombat(player, npc);
                message << "\nYou engage in combat with " << npc.getShortDescription() << "!\n";
            }

            message << "\n" << this->inflictDamage(npc);

            if (npc.getHealth() == 0) {
                message << this->winEvent(player, npc);
                return message.str();
            }

            message << this->inflictDamage(player);

            if (player.getHealth() == 0) {
                message << this->loseEvent(player, npc);
                return message.str();
            }

            auto combat_it = std::find_if(
                this->combatInstances.begin(),
                this->combatInstances.end(),
                [&player](const auto &combatState) {
                    return (combatState.attackerID == player.getId());
                }
            );

            if (combat_it != this->combatInstances.end()) {
                combat_it->endRound();
            }

            return message.str();

        } catch (const std::runtime_error &e) {
            message << "Invalid keyword.\n";
            return message.str();
        }
    }

    std::string
    CombatHandler::flee(const Connection &client) {
        std::ostringstream message;

        auto &player = this->accountHandler.getPlayerByClient(client);
        auto playerId = this->accountHandler.getPlayerIdByClient(client);

        if (!this->isInCombat(playerId)) {
            message << "You are in no danger to flee from.\n";
            return message.str();
        }

        auto roomId = this->accountHandler.getRoomIdByClient(client);
        auto doors = this->worldHandler.findRoom(roomId).getDoors();

        // Corner case where room has no doors
        if (doors.empty()) {
            // Player has halved chance of fleeing successfully
            std::bernoulli_distribution fleeChance{BASE_FLEE_CHANCE / 2};

            if (fleeChance(this->RNG)) {
                this->exitCombat(playerId);
                message << "You successfully flee from combat.\n";

            } else {
                auto opponentId = this->getOpponentId(playerId);
                auto &npc = this->worldHandler.findRoom(roomId).getNpcByUniqueId(opponentId);

                message << "\n"
                        << "You attempt to flee, but fail. ("
                        << (BASE_FLEE_CHANCE / 2 * 100) << "% chance of success)\n";

                message << this->inflictDamage(player);

                if (player.getHealth() == 0) {
                    message << this->loseEvent(player, npc);
                    return message.str();
                }

                auto combat_it = std::find_if(
                    this->combatInstances.begin(),
                    this->combatInstances.end(),
                    [&player](const auto &combatState) {
                        return (combatState.attackerID == player.getId());
                    }
                );

                if (combat_it != this->combatInstances.end()) {
                    combat_it->endRound();
                }
            }

            return message.str();
        }

        // Player has chance of fleeing to a random direction successfully
        std::bernoulli_distribution fleeChance(BASE_FLEE_CHANCE * doors.size());

        if (fleeChance(this->RNG)) {
            std::uniform_int_distribution<unsigned long> pickDoor(0, (doors.size() - 1));

            auto door = doors.at(pickDoor(this->RNG));
            auto direction = door.dir;
            auto destinationId = door.leadsTo;

            this->worldHandler.movePlayer(playerId, roomId, destinationId);
            this->accountHandler.setRoomIdByClient(client, destinationId);

            this->exitCombat(playerId);
            message << "You successfully flee to the " << direction << ".\n";

        } else {
            auto opponentId = this->getOpponentId(playerId);
            auto &npc = this->worldHandler.findRoom(roomId).getNpcByUniqueId(opponentId);

            message << "\n"
                    << "You attempt to flee, but fail. ("
                    << (BASE_FLEE_CHANCE * doors.size() * 100) << "% chance of success)\n";

            message << this->inflictDamage(player);

            if (player.getHealth() == 0) {
                message << this->loseEvent(player, npc);
                return message.str();
            }

            auto combat_it = std::find_if(
                combatInstances.begin(),
                combatInstances.end(),
                [&player](const auto &combatState) {
                    return (combatState.attackerID == player.getId());
                }
            );

            if (combat_it != this->combatInstances.end()) {
                combat_it->endRound();
            }
        }

        return message.str();
    }

    bool
    CombatHandler::isInCombat(const model::ID &id) {
        auto combat_it = std::find_if(
            this->combatInstances.begin(),
            this->combatInstances.end(),
            [&id](const auto &combatState) {
                return (combatState.attackerID == id || combatState.defenderID == id);
            }
        );

        return combat_it != this->combatInstances.end();
    };

    bool
    CombatHandler::areInCombat(const model::ID &combatantId1, const model::ID &combatantId2) {
        auto combat_it = std::find(
                this->combatInstances.begin(),
                this->combatInstances.end(),
                CombatInstance{combatantId1, combatantId2}
        );

        return combat_it != this->combatInstances.end();
    }

    void
    CombatHandler::exitCombat(const model::ID &combatantId1, const model::ID &combatantId2) {
        auto it = std::find(
            this->combatInstances.begin(),
            this->combatInstances.end(),
            CombatInstance{combatantId1, combatantId2});

        if (it != combatInstances.end()) {
            combatInstances.erase(it);
        }
    }

    void
    CombatHandler::exitCombat(const model::ID &id) {
        auto combat_it = std::find_if(
            this->combatInstances.begin(),
            this->combatInstances.end(),
            [&id](const auto &combatState) {
                return (combatState.attackerID == id || combatState.defenderID == id);
            }
        );

        if (combat_it != combatInstances.end()) {
            combatInstances.erase(combat_it);
        }
    }

    model::ID
    CombatHandler::getOpponentId(const model::ID &id) {
        auto combat_it = std::find_if(
            this->combatInstances.begin(),
            this->combatInstances.end(),
            [&id](const auto &combatInstance) {
                return (combatInstance.attackerID == id || combatInstance.defenderID == id);
            }
        );

        if (combat_it == this->combatInstances.end()) {
            return 0;
        } else if (combat_it->attackerID == id) {
            return combat_it->defenderID;
        } else {
            return combat_it->attackerID;
        }
    }

    void
    CombatHandler::replaceWithDecoy(const Player &player) {
        if (!isInCombat(player.getId())) {
            return;
        }

        auto playerId = player.getId();
        model::ID decoyId = -player.getId();
        std::vector<std::string> keywords = {boost::algorithm::to_lower_copy(player.getUsername())};
        std::vector<std::string> description = {"Upon closer inspection, you realize this is just a conjured decoy of " +
                                                player.getUsername() + "."};
        std::string shortDescription = "'" + player.getUsername() + "'";

        NPC decoy{decoyId, keywords, description, shortDescription, {}};
        decoy.setUniqueId(decoyId);

        auto client = this->accountHandler.getClientByPlayerId(playerId);
        auto roomId = this->accountHandler.getRoomIdByClient(client);
        this->worldHandler.findRoom(roomId).addNPC(decoy);

        for (auto &combatInstance : this->combatInstances) {
            if (combatInstance.attackerID == playerId) {
                combatInstance.attackerID = decoyId;
                break;
            };
        }
    };

    void
    CombatHandler::handleLogout(const Connection &client) {
        auto playerId = this->accountHandler.getPlayerIdByClient(client);

        if (this->isInCombat(playerId)) {
            auto roomId = this->accountHandler.getRoomIdByClient(client);
            auto npcId = this->getOpponentId(playerId);
            auto &npc = this->worldHandler.findRoom(roomId).getNpcByUniqueId(npcId);

            this->exitCombat(playerId);
            npc.setHealth(Character::STARTING_HEALTH);
        }
    }

    void
    CombatHandler::processCycle(std::deque<Message> &messages) {
        for (auto &combatInstance : this->combatInstances) {
            if (combatInstance.roundCyclesRemaining > 0) {
                combatInstance.decrement();

            } else {
                auto client = this->accountHandler.getClientByPlayerId(combatInstance.attackerID);

                if (client.id != AccountHandler::INVALID_ID) {
                    auto &player = this->accountHandler.getPlayerByClient(client);
                    auto roomId = this->accountHandler.getRoomIdByClient(client);
                    auto &npc = this->worldHandler.findRoom(roomId).getNpcByUniqueId(combatInstance.defenderID);

                    std::ostringstream message;
                    message << "\n" << this->inflictDamage(player);

                    if (player.getHealth() == 0) {
                        message << this->loseEvent(player, npc);
                    }

                    messages.push_back({client, message.str()});
                    combatInstance.endRound();

                } else {
                    model::ID roomId = 0;
                    // attacker was replaced by decoy
                    auto areas = this->worldHandler.getWorld().getAreas();

                    for (auto &area : areas) {

                        auto rooms = area.getRooms();
                        for (auto &room : rooms) {

                            auto npcs = room.getImmutableNpcs();
                            for (const auto &npc : npcs) {
                                if (npc.getId() == combatInstance.attackerID) {
                                    roomId = room.getId();
                                }
                            }
                        }
                    }

                    // decoy does not exist anywhere, process next combat instance
                    if (roomId == 0) {
                        continue;
                    }

                    auto &decoy= this->worldHandler.findRoom(roomId).getNpcByUniqueId(combatInstance.attackerID);
                    auto &npc = this->worldHandler.findRoom(roomId).getNpcByUniqueId(combatInstance.defenderID);

                    std::ostringstream message;
                    this->inflictDamage(decoy);

                    if (decoy.getHealth() == 0) {
                        // Delete decoy
                        this->exitCombat(decoy.getId());

                        auto &npcs = this->worldHandler.findRoom(roomId).getMutableNpcs();
                        auto npc_it = std::find_if(
                            npcs.begin(),
                            npcs.end(),
                            [&decoy](const auto &roomNpc) {
                                return roomNpc.getId() == decoy.getId();
                            });
                        npcs.erase(npc_it);
                        npc.setHealth(Character::STARTING_HEALTH);
                    }

                    combatInstance.endRound();

                }
            }
        }
    }
}