# ==================================================================
# === Post-Processing ==============================================
# ==================================================================
""" This script is an example of how to post-process the wing twist 
    optimization case to extract and analyze specific design iterations
    with pyLAVA.
    See the documentation for more details.
    
    you can run the script like this:
    mpirun -np 2 mbind.x -t 14 python run_post.py
"""

# === Import =======================================================
from mpi4py import MPI
import yaml
from pylava.pyLAVA_optimization import optLAVA
from pyoptsparse import History
# ==================================================================

# === Helper Functions =============================================
# helper-start
def runAnalysis(x, fileSuffix):
    # Deform the mesh to match these DVs
    solver._setDVs(x)

    # Solver the flow
    solver.AeroSolver(resetIter=True)

    # Write the optimized volume and surface
    solver.AeroSolver.writeVolGrid(f"volume.{fileSuffix}.x")
    solver.AeroSolver.writeVolSolution(f"volume.{fileSuffix}.q", "plot3dSoln", ['pressure'])
    solver.AeroSolver.writeSurface(f"surface.{fileSuffix}.plt",
                                fields=["iblank", "yplus", "cp", "cf", "cfx", "cfy", "cfz"])
# helper-end
# ==================================================================

# === Run ==========================================================

# Read the YAML file name for lava options
with open("lava.yaml", "r") as f:
    lava_options = yaml.safe_load(f)

# Read in the YAML file for optimization options
with open("opt.yaml", "r") as f:
    opt_options = yaml.safe_load(f)

# Instantiate the class
solver = optLAVA(lava_options, opt_options, comm=MPI.COMM_WORLD)

# hist-start
# Read the history file using pyOptSparse's History Class
hist = History("Output/history.hst", flag="r")
vals = hist.getValues()
# hist-end

# run-start
# Set the DVs to the first iteration
# and call the helper function above to solve the flow and
# write the geometry + solution
dvDict = {'wingffd_twist' : vals['wingffd_twist'][0]}
runAnalysis(dvDict, "initial")

# Set the DVs to the last iteration
# and call the helper function above to solve the flow and
# write the geometry + solution
dvDict = {'wingffd_twist' : vals['wingffd_twist'][-1]}
runAnalysis(dvDict, "optimized")
# run-end

# ==================================================================