#include "contour_creator.hh"
#include <algorithm>
#include <array>
#include <gdal/ogr_api.h>
#include "ogrsf_frmts.h"
#include <gdal/ogr_core.h>
#include <gdal/ogr_feature.h>
#include <gdal/ogr_geometry.h>
#include <iostream>
#include <ostream>
#include <vector>
#include <cstdint>
#include <gdal/gdal.h>
#include "gdal/gdal_priv.h"
#include <gdal/gdal_frmts.h>
#include <omp.h>
#include "argh.h"

std::vector<Point> produce_cellmap(HeightMap* heightmap, float z)
{
    int length = (heightmap->width-1)*(heightmap->height-1); // Defining length of vector
    uint8_t *cells = (uint8_t *) CPLMalloc(sizeof(uint8_t)*length);
    std::vector<Point> points; // Initiating a vector of points
    for (int i = 0; i<length; i++) {
        int y = i/(heightmap->width-1);
        int x = i%(heightmap->width-1);
        uint8_t result = (heightmap->get_pixel(x,y)>z)*8 + 
        (heightmap->get_pixel(x+1,y)>z)*4 +
        (heightmap->get_pixel(x+1,y+1)>z)*2 +
        (heightmap->get_pixel(x,y+1)>z);
        if (result != 0 and result != 15 ) {
            points.push_back(Point(x, y, result));
        };
    }
    return points;
}

bool is_in(int value, std::array<int, 8> array)
{
    return std::binary_search(array.begin(), array.end(), value);
}

std::vector<std::vector<Point>> create_lines(HeightMap* heightmap, int interval)
{
    int num_contours = (heightmap->max - heightmap->min)/interval;
    std::vector<std::vector<Point>> vector_contours;

    #pragma omp parallel
    {
        std::vector<std::vector<Point>> vec_private;
        #pragma omp for
        for (int i = 1; i <= num_contours; i++) 
        {
            auto points = produce_cellmap(heightmap, heightmap->min + interval*i);
            std::vector<Point> line;
            int x;
            int y;
            int points_allocated = 0;
            x = points[0].x;
            y = points[0].y;
            int current_case = points[0].mscase;
            points[0].allocated = true;
            line.push_back(points[0]);
            for (int j = 0; j < points.size(); j++){
                for (int k = 0; k < points.size(); k++){
                    Point* candidate = &points[k];
                    std::array<int,8> north{4, 5,6,7,8,9,10,11};
                    std::array<int,8> south{1,2,5,6,9,10,13,14};
                    std::array<int,8> east{2,3,4,5,10,11,12,13};
                    std::array<int,8> west{1,3,5,7,8,10,12,14};
                    std::vector<int> possible_connections;
                    switch (candidate->mscase) {
                        case 1:
                        case 14:
                            possible_connections.insert(possible_connections.end(), east.begin(), east.end());
                            possible_connections.insert(possible_connections.end(), north.begin(), north.end());
                        case 2:
                        case 13:
                            possible_connections.insert(possible_connections.end(), west.begin(), west.end());
                            possible_connections.insert(possible_connections.end(), south.begin(), south.end());
                        case 3:
                        case 12:
                            possible_connections.insert(possible_connections.end(), west.begin(), west.end());
                            possible_connections.insert(possible_connections.end(), east.begin(), east.end());
                        case 4:
                        case 11:
                            possible_connections.insert(possible_connections.end(), west.begin(), west.end());
                            possible_connections.insert(possible_connections.end(), north.begin(), north.end());
                        case 5:
                        case 10:
                            possible_connections.insert(possible_connections.end(), north.begin(), south.end());
                            possible_connections.insert(possible_connections.end(), south.begin(), south.end());
                            possible_connections.insert(possible_connections.end(), west.begin(), west.end());
                            possible_connections.insert(possible_connections.end(), east.begin(), east.end());
                        case 6:
                        case 9:
                            possible_connections.insert(possible_connections.end(), north.begin(), south.end());
                            possible_connections.insert(possible_connections.end(), south.begin(), south.end());
                        case 7:
                        case 8:
                            possible_connections.insert(possible_connections.end(), north.begin(), south.end());
                            possible_connections.insert(possible_connections.end(), east.begin(), east.end());
                    }

                    if (!candidate->allocated) {
                        if (candidate->x == x +1 && candidate->y == y && is_in(candidate->mscase, west) && is_in(current_case, east)) {
                            candidate->allocated = true;
                            x = candidate->x;
                            y = candidate->y;
                            current_case = candidate->mscase;
                            line.push_back(*candidate);
                            points_allocated++;
                            break;
                        }
                        else if (candidate->x == x -1 && candidate->y == y && is_in(candidate->mscase, east ) && is_in(current_case, west)) {
                            x = candidate->x;
                            y = candidate->y;
                            current_case = candidate->mscase;
                            candidate->allocated = true;
                            line.push_back(*candidate);
                            points_allocated++;
                            break;
                        }
                        else if (candidate->x == x && candidate->y == y + 1 && is_in(candidate->mscase, south) && is_in(current_case, north)) {
                            x = candidate->x;
                            y = candidate->y;
                            current_case = candidate->mscase;
                            candidate->allocated = true;
                            line.push_back(*candidate);
                            points_allocated++;
                            break;
                        }
                        else if (candidate->x == x && candidate->y == y - 1 && is_in(candidate->mscase, north) && is_in(current_case, south)) {
                            x = candidate->x;
                            y = candidate->y;
                            current_case = candidate->mscase;
                            candidate->allocated = true;
                            line.push_back(*candidate);
                            points_allocated++;
                            break;
                        }
                    }
                }
                if (j > points_allocated)
                {
                    vec_private.push_back(line);
                    line.clear();
                    //std::cout << points.size() << " ";
                    for (int k = 0; k < points.size(); k++){
                        Point* candidate = &points[k];
                        if (!candidate->allocated)
                        {
                            line.push_back(*candidate);
                            x = candidate->x;
                            y = candidate->y;
                            points_allocated++;
                            break;
                        }
                    }
                }
            }
        }
        #pragma omp critical

        vector_contours.insert(vector_contours.end(), vec_private.begin(), vec_private.end());
    }
        return vector_contours;
}

std::tuple<double, double> local_to_projected(double* geotransform, int x, int y)
{
    return {
        geotransform[1] * x + geotransform[2] * y +
        geotransform[1] * 0.5 + geotransform[2] * 0.5 + geotransform[0],
        geotransform[4] * x + geotransform[5] * y +
        geotransform[4] * 0.5 + geotransform[5] * 0.5 + geotransform[3]
        };
}

constexpr std::tuple<double, double, double, double> marching_squares_lookup(int mcase)
{
    switch (mcase) {
        case 1: return {0, 0.5, 0.5, 0};
        case 2: return {1, 0.5, 0.5, 0};
        case 3: return {0, 0.5, 1, 0.5};
        case 4: return {0.5, 1, 1, 0.5};
        case 5: return {0, 0, 0, 0}; //FIXME
        case 6: return {0.5, 1, 0.5, 0};
        case 7: return {0, 0.5, 0.5, 1};
        case 8: return {0, 0.5, 0.5, 1};
        case 9: return {0.5, 1, 0.5, 0};
        case 10: return {0, 0, 0, 0}; //FIXME
        case 11: return {0.5, 1, 1, 0.5};
        case 12: return {0, 0.5, 1, 0.5};
        case 13: return {0.5, 0, 1, 0.5};
        case 14: return {0, 0.5, 0.5, 0};
    };
}
void write_output_file(std::vector<std::vector<Point>> all_points, const char *filepath, HeightMap* heightmap)
{

    const char *pszDriverName = "GeoJSON";
    GDALDriver *poDriver;

    GDALAllRegister();

    poDriver = GetGDALDriverManager()->GetDriverByName(pszDriverName );
    if( poDriver == NULL )
    {
        printf( "%s driver not available.\n", pszDriverName );
        exit( 1 );
    }

    GDALDataset *poDS;

    poDS = poDriver->Create( filepath, 0, 0, 0, GDT_Unknown, NULL );
    if( poDS == NULL )
    {
        printf( "Creation of output file failed.\n" );
        exit( 1 );
    }

    OGRLayer *poLayer;

    poLayer = poDS->CreateLayer( "contours", &heightmap->reference_system, wkbLineString, NULL );
    if( poLayer == NULL )
    {
        printf( "Layer creation failed.\n" );
        exit( 1 );
    }

    OGRFieldDefn oField( "Name", OFTString );

    oField.SetWidth(32);

    if( poLayer->CreateField( &oField ) != OGRERR_NONE )
    {
        printf( "Creating Name field failed.\n" );
        exit( 1 );
    }
    int width = heightmap->width -1;
    int height = heightmap->height -1;
    int size = width * height;

    for (int j = 0; j < all_points.size(); j++)
    {
        OGRFeature *feature;
        OGRLineString *geometry = new OGRLineString();
        feature = OGRFeature::CreateFeature( poLayer->GetLayerDefn() );
        feature->SetField( "Name", j );

        std::vector<Point> points = all_points[j];

        for (int k = 0; k < points.size(); k++)
        {
            int x_raw = points[k].x;
            int y_raw = points[k].y;
            auto [x, y] = local_to_projected(heightmap->geotransform, x_raw, y_raw);
            //auto [x_bl, y_bl, x_tr, y_tr] = marching_squares_lookup(*(cellmap->cells + i) );
            geometry->setPoint(geometry->getNumPoints(),x ,y );
        }

        if ( feature->SetGeometry(geometry) != OGRERR_NONE)
        {
            printf( "Failed to set geometry.\n" );
            exit( 1 );
        }
        OGRGeometryFactory::destroyGeometry(geometry);
        if( poLayer->CreateFeature( feature ) != OGRERR_NONE )
        {
            printf( "Failed to create feature in shapefile.\n" );
            exit( 1 );
        }
        OGRFeature::DestroyFeature( feature );
    }
    GDALClose( poDS );
}

int main(int argc, const char* argv[])
{
    argh::parser cmdl(argv);

    if (cmdl[{ "-h", "--help" }])
    {
        std::cout << "Usage:\n"
            << "contour_creator [OPTIONS] <input_file>\n"
            << "-o; --output <FILENAME.geojson> - File to write output to (Default: contours.geojson)\n"
            << "-i; --interval <int> - Set the interval between contours (Default: 5)\n"
            << "-b; --blur - Blur the image\n"
            << "--stats - Print statistical information about the heightmap\n";
            exit(0);
    }
    int interval;
    cmdl({"-i", "--interval"}, 5) >> interval;
    if (interval <= 0)
        std::cerr << "Interval must be valid positive integer!" << "\n";
    
    std::string output_file;
    cmdl({"-o", "--output"}, "contours.geojson") >> output_file;

    std::string input_file;
    cmdl(1) >> input_file;
    

    HeightMap map(input_file.c_str());
    if (cmdl[{"-b", "--blur"}])
        map.blur(0.8);
    auto lines = create_lines(&map, interval);
    write_output_file(lines, output_file.c_str(), &map);
    std::cout << "Contours written to " << output_file << " 🗺️\n";
}