
##############################################
##Code to extract nodal results A,V,U in x,y,z
##############################################

# THINGS TO UPDATE:
#   NODES: list with numbers of nodes of interest
#   OUTPUT: file names for output of data (see lines 90 to 105)
# OUTPUT FORMAT:
#   Organized by rows: time
#                      1st node in list Ux (or Vx or Ax depending on the file)
#                      1st node in list Uy 
#                      1st node in list Uz 
#                      2nd node in list Ux
#                      2nd node in list Uy 
#                      .... and so on ...
# DELETE OUTPUT FILES BEFORE RUNNING (if they already exist)  
# CODE APPENDS TO EXISTING FILE does not overwrite. 
# mpco file must be open in STKO postprocessor

from PyMpc import *
from PyMpc import MpcOdbVirtualResult as vr
import time
from time import sleep
import traceback
import numpy as np
from time import sleep
from PySide2.QtCore import (
	QObject,
	Signal,
	Slot,
	QCoreApplication,
	QTimer,
	QThread,
	QEventLoop,
	)
from PySide2.QtWidgets import (
	QWidget,
	QDialog,
	QSizePolicy,
	QVBoxLayout,
	)

# clear terminal
App.clearTerminal()

# get document
doc = App.postDocument()

# get first database
if len(doc.databases) == 0:
	raise Exception("You need a database with ID = 1 for this test")
db = doc.getDatabase(1)

# create evaluation options
# here we want to extract data for all steps of the last stage
all_stages = db.getStageIDs()
last_stage = all_stages[-1]
all_steps = db.getStepIDs(last_stage)
all_times = db.getStepTimes(last_stage) # get all step times of the last stage
opt = MpcOdbVirtualResultEvaluationOptions()
opt.stage = last_stage

Disp = db.getNodalResult("Displacement", match = MpcOdb.Contains)
Acc = db.getNodalResult("Acceleration", match = MpcOdb.Contains)
Vel = db.getNodalResult("Velocity", match = MpcOdb.Contains)

# Nodal result. #Loop over timesteps of a single node, move to next node and repeat
# Loop over nodes
#array with nodes of interest (from stko model) !!!UPDATE ACCORDINGLY!!
NODES = [98,99,100,138,114,117,119,106,111,109,112,116,136,118,120,20,24,27,46,55,51,54,18,130,23,26,126,129,132] 
t0 = time.time()
for j in NODES: #Loop for NODES
	# Reset arrays to store data
	x = []; tt= []
	Ax = []; Ay = []; Az = []
	Vx = []; Vy = []; Vz = []
	Ux = []; Uy = []; Uz = []
	print(j)
	row = MpcOdbResultField.node(j) # the node field object is the row identifier in the result field
	# evaluate for each step at current node
	#with open("temp22.txt", 'a') as myfile:
	xt = [1,2,3,4,5,6,7,8,9,10]
	for i in all_steps: #Loop for timesep
		opt.step = i
		step_time = all_times[i-1]
		field = Disp.evaluate(opt)
		Uxi = field[row, 0] #0 = Ux, 1 = Uy, 2 = Uz (the column 0-based index, depends on the result. This is for nodal results)
		Uyi = field[row, 1]
		Uzi = field[row, 2]
		field = Acc.evaluate(opt)
		Axi = field[row, 0]
		Ayi = field[row, 1]
		Azi = field[row, 2]
		tt.append(step_time); x.append(i)
		Ux.append(Uxi);	Uy.append(Uyi);	Uz.append(Uzi)
		Ax.append(Axi);	Ay.append(Ayi);	Az.append(Azi)

        
	if j == NODES[0]: # Write time data to .txt file (only the 1st time)
		with open("RigidBlock_SWEEP_Damp5_wMass_LayeredSoil3_ABC_EfromG_D.txt", 'a') as myfileU:
			myfileU.write(', '.join(str(item) for item in tt)+'\n')
		with open("RigidBlock_SWEEP_Damp5_wMass_LayeredSoil3_ABC_EfromG_A.txt", 'a') as myfileA:
			myfileA.write(', '.join(str(item) for item in tt)+'\n')
	# Write array data to .txt file
	with open("RigidBlock_SWEEP_Damp5_wMass_LayeredSoil3_ABC_EfromG_D.txt", 'a') as myfileU:
		myfileU.write(', '.join(str(item) for item in Ux)+'\n')
		myfileU.write(', '.join(str(item) for item in Uy)+'\n')
		myfileU.write(', '.join(str(item) for item in Uz)+'\n')
	with open("RigidBlock_SWEEP_Damp5_wMass_LayeredSoil3_ABC_EfromG_A.txt", 'a') as myfileA:
		myfileA.write(', '.join(str(item) for item in Ax)+'\n')
		myfileA.write(', '.join(str(item) for item in Ay)+'\n')
		myfileA.write(', '.join(str(item) for item in Az)+'\n')
with open("NODES.txt", 'w') as myfile:   
		myfile.write(', '.join(str(item) for item in NODES)+'\n') 

t1 = time.time()
print(t1-t0)

print("DONE")