import numpy as num
import scipy.sparse.linalg
import scipy.linalg


def genPosDefSym(n, seed = None):
    a = num.random.randn(n,n)
    q,r = num.linalg.qr(a)
    eigs = num.random.rand(n)+0.5
    D = num.diag(eigs)
    retval = num.dot(num.dot(q,D), q.T)
    #make absolutely sure that the matrix we return is symmetric
    for i in range(n):
        for j in range(n):
            retval[i,j] = retval[j,i]
    return retval

def norm(x):
    return num.sqrt((x*x).sum())

class NegativeCurvatureException(Exception):
    pass

class CGStoppingCriterion(object):
    """
    This class decides when CG should stop. It also has to enforce
    minIters and maxIters. To make a new stopping criterion, you just
    need an object that defintes observeCGIterates and doWeStop. To
    understand what gets passed to observeCGIterates, you will need to
    look at the code for CG itself.
    """
    def __init__(self, maxIters, minIters=1, tol=5e-4, gapRatio=0.1, minGap = 10, minResidualNorm = 1e-10, log = None):
        self.CGCosts = []
        self.residualNorms = []
        self.maxIters = maxIters
        self.minIters = minIters
        self.tol = tol
        self.gapRatio = gapRatio
        self.minGap = minGap
        self.minResidualNorm = minResidualNorm
        self.log = log
        if self.log == None:
            self.log = open("/dev/null", 'w')
        
    def observeCGIterates(self, curCost, prevCost, beta, deltaNew, deltaOld, alpha, energy2Norm, r, b, x):
        if len(self.CGCosts) == 0:
            self.CGCosts.append(prevCost)
        self.CGCosts.append(curCost)
        self.residualNorms.append(norm(r))
    
    def doWeStop(self):
        if len(self.CGCosts) < 2: #doWeStop might be called before we have even done a single iteration
            return False
        i = len(self.CGCosts)-1 #current iteration, basically the number of times observeCGIterates has been called
        if i >= self.maxIters:
            print >>self.log, "CGStop: Max Iters"
            return True
        if self.residualNorms[-1] < self.minResidualNorm:
            print >>self.log, "CGStop: Small residual norm"
            return True
        testGap = max(self.minGap, int(num.ceil(i*self.gapRatio))) #how many iterations ago we look at CGCost
        if i < self.minIters or i < testGap:
            return False
        cost = self.CGCosts[-1]
        prevCost = self.CGCosts[-testGap]
        if prevCost < 0 and (cost-prevCost)/cost <= self.tol*testGap:
            print >>self.log, "CGStop: Rel. cost delta less than tolerance*testGap;",
            print >>self.log, "cost=%f, prevCost=%f, testGap=%f" % (cost, prevCost, testGap)
            return True
        return False
        
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):
    """
    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.
    """
    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 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)
        print i+1, curCost
        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):
    n = 1000
    A = genPosDefSym(n, seed = seed)
    b = num.random.randn(n,1)
    A=A.astype(num.float32)
    b=b.astype(num.float32)

    MInv = num.diag(1/A.diagonal())
    applyA = lambda v: num.dot(A,v)
    applyMInv = lambda v: num.dot(MInv,v)
    
    stoppingCriterion = CGStoppingCriterion(maxIters=1000, minIters=25, tol=5e-4, gapRatio=0.1, minGap = 10)
    cgIters = precondConjGrad(applyA, b, num.zeros((n,1)), applyMInv, stoppingCriterion)
    for j,x in enumerate(cgIters):
        print j+1, stoppingCriterion.CGCosts[-1]
        #if j == 5:
        #    kaw = x.copy()
    cgAns = x
    #print num.all(kaw == cgAns)
    print MInv.dtype
    ans = scipy.linalg.solve(A,b, sym_pos=True)
    print ans.shape
    print cgAns.shape

    #scipyAns, info = scipy.sparse.linalg.cg(A,b,num.zeros((n,1)), tol=0, maxiter=n, M = MInv)

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