#!/net/python/bin/python
import numpy as num
import pylab
from pprint import pprint

DEBUG = False

def unitvector(n,N):
    """
    returns an N-dimensional unit vector in direction n
    """
    assert n<N and n>=0
    v = num.zeros(N)
    v[n] = 1.0
    return v

def replace_highest(p,f_p,p_star,f_star):
    """
    returns p,f_p with p_star and f_star inserted, and p[-1] and f[-1] removed

     - f_p is a list of evaluations f(p) sorted in ascending order
     - p is a list of the corresponding vectors
     - p_star is a point to add to the list
     - f_star is the evaluation f(p_star)
    """
    j = num.searchsorted(f_p,f_star)
    p = p[:j] + [p_star] + p[j:-1]
    f_p = f_p[:j] + [f_star] + f_p[j:-1]
    return p,f_p

def minimize(f,p0,Nmax=1000,ftol = 1E-5,scale=1.0,quiet=True):
    n = len(p0)
    N = n+1 #number of points in the simplex

    p = [num.array(p0)] + [p0+unitvector(j,n) for j in range(n)]
    f_p = [f(pi) for pi in p]
    min_list = [p[0]]

    #sort p and f_p
    eval = [(f_p[j],p[j].tolist()) for j in range(n+1)]
    eval.sort()
    f_p = [e[0] for e in eval]
    p = [num.array(e[1]) for e in eval]

    #perform minimization
    i = 0

    while True:
        i+=1
        if i > Nmax:
            print "Maximum iterations reached! (%i)" % Nmax
            break

        old_min = p[0]

        #perform simplex step
        ####
        if DEBUG:
            print "------------\nSimplex Step"
            print '\t',p
            print '\t',f_p
        ####
        p_bar = sum(p[:-1]) * 1.0/n
        p_star = p_bar + scale*(p_bar-p[-1])
        f_star = f(p_star)

        if f_star <= f_p[0]:
            #this is better than the best point, so
            # try extrapolating two times farther
            p_star2 = p_bar + 2*scale*(p_bar - p[-1])
            f_star2 = f(p_star2)

            if f_star < f_star2:
                ####
                if DEBUG:
                    print "\tusing p*...",p_star
                ####
                j = num.searchsorted(f_p,f_star)
                p = p[:j] + [p_star] + p[j:-1]
                f_p = f_p[:j] + [f_star] + f_p[j:-1]
                
            else:
                ####
                if DEBUG:
                    print "\tusing p**...",p_star2
                ####
                j = num.searchsorted(f_p,f_star2)
                p = p[:j] + [p_star2] + p[j:-1]
                f_p = f_p[:j] + [f_star2] + f_p[j:-1]
                
        elif f_star >= f_p[1]:
            #f_star is worse than the second highest point,
            # so we will try a contraction
            p_dag = 0.5*scale*(p_bar+p[-1])
            f_dag = f(p_dag)
            
            if f_dag < f_p[0]: #f_dag < f(p_min)
                ####
                if DEBUG:
                    print "\tusing p+...",p_dag
                ####
                j = num.searchsorted(f_p,f_dag)
                p = p[:j] + [p_dag] + p[j:-1]
                f_p = f_p[:j] + [f_dag] + f_p[j:-1]
                
            else: #none of these options is better! - reduce the dimension
                ####
                if DEBUG:
                    print "\treducing dimension..."
                ####
                p_l = p[0]
                p = [0.5*(p_l+pi) for pi in p]
                #sort p and f_p
                eval = [(f_p[j],p[j].tolist()) for j in range(n+1)]
                eval.sort()
                f_p = [e[0] for e in eval]
                p = [num.array(e[1]) for e in eval]
        else:
            #we will use f_star
            ####
            if DEBUG:
                print "\tusing p* after all...",p_star
            ####
            #replace highest
            j = num.searchsorted(f_p,f_star)
            p = p[:j] + [p_star] + p[j:-1]
            f_p = f_p[:j] + [f_star] + f_p[j:-1]
            
        min_list.append(p[0])

        if ftol > f_p[-1]-f_p[0]:
            if not quiet:
                print "converged to f(p) = %.5g" % f_p[0]
                print "at (x1,x2) = ", p[0]
                print "   after %i iterations!" % i
                print "   ftol = %.2g" % ftol
                print "   scale factor = %.1f" % scale
            break
    
    return min_list,f_p[0],i

########################################################
# functions for debugging
if __name__=="__main__":
    def Rosenbrock(x):
        """
        returns Rosenbrock's function evaluated at x.
    
        x is an array [x1,x2,x3...] with an even number
        of entries
        """
        n = len(x)
        S = 0

        for i in range( int(n/2) ):
            j = 2*i #j will be an even number < n-1
            S += ( (1-x[j])**2 + 100*(x[j+1] - x[j]**2)**2 )

        return S
    
    def plotRosenbrock(xmin=-1.5,xmax=1.5,ymin=-0.5,ymax=1.5):
        """
        produces a contour plot of Rosenbrock's function
        """
        x1 = num.arange(xmin,xmax,0.01)
        x2 = num.arange(ymin,ymax,0.01)
        z = [[num.log10(Rosenbrock((i,j))) for i in x1] for j in x2]
        pylab.contour(x1,x2,z,70)

    def TestScale( p0=(-1.0,1.0) ):
        """
        tests the effect of the scale parameter on the
        convergence of minimization of Rosenbrock's Function
        from p0
        """
        print "scale\tNumber of Steps"
        for s in num.arange(0.1,2.1,0.1):
            min_list,f_min,nsteps = minimize(Rosenbrock,p0,\
                                             ftol=1E-6,scale=s)
            print "%.2f\t%i" % (s,nsteps)

    def Test_p0(N=1000):
        """
        tests N different random values of p0 for
        convergence of Rosenbrock's function,
        and reports on the convergence statistics
        """
        min_N = 1E6
        min_p0 = ()
        max_N = 0
        max_p0 = ()
        for i in range(N):
            p0 = ( -1.5 + 3*num.random.random(),-0.5+2*num.random.random() )
            min_list,f_min,nsteps = minimize(Rosenbrock,p0,ftol=1E-6,scale=1)
            if nsteps < min_N:
                min_N = nsteps
                min_p0 = p0
            if nsteps > max_N:
                max_N = nsteps
                max_p0 = p0

        print "Testing %i random starting points:" % N
        print "  Maximum: %i steps at (%.2f,%.2f)" % (max_N,max_p0[0],max_p0[1])
        print "  Minimum: %i steps at (%.2f,%.2f)" % (min_N,min_p0[0],min_p0[1])

    def plotConvergence(p0=(-1.0,1.0),ftol=1E-7,quiet=False):
        """
        plots the convergence path for the 2 dimensional
        Rosenbrock Function
        """
        min_list,f_min,nsteps = minimize(Rosenbrock,p0,ftol=ftol,quiet=quiet)
        plotRosenbrock()
        pylab.plot([m[0] for m in min_list],[m[1] for m in min_list],'-b.')


    #TestScale()
    #Test_p0()
    plotConvergence()
    pylab.show()