#include <vtkActor.h>
#include <vtkCamera.h>
#include <vtkCellArrayIterator.h>
#include <vtkCellData.h>
#include <vtkCleanPolyData.h>
#include <vtkDiscretizableColorTransferFunction.h>
#include <vtkImageData.h>
#include <vtkInteractorStyleTrackballCamera.h>
#include <vtkLODActor.h>
#include <vtkLookupTable.h>
#include <vtkMath.h>
#include <vtkPiecewiseFunction.h>
#include <vtkPolyDataMapper.h>
#include <vtkPolyDataWriter.h>
#include <vtkProperty.h>
#include <vtkRenderWindow.h>
#include <vtkRenderWindowInteractor.h>
#include <vtkRenderer.h>
#include <vtkSurfaceNets3D.h>
#include <vtkTimerLog.h>
#include <vtkXMLImageDataReader.h>

#include <array>
#include <map>
#include <set>
#include <sstream>

int main( int argc, char** argv )
{
    // check arguments part
    if( argc <= 3 ) {
        std::cout << "Need at least 3 arguments :" << std::endl;
        std::cout << " - Path to label-map.vti" << std::endl;
        std::cout << " - Extraction mode: \"all-at-once\" or \"one-by-one\"" << std::endl;
        std::cout << " - List of Label to extract, from 1 to 5" << std::endl;
        std::cout << "Examples :" << std::endl;
        std::cout << "SimpleSurfaceNets.exe C:/Path/to/label-map.vti one-by-one 2 3" << std::endl;
        std::cout << "SimpleSurfaceNets.exe C:/Path/to/label-map.vti all-at-once 1" << std::endl;
        std::cout << "SimpleSurfaceNets.exe C:/Path/to/label-map.vti all-at-once 1 2 3 4 5" << std::endl;
        return EXIT_FAILURE;
    }
    std::string l_option( argv[ 2 ] );
    if( l_option != "all-at-once" && l_option != "one-by-one" ) {
        std::cout << "Wrong second argument " << l_option << ", it should be : \"all-at-once\" or \"one-by-one\"" << std::endl;
        return EXIT_FAILURE;
    }

    // create the list of label to extract from the command line
    std::vector< unsigned char > l_label_to_extract;
    for( int i = 3; i < argc; i++ ) {
        std::string l_lab( argv[ i ] );
        if( l_lab != "1" && l_lab != "2" && l_lab != "3" && l_lab != "4" && l_lab != "5" ) {
            std::cout << "Wrong argument " << l_lab << ", it should be a label number between : \"1\" and \"5\"" << std::endl;
            return EXIT_FAILURE;
        }
        l_label_to_extract.push_back( atoi( l_lab.c_str() ) );
    }

    // load the volume from the disk (too lazy to perform some check)
    vtkNew< vtkXMLImageDataReader > l_reader;
    l_reader->SetFileName( argv[ 1 ] );
    l_reader->Update();

    // prepare a renderer
    vtkNew< vtkRenderer > l_renderer;
    l_renderer->SetBackground( 0.3, 0.3, 0.3 ); // gray background

    // prepare a render window
    vtkNew< vtkRenderWindow > l_render_window;
    l_render_window->AddRenderer( l_renderer );
    l_render_window->SetSize( 800, 600 );
    l_render_window->SetWindowName( l_option.c_str() );

    // prepare an interactor
    vtkNew< vtkInteractorStyleTrackballCamera > l_trackball;
    vtkNew< vtkRenderWindowInteractor > l_interactor;
    l_interactor->SetRenderWindow( l_render_window );
    l_interactor->SetInteractorStyle( l_trackball );

    // prepare a map of colors for each label
    std::map< unsigned char, std::array< double, 4 > > l_map_colors{ { 1, { 1, 0, 0, 0.1 } },
                                                                     { 2, { 0, 1, 0, 1 } },
                                                                     { 3, { 0, 0, 1, 1 } },
                                                                     { 4, { 1, 0, 1, 0.1 } },
                                                                     { 5, { 1, 1, 0, 0.1 } } };

    // set opacity to 0 for label that are not extracted
    for( auto& l_item : l_map_colors ) {
        if( std::find( l_label_to_extract.begin(), l_label_to_extract.end(), l_item.first ) == l_label_to_extract.end() ) {
            l_item.second[ 3 ] = 0;
        }
    }

    vtkNew< vtkTimerLog > l_total_timer;
    l_total_timer->StartTimer();

    // now set up the contour extraction algorithm
    if( l_option == "all-at-once" ) {
        // perform just one surface net pass for all selected label
        vtkNew< vtkSurfaceNets3D > l_sn;
        l_sn->SetInputConnection( l_reader->GetOutputPort() );
        l_sn->SetOutputMeshTypeToTriangles();
        for( int l_index = 0; l_index < l_label_to_extract.size(); l_index++ ) {
            l_sn->SetValue( l_index, l_label_to_extract[ l_index ] );
        }
        l_sn->Update();

        auto l_out = l_sn->GetOutput();

        auto l_data = l_out->GetCellData()->GetArray( "BoundaryLabels" );
        auto l_ptr = vtkUnsignedCharArray::SafeDownCast( l_data )->GetPointer( 0 );
        vtkIdType l_nb_tuple = l_data->GetNumberOfTuples();

        // extract one mesh
        for( auto l_it : l_map_colors ) {
            unsigned char l_label = l_it.first;
            // if the current label is not in the list, jump to the next one
            if( std::find( l_label_to_extract.begin(), l_label_to_extract.end(), l_label ) == l_label_to_extract.end() ) {
                continue;
            }

            vtkNew< vtkTimerLog > l_timer_extraction;
            l_timer_extraction->StartTimer();

            std::cout << "Going to create one new mesh for label : " << (int)l_label << std::endl;

            vtkNew< vtkPolyData > l_poly;
            l_poly->SetPoints( l_out->GetPoints() );
            l_poly->Allocate();

            vtkNew< vtkCellArray > l_array;
            auto l_cell_array = l_out->GetPolys();
            l_cell_array->InitTraversal();
            vtkIdType i = 0;
            vtkIdType l_npts = 0;
            const vtkIdType* l_pt_ids = nullptr;
            auto l_iter = vtk::TakeSmartPointer( l_cell_array->NewIterator() );
            for( l_iter->GoToFirstCell(); !l_iter->IsDoneWithTraversal(); l_iter->GoToNextCell() ) {
                // if the first OR the second component of the tuple for cell 'i' contains the label, let's keep it for that mesh
                if( l_ptr[ i ] == l_label || l_ptr[ i + 1 ] == l_label ) {
                    l_iter->GetCurrentCell( l_npts, l_pt_ids );
                    l_array->InsertNextCell( l_npts, l_pt_ids );
                }
                i += 2;
            }

            l_poly->SetPolys( l_array );

            l_timer_extraction->StopTimer();
            std::cout << "extraction took : " << l_timer_extraction->GetElapsedTime() << " seconds" << std::endl;

            // And add the single mesh corresponding to just one label to the renderer
            vtkNew< vtkPolyDataMapper > l_mapper;
            l_mapper->SetInputData( l_poly );
            l_mapper->SetScalarVisibility( false ); // no scalar visibility
            vtkNew< vtkActor > l_actor;
            l_actor->SetMapper( l_mapper );
            l_actor->GetProperty()->SetColor( l_it.second[ 0 ], l_it.second[ 1 ], l_it.second[ 2 ] ); // no lookup table here, handle color and opacity directly from the property
            l_actor->GetProperty()->SetOpacity( l_it.second[ 3 ] );
            l_renderer->AddViewProp( l_actor );
        }
    }

    else {
        // one by one mode :
        // Loop other selected label, extract it from the image and run the surface net for each selected label
        vtkNew< vtkImageData > l_image_copy;
        l_image_copy->DeepCopy( l_reader->GetOutput() );

        // iterate other all labels
        for( const auto& l_it : l_map_colors ) {
            unsigned char l_label = l_it.first;
            // if the current label is not in the list, jump to the next one
            if( std::find( l_label_to_extract.begin(), l_label_to_extract.end(), l_label ) == l_label_to_extract.end() ) {
                continue;
            }
            // use the copy image, and keep only the current label, others are set to 0
            auto l_data_ptr = static_cast< unsigned char* >( l_image_copy->GetScalarPointer() );
            auto l_data_ptr_reader = static_cast< unsigned char* >( l_reader->GetOutput()->GetScalarPointer() );
            vtkIdType l_nb_pts = l_image_copy->GetNumberOfPoints();
            for( vtkIdType i = 0; i < l_nb_pts; ++i ) {
                l_data_ptr[ i ] = ( l_data_ptr_reader[ i ] == l_label ) ? l_label : 0;
            }

            // Go Surface Nets !
            vtkNew< vtkSurfaceNets3D > l_sn;
            l_sn->SetInputData( l_image_copy );
            l_sn->SetValue( 0, l_label ); // Only for one label
            l_sn->Update();

            // And add the single mesh corresponding to just one label to the renderer
            vtkNew< vtkPolyDataMapper > l_mapper;
            l_mapper->SetInputData( l_sn->GetOutput() );
            l_mapper->SetScalarVisibility( false ); // no scalar visibility
            vtkNew< vtkActor > l_actor;
            l_actor->SetMapper( l_mapper );
            l_actor->GetProperty()->SetColor( l_it.second[ 0 ], l_it.second[ 1 ], l_it.second[ 2 ] ); // no lookup table here, handle color and opacity directly from the property
            l_actor->GetProperty()->SetOpacity( l_it.second[ 3 ] );
            l_renderer->AddViewProp( l_actor );
        }
    }

    l_renderer->ResetCamera();

    // Setup the camera <3
    l_renderer->GetActiveCamera()->SetViewUp( 0, 0, 1 );
    double l_distance = 1000;
    double l_center[ 3 ], l_pos[ 3 ];
    l_reader->GetOutput()->GetCenter( l_center );
    l_reader->GetOutput()->GetCenter( l_pos );
    l_pos[ 1 ] -= l_distance;
    double l_cam_angle = atan( ( l_center[ 2 ] * 2. ) / ( l_distance * 2. ) ) * 360.0 / vtkMath::Pi(); // some math to compute the camera view angle
    l_renderer->GetActiveCamera()->SetFocalPoint( l_center );
    l_renderer->GetActiveCamera()->SetPosition( l_pos );
    l_renderer->GetActiveCamera()->SetViewAngle( l_cam_angle + 1 );
    l_renderer->GetActiveCamera()->SetViewAngle( l_cam_angle + 1 );
    l_renderer->GetActiveCamera()->SetClippingRange( 0.01, 10000 );

    l_total_timer->StopTimer();

    std::cout << "Total time : " << l_total_timer->GetElapsedTime() << " seconds" << std::endl;

    l_interactor->Start();

    return EXIT_SUCCESS;
}