# This script gets stresses from a solution at selected times and outputs them in Ansys format

import os
import sys

file_name = mw.file_name()
our_dir = os.path.dirname(mw.file_name())
# Check for a solution before we start doing any real work
# If there's no solution, test_sol will be an empty list
test_sol = solution.all_elements()
if test_sol:
    stress_times = []
    stress_time = 100000
    while stress_time != "":
        stress_time = mw.input("Please enter a time step (blank to end)")
        try:
            # Check that a number was given. Anything other than a number will be treated as ending the input
            stress_time = float(stress_time)
            stress_times.append(stress_time)
        except:
            stress_time = ""
    stress_times.sort()
    # Make sure times are in ascending order
    ele_nodes = mw.all_nodes()
    # Get all the nodes
    steps = solution.time_steps()
    # Cast to list or the Python doesn't know what to do with it
    steps = list(steps)
    extra_nodes = False
    no_err = True
    first_step = True
    if not stress_times:
        # If no times provided, do nothing
        no_err = False
    for step in steps:
        for time in stress_times:
            if solution.time(step) >= time:
                # Times may not be exact, so check if the step is greater than or equal to a provided input time.
                # If not, do nothing. If so, treat it as a normal step and get all the temps.
                stress_times.remove(time)
                # Lists for normal and shear stresses
                s_xx = []; s_yy = []; s_zz = []
                s_xy = []; s_yz = []; s_xz = []
                # Assumed a transient input, so there will be time steps
                try:
                    # Add the stress values to their respective lists
                    for node in ele_nodes:
                        s_xx.append(solution.node_value("stressxx", node, step))
                        s_yy.append(solution.node_value("stressyy", node, step))
                        s_zz.append(solution.node_value("stresszz", node, step))
                        s_xy.append(solution.node_value("stressxy", node, step))
                        s_yz.append(solution.node_value("stressyz", node, step))
                        s_xz.append(solution.node_value("stresszx", node, step))
                except:
                    # If more nodes are added after the solution is run, the script will try and fail
                    # to get the stresses of the new nodes because they have no temps.
                    extra_nodes = True
                try:
                    # Write to file. One file per timestep
                    with open(str(file_name)[:-5] + "_Time_" + str(solution.time(step)) + ".stress", 'w') as f:
                        # If it's the first time step, make a new file. Else add to existing file
                        count = 0
                        for node in ele_nodes:
                            f.write(" " + "{:6d}".format(node) + " ")
                            # Format is "3, 273.15"
                            # First number is node ID, second is stress (not 100% on unit, believe pascals)
                            f.write("{:11.4e}".format(s_xx[count]))
                            # After the first stress, add the rest to the same line
                            f.write("{:11.4e}".format(s_yy[count]))
                            f.write("{:11.4e}".format(s_zz[count]))
                            f.write("{:11.4e}".format(s_xy[count]))
                            f.write("{:11.4e}".format(s_yz[count]))
                            f.write("{:11.4e}".format(s_xz[count]) + "\n")
                            count = count + 1
                except:
                    if no_err:
                        mw.message("Could not open a stress file! Is it already open?")
                        no_err = False
                        # Prevent the error message showing up many times if the file is already open
    if extra_nodes:
        mw.message("WARNING: Total nodes did not match solution nodes. Some nodes will not have a stress.")
    if no_err:
        mw.message("Done! Node stresses are in .stress files")
else:
    # If we get here, there is no solution
    mw.message("No solution found! Check that your solution ran and was saved.")