from __future__ import division
import numpy as np
from scipy.sparse.linalg import iterative  as iter
import pylab 
import pprint 
from scipy.optimize import nnls 

def PhiB(mux, muy, minVal, maxVal, x):
    phib = mux * np.abs( np.sum(np.log( x-minVal)) )
    #phib += np.abs( np.log(1. - maxVal / np.sum(x)) )
    return phib
    #phib += np.log std::log(maxVal - x.segment(ib*block, block).sum());
    

    #template < typename  Scalar >
    #Scalar PhiB2 (const Scalar& minVal, const Scalar& maxVal, const VectorXr x,
    #              const int& block, const int &nblocks) {
    #    Scalar phib  = std::abs((x.array() - minVal).log().sum());
    #    //phib      += std::abs((maxVal - x.array()).log().sum()*muy);
    #    for (int ib=0; ib<nblocks; ++ib) {
    #        //HiSegments(ib) = x.segment(ib*block, block).sum();
    #        phib += Scalar(block)*std::log(maxVal - x.segment(ib*block, block).sum());
    #    }
    #    return phib;
    #}

def logBarrier(A, b, T2Bins, x_0=0, xr=0, alpha=10, mu1=10, mu2=10, smooth=False, MAXITER=70, fignum=1000, sigma=1, callback=None):
    """Impliments a log barrier Tikhonov solution to a linear system of equations 
        Ax = b  s.t.  x_min < x < x_max. A log-barrier term is used for the constraint
    """
    # TODO input
    minVal = 0.0
    maxVal = 1e8

    Wd =    (np.eye(len(b)) / (sigma))        # Wd = eye( sigma )
    WdTWd = (np.eye(len(b)) / (sigma**2))     # Wd = eye( sigma )

    #print ("calculating AT WdT.Wd A ") # AT WdTWd A with known data matrix 
    #print (" A.shape" , np.shape(A), np.shape(b))
    ATWdTWdA = np.dot(A.conj().transpose(), np.dot( WdTWd, A ))     # TODO, implicit calculation instead?
    N = np.shape(A)[1]                        # number of model
    M = np.shape(A)[0]                        # number of data
    SIGMA = .25  #.25 #.125 # .25 #.01#3e-1
    EPSILON = 1e-35 # was 35 

    # reference model
    if np.size(xr) == 1:
        xr =  np.zeros(N)     
    
    # initial guess
    if np.size(x_0) == 1:
        x = 1e-10 + np.zeros(N)     
    else:
        x = 1e-10 + x_0
        
    # Construct model constraint base   
    #print ("constructing Phim")
    Phim_base = np.zeros( [N , N] ) 
    a1 = 1.0     # smallest
    
    # calculate largest term            
    D1 = 1./abs(T2Bins[1]-T2Bins[0])
    D2 = 1./abs(T2Bins[2]-T2Bins[1])
    #a2 = 1. #(1./(2.*D1+D2))    # smooth
    
    if smooth == "Both":
        #print ("SMOOTH")
        # Smooth model
        print ("Both small and smooth model")
        for ip in range(N):
            D1 = 0.
            D2 = 0.
            DMAX = 2./(T2Bins[1]-T2Bins[0])
            if ip > 0:
                D1 = 1./abs(T2Bins[ip]-T2Bins[ip-1])
            if ip < N-1:
                D2 = 1./abs(T2Bins[ip+1]-T2Bins[ip])
            if ip > 0:
                Phim_base[ip,ip-1] = -(D1)               # smooth in log space
            if ip == 0:
                Phim_base[ip,ip  ] =  (D1+D2) + a1       # Encourage a little low model, no a1
            elif ip == N-1:
                Phim_base[ip,ip  ] =  (D1+D2) + a1       # Penalize long decays
            else:
                Phim_base[ip,ip  ] =  (D1+D2) + a1       # Smooth and small
            if ip < N-1:
                Phim_base[ip,ip+1] = -(D2)               # smooth in log space

        #Phim_base /= np.max(Phim_base)
        #Phim_base += a1*np.eye(N)

    elif smooth == "Smooth":
        print ("Smooth model")
        for ip in range(N):
            if ip > 0:
                Phim_base[ip,ip-1] = -1    # smooth in log space
            if ip == 0:
                Phim_base[ip,ip  ] = 1.0   # Encourage a little low model
            elif ip == N-1:
                Phim_base[ip,ip  ] = 8.0   # Penalize long decays
            else:
                Phim_base[ip,ip  ] = 2.0   # Smooth and small
            if ip < N-1:
                Phim_base[ip,ip+1] = -1    # smooth in log space
        #print(Phim_base)    

    else: 
        print ("SMALLEST")
        # Smallest model
        for ip in range(N):
            Phim_base[ip,ip  ] = 1.
    
    Phi_m =  alpha*Phim_base
    WmTWm = Phim_base # np.dot(Phim_base, Phim_base.T)            
    b_pre = np.dot(A, x)

    phid = np.linalg.norm( np.dot(Wd, (b-b_pre)) )**2
    phim = np.linalg.norm( np.dot(Phim_base, (x-xr)) )**2

    mu2 = phim
    phib = PhiB(mu1, mu2, 0, 1e8, x) 
    mu1 = ((phid + alpha*phim) / phib)

    #print ("iteration", -1, mu1, mu2, phib, phid, phim, len(b))
    for i in range(MAXITER):
            
        b_pre = np.dot(A, x)
        phid = np.linalg.norm(np.dot(Wd, (b-b_pre)))**2
        # technically phim should not be on WmTWm matrix. So no need to square result. 
        phim = np.linalg.norm(np.dot(Phim_base, (x-xr)))#**2
        phib = PhiB(mu1, mu2, 0, 1e8, x) 
        Phi_m =  alpha*Phim_base
        
        mu1 = ((phid + alpha*phim) / phib) 

        WmTWm = Phim_base # np.dot(Phim_base, Phim_base.T)            
        #ztilde = x
        #print("ztilde", ztilde)
        phid_old = phid
        inner = 0
        First = True

        while ( (phib / (phid+alpha*phim)) > EPSILON  or First==True ):

            First = False
            # Log barrier, keep each element above minVal
            X1 = np.eye(N) * (x-minVal)**-1           
            X2 = np.eye(N) * (x-minVal)**-2         
            
            # Log barrier, keep sum below maxVal TODO normalize by component. Don't want to push all down  
            Y1 = np.eye(N) * (maxVal - np.sum(x))**-1           
            Y2 = np.eye(N) * (maxVal - np.sum(x))**-2         
            
            AA = ATWdTWdA + mu1*X2 + mu2*Y2 + Phi_m 
            M = np.eye( N ) * (1./np.diag(ATWdTWdA + mu1*X2 + mu2*Y2 + Phi_m))
        
            # Solve system (newton step) 
            b2 = np.dot(A.transpose(), np.dot(WdTWd, b-b_pre) ) + 2.*mu1*np.diag(X1) + 2.*mu2*np.diag(Y1) - alpha*np.dot(WmTWm,(x-xr))
            ztilde = iter.cg(AA, b2, M=M) # tol=1e-3*phid, maxiter=200, callback=callback) #, tol=1e-2, maxiter) #, x0=x) #, tol=ttol) #, M=M, x0=x)        
            h = (ztilde[0]) 
            
            # Solve system (direct solution) 
            #b2 = np.dot(A.conj().transpose(), np.dot(WdTWd, b)) + 2.*mu1*np.diag(X1) + 2.*mu2*np.diag(Y1) - alpha*np.dot(WmTWm,(x-xr))
            #ztilde = iter.cg(AA, b2, x0=x) #, tol=1e-2) #, x0=x) #, tol=ttol) #, M=M, x0=x)        
            #h = (ztilde[0].real - x) 

            # step size
            d = np.min( (1, 0.95 * np.min(x/np.abs(h+1e-120))) )
            
            ##########################################################
            # Update and fix any over/under stepping
            x = x+d*h 
            #x = np.max( ( (minVal+1e-120)*np.ones(N),  x+d*h), 0)
        
            # Determine mu steps to take
            s1 = mu1 * (np.dot(X2, ztilde[0].real) - 2.*np.diag(X1))
            s2 = mu2 * (np.dot(Y2, ztilde[0].real) - 2.*np.diag(Y1))

            # determine mu for next step
            mu1 = SIGMA/N * np.abs(np.dot(s1, x))
            mu2 = SIGMA/N * np.abs(np.dot(s2, x))
            
            b_pre = np.dot(A, x)
            phid = np.linalg.norm(np.dot(Wd, (b-b_pre)))**2
            phim = np.linalg.norm(np.dot(Phim_base, (x-xr)))#**2
            phib = PhiB(mu1, mu2, minVal, maxVal, x)
            inner += 1
        
        # determine alpha
        scale = (len(b)/phid)
        alpha *= scale  #**(1/6) 

        score = np.sqrt(phid/(len(b)+1.)) # unbiased  
 
        # check stopping criteria 
        if score < 1: 
            print ("*overshot* optimal solution found") #, alpha, score)
            break
        if score < 1.1: # or np.linalg.norm(x_old-x) < 1e-5 or phid > phid_old:
            #print ("overshoot")
            #alpha *= 10
            print ("optimal solution found") #, alpha, score)
            break
        if i > 10 and (np.sqrt(phid_old/(len(b)+1.)) - score) < 1e-2: # 1e-2
            print ("slow convergence") #, alpha, score, i, scale, score-np.sqrt(phid_old/len(b)))
            break
            
    print ( "alpha","phid", "iter", "search", "prior" )
    print ( alpha, score, i, scale, np.sqrt(phid_old/(len(b)+1)))

    return x