/*
Points to Voxels & Voxels to Points (Modified from SparseConv)
Written by Li Jiang
All Rights Reserved 2020.
*/

#include "voxelize.h"

/* ================================== voxelize_idx ================================== */
template <Int dimension>
void voxelize_idx(/* long N*4 */ at::Tensor coords, /* long M*4 */ at::Tensor output_coords,
                  /* Int N */ at::Tensor input_map, /* Int M*(maxActive+1) */ at::Tensor output_map, Int batchSize, Int mode){
    assert(coords.ndimension() == 2);
    // assert(coords.size(1) >= dimension and coords.size(1) <= dimension + 1);

    RuleBook voxelizeRuleBook;  // rule[1]: M voxels -> N points  output_map
    SparseGrids<dimension> inputSGs; // voxel_coords -> voxel_idx in M voxels      input_map: N points -> M voxels
    Int nActive = 0;

    Int maxActive = voxelize_inputmap<dimension>(inputSGs, input_map.data<Int>(), voxelizeRuleBook, nActive, coords.data<long>(), coords.size(0), coords.size(1), batchSize, mode);

    output_map.resize_({nActive, maxActive + 1});
    output_map.zero_();

    output_coords.resize_({nActive, coords.size(1)});
    output_coords.zero_();

    Int *oM = output_map.data<Int>();
    long *oC = output_coords.data<long>();
    voxelize_outputmap<dimension>(coords.data<long>(), oC, oM, &voxelizeRuleBook[1][0], nActive, maxActive);
}


template <Int dimension>
void voxelize_outputmap(long *coords, long *output_coords, Int *output_map, Int *rule, Int nOutputRows, Int maxActive){
    for(Int i = 0; i < nOutputRows; i++){
        for(Int j = 0; j <= maxActive; j++)
            output_map[j] = rule[j];
        Int inputIdx = rule[1];
        rule += (1 + maxActive);
        output_map += (1 + maxActive);

        long *coord = coords + inputIdx * (dimension + 1);
        long *output_coord  = output_coords + i * (dimension + 1);
        for(Int j = 0; j <= dimension; j++){
            output_coord[j] = coord[j];
        }
    }
}

//mode 0=guaranteed unique 1=last item(overwrite) 2=first item(keep) 3=sum, 4=mean
//input: coords
//output: SGs: one map for each batch: map from voxel_coord to voxel_idx(in M voxels)
//output: input_map: N, N points -> M voxels
//output: rules
//output: nActive
//output: maxActive
template <Int dimension>
Int voxelize_inputmap(SparseGrids<dimension> &SGs, Int *input_map, RuleBook &rules, Int &nActive, long *coords, Int nInputRows, Int nInputColumns, Int batchSize, Int mode){
    assert(nActive == 0);
    assert(rules.size() == 0);
    assert(SGs.size() == 0);

    SGs.resize(batchSize);
    Point<dimension> p;

    std::vector<std::vector<Int>> outputRows;
    if(nInputColumns == dimension){
        SGs.resize(1);
        auto &sg = SGs[0];
        for(Int i = 0; i < nInputRows; i++){
            for(Int j = 0; j < dimension; j++)
                p[j] = coords[j];
            coords += dimension;
            auto iter = sg.mp.find(p);
            if (iter == sg.mp.end()){
                sg.mp[p] = nActive++;
                outputRows.resize(nActive);
            }
            outputRows[sg.mp[p]].push_back(i);

            input_map[i] = sg.mp[p];
        }
    }
    else{  // nInputColumns == dimension + 1 (1 in index 0 for batchidx)
        Int batchIdx;
        for(Int i = 0; i < nInputRows; i++){
            batchIdx = coords[0];
            for(Int j = 0; j < dimension; j++)
                p[j] = coords[j + 1];
            coords += (dimension + 1);
            if(batchIdx + 1 >= (Int)SGs.size()){
                SGs.resize(batchIdx + 1);
            }
            auto &sg = SGs[batchIdx];
            auto iter = sg.mp.find(p);
            if(iter == sg.mp.end()){
                sg.mp[p] = nActive++;
                outputRows.resize(nActive);
            }
            outputRows[sg.mp[p]].push_back(i);

            input_map[i] = sg.mp[p];
        }
    }

    // Rulebook Format
    // rules[0][0] == mode
    // rules[0][1] == maxActive per spatial location (==1 for modes 0,1,2)
    // rules[0][2] == nInputRows
    // rules[0][3] == nOutputRows
    // rules[1]   nOutputRows x (1+maxActive)
    rules.resize(2);
    rules[0].push_back(mode);
    rules[0].push_back(1);
    rules[0].push_back(nInputRows);
    rules[0].push_back(outputRows.size());
    auto &rule = rules[1];
    if(mode == 0){
        assert(nInputRows == (Int)outputRows.size());
        for(Int i = 0; i < nActive; i++){
            rule.push_back(1);
            assert((Int)outputRows[i].size() == 1);
            rule.push_back(outputRows[i][0]);
        }
    }
    if(mode == 1){
        for(Int i = 0; i < nActive; i++){
            rule.push_back(1);
            rule.push_back(outputRows[i].front());
        }
    }
    if(mode == 2){
        for(Int i = 0; i < nActive; i++){
            rule.push_back(1);
            rule.push_back(outputRows[i].back());
        }
    }
    Int maxActive = 1;
    if(mode == 3 or mode == 4){
        for(auto &row: outputRows)
            maxActive = std::max(maxActive, (Int)row.size());
        rules[0][1] = maxActive;
        for(auto &row: outputRows){
            rule.push_back(row.size());
            for(auto &r: row)
                rule.push_back(r);
            rule.resize((rule.size() + maxActive) / (maxActive + 1) * (maxActive + 1));
        }
    }
    return maxActive;
}


/* ================================== voxelize ================================== */
template <typename T>
void voxelize_fp(/* cuda float N*C */ at::Tensor feats, // N * 3 -> M * 3 (N >= M)
              /* cuda float M*C */ at::Tensor output_feats,
              /* cuda Int M*(maxActive+1) */ at::Tensor output_map, Int mode, Int nActive, Int maxActive, Int nPlane){

    auto iF = feats.data<T>();
    auto oF = output_feats.data<T>();

    Int *rules = output_map.data<Int>();

    voxelize_fp_cuda<T>(nActive, maxActive, nPlane, iF, oF, rules, mode==4);
}

template <typename T>
void voxelize_bp(/* cuda float M*C */ at::Tensor d_output_feats, /* cuda float N*C */ at::Tensor d_feats, /* cuda Int M*(maxActive+1) */ at::Tensor output_map,
            Int mode, Int nActive, Int maxActive, Int nPlane){
    auto d_oF = d_output_feats.data<T>();
    auto d_iF = d_feats.data<T>();

    Int *rules = output_map.data<Int>();

    voxelize_bp_cuda<T>(nActive, maxActive, nPlane, d_oF, d_iF, rules, mode==4);
}

/* ================================== point_recover ================================== */
template <typename T>
void point_recover_fp(/* cuda float M*C */ at::Tensor feats, /* cuda float N*C */ at::Tensor output_feats, /* cuda Int M*(maxActive+1) */ at::Tensor idx_map,
                Int nActive, Int maxActive, Int nPlane){
    auto iF = feats.data<T>();
    auto oF = output_feats.data<T>();

    Int *rules = idx_map.data<Int>();

    voxelize_bp_cuda<T>(nActive, maxActive, nPlane, iF, oF, rules, false);
}


template <typename T>
void point_recover_bp(/* cuda float N*C */ at::Tensor d_output_feats, /* cuda float M*C */ at::Tensor d_feats,  /* cuda Int M*(maxActive+1) */ at::Tensor idx_map,
                Int nActive, Int maxActive, Int nPlane){
    auto d_oF = d_output_feats.data<T>();
    auto d_iF = d_feats.data<T>();

    Int *rules = idx_map.data<Int>();

    voxelize_fp_cuda<T>(nActive, maxActive, nPlane, d_oF, d_iF, rules, false);
}
