# ==================================================================
# === Optimization =================================================
# ==================================================================
""" This script is an example of how to setup a wing twist 
    optimization case using LAVA and OpenMDAO
    See the documentation for more details.
    
    you can run the script like this:
    mpirun -np 2 mbind.x -t 14 python run_opt.py
"""

# === Import =======================================================
from mpi4py import MPI
import yaml
import os
import numpy as np

# imports-start
import openmdao.api as om
from mphys import Multipoint, MPhysVariables
from mphys.scenarios.aerodynamic import ScenarioAerodynamic
from pylava.mphys.mphys_lava import LAVABuilder
from pygeo.mphys import OM_DVGEOCOMP
# imports-end
# ==================================================================

# === Model ========================================================
class Top(Multipoint):
    def setup(self) :

        # == LAVA Setup == 

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

        # builder-start
        # LAVA Builder
        lava_builder = LAVABuilder(lava_options, scenario="aerodynamic",
                                   lava_suffix='3D.NS1.INTEL.AVX2.OPENMPI-lib',
                                   design_surf_fams=["wall"],
                                   functionals=["cl","cd","cmy"])
        lava_builder.initialize(self.comm)
        # builder-end

        # mphys-setup-start
        # # == MPhys Setup ==

        # Keep DVs at the top level
        self.add_subsystem("dvs", om.IndepVarComp(), promotes=["*"])

        # Add mesh to subsystems
        self.add_subsystem('mesh', lava_builder.get_mesh_coordinate_subsystem())

        # Add the geometry component from pyGeo
        self.add_subsystem("geometry", OM_DVGEOCOMP(
            file=os.path.join("ffd","ffd.xyz"), type="ffd"))

        # We're cruisin
        self.mphys_add_scenario("cruise", ScenarioAerodynamic(aero_builder=lava_builder))

        # Ensure that the mesh and cruise coordinates are the same
        self.connect(
            f"mesh.{MPhysVariables.Aerodynamics.Surface.Mesh.COORDINATES}", 
            f"geometry.{MPhysVariables.Aerodynamics.Surface.Geometry.COORDINATES_INPUT}"
        )
        self.connect(
            f"geometry.{MPhysVariables.Aerodynamics.Surface.Geometry.COORDINATES_OUTPUT}", 
            "cruise.x_aero")
        # mphys-setup-end

    # configure-start
    def configure(self) :
        """
        Create the aero problem here
        """

        # Create the geometric DVs
        points = self.mesh.mphys_get_surface_mesh()

        # Add pointset to pyGeo's FFD
        self.geometry.nom_add_discipline_coords(MPhysVariables.Aerodynamics.Surface.Geometry, points)
        
        # == Geometry Setup ==

        # Create a reference axis
        nRefAxPts = self.geometry.nom_addRefAxis(name="c4", xFraction=0.25, alignIndex="k")
        nTwist = nRefAxPts

        # Setup the global twist variable
        def twist(val, geo) :
            for i in range(0, nRefAxPts) :
                geo.rot_z["c4"].coef[i] = val[i]

        self.geometry.nom_addGlobalDV(dvName="twist", value=np.zeros((nTwist)), func=twist)

        # Add the DVs to the independent vars and connect them
        self.dvs.add_output("twist", val=np.zeros((nTwist)))
        self.connect("twist", "geometry.twist")

        # Define the design variables for openMDAO
        self.add_design_var("twist", lower=-10.0, upper=10.0, scaler=1.0)

        # Add the aerodynamic objective and constraint
        self.add_objective("cruise.cd", scaler=1.0)
        self.add_constraint("cruise.cl", equals=0.5, scaler=1.0, linear=False)
        self.add_constraint("cruise.cmy", lower=-0.17, scaler=1.0, linear=False)
    # configure-end

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


# === Run ==========================================================
if __name__ == "__main__":
    
    # main-start
    # Instantiate the MDAO problem
    prob = om.Problem(allow_post_setup_reorder=False)
    prob.model = Top()
    
    # Set the optimizer
    prob.driver = om.ScipyOptimizeDriver()
    prob.driver.options["optimizer"] = "SLSQP"
    prob.driver.opt_settings = {
        "IFILE": "SLSQP.out"
    }
    
    # Attach recorder to the driver
    recorder = om.SqliteRecorder('cases.sql')
    prob.driver.add_recorder(recorder)

    prob.setup(mode="rev")
    prob.run_driver()
    # main-end

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