import numpy as np
from scipy.io import loadmat

from surropt.caballero import Caballero
from surropt.core.options.nlp import DockerNLPOptions
from surropt.core.procedures.output import Report
from tests_ import RESOURCES_PATH

from win32com.client import Dispatch


def open_connection():
    aspen = Dispatch('Apwn.Document')
    aspen.InitFromArchive2(str(RESOURCES_PATH / 'infill.bkp'))
    return aspen


def aspen_auto(x: np.ndarray, asp_srv):
    # set input values
    asp_srv.Tree.FindNode(r"\Data\Blocks\TOWER\Input\BASIS_RR").Value = x[0]
    asp_srv.Tree.FindNode(r"\Data\Blocks\TOWER\Input\D:F").Value = x[1]

    # run the results
    asp_srv.Engine.Run2()

    # convergence flag
    conv = asp_srv.Tree.FindNode(
        r"\Data\Results Summary\Run-Status\Output\UOSSTAT2").Value

    if conv == 8:
        # run converged, sample results
        status = True

        RR = asp_srv.Tree.FindNode("\Data\Blocks\TOWER\Output\MOLE_RR").Value
        DF = asp_srv.Tree.FindNode("\Data\Blocks\TOWER\Output\MOLE_DFR").Value
        D = asp_srv.Tree.FindNode("\Data\Streams\D\Output\TOT_FLOW").Value
        XB = asp_srv.Tree.FindNode(
            "\Data\Streams\B\Output\MOLEFRAC\MIXED\PROPENE").Value
        B = asp_srv.Tree.FindNode("\Data\Streams\B\Output\TOT_FLOW").Value
        QR = asp_srv.Tree.FindNode(
            "\Data\Blocks\TOWER\Output\REB_DUTY").Value * 4.1868
        L = asp_srv.Tree.FindNode("\Data\Blocks\TOWER\Output\MOLE_L1").Value
        V = asp_srv.Tree.FindNode("\Data\Blocks\TOWER\Output\MOLE_VN").Value
        F = asp_srv.Tree.FindNode("\Data\Streams\FEED\Output\TOT_FLOW").Value
        XD = asp_srv.Tree.FindNode(
            "\Data\Streams\D\Output\MOLEFRAC\MIXED\PROPENE").Value

        J = -(20 * D + (10 - 20 * XB) * B - 70 * QR)

        g = [0.995 - XD, QR - 80.0]

    else:
        status = False
        J = 0
        g = [1.0, 1.0]  # value above con_tol to ensure infeasibility

    results = {
        'status': status,
        'f': J,
        'g': g,
        'extras': []
    }
    return results


def main():
    # load initial doe
    m_c = loadmat(RESOURCES_PATH / 'column33pts.mat')

    # MVs
    x = m_c['doeBuild'][:, :2]

    # constraints
    xd_raw = m_c['doeBuild'][:, 4]
    qr_raw = m_c['doeBuild'][:, 7]

    g1 = 0.995 - xd_raw
    g2 = qr_raw - 80.0

    g = np.c_[g1, g2]

    # objective function
    f = m_c['doeBuild'][:, 6]

    # open the connection
    aspen = open_connection()

    # sampling function
    def model_function(x): return aspen_auto(x, asp_srv=aspen)

    # bounds
    lb = [7.0, 0.1]
    ub = [25.0, 0.9]

    # nlp server options
    nlp_opts = DockerNLPOptions(name='wsl-server',
                                server_url='http://localhost:5000')

    # report options object
    # this is the object where output generated by the caballero procedure is
    # reported to the user via terminal or plot window
    report_opts = Report(terminal=True, plot=False)

    caballero_obj = Caballero(x=x, g=g, f=f, model_function=model_function,
                              lb=lb, ub=ub, regression='poly1',
                              nlp_options=nlp_opts,
                              report_options=report_opts)

    caballero_obj.optimize()


if __name__ == "__main__":
    main()
