///////////////////////////////////////////////////////////////////////////
//
//  System:    Simplygon
//  File:      MappingImage.cpp
//  Language:  C++
//
//  Copyright (c) 2019 Microsoft. All rights reserved.
//
//  This is private property, and it is illegal to copy or distribute in
//  any form, without written authorization by the copyright owner(s).
//
///////////////////////////////////////////////////////////////////////////
//
//  #Description# 
//
//  A scene is loaded from a WaveFront file and the geometries are reduced.
//
//  The input objects are textured with a procedurally generated texture. 
//
//  A mapping image object is created for the reduced scene.
//  The mapping image is used to manually cast the original material onto the LOD.
//
//  The textures and optimized objects are saved to file.
//
///////////////////////////////////////////////////////////////////////////

#include "../Common/Example.h"
#include "BMPheader.h"
#include <map>
#include <sstream>
#include <string>

void RunTextureCastingWithMappingImage( const std::string& readFrom, const std::string& writeTo );

spImageData GenerateBrickImage( unsigned int image_width, unsigned int image_height );

float clamp( float value, float min, float max );
int clamp( int value, int min, int max );

int main( int argc, char* argv[] )
    {
    try
    {
        InitExample();

        // Set global variable. Using Orthonormal method for calculating
        // tangentspace.
        sg->SetGlobalSetting("DefaultTBNType", SG_TANGENTSPACEMETHOD_ORTHONORMAL);

        std::string assetPath = GetAssetPath();

        // Run the example code
        RunTextureCastingWithMappingImage(assetPath + "material_casting_test_object.obj", "object_textured");

        DeinitExample();
    }
    catch (const std::exception& ex)
    {
        std::cerr << ex.what() << std::endl;
        return -1;
    }

    return 0;
    }

void RunTextureCastingWithMappingImage( const std::string& readFrom, const std::string& writeTo )
    {
    //Read object and set output paths
    //Import geometries from a file and optimize them. 

    std::string output_geometry_filename = GetExecutablePath() + writeTo + ".obj";
    std::string bmp_output_name = GetExecutablePath() + writeTo + "_diffuse.bmp";

    //Load object from file
    spWavefrontImporter objReader = sg->CreateWavefrontImporter();
    objReader->SetImportFilePath( readFrom.c_str() );
    if( !objReader->RunImport() )
        throw std::exception("Failed to load input file!");

    //Get the scene from the importer
    spScene scene = objReader->GetScene();

    //Get the material table
    spMaterialTable original_materials = scene->GetMaterialTable();

    //Create a new texture and texture the original geometry with it
    uint input_texture_width = 512;
    uint input_texture_height = 512;

    spImageData input_texture = GenerateBrickImage( input_texture_width, input_texture_height );

    spUnsignedCharArray input_colors = SafeCast<IUnsignedCharArray>( input_texture->GetColors() );

    // Select all the SceneMesh nodes from the scene
    spSelectionSet allGeometriesSelectionSet = scene->GetSelectionSetTable()->GetSelectionSet( scene->SelectNodes( "ISceneMesh" ) );

    //Store a copy of the original GeometyData objects
    std::vector<spGeometryData> original_unchanged_geometries;

    //Map from the node with the geometry to the index in the copied original geometry vector
    std::map<spSceneMesh, uint> scene_mesh_node_to_index;

    for( uint i = 0; i < allGeometriesSelectionSet->GetItemCount(); ++i )
        {
        spSceneMesh sceneMesh = Cast<ISceneMesh>( scene->GetNodeByGUID( allGeometriesSelectionSet->GetItem( i ) ) );

        //Give the scene mesh a name if it doesn't have one
        rstring current_name = sceneMesh->GetName();
        if( current_name.IsEmpty() )
            {
            char new_geometry_name[256];
            sprintf( new_geometry_name, "Object%d", i );
            sceneMesh->SetName( new_geometry_name );
            }

        scene_mesh_node_to_index.insert( std::pair<spSceneMesh, uint>( sceneMesh, (uint)original_unchanged_geometries.size() ) );
        original_unchanged_geometries.push_back( sceneMesh->GetGeometry()->NewCopy( true ) );
        }

    //Reducer
    spReductionProcessor red = sg->CreateReductionProcessor();
    red->SetScene( scene );

    //Set the Repair Settings.
    spRepairSettings repair_settings = red->GetRepairSettings();
    repair_settings->SetWeldDist( 0.0f );
    repair_settings->SetTjuncDist( 0.0f );

    //Set the Reduction Settings.
    spReductionSettings reduction_settings = red->GetReductionSettings();

    //Texture coordinates will be calculated after the reduction is done, 
    //so there is no need for texture coordinates to be preserved.
    reduction_settings->SetTextureImportance( 0.0f );

    //Reduce to one tenth original triangles
    reduction_settings->SetTriangleRatio( 0.1f );

    // Set the Normal Calculation Settings.
    spNormalCalculationSettings normal_settings = red->GetNormalCalculationSettings();
    normal_settings->SetReplaceNormals( true );
    normal_settings->SetHardEdgeAngleInRadians( 3.14159f * 90.f / 180.0f );


    //Setup the mapping image properties
    uint texture_width = 512; // Should be a multiple of 256
    uint texture_height = 512; // Should be a multiple of 256
    uint multisampling = 4; // 256 should be divisible by this

    // Set the Image Mapping Settings for the reducer
    spMappingImageSettings mapping_settings = red->GetMappingImageSettings();

    // generate diffuse, specular, normal maps and custom channel later.
    mapping_settings->SetGenerateMappingImage( true );
    // Set to generate new texture coordinates.
    mapping_settings->SetGenerateTexCoords( true );
    // The higher the number, the fewer texture-borders.
    mapping_settings->SetParameterizerMaxStretch( 0.2f );
    // Buffer space for when texture is mip-mapped, so color values don't blend over.
    mapping_settings->SetGutterSpace( 4 );

    mapping_settings->SetWidth( texture_width );
    mapping_settings->SetHeight( texture_height );
    mapping_settings->SetMultisamplingLevel( multisampling );

    // Reduce the geometry
    red->RunProcessing();


    // Create a BMP to store the new texture in
    output_image new_texture;
    new_texture.setup( texture_width, texture_height );
    new_texture.set_mapping( 0, 0 );

    // Use the Mapping image to create a new texture for the reduced geometry

    // Get the mapping image object that contains information back to the original geometries
    spMappingImage mapping_image = red->GetMappingImage();

    // spMappingImageMeshData contains information needed to convert the mapping data to the individual
    // geometries from the collection of geometries
    spMappingImageMeshData mapping_image_mesh_data = mapping_image->GetMappingMeshData();

    // The spChunkedImageData contains per-texel mapping information back to the original geometries
    // as two fields:
    //
    //  - Original triangle id, from a geometry consisted of all the original geometry combined.
    //    This means that the local triangle IDs from each original geometries will not valid here.
    //    Instead, the triangle ID is a global ID from the combined geometry.
    //    To map back to local triangle IDs, the spMappingImageMeshData is used.
    //    
    //  - The original Barycentric coordinates
    //
    spChunkedImageData chunked_image_data = mapping_image->GetImageData();

    // Get the number of image data chunks
    uint x_chunks = chunked_image_data->GetXSize();
    uint y_chunks = chunked_image_data->GetYSize();

    // Loop the image data
    // First by chunk
    // Then by pixel
    // Then by sub-pixel (if multisampling > 1)
    for( uint x_chunk = 0; x_chunk < x_chunks; ++x_chunk )
        {
        for( uint y_chunk = 0; y_chunk < y_chunks; ++y_chunk )
            {
            // Lock and fetch a chunk of mapping image pixels
            spImageData current_chunk = chunked_image_data->LockChunk2D( x_chunk, y_chunk );

            // Get the mapping image fields "TriangleIds" & "BarycentricCoords"
            spRidArray triangle_ids = SafeCast<IRidArray>( current_chunk->GetField( "TriangleIds" ) );
            spUnsignedShortArray barycentric_coords = SafeCast<IUnsignedShortArray>( current_chunk->GetField( "BarycentricCoords" ) );

            // Continue with next iteration if any of the fields doesn't exist
            if( !triangle_ids || !barycentric_coords )
                continue;

            // Get the number of pixels in the chunk, excluding sub-pixels
            uint px_size = current_chunk->GetXSize() / multisampling;
            uint py_size = current_chunk->GetYSize() / multisampling;

            //Loop the pixels
            for( uint py = 0; py < py_size; ++py )
                {
                for( uint px = 0; px < px_size; ++px )
                    {
                    // Initialize pixel color
                    int pixel_r = 0;
                    int pixel_g = 0;
                    int pixel_b = 0;

                    // samples_used is used when multisampling to divide by the number of samples used per particular pixel
                    int samples_used = 0;

                    // Loop the sub-pixels (will only loop once if multisampling is 1)
                    for( uint py_multisample = 0; py_multisample < multisampling; ++py_multisample )
                        {
                        for( uint px_multisample = 0; px_multisample < multisampling; ++px_multisample )
                            {
                            // The one-dimensional index of the current sub-pixel
                            // index = column + row * column_size
                            rid sub_pixel_index = (px * multisampling + px_multisample) + (py * multisampling + py_multisample) * (px_size * multisampling);

                            // The global triangle ID in the combined geometry
                            rid global_triangle_id = triangle_ids->GetItem( sub_pixel_index );

                            // Continue with the next sub-pixel, should the current sub-pixel not map to a triangle
                            if( global_triangle_id < 0 )
                                continue;

                            // Get the corresponding barycentric coordinates in the original geometry
                            unsigned short barycentric_x = barycentric_coords->GetItem( sub_pixel_index * 2 + 0 );
                            unsigned short barycentric_y = barycentric_coords->GetItem( sub_pixel_index * 2 + 1 );
                            unsigned short barycentric_z = 65535 - barycentric_x - barycentric_y;

                            // Get the normalized barycentric coordinates
                            float barycentric_x_normalized = barycentric_x / float( 65535 );
                            float barycentric_y_normalized = barycentric_y / float( 65535 );
                            float barycentric_z_normalized = barycentric_z / float( 65535 );

                            // Declare the local triangle id variable
                            uint local_triangle_id = -1;

                            uint geometry_count = mapping_image_mesh_data->GetMappedGeometriesCount();

                            int mapped_geometry_id = -1;

                            // To find which original geometry is being mapped to using the global triangle id
                            // the mapping image mesh data is used.
                            for( uint g_id = 0; g_id < geometry_count; ++g_id )
                                {
                                int next_geometry_starting_triangle_id = -1;

                                // If g_id is the last geometry, then 
                                // keep next_geometry_starting_triangle_id as -1
                                if( g_id < geometry_count - 1 )
                                    {
                                    next_geometry_starting_triangle_id = mapping_image_mesh_data->GetStartTriangleIdOfGeometry( g_id + 1 );
                                    }

                                // If the global triangle id is below the next geometry's starting triangle id,
                                // we know that the current geomety (g_id) will contain the current triangle
                                //
                                // Also if next_geometry_starting_triangle_id is -1, since then we're at the last geometry
                                if( global_triangle_id < next_geometry_starting_triangle_id || next_geometry_starting_triangle_id == -1 )
                                    {
                                    mapped_geometry_id = g_id;

                                    uint starting_triangle_index = mapping_image_mesh_data->GetStartTriangleIdOfGeometry( g_id );

                                    local_triangle_id = global_triangle_id - starting_triangle_index;
                                    break;
                                    }
                                }

                            // Make sure the mapped-to geometry is valid
                            assert( mapped_geometry_id >= 0 && mapped_geometry_id < (int)geometry_count );

                            //////
                            // The scene_path can be retrieved from the mapping image
                            // the path corresponds to its place in the scene and the scene mesh name
                            rstring scene_path = mapping_image_mesh_data->GetScenePathOfGeometry( mapped_geometry_id );
                            //////

                            spSceneMesh sceneMesh = Cast<ISceneMesh>( scene->GetNodeFromPath( scene_path ) );

                            // Now we can get whatever information in the original geometry that we want
                            // In this example, we fetch the original texture coordinates
                            uint original_geometry_index = scene_mesh_node_to_index.find( sceneMesh )->second;
                            spRealArray texcoords = original_unchanged_geometries[original_geometry_index]->GetTexCoords( 0 );

                            int corner0 = local_triangle_id * 3 + 0;
                            int corner1 = local_triangle_id * 3 + 1;
                            int corner2 = local_triangle_id * 3 + 2;

                            float u_corner0 = texcoords->GetItem( corner0 * 2 + 0 );
                            float v_corner0 = texcoords->GetItem( corner0 * 2 + 1 );

                            float u_corner1 = texcoords->GetItem( corner1 * 2 + 0 );
                            float v_corner1 = texcoords->GetItem( corner1 * 2 + 1 );

                            float u_corner2 = texcoords->GetItem( corner2 * 2 + 0 );
                            float v_corner2 = texcoords->GetItem( corner2 * 2 + 1 );

                            // With the barycentric coordinates we can interpolate to a precise texture coordinate value in the triangle
                            float u = barycentric_x_normalized * u_corner0 + barycentric_y_normalized * u_corner1 + barycentric_z_normalized * u_corner2;
                            float v = barycentric_x_normalized * v_corner0 + barycentric_y_normalized * v_corner1 + barycentric_z_normalized * v_corner2;

                            // With the texture coords, we sample the original texture (the one we created in the beginning)

                            int sample_x = int( u * input_texture_width );
                            int sample_y = int( v * input_texture_height );

                            // Clamp 
                            sample_x = sample_x > (int)input_texture_width - 1 ? input_texture_width - 1 : sample_x;
                            sample_y = sample_y > (int)input_texture_height - 1 ? input_texture_height - 1 : sample_y;

                            // Fetch input texture color
                            unsigned char rc = input_colors->GetItem( (sample_x + sample_y * input_texture_width) * 3 + 0 );
                            unsigned char gc = input_colors->GetItem( (sample_x + sample_y * input_texture_width) * 3 + 1 );
                            unsigned char bc = input_colors->GetItem( (sample_x + sample_y * input_texture_width) * 3 + 2 );

                            // Accumulate colors in the pixel color variables
                            pixel_r += rc;
                            pixel_g += gc;
                            pixel_b += bc;

                            samples_used++;
                            }
                        }

                    // Normalize colors based on number of sub-pixels used
                    if( samples_used > 0 )
                        {
                        pixel_r /= samples_used;
                        pixel_g /= samples_used;
                        pixel_b /= samples_used;
                        }

                    // Set the output texture color
                    uint texture_px = x_chunk * px_size + px;
                    uint texture_py = y_chunk * py_size + py;

                    new_texture.set_pixel( texture_px, texture_py, 2, pixel_r );
                    new_texture.set_pixel( texture_px, texture_py, 1, pixel_g );
                    new_texture.set_pixel( texture_px, texture_py, 0, pixel_b );

                    }
                }

            // Unlock the chunk of image data
            chunked_image_data->UnlockChunk2D( x_chunk, y_chunk );

            }
        }

    // Store the new texture as a BMP
    new_texture.write_to_file( bmp_output_name.c_str() );

    // Add the new texture to the reduced geometry

    // Create new material table.
    spMaterialTable output_materials = sg->CreateMaterialTable();
    spTextureTable output_textures = sg->CreateTextureTable();

    //  Create new material for the table.
    spMaterial output_material = sg->CreateMaterial();
    output_material->SetName( "diffuse_brick_texture" );

    //Add the new material to the table
    output_materials->AddMaterial( output_material );

    //Add the diffuse texture to the diffuse channel in the output material
    {
    spTexture newTex = sg->CreateTexture();
    newTex->SetFilePath( bmp_output_name.c_str() );
    newTex->SetName( "Diffuse" );
    output_textures->AddTexture( newTex );
    spShadingTextureNode texNode = sg->CreateShadingTextureNode();
    texNode->SetTextureName( "Diffuse" );
    output_material->SetShadingNetwork( SG_MATERIAL_CHANNEL_DIFFUSE, texNode );
    }

    scene->GetMaterialTable()->Clear();
    scene->GetMaterialTable()->Copy( output_materials );
    scene->GetTextureTable()->Copy( output_textures );

    // Export the LOD
    spWavefrontExporter objexp = sg->CreateWavefrontExporter();

    objexp->SetExportFilePath( output_geometry_filename.c_str() );
    objexp->SetScene( scene );
    objexp->RunExport();
    }

spImageData GenerateBrickImage( unsigned int image_width, unsigned int image_height )
    {
    spImageData img = sg->CreateImageData();

    //Set imagedata dimensions
    img->Set2DSize( image_width, image_height );

    //Add rgb colors to imagedata
    img->AddColors( TYPES_ID_UCHAR, "RGB" );

    //Get the reference to the colors to array with tuple size 3, tuple count (image_width x image_height)
    spUnsignedCharArray colors = SafeCast<IUnsignedCharArray>( img->GetColors() );

    //Create bmp to store texture
    //Not needed, only used for debugging and displaying
    output_image bmpOut;
    bmpOut.setup( image_width, image_height );
    bmpOut.set_mapping( 0, 0 );

    //Brick texture parameters
    unsigned int brick_height = 15;
    unsigned int brick_length = 35;
    unsigned int brick_apart = 1;

    int brick_R = 165;
    int brick_G = 66;
    int brick_B = 25;

    int between_bricks_color = 20;
    int between_bricks_noise_color = 170;

    //Iterate over the image pixels and set the array values
    for( unsigned int y = 0; y < image_height; y++ )
        for( unsigned int x = 0; x < image_width; x++ )
            {
            //One dimensional index
            int i = image_width * y + x;

            //Generate random values to add noise to the texture
            int random1 = rand() % 50 - 25;
            int random2 = rand() % 50 - 25;
            int random3 = rand() % 50 - 25;

            //Every second row of bricks is offset half a brick length
            unsigned int brick_offset = (uint)(0.5 * brick_length) * ((uint)(y / brick_height) % 2);

            //Check if xy-coordinate is on or betweeen bricks
            bool between_bricks = ((x + brick_offset) % brick_length < brick_apart) || (y % brick_height < brick_apart);
            //Add some noise to the coordinates
            bool between_bricks_noise = ((x + random1 / 20 + brick_offset) % brick_length < brick_apart) || ((y + random2 / 20) % brick_height < brick_apart);

            //Brick colors to array
            colors->SetItem( i * 3 + 0, clamp( (int)(between_bricks_noise ? between_bricks_noise_color + random1 : (between_bricks ? between_bricks_color + random1 : brick_R + random1)), 0, 255 ) );
            colors->SetItem( i * 3 + 1, clamp( (int)(between_bricks_noise ? between_bricks_noise_color + random1 : (between_bricks ? between_bricks_color + random1 : brick_G + random2)), 0, 255 ) );
            colors->SetItem( i * 3 + 2, clamp( (int)(between_bricks_noise ? between_bricks_noise_color + random1 : (between_bricks ? between_bricks_color + random1 : brick_B + random3)), 0, 255 ) );

            //Copy colors to bmp
            bmpOut.set_pixel( x, y, 2, colors->GetItem( i * 3 + 0 ) );
            bmpOut.set_pixel( x, y, 1, colors->GetItem( i * 3 + 1 ) );
            bmpOut.set_pixel( x, y, 0, colors->GetItem( i * 3 + 2 ) );
            }

    //Store bmp
    bmpOut.write_to_file( "brick_texture.bmp" );
    return img;
    }

float clamp( float value, float min, float max )
    {
    return value < min ? min : (value > max ? max : value);
    }

int clamp( int value, int min, int max )
    {
    return value < min ? min : (value > max ? max : value);
    }