/*
 * This file is part of a C++ interface to the Radiative Transfer for Energetics (RTE)
 * and Rapid Radiative Transfer Model for GCM applications Parallel (RRTMGP).
 *
 * The original code is found at https://github.com/RobertPincus/rte-rrtmgp.
 *
 * Contacts: Robert Pincus and Eli Mlawer
 * email: rrtmgp@aer.com
 *
 * Copyright 2015-2020,  Atmospheric and Environmental Research and
 * Regents of the University of Colorado.  All right reserved.
 *
 * This C++ interface can be downloaded from https://github.com/Chiil/rrtmgp_cpp
 *
 * Contact: Chiel van Heerwaarden
 * email: chiel.vanheerwaarden@wur.nl
 *
 * Copyright 2020, Wageningen University & Research.
 *
 * Use and duplication is permitted under the terms of the
 * BSD 3-clause license, see http://opensource.org/licenses/BSD-3-Clause
 *
 */

#include <cmath>
#include <numeric>
#include <boost/algorithm/string.hpp>
// #include <xtensor/xarray.hpp>
// #include <xtensor/xadapt.hpp>
// #include "xtensor/xnoalias.hpp"

#include "Gas_concs.h"
#include "Gas_optics.h"
#include "Array.h"
#include "Optical_props.h"
#include "Source_functions.h"

#include "rrtmgp_kernels.h"
#define restrict __restrict__

namespace
{
    int find_index(
            const Array<std::string,1>& data, const std::string& value)
    {
        auto it = std::find(data.v().begin(), data.v().end(), value);
        if (it == data.v().end())
            return -1;
        else
            return it - data.v().begin() + 1;
    }

    template<typename TF>
    void reduce_minor_arrays(
                const Gas_concs<TF>& available_gases,
                const Array<std::string,1>& gas_names,
                const Array<std::string,1>& gas_minor,
                const Array<std::string,1>& identifier_minor,
                const Array<TF,3>& kminor_atm,
                const Array<std::string,1>& minor_gases_atm,
                const Array<int,2>& minor_limits_gpt_atm,
                const Array<int,1>& minor_scales_with_density_atm, // CvH: logical bool or int?
                const Array<std::string,1>& scaling_gas_atm,
                const Array<int,1>& scale_by_complement_atm, // CvH: bool or int
                const Array<int,1>& kminor_start_atm,

                Array<TF,3>& kminor_atm_red,
                Array<std::string,1>& minor_gases_atm_red,
                Array<int,2>& minor_limits_gpt_atm_red,
                Array<int,1>& minor_scales_with_density_atm_red, // CvH bool or int
                Array<std::string,1>& scaling_gas_atm_red,
                Array<int,1>& scale_by_complement_atm_red,
                Array<int,1>& kminor_start_atm_red)
    {
        int nm = minor_gases_atm.dim(1);
        int tot_g = 0;

        Array<int,1> gas_is_present({nm});

        for (int i=1; i<=nm; ++i)
        {
            const int idx_mnr = find_index(identifier_minor, minor_gases_atm({i}));

            // Search for
            std::string gas_minor_trimmed = gas_minor({idx_mnr});
            boost::trim(gas_minor_trimmed);

            gas_is_present({i}) = available_gases.exists(gas_minor_trimmed);
            if (gas_is_present({i}))
                tot_g += minor_limits_gpt_atm({2,i}) - minor_limits_gpt_atm({1,i}) + 1;
        }

        const int red_nm = std::accumulate(gas_is_present.v().begin(), gas_is_present.v().end(), 0);

        if (red_nm == nm)
        {
            kminor_atm_red = kminor_atm;
            minor_gases_atm_red = minor_gases_atm;
            minor_limits_gpt_atm_red = minor_limits_gpt_atm;
            minor_scales_with_density_atm_red = minor_scales_with_density_atm;
            scaling_gas_atm_red = scaling_gas_atm;
            scale_by_complement_atm_red = scale_by_complement_atm;
            kminor_start_atm_red = kminor_start_atm;
        }
        else
        {
            // Use a lambda function as the operation has to be repeated many times.
            auto resize_and_set = [&](auto& a_red, const auto& a)
            {
                a_red.set_dims({red_nm});
                int counter = 1;
                for (int i=1; i<=gas_is_present.dim(1); ++i)
                {
                    if (gas_is_present({i}))
                    {
                       a_red({counter}) = a({i});
                       ++counter;
                    }
                }
            };

            resize_and_set(minor_gases_atm_red, minor_gases_atm);
            resize_and_set(minor_scales_with_density_atm_red, minor_scales_with_density_atm);
            resize_and_set(scaling_gas_atm_red, scaling_gas_atm);
            resize_and_set(scale_by_complement_atm_red, scale_by_complement_atm);
            resize_and_set(kminor_start_atm_red, kminor_start_atm);

            minor_limits_gpt_atm_red.set_dims({2, red_nm});
            kminor_atm_red.set_dims({tot_g, kminor_atm.dim(2), kminor_atm.dim(3)});

            int icnt = 0;
            int n_elim = 0;
            for (int i=1; i<=nm; ++i)
            {
                int ng = minor_limits_gpt_atm({2,i}) - minor_limits_gpt_atm({1,i}) + 1;
                if (gas_is_present({i}))
                {
                    ++icnt;
                    minor_limits_gpt_atm_red({1,icnt}) = minor_limits_gpt_atm({1,i});
                    minor_limits_gpt_atm_red({2,icnt}) = minor_limits_gpt_atm({2,i});
                    kminor_start_atm_red({icnt}) = kminor_start_atm({i}) - n_elim;

                    for (int j=1; j<=ng; ++j)
                        for (int i2=1; i2<=kminor_atm.dim(2); ++i2)
                            for (int i3=1; i3<=kminor_atm.dim(3); ++i3)
                                kminor_atm_red({kminor_start_atm_red({icnt})+j-1,i2,i3}) =
                                        kminor_atm({kminor_start_atm({i})+j-1,i2,i3});
                }
                else
                    n_elim += ng;
            }
        }
    }

    void create_idx_minor(
            const Array<std::string,1>& gas_names,
            const Array<std::string,1>& gas_minor,
            const Array<std::string,1>& identifier_minor,
            const Array<std::string,1>& minor_gases_atm,
            Array<int,1>& idx_minor_atm)
    {
        Array<int,1> idx_minor_atm_out({minor_gases_atm.dim(1)});

        for (int imnr=1; imnr<=minor_gases_atm.dim(1); ++imnr)
        {
            // Find identifying string for minor species in list of possible identifiers (e.g. h2o_slf)
            const int idx_mnr = find_index(identifier_minor, minor_gases_atm({imnr}));

            // Find name of gas associated with minor species identifier (e.g. h2o)
            idx_minor_atm_out({imnr}) = find_index(gas_names, gas_minor({idx_mnr}));
        }

        idx_minor_atm = idx_minor_atm_out;
    }

    void create_idx_minor_scaling(
            const Array<std::string,1>& gas_names,
            const Array<std::string,1>& scaling_gas_atm,
            Array<int,1>& idx_minor_scaling_atm)
    {
        Array<int,1> idx_minor_scaling_atm_out({scaling_gas_atm.dim(1)});

        for (int imnr=1; imnr<=scaling_gas_atm.dim(1); ++imnr)
            idx_minor_scaling_atm_out({imnr}) = find_index(gas_names, scaling_gas_atm({imnr}));

        idx_minor_scaling_atm = idx_minor_scaling_atm_out;
    }

    void create_key_species_reduce(
            const Array<std::string,1>& gas_names,
            const Array<std::string,1>& gas_names_red,
            const Array<int,3>& key_species,
            Array<int,3>& key_species_red,
            Array<int,1>& key_species_present_init)
    {
        const int np = key_species.dim(1);
        const int na = key_species.dim(2);
        const int nt = key_species.dim(3);

        key_species_red.set_dims({key_species.dim(1), key_species.dim(2), key_species.dim(3)});
        key_species_present_init.set_dims({gas_names.dim(1)});

        for (int i=1; i<=key_species_present_init.dim(1); ++i)
            key_species_present_init({i}) = 1;

        for (int ip=1; ip<=np; ++ip)
            for (int ia=1; ia<=na; ++ia)
                for (int it=1; it<=nt; ++it)
                {
                    const int ks = key_species({ip,ia,it});
                    if (ks != 0)
                    {
                        const int ksr = find_index(gas_names_red, gas_names({ks}));
                        key_species_red({ip,ia,it}) = ksr;
                        if (ksr == -1)
                            key_species_present_init({ks}) = 0;
                    }
                    else
                        key_species_red({ip,ia,it}) = ks;
                }
    }

    void check_key_species_present_init(
            const Array<std::string,1>& gas_names,
            const Array<int,1>& key_species_present_init
            )
    {
        for (int i=1; i<=key_species_present_init.dim(1); ++i)
        {
            if (key_species_present_init({i}) == 0)
            {
                std::string error_message = "Gas optics: required gas " + gas_names({i}) + " is missing";
                throw std::runtime_error(error_message);
            }
        }
    }

    void create_flavor(
            const Array<int,3>& key_species,
            Array<int,2>& flavor)
    {
        Array<int,2> key_species_list({2, key_species.dim(3)*2});

        // Prepare list of key species.
        int i = 1;
        for (int ibnd=1; ibnd<=key_species.dim(3); ++ibnd)
            for (int iatm=1; iatm<=key_species.dim(1); ++iatm)
            {
                key_species_list({1,i}) = key_species({1,iatm,ibnd});
                key_species_list({2,i}) = key_species({2,iatm,ibnd});
                ++i;
            }

        // Rewrite single key_species pairs.
        for (int i=1; i<=key_species_list.dim(2); ++i)
        {
            if ( key_species_list({1,i}) == 0 && key_species_list({2,i}) == 0 )
            {
                key_species_list({1,i}) = 2;
                key_species_list({2,i}) = 2;
            }
        }

        // Count unique key species pairs.
        int iflavor = 0;
        for (int i=1; i<=key_species_list.dim(2); ++i)
        {
            bool pair_exists = false;
            for (int ii=1; ii<=i-1; ++ii)
            {
                if ( (key_species_list({1,i}) == key_species_list({1,ii})) &&
                     (key_species_list({2,i}) == key_species_list({2,ii})) )
                {
                    pair_exists = true;
                    break;
                }
            }
            if (!pair_exists)
                ++iflavor;
        }

        // Fill flavors.
        flavor.set_dims({2,iflavor});
        iflavor = 0;
        for (int i=1; i<=key_species_list.dim(2); ++i)
        {
            bool pair_exists = false;
            for (int ii=1; ii<=i-1; ++ii)
            {
                if ( (key_species_list({1,i}) == key_species_list({1,ii})) &&
                     (key_species_list({2,i}) == key_species_list({2,ii})) )
                {
                    pair_exists = true;
                    break;
                }
            }
            if (!pair_exists)
            {
                ++iflavor;
                flavor({1,iflavor}) = key_species_list({1,i});
                flavor({2,iflavor}) = key_species_list({2,i});
            }
        }
    }

    int key_species_pair2flavor(
            const Array<int,2>& flavor,
            const Array<int,1>& key_species_pair)
    {
        // Search for match.
        for (int iflav=1; iflav<=flavor.dim(2); ++iflav)
        {
            if ( key_species_pair({1}) == flavor({1, iflav}) &&
                 key_species_pair({2}) == flavor({2, iflav}) )
                return iflav;
        }

        // No match found.
        return -1;
    }

    void create_gpoint_flavor(
            const Array<int,3>& key_species,
            const Array<int,1>& gpt2band,
            const Array<int,2>& flavor,
            Array<int,2>& gpoint_flavor)
    {
        const int ngpt = gpt2band.dim(1);
        gpoint_flavor.set_dims({2,ngpt});

        for (int igpt=1; igpt<=ngpt; ++igpt)
            for (int iatm=1; iatm<=2; ++iatm)
            {
                int pair_1 = key_species( {1, iatm, gpt2band({igpt})} );
                int pair_2 = key_species( {2, iatm, gpt2band({igpt})} );

                // Rewrite species pair.
                Array<int,1> rewritten_pair({2});
                if (pair_1 == 0 && pair_2 == 0)
                {
                    rewritten_pair({1}) = 2;
                    rewritten_pair({2}) = 2;
                }
                else
                {
                    rewritten_pair({1}) = pair_1;
                    rewritten_pair({2}) = pair_2;
                }

                // Write the output.
                gpoint_flavor({iatm,igpt}) = key_species_pair2flavor(
                        flavor, rewritten_pair);
            }
    }

    /*
    template<typename TF>
    inline void reorder123x321_test(
            TF* restrict out, const TF* restrict in,
            const int d1, const int d2, const int d3)
    {
        const int jj_in = d1;
        const int kk_in = d1*d2;

        const int ii_out = d3*d2;
        const int jj_out = d3;

        for (int i=0; i<d1; ++i)
            for (int j=0; j<d2; ++j)
                #pragma GCC ivdep
                for (int k=0; k<d3; ++k)
                {
                    const int ijk_in  = i + j*jj_in  + k*kk_in ;
                    const int ijk_out = k + j*jj_out + i*ii_out;
                    out[ijk_out] = in[ijk_in];
                }
    }

    template<typename TF>
    inline void reorder123x321_test(
            TF* out, const TF* in,
            const size_t d1, const size_t d2, const size_t d3)
    {
        const size_t size = d1*d2*d3;
        const std::array<size_t, 3> in_shape = { d1, d2, d3 };
        const std::array<size_t, 3> out_shape = { d3, d2, d1 };

        const auto a_in = xt::adapt<xt::layout_type::column_major>(in, size, xt::no_ownership(), in_shape);
        auto a_out = xt::adapt<xt::layout_type::column_major>(out, size, xt::no_ownership(), out_shape);

        xt::noalias(a_out) = xt::transpose(a_out);
    }
    */
}

// IMPLEMENTATION OF CLASS FUNCTIONS.
// Constructor of longwave variant.
template<typename TF>
Gas_optics<TF>::Gas_optics(
        const Gas_concs<TF>& available_gases,
        const Array<std::string,1>& gas_names,
        const Array<int,3>& key_species,
        const Array<int,2>& band2gpt,
        const Array<TF,2>& band_lims_wavenum,
        const Array<TF,1>& press_ref,
        const TF press_ref_trop,
        const Array<TF,1>& temp_ref,
        const TF temp_ref_p,
        const TF temp_ref_t,
        const Array<TF,3>& vmr_ref,
        const Array<TF,4>& kmajor,
        const Array<TF,3>& kminor_lower,
        const Array<TF,3>& kminor_upper,
        const Array<std::string,1>& gas_minor,
        const Array<std::string,1>& identifier_minor,
        const Array<std::string,1>& minor_gases_lower,
        const Array<std::string,1>& minor_gases_upper,
        const Array<int,2>& minor_limits_gpt_lower,
        const Array<int,2>& minor_limits_gpt_upper,
        const Array<int,1>& minor_scales_with_density_lower,
        const Array<int,1>& minor_scales_with_density_upper,
        const Array<std::string,1>& scaling_gas_lower,
        const Array<std::string,1>& scaling_gas_upper,
        const Array<int,1>& scale_by_complement_lower,
        const Array<int,1>& scale_by_complement_upper,
        const Array<int,1>& kminor_start_lower,
        const Array<int,1>& kminor_start_upper,
        const Array<TF,2>& totplnk,
        const Array<TF,4>& planck_frac,
        const Array<TF,3>& rayl_lower,
        const Array<TF,3>& rayl_upper) :
            Optical_props<TF>(band_lims_wavenum, band2gpt),
            totplnk(totplnk),
            planck_frac(planck_frac)
{
    // Initialize the absorption coefficient array, including Rayleigh scattering
    // tables if provided.
    init_abs_coeffs(
            available_gases,
            gas_names, key_species,
            band2gpt, band_lims_wavenum,
            press_ref, temp_ref,
            press_ref_trop, temp_ref_p, temp_ref_t,
            vmr_ref,
            kmajor, kminor_lower, kminor_upper,
            gas_minor,identifier_minor,
            minor_gases_lower, minor_gases_upper,
            minor_limits_gpt_lower,
            minor_limits_gpt_upper,
            minor_scales_with_density_lower,
            minor_scales_with_density_upper,
            scaling_gas_lower, scaling_gas_upper,
            scale_by_complement_lower,
            scale_by_complement_upper,
            kminor_start_lower,
            kminor_start_upper,
            rayl_lower, rayl_upper);

    // Temperature steps for Planck function interpolation.
    // Assumes that temperature minimum and max are the same for the absorption coefficient grid and the
    // Planck grid and the Planck grid is equally spaced.
    totplnk_delta = (temp_ref_max - temp_ref_min) / (totplnk.dim(1)-1);
}

// Constructor of the shortwave variant.
template<typename TF>
Gas_optics<TF>::Gas_optics(
        const Gas_concs<TF>& available_gases,
        const Array<std::string,1>& gas_names,
        const Array<int,3>& key_species,
        const Array<int,2>& band2gpt,
        const Array<TF,2>& band_lims_wavenum,
        const Array<TF,1>& press_ref,
        const TF press_ref_trop,
        const Array<TF,1>& temp_ref,
        const TF temp_ref_p,
        const TF temp_ref_t,
        const Array<TF,3>& vmr_ref,
        const Array<TF,4>& kmajor,
        const Array<TF,3>& kminor_lower,
        const Array<TF,3>& kminor_upper,
        const Array<std::string,1>& gas_minor,
        const Array<std::string,1>& identifier_minor,
        const Array<std::string,1>& minor_gases_lower,
        const Array<std::string,1>& minor_gases_upper,
        const Array<int,2>& minor_limits_gpt_lower,
        const Array<int,2>& minor_limits_gpt_upper,
        const Array<int,1>& minor_scales_with_density_lower,
        const Array<int,1>& minor_scales_with_density_upper,
        const Array<std::string,1>& scaling_gas_lower,
        const Array<std::string,1>& scaling_gas_upper,
        const Array<int,1>& scale_by_complement_lower,
        const Array<int,1>& scale_by_complement_upper,
        const Array<int,1>& kminor_start_lower,
        const Array<int,1>& kminor_start_upper,
        const Array<TF,1>& solar_source_quiet,
        const Array<TF,1>& solar_source_facular,
        const Array<TF,1>& solar_source_sunspot,
        const TF tsi_default,
        const TF mg_default,
        const TF sb_default,
        const Array<TF,3>& rayl_lower,
        const Array<TF,3>& rayl_upper) :
            Optical_props<TF>(band_lims_wavenum, band2gpt)
{
    // Initialize the absorption coefficient array, including Rayleigh scattering
    // tables if provided.
    init_abs_coeffs(
            available_gases,
            gas_names, key_species,
            band2gpt, band_lims_wavenum,
            press_ref, temp_ref,
            press_ref_trop, temp_ref_p, temp_ref_t,
            vmr_ref,
            kmajor, kminor_lower, kminor_upper,
            gas_minor,identifier_minor,
            minor_gases_lower, minor_gases_upper,
            minor_limits_gpt_lower,
            minor_limits_gpt_upper,
            minor_scales_with_density_lower,
            minor_scales_with_density_upper,
            scaling_gas_lower, scaling_gas_upper,
            scale_by_complement_lower,
            scale_by_complement_upper,
            kminor_start_lower,
            kminor_start_upper,
            rayl_lower, rayl_upper);

    // Compute the solar source.
    this->solar_source_quiet = solar_source_quiet;
    this->solar_source_facular = solar_source_facular;
    this->solar_source_sunspot = solar_source_sunspot;

    this->solar_source.set_dims(solar_source_quiet.get_dims());

    set_solar_variability(mg_default, sb_default);
}

template<typename TF>
void Gas_optics<TF>::init_abs_coeffs(
        const Gas_concs<TF>& available_gases,
        const Array<std::string,1>& gas_names,
        const Array<int,3>& key_species,
        const Array<int,2>& band2gpt,
        const Array<TF,2>& band_lims_wavenum,
        const Array<TF,1>& press_ref,
        const Array<TF,1>& temp_ref,
        const TF press_ref_trop,
        const TF temp_ref_p,
        const TF temp_ref_t,
        const Array<TF,3>& vmr_ref,
        const Array<TF,4>& kmajor,
        const Array<TF,3>& kminor_lower,
        const Array<TF,3>& kminor_upper,
        const Array<std::string,1>& gas_minor,
        const Array<std::string,1>& identifier_minor,
        const Array<std::string,1>& minor_gases_lower,
        const Array<std::string,1>& minor_gases_upper,
        const Array<int,2>& minor_limits_gpt_lower,
        const Array<int,2>& minor_limits_gpt_upper,
        const Array<int,1>& minor_scales_with_density_lower,
        const Array<int,1>& minor_scales_with_density_upper,
        const Array<std::string,1>& scaling_gas_lower,
        const Array<std::string,1>& scaling_gas_upper,
        const Array<int,1>& scale_by_complement_lower,
        const Array<int,1>& scale_by_complement_upper,
        const Array<int,1>& kminor_start_lower,
        const Array<int,1>& kminor_start_upper,
        const Array<TF,3>& rayl_lower,
        const Array<TF,3>& rayl_upper)
{
    // Which gases known to the gas optics are present in the host model (available_gases)?
    std::vector<std::string> gas_names_to_use;

    for (const std::string &s : gas_names.v())
    {
        if (available_gases.exists(s))
            gas_names_to_use.push_back(s);
    }

    // Now the number of gases is the union of those known to the k-distribution and provided
    // by the host model.
    const int n_gas = gas_names_to_use.size();
    Array<std::string,1> gas_names_this(std::move(gas_names_to_use), {n_gas});
    this->gas_names = gas_names_this;

    // Initialize the gas optics object, keeping only those gases known to the
    // gas optics and also present in the host model.
    // Add an offset to the indexing to interface the negative ranging of fortran.
    Array<TF,3> vmr_ref_red({vmr_ref.dim(1), n_gas + 1, vmr_ref.dim(3)});
    vmr_ref_red.set_offsets({0, -1, 0});

    // Gas 0 is used in single-key species method, set to 1.0 (col_dry)
    for (int i1=1; i1<=vmr_ref_red.dim(1); ++i1)
        for (int i3=1; i3<=vmr_ref_red.dim(3); ++i3)
            vmr_ref_red({i1, 0, i3}) = vmr_ref({i1, 1, i3});

    for (int i=1; i<=n_gas; ++i)
    {
        int idx = find_index(gas_names, this->gas_names({i}));
        for (int i1=1; i1<=vmr_ref_red.dim(1); ++i1)
            for (int i3=1; i3<=vmr_ref_red.dim(3); ++i3)
                vmr_ref_red({i1, i, i3}) = vmr_ref({i1, idx+1, i3}); // CvH: why +1?
    }

    this->vmr_ref = std::move(vmr_ref_red);

    // Reduce minor arrays so variables only contain minor gases that are available.
    // Reduce size of minor Arrays.
    Array<std::string, 1> minor_gases_lower_red;
    Array<std::string, 1> scaling_gas_lower_red;
    Array<std::string, 1> minor_gases_upper_red;
    Array<std::string, 1> scaling_gas_upper_red;

    reduce_minor_arrays(
            available_gases,
            gas_names,
            gas_minor, identifier_minor,
            kminor_lower,
            minor_gases_lower,
            minor_limits_gpt_lower,
            minor_scales_with_density_lower,
            scaling_gas_lower,
            scale_by_complement_lower,
            kminor_start_lower,
            this->kminor_lower,
            minor_gases_lower_red,
            this->minor_limits_gpt_lower,
            this->minor_scales_with_density_lower,
            scaling_gas_lower_red,
            this->scale_by_complement_lower,
            this->kminor_start_lower);

    reduce_minor_arrays(
            available_gases,
            gas_names,
            gas_minor,
            identifier_minor,
            kminor_upper,
            minor_gases_upper,
            minor_limits_gpt_upper,
            minor_scales_with_density_upper,
            scaling_gas_upper,
            scale_by_complement_upper,
            kminor_start_upper,
            this->kminor_upper,
            minor_gases_upper_red,
            this->minor_limits_gpt_upper,
            this->minor_scales_with_density_upper,
            scaling_gas_upper_red,
            this->scale_by_complement_upper,
            this->kminor_start_upper);

    // Arrays not reduced by the presence, or lack thereof, of a gas
    this->press_ref = press_ref;
    this->temp_ref = temp_ref;
    this->kmajor = kmajor;

    // Create a new vector that consists of rayl_lower and rayl_upper stored in one variable.
    if (rayl_lower.size() > 0)
    {
        this->krayl.set_dims({rayl_lower.dim(1), rayl_lower.dim(2), rayl_lower.dim(3), 2});
        for (int i=0; i<rayl_lower.size(); ++i)
        {
            this->krayl.v()[i                    ] = rayl_lower.v()[i];
            this->krayl.v()[i + rayl_lower.size()] = rayl_upper.v()[i];
        }
    }

    // ---- post processing ----
    //  creates log reference pressure
    this->press_ref_log = this->press_ref;
    for (int i1=1; i1<=this->press_ref_log.dim(1); ++i1)
        this->press_ref_log({i1}) = std::log(this->press_ref_log({i1}));

    // log scale of reference pressure
    this->press_ref_trop_log = std::log(press_ref_trop);

    // Get index of gas (if present) for determining col_gas
    create_idx_minor(
            this->gas_names, gas_minor, identifier_minor, minor_gases_lower_red, this->idx_minor_lower);
    create_idx_minor(
            this->gas_names, gas_minor, identifier_minor, minor_gases_upper_red, this->idx_minor_upper);

    // Get index of gas (if present) that has special treatment in density scaling
    create_idx_minor_scaling(
            this->gas_names, scaling_gas_lower_red, this->idx_minor_scaling_lower);
    create_idx_minor_scaling(
            this->gas_names, scaling_gas_upper_red, this->idx_minor_scaling_upper);

    // Create flavor list.
    // Reduce (remap) key_species list; checks that all key gases are present in incoming
    Array<int, 3> key_species_red;
    Array<int, 1> key_species_present_init; // CvH bool or int?

    create_key_species_reduce(
            gas_names, this->gas_names, key_species, key_species_red, key_species_present_init);

    check_key_species_present_init(gas_names, key_species_present_init);

    // create flavor list
    create_flavor(key_species_red, this->flavor);

    // create gpoint flavor list
    create_gpoint_flavor(
            key_species_red, this->get_gpoint_bands(), this->flavor, this->gpoint_flavor);

    // minimum, maximum reference temperature, pressure -- assumes low-to-high ordering
    // for T, high-to-low ordering for p
    this->temp_ref_min = this->temp_ref({1});
    this->temp_ref_max = this->temp_ref({temp_ref.dim(1)});
    this->press_ref_min = this->press_ref({press_ref.dim(1)});
    this->press_ref_max = this->press_ref({1});

    // creates press_ref_log, temp_ref_delta
    this->press_ref_log_delta =
            (std::log(this->press_ref_min) - std::log(this->press_ref_max)) / (this->press_ref.dim(1) - 1);
    this->temp_ref_delta = (this->temp_ref_max - this->temp_ref_min) / (this->temp_ref.dim(1) - 1);

    // Which species are key in one or more bands?
    // this->flavor is an index into this->gas_names
    // if (allocated(this%is_key)) deallocate(this%is_key) ! Shouldn't ever happen...
    Array<int,1> is_key({get_ngas()}); // CvH bool, defaults to 0.?

    for (int j=1; j<=this->flavor.dim(2); ++j)
        for (int i=1; i<=this->flavor.dim(1); ++i)
        {
            if (this->flavor({i, j}) != 0)
                is_key({this->flavor({i, j})}) = true;
        }

    this->is_key = is_key;
}

template<typename TF>
void Gas_optics<TF>::set_solar_variability(
        const TF mg_index, const TF sb_index)
{
    constexpr TF a_offset = TF(0.1495954);
    constexpr TF b_offset = TF(0.00066696);

    for (int igpt=1; igpt<=this->solar_source_quiet.dim(1); ++igpt)
    {
        this->solar_source({igpt}) = this->solar_source_quiet({igpt})
                + (mg_index - a_offset) * this->solar_source_facular({igpt})
                + (sb_index - b_offset) * this->solar_source_sunspot({igpt});
    }

    // Scale solar source to input TSI value
    // if (present(tsi)) error_msg = this%set_tsi(tsi)
}

// Calculate the molecules of dry air.
template<typename TF>
void Gas_optics<TF>::get_col_dry(
        Array<TF,2>& col_dry, const Array<TF,2>& vmr_h2o,
        const Array<TF,2>& plev)
{
    // CvH: RRTMGP uses more accurate method based on latitude.
    constexpr TF g0 = 9.80665;

    constexpr TF avogad = 6.02214076e23;
    constexpr TF m_dry = 0.028964;
    constexpr TF m_h2o = 0.018016;

    Array<double,2> delta_plev({col_dry.dim(1), col_dry.dim(2)});
    Array<double,2> m_air     ({col_dry.dim(1), col_dry.dim(2)});

    for (int ilay=1; ilay<=col_dry.dim(2); ++ilay)
        for (int icol=1; icol<=col_dry.dim(1); ++icol)
            delta_plev({icol, ilay}) = std::abs(plev({icol, ilay}) - plev({icol, ilay+1}));

    for (int ilay=1; ilay<=col_dry.dim(2); ++ilay)
        for (int icol=1; icol<=col_dry.dim(1); ++icol)
            m_air({icol, ilay}) = (m_dry + m_h2o * vmr_h2o({icol, ilay})) / (1. + vmr_h2o({icol, ilay}));

    for (int ilay=1; ilay<=col_dry.dim(2); ++ilay)
        for (int icol=1; icol<=col_dry.dim(1); ++icol)
        {
            col_dry({icol, ilay}) = TF(10.) * delta_plev({icol, ilay}) * avogad / (TF(1000.)*m_air({icol, ilay})*TF(100.)*g0);
            col_dry({icol, ilay}) /= (TF(1.) + vmr_h2o({icol, ilay}));
        }
}

// Gas optics solver longwave variant.
template<typename TF>
void Gas_optics<TF>::gas_optics(
        const Array<TF,2>& play,
        const Array<TF,2>& plev,
        const Array<TF,2>& tlay,
        const Array<TF,1>& tsfc,
        const Gas_concs<TF>& gas_desc,
        std::unique_ptr<Optical_props_arry<TF>>& optical_props,
        Source_func_lw<TF>& sources,
        const Array<TF,2>& col_dry,
        const Array<TF,2>& tlev) const
{
    const int ncol = play.dim(1);
    const int nlay = play.dim(2);
    const int ngpt = this->get_ngpt();
    const int nband = this->get_nband();

    Array<int,2> jtemp({play.dim(1), play.dim(2)});
    Array<int,2> jpress({play.dim(1), play.dim(2)});
    Array<int,2> tropo({play.dim(1), play.dim(2)});
    Array<TF,6> fmajor({2, 2, 2, this->get_nflav(), play.dim(1), play.dim(2)});
    Array<int,4> jeta({2, this->get_nflav(), play.dim(1), play.dim(2)});

    // Gas optics.
    compute_gas_taus(
            ncol, nlay, ngpt, nband,
            play, plev, tlay, gas_desc,
            optical_props,
            jtemp, jpress, jeta, tropo, fmajor,
            col_dry);

    // External sources.
    source(
            ncol, nlay, nband, ngpt,
            play, plev, tlay, tsfc,
            jtemp, jpress, jeta, tropo, fmajor,
            sources, tlev);
}

// Gas optics solver shortwave variant.
template<typename TF>
void Gas_optics<TF>::gas_optics(
        const Array<TF,2>& play,
        const Array<TF,2>& plev,
        const Array<TF,2>& tlay,
        const Gas_concs<TF>& gas_desc,
        std::unique_ptr<Optical_props_arry<TF>>& optical_props,
        Array<TF,2>& toa_src,
        const Array<TF,2>& col_dry) const
{
    const int ncol = play.dim(1);
    const int nlay = play.dim(2);
    const int ngpt = this->get_ngpt();
    const int nband = this->get_nband();

    Array<int,2> jtemp({play.dim(1), play.dim(2)});
    Array<int,2> jpress({play.dim(1), play.dim(2)});
    Array<int,2> tropo({play.dim(1), play.dim(2)});
    Array<TF,6> fmajor({2, 2, 2, this->get_nflav(), play.dim(1), play.dim(2)});
    Array<int,4> jeta({2, this->get_nflav(), play.dim(1), play.dim(2)});

    // Gas optics.
    compute_gas_taus(
            ncol, nlay, ngpt, nband,
            play, plev, tlay, gas_desc,
            optical_props,
            jtemp, jpress, jeta, tropo, fmajor,
            col_dry);

    // External source function is constant.
    for (int igpt=1; igpt<=ngpt; ++igpt)
        for (int icol=1; icol<=ncol; ++icol)
            toa_src({icol, igpt}) = this->solar_source({igpt});
}

namespace rrtmgp_kernel_launcher
{
    template<typename TF> void zero_array(
            int ni, int nj, int nk, Array<TF,3>& array)
    {
        rrtmgp_kernels::zero_array_3D(&ni, &nj, &nk, array.ptr());
    }

    template<typename TF> void zero_array(
            int ni, int nj, int nk, int nl, Array<TF,4>& array)
    {
        rrtmgp_kernels::zero_array_4D(&ni, &nj, &nk, &nl, array.ptr());
    }

    template<typename TF>
    void interpolation(
            int ncol, int nlay,
            int ngas, int nflav, int neta, int npres, int ntemp,
            const Array<int,2>& flavor,
            const Array<TF,1>& press_ref_log,
            const Array<TF,1>& temp_ref,
            TF press_ref_log_delta,
            TF temp_ref_min,
            TF temp_ref_delta,
            TF press_ref_trop_log,
            const Array<TF,3>& vmr_ref,
            const Array<TF,2>& play,
            const Array<TF,2>& tlay,
            Array<TF,3>& col_gas,
            Array<int,2>& jtemp,
            Array<TF,6>& fmajor, Array<TF,5>& fminor,
            Array<TF,4>& col_mix,
            Array<int,2>& tropo,
            Array<int,4>& jeta,
            Array<int,2>& jpress)
    {
        rrtmgp_kernels::interpolation(
                &ncol, &nlay,
                &ngas, &nflav, &neta, &npres, &ntemp,
                const_cast<int*>(flavor.ptr()),
                const_cast<TF*>(press_ref_log.ptr()),
                const_cast<TF*>(temp_ref.ptr()),
                &press_ref_log_delta,
                &temp_ref_min,
                &temp_ref_delta,
                &press_ref_trop_log,
                const_cast<TF*>(vmr_ref.ptr()),
                const_cast<TF*>(play.ptr()),
                const_cast<TF*>(tlay.ptr()),
                col_gas.ptr(),
                jtemp.ptr(),
                fmajor.ptr(), fminor.ptr(),
                col_mix.ptr(),
                tropo.ptr(),
                jeta.ptr(),
                jpress.ptr());
    }

    template<typename TF>
    void compute_tau_absorption(
            int ncol, int nlay, int nband, int ngpt,
            int ngas, int nflav, int neta, int npres, int ntemp,
            int nminorlower, int nminorklower,
            int nminorupper, int nminorkupper,
            int idx_h2o,
            const Array<int,2>& gpoint_flavor,
            const Array<int,2>& band_lims_gpt,
            const Array<TF,4>& kmajor,
            const Array<TF,3>& kminor_lower,
            const Array<TF,3>& kminor_upper,
            const Array<int,2>& minor_limits_gpt_lower,
            const Array<int,2>& minor_limits_gpt_upper,
            const Array<int,1>& minor_scales_with_density_lower,
            const Array<int,1>& minor_scales_with_density_upper,
            const Array<int,1>& scale_by_complement_lower,
            const Array<int,1>& scale_by_complement_upper,
            const Array<int,1>& idx_minor_lower,
            const Array<int,1>& idx_minor_upper,
            const Array<int,1>& idx_minor_scaling_lower,
            const Array<int,1>& idx_minor_scaling_upper,
            const Array<int,1>& kminor_start_lower,
            const Array<int,1>& kminor_start_upper,
            Array<int,2>& tropo,
            Array<TF,4>& col_mix, Array<TF,6>& fmajor, Array<TF,5>& fminor,
            const Array<TF,2>& play, const Array<TF,2>& tlay, Array<TF,3>& col_gas,
            Array<int,4>& jeta, Array<int,2>& jtemp, Array<int,2>& jpress,
            Array<TF,3>& tau)
    {
        rrtmgp_kernels::compute_tau_absorption(
            &ncol, &nlay, &nband, &ngpt,
            &ngas, &nflav, &neta, &npres, &ntemp,
            &nminorlower, &nminorklower,
            &nminorupper, &nminorkupper,
            &idx_h2o,
            const_cast<int*>(gpoint_flavor.ptr()),
            const_cast<int*>(band_lims_gpt.ptr()),
            const_cast<TF*>(kmajor.ptr()),
            const_cast<TF*>(kminor_lower.ptr()),
            const_cast<TF*>(kminor_upper.ptr()),
            const_cast<int*>(minor_limits_gpt_lower.ptr()),
            const_cast<int*>(minor_limits_gpt_upper.ptr()),
            const_cast<int*>(minor_scales_with_density_lower.ptr()),
            const_cast<int*>(minor_scales_with_density_upper.ptr()),
            const_cast<int*>(scale_by_complement_lower.ptr()),
            const_cast<int*>(scale_by_complement_upper.ptr()),
            const_cast<int*>(idx_minor_lower.ptr()),
            const_cast<int*>(idx_minor_upper.ptr()),
            const_cast<int*>(idx_minor_scaling_lower.ptr()),
            const_cast<int*>(idx_minor_scaling_upper.ptr()),
            const_cast<int*>(kminor_start_lower.ptr()),
            const_cast<int*>(kminor_start_upper.ptr()),
            tropo.ptr(),
            col_mix.ptr(), fmajor.ptr(), fminor.ptr(),
            const_cast<TF*>(play.ptr()), const_cast<TF*>(tlay.ptr()), col_gas.ptr(),
            jeta.ptr(), jtemp.ptr(), jpress.ptr(),
            tau.ptr());
    }

    template<typename TF>
    void compute_tau_rayleigh(
            int ncol, int nlay, int nband, int ngpt,
            int ngas, int nflav, int neta, int npres, int ntemp,
            const Array<int,2>& gpoint_flavor,
            const Array<int,2>& band_lims_gpt,
            const Array<TF,4>& krayl,
            int idx_h2o, const Array<TF,2>& col_dry, const Array<TF,3>& col_gas,
            const Array<TF,5>& fminor, const Array<int,4>& jeta,
            const Array<int,2>& tropo, const Array<int,2>& jtemp,
            Array<TF,3>& tau_rayleigh)
    {
        rrtmgp_kernels::compute_tau_rayleigh(
                &ncol, &nlay, &nband, &ngpt,
                &ngas, &nflav, &neta, &npres, &ntemp,
                const_cast<int*>(gpoint_flavor.ptr()),
                const_cast<int*>(band_lims_gpt.ptr()),
                const_cast<TF*>(krayl.ptr()),
                &idx_h2o,
                const_cast<TF*>(col_dry.ptr()), const_cast<TF*>(col_gas.ptr()),
                const_cast<TF*>(fminor.ptr()), const_cast<int*>(jeta.ptr()),
                const_cast<int*>(tropo.ptr()), const_cast<int*>(jtemp.ptr()),
                tau_rayleigh.ptr());
    }

    template<typename TF>
    void reorder123x321(
            const Array<TF,3>& data,
            Array<TF,3>& data_out)
    {
        int dim1 = data.dim(1);
        int dim2 = data.dim(2);
        int dim3 = data.dim(3);
        rrtmgp_kernels::reorder_123x321_kernel(
                &dim1, &dim2, &dim3,
                const_cast<TF*>(data.ptr()),
                data_out.ptr());
    }

    template<typename TF>
    void combine_and_reorder_2str(
            int ncol, int nlay, int ngpt,
            const Array<TF,3>& tau_local, const Array<TF,3>& tau_rayleigh,
            Array<TF,3>& tau, Array<TF,3>& ssa, Array<TF,3>& g)
    {
        rrtmgp_kernels::combine_and_reorder_2str(
                &ncol, &nlay, &ngpt,
                const_cast<TF*>(tau_local.ptr()), const_cast<TF*>(tau_rayleigh.ptr()),
                tau.ptr(), ssa.ptr(), g.ptr());
    }

    template<typename TF>
    void compute_Planck_source(
            int ncol, int nlay, int nbnd, int ngpt,
            int nflav, int neta, int npres, int ntemp, int nPlanckTemp,
            const Array<TF,2>& tlay, const Array<TF,2>& tlev, const Array<TF,1>& tsfc, int sfc_lay,
            const Array<TF,6>& fmajor, const Array<int,4>& jeta, const Array<int,2>& tropo, const Array<int,2>& jtemp, const Array<int,2>& jpress,
            const Array<int,1>& gpoint_bands, const Array<int,2>& band_lims_gpt, const Array<TF,4>& pfracin, TF temp_ref_min,
            TF totplnk_delta, const Array<TF,2>& totplnk, const Array<int,2>& gpoint_flavor,
            Array<TF,2>& sfc_src, Array<TF,3>& lay_src, Array<TF,3>& lev_src_inc, Array<TF,3>& lev_src_dec,
            Array<TF,2>& sfc_src_jac)
    {
        rrtmgp_kernels::compute_Planck_source(
                &ncol, &nlay, &nbnd, &ngpt,
                &nflav, &neta, &npres, &ntemp, &nPlanckTemp,
                const_cast<TF*>(tlay.ptr()),
                const_cast<TF*>(tlev.ptr()),
                const_cast<TF*>(tsfc.ptr()),
                &sfc_lay,
                const_cast<TF*>(fmajor.ptr()),
                const_cast<int*>(jeta.ptr()),
                const_cast<int*>(tropo.ptr()),
                const_cast<int*>(jtemp.ptr()),
                const_cast<int*>(jpress.ptr()),
                const_cast<int*>(gpoint_bands.ptr()), const_cast<int*>(band_lims_gpt.ptr()), const_cast<TF*>(pfracin.ptr()), &temp_ref_min,
                &totplnk_delta, const_cast<TF*>(totplnk.ptr()), const_cast<int*>(gpoint_flavor.ptr()),
                sfc_src.ptr(), lay_src.ptr(), lev_src_inc.ptr(), lev_src_dec.ptr(),
                sfc_src_jac.ptr());
    }
}

template<typename TF>
void Gas_optics<TF>::compute_gas_taus(
        const int ncol, const int nlay, const int ngpt, const int nband,
        const Array<TF,2>& play,
        const Array<TF,2>& plev,
        const Array<TF,2>& tlay,
        const Gas_concs<TF>& gas_desc,
        std::unique_ptr<Optical_props_arry<TF>>& optical_props,
        Array<int,2>& jtemp, Array<int,2>& jpress,
        Array<int,4>& jeta,
        Array<int,2>& tropo,
        Array<TF,6>& fmajor,
        const Array<TF,2>& col_dry) const
{
    Array<TF,3> tau({ngpt, nlay, ncol});
    Array<TF,3> tau_rayleigh({ngpt, nlay, ncol});
    Array<TF,3> vmr({ncol, nlay, this->get_ngas()});
    Array<TF,3> col_gas({ncol, nlay, this->get_ngas()+1});
    col_gas.set_offsets({0, 0, -1});
    Array<TF,4> col_mix({2, this->get_nflav(), ncol, nlay});
    Array<TF,5> fminor({2, 2, this->get_nflav(), ncol, nlay});

    // CvH add all the checking...
    const int ngas = this->get_ngas();
    const int nflav = this->get_nflav();
    const int neta = this->get_neta();
    const int npres = this->get_npres();
    const int ntemp = this->get_ntemp();

    const int nminorlower = this->minor_scales_with_density_lower.dim(1);
    const int nminorklower = this->kminor_lower.dim(1);
    const int nminorupper = this->minor_scales_with_density_upper.dim(1);
    const int nminorkupper = this->kminor_upper.dim(1);

    for (int igas=1; igas<=ngas; ++igas)
    {
        const Array<TF,2>& vmr_2d = gas_desc.get_vmr(this->gas_names({igas}));

        // Fill array with constant value.
        if (vmr_2d.dim(1) == 1 && vmr_2d.dim(2) == 1)
        {
            const TF vmr_c = vmr_2d({1, 1});
            for (int ilay=1; ilay<=nlay; ++ilay)
                for (int icol=1; icol<=ncol; ++icol)
                    vmr({icol, ilay, igas}) = vmr_c;
        }
        // Fill array with constant profile.
        else if (vmr_2d.dim(1) == 1)
        {
            for (int ilay=1; ilay<=nlay; ++ilay)
            {
                const TF vmr_lay = vmr_2d({1, ilay});
                for (int icol=1; icol<=ncol; ++icol)
                    vmr({icol, ilay, igas}) = vmr_lay;
            }
        }
        // Fill array with full 2d data.
        else
        {
            for (int ilay=1; ilay<=nlay; ++ilay)
                for (int icol=1; icol<=ncol; ++icol)
                    vmr({icol, ilay, igas}) = vmr_2d({icol, ilay});
        }
    }

    // CvH: Assume that col_dry is provided.
    for (int ilay=1; ilay<=nlay; ++ilay)
        for (int icol=1; icol<=ncol; ++icol)
            col_gas({icol, ilay, 0}) = col_dry({icol, ilay});

    for (int igas=1; igas<=ngas; ++igas)
        for (int ilay=1; ilay<=nlay; ++ilay)
            for (int icol=1; icol<=ncol; ++icol)
                col_gas({icol, ilay, igas}) = vmr({icol, ilay, igas}) * col_dry({icol, ilay});

    // Call the fortran kernels
    rrtmgp_kernel_launcher::zero_array(ngpt, nlay, ncol, tau);

    rrtmgp_kernel_launcher::interpolation(
            ncol, nlay,
            ngas, nflav, neta, npres, ntemp,
            this->flavor,
            this->press_ref_log,
            this->temp_ref,
            this->press_ref_log_delta,
            this->temp_ref_min,
            this->temp_ref_delta,
            this->press_ref_trop_log,
            this->vmr_ref,
            play,
            tlay,
            col_gas,
            jtemp,
            fmajor, fminor,
            col_mix,
            tropo,
            jeta, jpress);

    int idx_h2o = -1;
    for (int i=1; i<=this->gas_names.dim(1); ++i)
        if (gas_names({i}) == "h2o")
        {
            idx_h2o = i;
            break;
        }

    if (idx_h2o == -1)
        throw std::runtime_error("idx_h2o cannot be found");

    rrtmgp_kernel_launcher::compute_tau_absorption(
            ncol, nlay, nband, ngpt,
            ngas, nflav, neta, npres, ntemp,
            nminorlower, nminorklower,
            nminorupper, nminorkupper,
            idx_h2o,
            this->gpoint_flavor,
            this->get_band_lims_gpoint(),
            this->kmajor,
            this->kminor_lower,
            this->kminor_upper,
            this->minor_limits_gpt_lower,
            this->minor_limits_gpt_upper,
            this->minor_scales_with_density_lower,
            this->minor_scales_with_density_upper,
            this->scale_by_complement_lower,
            this->scale_by_complement_upper,
            this->idx_minor_lower,
            this->idx_minor_upper,
            this->idx_minor_scaling_lower,
            this->idx_minor_scaling_upper,
            this->kminor_start_lower,
            this->kminor_start_upper,
            tropo,
            col_mix, fmajor, fminor,
            play, tlay, col_gas,
            jeta, jtemp, jpress,
            tau);

    bool has_rayleigh = (this->krayl.size() > 0);

    if (has_rayleigh)
    {
        rrtmgp_kernel_launcher::compute_tau_rayleigh(
                ncol, nlay, nband, ngpt,
                ngas, nflav, neta, npres, ntemp,
                this->gpoint_flavor,
                this->get_band_lims_gpoint(),
                this->krayl,
                idx_h2o, col_dry, col_gas,
                fminor, jeta, tropo, jtemp,
                tau_rayleigh);
    }

    combine_and_reorder(tau, tau_rayleigh, has_rayleigh, optical_props);
}

template<typename TF>
void Gas_optics<TF>::combine_and_reorder(
        const Array<TF,3>& tau,
        const Array<TF,3>& tau_rayleigh,
        const bool has_rayleigh,
        std::unique_ptr<Optical_props_arry<TF>>& optical_props) const
{
    int ncol = tau.dim(3);
    int nlay = tau.dim(2);
    int ngpt = tau.dim(1);

    if (!has_rayleigh)
    {
        // CvH for 2 stream and n-stream zero the g and ssa
        rrtmgp_kernel_launcher::reorder123x321(tau, optical_props->get_tau());
        // reorder123x321_test(optical_props->get_tau().ptr(), tau.ptr(), ngpt, nlay, ncol);

        // rrtmgp_kernel_launcher::zero_array(ngpt, nlay, ncol, optical_props->get_ssa());
        // rrtmgp_kernel_launcher::zero_array(ngpt, nlay, ncol, optical_props->get_g  ());
    }
    else
    {
        // In case of 1scl type
        // rrtmgp_kernel_launcher::reorder123x321(tau, optical_props->get_tau());

        // In case of 2str type
        rrtmgp_kernel_launcher::combine_and_reorder_2str(
                ncol, nlay, ngpt,
                tau, tau_rayleigh,
                optical_props->get_tau(), optical_props->get_ssa(), optical_props->get_g());
    }
}

template<typename TF>
void Gas_optics<TF>::source(
        const int ncol, const int nlay, const int nbnd, const int ngpt,
        const Array<TF,2>& play, const Array<TF,2>& plev,
        const Array<TF,2>& tlay, const Array<TF,1>& tsfc,
        const Array<int,2>& jtemp, const Array<int,2>& jpress,
        const Array<int,4>& jeta, const Array<int,2>& tropo,
        const Array<TF,6>& fmajor,
        Source_func_lw<TF>& sources,
        const Array<TF,2>& tlev) const
{
    // CvH Assume tlev is available.
    // Compute internal (Planck) source functions at layers and levels,
    // which depend on mapping from spectral space that creates k-distribution.
    const int nflav = this->get_nflav();
    const int neta = this->get_neta();
    const int npres = this->get_npres();
    const int ntemp = this->get_ntemp();
    const int nPlanckTemp = this->get_nPlanckTemp();
    auto gpoint_bands = this->get_gpoint_bands();
    auto band_lims_gpoint = this->get_band_lims_gpoint();

    Array<TF,3> lay_source_t({ngpt, nlay, ncol});
    Array<TF,3> lev_source_inc_t({ngpt, nlay, ncol});
    Array<TF,3> lev_source_dec_t({ngpt, nlay, ncol});
    Array<TF,2> sfc_source_t({ngpt, ncol});
    Array<TF,2> sfc_source_jac({ngpt, ncol});

    int sfc_lay = play({1, 1}) > play({1, nlay}) ? 1 : nlay;
    rrtmgp_kernel_launcher::compute_Planck_source(
            ncol, nlay, nbnd, ngpt,
            nflav, neta, npres, ntemp, nPlanckTemp,
            tlay, tlev, tsfc, sfc_lay,
            fmajor, jeta, tropo, jtemp, jpress,
            gpoint_bands, band_lims_gpoint, this->planck_frac, this->temp_ref_min,
            this->totplnk_delta, this->totplnk, this->gpoint_flavor,
            sfc_source_t, lay_source_t, lev_source_inc_t, lev_source_dec_t,
            sfc_source_jac);

    // CvH this transpose is super slow.
    for (int j=1; j<=sfc_source_t.dim(2); ++j)
        for (int i=1; i<=sfc_source_t.dim(1); ++i)
        {
            sources.get_sfc_source    ()({j, i}) = sfc_source_t  ({i, j});
            sources.get_sfc_source_jac()({j, i}) = sfc_source_jac({i, j});
        }

    rrtmgp_kernel_launcher::reorder123x321(lay_source_t, sources.get_lay_source());
    rrtmgp_kernel_launcher::reorder123x321(lev_source_inc_t, sources.get_lev_source_inc());
    rrtmgp_kernel_launcher::reorder123x321(lev_source_dec_t, sources.get_lev_source_dec());
    // reorder123x321_test(sources.get_lay_source    ().ptr(), lay_source_t    .ptr(), ngpt, nlay, ncol);
    // reorder123x321_test(sources.get_lev_source_inc().ptr(), lev_source_inc_t.ptr(), ngpt, nlay, ncol);
    // reorder123x321_test(sources.get_lev_source_dec().ptr(), lev_source_dec_t.ptr(), ngpt, nlay, ncol);
}

#ifdef FLOAT_SINGLE_RRTMGP
template class Gas_optics<float>;
#else
template class Gas_optics<double>;
#endif


#ifndef STAN_MATH_REV_MAT_FUN_QUAD_FORM_HPP
#define STAN_MATH_REV_MAT_FUN_QUAD_FORM_HPP

#include <boost/utility/enable_if.hpp>
#include <boost/type_traits.hpp>
#include <stan/math/rev/core.hpp>
#include <stan/math/prim/mat/fun/Eigen.hpp>
#include <stan/math/prim/mat/fun/typedefs.hpp>
#include <stan/math/rev/mat/fun/typedefs.hpp>
#include <stan/math/prim/mat/fun/value_of.hpp>
#include <stan/math/prim/mat/fun/quad_form.hpp>
#include <stan/math/prim/mat/err/check_multiplicable.hpp>
#include <stan/math/prim/mat/err/check_square.hpp>
#include <stan/math/prim/mat/err/check_symmetric.hpp>

namespace stan {
namespace math {

namespace {
template <typename Ta, int Ra, int Ca, typename Tb, int Rb, int Cb>
class quad_form_vari_alloc : public chainable_alloc {
 private:
  inline void compute(const Eigen::Matrix<double, Ra, Ca>& A,
                      const Eigen::Matrix<double, Rb, Cb>& B) {
    Eigen::Matrix<double, Cb, Cb> Cd(B.transpose() * A * B);
    for (int j = 0; j < C_.cols(); j++) {
      for (int i = 0; i < C_.rows(); i++) {
        if (sym_) {
          C_(i, j) = var(new vari(0.5 * (Cd(i, j) + Cd(j, i)), false));
        } else {
          C_(i, j) = var(new vari(Cd(i, j), false));
        }
      }
    }
  }

 public:
  quad_form_vari_alloc(const Eigen::Matrix<Ta, Ra, Ca>& A,
                       const Eigen::Matrix<Tb, Rb, Cb>& B,
                       bool symmetric = false)
      : A_(A), B_(B), C_(B_.cols(), B_.cols()), sym_(symmetric) {
    compute(value_of(A), value_of(B));
  }

  Eigen::Matrix<Ta, Ra, Ca> A_;
  Eigen::Matrix<Tb, Rb, Cb> B_;
  Eigen::Matrix<var, Cb, Cb> C_;
  bool sym_;
};

template <typename Ta, int Ra, int Ca, typename Tb, int Rb, int Cb>
class quad_form_vari : public vari {
 protected:
  inline void chainA(Eigen::Matrix<double, Ra, Ca>& A,
                     const Eigen::Matrix<double, Rb, Cb>& Bd,
                     const Eigen::Matrix<double, Cb, Cb>& adjC) {}
  inline void chainB(Eigen::Matrix<double, Rb, Cb>& B,
                     const Eigen::Matrix<double, Ra, Ca>& Ad,
                     const Eigen::Matrix<double, Rb, Cb>& Bd,
                     const Eigen::Matrix<double, Cb, Cb>& adjC) {}

  inline void chainA(Eigen::Matrix<var, Ra, Ca>& A,
                     const Eigen::Matrix<double, Rb, Cb>& Bd,
                     const Eigen::Matrix<double, Cb, Cb>& adjC) {
    Eigen::Matrix<double, Ra, Ca> adjA(Bd * adjC * Bd.transpose());
    for (int j = 0; j < A.cols(); j++) {
      for (int i = 0; i < A.rows(); i++) {
        A(i, j).vi_->adj_ += adjA(i, j);
      }
    }
  }
  inline void chainB(Eigen::Matrix<var, Rb, Cb>& B,
                     const Eigen::Matrix<double, Ra, Ca>& Ad,
                     const Eigen::Matrix<double, Rb, Cb>& Bd,
                     const Eigen::Matrix<double, Cb, Cb>& adjC) {
    Eigen::Matrix<double, Ra, Ca> adjB(Ad * Bd * adjC.transpose()
                                       + Ad.transpose() * Bd * adjC);
    for (int j = 0; j < B.cols(); j++)
      for (int i = 0; i < B.rows(); i++)
        B(i, j).vi_->adj_ += adjB(i, j);
  }

  inline void chainAB(Eigen::Matrix<Ta, Ra, Ca>& A,
                      Eigen::Matrix<Tb, Rb, Cb>& B,
                      const Eigen::Matrix<double, Ra, Ca>& Ad,
                      const Eigen::Matrix<double, Rb, Cb>& Bd,
                      const Eigen::Matrix<double, Cb, Cb>& adjC) {
    chainA(A, Bd, adjC);
    chainB(B, Ad, Bd, adjC);
  }

 public:
  quad_form_vari(const Eigen::Matrix<Ta, Ra, Ca>& A,
                 const Eigen::Matrix<Tb, Rb, Cb>& B, bool symmetric = false)
      : vari(0.0) {
    impl_ = new quad_form_vari_alloc<Ta, Ra, Ca, Tb, Rb, Cb>(A, B, symmetric);
  }

  virtual void chain() {
    Eigen::Matrix<double, Cb, Cb> adjC(impl_->C_.rows(), impl_->C_.cols());

    for (int j = 0; j < impl_->C_.cols(); j++)
      for (int i = 0; i < impl_->C_.rows(); i++)
        adjC(i, j) = impl_->C_(i, j).vi_->adj_;

    chainAB(impl_->A_, impl_->B_, value_of(impl_->A_), value_of(impl_->B_),
            adjC);
  }

  quad_form_vari_alloc<Ta, Ra, Ca, Tb, Rb, Cb>* impl_;
};
}  // namespace

template <typename Ta, int Ra, int Ca, typename Tb, int Rb, int Cb>
inline typename boost::enable_if_c<boost::is_same<Ta, var>::value
                                       || boost::is_same<Tb, var>::value,
                                   Eigen::Matrix<var, Cb, Cb> >::type
quad_form(const Eigen::Matrix<Ta, Ra, Ca>& A,
          const Eigen::Matrix<Tb, Rb, Cb>& B) {
  check_square("quad_form", "A", A);
  check_multiplicable("quad_form", "A", A, "B", B);

  quad_form_vari<Ta, Ra, Ca, Tb, Rb, Cb>* baseVari
      = new quad_form_vari<Ta, Ra, Ca, Tb, Rb, Cb>(A, B);

  return baseVari->impl_->C_;
}

template <typename Ta, int Ra, int Ca, typename Tb, int Rb>
inline typename boost::enable_if_c<
    boost::is_same<Ta, var>::value || boost::is_same<Tb, var>::value, var>::type
quad_form(const Eigen::Matrix<Ta, Ra, Ca>& A,
          const Eigen::Matrix<Tb, Rb, 1>& B) {
  check_square("quad_form", "A", A);
  check_multiplicable("quad_form", "A", A, "B", B);

  quad_form_vari<Ta, Ra, Ca, Tb, Rb, 1>* baseVari
      = new quad_form_vari<Ta, Ra, Ca, Tb, Rb, 1>(A, B);

  return baseVari->impl_->C_(0, 0);
}

}  // namespace math
}  // namespace stan
#endif


/*
  Tencent is pleased to support the open source community by making
  Plato available.
  Copyright (C) 2019 THL A29 Limited, a Tencent company.
  All rights reserved.

  Licensed under the BSD 3-Clause License (the "License"); you may
  not use this file except in compliance with the License. You may
  obtain a copy of the License at

  https://opensource.org/licenses/BSD-3-Clause

  Unless required by applicable law or agreed to in writing, software
  distributed under the License is distributed on an "AS IS" basis,
  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
  implied. See the License for the specific language governing
  permissions and limitations under the License.

  See the AUTHORS file for names of contributors.
*/

#ifndef __PLATO_ALGO_BNC_BADER_HPP__
#define __PLATO_ALGO_BNC_BADER_HPP__

#include <cstdint>
#include <cstdlib>
#include <unordered_map>
#include <algorithm>
#include <utility>
#include <vector>
#include <unordered_set>
#include <boost/format.hpp>
#include <boost/lockfree/queue.hpp>

#include "glog/logging.h"

#include "plato/util/perf.hpp"
#include "plato/util/atomic.hpp"
#include "plato/graph/graph.hpp"
#include "plato/engine/dualmode.hpp"

#include "cgm.hpp"

namespace plato { namespace algo {
/**
 * @brief bader option
 */
struct bader_opts_t {
  vid_t max_iteration_ = 0;
  int chosen_ = -1;
  float constant_ = 2;
};

/**
 * @brief
 * @tparam INCOMING
 * @tparam OUTGOING
 * @tparam T
 */
template <typename INCOMING, typename OUTGOING, typename T>
class bader_betweenness_t {
public:
  using partition_t = typename dualmode_detail::partition_traits<INCOMING, OUTGOING>::type;
  using betweenness_state_t  = dense_state_t<T, partition_t>;
  using label_state_t = dense_state_t<vid_t, partition_t>;
  using active_subset_t = bitmap_t<>;
  struct bader_msg_type_t {
    vid_t src_;
    T value_;
  };

  /**
   * @brief
   * @param engine
   * @param graph_info
   * @param opts
   */
  explicit bader_betweenness_t (
    dualmode_engine_t<INCOMING, OUTGOING> * engine,
    const graph_info_t &graph_info, const bader_opts_t& opts = bader_opts_t());
  /// \brief
  ~bader_betweenness_t();

  // return the betweenness centrality of a given vertex
  T get_betweenness_of(vid_t v_i) const;

  // compute the betweenness of all vertices
  void compute();
  /**
   * @brief
   * @tparam Callback
   * @param streams
   */
  template <typename Callback>
  void save(Callback&& callback);

  /**
   * @brief
   * @return
   */
  vid_t get_chosen() const { return chosen_; }

  /**
   * @brief
   * @return
   */
  std::unordered_set<vid_t>* get_samples() { return &samples_; }

  /**
   * @brief
   * @return
   */
  vid_t get_major_componnent_vertices() const { return major_component_vertices_; }

private:
  /**
   * @brief
   * @param root
   * @param post_proc
   */
  void epoch(vid_t root, std::function<void(const betweenness_state_t*, vid_t)> post_proc);

private:
  dualmode_engine_t<INCOMING, OUTGOING> * engine_;
  graph_info_t graph_info_;
  betweenness_state_t betweenness_;
  betweenness_state_t num_paths_;
  betweenness_state_t dependencies_;
  T sum_dependence_;
  T sum_dependence_max_;
  vid_t chosen_;
  float constant_;
  vid_t max_iteration_;
  std::unordered_set<vid_t> samples_;
  connected_component_t<INCOMING, OUTGOING> *cc_;
  label_state_t global_labels_;
  vid_t major_component_label_;
  vid_t major_component_vertices_;
  active_subset_t active_all_;
};

template <typename INCOMING, typename OUTGOING, typename T>
bader_betweenness_t<INCOMING, OUTGOING, T>::bader_betweenness_t(
  dualmode_engine_t<INCOMING, OUTGOING>* engine, const graph_info_t &graph_info,
  const bader_opts_t& opts)
  : engine_(engine),
    graph_info_(graph_info),
    betweenness_(graph_info.max_v_i_, engine->in_edges()->partitioner()),
    num_paths_(graph_info.max_v_i_, engine->in_edges()->partitioner()),
    dependencies_(graph_info.max_v_i_, engine->in_edges()->partitioner()),
    sum_dependence_(0),
    global_labels_(graph_info.max_v_i_, engine->in_edges()->partitioner()),
    active_all_(graph_info.max_v_i_ + 1) {
  cc_ = new connected_component_t<INCOMING, OUTGOING>(engine, graph_info);
  cc_->compute();
  major_component_label_ = cc_->get_major_label();
  major_component_vertices_ = cc_->get_major_vertices();
  constant_ = opts.constant_;
  sum_dependence_max_ = constant_ * major_component_vertices_;
  if (opts.max_iteration_ == 0) {
    max_iteration_ = major_component_vertices_;
  } else {
    max_iteration_ = opts.max_iteration_;
  }

  active_all_.fill();

  //init label
  auto local_label = cc_->get_labels();

  auto active_view_all = plato::create_active_v_view(engine_->out_edges()->partitioner()->self_v_view(), active_all_);
  active_view_all.template foreach<vid_t>([&](vid_t v_i){
    global_labels_[v_i] = (*local_label)[v_i];
    return 0;
  });

  auto& cluster_info = plato::cluster_info_t::get_instance();
  if (cluster_info.partition_id_ != 0) {
    vid_t v_begin = engine_->out_edges()->partitioner()->offset_[cluster_info.partition_id_];
    vid_t v_end = engine_->out_edges()->partitioner()->offset_[cluster_info.partition_id_+1];
    MPI_Send(&(global_labels_[v_begin]), v_end - v_begin, get_mpi_data_type<vid_t>(), 0, 0, MPI_COMM_WORLD);
  }
  else {
    for (int i = 1; i < cluster_info.partitions_; ++i ) {
      MPI_Status recv_status;
      vid_t v_begin = engine_->out_edges()->partitioner()->offset_[i];
      vid_t v_end = engine_->out_edges()->partitioner()->offset_[i+1];
      MPI_Recv(&(global_labels_[v_begin]), v_end - v_begin, get_mpi_data_type<vid_t>(), i, 0, MPI_COMM_WORLD, &recv_status);
    }
  }

  //get chosen vertex
  if (opts.chosen_ == -1) {
    if (cluster_info.partition_id_ == 0) {
      do {
        chosen_ = rand() % graph_info.vertices_;
      } while (global_labels_[chosen_] != major_component_label_);
    }
    MPI_Bcast(&chosen_, 1, get_mpi_data_type<vid_t>(), 0, MPI_COMM_WORLD);
    LOG(INFO) << "chosen: "  << chosen_ << std::endl;
  }
  else {
    chosen_ = opts.chosen_;
  }

}

template <typename INCOMING, typename OUTGOING, typename T>
bader_betweenness_t<INCOMING, OUTGOING, T>::~bader_betweenness_t() {
  delete cc_;
}

template <typename INCOMING, typename OUTGOING, typename T>
void bader_betweenness_t<INCOMING, OUTGOING, T>::compute() {
  auto continue_iterating = [&](vid_t iter) {
    return iter < max_iteration_ && sum_dependence_ < sum_dependence_max_;
  };

  // pick a new random vertex that is not used before
  // we use a set to record the vertices that are used.
  auto& cluster_info = plato::cluster_info_t::get_instance();
  auto gen_next_vertex = [&]() {
    vid_t nxt;
    if (cluster_info.partition_id_ == 0) {
      size_t curr_size = samples_.size();
      do {
        nxt = std::rand() % graph_info_.vertices_;
        if (global_labels_[nxt] == major_component_label_) {
          samples_.insert(nxt);
        }
      } while (curr_size == samples_.size());
    }
    MPI_Bcast(&nxt, 1, get_mpi_data_type<vid_t>(), 0, MPI_COMM_WORLD);
    if (cluster_info.partition_id_ != 0) samples_.insert(nxt);
    return nxt;
  };
  auto active_view_all = plato::create_active_v_view(engine_->out_edges()->partitioner()->self_v_view(), active_all_);
  auto accumulate = [&](const betweenness_state_t* dependencies, vid_t root) {
    active_view_all.template foreach<vid_t>([&](vid_t v_i){
      if (v_i != root) {
        T d = (*dependencies)[v_i];
        betweenness_[v_i] += d;
        if (v_i == chosen_) {
          sum_dependence_ += d;
        }
      }
      return 1;
    });

    int chosen_partition = engine_->out_edges()->partitioner()->get_partition_id(chosen_);
    MPI_Bcast(&sum_dependence_, 1, get_mpi_data_type<T>(), chosen_partition,MPI_COMM_WORLD);
    LOG(INFO) << "sum_deppendece: " << sum_dependence_ << " sum_dependence_max: " << sum_dependence_max_ << std::endl;
  };

  auto update_betweenness = [&](vid_t sample_size) {
    active_view_all.template foreach<vid_t>([&](vid_t v_i) {
      betweenness_[v_i] = major_component_vertices_ * betweenness_[v_i] / sample_size;
      if (std::isnan(betweenness_[v_i])) {
        betweenness_[v_i] = 0;
      }
      return 1;
    });
  };

  for (vid_t iter = 0; continue_iterating(iter); ++iter) {
    auto root = gen_next_vertex();
    LOG(INFO) << "next root: " << root << std::endl;
    epoch(root, accumulate);
    LOG_IF(INFO, !cluster_info.partition_id_ && (iter + 1) % 10 == 0)
    << boost::format("iter=%u/%u, sum_dependence[%u]=%.1f/%.1f\n") % (iter + 1) % max_iteration_ % chosen_ % sum_dependence_ % sum_dependence_max_;
  }

  vid_t sample_size = samples_.size();
  MPI_Bcast(&sample_size, 1, get_mpi_data_type<vid_t>(), 0, MPI_COMM_WORLD);

  update_betweenness(sample_size);
}

template <typename INCOMING, typename OUTGOING, typename T>
T bader_betweenness_t<INCOMING, OUTGOING, T>::get_betweenness_of(vid_t v_i) const {
  auto& cluster_info = plato::cluster_info_t::get_instance();
  int vtx_partition = engine_->out_edges()->partitioner()->get_partition_id(v_i);
  T ret = cluster_info.partition_id_ == vtx_partition ? betweenness_[v_i] : 0;
  MPI_Bcast(&ret, 1, get_mpi_data_type<T>(), vtx_partition, MPI_COMM_WORLD);
  return ret;
}

template <typename INCOMING, typename OUTGOING, typename T>
void bader_betweenness_t<INCOMING, OUTGOING, T>::epoch(
  vid_t root, std::function<void(const betweenness_state_t*, vid_t)> post_proc) {
  std::vector<active_subset_t*> levels;
  auto active_current = new active_subset_t(graph_info_.max_v_i_ + 1);
  auto visited = engine_->alloc_v_subset();
  vid_t actives = 1;
  visited.clear();
  visited.set_bit(root);
  active_current->clear();
  active_current->set_bit(root);
  levels.push_back(active_current);
  num_paths_.fill((T)0.0);
  num_paths_[root] = 1.0;

  using adj_unit_list_spec_t = typename INCOMING::adj_unit_list_spec_t;
  using pull_context_t = plato::template mepa_ag_context_t<bader_msg_type_t>;
  using pull_message_t = plato::template mepa_ag_message_t<bader_msg_type_t>;
  using push_context_t = plato::template mepa_bc_context_t<bader_msg_type_t>;
  for (int epoch_i = 0; actives > 0; ++epoch_i) {
    auto active_next = new active_subset_t(graph_info_.max_v_i_ + 1);
    active_next->clear();
    engine_->template foreach_edges<bader_msg_type_t, vid_t> (
      [&](const push_context_t& context, vid_t v_i) {
        context.send(bader_msg_type_t{ v_i, num_paths_[v_i] });
      },
      [&](int /*p_i*/, bader_msg_type_t& msg) {
        auto neighbours = engine_->out_edges()->neighbours(msg.src_);
        for (auto it = neighbours.begin_; neighbours.end_ != it; ++it) {
          vid_t dst = it->neighbour_;
          if (!visited.get_bit(dst)) {
            if (num_paths_[dst] == 0) {
              active_next->set_bit(dst);
            }
            write_add(&(num_paths_[dst]), msg.value_);
          }
        }
        return 0;
      },
      [&](const pull_context_t& context, vid_t v_i, const adj_unit_list_spec_t& adjs) {
        if (visited.get_bit(v_i)) return;
        T sum = 0;
        for (auto it = adjs.begin_; adjs.end_ != it; ++it) {
          vid_t src = it->neighbour_;
          if (active_current->get_bit(src)) {
            sum += num_paths_[src];
          }
        }
        if (sum > 0) {
          context.send(pull_message_t{ v_i, bader_msg_type_t{ v_i, sum } });
        }
      },
      [&](int, pull_message_t& msg) {
        if (!visited.get_bit(msg.v_i_)) {
          active_next->set_bit(msg.v_i_);
          write_add(&(num_paths_[msg.v_i_]), msg.message_.value_);
        }
        return 0;
      },
      *active_current
    );
    auto active_view = plato::create_active_v_view(engine_->out_edges()->partitioner()->self_v_view(), *active_next);
    actives = active_view.template foreach<vid_t> ([&](vid_t v_i) {
      visited.set_bit(v_i); return 1;
    });
    if (actives > 0) levels.push_back(active_next);
    active_current = active_next;
  }

  dependencies_.fill((T)0.0);
  visited.clear();
  while (levels.size() > 0) {
    auto active_view = plato::create_active_v_view(engine_->out_edges()->partitioner()->self_v_view(), *(levels.back()));
    actives = active_view.template foreach<vid_t> ([&](vid_t v_i) {
      visited.set_bit(v_i); return 1;
    });
    engine_->template foreach_edges<bader_msg_type_t, vid_t> (
      [&](const push_context_t& context, vid_t v_i) {
        T value = (dependencies_[v_i] + 1.0) / num_paths_[v_i];
        context.send(bader_msg_type_t{ v_i, value });
      },
      [&](int /*p_i*/, bader_msg_type_t& msg) {
        auto neighbours = engine_->out_edges()->neighbours(msg.src_);
        for (auto it = neighbours.begin_; neighbours.end_ != it; ++it) {
          vid_t dst = it->neighbour_;
          if (!visited.get_bit(dst)) {
            write_add(&(dependencies_[dst]), msg.value_ * num_paths_[dst]);
          }
        }
        return 0;
      },
      [&](const pull_context_t& context, vid_t v_i, const adj_unit_list_spec_t& adjs) {
        if (visited.get_bit(v_i)) return;
        T sum = 0;
        for (auto it = adjs.begin_; adjs.end_ != it; ++it) {
          vid_t src = it->neighbour_;
          if (levels.back()->get_bit(src)) {
            sum += (dependencies_[src] + 1.0) / num_paths_[src];
          }
        }
        if (sum > 0) {
          context.send(pull_message_t{ v_i, bader_msg_type_t{ v_i, sum } });
        }
      },
      [&](int, pull_message_t& msg) {
        if (!visited.get_bit(msg.v_i_)) {
          write_add(&(dependencies_[msg.v_i_]), msg.message_.value_ * num_paths_[msg.v_i_]);
        }
        return 0;
      },
      *(levels.back())
    );
    delete levels.back();
    levels.pop_back();
  }

  post_proc(&dependencies_, root);
}

template <typename INCOMING, typename OUTGOING, typename T>
template <typename Callback>
void bader_betweenness_t<INCOMING, OUTGOING, T>::save(Callback&& callback) {
  // traverse
  auto active_view_all = plato::create_active_v_view(engine_->out_edges()->partitioner()->self_v_view(), active_all_);
  active_view_all.template foreach<vid_t>([&] (vid_t v_i) {
    callback(v_i, betweenness_[v_i]);
    return 1;
  });
}

}
}
#endif


/**
 * @file
 * @copyright This code is licensed under the 3-clause BSD license.\n
 *            Copyright ETH Zurich, Laboratory of Physical Chemistry, Reiher Group.\n
 *            See LICENSE.txt for details.
 */

#include <Utils/GeometricDerivatives/AdiabaticModeLocalizer.h>
#include <pybind11/eigen.h>
#include <pybind11/pybind11.h>
#include <pybind11/stl.h>
#include <boost/optional.hpp>

using namespace Scine::Utils;

void init_adiabatic_mode_localizer(pybind11::module& m) {
  pybind11::class_<AdiabaticModesContainer> adiabaticModesContainer(m, "AdiabaticModesContainer", R"delim(
    A class storing the localized adiabatic properties.
    This class is an extension of the standard normal mode container, with modes being accessed by specifying the
    internal coordinate of interest.
    Note: Currently only interatomic stretching coordinates, e.g., bonds, are supported as internal coordinates.
    )delim");

  adiabaticModesContainer.def("add_mode", &AdiabaticModesContainer::addMode, pybind11::arg("internal_coordinate"),
                              pybind11::arg("mode"), pybind11::arg("force_constant"), R"delim(
    Adds a mode to the container.

    :param internal_coordinate: The internal coordinate to which this mode corresponds.
    :param mode: The mode to be added.
    :param force_constant: The adiabatic force constant of the mode.
    )delim");

  adiabaticModesContainer.def("get_mode", &AdiabaticModesContainer::getMode, pybind11::arg("internal_coordinate"), R"delim(
    Returns the localized vibrational mode corresponding to the given internal coordinate.

    :param: internal_coordinate The internal coordinate of interest.
    :return: The localized mode.
    )delim");

  adiabaticModesContainer.def("get_mode_as_molecular_trajectory", &AdiabaticModesContainer::getModeAsMolecularTrajectory,
                              pybind11::arg("internal_coordinate"), pybind11::arg("atoms"),
                              pybind11::arg("scaling_factor"), R"delim(
    Returns a molecular trajectory corresponding to a vibrational mode.

    :param internal_coordinate: The internal coordinate of interest.
    :param atoms: The atom collection of interest.
    :param scaling_factor: The scaling factor applied to the mode to obtain the maximum displacement.
    :return: The molecular trajectory representing the mode.
    )delim");

  adiabaticModesContainer.def("get_wave_numbers", &AdiabaticModesContainer::getWaveNumbers, R"delim(
    Gets the wave numbers corresponding to the localized modes.

    :return: Mapping between internal coordinates and adiabatic wave numbers.
    )delim");

  adiabaticModesContainer.def("get_force_constants", &AdiabaticModesContainer::getForceConstants, R"delim(
    Gets the force constants corresponding to the localized modes.

    :return: Mapping between internal coordinates and adiabatic force constants.
    )delim");

  pybind11::class_<AdiabaticModeLocalizer> adiabaticModeLocalizer(m, "AdiabaticModeLocalizer", R"delim(
    This class calculates the localized vibrational modes and adiabatic force constants corresponding to internal
    coordinates based on the local vibrational mode theory by Konkoli and Cremer. Adiabatic force constants are
    equivalent to the relaxed force constants derived from Decius' compliance matrix.

    Konkoli, Z.; Cremer, D. Int. J. Quantum Chem. 1998, 67 (1), 1–9.
    https://doi.org/10.1002/(SICI)1097-461X(1998)67:1<1::AID-QUA1>3.0.CO;2-Z.

    Implemented after:
    Kraka, E.; Zou, W.; Tao, Y. Wiley Interdiscip. Rev. Comput. Mol. Sci. 2020, e1480. https://doi.org/10.1002/wcms.1480.

    Note: Currently only interatomic stretching coordinates, e.g., bonds, are supported as internal coordinates.
    )delim");

  adiabaticModeLocalizer.def(pybind11::init<Eigen::MatrixXd, AtomCollection, BondOrderCollection, double>(),
                             pybind11::arg("hessian"), pybind11::arg("atoms"), pybind11::arg("bond_orders"),
                             pybind11::arg("bonding_threshold") = 0.5, R"delim(
    Constructor for the AdiabaticModeLocalizer class processing a bond order collection.
    All bonds with a bond order larger than bonding_threshold will be used as internal coordinates.

    :param hessian: The cartesian, not mass-weighted Hessian matrix.
    :param atoms: The atom collection of interest.
    :param bond_orders: The bond order collection of the structure of interest.
    :param: bonding_threshold The bond order threshold beyond which bonds are analyzed in terms of localized modes. (default: 0.5)
    )delim");

  adiabaticModeLocalizer.def(pybind11::init<Eigen::MatrixXd, AtomCollection, const std::vector<std::pair<int, int>>>(),
                             pybind11::arg("hessian"), pybind11::arg("atoms"), pybind11::arg("internal_coordinates"), R"delim(
    Constructor for the AdiabaticModeLocalizer class processing indices of atom pairs as internal coordinates.

    :param hessian: The cartesian, not mass-weighted Hessian matrix.
    :param atoms: The atom collection of interest.
    :param internal_coordinates:  The internal coordinates of interest.
    )delim");

  adiabaticModeLocalizer.def("localize_modes", &AdiabaticModeLocalizer::localizeModes, R"delim(
    Calculates the adiabatic localized modes corresponding to the given internal coordinates.

    References:
    Konkoli, Z.; Cremer, D. Int. J. Quantum Chem. 1998, 67 (1), 1–9.
    https://doi.org/10.1002/(SICI)1097-461X(1998)67:1<1::AID-QUA1>3.0.CO;2-Z.

    Kraka, E.; Zou, W.; Tao, Y. Wiley Interdiscip. Rev. Comput. Mol. Sci. 2020, e1480.
    https://doi.org/10.1002/wcms.1480.

    :return: A container with the localized adiabatic modes, force constants and wavenumbers.
    )delim");
}


#pragma once
#include <trajopt_utils/macros.h>
TRAJOPT_IGNORE_WARNINGS_PUSH
#include <Eigen/Dense>
#include <functional>
#include <memory>
TRAJOPT_IGNORE_WARNINGS_POP

/*
 * Numerical derivatives
 */

namespace sco
{
class ScalarOfVector;
class VectorOfVector;
class MatrixOfVector;
typedef std::shared_ptr<ScalarOfVector> ScalarOfVectorPtr;
typedef std::shared_ptr<VectorOfVector> VectorOfVectorPtr;
typedef std::shared_ptr<MatrixOfVector> MatrixOfVectorPtr;

class ScalarOfVector
{
public:
  virtual double operator()(const Eigen::VectorXd& x) const = 0;
  double call(const Eigen::VectorXd& x) const
  {
    return operator()(x);
  }
  virtual ~ScalarOfVector()
  {
  }
  typedef std::function<double(Eigen::VectorXd)> func;
  static ScalarOfVectorPtr construct(const func&);
  //  typedef VectorXd (*c_func)(const VectorXd&);
  //  static ScalarOfVectorPtr construct(const c_func&);
};
class VectorOfVector
{
public:
  virtual Eigen::VectorXd operator()(const Eigen::VectorXd& x) const = 0;
  Eigen::VectorXd call(const Eigen::VectorXd& x) const
  {
    return operator()(x);
  }
  virtual ~VectorOfVector()
  {
  }
  typedef std::function<Eigen::VectorXd(Eigen::VectorXd)> func;
  static VectorOfVectorPtr construct(const func&);
  //  typedef VectorXd (*c_func)(const VectorXd&);
  //  static VectorOfVectorPtr construct(const c_func&);
};
class MatrixOfVector
{
public:
  virtual Eigen::MatrixXd operator()(const Eigen::VectorXd& x) const = 0;
  Eigen::MatrixXd call(const Eigen::VectorXd& x) const
  {
    return operator()(x);
  }
  virtual ~MatrixOfVector()
  {
  }
  typedef std::function<Eigen::MatrixXd(Eigen::VectorXd)> func;
  static MatrixOfVectorPtr construct(const func&);
  //  typedef VectorMatrixXd (*c_func)(const VectorXd&);
  //  static MatrixOfVectorPtr construct(const c_func&);
};

Eigen::VectorXd calcForwardNumGrad(const ScalarOfVector& f, const Eigen::VectorXd& x, double epsilon);
Eigen::MatrixXd calcForwardNumJac(const VectorOfVector& f, const Eigen::VectorXd& x, double epsilon);
void calcGradAndDiagHess(const ScalarOfVector& f, const Eigen::VectorXd& x, double epsilon, double& y,
                         Eigen::VectorXd& grad, Eigen::VectorXd& hess);
void calcGradHess(ScalarOfVectorPtr f, const Eigen::VectorXd& x, double epsilon, double& y, Eigen::VectorXd& grad,
                  Eigen::MatrixXd& hess);
VectorOfVectorPtr forwardNumGrad(ScalarOfVectorPtr f, double epsilon);
MatrixOfVectorPtr forwardNumJac(VectorOfVectorPtr f, double epsilon);
}


/**
    anest.cc
    Purpose: Perform the Displacement Method of Analysis
    for structural analysis.

    @author Adolfo Correa
    @version 0.1 02/03/20
*/

#include<iostream>
#include<vector>
#include <armadillo>
#include <fstream>
#include <string>

using namespace std;
using namespace arma;

void preparq(int, int, int, int, int);
// void nos(int nno, int nnv, vector<double> x);

void nos(int nno, int nnv, vector<double> x, \
vector<double> y, vector<vector<double>> fa, \
vector<vector<int>> d, vector<vector<int>> v, \
vector<int> nv, int gln);

void testing_armadillo(void);

int main(int argc, char* argv[])
{
	int nno, nel, ngl, nnv, gln, nelt, nelp, \
	nll, ntc = nno = nel = ngl = nnv = gln = nelt \
	= nelp = nll = 0;

	vector<int> nv, te, nc;

	/// Vector 2D
	vector<vector<int>> no, g, v, d, il;
	vector<double> x, y, e, a, iz, ra, fp, f, u, \
	qx, qy;
	/// Vector 2D
	vector<vector<double>> fa, sg, esf;

/// Read input files
	preparq(nno, nel, ntc, ngl, gln);

	nno = nel = 4; ntc = 3; gln = 3; ngl = gln * nno;

// Dimensionar vetores e matrizes
// matrix.resize(M, vector<int>(N, default_value));

/**
* \brief Initialize the variables
*/
	d.resize(nno, vector<int>(gln, 0));
	v.resize(nno, vector<int>(gln, 0));
	no.resize(2, vector<int>(nel, 0));
	g.resize(nel, vector<int>(2*gln, 0));
	il.resize(nel, vector<int>(2*gln, 0));

	fa.resize(nno, vector<double>(gln, 0));
	esf.resize(6, vector<double>(nel, 0));
	x.resize(nno, 0);
	y.resize(nno, 0);
	nv.resize(nno, 0);
	te.resize(nel, 0);
	qx.resize(nel, 0);
	qy.resize(nel, 0);
	e.resize(ntc, 0);
	a.resize(ntc, 0);
	iz.resize(ntc, 0);
	fp.resize(nel, 0);
	nc.resize(nel, 0);

//	nos(nno, nnv, x);
	nos(nno, nnv, x, y, fa, d, v, nv, gln);

	/// Testing Armadillo library
	testing_armadillo();

	return 0;
}

void preparq(int nno, int nel, int ntc, int ngl, int gln)
{
	cout<<"Reading input files"<<endl;
	ifstream file;
	ofstream auxFile;
	string line, filename = "exemplo";
	file.open(filename);
	auxFile.open("anest.dat");

	// Read lines of input file
	for(int i = 0; !file.eof(); i++)
	{
		//file.get(line);
		getline(file, line);
		cout<<line<<endl;
		if(line.substr(0,1) != ";") auxFile<<line<<endl;
	}

	file.close();
	auxFile.close();

}

/**
 * Read the nodes coordinates of the structure from
 * the input file.
 *
 * @param[in] nno number of nodes
 * @param[in] gln degree of freedom per node
 * @param[out] x x-coordinates of the nodes
 * @param[out] y y-coordinates of the nodes
 */
void nos(int nno, int nnv, vector<double> x, \
vector<double> y, vector<vector<double>> fa, \
vector<vector<int>> d, vector<vector<int>> v, \
vector<int> nv, int gln)
{
	cout<<"Reading nos"<<endl;
}

/**
 * Create two matrices with random values
 * and multiply them
 *
 * @param void.
 * @return none, only print on the screen.
 */
void testing_armadillo(void)
{
	cout<<endl<<"Testing armadillo::mat"<<endl<<endl;

	mat A(4, 5, fill::randu); ///< matrix A filled with random unsigned integers
	mat B(4, 5, fill::randn); ///< matrix B filled with randoms integers

	cout<<"A = "<<A<<endl;
	cout<<"B = "<<B<<endl;

	/**
	  \f[
  	    C = A \cross B^T
	  \f]
	*/
	cout<<"A*B.t() = "<<A*B.t()<<endl; ///< Plot the product of the two matrices

}


/**************************************************************************\
 *   This file is part of CaSPER.                                          *
 *                                                                         *
 *   Copyright:                                                            *
 *   2009-2009 - Marco Correia <marco.v.correia@gmail.com>           *
 *                                                                         *
 *   Licensed under the Apache License, Version 2.0 (the "License");       *
 *   you may not use this file except in compliance with the License.      *
 *   You may obtain a copy of the License at                               *
 *            http://www.apache.org/licenses/LICENSE-2.0                   *
 *   Unless required by applicable law or agreed to in writing, software   *
 *   distributed under the License is distributed on an                    *
 *   "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND,          *
 *   either express or implied.                                            *
 *   See the License for the specific language governing permissions and   *
 *   limitations under the License.                                        *
 \*************************************************************************/

#include <boost/xpressive/xpressive.hpp>
#include <libxml++/libxml++.h>
#include <bindings/cpp/cpp.h>
#include <bindings/cpp/print.h>

#include <iostream>

#include "exprparser.h"

using namespace boost::xpressive;
using namespace regex_constants;
using namespace Casperbind::cpp;

ExprParser::ExprParser()
{
	integer	= +digit;
	boolean	= as_xpr("true") | "false";
	id		= alpha >> *alnum;

	eNeg = "neg(" >> by_ref(integerExpr) >> ")";
	eAbs = "abs(" >> by_ref(integerExpr) >> ")";
	eAdd = "add(" >> by_ref(integerExpr) >> ","
									 >> by_ref(integerExpr) >> ")";
	eSub	= "sub(" >> by_ref(integerExpr) >> ","
									 >> by_ref(integerExpr) >> ")";
	eMul	= "mul(" >> by_ref(integerExpr) >> ","
									 >> by_ref(integerExpr) >> ")";
	eDiv	= "div(" >> by_ref(integerExpr) >> ","
									 >> by_ref(integerExpr) >> ")";
	eMod	= "mod(" >> by_ref(integerExpr) >> ","
									 >> by_ref(integerExpr) >> ")";
	ePow	= "pow(" >> by_ref(integerExpr) >> ","
									 >> by_ref(integerExpr) >> ")";
	eMin	= "min(" >> by_ref(integerExpr) >> ","
									 >> by_ref(integerExpr) >> ")";
	eMax	= "max(" >> by_ref(integerExpr) >> ","
									 >> by_ref(integerExpr) >> ")";
	eIf	= "if(" >>  by_ref(boolExpr) >> ","
									>> by_ref(integerExpr) >> ","
									>> by_ref(integerExpr)>> ")";

	integerExpr = integer | id | eNeg | eAbs | eAdd | eSub
			| eMul | eDiv | eMod | ePow | eMin | eMax | eIf;

	pNot= "not(" >> by_ref(boolExpr) >> ")";
	pAnd= "and(" >> by_ref(boolExpr) >> "," >> by_ref(boolExpr) >> ")";
	pOr	= "or(" >> by_ref(boolExpr) >> "," >> by_ref(boolExpr) >> ")";
	pXOr= "xor(" >> by_ref(boolExpr) >> ","	>> by_ref(boolExpr) >> ")";
	pIff= "iff(" >> by_ref(boolExpr) >> ","	>> by_ref(boolExpr) >> ")";
	pEq	= "eq(" >> integerExpr >> "," >> integerExpr >> ")";
	pNe	= "ne(" >> integerExpr >> "," >> integerExpr >> ")";
	pGe	= "ge(" >> integerExpr >> "," >> integerExpr >> ")";
	pGt	= "gt(" >> integerExpr >> "," >> integerExpr >> ")";
	pLe	= "le(" >> integerExpr >> "," >> integerExpr >> ")";
	pLt	= "lt(" >> integerExpr >> "," >> integerExpr >> ")";

	boolExpr = boolean | pNot | pAnd | pOr | pXOr | pIff | pEq
			| pNe | pGe | pGt | pLe | pLt;
}

SharedSymbol ExprParser::parse(const std::string& s,
		const std::map<std::string,casperbind::cpp::SharedSymbol>& mpars) const
{
	smatch what;
    if (!regex_match( s, what, boolExpr))
      throw xmlpp::validity_error("expected boolean expression or predicate : "+s);

    assert(what.nested_results().size()==1);

    return parseBoolExpression(what,mpars);
}

SharedSymbol ExprParser::parseBoolExpression(const smatch& what,
		const std::map<std::string,casperbind::cpp::SharedSymbol>& mpars) const
{
	assert(what.nested_results().size()==1);

	smatch::nested_results_type::const_iterator begin = what.nested_results().begin();
	smatch::nested_results_type::const_iterator end = what.nested_results().end();

	if (begin->regex_id() == boolean.regex_id())
	{
		if (begin->str() == "true")
			return int(1);
		else
		if (begin->str() == "false")
			return int(0);
		assert(0);
	}
	else
	if (begin->regex_id() == pNot.regex_id())
	{
		SymbolArray pars(1);
		pars[0] = parseBoolExpression(*begin->nested_results().begin(),mpars);
		return Predicate(Predicate::pNot,pars);
	}
	else
	if (begin->regex_id() == pAnd.regex_id() or
		begin->regex_id() == pOr.regex_id() or
		begin->regex_id() == pXOr.regex_id() or
		begin->regex_id() == pIff.regex_id())
	{
		SymbolArray pars(2);
		pars[0] = parseBoolExpression(*begin->nested_results().begin(),mpars);
		pars[1] = parseBoolExpression(*++begin->nested_results().begin(),mpars);
		if (begin->regex_id() == pAnd.regex_id())
			return Predicate(Predicate::pAnd,pars);
		else
		if (begin->regex_id() == pOr.regex_id())
			return Predicate(Predicate::pOr,pars);
		else
		if (begin->regex_id() == pXOr.regex_id())
			return Predicate(Predicate::pXOr,pars);
		else
		if (begin->regex_id() == pIff.regex_id())
			return Predicate(Predicate::pIff,pars);
	}
	else
	if (begin->regex_id() == pEq.regex_id() or
		begin->regex_id() == pNe.regex_id() or
		begin->regex_id() == pGe.regex_id() or
		begin->regex_id() == pGt.regex_id() or
		begin->regex_id() == pLe.regex_id() or
		begin->regex_id() == pLt.regex_id())
	{
		SymbolArray pars(2);
		pars[0] = parseIntExpression(*begin->nested_results().begin(),mpars);
		pars[1] = parseIntExpression(*++begin->nested_results().begin(),mpars);
		if (begin->regex_id() == pEq.regex_id())
			return Predicate(Predicate::pEqual,pars);
		else
		if (begin->regex_id() == pNe.regex_id())
			return Predicate(Predicate::pDistinct,pars);
		else
		if (begin->regex_id() == pGe.regex_id())
			return Predicate(Predicate::pGreaterEqual,pars);
		else
		if (begin->regex_id() == pGt.regex_id())
			return Predicate(Predicate::pGreater,pars);
		else
		if (begin->regex_id() == pLe.regex_id())
			return Predicate(Predicate::pLessEqual,pars);
		else
		if (begin->regex_id() == pLt.regex_id())
			return Predicate(Predicate::pLess,pars);
	}
	assert(0);
}

SharedSymbol ExprParser::parseIntExpression(const smatch& what,
		const std::map<std::string,casperbind::cpp::SharedSymbol>& mpars) const
{
	assert(what.nested_results().size()==1);

	smatch::nested_results_type::const_iterator begin = what.nested_results().begin();
	smatch::nested_results_type::const_iterator end = what.nested_results().end();

	assert(what.nested_results().size()==1);

	if (begin->regex_id() == integer.regex_id())
		return int(atoi(begin->str().c_str()));
	else
	if (begin->regex_id() == id.regex_id())
	{
		std::map<std::string,casperbind::cpp::SharedSymbol>::const_iterator
				it = mpars.find(begin->str());
		if (it==mpars.end())
			throw xmlpp::validity_error("undeclared variable appearing in expression : "+begin->str());
		return it->second;
	}
	else
	if (begin->regex_id() == eNeg.regex_id() or
		begin->regex_id() == eAbs.regex_id())
	{
		SymbolArray pars(1);
		pars[0] = parseIntExpression(*begin->nested_results().begin(),mpars);
		if (begin->regex_id() == eNeg.regex_id())
			return Expression(Expression::eSym,pars);
		else
		if (begin->regex_id() == eAbs.regex_id())
			return Expression(Expression::eAbs,pars);
	}
	else
	if (begin->regex_id() == eAdd.regex_id() or
		begin->regex_id() == eSub.regex_id() or
		begin->regex_id() == eMul.regex_id() or
		begin->regex_id() == eDiv.regex_id() or
		begin->regex_id() == eMod.regex_id() or
		begin->regex_id() == ePow.regex_id() or
		begin->regex_id() == eMin.regex_id() or
		begin->regex_id() == eMax.regex_id())
	{
		SymbolArray pars(2);
		pars[0] = parseIntExpression(*begin->nested_results().begin(),mpars);
		pars[1] = parseIntExpression(*++begin->nested_results().begin(),mpars);
		if (begin->regex_id() == eAdd.regex_id())
			return Expression(Expression::eAdd,pars);
		else
		if (begin->regex_id() == eSub.regex_id())
			return Expression(Expression::eSub,pars);
		else
		if (begin->regex_id() == eMul.regex_id())
			return Expression(Expression::eMul,pars);
		else
		if (begin->regex_id() == eDiv.regex_id())
			return Expression(Expression::eDiv,pars);
		else
		if (begin->regex_id() == eMod.regex_id())
			return Expression(Expression::eMod,pars);
		else
		if (begin->regex_id() == ePow.regex_id())
			return Expression(Expression::ePow,pars);
		else
		if (begin->regex_id() == eMin.regex_id())
			return Expression(Expression::eMin,pars);
		else
		if (begin->regex_id() == eMax.regex_id())
			return Expression(Expression::eMax,pars);
	}
	else
	if (begin->regex_id() == eIf.regex_id())
	{
		SymbolArray pars(3);
		pars[0] = parseBoolExpression(*begin->nested_results().begin(),mpars);
		pars[1] = parseIntExpression(*++begin->nested_results().begin(),mpars);
		pars[2] = parseIntExpression(*++(++begin->nested_results().begin()),mpars);
		return Expression(Expression::eIfThenElse,pars);
	}

	assert(0);
}

#if 0
int main(int argc, char** argv)
{
	std::map<std::string,casperbind::cpp::SharedSymbol> mpars;
	mpars["V0"] = casperbind::cpp::Variable(casperbind::cpp::IntDomain(casperbind::cpp::IntRange(1,10)));
	mpars["V1"] = casperbind::cpp::Variable(casperbind::cpp::IntDomain(casperbind::cpp::IntRange(2,4)));

	std::string s(argv[1]);
	std::cout << "input string is [" << s << "]" << std::endl;
	ExprParser e;
	SharedSymbol o = e.parse(s,mpars);
	std::cout << o << std::endl;
}
#endif


/*
 * Copyright (c) 2020 Andrew Price
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *
 * 1. Redistributions of source code must retain the above copyright notice,
 *    this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 * 3. Neither the name of the copyright holder nor the names of its
 *    contributors may be used to endorse or promote products derived from
 *    this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
 * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
 * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
 * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
 * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
 * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
 * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
 * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 * POSSIBILITY OF SUCH DAMAGE.
 */

#include <Eigen/Eigenvalues>
#include <Eigen/SparseCore>
//#include <Spectra/GenEigsSolver.h>
//#include <Spectra/MatOp/SparseGenMatProd.h>
//#include <Spectra/SymEigsSolver.h>
//#include <Spectra/MatOp/SparseSymMatProd.h>

#include "mps_voxels/video_graph.h"
#include "mps_voxels/graph_matrix_utils.h"
#include <mps_msgs/SegmentGraph.h>

#include <boost/graph/graph_traits.hpp>
#include <boost/graph/connected_components.hpp>
#include <boost/graph/graphviz.hpp>

#include <ros/ros.h>

template <class WeightMap>
class edge_writer {
public:
	edge_writer(WeightMap w) : wm(w) {}
	template <class Edge>
	void operator()(std::ostream &out, const Edge& e) const {
		out << "[penwidth=\"" << wm[e]*4.0 << "\"]";
	}
private:
	WeightMap wm;
};

template <class WeightMap>
inline edge_writer<WeightMap>
make_edge_writer(WeightMap w) {
	return edge_writer<WeightMap>(w);
}

std::string node_name(const NodeProperties& np) {return "f" + std::to_string(np.t.toSec()) + "c" + std::to_string(np.leafID);}

void print(std::ostream& out, const VideoSegmentationGraph& G)
{
	out << "graph G {\n";

	std::map<ros::Time, std::vector<VideoSegmentationGraph::vertex_descriptor>> frames;
	for (VideoSegmentationGraph::vertex_descriptor vd : make_range(boost::vertices(G)))
	{
		frames[G[vd].t].push_back(vd);
	}

	for (const auto& frame : frames)
	{
		out << "subgraph cluster_" << frame.first << " {\n";
		out << "label=\"Frame " << frame.first << "\";\n";
		for (const VideoSegmentationGraph::vertex_descriptor vd : frame.second)
		{
			const NodeProperties& np = G[vd];
			out << node_name(np) << " [label=\"" << np.leafID << "\"];\n";
		}
		out << "}\n";
	}

	for (VideoSegmentationGraph::edge_descriptor ed : make_range(boost::edges(G)))
	{
		VideoSegmentationGraph::vertex_descriptor u = boost::source(ed, G);
		VideoSegmentationGraph::vertex_descriptor v = boost::target(ed, G);

		const EdgeProperties& ep = G[ed];
		const NodeProperties& npu = G[u];
		const NodeProperties& npv = G[v];
		out << node_name(npu) << "--" << node_name(npv) << " [penwidth=\"" << 4.0*ep.affinity << "\"];\n";
	}

	out << "}\n";
}

int main(int argc, char* argv[])
{
	ros::init(argc, argv, "test_graph_segmentation");
	ros::NodeHandle nh, pnh("~");

	ros::ServiceClient segmentClient = nh.serviceClient<mps_msgs::SegmentGraph>("/segment_graph");
	if (!segmentClient.waitForExistence(ros::Duration(3)))
	{
		ROS_ERROR("Segmentation server not connected.");
		return -1;
	}

	VideoSegmentationGraph G;

	std::map<std::pair<ros::Time, int>, VideoSegmentationGraph::vertex_descriptor> segmentToNode;
	std::vector<NodeProperties> nps{NodeProperties({ros::Time(0), 0}), NodeProperties({ros::Time(0), 1}), NodeProperties({ros::Time(0), 2}), NodeProperties({ros::Time(1), 0}), NodeProperties({ros::Time(1), 1})};
	for (const auto& np : nps)
	{
		// Add leaf i to graph
		// NB: labels start from 1, as 0 is background, so tree[0] = labels[1]
		VideoSegmentationGraph::vertex_descriptor v = boost::add_vertex(np, G);
		segmentToNode.insert({{np.t, np.leafID}, v});
	}

	boost::add_edge(segmentToNode[{ros::Time(0), 0}], segmentToNode[{ros::Time(0), 1}], {1.0/1.0}, G);
	boost::add_edge(segmentToNode[{ros::Time(0), 0}], segmentToNode[{ros::Time(0), 2}], {1.0/10.0}, G);
	boost::add_edge(segmentToNode[{ros::Time(0), 1}], segmentToNode[{ros::Time(0), 2}], {1.0/10.0}, G);
	boost::add_edge(segmentToNode[{ros::Time(1), 0}], segmentToNode[{ros::Time(1), 1}], {1.0/15.0}, G);
	boost::add_edge(segmentToNode[{ros::Time(0), 2}], segmentToNode[{ros::Time(1), 1}], {1.0/2.0}, G);
	boost::add_edge(segmentToNode[{ros::Time(0), 0}], segmentToNode[{ros::Time(1), 0}], {1.0/2.0}, G);
	boost::add_edge(segmentToNode[{ros::Time(0), 1}], segmentToNode[{ros::Time(1), 0}], {1.0/2.0}, G);

	print(std::cerr, G);//, boost::make_label_writer(boost::get(&NodeProperties::leafID, G)), make_edge_writer(boost::get(&EdgeProperties::affinity, G)));

	{
		Eigen::MatrixXd laplacian = getLaplacianNormalized(G);
		const int numCells = laplacian.rows();
		std::cerr << "Segmenting subgraph with " << numCells << " vertices (out of " << boost::num_vertices(G) << ")."
		          << std::endl;
		Eigen::SelfAdjointEigenSolver<Eigen::MatrixXd> solver;
		solver.compute(laplacian, Eigen::ComputeEigenvectors);
		Eigen::VectorXd eigenvalues = solver.eigenvalues();
		Eigen::MatrixXd eigenvectors = solver.eigenvectors();
		std::cerr << eigenvalues.transpose() << std::endl;
		std::cerr << eigenvectors << std::endl;
	}

//	{
//		Eigen::SparseMatrix<double> laplacian = getLaplacianSparseNormalized(G);
//		using SparseMatProd = Spectra::SparseGenMatProd<double>; // Spectra::SparseSymMatProd<double>
//		SparseMatProd op(laplacian);
//		Spectra::GenEigsSolver<double, Spectra::LARGEST_MAGN, SparseMatProd> eigs(&op, 3, 5);
//		eigs.init();
//		int nconv = eigs.compute();
//		if (eigs.info() == Spectra::SUCCESSFUL)
//		{
//			std::cerr << nconv << "\n" << eigs.eigenvectors() << std::endl;
//		}
//		else
//		{
//			ROS_WARN("Failed to compute eigenvectors. Clusterign aborted.");
//		}
//	}

	{
		Eigen::SparseMatrix<double> adj = getAdjacencySparse(G);
		mps_msgs::SegmentGraphRequest req;
		for (const auto& triplet : to_triplets(adj))
		{
			std::cerr << triplet.row() << "," << triplet.col() << "," << triplet.value() << std::endl;
			req.adjacency.row_index.push_back(triplet.row());
			req.adjacency.col_index.push_back(triplet.col());
			req.adjacency.value.push_back(triplet.value());
		}
		req.num_labels = 2;

		for (const std::string& alg : std::vector<std::string>{"spectral", "dbscan"})
		{
			req.algorithm = alg;

			mps_msgs::SegmentGraphResponse resp;
			segmentClient.call(req, resp);
			std::cerr << resp << std::endl;

		}
	}

	return 0;
}