/*---------------------------------------------------------------------------*\
  =========                 |
  \\      /  F ield         | OpenFOAM: The Open Source CFD Toolbox
   \\    /   O peration     |
    \\  /    A nd           | www.openfoam.com
     \\/     M anipulation  |
-------------------------------------------------------------------------------
    Copyright (C) 2016 OpenFOAM Foundation
    Copyright (C) 2018-2021 OpenCFD Ltd.
-------------------------------------------------------------------------------
License
    This file is part of OpenFOAM.

    OpenFOAM is free software: you can redistribute it and/or modify it
    under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    OpenFOAM is distributed in the hope that it will be useful, but WITHOUT
    ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
    FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
    for more details.

    You should have received a copy of the GNU General Public License
    along with OpenFOAM.  If not, see <http://www.gnu.org/licenses/>.

Application
    FoamToGraph

Description
    FoamToGraph By H.Tofighian

\*---------------------------------------------------------------------------*/
#include <torch/torch.h>
#include <torch/script.h>

#include "argList.H"
#include "timeSelector.H"
#include "volFields.H"

#include <stdio.h>
#include <stdlib.h>

using namespace Foam;

// * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * //

int main(int argc, char *argv[])
{
    timeSelector::addOptions();

    #include "setRootCase.H"
    #include "createTime.H"
    instantList timeDirs = timeSelector::select0(runTime, args);
    #include "createNamedMesh.H"


    // Create output directory in case directory
    fileName outputDir(runTime.rootPath()/runTime.globalCaseName()/"graph_data");
    Foam::mkDir(outputDir);
    Info<< "Saving graph data to: " << outputDir << nl << endl;

    forAll(timeDirs, timei)
    {
        runTime.setTime(timeDirs[timei], timei);

        Info<< "Time = " << runTime.timeName() << endl;
        Info<< "Reading field U\n" << endl;
        volVectorField U
        (
            IOobject
            (
                "U",
                runTime.timeName(),
                mesh,
                IOobject::MUST_READ,
                IOobject::NO_WRITE
            ),
            mesh
        );
        
        const label nCells = mesh.cells().size();
        
        // Extract node features: use the velocity components (x, y, z) from each cell.
        std::vector<float> node_feat;
        node_feat.reserve(3 * nCells);
        forAll(U, i)
        {
            node_feat.push_back(U[i].x());
            node_feat.push_back(U[i].y());
            node_feat.push_back(U[i].z());
        }
        // Create a LibTorch tensor for node features with shape [nCells, 3].
        torch::Tensor node_features = torch::from_blob(node_feat.data(), {static_cast<long>(nCells), 3}, torch::kFloat32).clone(); // Shape: [num_nodes, 3]
        
        // Build edge indices using mesh connectivity.
        // In OpenFOAM, each internal face connects two cells. The mesh provides
        // owner and neighbour lists; we add both directions for an undirected graph.
        const labelList& owner = mesh.owner();
        const labelList& neighbour = mesh.neighbour();
        const size_t numFaces = owner.size();
        const size_t numEdges = 2*numFaces;

        // Separate lists for source and target indices.
        std::vector<int64_t> edge_sources;
        std::vector<int64_t> edge_targets;

        // Reserve space for two directed edges per face.
        edge_sources.reserve(numEdges);
        edge_targets.reserve(numEdges);

        forAll(owner, i)
        {
            // Forward edge: owner -> neighbour.
            edge_sources.push_back(owner[i]);
            edge_targets.push_back(neighbour[i]);
            // Reverse edge: neighbour -> owner.
            edge_sources.push_back(neighbour[i]);
            edge_targets.push_back(owner[i]);
        }
        // Create a tensor for edge indices; stack two vectors.
        torch::Tensor edge_index = torch::stack({
            torch::from_blob(edge_sources.data(), {static_cast<long>(numEdges)}, torch::kInt64),
            torch::from_blob(edge_targets.data(), {static_cast<long>(numEdges)}, torch::kInt64)
        }).clone(); // Shape: [2, num_edges]
        
        // Extract node positions: cell centers (x, y, z) for each cell
        std::vector<float> node_pos;
        node_pos.reserve(3 * nCells);
        
        // Get cell centers
        const pointField& cellCenters = mesh.C();
        
        // Add each cell center coordinate to the positions vector
        forAll(cellCenters, i)
        {
            node_pos.push_back(cellCenters[i].x());
            node_pos.push_back(cellCenters[i].y());
            node_pos.push_back(cellCenters[i].z());
        }
        
        // Create a LibTorch tensor for node positions with shape [nCells, 3]
        torch::Tensor node_positions = torch::from_blob(node_pos.data(), {static_cast<long>(nCells), 3}, torch::kFloat32).clone(); // Shape: [num_nodes, 3]



        // Save tensors
        // Save node features
        {
            fileName nodeFile(outputDir/"node_features_" + runTime.timeName() + ".pt");
            auto node_bytes = torch::pickle_save(node_features);
            std::ofstream fout(nodeFile, std::ios::out | std::ios::binary);
            fout.write(node_bytes.data(), node_bytes.size());
            fout.close();
            Info<< "Saved node features to " << nodeFile << endl;
        }

        // Save edge indices
        {
            fileName edgeFile(outputDir/"edge_index_" + runTime.timeName() + ".pt");
            auto edge_bytes = torch::pickle_save(edge_index);
            std::ofstream fout(edgeFile, std::ios::out | std::ios::binary);
            fout.write(edge_bytes.data(), edge_bytes.size());
            fout.close();
            Info<< "Saved edge indices to " << edgeFile << endl;
        }
        // Save node positions
        {
            fileName posFile(outputDir/"node_positions_" + runTime.timeName() + ".pt");
            auto pos_bytes = torch::pickle_save(node_positions);
            std::ofstream fout(posFile, std::ios::out | std::ios::binary);
            fout.write(pos_bytes.data(), pos_bytes.size());
            fout.close();
            Info<< "Saved node positions to " << posFile << endl;
        }
    }
    
    Info << "Execution complete." << nl;

    return 0;
}


// ************************************************************************* //
