##############################################################################
# Copyright (c) 2024, Oak Ridge National Laboratory                          #
# All rights reserved.                                                       #
#                                                                            #
# This file is part of HydraGNN and is distributed under a BSD 3-clause      #
# license. For the licensing terms see the LICENSE file in the top-level     #
# directory.                                                                 #
#                                                                            #
# SPDX-License-Identifier: BSD-3-Clause                                      #
##############################################################################


##############################################################################
########################  INSTRUCTIONS TO RUN THE CODE #######################
##############################################################################

"""
1. Follow run instructions on HydraGNN Wiki page: https://github.com/ORNL/HydraGNN/wiki/Run
2. Copy the "Ensemble_of_models" folder with the models into "examples/ensemble_learning"
3. Copy the script "inference_example_with_dummy_data_object.py" into "examples/ensemble_learning"
"""

import json, os
from mpi4py import MPI
import argparse

import torch
from torch_geometric.data import Data
from torch_geometric.transforms import Distance

import hydragnn

try:
    from hydragnn.utils.adiosdataset import AdiosWriter, AdiosDataset
except ImportError:
    pass

from hydragnn.preprocess.utils import (
    RadiusGraph,
)

radius_graph = RadiusGraph(5.0, loop=False, max_num_neighbors=50)
transform_coordinates = Distance(norm=False, cat=False)

from ensemble_utils import model_ensemble

######################################################
###### PYTORCH GEOMETRIC DATA OBJECT DEFINITION ######
######################################################

# Define the dummy inputs
lattice_mat = torch.tensor([[10.0, 0.0, 0.0], [0.0, 10.0, 0.0], [0.0, 0.0, 10.0]], dtype=torch.float32)  # 3x3 supercell size matrix
positions = torch.tensor([[0.5, 0.5, 0.5], [1.5, 1.5, 1.5]], dtype=torch.float32)  # Nx3 positions matrix
atomic_numbers = torch.tensor([1, 6], dtype=torch.int64).unsqueeze(1)  # Atomic numbers as a 1D tensor
x = torch.cat([atomic_numbers, positions], dim=1)

# Creating the Data object
data = Data(
    supercell_size=lattice_mat,
    pos=positions,
    atomic_numbers=atomic_numbers.view(-1, 1),  # Reshape atomic_numbers to Nx1 tensor
    x=x,
)

data = radius_graph(data)
data = transform_coordinates(data)

######################################################
###### LOAD MODELS AND USE THEM FOR PREDICTIONS ######
######################################################


if __name__ == "__main__":
    ##################################################################################################################
    parser = argparse.ArgumentParser()
    print("gfm starting")
    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    parser.add_argument("--models_dir_folder", help="folder of trained models", type=str, default="./Ensemble_of_models")

    args = parser.parse_args()
    args.parameters = vars(args)
    ##################################################################################################################
    modeldirlists = args.models_dir_folder.split(",")
    assert len(modeldirlists) == 1 or len(modeldirlists) == 2
    if len(modeldirlists) == 1:
        modeldirlist = [os.path.join(args.models_dir_folder, name) for name in os.listdir(args.models_dir_folder) if
                        os.path.isdir(os.path.join(args.models_dir_folder, name))]
    else:
        modeldirlist = []
        for models_dir_folder in modeldirlists:
            modeldirlist.extend([os.path.join(models_dir_folder, name) for name in os.listdir(models_dir_folder) if
                                 os.path.isdir(os.path.join(models_dir_folder, name))])

    var_config = None
    for modeldir in modeldirlist:
        input_filename = os.path.join(modeldir, "config.json")
        with open(input_filename, "r") as f:
            config = json.load(f)
        if var_config is not None:
            assert var_config == config["NeuralNetwork"][
                "Variables_of_interest"], "Inconsistent variable config in %s" % input_filename
        else:
            var_config = config["NeuralNetwork"]["Variables_of_interest"]
    verbosity = config["Verbosity"]["level"]

    ##################################################################################################################
    # Always initialize for multi-rank training.
    comm_size, rank = hydragnn.utils.setup_ddp()
    ##################################################################################################################

    comm = MPI.COMM_WORLD

    ##################################################################################################################
    model_ens = model_ensemble(modeldirlist)
    model_ens = hydragnn.utils.get_distributed_model(model_ens, verbosity)
    model_ens.eval()
    ##################################################################################################################
    nheads = len(config["NeuralNetwork"]["Variables_of_interest"]["output_names"])

    pred_ens = model_ens(data)

    for model_id in range(len(modeldirlist)):
        print(f"Model {modeldirlist[model_id]} - Prediction of energy: \n", pred_ens[model_id][0])
        print(f"Model {modeldirlist[model_id]} - Prediction of forces: \n", pred_ens[model_id][1])