""" Unit test for the Brent, one variable nonlinear solver. """

import unittest


from openmdao.api import Group, Problem, Component, Brent, ScipyGMRES
from openmdao.test.util import assert_rel_error


class CompPart1(Component): 

    def __init__(self): 
        super(CompPart1, self).__init__()
        self.fd_options['force_fd'] = True

        self.add_param('a', val=1.)
        self.add_param('n', val=77.0/27.0)
        self.add_param('part2', val=0.)

        self.add_state('x', val=2., lower=0, upper=100)

    def solve_nonlinear(self, p, u, r): 
        pass

    def apply_nonlinear(self, p, u, r): 
        r['x'] = p['a'] * u['x']**p['n'] + p['part2'] #+ p['b'] * u['x'] - p['c']
        # print self.pathname, "ap_nl", p['part2'], p['a'], u['x'], p['n'], r['x']    


class CompPart2(Component): 

    def __init__(self): 
        super(CompPart2, self).__init__()
        self.fd_options['force_fd'] = True

        self.add_param('b', val=1.)
        self.add_param('c', val=10.)
        self.add_param('x', val=2.)
        self.add_output('part2', val=0.)

    def solve_nonlinear(self, p, u, r): 
        
        u['part2'] = p['b'] * p['x'] - p['c']
        # print self.pathname, "sp_nl", p['x'], p['c'], p['b'], u['part2']


class Combined(Group): 

    def __init__(self): 
        super(Combined, self).__init__()

        self.add('p1', CompPart1(), promotes=['*'])
        self.add('p2', CompPart2(), promotes=['*'])
        # self.add('i1', IndepVarComp('a', 1.), promotes=['*'])
        # self.add('i2', IndepVarComp('b', 1.), promotes=['*'])

        # self.nl_solver = Newton()
        self.nl_solver = Brent()
        # self.nl_solver.options['iprint'] = 1
        self.nl_solver.options['state_var'] = 'x'

        self.ln_solver = ScipyGMRES()

        self.set_order(('p1','p2'))


class BrentMultiCompTestCase(unittest.TestCase): 
    """test to make sure brent can converge multiple components 
    in a group with a single residual across them all
    """

    def test_multi_comp(self): 
        p = Problem()
        p.root = Combined()
        p.setup(check=False)
        p.run()

        assert_rel_error(self, p.root.unknowns['x'], 2.06720359226, .0001)
        assert_rel_error(self, p.root.resids['x'], 0, .0001)

if __name__ == "__main__": 

   unittest.main()