#include "IO/Writer.h"
#include "IO/IOHelpers.h"
#include "IO/MeshDatabase.h"
#include "IO/silo.h"
#include "common/MPI.h"
#include "common/Utilities.h"

#include <algorithm>
#include <memory>
#include <set>
#include <sys/stat.h>
#include <vector>


enum class Format { OLD, NEW, SILO, UNKNOWN };


/****************************************************
 * Recursively create the subdirectory               *
 ****************************************************/
static void recursiveMkdir( const std::string &path, mode_t mode )
{
    // Iterate through the root directories until we create the desired path
    for ( size_t pos = 0; pos < path.size(); ) {
        // slide backwards in string until next slash found
        pos++;
        for ( ; pos < path.size(); pos++ ) {
            if ( path[pos] == '/' || path[pos] == 92 )
                break;
        }
        // Create the temporary path
        auto path2 = path.substr( 0, pos );
        // Check if the temporary path exists
        struct stat status;
        int result = stat( path2.data(), &status );
        if ( result == 0 ) {
            // if there is a part of the path that already exists make sure it is really a directory
            if ( !S_ISDIR( status.st_mode ) ) {
                ERROR(
                    "Error in recursiveMkdir...\n"
                    "    Cannot create directories in path = " +
                    path +
                    "\n    because some intermediate item in path exists and is NOT a directory" );
            }
            continue;
        }
        // Create the directory and test the result
        result = mkdir( path2.data(), mode );
        if ( result != 0 ) {
            // Maybe another rank created the directory, check
            int result = stat( path2.data(), &status );
            if ( result != 0 && !S_ISDIR( status.st_mode ) )
                ERROR( "Error in Utilities::recursiveMkdir...\n"
                       "    Cannot create directory  = " +
                       path2 );
        }
    }
}


/****************************************************
 * Initialize the writer                             *
 ****************************************************/
static std::string global_IO_path;
static Format global_IO_format = Format::UNKNOWN;
void IO::initialize( const std::string &path, const std::string &format, bool append )
{
    if ( path.empty() )
        global_IO_path = ".";
    else
        global_IO_path = path;
    if ( format == "old" )
        global_IO_format = Format::OLD;
    else if ( format == "new" )
        global_IO_format = Format::NEW;
    else if ( format == "silo" )
        global_IO_format = Format::SILO;
    else
        ERROR( "Unknown format" );
    int rank = Utilities::MPI( MPI_COMM_WORLD ).getRank();
    if ( !append && rank == 0 ) {
        recursiveMkdir( path, S_IRWXU | S_IRGRP );
        std::string filename;
        if ( global_IO_format == Format::OLD || global_IO_format == Format::NEW )
            filename = global_IO_path + "/summary.LBM";
        else if ( global_IO_format == Format::SILO )
            filename = global_IO_path + "/LBM.visit";
        else
            ERROR( "Unknown format" );
        auto fid = fopen( filename.c_str(), "wb" );
        fclose( fid );
    }
}


// Write the mesh data in the original format
static std::vector<IO::MeshDatabase> writeMeshesOrigFormat(
    const std::vector<IO::MeshDataStruct> &meshData, const std::string &path, int rank )
{
    std::vector<IO::MeshDatabase> meshes_written;
    for ( size_t i = 0; i < meshData.size(); i++ ) {
        char domainname[100], filename[100], fullpath[200];
        sprintf( domainname, "%05i", rank );
        sprintf( filename, "%s.%05i", meshData[i].meshName.c_str(), rank );
        sprintf( fullpath, "%s/%s", path.c_str(), filename );
        FILE *fid = fopen( fullpath, "wb" );
        INSIST( fid != NULL, std::string( "Error opening file: " ) + fullpath );
        std::shared_ptr<IO::Mesh> mesh = meshData[i].mesh;
        IO::MeshDatabase mesh_entry;
        mesh_entry.name      = meshData[i].meshName;
        mesh_entry.type      = meshType( *mesh );
        mesh_entry.meshClass = meshData[i].mesh->className();
        mesh_entry.format    = IO::FileFormat::OLD;
        IO::DatabaseEntry domain;
        domain.name   = domainname;
        domain.file   = filename;
        domain.offset = 0;
        mesh_entry.domains.push_back( domain );
        if ( !meshData[i].vars.empty() ) {
            printf( "Warning: variables are not supported with this format (original)\n" );
            // for (size_t j=0; j<meshData[i].vars.size(); j++)
            //    mesh_entry.variables.push_back( meshData[i].vars[j]->name );
        }
        const std::string meshClass = mesh->className();
        if ( meshClass == "PointList" ) {
            // List of points
            std::shared_ptr<IO::PointList> pointlist =
                std::dynamic_pointer_cast<IO::PointList>( mesh );
            const std::vector<Point> &P = pointlist->points;
            for ( size_t i = 0; i < P.size(); i++ ) {
                double x[3];
                x[0] = P[i].x;
                x[1] = P[i].y;
                x[2] = P[i].z;
                fwrite( x, sizeof( double ), 3, fid );
            }
        } else if ( meshClass == "TriList" || meshClass == "TriMesh" ) {
            // Triangle mesh
            std::shared_ptr<IO::TriList> trilist = IO::getTriList( mesh );
            const std::vector<Point> &A          = trilist->A;
            const std::vector<Point> &B          = trilist->B;
            const std::vector<Point> &C          = trilist->C;
            for ( size_t i = 0; i < A.size(); i++ ) {
                double tri[9];
                tri[0] = A[i].x;
                tri[1] = A[i].y;
                tri[2] = A[i].z;
                tri[3] = B[i].x;
                tri[4] = B[i].y;
                tri[5] = B[i].z;
                tri[6] = C[i].x;
                tri[7] = C[i].y;
                tri[8] = C[i].z;
                fwrite( tri, sizeof( double ), 9, fid );
            }
        } else if ( meshClass == "DomainMesh" ) {
            // This format was never supported with the old format
        } else {
            ERROR( "Unknown mesh" );
        }
        fclose( fid );
        std::sort( mesh_entry.variables.begin(), mesh_entry.variables.end() );
        mesh_entry.variables.erase(
            std::unique( mesh_entry.variables.begin(), mesh_entry.variables.end() ),
            mesh_entry.variables.end() );
        meshes_written.push_back( mesh_entry );
    }
    return meshes_written;
}


// Create the database entry for the mesh data
static IO::MeshDatabase getDatabase(
    const std::string &filename, const IO::MeshDataStruct &mesh, IO::FileFormat format, int rank )
{
    char domainname[100];
    sprintf( domainname, "%s_%05i", mesh.meshName.c_str(), rank );
    // Create the MeshDatabase
    IO::MeshDatabase database;
    database.name      = mesh.meshName;
    database.type      = meshType( *( mesh.mesh ) );
    database.meshClass = mesh.mesh->className();
    database.format    = format;
    // Write the mesh
    IO::DatabaseEntry domain;
    domain.name   = domainname;
    domain.file   = filename;
    domain.offset = -1;
    database.domains.push_back( domain );
    // Write the variables
    for ( size_t i = 0; i < mesh.vars.size(); i++ ) {
        // Add basic variable info
        IO::VariableDatabase info;
        info.name = mesh.vars[i]->name;
        info.type = mesh.vars[i]->type;
        info.dim  = mesh.vars[i]->dim;
        database.variables.push_back( info );
        // Add domain variable info
        IO::DatabaseEntry variable;
        variable.name   = mesh.vars[i]->name;
        variable.file   = filename;
        variable.offset = -1;
        std::pair<std::string, std::string> key( domain.name, mesh.vars[i]->name );
        database.variable_data.insert(
            std::pair<std::pair<std::string, std::string>, IO::DatabaseEntry>( key, variable ) );
    }
    return database;
}


// Write a mesh (and variables) to a file
static IO::MeshDatabase write_domain( FILE *fid, const std::string &filename,
    const IO::MeshDataStruct &mesh, IO::FileFormat format, int rank )
{
    const int level = 0;
    // Create the MeshDatabase
    IO::MeshDatabase database = getDatabase( filename, mesh, format, rank );
    // Write the mesh
    IO::DatabaseEntry &domain      = database.domains[0];
    domain.offset                  = ftell( fid );
    std::pair<size_t, void *> data = mesh.mesh->pack( level );
    fprintf( fid, "Mesh: %s-%05i: %lu\n", mesh.meshName.c_str(), rank, data.first );
    fwrite( data.second, 1, data.first, fid );
    fprintf( fid, "\n" );
    delete[]( char * ) data.second;
    // Write the variables
    for ( size_t i = 0; i < mesh.vars.size(); i++ ) {
        ASSERT( mesh.vars[i]->type != IO::VariableType::NullVariable );
        std::pair<std::string, std::string> key( domain.name, mesh.vars[i]->name );
        auto &variable  = database.variable_data[key];
        variable.offset = ftell( fid );
        int dim         = mesh.vars[i]->dim;
        auto type       = getString( mesh.vars[i]->type );
        size_t N        = mesh.vars[i]->data.length();
        size_t N_mesh   = mesh.mesh->numberPointsVar( mesh.vars[i]->type );
        ASSERT( N == dim * N_mesh );
        fprintf( fid, "Var: %s-%05i-%s: %i, %s, %lu, %lu, double\n", database.name.c_str(), rank,
            variable.name.c_str(), dim, type.data(), N_mesh, N * sizeof( double ) );
        fwrite( mesh.vars[i]->data.data(), sizeof( double ), N, fid );
        fprintf( fid, "\n" );
    }
    return database;
}


#ifdef USE_SILO
// Write a PointList mesh (and variables) to a file
template<class TYPE>
static void writeSiloPointMesh(
    DBfile *fid, const IO::PointList &mesh, const std::string &meshname )
{
    const auto &points = mesh.getPoints();
    std::vector<TYPE> x( points.size() ), y( points.size() ), z( points.size() );
    for ( size_t i = 0; i < x.size(); i++ ) {
        x[i] = points[i].x;
        y[i] = points[i].y;
        z[i] = points[i].z;
    }
    const TYPE *coords[] = { x.data(), y.data(), z.data() };
    IO::silo::writePointMesh<TYPE>( fid, meshname, 3, points.size(), coords );
}
static void writeSiloPointList(
    DBfile *fid, const IO::MeshDataStruct &meshData, IO::MeshDatabase database )
{
    const IO::PointList &mesh  = dynamic_cast<IO::PointList &>( *meshData.mesh );
    const std::string meshname = database.domains[0].name;
    if ( meshData.precision == IO::DataType::Double ) {
        writeSiloPointMesh<double>( fid, mesh, meshname );
    } else if ( meshData.precision == IO::DataType::Float ) {
        writeSiloPointMesh<float>( fid, mesh, meshname );
    } else {
        ERROR( "Unsupported format" );
    }
    const auto &points = mesh.getPoints();
    std::vector<double> x( points.size() ), y( points.size() ), z( points.size() );
    for ( size_t i = 0; i < x.size(); i++ ) {
        x[i] = points[i].x;
        y[i] = points[i].y;
        z[i] = points[i].z;
    }
    const double *coords[] = { x.data(), y.data(), z.data() };
    IO::silo::writePointMesh( fid, meshname, 3, points.size(), coords );
    for ( size_t i = 0; i < meshData.vars.size(); i++ ) {
        const IO::Variable &var = *meshData.vars[i];
        if ( var.precision == IO::DataType::Double ) {
            IO::silo::writePointMeshVariable( fid, meshname, var.name, var.data );
        } else if ( var.precision == IO::DataType::Float ) {
            Array<float> data2( var.data.size() );
            data2.copy( var.data );
            IO::silo::writePointMeshVariable( fid, meshname, var.name, data2 );
        } else if ( var.precision == IO::DataType::Int ) {
            Array<int> data2( var.data.size() );
            data2.copy( var.data );
            IO::silo::writePointMeshVariable( fid, meshname, var.name, data2 );
        } else {
            ERROR( "Unsupported format" );
        }
    }
}
// Write a TriMesh mesh (and variables) to a file
template<class TYPE>
static void writeSiloTriMesh( DBfile *fid, const IO::TriMesh &mesh, const std::string &meshname )
{
    const auto &points = mesh.vertices->getPoints();
    std::vector<TYPE> x( points.size() ), y( points.size() ), z( points.size() );
    for ( size_t i = 0; i < x.size(); i++ ) {
        x[i] = points[i].x;
        y[i] = points[i].y;
        z[i] = points[i].z;
    }
    const TYPE *coords[] = { x.data(), y.data(), z.data() };
    const int *tri[]     = { mesh.A.data(), mesh.B.data(), mesh.C.data() };
    IO::silo::writeTriMesh<TYPE>( fid, meshname, 3, 2, points.size(), coords, mesh.A.size(), tri );
}
static void writeSiloTriMesh2( DBfile *fid, const IO::MeshDataStruct &meshData,
    const IO::TriMesh &mesh, IO::MeshDatabase database )
{
    const std::string meshname = database.domains[0].name;
    if ( meshData.precision == IO::DataType::Double ) {
        writeSiloTriMesh<double>( fid, mesh, meshname );
    } else if ( meshData.precision == IO::DataType::Float ) {
        writeSiloTriMesh<float>( fid, mesh, meshname );
    } else {
        ERROR( "Unsupported format" );
    }
    for ( size_t i = 0; i < meshData.vars.size(); i++ ) {
        const IO::Variable &var = *meshData.vars[i];
        if ( var.precision == IO::DataType::Double ) {
            IO::silo::writeTriMeshVariable( fid, 3, meshname, var.name, var.data, var.type );
        } else if ( var.precision == IO::DataType::Float ) {
            Array<float> data2( var.data.size() );
            data2.copy( var.data );
            IO::silo::writeTriMeshVariable( fid, 3, meshname, var.name, data2, var.type );
        } else if ( var.precision == IO::DataType::Int ) {
            Array<int> data2( var.data.size() );
            data2.copy( var.data );
            IO::silo::writeTriMeshVariable( fid, 3, meshname, var.name, data2, var.type );
        } else {
            ERROR( "Unsupported format" );
        }
    }
}
static void writeSiloTriMesh(
    DBfile *fid, const IO::MeshDataStruct &meshData, IO::MeshDatabase database )
{
    const IO::TriMesh &mesh = dynamic_cast<IO::TriMesh &>( *meshData.mesh );
    writeSiloTriMesh2( fid, meshData, mesh, database );
}
static void writeSiloTriList(
    DBfile *fid, const IO::MeshDataStruct &meshData, IO::MeshDatabase database )
{
    auto mesh = getTriMesh( meshData.mesh );
    writeSiloTriMesh2( fid, meshData, *mesh, database );
}
// Write a DomainMesh mesh (and variables) to a file
static void writeSiloDomainMesh(
    DBfile *fid, const IO::MeshDataStruct &meshData, IO::MeshDatabase database )
{
    const IO::DomainMesh &mesh = dynamic_cast<IO::DomainMesh &>( *meshData.mesh );
    RankInfoStruct info( mesh.rank, mesh.nprocx, mesh.nprocy, mesh.nprocz );
    std::array<double, 6> range = { info.ix * mesh.Lx / info.nx,
        ( info.ix + 1 ) * mesh.Lx / info.nx, info.jy * mesh.Ly / info.ny,
        ( info.jy + 1 ) * mesh.Ly / info.ny, info.kz * mesh.Lz / info.nz,
        ( info.kz + 1 ) * mesh.Lz / info.nz };
    std::array<int, 3> N        = { mesh.nx, mesh.ny, mesh.nz };
    auto meshname               = database.domains[0].name;
    IO::silo::writeUniformMesh<3>( fid, meshname, range, N );
    IO::silo::write<int>(
        fid, meshname + "_rankinfo", { mesh.rank, mesh.nprocx, mesh.nprocy, mesh.nprocz } );
    for ( size_t i = 0; i < meshData.vars.size(); i++ ) {
        const auto &var = *meshData.vars[i];
        if ( var.precision == IO::DataType::Double ) {
            IO::silo::writeUniformMeshVariable<3>( fid, meshname, N, var.name, var.data, var.type );
        } else if ( var.precision == IO::DataType::Float ) {
            Array<float> data2( var.data.size() );
            data2.copy( var.data );
            IO::silo::writeUniformMeshVariable<3>( fid, meshname, N, var.name, data2, var.type );
        } else if ( var.precision == IO::DataType::Int ) {
            Array<int> data2( var.data.size() );
            data2.copy( var.data );
            IO::silo::writeUniformMeshVariable<3>( fid, meshname, N, var.name, data2, var.type );
        } else {
            ERROR( "Unsupported format" );
        }
    }
}
// Write a mesh (and variables) to a file
static IO::MeshDatabase write_domain_silo( DBfile *fid, const std::string &filename,
    const IO::MeshDataStruct &mesh, IO::FileFormat format, int rank )
{
    // Create the MeshDatabase
    auto database = getDatabase( filename, mesh, format, rank );
    if ( database.meshClass == "PointList" ) {
        writeSiloPointList( fid, mesh, database );
    } else if ( database.meshClass == "TriMesh" ) {
        writeSiloTriMesh( fid, mesh, database );
    } else if ( database.meshClass == "TriList" ) {
        writeSiloTriList( fid, mesh, database );
    } else if ( database.meshClass == "DomainMesh" ) {
        writeSiloDomainMesh( fid, mesh, database );
    } else {
        ERROR( "Unknown mesh class" );
    }
    return database;
}
// Write the summary file for silo
std::pair<int, int> getSiloMeshType( const std::string &meshClass )
{
    int meshType = 0;
    int varType  = 0;
    if ( meshClass == "PointList" ) {
        meshType = DB_POINTMESH;
        varType  = DB_POINTVAR;
    } else if ( meshClass == "TriMesh" ) {
        meshType = DB_UCDMESH;
        varType  = DB_UCDVAR;
    } else if ( meshClass == "TriList" ) {
        meshType = DB_UCDMESH;
        varType  = DB_UCDVAR;
    } else if ( meshClass == "DomainMesh" ) {
        meshType = DB_QUAD_RECT;
        varType  = DB_QUADVAR;
    } else {
        ERROR( "Unknown mesh class" );
    }
    return std::make_pair( meshType, varType );
}
void writeSiloSummary(
    const std::vector<IO::MeshDatabase> &meshes_written, const std::string &filename )
{
    auto fid = IO::silo::open( filename, IO::silo::CREATE );
    for ( const auto &data : meshes_written ) {
        auto type = getSiloMeshType( data.meshClass );
        std::vector<int> meshTypes( data.domains.size(), type.first );
        std::vector<int> varTypes( data.domains.size(), type.second );
        std::vector<std::string> meshNames;
        for ( const auto &tmp : data.domains )
            meshNames.push_back( tmp.file + ":" + tmp.name );
        IO::silo::writeMultiMesh( fid, data.name, meshNames, meshTypes );
        for ( const auto &variable : data.variables ) {
            std::vector<std::string> varnames;
            for ( const auto &tmp : data.domains )
                varnames.push_back( tmp.file + ":" + variable.name );
            IO::silo::writeMultiVar( fid, variable.name, varnames, varTypes );
        }
    }
    IO::silo::close( fid );
}
#endif


// Write the mesh data in the new format
static std::vector<IO::MeshDatabase> writeMeshesNewFormat(
    const std::vector<IO::MeshDataStruct> &meshData, const std::string &path, IO::FileFormat format,
    int rank )
{
    std::vector<IO::MeshDatabase> meshes_written;
    char filename[100], fullpath[200];
    sprintf( filename, "%05i", rank );
    sprintf( fullpath, "%s/%s", path.c_str(), filename );
    FILE *fid = fopen( fullpath, "wb" );
    for ( size_t i = 0; i < meshData.size(); i++ ) {
        std::shared_ptr<IO::Mesh> mesh = meshData[i].mesh;
        meshes_written.push_back( write_domain( fid, filename, meshData[i], format, rank ) );
    }
    fclose( fid );
    return meshes_written;
}


// Write the mesh data to silo
static std::vector<IO::MeshDatabase> writeMeshesSilo(
    const std::vector<IO::MeshDataStruct> &meshData, const std::string &path, IO::FileFormat format,
    int rank )
{
#ifdef USE_SILO
    std::vector<IO::MeshDatabase> meshes_written;
    char filename[100], fullpath[200];
    sprintf( filename, "%05i.silo", rank );
    sprintf( fullpath, "%s/%s", path.c_str(), filename );
    auto fid = IO::silo::open( fullpath, IO::silo::CREATE );
    for ( size_t i = 0; i < meshData.size(); i++ ) {
        auto mesh = meshData[i].mesh;
        meshes_written.push_back( write_domain_silo( fid, filename, meshData[i], format, rank ) );
    }
    IO::silo::close( fid );
    return meshes_written;
#else
    NULL_USE( meshData );
    NULL_USE( path );
    NULL_USE( format );
    NULL_USE( rank );
    ERROR( "Application built without silo support" );
    return std::vector<IO::MeshDatabase>();
#endif
}


/****************************************************
 * Write the mesh data                               *
 ****************************************************/
void IO::writeData( const std::string &subdir, const std::vector<IO::MeshDataStruct> &meshData,
    const Utilities::MPI &comm )
{
    if ( global_IO_path.empty() )
        IO::initialize();
    PROFILE_START( "writeData" );
    int rank = Utilities::MPI( MPI_COMM_WORLD ).getRank();
    // Check the meshData before writing
    for ( const auto &data : meshData )
        ASSERT( data.check() );
    // Create the output directory
    std::string path = global_IO_path + "/" + subdir;
    recursiveMkdir( path, S_IRWXU | S_IRGRP );
    // Write the mesh files
    std::vector<IO::MeshDatabase> meshes_written;
    if ( global_IO_format == Format::OLD ) {
        // Write the original triangle format
        meshes_written = writeMeshesOrigFormat( meshData, path, rank );
    } else if ( global_IO_format == Format::NEW ) {
        // Write the new format (double precision)
        meshes_written = writeMeshesNewFormat( meshData, path, IO::FileFormat::NEW, rank );
    } else if ( global_IO_format == Format::SILO ) {
        // Write silo
        meshes_written = writeMeshesSilo( meshData, path, IO::FileFormat::SILO, rank );
    } else {
        ERROR( "Unknown format" );
    }
    // Gather a complete list of files on rank 0
    meshes_written = gatherAll( meshes_written, comm );
    // Write the summary files
    if ( rank == 0 ) {
        // Write the summary file for the current timestep
        char filename[200];
        sprintf( filename, "%s/LBM.summary", path.c_str() );
        write( meshes_written, filename );
// Write summary silo file if needed
#ifdef USE_SILO
        if ( global_IO_format == Format::SILO ) {
            sprintf( filename, "%s/summary.silo", path.c_str() );
            writeSiloSummary( meshes_written, filename );
        }
#endif
        // Add the timestep to the global summary file
        if ( global_IO_format == Format::OLD || global_IO_format == Format::NEW ) {
            auto filename = global_IO_path + "/summary.LBM";
            FILE *fid     = fopen( filename.c_str(), "ab" );
            fprintf( fid, "%s/\n", subdir.c_str() );
            fclose( fid );
        } else if ( global_IO_format == Format::SILO ) {
            auto filename = global_IO_path + "/LBM.visit";
            FILE *fid     = fopen( filename.c_str(), "ab" );
            fprintf( fid, "%s/summary.silo\n", subdir.c_str() );
            fclose( fid );
        } else {
            ERROR( "Unknown format" );
        }
    }
    PROFILE_STOP( "writeData" );
}