// Subdivide a 2D grid of 3D points using stencil-based subdivision.
// Usage: subdivide_grid {num_iter} {input_file}

#include <iostream>
#include <string>
#include <vector>
#include <array>

#include "grid2DPoints3D.h"
#include "grid2DIO.h"
#include "writeGrid2DOFF.h"


// helper
// return a * p (scalar multiply).
static std::array<double, 3> Scale(double a, const std::array<double, 3>& p) {
    return {a * p[0], a * p[1], a * p[2]};
}
// return p + q (component-wise add).
static std::array<double, 3> Add(const std::array<double, 3>& p,
                                  const std::array<double, 3>& q)
{
    return {p[0] + q[0], p[1] + q[1], p[2] + q[2]};
}

// main subdivision function

void SubdivideI_Grid2D(const GRID2D_POINTS3D& grid2D,
                        GRID2D_POINTS3D& new_grid2D)
{
    const int n = grid2D.NumRows();
    const int m = grid2D.NumCols();
    const int new_n = 2 * n - 1;
    const int new_m = 2 * m - 1;

    new_grid2D.SetDimensions(new_n, new_m);

    auto P = [&](int i, int j) { return grid2D.Coord(i, j); };

    auto Set = [&](int I, int J, const std::array<double, 3>& v) {
        new_grid2D.SetCoord(I, J, v[0], v[1], v[2]);
    };

    // (a) Face centers: new vertex (2i+1, 2j+1)
    for (int i = 0; i < n - 1; i++) {
        for (int j = 0; j < m - 1; j++) {
            auto v = Scale(0.25, Add(Add(P(i, j), P(i, j + 1)),
                                     Add(P(i + 1, j), P(i + 1, j + 1))));
            Set(2 * i + 1, 2 * j + 1, v);
        }
    }

    // (b) Horizontal edge midpoints: new vertex (2i, 2j+1)
    for (int i = 0; i < n; i++) {
        for (int j = 0; j < m - 1; j++) {
            std::array<double, 3> v;
            if (i == 0 || i == n - 1) {
                // midpoint of edge endpoints.
                v = Scale(0.5, Add(P(i, j), P(i, j + 1)));
            }
            else {
                v = Add(Scale(6.0 / 16.0, Add(P(i, j), P(i, j + 1))),
                        Scale(1.0 / 16.0, Add(Add(P(i - 1, j), P(i - 1, j + 1)),
                                              Add(P(i + 1, j), P(i + 1, j + 1)))));
            }
            Set(2 * i, 2 * j + 1, v);
        }
    }

    // (c) Vertical edge midpoints: new vertex (2i+1, 2j)
    for (int i = 0; i < n - 1; i++) {
        for (int j = 0; j < m; j++) {
            std::array<double, 3> v;
            if (j == 0 || j == m - 1) {
                // midpoint of edge endpoints.
                v = Scale(0.5, Add(P(i, j), P(i + 1, j)));
            }
            else {
                v = Add(Scale(6.0 / 16.0, Add(P(i, j), P(i + 1, j))),
                        Scale(1.0 / 16.0, Add(Add(P(i, j - 1), P(i, j + 1)),
                                              Add(P(i + 1, j - 1), P(i + 1, j + 1)))));
            }
            Set(2 * i + 1, 2 * j, v);
        }
    }

    // (d/e.iii/e.iv/e.v) Original vertices: new vertex (2i, 2j)
    for (int i = 0; i < n; i++) {
        for (int j = 0; j < m; j++) {
            const bool on_top    = (i == 0);
            const bool on_bottom = (i == n - 1);
            const bool on_left   = (j == 0);
            const bool on_right  = (j == m - 1);

            std::array<double, 3> v;

            if ((on_top || on_bottom) && (on_left || on_right)) {
                // corner (e.iii): copy directly.
                v = P(i, j);
            }
            else if (on_left || on_right) {
                // left/right boundary, non-corner (e.iv): (1/8) row-adjacent, (3/4) center.
                v = Add(Scale(3.0 / 4.0, P(i, j)),
                        Scale(1.0 / 8.0, Add(P(i - 1, j), P(i + 1, j))));
            }
            else if (on_top || on_bottom) {
                // top/bottom boundary, non-corner (e.v): (1/8) column-adjacent, (3/4) center.
                v = Add(Scale(3.0 / 4.0, P(i, j)),
                        Scale(1.0 / 8.0, Add(P(i, j - 1), P(i, j + 1))));
            }
            else {
                // interior vertex (d): 9/16 center, 3/32 edge neighbors, 1/64 diagonal neighbors.
                v = Scale(9.0 / 16.0, P(i, j));
                v = Add(v, Scale(3.0 / 32.0, Add(Add(P(i - 1, j), P(i + 1, j)),
                                                   Add(P(i, j - 1), P(i, j + 1)))));
                v = Add(v, Scale(1.0 / 64.0, Add(Add(P(i - 1, j - 1), P(i - 1, j + 1)),
                                                   Add(P(i + 1, j - 1), P(i + 1, j + 1)))));
            }
            Set(2 * i, 2 * j, v);
        }
    }
}


// perform num_iter subdivision by calling SubdivideI_Grid2D repeatedly
// num_iter == 0 -> a copy of grid2D
void SubdivideMulti_Grid2D(const GRID2D_POINTS3D& grid2D,
                             GRID2D_POINTS3D& result,
                             int num_iter)
{
    if (num_iter == 0) {
        result = grid2D;
        return;
    }

    GRID2D_POINTS3D current = grid2D;
    GRID2D_POINTS3D next;

    for (int iter = 0; iter < num_iter; iter++) {
        SubdivideI_Grid2D(current, next);
        current = next;
    }

    result = current;
}


// io
/// construct output name
static std::string CreateOutputFilename(const std::string& infilename, const int num_iter) {
    const std::string grid2d_ext = ".grid2D";
    if (infilename.size() > grid2d_ext.size() &&
        infilename.substr(infilename.size() - grid2d_ext.size()) == grid2d_ext)
    {
        return infilename.substr(0, infilename.size() - grid2d_ext.size())
               + "-subdiv-" + std::to_string(num_iter) + ".grid2D";
    }
    return infilename + "-subdiv-" + std::to_string(num_iter) + ".grid2D";
}
static std::string CreateOffOutputFilename(const std::string& infilename, const int num_iter) {
    const std::string grid2d_ext = ".grid2D";
    if (infilename.size() > grid2d_ext.size() &&
        infilename.substr(infilename.size() - grid2d_ext.size()) == grid2d_ext)
    {
        return infilename.substr(0, infilename.size() - grid2d_ext.size())
               + "-subdiv-" + std::to_string(num_iter) + ".off";
    }
    return infilename + "-subdiv-" + std::to_string(num_iter) + ".off";
}

void UsageMsg(const std::string& command_name) {
    std::cerr << "Usage: " << command_name
              << " [-off] {num_iter} {input_file}\n";
}

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

    // Parse optional -off flag.
    bool write_off = false;
    int arg_start = 1;
    if (argc > 1 && std::string(argv[1]) == "-off") {
        write_off = true;
        arg_start = 2;
    }

    if (argc - arg_start != 2) {
        UsageMsg(argv[0]);
        return 0;
    }

    int num_iter;
    try {
        num_iter = std::stoi(argv[arg_start]);
    }
    catch (const std::exception&) {
        std::cerr << "Error: num_iter must be an integer.\n";
        UsageMsg(argv[0]);
        return -1;
    }

    if (num_iter < 0) {
        std::cerr << "Error: num_iter must be non-negative.\n";
        return -1;
    }

    const std::string infilename = argv[arg_start + 1];
    GRID2D_POINTS3D grid2D;

    try {
        OpenReadGrid2DPoints3D(infilename, grid2D);
    }
    catch (const std::exception&) {
        std::cerr << "Exiting...\n";
        return -1;
    }

    GRID2D_POINTS3D result;
    SubdivideMulti_Grid2D(grid2D, result, num_iter);

    // Write grid2D output.
    const std::string outfilename = CreateOutputFilename(infilename, num_iter);
    const std::vector<std::string> comments = {
        "Subdivided grid. Iterations: " + std::to_string(num_iter),
        "Input: " + infilename
    };

    try {
        std::cout << "Writing " << outfilename << ".\n";
        OpenWriteGrid2DPoints3D(outfilename, result, comments);
    }
    catch (const std::exception&) {
        std::cerr << "Exiting...\n";
        return -1;
    }

    // Optionally write .off output.
    if (write_off) {
        const std::string off_outfilename = CreateOffOutputFilename(infilename, num_iter);
        WRITE_GRID2D_OFF writeOFF;
        const std::vector<std::string> off_comments = {
            "Subdivided grid in OFF format. Iterations: " + std::to_string(num_iter),
            "Input: " + infilename
        };
        try {
            std::cout << "Writing " << off_outfilename << ".\n";
            writeOFF.OpenAndWrite(off_outfilename, result, off_comments);
        }
        catch (const std::exception&) {
            std::cerr << "Exiting...\n";
            return -1;
        }
    }

    return 0;
}