"""Test for a feature doc showing how to use cache_linear_solution"""
from __future__ import division

from distutils.version import LooseVersion
import unittest
from copy import deepcopy
from six.moves import cStringIO

import numpy as np
import scipy
from scipy.sparse.linalg import gmres

import openmdao.api as om
from openmdao.utils.assert_utils import assert_rel_error


class CacheLinearTestCase(unittest.TestCase):

    def test_feature_cache_linear(self):

        from distutils.version import LooseVersion
        import numpy as np
        import scipy
        from scipy.sparse.linalg import gmres

        import openmdao.api as om


        class QuadraticComp(om.ImplicitComponent):
            """
            A Simple Implicit Component representing a Quadratic Equation.

            R(a, b, c, x) = ax^2 + bx + c

            Solution via Quadratic Formula:
            x = (-b + sqrt(b^2 - 4ac)) / 2a
            """

            def setup(self):
                self.add_input('a', val=1.)
                self.add_input('b', val=1.)
                self.add_input('c', val=1.)
                self.add_output('states', val=[0,0])

                self.declare_partials(of='*', wrt='*')

            def apply_nonlinear(self, inputs, outputs, residuals):
                a = inputs['a']
                b = inputs['b']
                c = inputs['c']
                x = outputs['states'][0]
                y = outputs['states'][1]

                residuals['states'][0] = a * x ** 2 + b * x + c
                residuals['states'][1] = a * y + b

            def solve_nonlinear(self, inputs, outputs):
                a = inputs['a']
                b = inputs['b']
                c = inputs['c']
                outputs['states'][0] = (-b + (b ** 2 - 4 * a * c) ** 0.5) / (2 * a)
                outputs['states'][1] = -b/a

            def linearize(self, inputs, outputs, partials):
                a = inputs['a'][0]
                b = inputs['b'][0]
                c = inputs['c'][0]
                x = outputs['states'][0]
                y = outputs['states'][1]

                partials['states', 'a'] = [[x**2],[y]]
                partials['states', 'b'] = [[x],[1]]
                partials['states', 'c'] = [[1.0],[0]]
                partials['states', 'states'] = [[2*a*x+b, 0],[0, a]]

                self.state_jac = np.array([[2*a*x+b, 0],[0, a]])

            def solve_linear(self, d_outputs, d_residuals, mode):

                if mode == 'fwd':
                    print("incoming initial guess", d_outputs['states'])
                    if LooseVersion(scipy.__version__) < LooseVersion("1.1"):
                        d_outputs['states'] = gmres(self.state_jac, d_residuals['states'], x0=d_outputs['states'])[0]
                    else:
                        d_outputs['states'] = gmres(self.state_jac, d_residuals['states'], x0=d_outputs['states'], atol='legacy')[0]
                elif mode == 'rev':
                    if LooseVersion(scipy.__version__) < LooseVersion("1.1"):
                        d_residuals['states'] = gmres(self.state_jac, d_outputs['states'], x0=d_residuals['states'])[0]
                    else:
                        d_residuals['states'] = gmres(self.state_jac, d_outputs['states'], x0=d_residuals['states'], atol='legacy')[0]

        p = om.Problem()
        p.driver = om.ScipyOptimizeDriver()
        p.driver.options['optimizer'] = 'SLSQP'

        indeps = p.model.add_subsystem('indeps', om.IndepVarComp(), promotes_outputs=['a', 'b', 'c'])
        indeps.add_output('a', 1.)
        indeps.add_output('b', 4.)
        indeps.add_output('c', 1.)
        p.model.add_subsystem('quad', QuadraticComp(), promotes_inputs=['a', 'b', 'c'], promotes_outputs=['states'])
        p.model.add_subsystem('obj', om.ExecComp('y = (x[1]-x[0])**2', x=np.ones(2)))
        p.model.connect('states', 'obj.x')

        p.model.add_design_var('a', cache_linear_solution=True)
        p.model.add_constraint('states', upper=10)
        p.model.add_objective('obj.y')

        p.setup(mode='fwd')
        p.run_driver()

        print(p['a'], p['b'], p['c'])
        print(p['states'])
        assert_rel_error(self, p['obj.y'], 0.25029766, 1e-3)


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

