""" This script is an example of how to setup airfoil shape optimization with pyLAVA.
    See the documentation for more details.
    
    you can run the script like this:
    mpirun -np 2 mbind.x -t 14 python run_opt.py
"""
# imports-start
# === Import ================================================================

# Auxiliary imports
from mpi4py import MPI
import yaml

# Core imports
from pylava.pyLAVA_optimization import optLAVA
from baseclasses import AeroProblem
import numpy as np
# ===========================================================================
# imports-end

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

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


with open("opt.yaml", "r") as f:
    opt_options = yaml.safe_load(f)
        
# set the output directory to be the same for run data and optimization data
outputDir = opt_options['output']['outputDir']
lava_options['curvsolver']['outputDir'] = outputDir
# modify the options to disable output during ducring the convergence of a single CFD solve.
lava_options['output']['restart'] = {'frequency': -1}
lava_options['output']['surface'] = [{'frequency': -1}]
lava_options['output']['volume']  = [{'frequency': -1}]
# loadyamls-end 


# init-start
# Instantiate the class
solver = optLAVA(lava_options, opt_options)

# Initialize the optimization problem
solver.initOptimization(opt_options)

# this gives a visual summary of the design variables, objective, and constraints
solver.optProb.printSparsity()
# init-end
# run-start

# Optimize
sol = solver()

# Check Solution
if (solver.comm.rank == 0):
    print(sol)

solver.AeroSolver.writeSurface(f"{outputDir}/surface.final.plt", fields=['cp', 'cfx'])
# run-end

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