import time
import numpy as num
import scipy.linalg
import gnumpy as gnp

from conjugateGrad_numpy import genPosDefSym, CGStoppingCriterion, NegativeCurvatureException

def CGObjectiveFromScratch(A,b,x):
    return 0.5*(x*A(x)).sum() - (b*x).sum()

def CGObjective(r, b, x):
    """
    b and x are a column vectors of the same shape
    
    r is b-A(x)
    """
    return 0.5*((-r-b)*x).sum()

def precondConjGrad(A, b, x0, invM, stoppingCriterion, residualRefreshPeriod = None, log = None):
    """
    The Untransformed Preconditioned Conjugate Gradient Algorithm
    Minimize f(x) = 0.5*x^{T} A x - b^T x
    by applying unpreconditioned CG to the system
    E^{-1} A E^{-T} x = E^{-1} b  where
    EE^{T} = M
    So M is the preconditioner which is why we need a function that
    can apply M^{-1} to vectors.

    A must be a function that consumes a column vector v and returns
    the result of premultiplying v by A.

    b is a column vector
    
    x0 is a column vector

    invM must be a function that consumes a column vector v and returns
    the result of premultiplying v by the inverse of M.
    """
    if log == None:
        log = open("/dev/null",'w')
    r = b-A(x0) #r is the current residual
    x = x0 #x is the current point
    d = invM(r) #d is the current search direction
    deltaNew = (r*d).sum()
    i = 0
    curCost = CGObjective(r, b, x)
    print >>log, "CG iter", i, curCost
    while not stoppingCriterion.doWeStop():
        Ad = A(d)
        energy2Norm = (d*Ad).sum()
        if energy2Norm <= 0:
            raise NegativeCurvatureException()
        alpha = deltaNew/energy2Norm
        x += alpha*d
        
        if residualRefreshPeriod != None and i > 0 and i % residualRefreshPeriod == 0: 
            r = b - A(x) #refresh residual in case floating point error accumulates too much
        else:
            r -= alpha*Ad
        prevCost = curCost
        curCost = CGObjective(r, b, x)
        #(i, self.num_cg, model_loss, cur_diff, est_loss, V_norm_I, V_norm_A, np.sqrt((b*b).sum())))
        print >>log, "CG iter", i+1, curCost, "|x|=", num.sqrt((x*x).sum()), "|x'Ax|=", num.sqrt((x*(b-r)).sum()), \
              "|b|=", num.sqrt((b*b).sum())
        s = invM(r)
        deltaOld = deltaNew
        deltaNew = (r*s).sum()
        beta = deltaNew/deltaOld
        d = s+beta*d
        i += 1
        stoppingCriterion.observeCGIterates(curCost, prevCost, beta, deltaNew, deltaOld, alpha, energy2Norm, r, b, x)
        yield x
    

def f(A,b,x):
    return 0.5*(x*num.dot(A,x)).sum() - (b*x).sum()

def test2(seed):
    num.random.seed(seed)
    n = 2000#1000
    #A = genPosDefSym(n, std=0.5, seed = seed)
    A = genPosDefSym(n, seed = seed)
    b = num.random.randn(n,1)
    A = gnp.garray(A)
    b = gnp.garray(b)

    MInv = num.diag(1/A.asarray().diagonal())
    MInv = gnp.garray(MInv)
    
    applyA = lambda v: gnp.dot(A,v)
    applyMInv = lambda v: gnp.dot(MInv,v)

    #stoppingCriterion = CGStoppingCriterion(maxIters=1000, minIters=25, tol=5e-4, gapRatio=0.1, minGap = 10)
    stoppingCriterion = CGStoppingCriterion(maxIters=1000, minIters=25, tol=0.0, gapRatio=0.1, minGap = 10)
    t = time.time()
    cgIters = precondConjGrad(applyA, b, gnp.zeros((n,1)), applyMInv, stoppingCriterion, residualRefreshPeriod = 50)
    for j,x in enumerate(cgIters):
        print j+1, stoppingCriterion.CGCosts[-1]
    cgAns = x
    print "Elapsed time (s):", time.time()-t
    ans = scipy.linalg.solve(A.asarray(),b.asarray(), sym_pos=True)
    print ans.shape
    print cgAns.shape

    print '"exact":', f(A.asarray(),b.asarray(),ans)
    print "my CG:", f(A.asarray(),b.asarray(),cgAns.asarray())
    print
    
    print '"exact":', num.linalg.norm(num.dot(A.asarray(),ans)-b.asarray())
    print "my CG:", num.linalg.norm(num.dot(A.asarray(),cgAns.asarray())-b.asarray())
    eigs1 = num.linalg.eigvalsh(A.asarray())
    
    print min(eigs1), max(eigs1)

def prep(seed):
    num.random.seed(seed)
    n = 2000#1000
    A = genPosDefSym(n, seed = seed)
    b = num.random.randn(n,1)
    A = gnp.garray(A)
    b = gnp.garray(b)
    MInv = num.diag(1/A.asarray().diagonal())
    MInv = gnp.garray(MInv)
    return A, b, MInv

def profileThis(A,b, MInv):
    applyA = lambda v: gnp.dot(A,v)
    applyMInv = lambda v: gnp.dot(MInv,v)
    
    stoppingCriterion = CGStoppingCriterion(maxIters=1000, minIters=25, tol=0.0, gapRatio=0.1, minGap = 10)
    t = time.time()
    cgIters = precondConjGrad(applyA, b, gnp.zeros((n,1)), applyMInv, stoppingCriterion, residualRefreshPeriod = 50)
    for j,x in enumerate(cgIters):
        print j+1, stoppingCriterion.CGCosts[-1]
    cgAns = x
    print "Elapsed time (s):", time.time()-t
    

if __name__ == "__main__":
    test2(9)
