#  ___________________________________________________________________________
#
#  Pyomo: Python Optimization Modeling Objects
#  Copyright (c) 2008-2022
#  National Technology and Engineering Solutions of Sandia, LLC
#  Under the terms of Contract DE-NA0003525 with National Technology and
#  Engineering Solutions of Sandia, LLC, the U.S. Government retains certain
#  rights in this software.
#  This software is distributed under the 3-clause BSD License.
#  ___________________________________________________________________________

import os.path
from pyomo.common.fileutils import this_file_dir, import_file
import pyomo.common.unittest as unittest
import pyomo.environ as pyo
from pyomo.common.dependencies import attempt_import
from pyomo.common.log import LoggingIntercept
from pyomo.opt import TerminationCondition
from io import StringIO
import logging

from pyomo.contrib.pynumero.dependencies import (
    numpy as np, numpy_available, scipy, scipy_available
)
from pyomo.common.dependencies.scipy import sparse as spa

if not (numpy_available and scipy_available):
    raise unittest.SkipTest("Pynumero needs scipy and numpy to run CyIpopt tests")

pandas, pandas_available = attempt_import(
    'pandas',
    'One of the tests below requires a recent version of pandas for'
    ' comparing with a tolerance.',
    minimum_version='1.1.0',
    defer_check=False)

from pyomo.contrib.pynumero.asl import AmplInterface
if not AmplInterface.available():
    raise unittest.SkipTest(
        "Pynumero needs the ASL extension to run CyIpopt tests")

import pyomo.contrib.pynumero.algorithms.solvers.cyipopt_solver as cyipopt_solver
if not cyipopt_solver.cyipopt_available:
    raise unittest.SkipTest(
        "PyNumero needs CyIpopt installed to run CyIpopt tests")
import cyipopt as cyipopt_core

example_dir = os.path.join(this_file_dir(), '..')

class TestPyomoCyIpoptSolver(unittest.TestCase):
    def test_status_maps(self):
        # verify that all status messages from cyipopy can be cleanly
        # mapped back to a Pyomo TerminationCondition
        for msg in cyipopt_core.STATUS_MESSAGES.values():
            self.assertIn(msg, cyipopt_solver._cyipopt_status_enum)
        for status in cyipopt_solver._cyipopt_status_enum.values():
            self.assertIn(status, cyipopt_solver._ipopt_term_cond)


class TestExamples(unittest.TestCase):
    def test_external_grey_box_react_example_maximize_cb_outputs(self):
        ex = import_file(os.path.join(example_dir, 'external_grey_box', 'react_example', 'maximize_cb_outputs.py'))
        m = ex.maximize_cb_outputs()
        self.assertAlmostEqual(pyo.value(m.reactor.inputs['sv']), 1.34381, places=3)
        self.assertAlmostEqual(pyo.value(m.reactor.outputs['cb']), 1072.4372, places=2)

    def test_external_grey_box_react_example_maximize_cb_outputs_scaling(self):
        ex = import_file(os.path.join(example_dir, 'external_grey_box', 'react_example', 'maximize_cb_ratio_residuals.py'))
        aoptions={'nlp_scaling_method': 'user-scaling',
                 'output_file': '_cyipopt-external-greybox-react-scaling.log',
                 'file_print_level':10}
        m = ex.maximize_cb_ratio_residuals_with_output_scaling(additional_options=aoptions)
        self.assertAlmostEqual(pyo.value(m.reactor.inputs['sv']), 1.26541996, places=3)
        self.assertAlmostEqual(pyo.value(m.reactor.inputs['cb']), 1071.7410089, places=2)
        self.assertAlmostEqual(pyo.value(m.reactor.outputs['cb_ratio']), 0.15190409266, places=3)

        with open('_cyipopt-external-greybox-react-scaling.log', 'r') as fd:
            solver_trace = fd.read()
        os.remove('_cyipopt-external-greybox-react-scaling.log')

        self.assertIn('nlp_scaling_method = user-scaling', solver_trace)
        self.assertIn('output_file = _cyipopt-external-greybox-react-scaling.log', solver_trace)
        self.assertIn('objective scaling factor = 1', solver_trace)
        self.assertIn('x scaling provided', solver_trace)
        self.assertIn('c scaling provided', solver_trace)
        self.assertIn('d scaling provided', solver_trace)
        self.assertIn('DenseVector "x scaling vector" with 7 elements:', solver_trace)
        self.assertIn('x scaling vector[    2]= 1.2000000000000000e+00', solver_trace)
        self.assertIn('x scaling vector[    7]= 1.7000000000000000e+00', solver_trace)
        self.assertIn('x scaling vector[    6]= 1.1000000000000001e+00', solver_trace)
        self.assertIn('x scaling vector[    1]= 1.3000000000000000e+00', solver_trace)
        self.assertIn('x scaling vector[    3]= 1.3999999999999999e+00', solver_trace)
        self.assertIn('x scaling vector[    4]= 1.5000000000000000e+00', solver_trace)
        self.assertIn('x scaling vector[    5]= 1.6000000000000001e+00', solver_trace)
        self.assertIn('DenseVector "c scaling vector" with 6 elements:', solver_trace)
        self.assertIn('c scaling vector[    1]= 4.2000000000000000e+01', solver_trace)
        self.assertIn('c scaling vector[    2]= 1.0000000000000001e-01', solver_trace)
        self.assertIn('c scaling vector[    3]= 2.0000000000000001e-01', solver_trace)
        self.assertIn('c scaling vector[    4]= 2.9999999999999999e-01', solver_trace)
        self.assertIn('c scaling vector[    5]= 4.0000000000000002e-01', solver_trace)
        self.assertIn('c scaling vector[    6]= 1.0000000000000000e+01', solver_trace)

    def test_external_grey_box_react_example_maximize_with_output(self):
        ex = import_file(os.path.join(example_dir, 'external_grey_box', 'react_example', 'maximize_cb_ratio_residuals.py'))
        m = ex.maximize_cb_ratio_residuals_with_output()
        self.assertAlmostEqual(pyo.value(m.reactor.inputs['sv']), 1.26541996, places=3)
        self.assertAlmostEqual(pyo.value(m.reactor.inputs['cb']), 1071.7410089, places=2)
        self.assertAlmostEqual(pyo.value(m.reactor.outputs['cb_ratio']), 0.15190409266, places=3)

    def test_external_grey_box_react_example_maximize_with_hessian_with_output(self):
        ex = import_file(os.path.join(example_dir, 'external_grey_box', 'react_example', 'maximize_cb_ratio_residuals.py'))
        m = ex.maximize_cb_ratio_residuals_with_hessian_with_output()
        self.assertAlmostEqual(pyo.value(m.reactor.inputs['sv']), 1.26541996, places=3)
        self.assertAlmostEqual(pyo.value(m.reactor.inputs['cb']), 1071.7410089, places=2)
        self.assertAlmostEqual(pyo.value(m.reactor.outputs['cb_ratio']), 0.15190409266, places=3)

    def test_external_grey_box_react_example_maximize_with_hessian_with_output_pyomo(self):
        ex = import_file(os.path.join(example_dir, 'external_grey_box', 'react_example', 'maximize_cb_ratio_residuals.py'))
        m = ex.maximize_cb_ratio_residuals_with_hessian_with_output_pyomo()
        self.assertAlmostEqual(pyo.value(m.sv), 1.26541996, places=3)
        self.assertAlmostEqual(pyo.value(m.cb), 1071.7410089, places=2)
        self.assertAlmostEqual(pyo.value(m.cb_ratio), 0.15190409266, places=3)

    def test_pyomo_react_example_maximize_with_obj(self):
        ex = import_file(os.path.join(example_dir, 'external_grey_box', 'react_example', 'maximize_cb_ratio_residuals.py'))
        m = ex.maximize_cb_ratio_residuals_with_obj()
        self.assertAlmostEqual(pyo.value(m.reactor.inputs['sv']), 1.26541996, places=3)
        self.assertAlmostEqual(pyo.value(m.reactor.inputs['cb']), 1071.7410089, places=2)
        self.assertAlmostEqual(pyo.value(m.obj), 0.15190409266, places=3)

    def test_external_grey_box_react_example_maximize_with_additional_pyomo_variables(self):
        ex = import_file(os.path.join(example_dir, 'external_grey_box', 'react_example', 'maximize_cb_ratio_residuals.py'))
        m = ex.maximize_cb_ratio_residuals_with_pyomo_variables()
        self.assertAlmostEqual(pyo.value(m.reactor.inputs['sv']), 1.26541996, places=3)
        self.assertAlmostEqual(pyo.value(m.reactor.inputs['cb']), 1071.7410089, places=2)
        self.assertAlmostEqual(pyo.value(m.cb_ratio), 0.15190409266, places=3)

    @unittest.skipIf(not pandas_available, "Test uses pandas for data")
    def test_parameter_estimation(self):
        data_fname = os.path.join(example_dir, 'external_grey_box', 'param_est', 'smalldata.csv')
        baseline = pandas.read_csv(data_fname)

        # test the data generator
        ex = import_file(os.path.join(example_dir, 'external_grey_box', 'param_est', 'generate_data.py'))
        df1 = ex.generate_data(5, 200, 5, 42)
        df2 = ex.generate_data_external(5, 200, 5, 42)
        pandas.testing.assert_frame_equal(df1, baseline, atol=1e-3)
        pandas.testing.assert_frame_equal(df2, baseline, atol=1e-3)

        # test the estimation
        ex = import_file(os.path.join(example_dir, 'external_grey_box', 'param_est', 'perform_estimation.py'))

        m = ex.perform_estimation_external(data_fname, solver_trace=False)
        self.assertAlmostEqual(pyo.value(m.UA), 204.43761, places=3)

        m = ex.perform_estimation_pyomo_only(data_fname, solver_trace=False)
        self.assertAlmostEqual(pyo.value(m.UA), 204.43761, places=3)

    def test_cyipopt_callbacks(self):
        ex = import_file(os.path.join(example_dir, 'callback', 'cyipopt_callback.py'))

        output = StringIO()
        with LoggingIntercept(output, 'pyomo', logging.INFO):
            ex.main()

        self.assertIn("Residuals for iteration 2",
                          output.getvalue().strip())

    @unittest.skipIf(not pandas_available, "pandas needed to run this example")
    def test_cyipopt_functor(self):
        ex = import_file(os.path.join(example_dir, 'callback', 'cyipopt_functor_callback.py'))
        df = ex.main()
        self.assertEqual(df.shape, (7, 5))
        # check one of the residuals
        s = df['ca_bal']
        self.assertAlmostEqual(s.iloc[6], 0, places=3)

    def test_cyipopt_callback_halt(self):
        ex = import_file(os.path.join(example_dir, 'callback', 'cyipopt_callback_halt.py'))
        status = ex.main()
        self.assertEqual(status.solver.termination_condition, TerminationCondition.userInterrupt)

