#!/usr/bin/env python
# encoding: utf-8
"""
Author(s): Matthew Loper

See LICENCE.txt for licensing and contact information.
"""

import time
from numpy import *
import unittest
from . import ch
from .optimization import minimize
from .ch import Ch
import numpy as np
from scipy.optimize import rosen, rosen_der
from .utils import row, col


visualize = False


def Rosen():
    
    args = {
        'x1': Ch(-120.),
        'x2': Ch(-100.)
    }
    r1 = Ch(lambda x1, x2 : (x2 - x1**2.) * 10., args)
    r2 = Ch(lambda x1 : x1 * -1. + 1, args)

    func = [r1, r2]
    
    return func, [args['x1'], args['x2']]

class Madsen(Ch):
    dterms = ('x',)
    def compute_r(self):
        x1 = self.x.r[0]
        x2 = self.x.r[1]
        result = np.array((
            x1**2 + x2**2 + x1 * x2,
            np.sin(x1),
            np.cos(x2)
        ))
        return result
        
    def compute_dr_wrt(self, wrt):
        if wrt is not self.x:
            return None
        jac = np.zeros((3,2))
        x1 = self.x.r[0]
        x2 = self.x.r[1]
        jac[0,0] = 2. * x1 + x2
        jac[0,1] = 2. * x2 + x1
        
        jac[1,0] = np.cos(x1)
        jac[1,1] = 0
        
        jac[2,0] = 0
        jac[2,1] = -np.sin(x2)

        return jac
    

    def set_and_get_r(self, x_in):
        self.x = Ch(x_in)
        return col(self.r)
    
    def set_and_get_dr(self, x_in):
        self.x = Ch(x_in)
        return self.dr_wrt(self.x)




class RosenCh(Ch):
    dterms = ('x',)
    def compute_r(self):
        
        result = np.array((rosen(self.x.r) ))
        
        return result

    def set_and_get_r(self, x_in):
        self.x = Ch(x_in)
        return col(self.r)
    
    def set_and_get_dr(self, x_in):
        self.x = Ch(x_in)
        return self.dr_wrt(self.x).flatten()
    

    def compute_dr_wrt(self, wrt):
        if wrt is self.x:
            if visualize:
                import matplotlib.pyplot as plt
                residuals = np.sum(self.r**2)
                print('------> RESIDUALS %.2e' % (residuals,))
                print('------> CURRENT GUESS %s' % (str(self.x.r),))
                plt.figure(123)
                
                if not hasattr(self, 'vs'):
                    self.vs = []
                    self.xs = []
                    self.ys = []
                self.vs.append(residuals)
                self.xs.append(self.x.r[0])
                self.ys.append(self.x.r[1])
                plt.clf();
                plt.subplot(1,2,1)
                plt.plot(self.vs)
                plt.subplot(1,2,2)
                plt.plot(self.xs, self.ys)
                plt.draw()


            return row(rosen_der(self.x.r))
    


class TestOptimization(unittest.TestCase):

    def test_dogleg_rosen(self):
        obj, freevars = Rosen()
        minimize(fun=obj, x0=freevars, method='dogleg', options={'maxiter': 337, 'disp': False})
        self.assertTrue(freevars[0].r[0]==1.)
        self.assertTrue(freevars[1].r[0]==1.)

    def test_dogleg_madsen(self):
        obj = Madsen(x = Ch(np.array((3.,1.))))
        minimize(fun=obj, x0=[obj.x], method='dogleg', options={'maxiter': 34, 'disp': False})
        self.assertTrue(np.sum(obj.r**2)/2 < 0.386599528247)
        
    @unittest.skip('negative sign in exponent screws with reverse mode')
    def test_bfgs_rosen(self):
        from .optimization import minimize_bfgs_lsq        
        obj, freevars = Rosen()        
        minimize_bfgs_lsq(obj=obj, niters=421, verbose=False, free_variables=freevars)        
        self.assertTrue(freevars[0].r[0]==1.)
        self.assertTrue(freevars[1].r[0]==1.)

    def test_bfgs_madsen(self):
        from .ch import SumOfSquares
        import scipy.optimize
        obj = Ch(lambda x : SumOfSquares(Madsen(x = x)) )
        
        def errfunc(x):
            obj.x = Ch(x)
            return obj.r
        
        def gradfunc(x):
            obj.x = Ch(x)
            return obj.dr_wrt(obj.x).ravel()
        
        x0 = np.array((3., 1.))

        # Optimize with built-in bfgs.
        # Note: with 8 iters, this actually requires 14 gradient evaluations.
        # This can be verified by setting "disp" to 1.
        #tm = time.time()
        x1 = scipy.optimize.fmin_bfgs(errfunc, x0, fprime=gradfunc, maxiter=8, disp=0)
        #print 'forward: took %.es' % (time.time() - tm,)
        self.assertLess(obj.r/2., 0.4)

        # Optimize with chumpy's minimize (which uses scipy's bfgs).
        obj.x = x0
        minimize(fun=obj, x0=[obj.x], method='bfgs', options={'maxiter': 8, 'disp': False})
        self.assertLess(obj.r/2., 0.4)

    def test_nested_select(self):
        def beales(x, y):
            e1 = 1.5 - x + x*y
            e2 = 2.25 - x  + x*(y**2)
            e3 = 2.625 - x + x*(y**3)
            return {'e1': e1, 'e2': e2, 'e3': e3}

        x1 = ch.zeros(10)
        y1 = ch.zeros(10)

        # With a single select this worked
        minimize(beales(x1, y1), x0=[x1[1:4], y1], method='dogleg', options={'disp': False})

        x2 = ch.zeros(10)
        y2 = ch.zeros(10)

        # But this used to raise `AttributeError: 'Select' object has no attribute 'x'`
        minimize(beales(x2, y2), x0=[x2[1:8][:3], y2], method='dogleg', options={'disp': False})
        np.testing.assert_array_equal(x1, x2)
        np.testing.assert_array_equal(y1, y2)


suite = unittest.TestLoader().loadTestsFromTestCase(TestOptimization)

if __name__ == '__main__':
    
    if False: # show rosen 
        import matplotlib.pyplot as plt
        visualize = True
        plt.ion()
        unittest.main()
        import pdb; pdb.set_trace()
    else:
        unittest.main()