"""Tests the `debug_print` option for Nonlinear solvers."""

from __future__ import division, print_function

import os
import re
import shutil
import tempfile

import unittest
from distutils.version import LooseVersion

import numpy as np

from openmdao.core.problem import Problem
from openmdao.core.indepvarcomp import IndepVarComp

from openmdao.solvers.nonlinear.nonlinear_block_gs import NonlinearBlockGS
from openmdao.solvers.nonlinear.nonlinear_block_jac import NonlinearBlockJac
from openmdao.solvers.nonlinear.newton import NewtonSolver

from openmdao.test_suite.test_examples.test_circuit_analysis import Circuit

from openmdao.utils.assert_utils import assert_rel_error
from openmdao.utils.general_utils import run_model
from openmdao.utils.general_utils import printoptions

from parameterized import parameterized


nonlinear_solvers = [
    NonlinearBlockGS,
    NonlinearBlockJac,
    NewtonSolver
]


class TestNonlinearSolvers(unittest.TestCase):
    def setUp(self):
        # perform test in temporary directory
        self.startdir = os.getcwd()
        self.tempdir = tempfile.mkdtemp(prefix='test_solver')
        os.chdir(self.tempdir)

        # iteration coordinate, file name and variable data are common for all tests
        coord = 'rank0:root._solve_nonlinear|0|NLRunOnce|0|circuit._solve_nonlinear|0'
        filename = coord.replace('._solve_nonlinear', '')
        self.filename = re.sub('[^0-9a-zA-Z]', '_', filename) + '.dat'

        self.expected_data = '\n'.join([
            "",
            "# Inputs and outputs at start of iteration '%s':" % coord,
            "",
            "# nonlinear inputs",
            "{'circuit.D1.V_in': array([ 1.]),",
            " 'circuit.D1.V_out': array([ 0.]),",
            " 'circuit.R1.V_in': array([ 1.]),",
            " 'circuit.R1.V_out': array([ 0.]),",
            " 'circuit.R2.V_in': array([ 1.]),",
            " 'circuit.R2.V_out': array([ 1.]),",
            " 'circuit.n1.I_in:0': array([ 0.1]),",
            " 'circuit.n1.I_out:0': array([ 1.]),",
            " 'circuit.n1.I_out:1': array([ 1.]),",
            " 'circuit.n2.I_in:0': array([ 1.]),",
            " 'circuit.n2.I_out:0': array([ 1.])}",
            "",
            "# nonlinear outputs",
            "{'circuit.D1.I': array([ 1.]),",
            " 'circuit.R1.I': array([ 1.]),",
            " 'circuit.R2.I': array([ 1.]),",
            " 'circuit.n1.V': array([ 10.]),",
            " 'circuit.n2.V': array([ 0.001])}",
            ""
        ])

    def tearDown(self):
        # clean up the temporary directory
        os.chdir(self.startdir)
        try:
            shutil.rmtree(self.tempdir)
        except OSError:
            pass

    @parameterized.expand([
        [solver.__name__, solver] for solver in nonlinear_solvers
    ])
    def test_solver_debug_print(self, name, solver):
        p = Problem()
        model = p.model

        model.add_subsystem('ground', IndepVarComp('V', 0., units='V'))
        model.add_subsystem('source', IndepVarComp('I', 0.1, units='A'))
        model.add_subsystem('circuit', Circuit())

        model.connect('source.I', 'circuit.I_in')
        model.connect('ground.V', 'circuit.Vg')

        p.setup()

        nl = model.circuit.nonlinear_solver = solver()

        nl.options['debug_print'] = True

        # suppress solver output for test
        nl.options['iprint'] = model.circuit.linear_solver.options['iprint'] = -1

        # set some poor initial guesses so that we don't converge
        p['circuit.n1.V'] = 10.
        p['circuit.n2.V'] = 1e-3

        opts = {}
        # formatting has changed in numpy 1.14 and beyond.
        if LooseVersion(np.__version__) >= LooseVersion("1.14"):
            opts["legacy"] = '1.13'

        with printoptions(**opts):
            # run the model and check for expected output file
            output = run_model(p)

        expected_output = '\n'.join([
            self.expected_data,
            "Inputs and outputs at start of iteration "
            "have been saved to '%s'.\n" % self.filename
        ])

        self.assertEqual(output, expected_output)

        with open(self.filename, 'r') as f:
            self.assertEqual(f.read(), self.expected_data)

    def test_solver_debug_print_feature(self):
        from openmdao.api import Problem, IndepVarComp, NewtonSolver
        from openmdao.test_suite.test_examples.test_circuit_analysis import Circuit

        p = Problem()
        model = p.model

        model.add_subsystem('ground', IndepVarComp('V', 0., units='V'))
        model.add_subsystem('source', IndepVarComp('I', 0.1, units='A'))
        model.add_subsystem('circuit', Circuit())

        model.connect('source.I', 'circuit.I_in')
        model.connect('ground.V', 'circuit.Vg')

        p.setup()

        nl = model.circuit.nonlinear_solver = NewtonSolver()

        nl.options['iprint'] = 2
        nl.options['debug_print'] = True

        # set some poor initial guesses so that we don't converge
        p['circuit.n1.V'] = 10.
        p['circuit.n2.V'] = 1e-3

        opts = {}
        # formatting has changed in numpy 1.14 and beyond.
        if LooseVersion(np.__version__) >= LooseVersion("1.14"):
            opts["legacy"] = '1.13'

        with printoptions(**opts):
            # run the model
            p.run_model()

        with open('rank0_root_0_NLRunOnce_0_circuit_0.dat', 'r') as f:
            self.assertEqual(f.read(), self.expected_data)


if __name__ == "__main__":
    unittest.main()
