import json
import imas,os,datetime,sys
import getpass
import numpy as np
from imas import imasdef
 

db = 'data_access_tutorial'
shot=99357 
run=1

#creates the Data Entry object 'data_entry', a kind of handler of the pulse file with sho, run, belonging to database 'data_access_tutorial' of the current user, using the MDS+ backend
data_entry = imas.DBEntry(imasdef.MDSPLUS_BACKEND, db, shot, run, user_name=getpass.getuser())
#data_entry = imas.DBEntry(imasdef.HDF5_BACKEND, db, shot, run, user_name=getpass.getuser())
# Open save data_entry to data base

# Tries to open data_entry
op = data_entry.open()
#if fails, creates one
op = data_entry.open()
#open() and create(0 return a tuple (x_int,y_int) where x<0, y>0, x number of failures, y number of successes in the current session. 
#if open fails, create data_entry
if op[0]<0:
     cp=data_entry.create()
     if cp[0]==0:
        print("data entry created")
elif op[0]==0:
    print("data entry opened")	

# Fetched data Open file with fetched data, the data is not numpy array format
with open("data/data_jet_hrts_99357.json") as json_file:
    hrts=json.load(json_file)

x_coord=np.array(hrts['TE']['x'])
# no. of space points
nb_points = len(x_coord)
#no of time slices
nb_slices=len(hrts['TE']['time'])

#here, we can perform some read/write operations using the get/put() operations
#...

#creating the 'thomson_scattering' auxiliary IDS and initializing it
thomson = imas.thomson_scattering() #creates a 'thomson scattering' IDS
thomson.ids_properties.homogeneous_time=1 #setting the homogeneous time (mandatory)	
thomson.ids_properties.comment='IDS created for testing the IMAS Data Access layer' #setting the ids_properties.comment attribute
#thomson.time=np.array([0.]) #the time(vector) basis must be not empty if homogeneous_time==1 otherwise an error will occur at runtime

# since all data is available we can save whole time vector at once
thomson.time=np.array(hrts['TE']['time'])

thomson.ids_properties.creation_date = datetime.datetime.now().strftime("%y-%m-%d")

%m-%d")
# the number of channel corresponds to number of data points
thomson.channel.resize(nb_points)

for j in range(nb_points):
    thomson.channel[j].t_e.data.resize(1)
    thomson.channel[j].t_e.data_error_upper.resize(1)
    thomson.channel[j].n_e.data.resize(1)
    thomson.channel[j].n_e.data_error_upper.resize(1)

# python interface accepts only numpy arrays to be saved in ids
te_data = np.array(hrts['TE']['data']) #2D
dte_data = np.array(hrts['DTE']['data'])#2D
ne_data = np.array(hrts['NE']['data'])#2D
dne_data = np.array(hrts['DNE']['data'])#2D	
z_data = np.array(hrts['Z']['data']) #1D
r_data = np.array(hrts['TE']['x']) #1D
	
for j in range(nb_points):
    thomson.channel[j].position.r=r_data[j]
    thomson.channel[j].position.z=z_data[j] 
    thomson.channel[j].t_e.data=te_data[:,j]
    thomson.channel[j].t_e.data_error_upper=dte_data[:,j]
    thomson.channel[j].n_e.data=ne_data[:,j]
    thomson.channel[j].n_e.data_error_upper=dne_data[:,j]

sequence = hrts['TE']['seq'] 
data_entry.put(thomson,seq)# the last number is the occurence which can be used to store the data sequence number
#closing the Data Entry
data_entry.close()