// ======================================================================================
// Copyright 2017 State Key Laboratory of Remote Sensing Science, 
// Institute of Remote Sensing Science and Engineering, Beijing Normal University

// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at

//     http://www.apache.org/licenses/LICENSE-2.0

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

#include "CSF.h"

// #######################################################################################
class Rasterization
{
public:
    Rasterization() {}
    ~Rasterization() {}

    // for a cloth particle, if no corresponding lidar point are found.
    // the heightval are set as its neighbor's
    double static findHeightValByNeighbor(Particle *p)
    {
        std::queue<Particle *>  nqueue;
        std::vector<Particle *> pbacklist;

        size_t neiborsize = p->neighborsList.size();

        for(size_t i=0; i < neiborsize; i++)
        {
            p->isVisited = true;
            nqueue.push(p->neighborsList[i]);
        }

        // iterate over the nqueue
        while( !nqueue.empty() )
        {
            Particle *pneighbor = nqueue.front();
            nqueue.pop();
            pbacklist.push_back(pneighbor);

            if( pneighbor->nearestPointHeight > MIN_INF )
            {
                for (std::size_t i = 0; i < pbacklist.size(); i++)
                    pbacklist[i]->isVisited = false;

                while (!nqueue.empty()) {
                    Particle *pp = nqueue.front();
                    pp->isVisited = false;
                    nqueue.pop();
                }

                return pneighbor->nearestPointHeight;
            }
            else
            {
                size_t nsize = pneighbor->neighborsList.size();

                for(size_t i=0; i<nsize; i++)
                {
                    Particle *ptmp = pneighbor->neighborsList[i];

                    if( !ptmp->isVisited )
                    {
                        ptmp->isVisited = true;
                        nqueue.push(ptmp);
                    }
                }
            }
        }

        return MIN_INF;
    }

    double static findHeightValByScanline(Particle *p, Cloth &cloth)
    {
        int xpos = p->pos_x;
        int ypos = p->pos_y;

        for (int i = xpos + 1; i < cloth.num_particles_width; i++) {
            double crresHeight = cloth.getParticle(i, ypos)->nearestPointHeight;

            if (crresHeight > MIN_INF)
                return crresHeight;
        }

        for (int i = xpos - 1; i >= 0; i--) {
            double crresHeight = cloth.getParticle(i, ypos)->nearestPointHeight;

            if (crresHeight > MIN_INF)
                return crresHeight;
        }

        for (int j = ypos - 1; j >= 0; j--) {
            double crresHeight = cloth.getParticle(xpos, j)->nearestPointHeight;

            if (crresHeight > MIN_INF)
                return crresHeight;
        }

        for (int j = ypos + 1; j < cloth.num_particles_height; j++) {
            double crresHeight = cloth.getParticle(xpos, j)->nearestPointHeight;

            if (crresHeight > MIN_INF)
                return crresHeight;
        }

        return findHeightValByNeighbor(p);
    }

    void static RasterTerrian(Cloth &cloth, CSG_PointCloud &pc, std::vector<double> & heightVal)
    {
        #define SQUARE_DIST(x1, y1, x2, y2) (((x1) - (x2)) * ((x1) - (x2)) + ((y1) - (y2)) * ((y1) - (y2)))

        for(sLong i=0; i <pc.Get_Count(); i++)
        {
            double pc_x = pc.Get_X(i);
            double pc_z = pc.Get_Y(i);

            double deltaX = pc_x - cloth.origin_pos.f[0];
            double deltaZ = pc_z - cloth.origin_pos.f[2];
            int    col    = int(deltaX / cloth.step_x + 0.5);
            int    row    = int(deltaZ / cloth.step_y + 0.5);

            if ((col >= 0) && (row >= 0)) {
                Particle *pt = cloth.getParticle(col, row);
                pt->correspondingLidarPointList.push_back(i);
                double pc2particleDist = SQUARE_DIST(
                    pc_x, pc_z,
                    pt->getPos().f[0],
                    pt->getPos().f[2]
                );

                if (pc2particleDist < pt->tmpDist) {
                    pt->tmpDist            = pc2particleDist;
                    pt->nearestPointHeight = -pc.Get_Z(i);
                    pt->nearestPointIndex  = i;
                }
            }
        }
        heightVal.resize(cloth.getSize());

        // #ifdef CSF_USE_OPENMP
        // #pragma omp parallel for
        // #endif
        for (int i = 0; i < cloth.getSize(); i++) {
            Particle *pcur          = cloth.getParticle1d(i);
            double    nearestHeight = pcur->nearestPointHeight;

            if (nearestHeight > MIN_INF) {
                heightVal[i] = nearestHeight;
            } else {
                heightVal[i] = findHeightValByScanline(pcur, cloth);
            }
        }
    }
};

// #include "Rasterization.h"
// #######################################################################################


// #######################################################################################
// #include "c2cdist.h"
class c2cdist
{
public:

    c2cdist(double threshold) : class_treshold(threshold) {}

    void calCloud2CloudDist(Cloth &cloth, CSG_PointCloud &pc, std::vector<int>& groundIndexes, std::vector<int>& offGroundIndexes)
    {
        groundIndexes.resize(0);
        offGroundIndexes.resize(0);

        for(sLong i=0; i<pc.Get_Count(); i++)
        {
            double pc_x = pc.Get_X(i);
            double pc_z = pc.Get_Y(i);

            double deltaX = pc_x - cloth.origin_pos.f[0];
            double deltaZ = pc_z - cloth.origin_pos.f[2];

            int col0 = int(deltaX / cloth.step_x);
            int row0 = int(deltaZ / cloth.step_y);
            int col1 = col0 + 1;
            int row1 = row0;
            int col2 = col0 + 1;
            int row2 = row0 + 1;
            int col3 = col0;
            int row3 = row0 + 1;

            double subdeltaX = (deltaX - col0 * cloth.step_x) / cloth.step_x;
            double subdeltaZ = (deltaZ - row0 * cloth.step_y) / cloth.step_y;

            double fxy
                = cloth.getParticle(col0, row0)->pos.f[1] * (1 - subdeltaX) * (1 - subdeltaZ) +
                cloth.getParticle(col3, row3)->pos.f[1] * (1 - subdeltaX) * subdeltaZ +
                cloth.getParticle(col2, row2)->pos.f[1] * subdeltaX * subdeltaZ +
                cloth.getParticle(col1, row1)->pos.f[1] * subdeltaX * (1 - subdeltaZ);
            double height_var = fxy - (-pc.Get_Z(i));

            if (std::fabs(height_var) < class_treshold) {
                groundIndexes.push_back(i);
            } else {
                offGroundIndexes.push_back(i);
            }
        }
    }

private:

    double class_treshold;

};

// #include "c2cdist.h"
// #######################################################################################


CSF::CSF()
{
    params.bSlopeSmooth     = true;
    params.time_step        = 0.65;
    params.class_threshold  = 0.5;
    params.height_threshold = 0.3;
    params.cloth_resolution = 1;
    params.rigidness        = 3;
    params.iterations       = 500;
    this->index             = 0;
}

CSF::CSF(int index)
{
    params.bSlopeSmooth     = true;
    params.time_step        = 0.65;
    params.class_threshold  = 0.5;
    params.height_threshold = 0.3;
    params.cloth_resolution = 1;
    params.rigidness        = 3;
    params.iterations       = 500;
    this->index             = index;
}

Cloth CSF::do_cloth()
{
    // Terrain
    //std::cout << "[" << this->index << "] Configuring terrain..." << std::endl;
    //csf::Point bbMin, bbMax;
    //point_cloud.computeBoundingBox(bbMin, bbMax);
    //CSG_Point_3D bbMin(pPoints->Get_Extent().xMin, pPoints->Get_Extent().yMin, pPoints->Get_ZMin());
    //CSG_Point_3D bbMax(pPoints->Get_Extent().xMax, pPoints->Get_Extent().yMax, pPoints->Get_ZMax());
    CSG_Point_3D bbMin(pPoints->Get_Extent().xMin, -pPoints->Get_ZMax(), pPoints->Get_Extent().yMin);
    CSG_Point_3D bbMax(pPoints->Get_Extent().xMax, -pPoints->Get_ZMin(), pPoints->Get_Extent().yMax);
    //std::cout << "[" << this->index << "]  - bbMin: " << bbMin.x << " " << bbMin.y << " " << bbMin.z << std::endl;
    //std::cout << "[" << this->index << "]  - bbMax: " << bbMax.x << " " << bbMax.y << " " << bbMax.z << std::endl;

    double cloth_y_height = 0.05;

    int clothbuffer_d = 2;
    Vec3 origin_pos(
        bbMin.x - clothbuffer_d *params.cloth_resolution,
        bbMax.y + cloth_y_height,
        bbMin.z - clothbuffer_d *params.cloth_resolution
    );

    int width_num = static_cast<int>(
        std::floor((bbMax.x - bbMin.x) / params.cloth_resolution)
    ) + 2 * clothbuffer_d;

    int height_num = static_cast<int>(
        std::floor((bbMax.z - bbMin.z) / params.cloth_resolution)
    ) + 2 * clothbuffer_d;

    std::cout << "[" << this->index << "] Configuring cloth..." << std::endl;
    std::cout << "[" << this->index << "]  - width: " << width_num << " "
         << "height: " << height_num << std::endl;

    Cloth cloth(
        origin_pos,
        width_num,
        height_num,
        params.cloth_resolution,
        params.cloth_resolution,
        params.height_threshold,
        9999,
        params.rigidness,
        params.time_step
    );

    std::cout << "[" << this->index << "] Rasterizing..." << std::endl;
    Rasterization::RasterTerrian(cloth, *pPoints, cloth.getHeightvals());

    double time_step2 = params.time_step * params.time_step;
    double gravity    = 0.2;

    std::cout << "[" << this->index << "] Simulating..." << std::endl;
    cloth.addForce(Vec3(0, -gravity, 0) * time_step2);

    // boost::progress_display pd(params.iterations);
    for(int i=0; i<params.iterations && SG_UI_Process_Set_Progress(i, params.iterations); i++)
    {
        double maxDiff = cloth.timeStep();
        cloth.terrCollision();
		//params.class_threshold / 100
        if ((maxDiff != 0) && (maxDiff < 0.005)) {
            // early stop
            break;
        }
        // pd++;
    }

    if (params.bSlopeSmooth) {
        std::cout << "[" << this->index << "]  - post handle..." << std::endl;
        cloth.movableFilter();
    }

    return cloth;
}

std::vector<double> CSF::do_cloth_export() {
    auto cloth = do_cloth();
    return cloth.toVector();
}

void CSF::do_filtering(std::vector<int>& groundIndexes, std::vector<int>& offGroundIndexes)//, bool exportCloth)
{
    auto cloth = do_cloth();
    //if (exportCloth)
    //    cloth.saveToFile();
    c2cdist c2c(params.class_threshold);
    c2c.calCloud2CloudDist(cloth, *pPoints, groundIndexes, offGroundIndexes);
}


// #######################################################################################
// #include "cloth.h"

Cloth::Cloth(const Vec3& _origin_pos,
    int         _num_particles_width,
    int         _num_particles_height,
    double      _step_x,
    double      _step_y,
    double      _smoothThreshold,
    double      _heightThreshold,
    int         rigidness,
    double      time_step)
    : constraint_iterations(rigidness),
    time_step(time_step),
    smoothThreshold(_smoothThreshold),
    heightThreshold(_heightThreshold),
    origin_pos(_origin_pos),
    step_x(_step_x),
    step_y(_step_y),
    num_particles_width(_num_particles_width),
    num_particles_height(_num_particles_height) {
    // I am essentially using this vector as an array with room for
    // num_particles_width*num_particles_height particles
    particles.resize(num_particles_width * num_particles_height);

    double time_step2 = time_step * time_step;

    // creating particles in a grid of particles from (0,0,0) to
    // (width,-height,0) creating particles in a grid
    for (int i = 0; i < num_particles_width; i++) {
        for (int j = 0; j < num_particles_height; j++) {
            Vec3 pos(origin_pos.f[0] + i *step_x,
                origin_pos.f[1],
                origin_pos.f[2] + j *step_y);

            // insert particle in column i at j'th row
            particles[j * num_particles_width + i]       = Particle(pos, time_step2);
            particles[j * num_particles_width + i].pos_x = i;
            particles[j * num_particles_width + i].pos_y = j;
        }
    }

    // Connecting immediate neighbor particles with constraints
    // (distance 1 and sqrt(2) in the grid)
    for (int x = 0; x < num_particles_width; x++) {
        for (int y = 0; y < num_particles_height; y++) {
            if (x < num_particles_width - 1)
                makeConstraint(getParticle(x, y), getParticle(x + 1, y));

            if (y < num_particles_height - 1)
                makeConstraint(getParticle(x, y), getParticle(x, y + 1));

            if ((x < num_particles_width - 1) && (y < num_particles_height - 1))
                makeConstraint(getParticle(x, y), getParticle(x + 1, y + 1));

            if ((x < num_particles_width - 1) && (y < num_particles_height - 1))
                makeConstraint(getParticle(x + 1, y), getParticle(x, y + 1));
        }
    }

    // Connecting secondary neighbors with constraints (distance 2 and sqrt(4) in the grid)
    for (int x = 0; x < num_particles_width; x++) {
        for (int y = 0; y < num_particles_height; y++) {
            if (x < num_particles_width - 2)
                makeConstraint(getParticle(x, y), getParticle(x + 2, y));

            if (y < num_particles_height - 2)
                makeConstraint(getParticle(x, y), getParticle(x, y + 2));

            if ((x < num_particles_width - 2) && (y < num_particles_height - 2))
                makeConstraint(getParticle(x, y), getParticle(x + 2, y + 2));

            if ((x < num_particles_width - 2) && (y < num_particles_height - 2))
                makeConstraint(getParticle(x + 2, y), getParticle(x, y + 2));
        }
    }
}

double Cloth::timeStep() {
    int particleCount = static_cast<int>(particles.size());
    #pragma omp parallel for
    for (int i = 0; i < particleCount; i++) {
        particles[i].timeStep();
    }

    #pragma omp parallel for
    for (int j = 0; j < particleCount; j++) {
        particles[j].satisfyConstraintSelf(constraint_iterations);
    }

    double maxDiff = 0;

    for (int i = 0; i < particleCount; i++) {
        if (particles[i].isMovable()) {
            double diff = fabs(particles[i].old_pos.f[1] - particles[i].pos.f[1]);

            if (diff > maxDiff)
                maxDiff = diff;
        }
    }

    return maxDiff;
}

void Cloth::addForce(const Vec3 direction) {
    for (std::size_t i = 0; i < particles.size(); i++) {
        particles[i].addForce(direction);
    }
}

void Cloth::terrCollision() {
    int particleCount = static_cast<int>(particles.size());

    #pragma omp parallel for
    for (int i = 0; i < particleCount; i++) {
        Vec3 v = particles[i].getPos();

        if (v.f[1] < heightvals[i]) {
            particles[i].offsetPos(Vec3(0, heightvals[i] - v.f[1], 0));
            particles[i].makeUnmovable();
        }
    }
}

void Cloth::movableFilter() {
    std::vector<Particle> tmpParticles;

    for (int x = 0; x < num_particles_width; x++) {
        for (int y = 0; y < num_particles_height; y++) {
            Particle *ptc = getParticle(x, y);

            if (ptc->isMovable() && !ptc->isVisited) {
                std::queue<int> que;
                std::vector<XY> connected; // store the connected component
                std::vector<std::vector<int> > neibors;
                int sum   = 1;
                int index = y * num_particles_width + x;

                // visit the init node
                connected.push_back(XY(x, y));
                particles[index].isVisited = true;

                // enqueue the init node
                que.push(index);

                while (!que.empty()) {
                    Particle *ptc_f = &particles[que.front()];
                    que.pop();
                    int cur_x = ptc_f->pos_x;
                    int cur_y = ptc_f->pos_y;
                    std::vector<int> neibor;

                    if (cur_x > 0) {
                        Particle *ptc_left = getParticle(cur_x - 1, cur_y);

                        if (ptc_left->isMovable()) {
                            if (!ptc_left->isVisited) {
                                sum++;
                                ptc_left->isVisited = true;
                                connected.push_back(XY(cur_x - 1, cur_y));
                                que.push(num_particles_width * cur_y + cur_x - 1);
                                neibor.push_back(sum - 1);
                                ptc_left->c_pos = sum - 1;
                            } else {
                                neibor.push_back(ptc_left->c_pos);
                            }
                        }
                    }

                    if (cur_x < num_particles_width - 1) {
                        Particle *ptc_right = getParticle(cur_x + 1, cur_y);

                        if (ptc_right->isMovable()) {
                            if (!ptc_right->isVisited) {
                                sum++;
                                ptc_right->isVisited = true;
                                connected.push_back(XY(cur_x + 1, cur_y));
                                que.push(num_particles_width * cur_y + cur_x + 1);
                                neibor.push_back(sum - 1);
                                ptc_right->c_pos = sum - 1;
                            } else {
                                neibor.push_back(ptc_right->c_pos);
                            }
                        }
                    }

                    if (cur_y > 0) {
                        Particle *ptc_bottom = getParticle(cur_x, cur_y - 1);

                        if (ptc_bottom->isMovable()) {
                            if (!ptc_bottom->isVisited) {
                                sum++;
                                ptc_bottom->isVisited = true;
                                connected.push_back(XY(cur_x, cur_y - 1));
                                que.push(num_particles_width * (cur_y - 1) + cur_x);
                                neibor.push_back(sum - 1);
                                ptc_bottom->c_pos = sum - 1;
                            } else {
                                neibor.push_back(ptc_bottom->c_pos);
                            }
                        }
                    }

                    if (cur_y < num_particles_height - 1) {
                        Particle *ptc_top = getParticle(cur_x, cur_y + 1);

                        if (ptc_top->isMovable()) {
                            if (!ptc_top->isVisited) {
                                sum++;
                                ptc_top->isVisited = true;
                                connected.push_back(XY(cur_x, cur_y + 1));
                                que.push(num_particles_width * (cur_y + 1) + cur_x);
                                neibor.push_back(sum - 1);
                                ptc_top->c_pos = sum - 1;
                            } else {
                                neibor.push_back(ptc_top->c_pos);
                            }
                        }
                    }
                    neibors.push_back(neibor);
                }

                if (sum > MAX_PARTICLE_FOR_POSTPROCESSIN) {
                    std::vector<int> edgePoints = findUnmovablePoint(connected);
                    handle_slop_connected(edgePoints, connected, neibors);
                }
            }
        }
    }
}

std::vector<int> Cloth::findUnmovablePoint(std::vector<XY> connected) {
    std::vector<int> edgePoints;

    for (std::size_t i = 0; i < connected.size(); i++) {
        int x         = connected[i].x;
        int y         = connected[i].y;
        int index     = y * num_particles_width + x;
        Particle *ptc = getParticle(x, y);

        if (x > 0) {
            Particle *ptc_x = getParticle(x - 1, y);

            if (!ptc_x->isMovable()) {
                int index_ref = y * num_particles_width + x - 1;

                if ((fabs(heightvals[index] - heightvals[index_ref]) < smoothThreshold) &&
                    (ptc->getPos().f[1] - heightvals[index] < heightThreshold)) {
                    Vec3 offsetVec = Vec3(0, heightvals[index] - ptc->getPos().f[1], 0);
                    particles[index].offsetPos(offsetVec);
                    ptc->makeUnmovable();
                    edgePoints.push_back(i);
                    continue;
                }
            }
        }

        if (x < num_particles_width - 1) {
            Particle *ptc_x = getParticle(x + 1, y);

            if (!ptc_x->isMovable()) {
                int index_ref = y * num_particles_width + x + 1;

                if ((fabs(heightvals[index] - heightvals[index_ref]) < smoothThreshold) &&
                    (ptc->getPos().f[1] - heightvals[index] < heightThreshold)) {
                    Vec3 offsetVec = Vec3(0, heightvals[index] - ptc->getPos().f[1], 0);
                    particles[index].offsetPos(offsetVec);
                    ptc->makeUnmovable();
                    edgePoints.push_back(i);
                    continue;
                }
            }
        }

        if (y > 0) {
            Particle *ptc_y = getParticle(x, y - 1);

            if (!ptc_y->isMovable()) {
                int index_ref = (y - 1) * num_particles_width + x;

                if ((fabs(heightvals[index] - heightvals[index_ref]) < smoothThreshold) &&
                    (ptc->getPos().f[1] - heightvals[index] < heightThreshold)) {
                    Vec3 offsetVec = Vec3(0, heightvals[index] - ptc->getPos().f[1], 0);
                    particles[index].offsetPos(offsetVec);
                    ptc->makeUnmovable();
                    edgePoints.push_back(i);
                    continue;
                }
            }
        }

        if (y < num_particles_height - 1) {
            Particle *ptc_y = getParticle(x, y + 1);

            if (!ptc_y->isMovable()) {
                int index_ref = (y + 1) * num_particles_width + x;

                if ((fabs(heightvals[index] - heightvals[index_ref]) < smoothThreshold) &&
                    (ptc->getPos().f[1] - heightvals[index] < heightThreshold)) {
                    Vec3 offsetVec = Vec3(0, heightvals[index] - ptc->getPos().f[1], 0);
                    particles[index].offsetPos(offsetVec);
                    ptc->makeUnmovable();
                    edgePoints.push_back(i);
                    continue;
                }
            }
        }
    }

    return edgePoints;
}

void Cloth::handle_slop_connected(std::vector<int> edgePoints, std::vector<XY> connected, std::vector<std::vector<int> > neibors) {
    std::vector<bool> visited;

    for (std::size_t i = 0; i < connected.size(); i++) visited.push_back(false);

    std::queue<int> que;

    for (std::size_t i = 0; i < edgePoints.size(); i++) {
        que.push(edgePoints[i]);
        visited[edgePoints[i]] = true;
    }

    while (!que.empty()) {
        int index = que.front();
        que.pop();

        int index_center = connected[index].y * num_particles_width + connected[index].x;

        for (std::size_t i = 0; i < neibors[index].size(); i++) {
            int index_neibor = connected[neibors[index][i]].y * num_particles_width + connected[neibors[index][i]].x;

            if ((fabs(heightvals[index_center] - heightvals[index_neibor]) < smoothThreshold) &&
                (fabs(particles[index_neibor].getPos().f[1] - heightvals[index_neibor]) < heightThreshold)) {
                Vec3 offsetVec = Vec3(0, heightvals[index_neibor] - particles[index_neibor].getPos().f[1], 0);
                particles[index_neibor].offsetPos(offsetVec);
                particles[index_neibor].makeUnmovable();

                if (visited[neibors[index][i]] == false) {
                    que.push(neibors[index][i]);
                    visited[neibors[index][i]] = true;
                }
            }
        }
    }
}

std::vector<double> Cloth::toVector() {
    std::vector<double> clothCoordinates;
    clothCoordinates.reserve(particles.size()*3);
    for(auto& particle : particles) {
        clothCoordinates.push_back(particle.getPos().f[0]);
        clothCoordinates.push_back(particle.getPos().f[2]);
        clothCoordinates.push_back(-particle.getPos().f[1]);
    }
    return clothCoordinates;
}
