#!/usr/bin/python

import atomization, setupgen
import sys, pickle, random
from LinearAlgebra import inverse
import Numeric as N

from datetime import datetime, timedelta

N_MAX = 100
ALPHA = 1.
BETA = .5
GAMMA = 2.


"""
Default test function with one minimum at (1,2,3, ...).
The minimum is exactly 42. Takes a list of coordinates as an argument
and returns a number.
"""
def standardFunction(p):
    y = 42
    for i in range(len(p)):
        y += (p[i]-(i+1))**2
    return y

"""
Performs the 'amoeba'-like downhill simplex method in ndim dimensions.

p: a list of (ndim+1) vectors each with ndim coordinates,
corresponding to the vertices of the simplex.

y: a list of function values evaluated at the vertices, ordered consistently with the vertices in p. y thus must have length (ndim+1) as well

ndim: the dimension count of the space in question. Of course this variable is mostly for show since it's not really necessary in python

fTolerance: fractional tolerance used to evaluate convergence criterion

function: the function to be minimized. The function must take exactly
ndim parameters, each parameter being one number

maxIterations: the maximal number of iterations to be performed before
returning, in case convergence is slow

Returns the number of times the function has been evaluated during the
procedure.

After invocation the argument lists p and y will have been modified to contain
the simplex vertices and associated function values at termination of the
procedure.

"""
def amoeba(p, y, ndim, fTolerance, function=standardFunction, out=sys.stdout, dump='lastdump.dump.pckl'):

    mpts = ndim + 1
    evaluationCount = 0
    #This is probably the coordinate sum, i.e.
    #it probably has to do with the geometric center of the simplex
    psum = getpsum(p)

    while True:
        print >> out, 'Points:', p
        print >> out, 'yValues:', y
        print >> out, 'EvalCount:',evaluationCount
        print >> out
        out.flush()

        #Write current points to file for recovery if something goes wrong
        pickleDump((p,y),dump)
        
        iLow = 0 #index of lowest value
        iHigh = None #index of highest value
        i2ndHigh = None #index of second highest value
        if y[0] > y[1]:
            (iHigh, i2ndHigh) = (0, 1)
        else:
            (iHigh, i2ndHigh) = (1, 0)

        #Loop through vertices to find index values for highest/lowest entries
        for i in range(mpts):
            if y[i] < y[iLow]:
                iLow = i
            if y[i] > y[iHigh]:
                i2ndHigh = iHigh
                iHigh = i
            elif y[i] > y[i2ndHigh]:
                if i != iHigh:
                    i2ndHigh = i

        #Things should be floats already, but it's good to be safe
        relDeviation = float(abs(y[iHigh] - y[iLow]))/abs(y[iHigh]+y[iLow])

        print >> out,'Rel. deviation', relDeviation
        out.flush()
        
        if relDeviation < fTolerance:
            break

        if evaluationCount >= N_MAX:
            print '=== Max evaluation count',N_MAX,'exceeded, terminating! ==='
            #Some would call this an error, but we'll just return
            #as if nothing has happened
            break

        yTry = amotry(p, y, psum, ndim, function, iHigh, -ALPHA)
        evaluationCount += 1

        if yTry <= y[iLow]:
            yTry = amotry(p, y, psum, ndim, function, iHigh, GAMMA)
            evaluationCount += 1
        elif yTry >= y[i2ndHigh]:
            ySave = y[iHigh]
            yTry = amotry(p, y, psum, ndim, function, iHigh, BETA)
            evaluationCount += 1
            if yTry >= ySave:
                for i in range(mpts):
                    if i != iLow:
                        for j in range(ndim):
                            psum[j] = .5 * (p[i][j] + p[iLow][j])
                            p[i][j] = psum[j]
                        y[i] = function(psum)
                evaluationCount += ndim
                psum = getpsum(p)

    return evaluationCount
                            
"""
Extrapolates through or partway to simplex face, possibly finding a better
vertex
"""
def amotry(p, y, psum, ndim, function, iHigh, factor):
    #Wonder what these 'factors' do exactly
    factor1 = (1. - factor)/ndim
    factor2 = factor1 - factor

    pTry = [psum[j]*factor1 - p[iHigh][j]*factor2 for j in range(ndim)]
    
    yTry = function(pTry)

    if yTry < y[iHigh]:
        y[iHigh] = yTry
        for j in range(ndim):
            psum[j] += pTry[j] - p[iHigh][j]
            p[iHigh][j] = pTry[j]
    
    return yTry


"""
Given a list of (ndim+1) vectors each with ndim coordinates,
returns the list of coordinate sums across vectors,
i.e. the n'th element is the sum of the n'th coordinates of all vectors in p
"""
def getpsum(p):
    mpts = len(p)
    ndim = mpts - 1
    
    psum = array(ndim)
    for i in range(ndim):
        psum[i] = sum([q[i] for q in p])

    return psum


"""
Returns a list of vertex coordinates forming a regular simplex around the
designated center, where the size argument is the max vertex-center distance.

This method simply generates a random simplex, and may fail to do so at a
very small probability (if randomly generated vectors are linearly dependent)
"""
def getInitialPoints(center=[0,0], size=1, seed=0):
    ndim = len(center)
    mpts = ndim + 1
    r = random.Random(seed)

    points = array(mpts)
    for i in range(ndim+1):
        points[i] = [(r.random()-.5)*size+center[j] for j in range(ndim)]
        
    return points

"""
Runs the amoeba optimization function with sensible values
"""
def smalltest():
    f = standardFunction
    p = getInitialPoints([7,3,2,6,3])
    print 'Initial points gotten'
    print 'Mapping p through f'
    y = map(f, p)
    print 'Done mapping'
    ndim = len(p)-1
    fTolerance = .000001

    amoeba(p, y, ndim, fTolerance, f)

    print 'Done!'
    #print 'p', p
    #print 'y', y
    print 'p[0]',p[0]

"""
The badness function is enclosed in a class in order to make it possible for
it to 'remember' a setup name
"""
class SetupEvaluator:

    def __init__(self, setup):
        setup = 'opt.'+setup
        self.setup = setup
        self.generator = setupgen.SetupGenerator(setup)

    """
    Runs a full test of a given GPAW setup
    """
    def badness(self, args):
        refEnergy = -10.55
        refDist = 1.102

        try:
            self.generator.f(args) #new setup
            print 'New setup created'

            overallBadness = 0

            startTime = datetime.now()

            print 'Calculating atomization energy'
            energyBadness = 1/.05**2 #badness == 1 for deviation == .05 eV
            (c1,c2) = atomization.atomizationCalculators(out=None,
                                                         setup=self.setup)
            Ea = atomization.calcEnergy(c1,c2,a=6.0)
            print 'Energy',Ea
            db = energyBadness * (Ea - refEnergy)**2
            print 'Energy badness',db
            overallBadness += db

            print 'Calculating bond length'
            d = bondLength(self.setup)
            distanceBadness = 1./.005**2
            db = distanceBadness * (d - refDist)**2
            overallBadness += db
            print 'Bond length',d
            print 'Bond length badness',db

            print 'Calculating energy fluctuation amplitude'
            DE = energyFluctuationTest(self.setup)
            energyFluctuationBadness = 1./.005**2
            db = energyFluctuationBadness * DE**2
            overallBadness += db
            print 'Fluctuation magnitude',DE
            print 'Fluctuation badness',db

            print 'Calculating convergence rate'
            hVar = convergenceTest(self.setup)
            convergenceBadness = 1./.2**2
            db = convergenceBadness * hVar
            overallBadness += db
            print 'Energy difference',hVar
            print 'Energy difference badness',db

            print 'Calculating temporal badness'
            timeBadness = 1./20**2 #20 seconds --> badness == 1
            dt = (datetime.now() - startTime).seconds
            db =  timeBadness * dt**2
            overallBadness += db
            print 'Time',dt
            print 'Time badness',db

            print 'Overall badness',overallBadness

        except KeyboardInterrupt:
            raise KeyboardInterrupt #Don't ignore keyboard interrupts
        #except:
        #    return 10000.

        return overallBadness

"""
Returns the bond length. Calculates energy at three locations around the
reference bond length, interpolates with a 2nd degree polynomial and returns
the minimum of this polynomial which would be roughly equal to the bond length
without engaging in a large whole relaxation test
"""
def bondLength(setup):
    print 'Distance test'
    d0 = 1.102
    dd = ( .2 / 140. )**.5 #around .04 A. Bond properties correspond to
    #an energy of E = .5 k x**2 with k = 140 eV/A**2
    #If we want .1 eV deviation then the above dd should be used
    calc = atomization.MakeCalculator(atomization.molN.nbands2,
                                      out=None, setup=setup)
    D = [d0-dd, d0, d0+dd]
    #Calculate energies at the three points
    E = [atomization.energyAtDistance(d, calc=calc, a=5.5) for d in D]
    print 'Distances',D
    print 'Energies',E
    print
    #Now find parabola and determine minimum

    x = N.array(D)
    y = N.array(E)

    A = N.transpose(N.array([x**0, x**1, x**2]))
    c = N.dot(inverse(A), y)
    print 'Coordinates',c

    X = - c[1] / (2.*c[2]) # "-b/(2a)"
    print 'Bond length',X

    #print c
    #print N.dot(A, c) - y
    
    return X

"""
Returns whatever was written by pickledump
"""
def pickleLoad(fileName):
    content = pickle.load(open(fileName))
    return content

"""
Writes the data to the file, deleting any other content
"""
def pickleDump(data,fileName):
    pickle.dump(data, open(fileName, 'w'))

"""
Runs a complete optimization.

All files pertaining to the test will have file names containing
the testName parameter.

fTolerance is the termination tolerance of the
amoeba function.

If initData is specified it should be the file name of
a pickle file containing a set of points and values (p,y) that are
compatible with the amoeba function as specified in its documentation.
"""
def seriousTest(testName='test', fTolerance = .002, ndims=2, initData=None):

    fileName=testName+'.log'

    outFile = open(fileName, 'w')
    evaluator = SetupEvaluator(testName)
    
    f = evaluator.badness
    sensibledata = [1.1, 1.1, 1.1, 2.2, 0.4]
    #Defaults for new_nitrogen_setup. Warning: hardcoded here,
    #be sure to change this if new_nitrogen_setup is changed
    
    if initData is None:
        dumpFile = testName+'.dump.pckl'

        p = getInitialPoints(sensibledata[:ndims], size=.1)
        print >> outFile, 'Initial points gotten'
        print >> outFile, 'Points',p
        print >> outFile, 'Mapping p through f'
        outFile.flush()
        y = map(f, p)
        print >> outFile, 'Done mapping, dumping to file'
        print >> outFile, 'y values:',y
        pickleDump((p,y), dumpFile)
    else:
        (p,y) = pickleLoad(dumpFile)
        print >> outFile, 'Initial values loaded from',dumpFile
        print >> outFile, 'p: ',p
        print >> outFile, 'y: ',y

    outFile.flush()

    ndim = len(p)-1

    print >> outFile, 'Amoeba commencing'
    print >> outFile, '-----------------'
    outFile.flush()

    amoeba(p,y,ndim,fTolerance, f, out=outFile, dump=dumpFile)

"""
Plots the energy of a N2 moleculeas a function of different resolutions
(h-values) and returns the maximal difference
"""
def convergenceTest(setup):
    A = atomization
    h = [.15, .17, .20]
    calc = [A.MakeCalculator(A.molN.nbands2, out=None, h=h0,
                             setup=setup) for h0 in h]
    E = [A.energyAtDistance(A.molN.d, calc=c, a=4.) for c in calc]

    print 'h',h
    print 'E',E
    
    return max(E) - min(E)

"""
Returns the difference between the energy of a nitrogen molecule at the center
of the unit cell and the energy of one translated by h/2 along the z axis.
"""
def energyFluctuationTest(setup):
    A = atomization
    h = .2
    calc = A.MakeCalculator(A.molN.nbands2, out=None, setup=setup, h=h)
    d = A.molN.d
    E1 = A.energyAtDistance(d, calc=calc, a=4.)
    E2 = A.energyAtDistance(d, calc=calc, a=4., dislocation=(0.,0.,h/2.))

    return E2-E1

"""
'[None]*length' looks slightly silly so let's initialise our lists with a
nicely named function instead
"""
def array(length):
    return [None]*length