#ifndef PATCHWORK_H
#define PATCHWORK_H

#include <Eigen/Dense>
#include <boost/format.hpp>

#define MARKER_Z_VALUE -2.2
#define UPRIGHT_ENOUGH 0.55
#define FLAT_ENOUGH 0.2
#define TOO_HIGH_ELEVATION 0.0
#define TOO_TILTED 1.0

#define NUM_HEURISTIC_MAX_PTS_IN_PATCH 3000

using Eigen::MatrixXf;
using Eigen::JacobiSVD;
using Eigen::VectorXf;

using namespace std;

static int64_t gtm() {
	struct timeval tm;
	gettimeofday(&tm, 0);
	// return ms
	int64_t re = (((int64_t) tm.tv_sec) * 1000 * 1000 + tm.tv_usec);
	return re;
}

/*
    @brief PathWork ROS Node.
*/
template <typename PointT>
bool point_z_cmp(PointT a, PointT b) {
    return a.z < b.z;
}

template <typename PointT>
class PatchWork {

public:
    typedef std::vector<pcl::PointCloud<PointT> > Ring;
    typedef std::vector<Ring> Zone;

    PatchWork(){

	num_sectors_each_zone_ = {16, 32 ,54, 32};
	num_rings_each_zone_ =  {2, 4, 4, 4};
	elevation_thr_ = {-1.2, -0.9984, -0.851, -0.605};
	flatness_thr_ = {0.0, 0.000125, 0.000185, 0.000185};

        if (num_zones_ != 4 || num_sectors_each_zone_.size() != num_rings_each_zone_.size()) {
            throw invalid_argument("Some parameters are wrong! Check the num_zones and num_rings/sectors_each_zone");
        }
        if (elevation_thr_.size() != flatness_thr_.size()) {
            throw invalid_argument("Some parameters are wrong! Check the elevation/flatness_thresholds");
        }

        cout << (boost::format("Num. sectors: %d, %d, %d, %d") % num_sectors_each_zone_[0] % num_sectors_each_zone_[1] %
                 num_sectors_each_zone_[2] %
                 num_sectors_each_zone_[3]).str() << endl;
        cout << (boost::format("Num. rings: %01d, %01d, %01d, %01d") % num_rings_each_zone_[0] %
                 num_rings_each_zone_[1] %
                 num_rings_each_zone_[2] %
                 num_rings_each_zone_[3]).str() << endl;
        cout << (boost::format("elevation_thr_: %0.4f, %0.4f, %0.4f, %0.4f ") % elevation_thr_[0] % elevation_thr_[1] %
                 elevation_thr_[2] %
                 elevation_thr_[3]).str() << endl;
        cout << (boost::format("flatness_thr_: %0.4f, %0.4f, %0.4f, %0.4f ") % flatness_thr_[0] % flatness_thr_[1] %
                 flatness_thr_[2] %
                 flatness_thr_[3]).str() << endl;
        num_rings_of_interest_ = elevation_thr_.size();

        revert_pc.reserve(NUM_HEURISTIC_MAX_PTS_IN_PATCH);
        ground_pc_.reserve(NUM_HEURISTIC_MAX_PTS_IN_PATCH);
        non_ground_pc_.reserve(NUM_HEURISTIC_MAX_PTS_IN_PATCH);
        regionwise_ground_.reserve(NUM_HEURISTIC_MAX_PTS_IN_PATCH);
        regionwise_nonground_.reserve(NUM_HEURISTIC_MAX_PTS_IN_PATCH);

        min_range_z2_ = (7 * min_range_ + max_range_) / 8.0;
        min_range_z3_ = (3 * min_range_ + max_range_) / 4.0;
        min_range_z4_ = (min_range_ + max_range_) / 2.0;

        min_ranges = {min_range_, min_range_z2_, min_range_z3_, min_range_z4_};
        ring_sizes = {(min_range_z2_ - min_range_) / num_rings_each_zone_.at(0),
                      (min_range_z3_ - min_range_z2_) / num_rings_each_zone_.at(1),
                      (min_range_z4_ - min_range_z3_) / num_rings_each_zone_.at(2),
                      (max_range_ - min_range_z4_) / num_rings_each_zone_.at(3)};
        sector_sizes = {2 * M_PI / num_sectors_each_zone_.at(0), 2 * M_PI / num_sectors_each_zone_.at(1),
                        2 * M_PI / num_sectors_each_zone_.at(2),
                        2 * M_PI / num_sectors_each_zone_.at(3)};
        cout << "INITIALIZATION COMPLETE" << endl;

        for (int iter = 0; iter < num_zones_; ++iter) {
            Zone z;
            initialize_zone(z, num_sectors_each_zone_.at(iter), num_rings_each_zone_.at(iter));
            ConcentricZoneModel_.push_back(z);
        }
    }

    void estimate_ground(
            const pcl::PointCloud<PointT> &cloudIn,
            pcl::PointCloud<PointT> &cloudOut,
            pcl::PointCloud<PointT> &cloudNonground,
            double &time_taken);

private:

    int num_iter_ = 3;
    int num_lpr_ = 20;
    int num_min_pts_ = 10;
    int num_rings_ = 30;
    int num_sectors_ = 108;
    int num_zones_ = 4;
    int num_rings_of_interest_ = 2.7;

    double sensor_height_ = 1.723;
    double th_seeds_ = 0.5;
    double th_dist_ = 0.125;
    double max_range_ = 80.0;
    double min_range_ = 2.7;
    double uprightness_thr_ = 0.707;
    double adaptive_seed_selection_margin_ = -1.1;
    double min_range_z2_; // 12.3625
    double min_range_z3_; // 22.025
    double min_range_z4_; // 41.35

    bool verbose_ = false;

    float d_;
    MatrixXf normal_;
    VectorXf singular_values_;
    float th_dist_d_;
    Eigen::Matrix3f cov_;
    Eigen::Vector4f pc_mean_;
    double ring_size;
    double sector_size;
    // For visualization
    bool visualize_;

    vector<int> num_sectors_each_zone_;
    vector<int> num_rings_each_zone_;

    vector<double> sector_sizes;
    vector<double> ring_sizes;
    vector<double> min_ranges;
    vector<double> elevation_thr_;
    vector<double> flatness_thr_;

    vector<Zone> ConcentricZoneModel_;

    pcl::PointCloud<PointT> revert_pc, reject_pc;
    pcl::PointCloud<PointT> ground_pc_;
    pcl::PointCloud<PointT> non_ground_pc_;

    pcl::PointCloud<PointT> regionwise_ground_;
    pcl::PointCloud<PointT> regionwise_nonground_;

    void initialize_zone(Zone &z, int num_sectors, int num_rings);

    void flush_patches_in_zone(Zone &patches, int num_sectors, int num_rings);

    double calc_principal_variance(const Eigen::Matrix3f &cov, const Eigen::Vector4f &centroid);

    double xy2theta(const double &x, const double &y);

    double xy2radius(const double &x, const double &y);

    void pc2czm(const pcl::PointCloud<PointT> &src, std::vector<Zone> &czm);

    void estimate_plane_(const pcl::PointCloud<PointT> &ground);

    void extract_piecewiseground(
            const int zone_idx, const pcl::PointCloud<PointT> &src,
            pcl::PointCloud<PointT> &dst,
            pcl::PointCloud<PointT> &non_ground_dst);

    void estimate_plane_(const int zone_idx, const pcl::PointCloud<PointT> &ground);

    void extract_initial_seeds_(
            const int zone_idx, const pcl::PointCloud<PointT> &p_sorted,
            pcl::PointCloud<PointT> &init_seeds);

};

template<typename PointT> inline
void PatchWork<PointT>::initialize_zone(Zone &z, int num_sectors, int num_rings) {
    z.clear();
    pcl::PointCloud<PointT> cloud;
    cloud.reserve(1000);
    Ring ring;
    for (int i = 0; i < num_sectors; i++) {
        ring.emplace_back(cloud);
    }
    for (int j = 0; j < num_rings; j++) {
        z.emplace_back(ring);
    }
}

template<typename PointT> inline
void PatchWork<PointT>::flush_patches_in_zone(Zone &patches, int num_sectors, int num_rings) {
    for (int i = 0; i < num_sectors; i++) {
        for (int j = 0; j < num_rings; j++) {
            if (!patches[j][i].points.empty()) patches[j][i].points.clear();
        }
    }
}

template<typename PointT> inline
void PatchWork<PointT>::estimate_plane_(const pcl::PointCloud<PointT> &ground) {
    pcl::computeMeanAndCovarianceMatrix(ground, cov_, pc_mean_);
    // Singular Value Decomposition: SVD
    Eigen::JacobiSVD<Eigen::MatrixXf> svd(cov_, Eigen::DecompositionOptions::ComputeFullU);
    singular_values_ = svd.singularValues();

    // use the least singular vector as normal
    normal_ = (svd.matrixU().col(2));
    // mean ground seeds value
    Eigen::Vector3f seeds_mean = pc_mean_.head<3>();

    // according to normal.T*[x,y,z] = -d
    d_ = -(normal_.transpose() * seeds_mean)(0, 0);
    // set distance threhold to `th_dist - d`
    th_dist_d_ = th_dist_ - d_;
}

template<typename PointT> inline
void PatchWork<PointT>::extract_initial_seeds_(
        const int zone_idx, const pcl::PointCloud<PointT> &p_sorted,
        pcl::PointCloud<PointT> &init_seeds) {
    init_seeds.points.clear();

    // LPR is the mean of low point representative
    double sum = 0;
    int cnt = 0;

    int init_idx = 0;
    if (zone_idx == 0) {
        for (int i = 0; i < p_sorted.points.size(); i++) {
            if (p_sorted.points[i].z < adaptive_seed_selection_margin_ * sensor_height_) {
                ++init_idx;
            } else {
                break;
            }
        }
    }

    // Calculate the mean height value.
    for (int i = init_idx; i < p_sorted.points.size() && cnt < num_lpr_; i++) {
        sum += p_sorted.points[i].z;
        cnt++;
    }
    double lpr_height = cnt != 0 ? sum / cnt : 0;// in case divide by 0

    // iterate pointcloud, filter those height is less than lpr.height+th_seeds_
    for (int i = 0; i < p_sorted.points.size(); i++) {
        if (p_sorted.points[i].z < lpr_height + th_seeds_) {
            init_seeds.points.push_back(p_sorted.points[i]);
        }
    }
}

/*
    @brief Velodyne pointcloud callback function. The main GPF pipeline is here.
    PointCloud SensorMsg -> Pointcloud -> z-value sorted Pointcloud
    ->error points removal -> extract ground seeds -> ground plane fit mainloop
*/

template<typename PointT> inline
void PatchWork<PointT>::estimate_ground(
        const pcl::PointCloud<PointT> &cloud_in,
        pcl::PointCloud<PointT> &cloud_out,
        pcl::PointCloud<PointT> &cloud_nonground,
        double &time_taken) {

    static double start, t0, t1, t2, end;

    double t_total_ground = 0.0;
    double t_total_estimate = 0.0;
    // 1.Msg to pointcloud
    pcl::PointCloud<PointT> laserCloudIn;
    laserCloudIn = cloud_in;

    start = gtm();

    // 2.Sort on Z-axis value.
    sort(laserCloudIn.points.begin(), laserCloudIn.end(), point_z_cmp<PointT>);

    t0 = gtm();
    // 3.Error point removal
    // As there are some error mirror reflection under the ground,
    // here regardless point under 1.8* sensor_height
    // Sort point according to height, here uses z-axis in default
    auto it = laserCloudIn.points.begin();
    for (int i = 0; i < laserCloudIn.points.size(); i++) {
        if (laserCloudIn.points[i].z < -1.8 * sensor_height_) {
            it++;
        } else {
            break;
        }
    }
    laserCloudIn.points.erase(laserCloudIn.points.begin(), it);

    t1 = gtm();
    // 4. pointcloud -> regionwise setting
    for (int k = 0; k < num_zones_; ++k) {
        flush_patches_in_zone(ConcentricZoneModel_[k], num_sectors_each_zone_[k], num_rings_each_zone_[k]);
    }
    pc2czm(laserCloudIn, ConcentricZoneModel_);

    t2 =gtm();

    cloud_out.clear();
    cloud_nonground.clear();
    revert_pc.clear();
    reject_pc.clear();

    int concentric_idx = 0;
    for (int k = 0; k < num_zones_; ++k) {
        auto zone = ConcentricZoneModel_[k];
        for (uint16_t ring_idx = 0; ring_idx < num_rings_each_zone_[k]; ++ring_idx) {
            for (uint16_t sector_idx = 0; sector_idx < num_sectors_each_zone_[k]; ++sector_idx) {
                if (zone[ring_idx][sector_idx].points.size() > num_min_pts_) {
                    double t_tmp0 = gtm();
                    extract_piecewiseground(k, zone[ring_idx][sector_idx], regionwise_ground_, regionwise_nonground_);
                    double t_tmp1 = gtm();
                    t_total_ground += t_tmp1 - t_tmp0;

                    // Status of each patch
                    // used in checking uprightness, elevation, and flatness, respectively
                    const double ground_z_vec = abs(normal_(2, 0));
                    const double ground_z_elevation = pc_mean_(2, 0);
                    const double surface_variable =
                            singular_values_.minCoeff() /
                            (singular_values_(0) + singular_values_(1) + singular_values_(2));

                    double t_tmp2 = gtm();
                    if (ground_z_vec < uprightness_thr_) {
                        // All points are rejected
                        cloud_nonground += regionwise_ground_;
                        cloud_nonground += regionwise_nonground_;
                    } else { // satisfy uprightness
                        if (concentric_idx < num_rings_of_interest_) {
                            if (ground_z_elevation > elevation_thr_[ring_idx + 2 * k]) {
                                if (flatness_thr_[ring_idx + 2 * k] > surface_variable) {
                                    if (verbose_) {
                                        std::cout << "\033[1;36m[Flatness] Recovery operated. Check "
                                                  << ring_idx + 2 * k
                                                  << "th param. flatness_thr_: " << flatness_thr_[ring_idx + 2 * k]
                                                  << " > "
                                                  << surface_variable << "\033[0m" << std::endl;
                                        revert_pc += regionwise_ground_;
                                    }
                                    cloud_out += regionwise_ground_;
                                    cloud_nonground += regionwise_nonground_;
                                } else {
                                    if (verbose_) {
                                        std::cout << "\033[1;34m[Elevation] Rejection operated. Check "
                                                  << ring_idx + 2 * k
                                                  << "th param. of elevation_thr_: " << elevation_thr_[ring_idx + 2 * k]
                                                  << " < "
                                                  << ground_z_elevation << "\033[0m" << std::endl;
                                        reject_pc += regionwise_ground_;
                                    }
                                    cloud_nonground += regionwise_ground_;
                                    cloud_nonground += regionwise_nonground_;
                                }
                            } else {
                                cloud_out += regionwise_ground_;
                                cloud_nonground += regionwise_nonground_;
                            }
                        } else {
                            cloud_out += regionwise_ground_;
                            cloud_nonground += regionwise_nonground_;
                        }
                    }
                    double t_tmp3 = gtm();
                    t_total_estimate += t_tmp3 - t_tmp2;
                }
            }
            ++concentric_idx;
        }
    }
    end =gtm();
    time_taken = end - start;
//    ofstream time_txt("/home/shapelim/patchwork_time_anal.txt", std::ios::app);
//    time_txt<<t0 - start<<" "<<t1 - t0 <<" "<<t2-t1<<" "<<t_total_ground<< " "<<t_total_estimate<<"\n";
//    time_txt.close();

}

template<typename PointT> inline
double PatchWork<PointT>::calc_principal_variance(const Eigen::Matrix3f &cov, const Eigen::Vector4f &centroid) {
    double angle = atan2(centroid(1, 0), centroid(0, 0)); // y, x
    double c = cos(angle);
    double s = sin(angle);
    double var_x_prime = c * c * cov(0, 0) + s * s * cov(1, 1) + 2 * c * s * cov(0, 1);
    double var_y_prime = s * s * cov(0, 0) + c * c * cov(1, 1) - 2 * c * s * cov(0, 1);
    return max(var_x_prime, var_y_prime);
}

template<typename PointT> inline
double PatchWork<PointT>::xy2theta(const double &x, const double &y) { // 0 ~ 2 * PI
    if (y >= 0) {
        return atan2(y, x); // 1, 2 quadrant
    } else {
        return 2 * M_PI + atan2(y, x);// 3, 4 quadrant
    }
}

template<typename PointT> inline
double PatchWork<PointT>::xy2radius(const double &x, const double &y) {
    return sqrt(pow(x, 2) + pow(y, 2));
}

template<typename PointT> inline
void PatchWork<PointT>::pc2czm(const pcl::PointCloud<PointT> &src, std::vector<Zone> &czm) {

    for (auto const &pt : src.points) {
        int ring_idx, sector_idx;
        double r = xy2radius(pt.x, pt.y);
        if ((r <= max_range_) && (r > min_range_)) {
            double theta = xy2theta(pt.x, pt.y);

            if (r < min_range_z2_) { // In First rings
                ring_idx = min(static_cast<int>(((r - min_range_) / ring_sizes[0])), num_rings_each_zone_[0] - 1);
                sector_idx = min(static_cast<int>((theta / sector_sizes[0])), num_sectors_each_zone_[0] - 1);
                czm[0][ring_idx][sector_idx].points.emplace_back(pt);
            } else if (r < min_range_z3_) {
                ring_idx = min(static_cast<int>(((r - min_range_z2_) / ring_sizes[1])), num_rings_each_zone_[1] - 1);
                sector_idx = min(static_cast<int>((theta / sector_sizes[1])), num_sectors_each_zone_[1] - 1);
                czm[1][ring_idx][sector_idx].points.emplace_back(pt);
            } else if (r < min_range_z4_) {
                ring_idx = min(static_cast<int>(((r - min_range_z3_) / ring_sizes[2])), num_rings_each_zone_[2] - 1);
                sector_idx = min(static_cast<int>((theta / sector_sizes[2])), num_sectors_each_zone_[2] - 1);
                czm[2][ring_idx][sector_idx].points.emplace_back(pt);
            } else { // Far!
                ring_idx = min(static_cast<int>(((r - min_range_z4_) / ring_sizes[3])), num_rings_each_zone_[3] - 1);
                sector_idx = min(static_cast<int>((theta / sector_sizes[3])), num_sectors_each_zone_[3] - 1);
                czm[3][ring_idx][sector_idx].points.emplace_back(pt);
            }
        }

    }
}

// For adaptive
template<typename PointT> inline
void PatchWork<PointT>::extract_piecewiseground(
        const int zone_idx, const pcl::PointCloud<PointT> &src,
        pcl::PointCloud<PointT> &dst,
        pcl::PointCloud<PointT> &non_ground_dst) {
    // 0. Initialization
    if (!ground_pc_.empty()) ground_pc_.clear();
    if (!dst.empty()) dst.clear();
    if (!non_ground_dst.empty()) non_ground_dst.clear();
    // 1. set seeds!

    extract_initial_seeds_(zone_idx, src, ground_pc_);
    // 2. Extract ground
    for (int i = 0; i < num_iter_; i++) {
        estimate_plane_(ground_pc_);
        ground_pc_.clear();

        //pointcloud to matrix
        Eigen::MatrixXf points(src.points.size(), 3);
        int j = 0;
        for (auto &p:src.points) {
            points.row(j++) << p.x, p.y, p.z;
        }
        // ground plane model
        Eigen::VectorXf result = points * normal_;
        // threshold filter
        for (int r = 0; r < result.rows(); r++) {
            if (i < num_iter_ - 1) {
                if (result[r] < th_dist_d_) {
                    ground_pc_.points.push_back(src[r]);
                }
            } else { // Final stage
                if (result[r] < th_dist_d_) {
                    dst.points.push_back(src[r]);
                } else {
                    if (i == num_iter_ - 1) {
                        non_ground_dst.push_back(src[r]);
                    }
                }
            }
        }
    }
}

#endif


/*
 *            Copyright 2009-2017 The VOTCA Development Team
 *                       (http://www.votca.org)
 *
 *      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.
 *
 */

// Overload of uBLAS prod function with MKL/GSL implementations
#include <votca/tools/linalg.h>
#include <votca/xtp/qmiter.h>
#include <sys/stat.h>
#include <boost/algorithm/string.hpp>
#include <boost/format.hpp>
#include <votca/tools/constants.h>
#include <votca/ctp/apolarsite.h>

using boost::format;

namespace votca {
    namespace xtp {

void QMMIter::UpdateMPSFromGDMA(std::vector<std::vector<double> > &multipoles, std::vector< ctp::PolarSeg* > &psegs) {

            for (unsigned int i = 0, qac = 0; i < psegs.size(); ++i) {
                ctp::PolarSeg *pseg = psegs[i];
                for (unsigned int j = 0; j < pseg->size(); ++j, ++qac) {

                    // Retrieve multipole info of this atom
                    std::vector<double> update = multipoles[qac];
                    while (update.size() < 9) update.push_back(0.0);

                    // Convert e*(a_0)^k to e*(nm)^k where k = rank

                    for (int m = 1; m < 4; m++) {
                        update[m] *= pow(tools::conv::bohr2nm, 1);
                    }
                    for (int m = 4; m < 9; m++) {
                        update[m] *= pow(tools::conv::bohr2nm, 2);
                    }

                    ctp::APolarSite *aps = (*pseg)[j];

                    aps->setQs(update, 0);

                }
            }

            return;
        }

        void QMMIter::UpdatePosChrgFromQMAtoms(std::vector< ctp::QMAtom* > &qmatoms,
                std::vector< ctp::PolarSeg* > &psegs) {

            double AA_to_NM = tools::conv::ang2nm;

            double dR_RMS = 0.0;
            double dQ_RMS = 0.0;
            double dQ_SUM = 0.0;

            for (unsigned int i = 0, qac = 0; i < psegs.size(); ++i) {
                ctp::PolarSeg *pseg = psegs[i];
                for (unsigned int j = 0; j < pseg->size(); ++j, ++qac) {

                    // Retrieve info from ctp::QMAtom
                    ctp::QMAtom *qmatm = qmatoms[qac];
                    vec upd_r = qmatm->getPos();
                    upd_r *= AA_to_NM;
                    double upd_Q00 = qmatm->charge;

                    // Compare to previous r, Q00
                    ctp::APolarSite *aps = (*pseg)[j];
                    vec old_r = aps->getPos();
                    double old_Q00 = aps->getQ00();
                    double dR = abs(upd_r - old_r);
                    double dQ00 = upd_Q00 - old_Q00;

                    dR_RMS += dR*dR;
                    dQ_RMS += dQ00*dQ00;
                    dQ_SUM += dQ00;

                    // Forward updated r, Q00 to APS
                    aps->setPos(upd_r);
                    aps->setQ00(upd_Q00, 0);
                }
            }

            // cout << " dR_RMS " << dR_RMS << "dQ RMS " << dQ_RMS << endl;
            dR_RMS /= qmatoms.size();
            dQ_RMS /= qmatoms.size();
            dR_RMS = sqrt(dR_RMS);
            dQ_RMS = sqrt(dQ_RMS);

            // cout << " dR_RMS " << dR_RMS << "dQ RMS " << dQ_RMS << endl;

            this->setdRdQ(dR_RMS, dQ_RMS, dQ_SUM);
            return;
        }

        void QMMIter::setdRdQ(double dR_RMS, double dQ_RMS, double dQ_SUM) {

            _hasdRdQ = true;
            _dR_RMS = dR_RMS;
            _dQ_RMS = dQ_RMS;
            _dQ_SUM = dQ_SUM;
            return;
        }

        void QMMIter::setQMSF(double energy_QM, double energy_SF, double energy_GWBSE) {

            _hasQM = true;
            _e_QM = energy_QM;
            _e_SF = energy_SF;

            _hasGWBSE = true;
            _e_GWBSE = energy_GWBSE;

            return;
        }

        void QMMIter::setE_FM(double ef00, double ef01, double ef02,
                double ef11, double ef12, double em0, double em1, double em2, double efm) {

            _hasMM = true;
            _ef_00 = ef00;
            _ef_01 = ef01;
            _ef_02 = ef02;
            _ef_11 = ef11;
            _ef_12 = ef12;
            _em_0_ = em0;
            _em_1_ = em1;
            _em_2_ = em2;
            _e_fm_ = efm;
            return;
        }

        double QMMIter::getMMEnergy() {

            assert(_hasMM);
            return _ef_11 + _ef_12 + _em_1_ + _em_2_;
        }

        double QMMIter::getQMMMEnergy() {

            assert(_hasQM && _hasMM && _hasGWBSE);
            return _e_QM + + _e_GWBSE + _ef_11 + _ef_12 + _em_1_ + _em_2_;
        }

    }
}


/*=========================================================================

  Program:   Visualization Toolkit
  Module:    vtkPBGLRMATGraphSource.cxx

  Copyright (c) Ken Martin, Will Schroeder, Bill Lorensen
  All rights reserved.
  See Copyright.txt or http://www.kitware.com/Copyright.htm for details.

     This software is distributed WITHOUT ANY WARRANTY; without even
     the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
     PURPOSE.  See the above copyright notice for more information.

=========================================================================*/
/*-------------------------------------------------------------------------
  Copyright 2008 Sandia Corporation.
  Under the terms of Contract DE-AC04-94AL85000 with Sandia Corporation,
  the U.S. Government retains certain rights in this software.
-------------------------------------------------------------------------*/
/*
 * Copyright (C) 2008 The Trustees of Indiana University.
 * Use, modification and distribution is subject to the Boost Software
 * License, Version 1.0. (See http://www.boost.org/LICENSE_1_0.txt)
 */
#include "vtkPBGLRMATGraphSource.h"

#include "vtkBlockDistribution.h"
#include "vtkCellData.h"
#include "vtkExecutive.h"
#include "vtkFloatArray.h"
#include "vtkGraph.h"
#include "vtkIdTypeArray.h"
#include "vtkInformation.h"
#include "vtkMath.h"
#include "vtkMPI.h"
#include "vtkMutableDirectedGraph.h"
#include "vtkMutableUndirectedGraph.h"
#include "vtkObjectFactory.h"
#include "vtkPBGLDistributedGraphHelper.h"
#include "vtkPointData.h"
#include "vtkSmartPointer.h"

#include <vtksys/stl/set>
#include <vtksys/stl/algorithm>
#include <vtksys/stl/functional>

#include <boost/mpi/communicator.hpp>
#include <boost/mpi/collectives/scan.hpp>

vtkStandardNewMacro(vtkPBGLRMATGraphSource);

// ----------------------------------------------------------------------
vtkPBGLRMATGraphSource::vtkPBGLRMATGraphSource()
{
  int rank;
  MPI_Comm_rank(MPI_COMM_WORLD, &rank);
  this->NumberOfVertices = 128;
  this->NumberOfEdges = 512;
  this->A = 0.25;
  this->B = 0.25;
  this->C = 0.25;
  this->D = 0.25;
  this->IncludeEdgeWeights = false;
  this->AllowSelfLoops = false;
  this->GeneratePedigreeIds = true;
  this->VertexPedigreeIdArrayName = 0;
  this->SetVertexPedigreeIdArrayName("vertex id");
  this->EdgePedigreeIdArrayName = 0;
  this->SetEdgePedigreeIdArrayName("edge id");
  this->EdgeWeightArrayName = 0;
  this->SetEdgeWeightArrayName("edge weight");
  this->Seed = 1177 + 17 * rank;
  this->SetNumberOfInputPorts(0);
  this->SetNumberOfOutputPorts(1);
}

// ----------------------------------------------------------------------

vtkPBGLRMATGraphSource::~vtkPBGLRMATGraphSource()
{
  this->SetVertexPedigreeIdArrayName(0);
  this->SetEdgePedigreeIdArrayName(0);
  this->SetEdgeWeightArrayName(0);
}

// ----------------------------------------------------------------------

void
vtkPBGLRMATGraphSource::PrintSelf(ostream& os, vtkIndent indent)
{
  this->Superclass::PrintSelf(os, indent);
  os << indent << "NumberOfVertices: " << this->NumberOfVertices << endl;
  os << indent << "NumberOfEdges: " << this->NumberOfEdges << endl;
  os << indent << "Probabilities: " << this->A << ", " << this->B << ", "
     << this->C << ", " << this->D << endl;
  os << indent << "IncludeEdgeWeights: " << this->IncludeEdgeWeights << endl;
  os << indent << "AllowSelfLoops: " << this->AllowSelfLoops << endl;
  os << indent << "GeneratePedigreeIds: " << this->GeneratePedigreeIds << endl;
  os << indent << "VertexPedigreeIdArrayName: "
    << (this->VertexPedigreeIdArrayName ? this->VertexPedigreeIdArrayName : "(null)") << endl;
  os << indent << "EdgePedigreeIdArrayName: "
    << (this->EdgePedigreeIdArrayName ? this->EdgePedigreeIdArrayName : "(null)") << endl;
  os << indent << "EdgeWeightArrayName: "
    << (this->EdgeWeightArrayName ? this->EdgeWeightArrayName : "(null)") << endl;
  os << indent << "Seed: " << this->Seed << endl;
}

// ----------------------------------------------------------------------
void vtkPBGLRMATGraphSource::SetNumberOfVertices(vtkIdType value)
{
  vtkIdType mask = (vtkIdType) 1 << ((sizeof(vtkIdType) * CHAR_BIT) - 2);
  while (mask != 0)
    {
    if (value & mask)
      {
      // We've found the most significant 1-bit.
      if (value & (mask >> (vtkIdType) 1))
        {
        // Round up to the next power of two.
        mask = mask << (vtkIdType) 1;
        }
      break;
      }

    mask = mask >> (vtkIdType) 1;
    }

  this->NumberOfVertices = mask;
}

// ----------------------------------------------------------------------
void
vtkPBGLRMATGraphSource::SetProbabilities(double A, double B, double C, double D)
{
  if (A + B + C + D != 1.0)
    {
    vtkErrorMacro("R-MAT probabilities do not add up to 1.0.");
    return;
    }

  this->A = A;
  this->B = B;
  this->C = C;
  this->D = D;
}

// ----------------------------------------------------------------------
void
vtkPBGLRMATGraphSource::GetProbabilities(double *A, double *B, double *C, double *D)
{
  if (A)
    {
    *A = this->A;
    }
  if (B)
    {
    *B = this->B;
    }
  if (C)
    {
    *C = this->C;
    }
  if (D)
    {
    *D = this->D;
    }
}

// ----------------------------------------------------------------------

int
vtkPBGLRMATGraphSource::RequestData(
  vtkInformation*,
  vtkInformationVector**,
  vtkInformationVector *outputVector)
{
  int myRank;
  int numProcs;
  MPI_Comm_rank(MPI_COMM_WORLD, &myRank);
  MPI_Comm_size(MPI_COMM_WORLD, &numProcs);

  // Seed the random number generator so we can produce repeatable results
  vtkMath::RandomSeed(this->Seed);

  // Create a mutable, directed graph.
  vtkSmartPointer<vtkMutableDirectedGraph> dirBuilder =
    vtkSmartPointer<vtkMutableDirectedGraph>::New();

  // Create a Parallel BGL distributed graph helper
  vtkSmartPointer<vtkPBGLDistributedGraphHelper> helper
    = vtkSmartPointer<vtkPBGLDistributedGraphHelper>::New();

  // Hook the distributed graph helper into the graph to make it a
  // distributed graph.
  dirBuilder->SetDistributedGraphHelper(helper);

  // Vertex distribution.
  vtkBlockDistribution distribution(this->NumberOfVertices, numProcs);

  // Add NumberOfVertices new vertices.
  vtkIdType myNumberOfVertices = distribution.GetBlockSize(myRank);
  vtkIdType myStartVertex = distribution.GetFirstGlobalIndexOnProcessor(myRank);
  vtkIdType myEndVertex = myStartVertex + myNumberOfVertices;
  for (vtkIdType i = 0; i < myNumberOfVertices; ++i)
    {
    dirBuilder->AddVertex();
    }

  // Make sure everyone has added their own local vertices.
  helper->Synchronize();

  vtkIdType MaxEdges;
  if (this->AllowSelfLoops)
    {
    MaxEdges = this->NumberOfVertices * this->NumberOfVertices;
    }
  else
    {
    MaxEdges = (this->NumberOfVertices * (this->NumberOfVertices-1)) / 2;
    }

  if (this->NumberOfEdges > MaxEdges)
    {
    this->NumberOfEdges = MaxEdges;
    }

  vtkIdType avgNumberOfEdges = this->NumberOfEdges / numProcs;
  vtkIdType myNumberOfEdges = avgNumberOfEdges;
  if (myRank < this->NumberOfEdges % numProcs)
    {
    ++myNumberOfEdges;
    }

  vtkIdType numLevels = (vtkIdType) log2(this->NumberOfVertices) + 1;
  double AB = this->A + this->B;
  double CNorm = this->C / (this->C + this->D);
  double ANorm = this->A / (this->A + this->B);
  for (vtkIdType i = 0; i < myNumberOfEdges; i++)
    {
    bool newEdgeFound = false;
    while (!newEdgeFound)
      {
      vtkIdType s = 0;
      vtkIdType t = 0;

      for (vtkIdType level = 1; level < numLevels; ++level)
        {
        bool sBit = vtkMath::Random() > (this->A + this->B);
        bool tBit
          = (vtkMath::Random() > (CNorm * (sBit ? 1 : 0)
                                  + ANorm * (sBit ? 1 : 0))) ? 1 : 0;
        s |= ((vtkIdType) 1 << (level-1)) * (sBit ? 1 : 0);
        t |= ((vtkIdType) 1 << (level-1)) * (tBit ? 1 : 0);
        }

      if (s == t && (!this->AllowSelfLoops))
        {
        continue;
        }

      // TODO: We should apply some permutation to the s and t vertex
      // numbers, so that we get some kind of randomized distribution
      // of the vertices. Otherwise, we'll have a severely imbalanced
      // distribution, since the high-degree vertices are likely to
      // all end up on the lower-numbered ranks. Note that this
      // doesn't change the block distribution (computed below); it
      // sits on top of the block distribution.

      vtkIdType sVertex
        = helper->MakeDistributedId(distribution.GetProcessorOfElement(s),
                                    distribution.GetLocalIndexOfElement(s));
      vtkIdType tVertex
        = helper->MakeDistributedId(distribution.GetProcessorOfElement(t),
                                    distribution.GetLocalIndexOfElement(t));

      vtkDebugMacro(<<"Adding edge " << s << " to " << t);
      dirBuilder->LazyAddEdge(sVertex, tVertex);
      newEdgeFound = true;
      }
    }

  // Make sure everybody has added their edges and back-edges.
  helper->Synchronize();

  // Copy the structure into the output.
  vtkGraph *output = vtkGraph::GetData(outputVector);
  if (!output->CheckedShallowCopy(dirBuilder))
    {
    vtkErrorMacro(<<"Invalid structure.");
    return 0;
    }

  if (this->IncludeEdgeWeights)
    {
    if (!this->EdgeWeightArrayName)
      {
      vtkErrorMacro("When generating edge weights, "
        << "edge weights array name must be defined.");
      return 0;
      }
    vtkFloatArray *weights = vtkFloatArray::New();
    weights->SetName(this->EdgeWeightArrayName);
    for (vtkIdType i = 0; i < output->GetNumberOfEdges(); ++i)
      {
      weights->InsertNextValue(vtkMath::Random());
      }
    output->GetEdgeData()->AddArray(weights);
    weights->Delete();
    }

  if (this->GeneratePedigreeIds)
    {
    if (!this->VertexPedigreeIdArrayName || !this->EdgePedigreeIdArrayName)
      {
      vtkErrorMacro("When generating pedigree ids, "
        << "vertex and edge pedigree id array names must be defined.");
      return 0;
      }
    vtkIdType numVert = output->GetNumberOfVertices();
    vtkSmartPointer<vtkIdTypeArray> vertIds =
      vtkSmartPointer<vtkIdTypeArray>::New();
    vertIds->SetName(this->VertexPedigreeIdArrayName);
    vertIds->SetNumberOfTuples(numVert);
    for (vtkIdType i = 0; i < numVert; ++i)
      {
      vertIds->SetValue(i, myStartVertex + i);
      }
    output->GetVertexData()->SetPedigreeIds(vertIds);

    // Attach a distribution function to the helper that maps these
    // global vertex numbers back to .  TODO: Actually do it :)

    // Figure out how many edges come before us in the graph.
    vtkIdType numEdge = output->GetNumberOfEdges();
    boost::mpi::communicator world;
    vtkIdType myStartEdge
      = boost::mpi::scan(world, numEdge, std::plus<vtkIdType>()) - numEdge;

    vtkSmartPointer<vtkIdTypeArray> edgeIds =
      vtkSmartPointer<vtkIdTypeArray>::New();
    edgeIds->SetName(this->EdgePedigreeIdArrayName);
    edgeIds->SetNumberOfTuples(numEdge);
    for (vtkIdType i = 0; i < numEdge; ++i)
      {
      edgeIds->SetValue(i, myStartEdge + i);
      }
    output->GetEdgeData()->SetPedigreeIds(edgeIds);
    }

  return 1;
}

//----------------------------------------------------------------------------
int vtkPBGLRMATGraphSource::RequestDataObject(
  vtkInformation*,
  vtkInformationVector**,
  vtkInformationVector* )
{
  vtkDataObject *current = this->GetExecutive()->GetOutputData(0);
  if (!current || !vtkDirectedGraph::SafeDownCast(current))
    {
    vtkGraph *output = 0;
    output = vtkDirectedGraph::New();

    this->GetExecutive()->SetOutputData(0, output);
    output->Delete();
    }

  return 1;
}


//=======================================================================
// Copyright 2007 Aaron Windsor
//
// Distributed under the Boost Software License, Version 1.0. (See
// accompanying file LICENSE_1_0.txt or copy at
// http://www.boost.org/LICENSE_1_0.txt)
//=======================================================================

#ifndef __BOYER_MYRVOLD_PLANAR_TEST_HPP__
#define __BOYER_MYRVOLD_PLANAR_TEST_HPP__

#include <boost/graph/planar_detail/boyer_myrvold_impl.hpp>
#include <boost/parameter.hpp>
#include <boost/type_traits.hpp>
#include <boost/mpl/bool.hpp>

namespace boost
{

  struct no_kuratowski_subgraph_isolation {};
  struct no_planar_embedding {};

  namespace boyer_myrvold_params
  {

    BOOST_PARAMETER_KEYWORD(tag, graph)
    BOOST_PARAMETER_KEYWORD(tag, embedding)
    BOOST_PARAMETER_KEYWORD(tag, kuratowski_subgraph)
    BOOST_PARAMETER_KEYWORD(tag, vertex_index_map)
    BOOST_PARAMETER_KEYWORD(tag, edge_index_map)

    typedef parameter::parameters< parameter::required<tag::graph>,
                                   tag::embedding,
                                   tag::kuratowski_subgraph,
                                   tag::vertex_index_map,
                                   tag::edge_index_map
                                   > boyer_myrvold_params_t;

    namespace core
    {

      template <typename ArgumentPack>
      bool dispatched_boyer_myrvold(ArgumentPack const& args,
                                    mpl::true_,
                                    mpl::true_
                                    )
      {
        //Dispatch for no planar embedding, no kuratowski subgraph isolation

        typedef typename remove_const<
            typename parameter::value_type<ArgumentPack, tag::graph>::type
        >::type graph_t;

        typedef typename property_map<
            graph_t,
            vertex_index_t
        >::const_type vertex_default_index_map_t;

        typedef typename parameter::value_type<
            ArgumentPack,
            tag::vertex_index_map,
            vertex_default_index_map_t
        >::type vertex_index_map_t;

        graph_t const& g = args[graph];
        vertex_default_index_map_t v_d_map = get(vertex_index, g);
        vertex_index_map_t v_i_map = args[vertex_index_map | v_d_map];
        boyer_myrvold_impl
          <graph_t,
           vertex_index_map_t,
           graph::detail::no_old_handles,
           graph::detail::no_embedding
          >
          planarity_tester(g, v_i_map);

        return planarity_tester.is_planar() ? true : false;
      }

      template <typename ArgumentPack>
      bool dispatched_boyer_myrvold(ArgumentPack const& args,
                                    mpl::true_,
                                    mpl::false_
                                    )
      {
        //Dispatch for no planar embedding, kuratowski subgraph isolation
        typedef typename remove_const<
            typename parameter::value_type<ArgumentPack, tag::graph>::type
        >::type graph_t;

        typedef typename property_map<
            graph_t,
            vertex_index_t
        >::const_type vertex_default_index_map_t;

        typedef typename parameter::value_type<
            ArgumentPack,
            tag::vertex_index_map,
            vertex_default_index_map_t
        >::type vertex_index_map_t;

        typedef typename property_map<
            graph_t,
            edge_index_t
        >::const_type edge_default_index_map_t;

        typedef typename parameter::value_type<
            ArgumentPack,
            tag::edge_index_map,
            edge_default_index_map_t
        >::type edge_index_map_t;

        graph_t const& g = args[graph];
        vertex_default_index_map_t v_d_map = get(vertex_index, g);
        vertex_index_map_t v_i_map = args[vertex_index_map | v_d_map];
        edge_default_index_map_t e_d_map = get(edge_index, g);
        edge_index_map_t e_i_map = args[edge_index_map | e_d_map];
        boyer_myrvold_impl
          <graph_t,
           vertex_index_map_t,
           graph::detail::store_old_handles,
           graph::detail::no_embedding
          >
          planarity_tester(g, v_i_map);

        if (planarity_tester.is_planar())
          return true;
        else
          {
            planarity_tester.extract_kuratowski_subgraph
              (args[kuratowski_subgraph], e_i_map);
            return false;
          }
      }

      template <typename ArgumentPack>
      bool dispatched_boyer_myrvold(ArgumentPack const& args,
                                    mpl::false_,
                                    mpl::true_
                                    )
      {
        //Dispatch for planar embedding, no kuratowski subgraph isolation
        typedef typename remove_const<
            typename parameter::value_type<ArgumentPack, tag::graph>::type
        >::type graph_t;

        typedef typename property_map<
            graph_t,
            vertex_index_t
        >::const_type vertex_default_index_map_t;

        typedef typename parameter::value_type<
            ArgumentPack,
            tag::vertex_index_map,
            vertex_default_index_map_t
        >::type vertex_index_map_t;

        graph_t const& g = args[graph];
        vertex_default_index_map_t v_d_map = get(vertex_index, g);
        vertex_index_map_t v_i_map = args[vertex_index_map | v_d_map];
        boyer_myrvold_impl
          <graph_t,
           vertex_index_map_t,
           graph::detail::no_old_handles,
#ifdef BOOST_GRAPH_PREFER_STD_LIB
           graph::detail::std_list
#else
           graph::detail::recursive_lazy_list
#endif
          >
          planarity_tester(g, v_i_map);

        if (planarity_tester.is_planar())
          {
            planarity_tester.make_edge_permutation(args[embedding]);
            return true;
          }
        else
          return false;
      }

      template <typename ArgumentPack>
      bool dispatched_boyer_myrvold(ArgumentPack const& args,
                                    mpl::false_,
                                    mpl::false_
                                    )
      {
        //Dispatch for planar embedding, kuratowski subgraph isolation
        typedef typename remove_const<
            typename parameter::value_type<ArgumentPack, tag::graph>::type
        >::type graph_t;

        typedef typename property_map<
            graph_t,
            vertex_index_t
        >::const_type vertex_default_index_map_t;

        typedef typename parameter::value_type<
            ArgumentPack,
            tag::vertex_index_map,
            vertex_default_index_map_t
        >::type vertex_index_map_t;

        typedef typename property_map<
            graph_t,
            edge_index_t
        >::const_type edge_default_index_map_t;

        typedef typename parameter::value_type<
            ArgumentPack,
            tag::edge_index_map,
            edge_default_index_map_t
        >::type edge_index_map_t;

        graph_t const& g = args[graph];
        vertex_default_index_map_t v_d_map = get(vertex_index, g);
        vertex_index_map_t v_i_map = args[vertex_index_map | v_d_map];
        edge_default_index_map_t e_d_map = get(edge_index, g);
        edge_index_map_t e_i_map = args[edge_index_map | e_d_map];
        boyer_myrvold_impl
          <graph_t,
          vertex_index_map_t,
          graph::detail::store_old_handles,
#ifdef BOOST_BGL_PREFER_STD_LIB
           graph::detail::std_list
#else
           graph::detail::recursive_lazy_list
#endif
          >
          planarity_tester(g, v_i_map);

        if (planarity_tester.is_planar())
          {
            planarity_tester.make_edge_permutation(args[embedding]);
            return true;
          }
        else
          {
            planarity_tester.extract_kuratowski_subgraph
              (args[kuratowski_subgraph], e_i_map);
            return false;
          }
      }

      template <typename ArgumentPack>
      bool boyer_myrvold_planarity_test(ArgumentPack const& args)
      {

        typedef typename parameter::binding
          < ArgumentPack,
            tag::kuratowski_subgraph,
            const no_kuratowski_subgraph_isolation&
          >::type
          kuratowski_arg_t;

        typedef typename parameter::binding
          < ArgumentPack,
            tag::embedding,
            const no_planar_embedding&
          >::type
          embedding_arg_t;

         return dispatched_boyer_myrvold
           (args,
            boost::is_same
              <embedding_arg_t, const no_planar_embedding&>(),
            boost::is_same
              <kuratowski_arg_t, const no_kuratowski_subgraph_isolation&>()
            );
      }

    } //namespace core

  } //namespace boyer_myrvold_params

  template <typename A0>
  bool boyer_myrvold_planarity_test(A0 const& arg0)
  {
    return boyer_myrvold_params::core::boyer_myrvold_planarity_test
      (boyer_myrvold_params::boyer_myrvold_params_t()(arg0));
  }

  template <typename A0, typename A1>
  //  bool boyer_myrvold_planarity_test(A0 const& arg0, A1 const& arg1)
  bool boyer_myrvold_planarity_test(A0 const& arg0, A1 const& arg1)
  {
    return boyer_myrvold_params::core::boyer_myrvold_planarity_test
      (boyer_myrvold_params::boyer_myrvold_params_t()(arg0,arg1));
  }

  template <typename A0, typename A1, typename A2>
  bool boyer_myrvold_planarity_test(A0 const& arg0,
                                    A1 const& arg1,
                                    A2 const& arg2
                                    )
  {
    return boyer_myrvold_params::core::boyer_myrvold_planarity_test
      (boyer_myrvold_params::boyer_myrvold_params_t()(arg0,arg1,arg2));
  }

  template <typename A0, typename A1, typename A2, typename A3>
  bool boyer_myrvold_planarity_test(A0 const& arg0,
                                    A1 const& arg1,
                                    A2 const& arg2,
                                    A3 const& arg3
                                    )
  {
    return boyer_myrvold_params::core::boyer_myrvold_planarity_test
      (boyer_myrvold_params::boyer_myrvold_params_t()(arg0,arg1,arg2,arg3));
  }

  template <typename A0, typename A1, typename A2, typename A3, typename A4>
  bool boyer_myrvold_planarity_test(A0 const& arg0,
                                    A1 const& arg1,
                                    A2 const& arg2,
                                    A3 const& arg3,
                                    A4 const& arg4
                                    )
  {
    return boyer_myrvold_params::core::boyer_myrvold_planarity_test
      (boyer_myrvold_params::boyer_myrvold_params_t()
       (arg0,arg1,arg2,arg3,arg4)
       );
  }

}

#endif //__BOYER_MYRVOLD_PLANAR_TEST_HPP__


#ifndef MALABAR_LAYERS_HPP_
#define MALABAR_LAYERS_HPP_

/**
 *  @file
 *  @brief This header contains all malabar specific layers which have not
 *         yet merged back to upstream caffe.
 *
 *  Layers which are already available in caffe and have been changed
 *  or improved for the malabar project should be mentioned here as well.
 *
 *  Caffe Layer Improvements:
 *   TODO files and layers which have been improved should be listed here
 *
 */

#include <string>
#include <utility>
#include <vector>

#include <boost/random/mersenne_twister.hpp>

#include "caffe/layer.hpp"
#include "caffe/layers/loss_layer.hpp"
#include "caffe/layers/base_data_layer.hpp"
#include "caffe/layers/concat_layer.hpp"

#include "caffe/util/db.hpp"

#include "caffe/proto/caffe.pb.h"

#include "caffe/layers/neuron_layer.hpp"
#include "caffe/layers/softmax_layer.hpp"

namespace caffe {

  /**
   * Permutohedral Feature layer
   *
   */
  template <typename Dtype>
  class PixelFeatureLayer : public Layer<Dtype> {
  public:
    explicit PixelFeatureLayer(const LayerParameter& param)
      : Layer<Dtype>(param) {}
    virtual void LayerSetUp(const vector<Blob<Dtype>*>& bottom,
                            const vector<Blob<Dtype>*>& top);
    virtual void Reshape(const vector<Blob<Dtype>*>& bottom,
                         const vector<Blob<Dtype>*>& top);

    virtual inline const char* type() const { return "PixelFeature"; }
    virtual inline int ExactNumBottomBlobs() const { return 1; }
    virtual inline int ExactNumTopBlobs() const { return 1; }

  protected:
    virtual void Forward_cpu(const vector<Blob<Dtype>*>& bottom,
        const vector<Blob<Dtype>*>& top);
    virtual void Backward_cpu(const vector<Blob<Dtype>*>& top,
        const vector<bool>& propagate_down, const vector<Blob<Dtype>*>& bottom);

    virtual void Forward_gpu(const vector<Blob<Dtype>*>& bottom,
        const vector<Blob<Dtype>*>& top);
    virtual void Backward_gpu(const vector<Blob<Dtype>*>& top,
        const vector<bool>& propagate_down, const vector<Blob<Dtype>*>& bottom);

    // Permutohedral Feature parameters
    int count_;
    int num_;
    int channels_;
    int height_, width_;
    bool ran_once;
  };

  /**
   * Matrix multiplication layer.
   *
   */
  template <typename Dtype>
  class MatMulLayer : public Layer<Dtype> {
  public:
    explicit MatMulLayer(const LayerParameter& param)
      : Layer<Dtype>(param) {
      LayerParameter tmp_param;
      tmp_param.mutable_concat_param()->set_concat_dim(3);
      concat_layer_.reset(new ConcatLayer<Dtype>(tmp_param));
    }
    virtual void LayerSetUp(const vector<Blob<Dtype>*>& bottom,
                            const vector<Blob<Dtype>*>& top);
    virtual void Reshape(const vector<Blob<Dtype>*>& bottom,
                         const vector<Blob<Dtype>*>& top);

    virtual inline const char* type() const { return "MatMul"; }
    virtual inline int MinNumBottomBlobs() const { return 2; }
    virtual inline int ExactNumTopBlobs() const { return 1; }

  protected:
    virtual void Forward_cpu(const vector<Blob<Dtype>*>& bottom,
                             const vector<Blob<Dtype>*>& top);
    virtual void Forward_gpu(const vector<Blob<Dtype>*>& bottom,
                             const vector<Blob<Dtype>*>& top);
    virtual void Backward_cpu(const vector<Blob<Dtype>*>& top,
                              const vector<bool>& propagate_down,
                              const vector<Blob<Dtype>*>& bottom);
    virtual void Backward_gpu(const vector<Blob<Dtype>*>& top,
                              const vector<bool>& propagate_down,
                              const vector<Blob<Dtype>*>& bottom);
    Blob<Dtype> tmp_k_;
    boost::shared_ptr<ConcatLayer<Dtype> > concat_layer_;
  };

  /**
   * Matrix multiplication layer - 2.
   *
   */
  template <typename Dtype>
  class MatMul2Layer : public Layer<Dtype> {
  public:
    explicit MatMul2Layer(const LayerParameter& param)
      : Layer<Dtype>(param) {};
    virtual void LayerSetUp(const vector<Blob<Dtype>*>& bottom,
                            const vector<Blob<Dtype>*>& top);
    virtual void Reshape(const vector<Blob<Dtype>*>& bottom,
                         const vector<Blob<Dtype>*>& top);

    virtual inline const char* type() const { return "MatMul2"; }
    virtual inline int ExactNumBottomBlobs() const { return 2; }
    virtual inline int ExactNumTopBlobs() const { return 1; }

  protected:
    virtual void Forward_cpu(const vector<Blob<Dtype>*>& bottom,
                             const vector<Blob<Dtype>*>& top);
    virtual void Forward_gpu(const vector<Blob<Dtype>*>& bottom,
                             const vector<Blob<Dtype>*>& top);
    virtual void Backward_cpu(const vector<Blob<Dtype>*>& top,
                              const vector<bool>& propagate_down,
                              const vector<Blob<Dtype>*>& bottom);
    virtual void Backward_gpu(const vector<Blob<Dtype>*>& top,
                              const vector<bool>& propagate_down,
                              const vector<Blob<Dtype>*>& bottom);
    int M_;
    int K_;
    int N_;
    int num_kernels_;
    int channels_;
    bool bias_term_;
    Blob<Dtype> bias_multiplier_;
  };

/**
 * @brief Computes a product of two input Blobs, with the shape of the
 *        latter Blob "broadcast" to match the shape of the former.
 *        Equivalent to tiling the latter Blob, then computing the elementwise
 *        product.
 * Taken from - caffe ://github.com/jeffdonahue/caffe/commit/b4a5b6abb6272207da83bacde448fa3b2d4c7793?diff=split
 */
template <typename Dtype>
class ScalarLayer: public Layer<Dtype> {
 public:
  explicit ScalarLayer(const LayerParameter& param)
      : Layer<Dtype>(param) {}
  virtual void Reshape(const vector<Blob<Dtype>*>& bottom,
      const vector<Blob<Dtype>*>& top);

  virtual inline const char* type() const { return "Scalar"; }
  virtual inline int ExactNumBottomBlobs() const { return 2; }
  virtual inline int ExactNumTopBlobs() const { return 1; }

 protected:
  /**
   * In the below shape specifications, @f$ i @f$ denotes the value of the
   * `axis` field given by `this->layer_param_.scalar_param().axis()`, after
   * canonicalization (i.e., conversion from negative to positive index,
   * if applicable).
   *
   * @param bottom input Blob vector (length 2)
   *   -# @f$ (d_0 \times ... \times
   *           d_i \times ... \times d_j \times ... \times d_n) @f$
   *      the first factor @f$ x @f$
   *   -# @f$ (d_i \times ... \times d_j) @f$
   *      the second factor @f$ y @f$
   * @param top output Blob vector (length 1)
   *   -# @f$ (d_0 \times ... \times
   *           d_i \times ... \times d_j \times ... \times d_n) @f$
   *      the product @f$ z = x y @f$ computed after "broadcasting" y.
   *      Equivalent to tiling @f$ y @f$ to have the same shape as @f$ x @f$,
   *      then computing the elementwise product.
   */
  virtual void Forward_cpu(const vector<Blob<Dtype>*>& bottom,
      const vector<Blob<Dtype>*>& top);
  virtual void Forward_gpu(const vector<Blob<Dtype>*>& bottom,
      const vector<Blob<Dtype>*>& top);
  virtual void Backward_cpu(const vector<Blob<Dtype>*>& top,
      const vector<bool>& propagate_down, const vector<Blob<Dtype>*>& bottom);
  virtual void Backward_gpu(const vector<Blob<Dtype>*>& top,
      const vector<bool>& propagate_down, const vector<Blob<Dtype>*>& bottom);

  Blob<Dtype> sum_multiplier_;
  Blob<Dtype> sum_result_;
  int axis_;
  int outer_dim_, scalar_dim_, inner_dim_;
};

template <typename Dtype>
class Scalar2Layer: public Layer<Dtype> {
 public:
  explicit Scalar2Layer(const LayerParameter& param)
      : Layer<Dtype>(param) {}
  virtual void Reshape(const vector<Blob<Dtype>*>& bottom,
      const vector<Blob<Dtype>*>& top);

  virtual inline const char* type() const { return "Scalar2"; }
  virtual inline int ExactNumBottomBlobs() const { return 2; }
  virtual inline int ExactNumTopBlobs() const { return 1; }

 protected:
  virtual void Forward_cpu(const vector<Blob<Dtype>*>& bottom,
      const vector<Blob<Dtype>*>& top);
  virtual void Forward_gpu(const vector<Blob<Dtype>*>& bottom,
      const vector<Blob<Dtype>*>& top);
  virtual void Backward_cpu(const vector<Blob<Dtype>*>& top,
      const vector<bool>& propagate_down, const vector<Blob<Dtype>*>& bottom);
  virtual void Backward_gpu(const vector<Blob<Dtype>*>& top,
      const vector<bool>& propagate_down, const vector<Blob<Dtype>*>& bottom);

  Blob<Dtype> tmp_diff_;
  Blob<Dtype> tmp_sum_;
};

template <typename Dtype>
class Scalar3Layer: public Layer<Dtype> {
 public:
  explicit Scalar3Layer(const LayerParameter& param)
      : Layer<Dtype>(param) {}
  virtual void Reshape(const vector<Blob<Dtype>*>& bottom,
      const vector<Blob<Dtype>*>& top);

  virtual inline const char* type() const { return "Scalar3"; }
  virtual inline int ExactNumBottomBlobs() const { return 2; }
  virtual inline int ExactNumTopBlobs() const { return 1; }

 protected:
  virtual void Forward_cpu(const vector<Blob<Dtype>*>& bottom,
      const vector<Blob<Dtype>*>& top);
  virtual void Forward_gpu(const vector<Blob<Dtype>*>& bottom,
      const vector<Blob<Dtype>*>& top);
  virtual void Backward_cpu(const vector<Blob<Dtype>*>& top,
      const vector<bool>& propagate_down, const vector<Blob<Dtype>*>& bottom);
  virtual void Backward_gpu(const vector<Blob<Dtype>*>& top,
      const vector<bool>& propagate_down, const vector<Blob<Dtype>*>& bottom);

  Blob<Dtype> tmp_sum_;
};

/**
 * @brief Computes a product of input Blob and parameter blob, with the shape of the
 *        parameter Blob "broadcast" to match the shape of the former.
 *        Equivalent to tiling the parameter, then computing the elementwise
 *        product.
 * Adapted from - caffe ://github.com/jeffdonahue/caffe/commit/b4a5b6abb6272207da83bacde448fa3b2d4c7793?diff=split
 */
template <typename Dtype>
class ScalarConvLayer: public Layer<Dtype> {
 public:
  explicit ScalarConvLayer(const LayerParameter& param)
      : Layer<Dtype>(param) {}
  virtual void LayerSetUp(const vector<Blob<Dtype>*>& bottom,
          const vector<Blob<Dtype>*>& top);
  virtual void Reshape(const vector<Blob<Dtype>*>& bottom,
      const vector<Blob<Dtype>*>& top);

  virtual inline const char* type() const { return "ScalarConv"; }
  virtual inline int ExactNumBottomBlobs() const { return 1; }
  virtual inline int ExactNumTopBlobs() const { return 1; }

 protected:
  virtual void Forward_cpu(const vector<Blob<Dtype>*>& bottom,
      const vector<Blob<Dtype>*>& top);
  virtual void Forward_gpu(const vector<Blob<Dtype>*>& bottom,
      const vector<Blob<Dtype>*>& top);
  virtual void Backward_cpu(const vector<Blob<Dtype>*>& top,
      const vector<bool>& propagate_down, const vector<Blob<Dtype>*>& bottom);
  virtual void Backward_gpu(const vector<Blob<Dtype>*>& top,
      const vector<bool>& propagate_down, const vector<Blob<Dtype>*>& bottom);

  Blob<Dtype> sum_multiplier_;
  Blob<Dtype> sum_result_;
  int axis_;
  int outer_dim_, scalar_dim_, inner_dim_;
};

/**
 * @brief Computes @f$ y = exp ^ {\alpha x + \beta} @f$,
 *        as specified by the parameter @f$ \alpha @f$, shift @f$ \beta @f$.
 */
template <typename Dtype>
class ExpScaleLayer : public NeuronLayer<Dtype> {
 public:
  /**
   * @param param provides ExpParameter exp_scale_param,
   *     with ExpScaleLayer options:
   *   - shift (\b optional, default 0) the shift @f$ \beta @f$
   */
  explicit ExpScaleLayer(const LayerParameter& param)
      : NeuronLayer<Dtype>(param) {}
  virtual void LayerSetUp(const vector<Blob<Dtype>*>& bottom,
      const vector<Blob<Dtype>*>& top);

  virtual inline const char* type() const { return "ExpScale"; }

 protected:
  /**
   * @param bottom input Blob vector (length 1)
   *   -# @f$ (N \times C \times H \times W) @f$
   *      the inputs @f$ x @f$
   * @param top output Blob vector (length 1)
   *   -# @f$ (N \times C \times H \times W) @f$
   *      the computed outputs @f$
   *        y = e ^ {\alpha x + \beta}
   *      @f$
   */
  virtual void Forward_cpu(const vector<Blob<Dtype>*>& bottom,
      const vector<Blob<Dtype>*>& top);
  virtual void Forward_gpu(const vector<Blob<Dtype>*>& bottom,
      const vector<Blob<Dtype>*>& top);

  /**
   * @brief Computes the error gradient w.r.t. the exp inputs and
   * scale parameter.
   */
  virtual void Backward_cpu(const vector<Blob<Dtype>*>& top,
      const vector<bool>& propagate_down, const vector<Blob<Dtype>*>& bottom);
  virtual void Backward_gpu(const vector<Blob<Dtype>*>& top,
      const vector<bool>& propagate_down, const vector<Blob<Dtype>*>& bottom);

  Dtype outer_scale_;
};

/**
 * @brief Computes @f$ y = \alpha x  @f$,
 *        as specified by the learnable parameter @f$ \alpha @f$.
 */
template <typename Dtype>
class Scale2Layer : public NeuronLayer<Dtype> {
 public:
  explicit Scale2Layer(const LayerParameter& param)
      : NeuronLayer<Dtype>(param) {}
  virtual void LayerSetUp(const vector<Blob<Dtype>*>& bottom,
      const vector<Blob<Dtype>*>& top);

  virtual inline const char* type() const { return "Scale2"; }

 protected:
  virtual void Forward_cpu(const vector<Blob<Dtype>*>& bottom,
      const vector<Blob<Dtype>*>& top);
  virtual void Forward_gpu(const vector<Blob<Dtype>*>& bottom,
      const vector<Blob<Dtype>*>& top);

  /**
   * @brief Computes the error gradient w.r.t. the inputs and
   * scale parameter.
   */
  virtual void Backward_cpu(const vector<Blob<Dtype>*>& top,
      const vector<bool>& propagate_down, const vector<Blob<Dtype>*>& bottom);
  virtual void Backward_gpu(const vector<Blob<Dtype>*>& top,
      const vector<bool>& propagate_down, const vector<Blob<Dtype>*>& bottom);
};

template <typename Dtype>
class SmearLayer : public Layer<Dtype> {
 public:
  explicit SmearLayer(const LayerParameter& param)
      : Layer<Dtype>(param) {}
  virtual void LayerSetUp(const vector<Blob<Dtype>*>& bottom,
    const vector<Blob<Dtype>*>& top);
  virtual void Reshape(const vector<Blob<Dtype>*>& bottom,
      const vector<Blob<Dtype>*>& top);

  virtual inline const char* type() const { return "Smear"; }
  virtual inline int ExactNumBottomBlobs() const { return 2; }
  virtual inline int ExactNumTopBlobs() const { return 1; }

 protected:
  virtual void Forward_cpu(const vector<Blob<Dtype>*>& bottom,
      const vector<Blob<Dtype>*>& top);
  virtual void Forward_gpu(const vector<Blob<Dtype>*>& bottom,
      const vector<Blob<Dtype>*>& top);
  virtual void Backward_cpu(const vector<Blob<Dtype>*>& top,
      const vector<bool>& propagate_down, const vector<Blob<Dtype>*>& bottom);
  virtual void Backward_gpu(const vector<Blob<Dtype>*>& top,
      const vector<bool>& propagate_down, const vector<Blob<Dtype>*>& bottom);

  int num_;
  int channels_;

  int out_height_;
  int out_width_;

  int outer_num_;
  int inner_num_;

  Dtype ignore_idx_value_;
  Dtype ignore_feature_value_;
};

template <typename Dtype>
class PdistLayer : public Layer<Dtype> {
 public:
  explicit PdistLayer(const LayerParameter& param)
      : Layer<Dtype>(param) {}
  virtual void LayerSetUp(const vector<Blob<Dtype>*>& bottom,
    const vector<Blob<Dtype>*>& top);
  virtual void Reshape(const vector<Blob<Dtype>*>& bottom,
      const vector<Blob<Dtype>*>& top);

  virtual inline const char* type() const { return "Pdist"; }
  virtual inline int ExactNumBottomBlobs() const { return 2; }
  virtual inline int ExactNumTopBlobs() const { return 1; }

 protected:
  virtual void Forward_cpu(const vector<Blob<Dtype>*>& bottom,
      const vector<Blob<Dtype>*>& top);
  virtual void Forward_gpu(const vector<Blob<Dtype>*>& bottom,
      const vector<Blob<Dtype>*>& top);
  virtual void Backward_cpu(const vector<Blob<Dtype>*>& top,
      const vector<bool>& propagate_down, const vector<Blob<Dtype>*>& bottom);
  virtual void Backward_gpu(const vector<Blob<Dtype>*>& top,
      const vector<bool>& propagate_down, const vector<Blob<Dtype>*>& bottom);

  int num_;
  int channels_;

  int out_height_;
  int out_width_;

  float ignore_value_;
  float scale_value_;
};

template <typename Dtype>
class SpixelFeatureLayer : public Layer<Dtype> {
 public:
  explicit SpixelFeatureLayer(const LayerParameter& param)
      : Layer<Dtype>(param) {}
  virtual void LayerSetUp(const vector<Blob<Dtype>*>& bottom,
    const vector<Blob<Dtype>*>& top);
  virtual void Reshape(const vector<Blob<Dtype>*>& bottom,
      const vector<Blob<Dtype>*>& top);

  virtual inline const char* type() const { return "SpixelFeature"; }
  virtual inline int ExactNumBottomBlobs() const { return 2; }

 protected:
  virtual void Forward_cpu(const vector<Blob<Dtype>*>& bottom,
      const vector<Blob<Dtype>*>& top);
  virtual void Forward_gpu(const vector<Blob<Dtype>*>& bottom,
      const vector<Blob<Dtype>*>& top);
  virtual void Backward_cpu(const vector<Blob<Dtype>*>& top,
      const vector<bool>& propagate_down, const vector<Blob<Dtype>*>& bottom);
  virtual void Backward_gpu(const vector<Blob<Dtype>*>& top,
      const vector<bool>& propagate_down, const vector<Blob<Dtype>*>& bottom);

  int num_;
  int in_channels_;
  int height_;
  int width_;
  int out_channels_;
  int max_spixels_;

  float rgbxy_rgb_scale_;
  float rgbxy_xy_scale_;
  float xy_scale_;
  float rgb_scale_;
  float ignore_idx_value_;
  float ignore_feature_value_;

  Blob<Dtype> spixel_counts_;
};

template <typename Dtype>
class TransposeLayer: public Layer<Dtype> {
 public:
  explicit TransposeLayer(const LayerParameter& param)
      : Layer<Dtype>(param) {}
  virtual void LayerSetUp(const vector<Blob<Dtype>*>& bottom,
          const vector<Blob<Dtype>*>& top);
  virtual void Reshape(const vector<Blob<Dtype>*>& bottom,
      const vector<Blob<Dtype>*>& top);

  virtual inline const char* type() const { return "Transpose"; }
  virtual inline int ExactNumBottomBlobs() const { return 1; }
  virtual inline int ExactNumTopBlobs() const { return 1; }

 protected:
  virtual void Forward_cpu(const vector<Blob<Dtype>*>& bottom,
      const vector<Blob<Dtype>*>& top);
  virtual void Forward_gpu(const vector<Blob<Dtype>*>& bottom,
      const vector<Blob<Dtype>*>& top);
  virtual void Backward_cpu(const vector<Blob<Dtype>*>& top,
      const vector<bool>& propagate_down, const vector<Blob<Dtype>*>& bottom);
  virtual void Backward_gpu(const vector<Blob<Dtype>*>& top,
      const vector<bool>& propagate_down, const vector<Blob<Dtype>*>& bottom);

  int num_, channels_, height_, width_;
};

  /**
   * @brief Implements the Bilateral Convolution as brute-force but exact.
   *        K*x with K=exp(-D)
   *
   * TODO(pgehler): thorough documentation for Forward, Backward, and proto params.
   */
  template <typename Dtype>
  class BilateralBruteForceLayer : public Layer<Dtype> {
   public:
    explicit BilateralBruteForceLayer(const LayerParameter& param)
      : Layer<Dtype>(param) {};
    virtual void LayerSetUp(const vector<Blob<Dtype>*>& bottom,
      const vector<Blob<Dtype>*>& top);
    virtual void Reshape(const vector<Blob<Dtype>*>& bottom,
        const vector<Blob<Dtype>*>& top);

    virtual inline const char* type() const { return "BilateralBruteForce"; }
    virtual inline int ExactNumBottomBlobs() const { return 4; }
    virtual inline int ExactNumTopBlobs() const { return 1; }

   protected:
    virtual void Forward_cpu(const vector<Blob<Dtype>*>& bottom,
        const vector<Blob<Dtype>*>& top);
    virtual void Forward_gpu(const vector<Blob<Dtype>*>& bottom,
        const vector<Blob<Dtype>*>& top);
    virtual void Backward_cpu(const vector<Blob<Dtype>*>& top,
        const vector<bool>& propagate_down, const vector<Blob<Dtype>*>& bottom);
    virtual void Backward_gpu(const vector<Blob<Dtype>*>& top,
        const vector<bool>& propagate_down, const vector<Blob<Dtype>*>& bottom);

    int num_;
    int channels_;
    int scales_;
    int feature_size_;
    int in_height_, in_width_;
    int out_height_, out_width_;

    // we need this since we have to accumulate gradients for several scalse
    Blob<Dtype> bb_bottom_data_; // pixel
    Blob<Dtype> bb_bottom_f1_;   // pixel
    Blob<Dtype> bb_bottom_f2_;   // superpixel
    Blob<Dtype> bb_bottom_s_;    // scales

    Blob<Dtype> bb_top_pdist_;
    Blob<Dtype> bb_top_scalar_;
    Blob<Dtype> bb_top_softmax_;
    Blob<Dtype> bb_top_matmul_;

    // scalar needs only the top for the backward we can exploit that
    // Blob<Dtype> bb_pdist_full_;
    Blob<Dtype> bb_scalar_full_;
    Blob<Dtype> bb_softmax_full_;

    std::vector<Blob<Dtype>*> tmp_bottom_pdist_;
    std::vector<Blob<Dtype>*> tmp_top_pdist_;
    boost::shared_ptr<PdistLayer<Dtype> > pdist_layer_;
    std::vector<Blob<Dtype>*> tmp_bottom_scalar_;
    std::vector<Blob<Dtype>*> tmp_top_scalar_;
    boost::shared_ptr<Scalar3Layer<Dtype> > scalar_layer_;
    std::vector<Blob<Dtype>*> tmp_bottom_softmax_;
    std::vector<Blob<Dtype>*> tmp_top_softmax_;
    boost::shared_ptr<SoftmaxLayer<Dtype> > softmax_layer_;
    std::vector<Blob<Dtype>*> tmp_bottom_matmul_;
    std::vector<Blob<Dtype>*> tmp_top_matmul_;
    boost::shared_ptr<MatMul2Layer<Dtype> > matmul_layer_;
  };

}  // namespace caffe

#endif  // MALABAR_LAYERS_HPP_


/*ckwg +29
 * Copyright 2014-2020 by Kitware, Inc.
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *
 *  * Redistributions of source code must retain the above copyright notice,
 *    this list of conditions and the following disclaimer.
 *
 *  * 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.
 *
 *  * Neither name of Kitware, Inc. nor the names of any 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 AUTHORS 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.
 */

/**
 * \file
 * \brief Implementation of Necker reversal functions
 */

#include "necker_reverse.h"
#include <Eigen/Geometry>
#include <Eigen/SVD>

namespace kwiver {
namespace arrows {
namespace mvg {

/// Compute a plane passing through the landmarks
vital::vector_4d
landmark_plane(const vital::landmark_map::map_landmark_t& landmarks)
{
  using namespace kwiver::vital;
  // compute the landmark location mean and covariance
  vector_3d lc(0.0, 0.0, 0.0);
  matrix_3x3d covar = matrix_3x3d::Zero();
  for (auto const& p : landmarks)
  {
    vector_3d pt = p.second->loc();
    lc += pt;
    covar += pt * pt.transpose();
  }
  const double num_lm = static_cast<double>(landmarks.size());
  lc /= num_lm;
  covar /= num_lm;
  covar -= lc * lc.transpose();

  // the plane will pass through the landmark centeroid (lc)
  // and have a normal vector aligned with the smallest eigenvector of covar
  Eigen::JacobiSVD<matrix_3x3d> svd(covar, Eigen::ComputeFullV);
  vector_3d axis = svd.matrixV().col(2);
  return vector_4d(axis.x(), axis.y(), axis.z(), -(lc.dot(axis)));
}

/// Mirror landmarks about the specified plane
vital::landmark_map_sptr
mirror_landmarks(vital::landmark_map const& landmarks,
                 vital::vector_4d const& plane)
{
  using namespace kwiver::vital;
  landmark_map::map_landmark_t new_lms;
  const vector_3d axis(plane.x(), plane.y(), plane.z());
  const double d = plane[3];
  // mirror landmark locations about the mirroring plane
  for (auto const& p : landmarks.landmarks())
  {
    vector_3d v = p.second->loc();
    v -= 2.0 * (v.dot(axis) + d) * axis;
    auto new_lm = std::make_shared<vital::landmark_d>(*p.second);
    new_lm->set_loc(v);
    new_lms[p.first] = new_lm;
  }
  return std::make_shared<simple_landmark_map>(new_lms);
}

/// Compute the Necker reversal of a camera in place
void
necker_reverse_inplace(vital::simple_camera_perspective& camera,
                       vital::vector_4d const& plane)
{
  using namespace kwiver::vital;
  const vector_3d axis(plane.x(), plane.y(), plane.z());
  const double d = plane[3];
  const rotation_d Ra180(vector_4d(axis.x(), axis.y(), axis.z(), 0.0));
  static const rotation_d Rz180(vector_4d(0.0, 0.0, 1.0, 0.0));

  // extract the camera center
  const vital::vector_3d cc = camera.center();
  // extract the camera principal axis
  const vital::vector_3d pa = camera.rotation().matrix().row(2);
  // compute the distance from cc along pa until intersection with
  // the mirroring plane of the points
  const double dist = -(cc.dot(axis) + d) / pa.dot(axis);
  // compute the ground point where the principal axis
  // intersects the mirroring plane
  const vital::vector_3d gp = cc + dist * pa;
  // rotate the camera center 180 degrees about the mirroring plane normal
  // axis centered at gp, also rotate the camera 180 about its principal axis
  camera.set_center(Ra180 * (cc - gp) + gp);
  camera.set_rotation(Rz180 * camera.rotation() * Ra180);
}

/// Compute the Necker reversal of the cameras
vital::camera_map_sptr
necker_reverse(vital::camera_map const& cameras,
               vital::vector_4d const& plane)
{
  using namespace kwiver::vital;
  camera_map::map_camera_t cams;
  // flip cameras around
  for (auto & p : cameras.cameras())
  {
    // make a clone of this camera as a simple_camera_perspective
    auto flipped = std::make_shared<vital::simple_camera_perspective>(
      dynamic_cast<vital::simple_camera_perspective&>(*p.second));

    necker_reverse_inplace(*flipped, plane);
    cams[p.first] = flipped;
  }
  return std::make_shared<simple_camera_map>(cams);
}

/// Compute an approximate Necker reversal of cameras and landmarks
void
necker_reverse(vital::camera_map_sptr& cameras,
               vital::landmark_map_sptr& landmarks,
               bool reverse_landmarks)
{
  using namespace kwiver::vital;
  const vector_4d plane = landmark_plane(landmarks->landmarks());

  // reverse the cameras
  cameras = necker_reverse(*cameras, plane);

  if (reverse_landmarks)
  {
    // mirror the landmarks
    landmarks = mirror_landmarks(*landmarks, plane);
  }
}

} // end namespace mvg
} // end namespace arrows
} // end namespace kwiver


#ifdef USE_OPENCV
#include <opencv2/core/core.hpp>

#include <fstream>  // NOLINT(readability/streams)
#include <iostream>  // NOLINT(readability/streams)
#include <string>
#include <utility>
#include <vector>
#include <boost/algorithm/string.hpp>
#include <boost/thread.hpp>
#include <opencv2/core/core.hpp>
#include <opencv2/imgproc/imgproc.hpp>
#include <opencv2/highgui/highgui.hpp>
#include <boost/random/uniform_real_distribution.hpp>

#include "caffe/layers/bb3txt_data_layer.hpp"
#include "caffe/util/benchmark.hpp"
#include "caffe/util/io.hpp"
#include "caffe/util/math_functions.hpp"
#include "caffe/util/rng.hpp"

// The maximum number of bounding boxes (annotations) in one image - we set the label blob shape according
// to this number
#define MAX_NUM_BBS_PER_IMAGE 20

namespace caffe {

namespace {

    /**
     * @brief Computes the number of bounding boxes in the annotation
     * @param labels Annotation of one image (dimensions MAX_NUM_BBS_PER_IMAGE x 12 x 1 x 1)
     * @param b Image id in the batch
     * @return Number of bounding boxes in this label
     */
    template <typename Dtype>
    int numBBs (const Blob<Dtype> &labels)
    {
        for (int i = 0; i < labels.shape(1); ++i)
        {
            // Data are stored like this [label, ...]
            const Dtype *data = labels.cpu_data() + labels.offset(i);
            // If the label is -1, there are no more bounding boxes
            if (data[0] == Dtype(-1.0f)) return i;
        }

        return labels.shape(1);
    }

}

template <typename Dtype>
BB3TXTDataLayer<Dtype>::BB3TXTDataLayer (const LayerParameter &param)
    : BasePrefetchingDataLayer<Dtype>(param),
      InternalThreadpool(std::min(int(std::max(int(boost::thread::hardware_concurrency()/2), 1)),
                                  int(param.image_data_param().batch_size())))
{
}

template <typename Dtype>
BB3TXTDataLayer<Dtype>::~BB3TXTDataLayer<Dtype> ()
{
    this->StopInternalThreadpool();
    this->StopInternalThread();
}

template <typename Dtype>
void BB3TXTDataLayer<Dtype>::DataLayerSetUp (const vector<Blob<Dtype>*> &bottom,
                                            const vector<Blob<Dtype>*> &top)
{
    CHECK(this->layer_param_.has_bbtxt_param()) << "BBTXTParam is mandatory!";
    CHECK(this->layer_param_.bbtxt_param().has_height()) << "Height must be set!";
    CHECK(this->layer_param_.bbtxt_param().has_width()) << "Width must be set!";
    CHECK(this->layer_param_.bbtxt_param().has_reference_size_min()) << "Min reference size must be set!";
    CHECK(this->layer_param_.bbtxt_param().has_reference_size_max()) << "Max reference size must be set!";
    CHECK_LT(this->layer_param_.bbtxt_param().reference_size_min(),
             this->layer_param_.bbtxt_param().reference_size_max()) << "Min reference must be lower than max";

    const int height     = this->layer_param_.bbtxt_param().height();
    const int width      = this->layer_param_.bbtxt_param().width();
    const int batch_size = this->layer_param_.image_data_param().batch_size();

    this->_rng.reset(new Caffe::RNG(caffe_rng_rand()));

    // Load the BBTXT file with 2D bounding box annotations
    this->_loadBB3TXTFile();
    this->_i_global = 0;

    CHECK(!this->_images.empty()) << "The given BBTXT file is empty!";
    LOG(INFO) << "There are " << this->_images.size() << " images in the dataset.";

    // This is the shape of the input blob
    std::vector<int> top_shape = {batch_size, 3, height, width};
    this->transformed_data_.Reshape(top_shape);  // For prefetching
    top[0]->Reshape(top_shape);

    // Label blob
    std::vector<int> label_shape = {batch_size, MAX_NUM_BBS_PER_IMAGE, 12};
    this->transformed_label_.Reshape(label_shape);  // For prefetching
    top[1]->Reshape(label_shape);

    // Initialize prefetching
    // We also have to reshape the prefetching blobs to the correct batch size
    for (int i = 0; i < this->prefetch_.size(); ++i)
    {
        this->prefetch_[i]->data_.Reshape(top_shape);
        this->prefetch_[i]->label_.Reshape(label_shape);
    }

    this->StartInternalThreadpool();
}

template <typename Dtype>
void BB3TXTDataLayer<Dtype>::load_batch (Batch<Dtype> *batch)
{
    // This function is called on a prefetch thread

    CHECK(batch->data_.count());
    CHECK(this->transformed_data_.count());

    const int batch_size = this->layer_param_.image_data_param().batch_size();

    Dtype* prefetch_data  = batch->data_.mutable_cpu_data();
    Dtype* prefetch_label = batch->label_.mutable_cpu_data();

    this->transformed_data_.set_cpu_data(prefetch_data);
    this->transformed_label_.set_cpu_data(prefetch_label);

    this->_num_processed.reset();
    for (int b = 0; b < batch_size; ++b) this->_b_queue.push(b);
    this->_num_processed.waitToCount(batch_size);
}

template <typename Dtype>
void BB3TXTDataLayer<Dtype>::InternalThreadpoolEntry (int t)
{
    // This method runs on each thread of the internal threadpool

    std::cout << "============================= STARTING DATA THREAD " << t << std::endl;

    try {
        while (!this->must_stopt(t))
        {
            int b = this->_b_queue.pop();

            // Get index of image and bounding box we will crop
            SelectedBB<Dtype> selbb = this->_getImageAndBB();

            cv::Mat cv_img = cv::imread(selbb.filename, CV_LOAD_IMAGE_COLOR);
            CHECK(cv_img.data) << "Could not open " << selbb.filename;

            // Copy the annotation - we really have to copy it because it will be altered during image
            // transformations like cropping or scaling
            caffe_copy(selbb.label->count(), selbb.label->cpu_data(),
                       this->transformed_label_.mutable_cpu_data() + this->transformed_label_.offset(b));

            // We select a bounding box from the image and then make a crop such that the bounding box is
            // inside of it and it has the reference size (Training - we select a random bounding box to crop
            // from each image, Testing - crop all bounding boxes from the image - this way we ensure
            // the test set is always the same)
            this->_cropAndTransform(cv_img, b, selbb.bb_id);

            // Raise the counter on processed images
            this->_num_processed.increase();

        }
    } catch (boost::thread_interrupted&) {
        // Interrupted exception is expected on shutdown
    }

    std::cout << "============================= KILLING DATA THREAD " << t << std::endl;
}

// -----------------------------------------  PROTECTED METHODS  ----------------------------------------- //

template <typename Dtype>
void BB3TXTDataLayer<Dtype>::_loadBB3TXTFile ()
{
    const std::string& source = this->layer_param_.image_data_param().source();

    std::ifstream infile(source.c_str(), std::ios::in);
    CHECK(infile.is_open()) << "BB3TXT file '" << source << "' could not be opened!";

    std::string line;
    std::vector<std::string> data;
    std::string current_filename = "";
    int i = 0;

    // Read the whole file and create entries in the _images for all images
    while (std::getline(infile, line))
    {
        // Split the line - entries separated by space:
        // [filename label confidence fblx fbly fbrx fbry rblx rbly ftly]
        boost::split(data, line, boost::is_any_of(" "));
        CHECK_EQ(data.size(), 14) << "Line '" << line << "' corrupted!";

        if (current_filename != data[0])
        {
            // This is a label to a new image
            if (this->_images.size() > 0 && i < MAX_NUM_BBS_PER_IMAGE)
            {
                // Finalize the last annotation - we put -1 as next bounding box label to signalize
                // the end - this is because each image can have a different number of bounding boxes
                int offset = this->_images.back().second->offset(i);
                Dtype* bb3_position = this->_images.back().second->mutable_cpu_data() + offset;
                bb3_position[0] = Dtype(-1.0f);
            }

            CHECK(boost::filesystem::exists(data[0])) << "File '" << data[0] << "' not found!";

            // Create new image entry
            this->_images.emplace_back(data[0], std::make_shared<Blob<Dtype>>(MAX_NUM_BBS_PER_IMAGE, 12, 1, 1));
            i = 0;
            current_filename = data[0];
        }

        // Write the bounding box info into the blob
        if (i < MAX_NUM_BBS_PER_IMAGE)
        {
            int offset = this->_images.back().second->offset(i);
            Dtype* bb3_position = this->_images.back().second->mutable_cpu_data() + offset;

            bb3_position[0] = Dtype(std::stof(data[1])); // label
            // xmin, ymin, xmax, ymax, fblx, fbly, fbrx, fbry, rblx, rbly, ftly
            for (int p = 1; p < 12; ++p) bb3_position[p] = Dtype(std::stof(data[p+2]));

            this->_indices.emplace_back(this->_images.size()-1, i);
            i++;
        }
        else
        {
            LOG(WARNING) << "Skipping bb - max number of bounding boxes per image reached.";
        }
    }

    // Close the last annotation
    if (i < MAX_NUM_BBS_PER_IMAGE)
    {
        int offset = this->_images.back().second->offset(i);
        Dtype* bb3_position = this->_images.back().second->mutable_cpu_data() + offset;
        bb3_position[0] = Dtype(-1.0f);
    }
}

template <typename Dtype>
void BB3TXTDataLayer<Dtype>::_shuffleBoundingBoxes ()
{
    caffe::rng_t* prefetch_rng = static_cast<caffe::rng_t*>(_rng->generator());
    shuffle(this->_indices.begin(), this->_indices.end(), prefetch_rng);
}

template <typename Dtype>
void BB3TXTDataLayer<Dtype>::_cropAndTransform (const cv::Mat &cv_img, int b, int bb_id)
{
    CHECK_EQ(cv_img.channels(), 3) << "Image must have 3 color channels";

    // Crop this bounding box from the image
    cv::Mat cv_img_cropped;
    this->_cropBBFromImage(cv_img, cv_img_cropped, b, bb_id);

    // Mirror
    // We cannot do it here because we do not have all the needed coordinates - we have to do it externaly

    // Rotate

    // Copy the cropped and transformed image to the input blob
    this->_applyPixelTransformationsAndCopyOut(cv_img_cropped, b);
}

template <typename Dtype>
void BB3TXTDataLayer<Dtype>::_cropBBFromImage (const cv::Mat &cv_img, cv::Mat &cv_img_cropped_out,
                                              int b, int bb_id)
{
    // Input dimensions of the network
    const int height             = this->layer_param_.bbtxt_param().height();
    const int width              = this->layer_param_.bbtxt_param().width();
    const int reference_size_min = this->layer_param_.bbtxt_param().reference_size_min();
    const int reference_size_max = this->layer_param_.bbtxt_param().reference_size_max();
    caffe::rng_t* rng            = static_cast<caffe::rng_t*>(this->_rng->generator());

    // Get dimensions of the bounding box - format [label, xmin, ymin, xmax, ymax]
    const Dtype *bb_data = this->transformed_label_.cpu_data() + this->transformed_label_.offset(b, bb_id);
    const Dtype x = bb_data[1];
    const Dtype y = bb_data[2];
    const Dtype w = bb_data[3] - bb_data[1];
    const Dtype h = bb_data[4] - bb_data[2];

    int reference_size;
    if (this->phase_ == TRAIN)
    {
        // Choose the reference bounding box size within the given range

        // IMPORTANT! We need larger values to be less probable than small - because an accumulator, which
        // detects larger objects spans across more "integer" sizes. For example accumulator x2 can detect
        // bbs in [25,60], then x4 detects bbs from [50,120], which has double the size of the interval. The
        // random distribution needs to be augmented accordingly - we will project the bb sizes to log space
        double rsl_min = std::log(double(reference_size_min));
        double rsl_max = std::log(double(reference_size_max));

        boost::random::uniform_real_distribution<double> dists(0.0, 1.0);
        reference_size = std::round(std::exp(rsl_min + (dists(*rng) * (rsl_max-rsl_min))));
    }
    else
    {
        // Testing phase - keep the original size
        reference_size = std::max(w, h);

        if (reference_size < reference_size_min) reference_size = reference_size_min;
        if (reference_size > reference_size_max) reference_size = reference_size_max;
    }

    const Dtype size      = std::max(w, h);
    const int crop_width  = std::ceil(double(width) / reference_size * size);
    const int crop_height = std::ceil(double(height) / reference_size * size);

    // Select a random position of the crop, but it has to fully contain the bounding box
    int crop_x, crop_y;
    if (this->phase_ == TRAIN)
    {
        // In training we want to create a random crop around the bounding box
        boost::random::uniform_int_distribution<> distx(x+w-crop_width, x);
        boost::random::uniform_int_distribution<> disty(y+h-crop_height, y);
        crop_x = distx(*rng);
        crop_y = disty(*rng);
    }
    else
    {
        // In testing we want to always crop the same crop. Thus we will place the bounding box to the center
        // of the crop
        crop_x = x + w/2 - crop_width/2;
        crop_y = y + h/2 - crop_height/2;
    }

    // Now if the crop spans outside the image we have to pad the image
    int border_left   = (crop_x < 0) ? -crop_x : 0;
    int border_top    = (crop_y < 0) ? -crop_y : 0;
    int border_right  = (crop_x+crop_width  > cv_img.cols) ? (crop_x+crop_width  - cv_img.cols) : 0;
    int border_bottom = (crop_y+crop_height > cv_img.rows) ? (crop_y+crop_height - cv_img.rows) : 0;

    // -- CROP THE IMAGE -- //
    // Because we want to save memory we will not use the copyMakeBorder function to create the border, but
    // instead we will compute an intersection of the crop with the image, crop that intersection, scale it
    // and then place it on black background to the position in which it would be cropped from

    // Intersection of the crop with the image
    int ints_width  = crop_width - border_left - border_right;
    int ints_height = crop_height - border_top - border_bottom;
    cv::Rect intersection(crop_x+border_left, crop_y+border_top, ints_width, ints_height);

    // Scale the intersection according to the crip scaling ratio (do not round here because the crop could
    // escape from the network input size - conversion to int will floor it)
    int ints_width_scaled  = ints_width / size * reference_size;
    int ints_height_scaled = ints_height / size * reference_size;
    int border_left_scaled = border_left / size * reference_size;  // x coordinate of the crop
    int border_top_scaled  = border_top / size * reference_size;   // y coordinate of the crop
    // Corrections for imprecisions - we have to keep the crop inside of the network input image dimensions
    if (ints_width_scaled > width) ints_width_scaled = width;
    if (ints_height_scaled > height) ints_height_scaled = height;
    if (border_left_scaled+ints_width_scaled > width) border_left_scaled = width - ints_width_scaled;
    if (border_top_scaled+ints_height_scaled > height) border_top_scaled = height - ints_height_scaled;

    CHECK_GT(ints_width_scaled, 0) << "Crop does not intersect the image";
    CHECK_GT(ints_height_scaled, 0) << "Crop does not intersect the image";
    CHECK_LE(ints_width_scaled, width) << "Crop larger than width: " << ints_width_scaled;
    CHECK_LE(ints_height_scaled, height) << "Crop larger than height: " << ints_height_scaled;
    CHECK_LE(border_left_scaled+ints_width_scaled, width) << "Moved crop does not fit in the image";
    CHECK_LE(border_top_scaled+ints_height_scaled, height) << "Moved crop does not fit in the image";

    // Crop and scale down the cropped intersection
    cv::Mat cv_img_cropped_scaled;
    cv::resize(cv_img(intersection), cv_img_cropped_scaled, cv::Size(ints_width_scaled, ints_height_scaled));

    // Initialize the network input with black
    cv_img_cropped_out = cv::Mat::zeros(height, width, CV_8UC3);

    // Place the crop onto black canvas of the size of the network input image
    cv::Mat cv_img_subcrop = cv_img_cropped_out(cv::Rect(border_left_scaled, border_top_scaled,
                                                         ints_width_scaled, ints_height_scaled));
    cv_img_cropped_scaled.copyTo(cv_img_subcrop);

    // -- UPDATE THE COORDINATES OF THE BOUNDING BOXES -- //
    // Update the bounding box coordinates - we need to update all annotations
    Dtype x_scaling   = float(width) / crop_width;
    Dtype y_scaling   = float(height) / crop_height;
    for (int i = 0; i < MAX_NUM_BBS_PER_IMAGE; ++i)
    {
        // Data are stored like this [label, xmin, ymin, xmax, ymax, fblx, fbly, fbrx, fbry, rblx, rbly, ftly]
        Dtype *data = this->transformed_label_.mutable_cpu_data() + this->transformed_label_.offset(b, i);

        if (data[0] == Dtype(-1.0f)) break;

        // Align with x, y of the crop and scale
        data[1]  = (data[1]-crop_x) * x_scaling; // xmin
        data[2]  = (data[2]-crop_y) * y_scaling; // ymin
        data[3]  = (data[3]-crop_x) * x_scaling; // xmax
        data[4]  = (data[4]-crop_y) * y_scaling; // ymax
        data[5]  = (data[5]-crop_x) * x_scaling; // fblx
        data[6]  = (data[6]-crop_y) * y_scaling; // fbly
        data[7]  = (data[7]-crop_x) * x_scaling; // fbrx
        data[8]  = (data[8]-crop_y) * y_scaling; // fbry
        data[9]  = (data[9]-crop_x) * x_scaling; // rblx
        data[10] = (data[10]-crop_y) * y_scaling; // rbly
        data[11] = (data[11]-crop_y) * y_scaling; // ftly

//        cv::rectangle(cv_img_cropped_out, cv::Rect(data[1], data[2], data[3]-data[1], data[4]-data[2]), cv::Scalar(0,0,255), 2);
    }
//    static int imi = 0;
//    if(this->phase_ == TRAIN) cv::imwrite("cropped" + std::to_string(imi++) + ".png", cv_img_cropped_out);
}

template <typename Dtype>
void BB3TXTDataLayer<Dtype>::_applyPixelTransformationsAndCopyOut (const cv::Mat &cv_img_cropped, int b)
{
    const int width  = cv_img_cropped.cols;
    const int height = cv_img_cropped.rows;

    CHECK(cv_img_cropped.data) << "Something went wrong with cropping!";
    CHECK_EQ(height, this->transformed_data_.shape(2)) << "Wrong crop height! Does not match network!";
    CHECK_EQ(width,  this->transformed_data_.shape(3)) << "Wrong crop width! Does not match network!";

    // Distortions of the image color space - during testing those will be 0
    Dtype exposure         = Dtype(0.0f);
    std::vector<Dtype> bgr = {Dtype(0.0f), Dtype(0.0f), Dtype(0.0f)};
    cv::Mat noise          = cv::Mat::zeros(height, width, CV_32FC3);
    if (this->phase_ == TRAIN)
    {
        // Apply random transformations
        caffe::rng_t* rng = static_cast<caffe::rng_t*>(this->_rng->generator());

        boost::random::uniform_int_distribution<> diste(-40, 50);  // Exposure
        exposure = diste(*rng);

        boost::random::uniform_int_distribution<> distbgr(-30, 30);  // Hue
        bgr[0] = distbgr(*rng);
        bgr[1] = distbgr(*rng);
        bgr[2] = distbgr(*rng);

        boost::random::uniform_int_distribution<> distn(0, 25);  // Noise standard deviation
        cv::randn(noise, 0, distn(*rng));
    }

    // Normalize to 0 mean and unit variance and copy the image to the transformed_image
    // + apply exposure, noise, hue, saturation, ...
    Dtype* transformed_data = this->transformed_data_.mutable_cpu_data() + this->transformed_data_.offset(b);
    for (int i = 0; i < height; ++i)
    {
        const uchar* ptr  = cv_img_cropped.ptr<uchar>(i);
        const float* ptrn = noise.ptr<float>(i);
        int img_index = 0;  // Index in the cv_img_cropped
        for (int j = 0; j < width; ++j)
        {
            for (int c = 0; c < 3; ++c)
            {
                const int top_index = (c * height + i) * width + j;
                // Apply exposure change, hue distortion, noise
                const Dtype val = Dtype(ptr[img_index]) + exposure + bgr[c] + Dtype(ptrn[img_index]);
                // Zero mean and unit variance
                transformed_data[top_index] = (std::max(Dtype(0.0f), std::min(Dtype(255.0f), val))
                        - Dtype(128.0f)) / Dtype(128.0f);
                img_index++;
            }
        }
    }

//    std::vector<cv::Mat> chnls;
//    chnls.push_back(cv::Mat(height, width, CV_32FC1, transformed_image.mutable_cpu_data()+transformed_image.offset(0,0)));
//    chnls.push_back(cv::Mat(height, width, CV_32FC1, transformed_image.mutable_cpu_data()+transformed_image.offset(0,1)));
//    chnls.push_back(cv::Mat(height, width, CV_32FC1, transformed_image.mutable_cpu_data()+transformed_image.offset(0,2)));
//    cv::Mat img; cv::merge(chnls, img);
//    img *= 128.0f;
//    img += cv::Scalar(128.0f,128.0f,128.0f);
//    cv::Mat img_u; img.convertTo(img_u, CV_8UC3);
//    static int imi = 0;
//    cv::imwrite("transfromed" + std::to_string(imi++) + ".png", img_u);
}

template <typename Dtype>
SelectedBB<Dtype> BB3TXTDataLayer<Dtype>::_getImageAndBB()
{
    std::lock_guard<std::mutex> lock(this->_i_global_mtx);

    // Get image and bounding box index
    auto indices = this->_indices[this->_i_global++];

    if (this->_i_global >= this->_indices.size())
    {
        this->_i_global = 0;  // Restart
        if (this->phase_ == TRAIN) this->_shuffleBoundingBoxes();
    }

    SelectedBB<Dtype> sel;
    sel.filename = this->_images[indices.first].first;
    sel.label    = this->_images[indices.first].second;
    sel.bb_id    = indices.second;

    return sel;
}

// ----------------------------------------  LAYER INSTANTIATION  ---------------------------------------- //

INSTANTIATE_CLASS(BB3TXTDataLayer);
REGISTER_LAYER_CLASS(BB3TXTData);

}  // namespace caffe
#endif  // USE_OPENCV