import adios2
from mpi4py import MPI
import numpy

comm = MPI.COMM_WORLD
rank = comm.Get_rank()
size = comm.Get_size()

adios = adios2.ADIOS(MPI.COMM_WORLD)
bpIO = adios.DeclareIO("BPFile")
bpIO.SetEngine('bp4')

bpIO.DefineAttribute("Fides_Data_Model", "xgc")

meshVars = ["psi", "n_n", "n_t"]
diagVars = ["psi_mks", "i_gc_density_1d", "i_perp_temperature_df_1d", "i_parallel_mean_en_df_1d"]
dVars = ["nphi", "dpot", "pot0", "potm0"]

varList = meshVars + diagVars + dVars
bpIO.DefineAttribute("Fides_Variable_List", varList)

varAssoc = []
varSources = []
arrayTypes = []
for i in range(len(varList)):
    if i < len(meshVars):
        varAssoc.append("field_data")
        varSources.append("mesh")
        arrayTypes.append("basic")
    elif i < len(meshVars) + len(diagVars):
        varAssoc.append("field_data")
        varSources.append("diag")
        arrayTypes.append("basic")
    else:
        varSources.append("3d")
        if varList[i] == "dpot":
            varAssoc.append("points")
            arrayTypes.append("xgc_field")
        else:
            varAssoc.append("field_data")
            arrayTypes.append("basic")

bpIO.DefineAttribute("Fides_Variable_Associations", varAssoc)
bpIO.DefineAttribute("Fides_Variable_Sources", varSources)
bpIO.DefineAttribute("Fides_Variable_Array_Types", arrayTypes)

bpFileWriter = bpIO.Open("xgc-attr.bp", adios2.Mode.Write)
bpFileWriter.Close()
