# ==================================================================
# === 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 os
import openmdao.api as om
from run_om_opt import Top
# ==================================================================

# == Helper Functions ==============================================
# analysis-start
def runAnalysis(prob, case, fileSuffix) :

    # Get the output directory
    out_dir = prob.get_outputs_dir()

    # Load the case and run the model
    prob.load_case(case)
    prob.run_model()

    # Get lava and write the surface output
    lava = prob.model.cruise.coupling.aero_solver
    lava.writeSurface(os.path.join(out_dir,f"surface.{fileSuffix}.plt"), 
        fields=["iblank", "yplus", "cp", "cf", "cfx", "cfy", "cfz"])
# analysis-end

# print-case-start
def printCase(label, case):
    print(f"\n{label}")
    print("-" * len(label))

    print("Objectives:")
    for name, val in case.get_objectives().items():
        print(f"  {name}: {val[0]:.8f}")

    print("Constraints:")
    for name, val in case.get_constraints().items():
        print(f"  {name}: {val[0]:.8f}")
# print-case-end
# ==================================================================

# === Main Script ==================================================
if __name__ == "__main__":
    
    # Instantiate the MDAO problem
    prob = om.Problem()
    prob.model = Top()
    prob.setup(mode="rev")

    # reader-start
    cr = om.CaseReader('run_om_opt_out/cases.sql')
    # reader-end

    # run-start
    # Get the first function eval
    case = cr.get_case(0)
    runAnalysis(prob, case, "initial")

    # Get the last function eval
    case = cr.get_case(-1)
    runAnalysis(prob, case, "optimized")
    # run-end
    
    # print-start
    # Print the initial and optimized objectives
    if MPI.COMM_WORLD.rank == 0 :
        printCase("INITIAL CASE", cr.get_case(0))
        printCase("OPTIMIZED CASE", cr.get_case(-1))
    # print-end
# ==================================================================