//
//   Copyright (c) 2004, 2005, 2006, 2007   Troy D. Straszheim
//
//   This file is part of IceTray.
//
//   IceTray is free software; you can redistribute it and/or modify
//   it under the terms of the GNU General Public License as published by
//   the Free Software Foundation; either version 3 of the License, or
//   (at your option) any later version.
//
//   IceTray is distributed in the hope that it will be useful,
//   but WITHOUT ANY WARRANTY; without even the implied warranty of
//   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
//   GNU General Public License for more details.
//
//   You should have received a copy of the GNU General Public License
//   along with this program.  If not, see <http://www.gnu.org/licenses/>.
//

#include <boost/preprocessor.hpp>
#include <vector>

#include <dataclasses/I3Constants.h>

using namespace boost::python;

#define I3_CONSTANT_NAMES			\
  (c)(n_ice_phase)(n_ice_group)(n_ice)		\
  (theta_cherenkov)(c_ice)(pi)(e)(NA)		\
  (SurfaceElev)(OriginElev)(zIceTop)		\
  (Coordinate_shift_x)(Coordinate_shift_y)	\
  (Coordinate_shift_z)(dt_window_l)(dt_window_h)

#define I3CONSTANT_DEF(r,data,t) \
  .def_readonly(BOOST_PP_STRINGIZE(t), BOOST_PP_CAT(I3Constants::, t))

// dummy class just used as a namespace.
namespace {
  struct dummy { };
}

void register_I3Constants()
{
  class_<dummy>("I3Constants")
    .def_readonly("c", I3Constants::c, "Speed of light in vacuum")
    .def_readonly("n_ice_phase", I3Constants::n_ice_phase, "the average phase velocity for wavelength of 400nm")
    .def_readonly("n_ice_group", I3Constants::n_ice_group, "avg group velocity for 400nm")
    .def_readonly("n_ice", I3Constants::n_ice, "index of refraction of ice")
    .def_readonly("theta_cherenkov", I3Constants::theta_cherenkov, "cherenkov angle in ice")
    .def_readonly("c_ice", I3Constants::c_ice, "c in ice")
    .def_readonly("pi", I3Constants::pi, "what is pi doing duplicated here")
    .def_readonly("e", I3Constants::e, "OY-ler's number.  awww yeah you know how we do.")
    .def_readonly("NA", I3Constants::NA, "avogadro's number")
    .def_readonly("SurfaceElev", I3Constants::SurfaceElev, "elevation of surface (amanda hole 4)")
    .def_readonly("OriginElev", I3Constants::OriginElev, "elev of icecube origin, by definition")
    .def_readonly("zIceTop", I3Constants::zIceTop, " Z-coordinate of IceTop (Origin Depth)"
		  "Obtained from DEFINED elevation of origin and MEASURED surface elevation")
    .def_readonly("Coordinate_shift_x", I3Constants::Coordinate_shift_x, "conversion between icecube and amanda")
    .def_readonly("Coordinate_shift_y", I3Constants::Coordinate_shift_y, "conversion between icecube and amanda")
    .def_readonly("Coordinate_shift_z", I3Constants::Coordinate_shift_z, "conversion between icecube and amanda")
    .def_readonly("dt_window_l", I3Constants::dt_window_l, "default value for time residuals")
    .def_readonly("dt_window_h", I3Constants::dt_window_h, "default value for time residuals")
    .def( freeze() )
    ;
}


/*
 Copyright (C) 2019 Quaternion Risk Management Ltd
 All rights reserved.

 This file is part of ORE, a free-software/open-source library
 for transparent pricing and risk analysis - http://opensourcerisk.org

 ORE is free software: you can redistribute it and/or modify it
 under the terms of the Modified BSD License.  You should have received a
 copy of the license along with this program.
 The license is also available online at <http://opensourcerisk.org>

 This program is distributed on the basis that it will form a useful
 contribution to risk analytics and model standardisation, but WITHOUT
 ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
 FITNESS FOR A PARTICULAR PURPOSE. See the license for more details.
*/
#include <boost/date_time.hpp>
#include <boost/make_shared.hpp>
#include <ql/cashflows/cashflows.hpp>
#include <ql/cashflows/coupon.hpp>
#include <ql/cashflows/simplecashflow.hpp>
#include <ql/event.hpp>
#include <ql/quotes/compositequote.hpp>
#include <ql/termstructures/credit/flathazardrate.hpp>
#include <ql/termstructures/yield/zerospreadedtermstructure.hpp>
#include <qle/pricingengines/discountingforwardbondengine.hpp>

namespace QuantExt {

DiscountingForwardBondEngine::DiscountingForwardBondEngine(
    const Handle<YieldTermStructure>& discountCurve, const Handle<YieldTermStructure>& incomeCurve,
    const Handle<YieldTermStructure>& bondReferenceYieldCurve, const Handle<Quote>& bondSpread,
    const Handle<DefaultProbabilityTermStructure>& bondDefaultCurve, const Handle<Quote>& bondRecoveryRate,
    Period timestepPeriod, boost::optional<bool> includeSettlementDateFlows, const Date& settlementDate,
    const Date& npvDate)
    : discountCurve_(discountCurve), incomeCurve_(incomeCurve), bondReferenceYieldCurve_(bondReferenceYieldCurve),
      bondSpread_(bondSpread), bondDefaultCurve_(bondDefaultCurve), bondRecoveryRate_(bondRecoveryRate),
      timestepPeriod_(timestepPeriod), includeSettlementDateFlows_(includeSettlementDateFlows),
      settlementDate_(settlementDate), npvDate_(npvDate) {

    bondReferenceYieldCurve_ =
        bondSpread_.empty() ? bondReferenceYieldCurve
                            : Handle<YieldTermStructure>(
                                  boost::make_shared<ZeroSpreadedTermStructure>(bondReferenceYieldCurve, bondSpread_));
    registerWith(discountCurve_);           // curve for discounting of the forward derivative contract. OIS, usually.
    registerWith(incomeCurve_);             // this is a curve for compounding of the bond
    registerWith(bondReferenceYieldCurve_); // this is the bond reference curve, for discounting, usually RePo
    registerWith(bondSpread_);
    registerWith(bondDefaultCurve_);
    registerWith(bondRecoveryRate_);
}

void DiscountingForwardBondEngine::calculate() const {
    // Do some checks on data
    QL_REQUIRE(!discountCurve_.empty(), "discounting term structure handle is empty");
    QL_REQUIRE(!incomeCurve_.empty(), "income term structure handle is empty");
    QL_REQUIRE(!bondReferenceYieldCurve_.empty(), "bond rerference term structure handle is empty");

    Date npvDate = npvDate_; // this is today when the valuation occurs
    if (npvDate == Null<Date>()) {
        npvDate = (*discountCurve_)->referenceDate();
    }
    Date settlementDate = settlementDate_; //
    if (settlementDate == Null<Date>()) {
        settlementDate = (*discountCurve_)->referenceDate();
    }

    Date maturityDate =
        arguments_.fwdMaturityDate; // this is the date when the forward is executed, i.e. cash and bond change hands

    Real cmpPayment = arguments_.compensationPayment;
    if (cmpPayment == Null<Real>()) {
        cmpPayment = 0.0;
    }
    Date cmpPaymentDate = arguments_.compensationPaymentDate;
    if (cmpPaymentDate == Null<Date>()) {
        cmpPaymentDate = npvDate;
    }

    // in case that the premium payment has occured in the past, we set the amount to 0. the date itself is set to the
    // npvDate to have a valid date for "discounting"
    Date cmpPaymentDate_use = cmpPaymentDate >= npvDate ? cmpPaymentDate : maturityDate;
    cmpPayment = cmpPaymentDate >= npvDate ? cmpPayment : 0.0; // premium cashflow is not relevant for npv if in the
                                                               // past

    // initialize
    results_.value = 0.0;               // this is today's npv of the forward contract
    results_.underlyingSpotValue = 0.0; // this is today's value of the "restricted bond". Restricted means that only
                                        // cashflows after maturity are taken into account.
    results_.forwardValue = 0.0;        // this the value of the forward contract just before maturity

    bool dirty = arguments_.settlementDirty;

    results_.underlyingSpotValue = calculateBondNpv(npvDate, maturityDate); // cashflows before maturity will be ignored

    boost::tie(results_.forwardValue, results_.value) = calculateForwardContractPresentValue(
        results_.underlyingSpotValue, cmpPayment, npvDate, maturityDate, cmpPaymentDate_use, dirty);
}

Real DiscountingForwardBondEngine::calculateBondNpv(Date npvDate, Date computeDate) const {
    Real npvValue = 0.0;
    Size numCoupons = 0;
    bool hasLiveCashFlow = false;

    // handle case where we wish to price simply with benchmark curve and scalar security spread
    // i.e. credit curve term structure (and recovery) have not been specified
    // we set the default probability and recovery rate to zero in this instance (issuer credit worthiness already
    // captured within security spread)
    boost::shared_ptr<DefaultProbabilityTermStructure> creditCurvePtr =
        bondDefaultCurve_.empty()
            ? boost::make_shared<QuantLib::FlatHazardRate>(npvDate, 0.0, bondReferenceYieldCurve_->dayCounter())
            : bondDefaultCurve_.currentLink();
    Rate recoveryVal = bondRecoveryRate_.empty() ? 0.0 : bondRecoveryRate_->value(); // setup default bond recovery rate

    // load the shared pointer into bd
    boost::shared_ptr<Bond> bd = arguments_.underlying;
    for (Size i = 0; i < bd->cashflows().size(); i++) {
        // Recovery amount is computed over the whole time interval (npvDate,maturityOfBond)

        if (bd->cashflows()[i]->hasOccurred(
                computeDate, includeSettlementDateFlows_)) // Cashflows before computeDate not relevant for npv
            continue;

        /* The amount recovered in the case of default is the recoveryrate*Notional*Probability of
           Default; this is added to the NPV value. For coupon bonds the coupon periods are taken
           as the timesteps for integrating over the probability of default.
        */
        boost::shared_ptr<Coupon> coupon = boost::dynamic_pointer_cast<Coupon>(bd->cashflows()[i]);

        if (coupon) {
            numCoupons++;
            Date startDate = coupon->accrualStartDate();
            Date endDate = coupon->accrualEndDate();
            Date effectiveStartDate = (startDate <= computeDate && computeDate <= endDate) ? computeDate : startDate;
            Date defaultDate = effectiveStartDate + (endDate - effectiveStartDate) / 2;
            Probability P = creditCurvePtr->defaultProbability(effectiveStartDate, endDate);

            npvValue += coupon->nominal() * recoveryVal * P * bondReferenceYieldCurve_->discount(defaultDate);
        }

        hasLiveCashFlow = true; // check if a cashflow  is available after the date of valuation.

        // Coupon value is discounted future payment times the survival probability
        Probability S = creditCurvePtr->survivalProbability(bd->cashflows()[i]->date()) /
                        creditCurvePtr->survivalProbability(computeDate);
        npvValue += bd->cashflows()[i]->amount() * S * bondReferenceYieldCurve_->discount(bd->cashflows()[i]->date());
    }

    // the ql instrument might not yet be expired and still have not anything to value if
    // the computeDate > evaluation date
    if (!hasLiveCashFlow)
        return 0.0;

    if (bd->cashflows().size() > 1 && numCoupons == 0) {
        QL_FAIL("DiscountingForwardBondEngine does not support bonds with multiple cashflows but no coupons");
    }

    boost::shared_ptr<Coupon> firstCoupon = boost::dynamic_pointer_cast<Coupon>(bd->cashflows()[0]);
    if (firstCoupon) {
        Date startDate = computeDate; // face value recovery starting with computeDate
        while (startDate < bd->cashflows()[0]->date()) {
            Date stepDate = startDate + timestepPeriod_;
            Date endDate = (stepDate > bd->cashflows()[0]->date()) ? bd->cashflows()[0]->date() : stepDate;
            Date defaultDate = startDate + (endDate - startDate) / 2;
            Probability P = creditCurvePtr->defaultProbability(startDate, endDate);

            npvValue += firstCoupon->nominal() * recoveryVal * P * bondReferenceYieldCurve_->discount(defaultDate);
            startDate = stepDate;
        }
    }
    /* If there are no coupon, as in a Zero Bond, we must integrate over the entire period from npv date to
       maturity. The timestepPeriod specified is used as provide the steps for the integration. This only applies
       to bonds with 1 cashflow, identified as a final redemption payment.
    */
    if (bd->cashflows().size() == 1) {
        boost::shared_ptr<Redemption> redemption = boost::dynamic_pointer_cast<Redemption>(bd->cashflows()[0]);
        if (redemption) {
            Date startDate = computeDate;
            while (startDate < redemption->date()) {
                Date stepDate = startDate + timestepPeriod_;
                Date endDate = (stepDate > redemption->date()) ? redemption->date() : stepDate;
                Date defaultDate = startDate + (endDate - startDate) / 2;
                Probability P = creditCurvePtr->defaultProbability(startDate, endDate);

                npvValue += redemption->amount() * recoveryVal * P * bondReferenceYieldCurve_->discount(defaultDate);
                startDate = stepDate;
            }
        }
    }
    return npvValue * arguments_.bondNotional;
}

boost::tuple<Real, Real> DiscountingForwardBondEngine::calculateForwardContractPresentValue(
    Real spotValue, Real cmpPayment, Date npvDate, Date computeDate, Date cmpPaymentDate, bool dirty) const {
    // here we go with the true forward computation
    Real forwardBondValue = 0.0;
    Real forwardContractPresentValue = 0.0;
    Real forwardContractForwardValue = 0.0;

    // handle case where we wish to price simply with benchmark curve and scalar security spread
    // i.e. credit curve term structure (and recovery) have not been specified
    // we set the default probability and recovery rate to zero in this instance (issuer credit worthiness already
    // captured within security spread)
    boost::shared_ptr<DefaultProbabilityTermStructure> creditCurvePtr =
        bondDefaultCurve_.empty()
            ? boost::make_shared<QuantLib::FlatHazardRate>(npvDate, 0.0, bondReferenceYieldCurve_->dayCounter())
            : bondDefaultCurve_.currentLink();
    Rate recoveryVal = bondRecoveryRate_.empty() ? 0.0 : bondRecoveryRate_->value(); // setup default bond recovery rate

    // load the shared pointer into bd
    boost::shared_ptr<Bond> bd = arguments_.underlying;

    // the case of dirty strike corresponds here to an accrual of 0.0. This will be convenient in the code.
    Real accruedAmount =
        dirty ? 0.0 : bd->accruedAmount(computeDate) * bd->notional(computeDate) / 100 * arguments_.bondNotional;

    /* Discounting and compounding, taking account of possible bond default before delivery*/

    forwardBondValue =
        spotValue / (incomeCurve_->discount(computeDate)); // compounding to date of maturity of forward contract

    // Subtract strike at maturity. Regarding accrual (i.e. strike is given clean vs dirty) there are two
    // cases: long or short.

    // Long: forwardBondValue - strike_dirt = (forwardBondValue - accrual) - strike_clean
    // Short: strike_dirt - forwardBondValue = strike_clean - (forwardBondValue - accrual)
    // In total:
    forwardContractForwardValue = (*arguments_.payoff)(forwardBondValue - accruedAmount);

    // forwardContractPresentValue adjusted for potential default before computeDate:
    forwardContractPresentValue =
        forwardContractForwardValue * (discountCurve_->discount(computeDate)) *
            creditCurvePtr->survivalProbability(computeDate) -
        cmpPayment *
            (discountCurve_->discount(cmpPaymentDate)); // The forward is a derivative. We use "OIS curve" to discount.
                                                        // We subtract the potential payment due to Premium.

    // Take account of face value recovery:
    // A) Recovery for time period when coupons are present
    for (Size i = 0; i < bd->cashflows().size(); i++) {
        if (bd->cashflows()[i]->hasOccurred(
                npvDate, includeSettlementDateFlows_)) // Cashflows before npvDate not relevant for npv
            continue;
        if (bd->cashflows()[i]->date() >=
            computeDate) // Cashflows after computeDate do not fall into the forward period
            continue;
        boost::shared_ptr<Coupon> coupon = boost::dynamic_pointer_cast<Coupon>(bd->cashflows()[i]);

        if (coupon) {
            Date startDate = coupon->accrualStartDate();
            Date endDate = coupon->accrualEndDate();
            Date effectiveStartDate = (startDate <= npvDate && npvDate <= endDate) ? npvDate : startDate;
            Date effectiveEndDate = (startDate <= computeDate && computeDate <= endDate) ? computeDate : endDate;
            Date defaultDate = effectiveStartDate + (effectiveEndDate - effectiveStartDate) / 2;
            Probability P = creditCurvePtr->defaultProbability(effectiveStartDate, effectiveEndDate);

            forwardContractPresentValue +=
                (*arguments_.payoff)(coupon->nominal() * arguments_.bondNotional * recoveryVal - accruedAmount) * P *
                (discountCurve_->discount(defaultDate));
        }
    }

    // B) Recovery for time period before coupons are present
    boost::shared_ptr<Coupon> firstCoupon = boost::dynamic_pointer_cast<Coupon>(bd->cashflows()[0]);
    if (firstCoupon) {
        Date startDate = npvDate; // face value recovery starting with npvDate
        Date stopDate = std::min(bd->cashflows()[0]->date(), computeDate);
        while (startDate < stopDate) {
            Date stepDate = startDate + timestepPeriod_;
            Date endDate = (stepDate > stopDate) ? stopDate : stepDate;
            Date defaultDate = startDate + (endDate - startDate) / 2;
            Probability P = creditCurvePtr->defaultProbability(startDate, endDate);

            forwardContractPresentValue +=
                (*arguments_.payoff)(firstCoupon->nominal() * arguments_.bondNotional * recoveryVal - accruedAmount) *
                P * (discountCurve_->discount(defaultDate));
            startDate = stepDate;
        }
    }
    // C) ZCB
    /* If there are no coupon, as in a Zero Bond, we must integrate over the entire period from npv date to
       maturity. The timestepPeriod specified is used as provide the steps for the integration. This only applies
       to bonds with 1 cashflow, identified as a final redemption payment.
    */
    if (bd->cashflows().size() == 1) {
        boost::shared_ptr<Redemption> redemption = boost::dynamic_pointer_cast<Redemption>(bd->cashflows()[0]);
        if (redemption) {
            Date startDate = npvDate;
            while (startDate < redemption->date()) {
                Date stepDate = startDate + timestepPeriod_;
                Date endDate = (stepDate > redemption->date()) ? redemption->date() : stepDate;
                Date defaultDate = startDate + (endDate - startDate) / 2;
                Probability P = creditCurvePtr->defaultProbability(startDate, endDate);

                forwardContractPresentValue +=
                    (*arguments_.payoff)(redemption->amount() * arguments_.bondNotional * recoveryVal - accruedAmount) *
                    P * (discountCurve_->discount(defaultDate));
                startDate = stepDate;
            }
        }
    }

    return boost::make_tuple(forwardContractForwardValue, forwardContractPresentValue);
}

} // namespace QuantExt


// -----------------------------------------------------------------------
//
// Copyright (C) 2020  - David Fernández Castellanos
//
// This file is part of the MEPLS software. You can use it, redistribute
// it, and/or modify it under the terms of the Creative Commons Attribution
// 4.0 International Public License. The full text of the license can be
// found in the file LICENSE at the top level of the MEPLS distribution.
//
// -----------------------------------------------------------------------

#include <example.h>
#include <mepls/utils.h>
#include <mepls/solver.h>
#include <mepls/system.h>
#include <mepls/event.h>
#include <mepls/history.h>

// MEPLS built-in dynamics
#include <mepls/dynamics.h>

// to parse command line arguments
#include <cmdparser.hpp>

// to parse input parameters files
#include <deal.II/base/parameter_handler.h>
#include <deal.II/base/path_search.h>

// to save output data in JSON format
#include <boost/property_tree/json_parser.hpp>

struct Parameters
{
	// these are the parameters and we will use in the simulation, initialized with some default values
	unsigned int seed = 1234567;
	unsigned int Nx = 32;
	unsigned int Ny = 32;
	double G = 30.;
	double nu = 0.3;
	double gamma = 0.05;
	double k = 6.;
	double strain_limit = 0.05;
	double lambda_init = 1.;
	double lambda_renew = 1.;
	std::string filename = "out.json";

	void declare_entries(dealii::ParameterHandler &prm)
	{
		// We declare the entries of the parameters text file. Each entry matches the name of a
		// simulation parameters (although it doesn't need to) and has a default value. We use as
		// the default value is the same value of the variables declared above

		prm.enter_subsection("Section1");

		prm.declare_entry("seed", mepls::utils::str::to_string(seed), dealii::Patterns::Integer(0),
						  "");
		prm.declare_entry("Nx", mepls::utils::str::to_string(Nx), dealii::Patterns::Integer(0), "");
		prm.declare_entry("Ny", mepls::utils::str::to_string(Ny), dealii::Patterns::Integer(0), "");
		prm.declare_entry("G", mepls::utils::str::to_string(G), dealii::Patterns::Double(0.0), "");
		prm.declare_entry("nu", mepls::utils::str::to_string(nu), dealii::Patterns::Double(0.0),
						  "");
		prm.declare_entry("gamma", mepls::utils::str::to_string(gamma),
						  dealii::Patterns::Double(0.0), "");
		prm.declare_entry("strain_limit", mepls::utils::str::to_string(strain_limit),
						  dealii::Patterns::Double(0.0), "");
		prm.declare_entry("lambda_init", mepls::utils::str::to_string(lambda_init),
						  dealii::Patterns::Double(0.0), "");
		prm.declare_entry("lambda_renew", mepls::utils::str::to_string(lambda_renew),
						  dealii::Patterns::Double(0.0), "");
		prm.declare_entry("k", mepls::utils::str::to_string(k), dealii::Patterns::Double(0.0), "");
		prm.declare_entry("filename", filename, dealii::Patterns::FileName(), "");

		prm.leave_subsection();
	}

	void load_entries(dealii::ParameterHandler &prm)
	{
		prm.enter_subsection("Section1");

		// We define how each parameter gets its value from a parameters file entry

		seed = prm.get_integer("seed");
		Nx = prm.get_integer("Nx");
		Ny = prm.get_integer("Ny");
		G = prm.get_double("G");
		nu = prm.get_double("nu");
		gamma = prm.get_double("gamma");
		k = prm.get_double("k");
		strain_limit = prm.get_double("strain_limit");
		lambda_init = prm.get_double("lambda_init");
		lambda_renew = prm.get_double("lambda_renew");
		filename = prm.get("filename");

		prm.leave_subsection();
	}

	void load_file(const std::string &filename)
	{
		// load a parameters file

		dealii::ParameterHandler prm;
		declare_entries(prm);
		prm.parse_input(filename);
		load_entries(prm);
	}

	void generate_file(const std::string &filename)
	{
		// generate a parameters file template

		std::ofstream outfile(filename);
		dealii::ParameterHandler prm;
		declare_entries(prm);
		prm.print_parameters(outfile, dealii::ParameterHandler::OutputStyle::Text);
	}

};

template<int dim>
void write_data(
	const mepls::history::History<dim> &sim_history, const Parameters &p)
{
	boost::property_tree::ptree data_tree;

	// write some metadata, such as a simulation name and description
	data_tree.put("Name", "Step4");
	data_tree.put("Description", "System driven in the athermal quasistatic limit");

	// write the simulation parameters
	data_tree.put("Parameters.dim", 2);
	data_tree.put("Parameters.seed", p.seed);
	data_tree.put("Parameters.Nx", p.Nx);
	data_tree.put("Parameters.Ny", p.Ny);
	data_tree.put("Parameters.G", p.G);
	data_tree.put("Parameters.nu", p.nu);
	data_tree.put("Parameters.gamma", p.gamma);
	data_tree.put("Parameters.lambda_renew", p.lambda_renew);
	data_tree.put("Parameters.lambda_init", p.lambda_init);
	data_tree.put("Parameters.k", p.k);
	data_tree.put("Parameters.strain_limit", p.strain_limit);

	// write some metadata, such as a simulation name and description
	data_tree.put("Name", "Step4");
	data_tree.put("Description", "System driven in the athermal quasistatic limit");

	// write the simulation parameters
	data_tree.put("Parameters.dim", 2);
	data_tree.put("Parameters.seed", p.seed);
	data_tree.put("Parameters.Nx", p.Nx);
	data_tree.put("Parameters.Ny", p.Ny);
	data_tree.put("Parameters.G", p.G);
	data_tree.put("Parameters.nu", p.nu);
	data_tree.put("Parameters.gamma", p.gamma);
	data_tree.put("Parameters.lambda_renew", p.lambda_renew);
	data_tree.put("Parameters.lambda_init", p.lambda_init);
	data_tree.put("Parameters.k", p.k);
	data_tree.put("Parameters.strain_limit", p.strain_limit);

	// We write the event histories with CSV format to a string.
	// Here, we write only the columns of interest, but there are more available (see the
	// documentation).
	std::ostringstream plastic_events_csv;
	plastic_events_csv << "index,element,eigenstrain_00,eigenstrain_11,eigenstrain_01\n";
	for(auto &row : sim_history.plastic)
		plastic_events_csv << row.index << "," << row.element << "," << row.eigenstrain_00 << ","
						   << row.eigenstrain_11 << "," << row.eigenstrain_01 << "\n";

	data_tree.put("Data.plastic_events", plastic_events_csv.str());

	std::ostringstream driving_events_csv;
	driving_events_csv << "index,dext_stress,dtotal_strain\n";
	for(auto &row : sim_history.driving)
		driving_events_csv << row.index << "," << row.dext_stress << "," << row.dtotal_strain
						   << "\n";

	data_tree.put("Data.driving_events", driving_events_csv.str());

	std::ostringstream macro_evolution_csv;
	macro_evolution_csv << "index,ext_stress,total_strain,time,av_vm_stress,av_vm_plastic_strain\n";
	for(auto &row : sim_history.macro_evolution)
		macro_evolution_csv << row.index << "," << row.ext_stress << "," << row.total_strain << ","
							<< row.time << "," << row.av_vm_stress << ","
							<< row.av_vm_plastic_strain << "\n";

	data_tree.put("Data.macro_evolution", macro_evolution_csv.str());

	// write the data tree to a JSON file
	std::ofstream output_file(p.filename);
	boost::property_tree::json_parser::write_json(output_file, data_tree);
	output_file.close();
}

void run(const Parameters &p)
{
	//----- SETUP ------

	// we do the same as in the previous tutorial, but this time we use the
	// parameters from the input Parameters object

	constexpr unsigned dim = 2;
	std::mt19937 generator(p.seed);

	dealii::SymmetricTensor<4, dim> C = mepls::utils::tensor::make_isotropic_stiffness<dim>(p.G,
																							p.nu);

	mepls::element::Vector<dim> elements;

	// we also store the elements into a vector that knows the derived class, so we can access
	// the example::element::Scalar<dim>::conf struct later (this member cannot be accessed
	// through the pointers to the base mepls::element::Element<dim>)
	std::vector<example::element::Scalar<dim> *> elements_scalar;

	for(double n = 0; n < p.Nx * p.Ny; ++n)
	{
		example::element::Scalar<dim>::Config conf;
		conf.number = n;
		conf.gamma = p.gamma;
		conf.lambda = p.lambda_init;
		conf.k = p.k;

		auto element = new example::element::Scalar<dim>(conf, generator);
		element->C(C);

		elements_scalar.push_back(element);
		elements.push_back(element);
	}

	mepls::elasticity_solver::LeesEdwards<dim> solver(p.Nx, p.Ny);
	for(auto &element : elements)
		solver.set_elastic_properties(element->number(), element->C());
	solver.setup_and_assembly();

	mepls::element::calculate_ext_stress_coefficients(elements, solver);
	mepls::element::calculate_local_stress_coefficients_central(elements, solver);

	mepls::system::Standard<dim> system(elements, solver, generator);

	mepls::history::History<dim> sim_history("Simulation_history");

	system.set_history(sim_history);

	// when the elements were created, their slip systems were initialized with thresholds from a
	// Weibull distribution with scale lambda_init. We change it now to lambda_renew. Therefore,
	// when plastic deformation occurs, the renew local thresholds will have a different
	// average
	for(auto &element : elements_scalar)
		element->conf.lambda = p.lambda_renew;

	//----- DYNAMICS ------

	// This object will allow us to check for different conditions by which
	// the simulation might stop. Different conditions might be checked, but as long as one of
	// them evaluates to false, when continue_simulation() is called it will return false
	mepls::utils::ContinueSimulation continue_simulation;

	sim_history.add_macro(system);

	// run while the continue_simulation object says so
	while(continue_simulation())
	{
		std::cout << system.macrostate["total_strain"] << " " << system.macrostate["ext_stress"]
				  << std::endl;

		// these are the same dynamics we implemented in the previous tutorial, but using MEPLS built-in

		// apply an external strain increment of 0.01%
		mepls::dynamics::finite_extremal_dynamics_step(1e-4, system);
		sim_history.add_macro(system);

		// perform and avalanche of slip events. By passing the continue_simulation object,
		// the relaxation function can set its own condition for stopping the simulation.
		// Specifically, it will check if the avalanche size overcomes a certain maximum upper limit
		mepls::dynamics::relaxation(system, continue_simulation);
		sim_history.add_macro(system);

		// check if the strain has reached the strain limit. If it has, the next time continue_simulation() is called
		// it will return false, so we will exit the main loop
		continue_simulation(system.macrostate["total_strain"] < p.strain_limit,
							"total strain limit reached");

	}

	// print the message of the stopping condition that was met
	std::cout << continue_simulation << std::endl;

	for(auto &element : elements)
		delete element;

	write_data(sim_history, p);
}

int main(int argc, char *argv[])
{
	// Read the command line arguments. We define the -f 'filename' to pass the
	// path to the parameters file
	cli::Parser parser(argc, argv);
	parser.set_optional<std::string>("f", "file", "./default.prm",
									 "Name of the input configuration file");
	parser.run_and_exit_if_error();
	// you can check https://github.com/FlorianRappl/CmdParser for a further documentation of cli::Parser

	// Create the parameters object
	Parameters p;

	// We try to load the parameters file, but if it doesn't exist, we generate a new one with the name
	// default.prm and default values
	try
	{
		p.load_file(parser.get<std::string>("f"));
	}
	catch(dealii::PathSearch::ExcFileNotFound &)
	{
		p.generate_file(parser.get<std::string>("f"));
		std::cout << "Configuration file " << parser.get<std::string>("f") << " created"
				  << std::endl;
		return 1;
	}

	// run the simulation
	run(p);

	return 0;
}


// Copyright Yamaha 2021
// MIT License
// https://github.com/yamaha-bps/ros2_pcl_utils/blob/master/LICENSE

#include "pcl_utils/feature.hpp"

#include <Eigen/Core>

#include <memory>
#include <utility>
#include <vector>

#include "pcl_utils/pcl_iterator.hpp"

namespace cbr
{

struct ScanInfo
{
  uint32_t idx = 0;  // index in scan
  bool ok = false;   // if eligible to be a feature
  float cvalue = 0;  // estimate of smoothness
};

std::pair<sensor_msgs::msg::PointCloud2::UniquePtr, sensor_msgs::msg::PointCloud2::UniquePtr>
pcl_features(const sensor_msgs::msg::PointCloud2 & msg, const PclFeatureParams & prm)
{
  const uint32_t N = msg.height * msg.width;
  const uint32_t W = prm.window;

  if (N < 2 * W + 1) {
    return {nullptr, nullptr};  // too few points to process
  }

  //////////////////////////
  // STEP 1: VALID POINTS //
  //////////////////////////

  std::vector<bool> valid(N, false);

  std::vector<bool> ang_discont(N - 1, false);  // angle discontinuity
  std::vector<int8_t> dist_discont(N - 1, 0);   // distance discontinuity
  std::vector<bool> inc_ang_viol(N - 1, 0);     // incidence angle bound violation

  const auto cos_ang_disc_thresh = std::cos(prm.ang_disc_thresh);
  const auto sin_ang_inc_thresh = std::sin(prm.ang_inc_thresh);

  PclIterator it(msg);
  for (auto i = 0u; i != N - 1; ++i, ++it) {
    const auto p0 = *it;
    const auto p1 = *(it + 1);
    const auto r0 = p0.norm();
    const auto r1 = p1.norm();
    const auto r01 = (p1 - p0).norm();

    // incidence angle difference from 90 (perpendicular incidence)
    //  | alpha - pi/2 | <= b
    //  | cos(alpha) | <= sin(b)
    if (std::abs<float>(p0.dot(p1 - p0)) / (r0 * r01) > sin_ang_inc_thresh) {
      inc_ang_viol[i] = true;
    }

    // angular discontinuity
    // cos alpha = dot(p0, p1) / ||p1|| ||p2||
    // |alpha| > b  <==>  cos(alpha) < cos(b) <==>  dot(p0, p1) / (||p1|| ||p2||) < cos(b)
    if (p0.dot(p1) / (r0 * r1) < cos_ang_disc_thresh) {
      ang_discont[i] = true;
    }

    // distance discontinuity
    if (2 * r01 > prm.rel_disc_thresh * (r0 + r1)) {
      if (r0 < r1) {
        dist_discont[i] = -1;  // p1 further away than p0
      } else {
        dist_discont[i] = 1;   // p1 closer than p0
      }
    }
  }

  // sliding window count of number of angular discontinuities
  uint32_t num_ang_discont = 0;
  for (auto i = 0u; i < 2u * W; ++i) {
    num_ang_discont += uint32_t(ang_discont[i]);
  }

  // sliding window count of distance discontinuities
  uint32_t num_dist_discont_neg_left = 0;  // count decreasing on the left
  uint32_t num_dist_discont_pos_rght = 0;  // count increasing on the right
  for (auto i = 0u; i < W; ++i) {
    num_dist_discont_neg_left += uint32_t(dist_discont[i] == -1);
  }
  for (auto i = W; i < 2 * W; ++i) {
    num_dist_discont_pos_rght += uint32_t(dist_discont[i] == 1);
  }

  it = PclIterator(msg) + W;
  for (auto i = W; i + W < N; ++i, ++it) {
    valid[i] = true;

    if (num_dist_discont_neg_left > 0 || num_dist_discont_pos_rght > 0) {
      // distance discontinuity condition (robustly visible)
      valid[i] = false;
    } else if (num_ang_discont > 0) {
      // no angular discontinuities in window
      valid[i] = false;
    } else if (inc_ang_viol[i - 1] && inc_ang_viol[i]) {
      // incidence angle condition
      valid[i] = false;
    } else if (it.intensity() < prm.min_intensity || it.intensity() > prm.max_intensity) {
      // intensity condition
      valid[i] = false;
    } else if (std::atan2((*it).tail<2>().norm(), (*it).x()) > prm.max_angle) {
      // angle condition
      valid[i] = false;
    } else {
      // depth condition
      auto norm = (*it).norm();
      if (norm < prm.min_depth || norm > prm.max_depth) {
        valid[i] = false;
      }
    }

    // update sliding window counters
    num_ang_discont -= ang_discont[i - W];
    num_ang_discont += ang_discont[i + W];
    num_dist_discont_neg_left -= uint32_t(dist_discont[i - W] == -1);
    num_dist_discont_neg_left += uint32_t(dist_discont[i] == -1);
    num_dist_discont_pos_rght -= uint32_t(dist_discont[i] == 1);
    num_dist_discont_pos_rght += uint32_t(dist_discont[i + W] == 1);
  }

  //////////////////////
  // STEP 2: FEATURES //
  //////////////////////

  // cumulative vector sum in window
  Eigen::Vector3f csum(0, 0, 0);

  std::vector<ScanInfo> cands;

  // fast iterator: point one past right end of interval
  auto i_f = 0u;
  auto it_f = PclIterator(msg);
  for (; i_f != 2 * W + 1; ++i_f, ++it_f) {
    csum += *it_f;
  }

  // mid iterator: points to active point
  auto it_m = PclIterator(msg) + W;

  // slow iterator: points to left end of interval
  auto i_s = 0u;
  auto it_s = PclIterator(msg);

  // calculate c-values:
  //  c = || (2*W + 1) * p - \sum_i p_i || / ( 2 * W * ||p|| )
  //    = || 2 * W * p - \sum_{i != 0} p_i || / (2 * W * ||p||  )
  //    = || p  -  (1/2W) \sum_{i != 0} p_i || / || p ||

  // if c-value is small it is likely a planar point (p close to average of surrounding pts)
  // if c-value is large it is likely an edge point (p far from average of surrounding pts)
  for (; i_f != N; ++i_s, ++i_f, ++it_s, ++it_m, ++it_f) {
    uint32_t i_mid = i_s + W;

    if (valid[i_mid]) {
      cands.push_back(
        {i_mid, true, ((2 * W + 1) * (*it_m) - csum).norm() / (2 * W * (*it_m).norm())}
      );
    }

    csum -= *it_s;
    csum += *it_f;
  }

  // sort candidates by c-value
  std::sort(
    cands.begin(), cands.end(), [](const auto & item1, const auto & item2) {
      return item1.cvalue < item2.cvalue;
    });

  // extract features

  // edges from back
  auto edges = std::make_unique<sensor_msgs::msg::PointCloud2>();
  edges->fields = msg.fields;
  edges->point_step = msg.point_step;
  edges->is_bigendian = msg.is_bigendian;
  edges->header = msg.header;

  for (auto it = cands.crbegin(); it != cands.crend() && it->cvalue > prm.edge_thresh; ++it) {
    auto &[idx, ok, cvalue] = *it;
    if (valid[idx]) {
      std::copy(
        msg.data.begin() + idx * msg.point_step,
        msg.data.begin() + (idx + 1) * msg.point_step,
        std::back_insert_iterator(edges->data)
      );
      for (auto idx_i = idx - W; idx_i != idx + W + 1; ++idx_i) {
        valid[idx_i] = false;  // mark surrounding as taken
      }
    }
  }

  edges->width = edges->data.size() / edges->point_step;
  edges->row_step = edges->width;
  edges->height = 1;
  edges->is_dense = 1;

  // planar from front
  auto planar = std::make_unique<sensor_msgs::msg::PointCloud2>();
  planar->fields = msg.fields;
  planar->point_step = msg.point_step;
  planar->is_bigendian = msg.is_bigendian;
  planar->header = msg.header;

  for (auto it = cands.cbegin(); it != cands.cend() && it->cvalue < prm.plane_thresh; ++it) {
    auto &[idx, ok, cvalue] = *it;
    if (valid[idx]) {
      std::copy(
        msg.data.begin() + idx * msg.point_step,
        msg.data.begin() + (idx + 1) * msg.point_step,
        std::back_insert_iterator(planar->data)
      );
      for (auto idx_i = idx - W; idx_i != idx + W + 1; ++idx_i) {
        valid[idx_i] = false;  // mark surrounding as taken
      }
    }
  }

  planar->width = planar->data.size() / planar->point_step;
  planar->row_step = planar->width;
  planar->height = 1;
  planar->is_dense = 1;

  return {std::move(edges), std::move(planar)};
}

}  // namespace cbr


/*
 ============================================================================
 Name        : INSIGHTv3_mpi.cpp
 Author      : Jan Mikelson
 Version     :
 Copyright   : 	This free software is available under the Creative Commons Attribution Share Alike License.
 You are permitted to use, redistribute and adapt this software as long as appropriate credit
 is given to the original author, and all derivative works are distributed under the same
 license or a compatible one.
 For more information, visit http://creativecommons.org/licenses/by-sa/3.0/ or send a letter to
 Creative Commons, 171 2nd Street, Suite 300, San Francisco, California, 94105, USA.
 Description : 	Parallel version of the INSIGHT algorithm. Further details can be found in Lillacci & Khammash 2013
 ============================================================================
 */

#include <ctime>
#include <vector>

#include <stdio.h>  /* defines FILENAME_MAX */
#ifdef WINDOWS
#include <direct.h>
#define GetCurrentDir _getcwd
#else
#include <unistd.h>
#define GetCurrentDir getcwd
#endif

char cCurrentPath[FILENAME_MAX];

#include <boost/program_options.hpp>
#include <boost/smart_ptr/make_shared.hpp>
#include <boost/ref.hpp>
#include <Eigen/Dense>
#include <mpi.h>

#include "INSIGHTv3.h"
#include "BirthDeathModel.h"
#include "InsightAlgorithm.h"
#include "InsightParticleSamplerFactory.h"
#include "IllegalArgumentException.h"
#include "ParticleEvaluatorInsight.h"
#include "UniformSamplerFactory.h"
#include "ScaledUniformSamplerFactory.h"
#include "ProblemFileReader.h"
#include "ModelDescription.h"
#include "UniformSampler.h"
#include "InsightIterationMpi.h"
#include "InsightMpiWorker.h"
#include "PrevPopReader.h"

using namespace INSIGHTv3;

//Default values
const static std::string PROBLEM_FILE = "";
const static int S = 3;
const static double FINAL_TOLERANCE = 0.03;
const static int NUM_PARTICLES = 1000;
const static std::string OUTPUT_FILE = "";
const static std::string PREVIOUS_POPULATION_FILE = "";
const static bool IGNORE_WEIGHTS = false;
const static int PRINT = 20;
const static double BETA = 0.01;
const static double TOL_KOLMOGOROV = 0.0000001;
const static double KAPPA_KOLMOGOROV = 0.005224; // kappa value for M = 80000 and beta = 0.01
const static double TOL_KAPPA_KOLMOGOROV = 0.01;

std::string problem_file_name;
int s;
double final_tolerance;
int num_particles;
std::string output_file_name;
std::string prev_pop_file;
bool ignore_weights;
int print;
double beta;
double tolerance_kolmogorov;
double kappa_kolmogorov;
double tol_kappa_kolmovorog;

namespace po = boost::program_options;

void handleOptions(int argc, char * argv[]);
void printInfo(ModelDescription& desc);

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

	handleOptions(argc, argv);

	clock_t tic = clock();
	MPI::Init(argc, argv);
	int my_rank = MPI::COMM_WORLD.Get_rank();
	int num_tasks = MPI::COMM_WORLD.Get_size();

	if (num_tasks == 1) {
		std::cerr
				<< "When using mpi version of INSIGHT at least two processes must be started!"
				<< std::endl;
		abort();
	}

	int seed = time(NULL) + my_rank;
#ifdef DEBUG
	seed = 2;
#endif
	RngPtr rng = boost::make_shared<RandomNumberGenerator>(seed);

	ProblemFileReader reader;
	ModelDescription desc = reader.read(rng, problem_file_name);

	KolmogorovComputerPtr kolmogorov_computer = boost::make_shared<
			KolmogorovComputer>(*desc.original_data, beta, tolerance_kolmogorov,
			kappa_kolmogorov);
	evaluator_ptr evaluator = boost::make_shared<ParticleEvaluatorInsight>(
			boost::ref(*desc.model), *desc.times, boost::ref(rng),
			kolmogorov_computer); // need to use boost::ref, see boost
	sampler_ptr prior_sampler = boost::make_shared<UniformSampler>(
			boost::ref(rng), *desc.lower_bounds_parameter,
			*desc.upper_bounds_parameter);

	if (my_rank == 0) {
		printInfo(desc);
		writer_ptr writer = boost::make_shared<InsightOutputWriter>(
				output_file_name);

		sampler_factory_ptr factory = boost::make_shared<
				ScaledUniformSamplerFactory>(boost::ref(rng),
				desc.model->num_params, *desc.lower_bounds_parameter,
				*desc.upper_bounds_parameter, ignore_weights);

		InsightIterationMpi iteration(num_tasks, prior_sampler);
		ToleranceProviderPtr tolerance_provider = boost::make_shared<
				FixedSequenceToleranceProvider>(s, desc.model->model,
				kolmogorov_computer);

		sampler_ptr first_sampler;
		std::vector<double> prev_acceptance_rates;
		std::vector<double> prev_tolerances;
		std::vector<int> prev_num_simulations;
		if (prev_pop_file.length() > 0) {
			PrevPopReader prev_pop_reader;
			ParticleSet particle_set;
			prev_pop_reader.readPreviousPop(prev_pop_file,
					desc.model->num_params, &particle_set,
					&prev_acceptance_rates, &prev_tolerances,
					&prev_num_simulations);
			first_sampler = factory->createSampler(particle_set,
					kolmogorov_computer->getThresholdForS(s), s);
		} else {
			first_sampler = prior_sampler;
		}
		InsightAlgorithm algorithm(iteration, factory, evaluator, writer,
				tolerance_provider, final_tolerance, num_particles,
				desc.model->num_params, prev_acceptance_rates, prev_tolerances,
				prev_num_simulations);
		IterationLogger logger(print);

		algorithm.run(first_sampler, logger);

		MPI_INSTRUCTION instruction(DIETAG);
		for (int rank = 1; rank < num_tasks; rank++) {
			MPI::COMM_WORLD.Send(&instruction, 1, MPI::INT, rank, INSTRUCTION);
		}
	} else {
		InsightMpiWorker worker(my_rank, desc.model->num_params, evaluator,
				prior_sampler);
		worker.run();
	}

	clock_t toc = clock();
	if (my_rank == 0) {
		std::cout << "\n\nSMC inference complete." << std::endl;
		std::cout << "\nElapsed time: "
				<< ((double) (toc - tic)) / CLOCKS_PER_SEC << "sec.\n"
				<< std::endl;
	}
	MPI::Finalize();

	return 0;
}

void handleOptions(int argc, char * argv[]) {

	po::options_description desc("Allowed options");
	desc.add_options()("help", "produce help message")("problem_file",
			po::value<std::string>(&problem_file_name)->default_value(
					PROBLEM_FILE),
			"problem file. A problem file must always be provided!")(
			"number_simulations,s", po::value<int>(&s)->default_value(S),
			"starting number of simulations")("final tolerance,e",
			po::value<double>(&final_tolerance)->default_value(FINAL_TOLERANCE),
			"the final tolerance for the algorithm")("pop_size,n",
			po::value<int>(&num_particles)->default_value(NUM_PARTICLES),
			"define the population size")("output_file,O",
			po::value<std::string>(&output_file_name)->default_value(
					OUTPUT_FILE), "output file")("previous_pop_file,P",
			po::value<std::string>(&prev_pop_file)->default_value(
					PREVIOUS_POPULATION_FILE), "previous population file")(
			"ignore_weights,i",
			po::value<bool>(&ignore_weights)->default_value(IGNORE_WEIGHTS),
			"ignore weights")("print,p",
			po::value<int>(&print)->default_value(PRINT),
			"define the frequency of print outs")("beta,b",
			po::value<double>(&beta)->default_value(BETA),
			"the confidence of the kolmogorov test for the rejection rule")(
			"tol_kol,k",
			po::value<double>(&tolerance_kolmogorov)->default_value(
					TOL_KOLMOGOROV),
			"the tolerance for the algorithm to compute the inverse of the kolmogorov distribution")(
			"kappa,K",
			po::value<double>(&kappa_kolmogorov)->default_value(
					KAPPA_KOLMOGOROV),
			"Kappa for the computation of the Kolmogorov distance")(
			"tol_kol_kap,T",
			po::value<double>(&tol_kappa_kolmovorog)->default_value(
					TOL_KAPPA_KOLMOGOROV),
			"the tolerance for the algorithm to compute kappa for the computation of the Kolmogorov distance");

	po::positional_options_description p;
	p.add("problem_file", -1);

	po::variables_map vm;
	po::store(po::parse_command_line(argc, argv, desc), vm);
	po::store(
			po::command_line_parser(argc, argv).options(desc).positional(p).run(),
			vm);
	po::notify(vm);

	if (vm.count("help")) {
		std::cout << desc << "\n";
		abort();
	}

	if (problem_file_name.size() == 0) {
		std::cerr << "Problem file is required" << std::endl;
		abort();
	}

	if (output_file_name.size() == 0) {
		GetCurrentDir(cCurrentPath, sizeof(cCurrentPath));
		output_file_name = std::string(cCurrentPath) + "/resutls.txt";
	}

}

void printInfo(ModelDescription& desc) {

	std::cout << "\nThis is INSIGHT v3.0\n" << std::endl;
	std::cout << "Problem file path:" << std::endl;
	std::cout << problem_file_name << std::endl;
	std::cout << "\nThe new particle population will be saved in:" << std::endl;
	std::cout << output_file_name << std::endl;

	std::cout << "\nWill estimate " << desc.model->model_name << " using "
			<< num_particles << " SMC particles." << std::endl;
	std::cout << "\nThe data is assumed to have been created using "
			<< desc.model->num_outputs << " outputs," << std::endl;
	std::cout << "measured at " << desc.times->size() << " timepoints."
			<< std::endl;

	if (prev_pop_file.length() > 0) {
		std::cout << "A previous population is loaded from file: " << std::endl;
		std::cout << prev_pop_file << std::endl;
	}
	std::cout << "To estimate the density of each population "
			<< "Kernel Density Estimation with uniform Kernel" << std::endl;
	std::cout << "is used." << std::endl;
	if (ignore_weights) {
		std::cout << std::endl << "Particle weights will be ignored by sampler!"
				<< std::endl;
	}
	std::cout << "\nThe parallel version of the algorithm is used."
			<< std::endl;
}


#include "Poisson.h"
#include "config.h"
#include "kernels/poisson/init.h"
#include "kernels/poisson/kernel.h"
#include "kernels/poisson/tensor.h"

#include "basis/WarpAndBlend.h"
#include "form/BC.h"
#include "form/DGCurvilinearCommon.h"
#include "form/InverseInequality.h"
#include "form/RefElement.h"
#include "geometry/Curvilinear.h"
#include "quadrules/SimplexQuadratureRule.h"
#include "tensor/EigenMap.h"
#include "util/LinearAllocator.h"

#include <Eigen/Core>
#include <Eigen/LU>
#include <cassert>

namespace tensor = tndm::poisson::tensor;
namespace init = tndm::poisson::init;
namespace kernel = tndm::poisson::kernel;

namespace tndm {

Poisson::Poisson(std::shared_ptr<Curvilinear<DomainDimension>> cl, functional_t<1> K,
                 DGMethod method)
    : DGCurvilinearCommon<DomainDimension>(std::move(cl), MinQuadOrder()), method_(method),
      space_(PolynomialDegree, ALIGNMENT),
      materialSpace_(PolynomialDegree, WarpAndBlendFactory<DomainDimension>(), ALIGNMENT),
      fun_K(make_volume_functional(std::move(K))), fun_force(zero_volume_function),
      fun_dirichlet(zero_facet_function), fun_slip(zero_facet_function) {

    Minv_ = space_.inverseMassMatrix();
    E_Q = space_.evaluateBasisAt(volRule.points());
    E_Q_T = space_.evaluateBasisAt(volRule.points(), {1, 0});
    Dxi_Q = space_.evaluateGradientAt(volRule.points());

    negative_E_Q_T = Managed<Matrix<double>>(E_Q_T.shape(), std::size_t{ALIGNMENT});
    EigenMap(negative_E_Q_T) = -EigenMap(E_Q_T);

    for (std::size_t f = 0; f < DomainDimension + 1u; ++f) {
        auto points = cl_->facetParam(f, fctRule.points());
        E_q.emplace_back(space_.evaluateBasisAt(points));
        E_q_T.emplace_back(space_.evaluateBasisAt(points, {1, 0}));
        Dxi_q.emplace_back(space_.evaluateGradientAt(points));
        Dxi_q_120.emplace_back(space_.evaluateGradientAt(points, {1, 2, 0}));
        matE_q_T.emplace_back(materialSpace_.evaluateBasisAt(points, {1, 0}));

        negative_E_q_T.emplace_back(space_.evaluateBasisAt(points, {1, 0}));
        auto E = EigenMap(negative_E_q_T.back());
        E = -E;
    }

    matE_Q_T = materialSpace_.evaluateBasisAt(volRule.points(), {1, 0});
    matDxi_Q = materialSpace_.evaluateGradientAt(volRule.points());
}

void Poisson::compute_mass_matrix(std::size_t elNo, double* M) const {
    kernel::massMatrix mm;
    mm.E_Q = E_Q.data();
    mm.J_Q = vol[elNo].get<AbsDetJ>().data();
    mm.M = M;
    mm.W = volRule.weights().data();
    mm.execute();
}

void Poisson::compute_inverse_mass_matrix(std::size_t elNo, double* Minv) const {
    compute_mass_matrix(elNo, Minv);

    auto J_Q = vol[elNo].get<AbsDetJ>();
    alignas(ALIGNMENT) double Jinv_Q[tensor::Jinv_Q::size()] = {};
    for (unsigned q = 0; q < tensor::Jinv_Q::Shape[0]; ++q) {
        Jinv_Q[q] = 1.0 / J_Q[q];
    }

    kernel::MinvWA wa;
    wa.E_Q = E_Q.data();
    wa.Jinv_Q = Jinv_Q;
    wa.MinvRef = Minv_.data();
    wa.MinvWA = Minv;
    wa.W = volRule.weights().data();
    wa.execute();
}

void Poisson::compute_K_Dx_q(std::size_t fctNo, FacetInfo const& info,
                             std::array<double*, 2> K_Dx_q) const {
    kernel::K_Dx_q dx;
    for (int i = 0; i < 2; ++i) {
        if (K_Dx_q[i]) {
            auto JInv = (i == 1) ? fct[fctNo].get<JInv1>() : fct[fctNo].get<JInv0>();
            dx.G_q = JInv.data()->data();
            dx.matE_q_T = matE_q_T[info.localNo[i]].data();
            dx.K = material[info.up[i]].get<K>().data();
            dx.K_Dx_q(0) = K_Dx_q[i];
            dx.Dxi_q(0) = Dxi_q[info.localNo[i]].data();
            dx.execute();
        }
    }
}

void Poisson::compute_K_q(std::size_t fctNo, FacetInfo const& info,
                          std::array<double*, 2> K_q) const {
    kernel::K_q kw;
    for (int i = 0; i < 2; ++i) {
        if (K_q[i]) {
            kw.matE_q_T = matE_q_T[info.localNo[i]].data();
            kw.K = material[info.up[i]].get<K>().data();
            kw.K_q(0) = K_q[i];
            kw.execute();
        }
    }
}

void Poisson::begin_preparation(std::size_t numElements, std::size_t numLocalElements,
                                std::size_t numLocalFacets) {
    base::begin_preparation(numElements, numLocalElements, numLocalFacets);

    material.setStorage(
        std::make_shared<material_vol_t>(numElements * materialSpace_.numBasisFunctions()), 0u,
        numElements, materialSpace_.numBasisFunctions());

    volPre.setStorage(std::make_shared<vol_pre_t>(numElements * volRule.size()), 0u,
                      numLocalElements, volRule.size());

    fctPre.setStorage(std::make_shared<fct_pre_t>(numLocalFacets * fctRule.size()), 0u,
                      numLocalFacets, fctRule.size());

    penalty_.resize(numLocalFacets);
}

void Poisson::prepare_volume(std::size_t elNo, LinearAllocator<double>& scratch) {
    base::prepare_volume(elNo, scratch);

    auto Kfield = material[elNo].get<K>().data();
    alignas(ALIGNMENT) double K_Q_raw[tensor::K_Q::size()];
    auto K_Q = Matrix<double>(K_Q_raw, 1, volRule.size());
    fun_K(elNo, K_Q);

    alignas(ALIGNMENT) double Mmem[tensor::matM::size()];
    kernel::project_K_lhs krnl_lhs;
    krnl_lhs.matE_Q_T = matE_Q_T.data();
    krnl_lhs.J_Q = vol[elNo].get<AbsDetJ>().data();
    krnl_lhs.matM = Mmem;
    krnl_lhs.W = volRule.weights().data();
    krnl_lhs.execute();

    kernel::project_K_rhs krnl_rhs;
    krnl_rhs.matE_Q_T = matE_Q_T.data();
    krnl_rhs.J_Q = vol[elNo].get<AbsDetJ>().data();
    krnl_rhs.K = Kfield;
    krnl_rhs.K_Q = K_Q_raw;
    krnl_rhs.W = volRule.weights().data();
    krnl_rhs.execute();

    using MMap = Eigen::Map<Eigen::Matrix<double, tensor::matM::Shape[0], tensor::matM::Shape[1]>,
                            Eigen::Unaligned,
                            Eigen::OuterStride<init::matM::Stop[0] - init::matM::Start[0]>>;
    using KMap = Eigen::Map<Eigen::Matrix<double, tensor::K::Shape[0], 1>, Eigen::Unaligned,
                            Eigen::InnerStride<1>>;

    auto K_eigen = KMap(Kfield);
    K_eigen = MMap(Mmem).fullPivLu().solve(K_eigen);
}

void Poisson::prepare_skeleton(std::size_t fctNo, FacetInfo const& info,
                               LinearAllocator<double>& scratch) {
    base::prepare_skeleton(fctNo, info, scratch);

    for (int side = 0; side < 2; ++side) {
        kernel::K_G_q k;
        k.G_q = side == 1 ? fct[fctNo].get<JInv1>().data()->data()
                          : fct[fctNo].get<JInv0>().data()->data();
        k.K = material[info.up[side]].get<K>().data();
        k.K_G_q(0) = side == 1 ? fctPre[fctNo].get<KJInv1>().data()->data()
                               : fctPre[fctNo].get<KJInv0>().data()->data();
        k.matE_q_T = matE_q_T[info.localNo[side]].data();
        k.execute();
    }
}

void Poisson::prepare_boundary(std::size_t fctNo, FacetInfo const& info,
                               LinearAllocator<double>& scratch) {
    base::prepare_boundary(fctNo, info, scratch);

    kernel::K_G_q k;
    k.G_q = fct[fctNo].get<JInv0>().data()->data();
    k.K = material[info.up[0]].get<K>().data();
    k.K_G_q(0) = fctPre[fctNo].get<KJInv0>().data()->data();
    k.matE_q_T = matE_q_T[info.localNo[0]].data();
    k.execute();
}

void Poisson::prepare_volume_post_skeleton(std::size_t elNo, LinearAllocator<double>& scratch) {
    base::prepare_volume_post_skeleton(elNo, scratch);

    auto Kfield = material[elNo].get<K>().data();

    kernel::J_W_K_Q krnl;
    krnl.J_W_K_Q = volPre[elNo].get<AbsDetJWK>().data()->data();
    krnl.J_Q = vol[elNo].get<AbsDetJ>().data();
    krnl.K = Kfield;
    krnl.matE_Q_T = matE_Q_T.data();
    krnl.W = volRule.weights().data();
    krnl.execute();
}

void Poisson::prepare_penalty(std::size_t fctNo, FacetInfo const& info, LinearAllocator<double>&) {
    auto const p = [&](int side) {
        auto Kfield = material[info.up[side]].get<K>().data();
        auto k0 = *std::min_element(Kfield, Kfield + materialSpace_.numBasisFunctions());
        auto k1 = *std::max_element(Kfield, Kfield + materialSpace_.numBasisFunctions());
        constexpr double c_N_1 = InverseInequality<Dim>::trace_constant(PolynomialDegree - 1);
        return (Dim + 1) * c_N_1 * (area_[fctNo] / volume_[info.up[side]]) * (k1 * k1 / k0);
    };

    if (info.up[0] != info.up[1]) {
        penalty_[fctNo] = (p(0) + p(1)) / 4.0;
    } else {
        penalty_[fctNo] = p(0);
    }
}

bool Poisson::assemble_volume(std::size_t elNo, Matrix<double>& A00,
                              LinearAllocator<double>& scratch) const {
    alignas(ALIGNMENT) double Dx_Q[tensor::Dx_Q::size()];

    assert(volRule.size() == tensor::W::Shape[0]);
    assert(Dxi_Q.shape(0) == tensor::Dxi_Q::Shape[0]);
    assert(Dxi_Q.shape(1) == tensor::Dxi_Q::Shape[1]);
    assert(Dxi_Q.shape(2) == tensor::Dxi_Q::Shape[2]);

    kernel::Dx_Q dx;
    dx.Dx_Q = Dx_Q;
    dx.Dxi_Q = Dxi_Q.data();
    dx.G_Q = vol[elNo].get<JInv>().data()->data();
    dx.execute();

    kernel::assembleVolume krnl;
    krnl.A = A00.data();
    krnl.Dx_Q = Dx_Q;
    krnl.K = material[elNo].get<K>().data();
    krnl.matE_Q_T = matE_Q_T.data();
    krnl.J_Q = vol[elNo].get<AbsDetJ>().data();
    krnl.W = volRule.weights().data();
    krnl.execute();
    return true;
}

bool Poisson::assemble_skeleton(std::size_t fctNo, FacetInfo const& info, Matrix<double>& A00,
                                Matrix<double>& A01, Matrix<double>& A10, Matrix<double>& A11,
                                LinearAllocator<double>& scratch) const {
    assert(fctRule.size() == tensor::w::Shape[0]);
    assert(E_q[0].shape(0) == tensor::E_q::Shape[0][0]);
    assert(E_q[0].shape(1) == tensor::E_q::Shape[0][1]);
    assert(Dxi_q[0].shape(0) == tensor::Dxi_q::Shape[0][0]);
    assert(Dxi_q[0].shape(1) == tensor::Dxi_q::Shape[0][1]);
    assert(Dxi_q[0].shape(2) == tensor::Dxi_q::Shape[0][2]);

    alignas(ALIGNMENT) double K_Dx_q0[tensor::K_Dx_q::size(0)];
    alignas(ALIGNMENT) double K_Dx_q1[tensor::K_Dx_q::size(1)];
    auto K_Dx_q = std::array<double*, 2>{K_Dx_q0, K_Dx_q1};
    compute_K_Dx_q(fctNo, info, K_Dx_q);

    alignas(ALIGNMENT) double L_q[2][std::max(tensor::L_q::size(0), tensor::L_q::size(1))];

    if (method_ == DGMethod::BR2) {
        alignas(ALIGNMENT) double Lift0[tensor::Lift::size(0)];
        alignas(ALIGNMENT) double Lift1[tensor::Lift::size(1)];
        alignas(ALIGNMENT) double Minv[2][tensor::M::size()];
        for (int i = 0; i < 2; ++i) {
            compute_inverse_mass_matrix(info.up[i], Minv[i]);
        }

        alignas(ALIGNMENT) double K_q0[tensor::K_q::size(0)];
        alignas(ALIGNMENT) double K_q1[tensor::K_q::size(1)];
        auto K_q = std::array<double*, 2>{K_q0, K_q1};
        compute_K_q(fctNo, info, K_q);

        kernel::lift_skeleton lift;
        lift.Lift(0) = Lift0;
        lift.Lift(1) = Lift1;
        lift.n_q = fct[fctNo].get<Normal>().data()->data();
        lift.w = fctRule.weights().data();
        for (int i = 0; i < 2; ++i) {
            lift.K_q(i) = K_q[i];
            lift.L_q(i) = L_q[i];
            lift.Minv(i) = Minv[i];
            lift.E_q(i) = E_q[info.localNo[i]].data();
        }
        lift.execute(0);
        lift.execute(1);
    } else { // IP
        kernel::lift_ip lift;
        lift.nl_q = fct[fctNo].get<NormalLength>().data();
        for (int i = 0; i < 2; ++i) {
            lift.L_q(i) = L_q[i];
            lift.E_q(i) = E_q[info.localNo[i]].data();
        }
        lift.execute(0);
        lift.execute(1);
    }

    kernel::assembleSurface assemble;
    assemble.c00 = -0.5;
    assemble.c01 = -assemble.c00;
    assemble.c10 = epsilon * 0.5;
    assemble.c11 = -assemble.c10;
    assemble.c20 = penalty(fctNo);
    assemble.c21 = -assemble.c20;
    assemble.a(0, 0) = A00.data();
    assemble.a(0, 1) = A01.data();
    assemble.a(1, 0) = A10.data();
    assemble.a(1, 1) = A11.data();
    for (int i = 0; i < 2; ++i) {
        assemble.K_Dx_q(i) = K_Dx_q[i];
        assemble.E_q(i) = E_q[info.localNo[i]].data();
        assemble.L_q(i) = L_q[i];
    }
    assemble.n_q = fct[fctNo].get<Normal>().data()->data();
    assemble.w = fctRule.weights().data();
    assemble.execute(0, 0);
    assemble.execute(0, 1);
    assemble.execute(1, 0);
    assemble.execute(1, 1);

    return true;
}

bool Poisson::assemble_boundary(std::size_t fctNo, FacetInfo const& info, Matrix<double>& A00,
                                LinearAllocator<double>& scratch) const {
    if (info.bc == BC::Natural) {
        return false;
    }

    assert(fctRule.size() == tensor::w::Shape[0]);
    assert(E_q[0].shape(0) == tensor::E_q::Shape[0][0]);
    assert(E_q[0].shape(1) == tensor::E_q::Shape[0][1]);
    assert(Dxi_q[0].shape(0) == tensor::Dxi_q::Shape[0][0]);
    assert(Dxi_q[0].shape(1) == tensor::Dxi_q::Shape[0][1]);
    assert(Dxi_q[0].shape(2) == tensor::Dxi_q::Shape[0][2]);

    alignas(ALIGNMENT) double L0[tensor::L_q::size(0)];
    if (method_ == DGMethod::BR2) {
        alignas(ALIGNMENT) double Lift0[tensor::Lift::size(0)];
        alignas(ALIGNMENT) double Minv0[tensor::M::size()];
        compute_inverse_mass_matrix(info.up[0], Minv0);

        alignas(ALIGNMENT) double K_q[tensor::K_q::size(0)];
        compute_K_q(fctNo, info, {K_q, nullptr});

        kernel::lift_boundary lift;
        lift.Lift(0) = Lift0;
        lift.K_q(0) = K_q;
        lift.L_q(0) = L0;
        lift.Minv(0) = Minv0;
        lift.E_q(0) = E_q[info.localNo[0]].data();
        lift.n_q = fct[fctNo].get<Normal>().data()->data();
        lift.w = fctRule.weights().data();
        lift.execute();
    } else { // IP
        kernel::lift_ip lift;
        lift.nl_q = fct[fctNo].get<NormalLength>().data();
        lift.L_q(0) = L0;
        lift.E_q(0) = E_q[info.localNo[0]].data();
        lift.execute(0);
    }

    alignas(ALIGNMENT) double K_Dx_q0[tensor::K_Dx_q::size(0)];
    compute_K_Dx_q(fctNo, info, {K_Dx_q0, nullptr});

    kernel::assembleSurface assemble;
    assemble.c00 = -1.0;
    assemble.c10 = epsilon;
    assemble.c20 = penalty(fctNo);
    assemble.a(0, 0) = A00.data();
    assemble.K_Dx_q(0) = K_Dx_q0;
    assemble.E_q(0) = E_q[info.localNo[0]].data();
    assemble.L_q(0) = L0;
    assemble.n_q = fct[fctNo].get<Normal>().data()->data();
    assemble.w = fctRule.weights().data();
    assemble.execute(0, 0);
    return true;
}

bool Poisson::rhs_volume(std::size_t elNo, Vector<double>& B,
                         LinearAllocator<double>& scratch) const {
    assert(tensor::b::Shape[0] == tensor::A::Shape[0]);

    alignas(ALIGNMENT) double F_Q_raw[tensor::F_Q::size()];
    assert(tensor::F_Q::size() == volRule.size());
    auto F_Q = Matrix<double>(F_Q_raw, 1, tensor::F_Q::Shape[0]);
    fun_force(elNo, F_Q);

    kernel::rhsVolume rhs;
    rhs.E_Q = E_Q.data();
    rhs.F_Q = F_Q_raw;
    rhs.J_Q = vol[elNo].get<AbsDetJ>().data();
    rhs.W = volRule.weights().data();
    rhs.b = B.data();
    rhs.execute();
    return true;
}

bool Poisson::bc_skeleton(std::size_t fctNo, BC bc, double f_q_raw[]) const {
    assert(tensor::f_q::size() == fctRule.size());
    auto f_q = Matrix<double>(f_q_raw, 1, tensor::f_q::Shape[0]);
    if (bc == BC::Fault) {
        fun_slip(fctNo, f_q, false);
    } else if (bc == BC::Dirichlet) {
        fun_dirichlet(fctNo, f_q, false);
    } else {
        return false;
    }
    return true;
}
bool Poisson::bc_boundary(std::size_t fctNo, BC bc, double f_q_raw[]) const {
    assert(tensor::f_q::size() == fctRule.size());
    auto f_q = Matrix<double>(f_q_raw, 1, tensor::f_q::Shape[0]);
    if (bc == BC::Fault) {
        fun_slip(fctNo, f_q, true);
        for (std::size_t q = 0; q < tensor::f_q::Shape[0]; ++q) {
            f_q(0, q) *= 0.5;
        }
    } else if (bc == BC::Dirichlet) {
        fun_dirichlet(fctNo, f_q, true);
    } else {
        return false;
    }
    return true;
}

bool Poisson::rhs_skeleton(std::size_t fctNo, FacetInfo const& info, Vector<double>& B0,
                           Vector<double>& B1, LinearAllocator<double>& scratch) const {
    alignas(ALIGNMENT) double f_q_raw[tensor::f_q::size()];
    if (!bc_skeleton(fctNo, info.bc, f_q_raw)) {
        return false;
    }

    alignas(ALIGNMENT) double f_lifted_q[tensor::f_lifted_q::size()];
    if (method_ == DGMethod::BR2) {
        alignas(ALIGNMENT) double f_lifted0[tensor::f_lifted::size(0)];
        alignas(ALIGNMENT) double f_lifted1[tensor::f_lifted::size(1)];
        alignas(ALIGNMENT) double Minv[2][tensor::M::size()];
        compute_inverse_mass_matrix(info.up[0], Minv[0]);
        compute_inverse_mass_matrix(info.up[1], Minv[1]);

        alignas(ALIGNMENT) double K_q0[tensor::K_q::size(0)];
        alignas(ALIGNMENT) double K_q1[tensor::K_q::size(1)];
        auto K_q = std::array<double*, 2>{K_q0, K_q1};
        compute_K_q(fctNo, info, K_q);

        kernel::rhs_lift_skeleton lift;
        for (int i = 0; i < 2; ++i) {
            lift.E_q(i) = E_q[info.localNo[i]].data();
            lift.K_q(i) = K_q[i];
            lift.Minv(i) = Minv[i];
        }
        lift.n_q = fct[fctNo].get<Normal>().data()->data();
        lift.f_q = f_q_raw;
        lift.f_lifted(0) = f_lifted0;
        lift.f_lifted(1) = f_lifted1;
        lift.f_lifted_q = f_lifted_q;
        lift.w = fctRule.weights().data();
        lift.execute();
    } else { // IP
        kernel::rhs_lift_ip lift;
        lift.nl_q = fct[fctNo].get<NormalLength>().data();
        lift.f_q = f_q_raw;
        lift.f_lifted_q = f_lifted_q;
        lift.execute();
    }

    alignas(ALIGNMENT) double K_Dx_q0[tensor::K_Dx_q::size(0)];
    alignas(ALIGNMENT) double K_Dx_q1[tensor::K_Dx_q::size(1)];
    compute_K_Dx_q(fctNo, info, {K_Dx_q0, K_Dx_q1});

    kernel::rhsFacet rhs;
    rhs.b = B0.data();
    rhs.c10 = 0.5 * epsilon;
    rhs.c20 = penalty(fctNo);
    rhs.f_q = f_q_raw;
    rhs.f_lifted_q = f_lifted_q;
    rhs.n_q = fct[fctNo].get<Normal>().data()->data();
    rhs.w = fctRule.weights().data();
    rhs.K_Dx_q(0) = K_Dx_q0;
    rhs.E_q(0) = E_q[info.localNo[0]].data();
    rhs.execute();

    rhs.b = B1.data();
    rhs.c20 *= -1.0;
    rhs.K_Dx_q(0) = K_Dx_q1;
    rhs.E_q(0) = E_q[info.localNo[1]].data();
    rhs.execute();

    return true;
}

bool Poisson::rhs_boundary(std::size_t fctNo, FacetInfo const& info, Vector<double>& B0,
                           LinearAllocator<double>& scratch) const {
    alignas(ALIGNMENT) double f_q_raw[tensor::f_q::size()];
    if (!bc_boundary(fctNo, info.bc, f_q_raw)) {
        return false;
    }

    alignas(ALIGNMENT) double f_lifted_q[tensor::f_lifted_q::size()];
    if (method_ == DGMethod::BR2) {
        alignas(ALIGNMENT) double f_lifted0[tensor::f_lifted::size(0)];
        alignas(ALIGNMENT) double M0[tensor::M::size()];
        compute_inverse_mass_matrix(info.up[0], M0);

        alignas(ALIGNMENT) double K_q[tensor::K_q::size(0)];
        compute_K_q(fctNo, info, {K_q, nullptr});

        kernel::rhs_lift_boundary lift;
        lift.E_q(0) = E_q[info.localNo[0]].data();
        lift.n_q = fct[fctNo].get<Normal>().data()->data();
        lift.K_q(0) = K_q;
        lift.Minv(0) = M0;
        lift.f_q = f_q_raw;
        lift.f_lifted(0) = f_lifted0;
        lift.f_lifted_q = f_lifted_q;
        lift.w = fctRule.weights().data();
        lift.execute();
    } else { // IP
        kernel::rhs_lift_ip lift;
        lift.nl_q = fct[fctNo].get<NormalLength>().data();
        lift.f_q = f_q_raw;
        lift.f_lifted_q = f_lifted_q;
        lift.execute();
    }

    alignas(ALIGNMENT) double K_Dx_q0[tensor::K_Dx_q::size(0)];
    compute_K_Dx_q(fctNo, info, {K_Dx_q0, nullptr});

    kernel::rhsFacet rhs;
    rhs.b = B0.data();
    rhs.c10 = epsilon;
    rhs.c20 = penalty(fctNo);
    rhs.f_q = f_q_raw;
    rhs.f_lifted_q = f_lifted_q;
    rhs.n_q = fct[fctNo].get<Normal>().data()->data();
    rhs.w = fctRule.weights().data();
    rhs.K_Dx_q(0) = K_Dx_q0;
    rhs.E_q(0) = E_q[info.localNo[0]].data();
    rhs.execute();
    return true;
}

void Poisson::apply(std::size_t elNo, mneme::span<SideInfo> info, Vector<double const> const& x_0,
                    std::array<Vector<double const>, NumFacets> const& x_n,
                    Vector<double>& y_0) const {

    alignas(ALIGNMENT) double Dx_Q[tensor::Dx_Q::size()];
    kernel::apply_volume av;
    av.Dx_Q = Dx_Q;
    av.Dxi_Q = Dxi_Q.data();
    av.G_Q = vol[elNo].get<JInv>().data()->data();
    av.J_W_K_Q = volPre[elNo].get<AbsDetJWK>().data()->data();
    av.U = x_0.data();
    av.U_new = y_0.data();
    av.execute();

    alignas(ALIGNMENT) double n_q_flipped[tensor::n_q::size()];
    alignas(ALIGNMENT) double n_unit_q_flipped[tensor::n_unit_q::size()];
    for (std::size_t f = 0; f < NumFacets; ++f) {
        bool is_skeleton_face = elNo != info[f].lid;
        bool is_fault_or_dirichlet = info[f].bc == BC::Fault || info[f].bc == BC::Dirichlet;

        auto fctNo = info[f].fctNo;
        double const* n_q = fct[fctNo].get<Normal>().data()->data();
        double const* n_unit_q = fct[fctNo].get<UnitNormal>().data()->data();
        double const* K_G_q0 = fctPre[fctNo].get<KJInv0>().data()->data();
        double const* K_G_q1 = fctPre[fctNo].get<KJInv1>().data()->data();
        if (is_skeleton_face && info[f].side == 1) {
            std::swap(K_G_q0, K_G_q1);

            for (int i = 0; i < tensor::n_q::size(); ++i) {
                n_q_flipped[i] = -n_q[i];
            }
            n_q = n_q_flipped;
            for (int i = 0; i < tensor::n_unit_q::size(); ++i) {
                n_unit_q_flipped[i] = -n_unit_q[i];
            }
            n_unit_q = n_unit_q_flipped;
        }

        alignas(ALIGNMENT) double u_hat_q[tensor::u_hat_q::size()] = {};
        alignas(ALIGNMENT) double sigma_hat_q[tensor::sigma_hat_q::size()] = {};
        if (info[f].bc == BC::None || (is_skeleton_face && is_fault_or_dirichlet)) {
            kernel::flux_u_skeleton fu;
            fu.negative_E_q_T(0) = negative_E_q_T[f].data();
            fu.E_q_T(1) = E_q_T[info[f].localNo].data();
            fu.U = x_0.data();
            fu.U_ext = x_n[f].data();
            fu.u_hat_q = u_hat_q;
            fu.execute();

            kernel::flux_sigma_skeleton fs;
            fs.c00 = -penalty(fctNo);
            fs.Dxi_q_120(0) = Dxi_q_120[f].data();
            fs.Dxi_q_120(1) = Dxi_q_120[info[f].localNo].data();
            fs.E_q_T(0) = E_q_T[f].data();
            fs.negative_E_q_T(1) = negative_E_q_T[info[f].localNo].data();
            fs.K_G_q(0) = K_G_q0;
            fs.K_G_q(1) = K_G_q1;
            fs.U = x_0.data();
            fs.U_ext = x_n[f].data();
            fs.n_unit_q = n_unit_q;
            fs.sigma_hat_q = sigma_hat_q;
            fs.execute();
        } else if (is_fault_or_dirichlet) {
            kernel::flux_u_boundary fu;
            fu.U = x_0.data();
            fu.u_hat_q = u_hat_q;
            fu.negative_E_q_T(0) = negative_E_q_T[f].data();
            fu.execute();

            kernel::flux_sigma_boundary fs;
            fs.c00 = -penalty(fctNo);
            fs.Dxi_q_120(0) = Dxi_q_120[f].data();
            fs.E_q_T(0) = E_q_T[f].data();
            fs.K_G_q(0) = K_G_q0;
            fs.U = x_0.data();
            fs.n_unit_q = n_unit_q;
            fs.sigma_hat_q = sigma_hat_q;
            fs.execute();
        } else {
            continue;
        }

        kernel::apply_facet af;
        af.Dxi_q(0) = Dxi_q[f].data();
        af.E_q(0) = E_q[f].data();
        af.K_G_q(0) = K_G_q0;
        af.n_q = n_q;
        af.sigma_hat_q = sigma_hat_q;
        af.u_hat_q = u_hat_q;
        af.U_new = y_0.data();
        af.w = fctRule.weights().data();
        af.execute();
    }
}

std::size_t Poisson::flops_apply(std::size_t elNo, mneme::span<SideInfo> info) const {
    std::size_t flops = kernel::apply_volume::HardwareFlops;
    for (std::size_t f = 0; f < NumFacets; ++f) {
        bool is_skeleton_face = elNo != info[f].lid;
        bool is_fault_or_dirichlet = info[f].bc == BC::Fault || info[f].bc == BC::Dirichlet;
        if (info[f].bc == BC::None || (is_skeleton_face && is_fault_or_dirichlet)) {
            flops += kernel::flux_u_skeleton::HardwareFlops;
            flops += kernel::flux_sigma_skeleton::HardwareFlops;
        } else if (is_fault_or_dirichlet) {
            flops += kernel::flux_u_boundary::HardwareFlops;
            flops += kernel::flux_sigma_boundary::HardwareFlops;
        } else {
            continue;
        }
        flops += kernel::apply_facet::HardwareFlops;
    }
    return flops;
}

void Poisson::coefficients_volume(std::size_t elNo, Matrix<double>& C,
                                  LinearAllocator<double>&) const {
    auto const coeff_K = material[elNo].get<K>();
    assert(coeff_K.size() == C.shape(0));
    for (std::size_t i = 0; i < coeff_K.size(); ++i) {
        C(i, 0) = coeff_K[i];
    }
}

TensorBase<Matrix<double>> Poisson::tractionResultInfo() const {
    return TensorBase<Matrix<double>>(tensor::grad_u::Shape[0], tensor::grad_u::Shape[1]);
}

void Poisson::traction_skeleton(std::size_t fctNo, FacetInfo const& info, Vector<double const>& u0,
                                Vector<double const>& u1, Matrix<double>& result) const {
    assert(result.size() == tensor::grad_u::size());

    alignas(ALIGNMENT) double f_q_raw[tensor::f_q::size()];
    bc_skeleton(fctNo, info.bc, f_q_raw);

    alignas(ALIGNMENT) double K_Dx_q0[tensor::K_Dx_q::size(0)];
    alignas(ALIGNMENT) double K_Dx_q1[tensor::K_Dx_q::size(1)];
    compute_K_Dx_q(fctNo, info, {K_Dx_q0, K_Dx_q1});

    kernel::grad_u krnl;
    krnl.c00 = -penalty(fctNo);
    krnl.K_Dx_q(0) = K_Dx_q0;
    krnl.K_Dx_q(1) = K_Dx_q1;
    krnl.E_q(0) = E_q[info.localNo[0]].data();
    krnl.E_q(1) = E_q[info.localNo[1]].data();
    krnl.f_q = f_q_raw;
    krnl.grad_u = result.data();
    krnl.n_unit_q = fct[fctNo].get<UnitNormal>().data()->data();
    krnl.u(0) = u0.data();
    krnl.u(1) = u1.data();
    krnl.execute();
}

void Poisson::traction_boundary(std::size_t fctNo, FacetInfo const& info, Vector<double const>& u0,
                                Matrix<double>& result) const {
    assert(result.size() == tensor::grad_u::size());

    alignas(ALIGNMENT) double f_q_raw[tensor::f_q::size()];
    bc_boundary(fctNo, info.bc, f_q_raw);

    alignas(ALIGNMENT) double K_Dx_q0[tensor::K_Dx_q::size(0)];
    compute_K_Dx_q(fctNo, info, {K_Dx_q0, nullptr});

    kernel::grad_u_bnd krnl;
    krnl.c00 = -penalty(fctNo);
    krnl.K_Dx_q(0) = K_Dx_q0;
    krnl.E_q(0) = E_q[info.localNo[0]].data();
    krnl.f_q = f_q_raw;
    krnl.grad_u = result.data();
    krnl.n_unit_q = fct[fctNo].get<UnitNormal>().data()->data();
    krnl.u(0) = u0.data();
    krnl.execute();
}

} // namespace tndm


/**
 * Copyright (c) 2011, 2012
 * Claudio Kopper <claudio.kopper@icecube.wisc.edu>
 * and the IceCube Collaboration <http://www.icecube.wisc.edu>
 *
 * Permission to use, copy, modify, and/or distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY
 * SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION
 * OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN
 * CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 *
 *
 * $Id: I3CLSimModuleHelper.cxx 178168 2019-12-19 21:40:33Z jvansanten $
 *
 * @file I3CLSimModuleHelper.cxx
 * @version $Revision: 178168 $
 * @date $Date: 2019-12-19 14:40:33 -0700 (Thu, 19 Dec 2019) $
 * @author Claudio Kopper
 */

#ifndef __STDC_FORMAT_MACROS
#define __STDC_FORMAT_MACROS
#endif
#include <inttypes.h>

#include "clsim/I3CLSimModuleHelper.h"

#include "clsim/function/I3CLSimFunctionConstant.h"
#include "clsim/function/I3CLSimFunctionFromTable.h"
#include "clsim/function/I3CLSimFunctionDeltaPeak.h"
#include "clsim/random_value/I3CLSimRandomValueInterpolatedDistribution.h"
#include "clsim/random_value/I3CLSimRandomValueWlenCherenkovNoDispersion.h"
#include "clsim/random_value/I3CLSimRandomValueConstant.h"

#include <boost/foreach.hpp>
#include <boost/algorithm/string.hpp>
#include <boost/variant/get.hpp>

#include "dataclasses/physics/I3MCTree.h"
#include "dataclasses/physics/I3MCTreeUtils.h"

namespace I3CLSimModuleHelper {

    namespace {
        double CherenkovYieldDistribution(double wlen, I3CLSimMediumPropertiesConstPtr mediumProperties, double beta=1.)
        {
            I3CLSimFunctionConstPtr nPhaseDist =
            mediumProperties->GetPhaseRefractiveIndex(0); // this assumes the refractive index does not change between layers

            if (!nPhaseDist->HasNativeImplementation())
                log_fatal("The refractive index distribution needs a native implementation to be usable!");

            const double nPhase = nPhaseDist->GetValue(wlen);

            return (2.*M_PI/(137.*(wlen*wlen)))*(1. - 1./ ( std::pow(beta*nPhase,2.) ) ); // dN/dxdwlen
        }

        // the normalization will not be correct here
        double CherenkovYieldDistributionNoDispersion(double wlen)
        {
            return 1./(wlen*wlen); // dN/dxdwlen
        }

    };

    I3CLSimRandomValueConstPtr
    makeWavelengthGenerator(I3CLSimFunctionConstPtr unbiasedSpectrum,
                            I3CLSimFunctionConstPtr wavelengthGenerationBias,
                            I3CLSimMediumPropertiesConstPtr mediumProperties)
    {
        {
            // special handling for delta peaks
            I3CLSimFunctionDeltaPeakConstPtr deltaPeak =
            boost::dynamic_pointer_cast<const I3CLSimFunctionDeltaPeak>(unbiasedSpectrum);
            if (deltaPeak) {
                const double peakPosition = deltaPeak->GetPeakPosition();

                return I3CLSimRandomValueConstantConstPtr
                (new I3CLSimRandomValueConstant(peakPosition));
            }
        }

        // if we get here, it's not a delta peak

        double minWlen = unbiasedSpectrum->GetMinWlen();
        double maxWlen = unbiasedSpectrum->GetMaxWlen();

        // check if the spectrum is from a tabulated distribution (instead of
        // a parameterized one)
        I3CLSimFunctionFromTableConstPtr unbiasedSpectrumFromTable;
        unbiasedSpectrumFromTable =
        boost::dynamic_pointer_cast<const I3CLSimFunctionFromTable>(unbiasedSpectrum);

        if (!unbiasedSpectrumFromTable) {
            // do not clip wavelengths if they are from a tabulated distribution. In that case,
            // re-use the entire table binning
            if (mediumProperties->GetMinWavelength() > minWlen) minWlen=mediumProperties->GetMinWavelength();
            if (mediumProperties->GetMaxWavelength() < maxWlen) maxWlen=mediumProperties->GetMaxWavelength();
        }

        const double wlenRange = maxWlen-minWlen;
        if (wlenRange <= 0.) log_fatal("Internal error, wavelength range <= 0!");

        if (wavelengthGenerationBias->GetMinWlen() > minWlen)
            log_fatal("wavelength generation bias has to have a wavelength range larger or equal to the spectrum wavelength range!");
        if (wavelengthGenerationBias->GetMaxWlen() < maxWlen)
            log_fatal("wavelength generation bias has to have a wavelength range larger or equal to the spectrum wavelength range!");

        // Check if the spectrum values are from a tabulated distribution.
        // If yes, use the table binning, if not make up a binning.
        if (unbiasedSpectrumFromTable)
        {
            std::size_t wlenPoints = unbiasedSpectrumFromTable->GetNumEntries();
            const double firstWlen = unbiasedSpectrumFromTable->GetFirstWavelength();

            std::vector<double> spectrum(wlenPoints, NAN);
            std::vector<double> wavelengths(wlenPoints, NAN);
            for (std::size_t i=0;i<wlenPoints;++i)
            {
                const double wavelength = unbiasedSpectrumFromTable->GetEntryWavelength(i);
                const double entry = unbiasedSpectrumFromTable->GetEntryValue(i);
                const double bias = wavelengthGenerationBias->GetValue(wavelength);

                spectrum[i] = bias * entry;
                wavelengths[i] = wavelength;
            }

            if (unbiasedSpectrumFromTable->GetInEqualSpacingMode()) {
                const double wlenStep = unbiasedSpectrumFromTable->GetWavelengthStepping();

                return I3CLSimRandomValueInterpolatedDistributionConstPtr
                (new I3CLSimRandomValueInterpolatedDistribution(firstWlen,
                                                                wlenStep,
                                                                spectrum));
            } else {
                // slightly less efficient if non-equally spaced
                return I3CLSimRandomValueInterpolatedDistributionConstPtr
                (new I3CLSimRandomValueInterpolatedDistribution(wavelengths,
                                                                spectrum));
            }
        }
        else
        {
            // use a pre-defined binning of 10ns (an arbitrary value..)

            std::size_t wlenPoints = static_cast<std::size_t>(wlenRange/(10.*I3Units::nanometer))+2;
            const double firstWlen = minWlen;
            const double wlenStep = wlenRange/static_cast<double>(wlenPoints-1);

            std::vector<double> spectrum(wlenPoints, NAN);
            for (std::size_t i=0;i<wlenPoints;++i)
            {
                const double wavelength = firstWlen + static_cast<double>(i)*wlenStep;
                const double entry = unbiasedSpectrum->GetValue(wavelength);
                const double bias = wavelengthGenerationBias->GetValue(wavelength);

                spectrum[i] = bias * entry;
            }

            return I3CLSimRandomValueInterpolatedDistributionConstPtr
            (new I3CLSimRandomValueInterpolatedDistribution(firstWlen,
                                                            wlenStep,
                                                            spectrum));
        }
    }

    I3CLSimRandomValueConstPtr
    makeCherenkovWavelengthGenerator(I3CLSimFunctionConstPtr wavelengthGenerationBias,
                                     bool generateCherenkovPhotonsWithoutDispersion,
                                     I3CLSimMediumPropertiesConstPtr mediumProperties)
    {
        const double minWlen = mediumProperties->GetMinWavelength();
        const double maxWlen = mediumProperties->GetMaxWavelength();
        const double wlenRange = maxWlen-minWlen;
        if (wlenRange <= 0.) log_fatal("Internal error, wavelength range <= 0!");

        if (wavelengthGenerationBias->GetMinWlen() > minWlen)
            log_fatal("wavelength generation bias has to have a wavelength range larger or equal to the medium property range!");
        if (wavelengthGenerationBias->GetMaxWlen() < maxWlen)
            log_fatal("wavelength generation bias has to have a wavelength range larger or equal to the medium property range!");

        bool noBias=false;
        //bool biasIsConstant=false;

        {
            I3CLSimFunctionConstantConstPtr wavelengthGenerationBiasConstant =
            boost::dynamic_pointer_cast<const I3CLSimFunctionConstant>(wavelengthGenerationBias);

            if (wavelengthGenerationBiasConstant)
            {
                //biasIsConstant=true;

                if ( std::abs(wavelengthGenerationBiasConstant->GetValue((minWlen+maxWlen)/2.)-1.) < 1e-10 )
                    noBias=true;
            }
        }

        I3CLSimFunctionFromTableConstPtr wavelengthGenerationBiasFromTable;
        wavelengthGenerationBiasFromTable =
        boost::dynamic_pointer_cast<const I3CLSimFunctionFromTable>(wavelengthGenerationBias);

        if ((!noBias) && (generateCherenkovPhotonsWithoutDispersion))
        {
            log_warn("**********");
            log_warn(" Using the \"GenerateCherenkovPhotonsWithoutDispersion\" option");
            log_warn(" with a biased photon spectrum generation does not yield a performance");
            log_warn(" increase. You might consider turning this option off to get a better");
            log_warn(" approximation of the Cherenkov spectrum.");
            log_warn("**********");
        }

        // Check if the bias values are from a tabulated distribution.
        // If yes, use the table binning, if not make up a binning.
        if (wavelengthGenerationBiasFromTable)
        {
            std::size_t wlenPoints = wavelengthGenerationBiasFromTable->GetNumEntries();
            const double firstWlen = wavelengthGenerationBiasFromTable->GetFirstWavelength();

            std::vector<double> spectrum(wlenPoints, NAN);
            std::vector<double> wavelengths(wlenPoints, NAN);
            for (std::size_t i=0;i<wlenPoints;++i)
            {
                const double wavelength = wavelengthGenerationBiasFromTable->GetEntryWavelength(i);
                const double bias = wavelengthGenerationBiasFromTable->GetEntryValue(i);

                if (generateCherenkovPhotonsWithoutDispersion)
                {
                    spectrum[i] =
                    bias * CherenkovYieldDistributionNoDispersion(wavelength);
                }
                else
                {
                    spectrum[i] =
                    bias * CherenkovYieldDistribution(wavelength, mediumProperties);
                }

                wavelengths[i] = wavelength;
            }

            if (wavelengthGenerationBiasFromTable->GetInEqualSpacingMode()) {
                const double wlenStep = wavelengthGenerationBiasFromTable->GetWavelengthStepping();

                return I3CLSimRandomValueInterpolatedDistributionConstPtr
                (new I3CLSimRandomValueInterpolatedDistribution(firstWlen,
                                                                wlenStep,
                                                                spectrum));
            } else {
                // slightly less efficient if non-equally spaced
                return I3CLSimRandomValueInterpolatedDistributionConstPtr
                (new I3CLSimRandomValueInterpolatedDistribution(wavelengths,
                                                                spectrum));
            }
        }
        else if ((noBias) && (generateCherenkovPhotonsWithoutDispersion))
        {
            return I3CLSimRandomValueWlenCherenkovNoDispersionConstPtr
            (new I3CLSimRandomValueWlenCherenkovNoDispersion(minWlen, maxWlen));
        }
        else
        {
            std::size_t wlenPoints = static_cast<std::size_t>(wlenRange/(10.*I3Units::nanometer))+2;
            const double firstWlen = minWlen;
            const double wlenStep = wlenRange/static_cast<double>(wlenPoints-1);

            std::vector<double> spectrum(wlenPoints, NAN);
            for (std::size_t i=0;i<wlenPoints;++i)
            {
                const double wavelength = firstWlen + static_cast<double>(i)*wlenStep;
                const double bias = wavelengthGenerationBias->GetValue(wavelength);

                if (generateCherenkovPhotonsWithoutDispersion)
                {
                    spectrum[i] =
                    bias * CherenkovYieldDistributionNoDispersion(wavelength);
                }
                else
                {
                    spectrum[i] =
                    bias * CherenkovYieldDistribution(wavelength, mediumProperties);
                }
            }

            return I3CLSimRandomValueInterpolatedDistributionConstPtr
            (new I3CLSimRandomValueInterpolatedDistribution(firstWlen,
                                                            wlenStep,
                                                            spectrum));
        }

    }

    I3CLSimStepToPhotonConverterOpenCLPtr initializeOpenCL(const I3CLSimOpenCLDevice &device,
                                                           I3RandomServicePtr rng,
                                                           I3CLSimSimpleGeometryConstPtr geometry,
                                                           I3CLSimMediumPropertiesConstPtr medium,
                                                           I3CLSimFunctionConstPtr wavelengthGenerationBias,
                                                           const std::vector<I3CLSimRandomValueConstPtr> &wavelengthGenerators,
                                                           bool enableDoubleBuffering,
                                                           bool doublePrecision,
                                                           bool stopDetectedPhotons,
                                                           bool saveAllPhotons,
                                                           double saveAllPhotonsPrescale,
                                                           double fixedNumberOfAbsorptionLengths,
                                                           double pancakeFactor,
                                                           uint32_t photonHistoryEntries,
                                                           uint32_t limitWorkgroupSize)
    {
        I3CLSimStepToPhotonConverterOpenCLPtr conv(new I3CLSimStepToPhotonConverterOpenCL(rng, device.GetUseNativeMath()));

        conv->SetDevice(device);

        conv->SetWlenGenerators(wavelengthGenerators);
        conv->SetWlenBias(wavelengthGenerationBias);

        conv->SetMediumProperties(medium);
        conv->SetGeometry(geometry);

        conv->SetEnableDoubleBuffering(enableDoubleBuffering);
        conv->SetDoublePrecision(doublePrecision);
        conv->SetStopDetectedPhotons(stopDetectedPhotons);
        conv->SetSaveAllPhotons(saveAllPhotons);
        conv->SetSaveAllPhotonsPrescale(saveAllPhotonsPrescale);

        conv->SetFixedNumberOfAbsorptionLengths(fixedNumberOfAbsorptionLengths);
        conv->SetDOMPancakeFactor(pancakeFactor);

        conv->SetPhotonHistoryEntries(photonHistoryEntries);

        conv->Compile();
        //log_trace("%s", conv.GetFullSource().c_str());

        std::size_t maxWorkgroupSize = conv->GetMaxWorkgroupSize();
        if (limitWorkgroupSize!=0) {
            maxWorkgroupSize = std::min(static_cast<std::size_t>(limitWorkgroupSize), maxWorkgroupSize);
        }

        conv->SetWorkgroupSize(maxWorkgroupSize);
        const std::size_t workgroupSize = conv->GetWorkgroupSize();

        // use approximately the given number of work items, convert to a multiple of the workgroup size
        std::size_t maxNumWorkitems = (static_cast<std::size_t>(device.GetApproximateNumberOfWorkItems())/workgroupSize)*workgroupSize;
        if (maxNumWorkitems==0) maxNumWorkitems=workgroupSize;

        conv->SetMaxNumWorkitems(maxNumWorkitems);

        log_info("Using OpenCL device: platform=%s device=%s isgpu=%d",
                  device.GetPlatformName().c_str(),
                  device.GetDeviceName().c_str(),
                  device.IsGPU());
        log_info("maximum workgroup size is %zu", maxWorkgroupSize);
        log_info("configured workgroup size is %zu", workgroupSize);
        if (maxNumWorkitems != device.GetApproximateNumberOfWorkItems()) {
            log_info("maximum number of work items is %zu (user configured was %" PRIu32 ")", maxNumWorkitems, device.GetApproximateNumberOfWorkItems());
        } else {
            log_debug("maximum number of work items is %zu (user configured was %" PRIu32 ")", maxNumWorkitems, device.GetApproximateNumberOfWorkItems());
        }

        conv->Initialize();

        return conv;
    }

}