import numpy as np
import scipy.stats as ss
import random


#% Adaptive MCMC, using Haario et al:
#% https://link.springer.com/article/10.1007/s11222-008-9110-y
#% and
#% http://probability.ca/jeff/ftpdir/adaptex.pdf
## *********************************

#% F:          Function giving log-posterior density for a parameter set x
#% x0:         Initial value of parameter set x
#% n:          Number of iterations
#% cov0:       Initial covariance matrix
#% fac:        Scaling factor for covariance matrix. Set fac = 1 for default
#% fixinds:    Elements of x that should be held fixed. Set fixinds = [] for full MCMC
#% blockinds:  Number of 'epi parameters' (e.g. beta, X2) in x, if we want to vary epi and non-epi parameters as independent 'blocks'. Set blockinds = [] if we want a full covariance matrix
#% displ:      Structure with display options. Set displ = true to show progress

def MCMC_adaptive(F, x0, n, sigma, fixinds, blockind, cov0, displ):
    xsto=None; outsto=None; history=None; accept_rate=None
    d = len(x0)
    b = 0.05
    sd = sigma*2.4*2.4/d
    if len(fixinds)!=0:
        inds=fixinds[0,:]
        vals = fixinds[1,:]
    else:
        inds=[]
        vals=[]

    #% Checks on the initial covariance matrix
    if len(cov0) ==0:
        cov0 = np.eye(d)
        cov0[inds,:] = 0
        cov0[:,inds] = 0
        if blockind!=[]:
            cov0[0:blockind,blockind:d] = 0
            cov0[blockind:d,0:blockind] = 0
    #end

    #% Initiate the output matrices
    xsto = np.zeros((d,n)); outsto = np.zeros(n)#xsto = zeros(d,n); outsto = zeros(1,n);
    history = np.zeros((d+1,n))#history = zeros(d+1,n);                                                    % Rows: 1:d Proposed values 4. Accept or reject

    xsto[:,0] = x0[:]; xbar = xsto#xsto(:,1) = x0(:); xbar = xsto;
    FX = F(x0); outsto[1] = FX#FX = F(x0); outsto(1) = FX;
    acc = 0

    current_mean = x0
    current_cov = cov0

    #if displ; figure; end
    # % --- Start the MCMC loop -------------------------------------------------
    for t in range(1,n-1):
        X = xsto[:,t-1]
        #% --- Make a proposal from the distribution
        Y0 =  ss.multivariate_normal(mean=X, cov=0.1*0.1*cov0*sigma/d).rvs()#Y0 = mvnrnd(X,0.1^2*cov0*sigma/d) #??
        if t < 2*d:
           Y = np.maximum(Y0,0); Y[inds] = vals#Y = max(Y0,0); Y(inds) = vals;
        else:
           ind0 = 0; ind1 = t-1
           covmat = np.cov(xsto[:,ind0:ind1].transpose()) #covmat = cov(xsto(:,ind0:ind1)');
                        
           covmat[inds,:] = 0; covmat[:,inds] = 0 #covmat(inds,:) = 0; covmat(:,inds) = 0
           if blockind!=[]:
            covmat[0:(blockind-1),blockind:] = 0; covmat[blockind:,0:(blockind-1)] = 0#covmat(1:blockind,(blockind+1:end)) = 0; covmat((blockind+1):end,1:blockind) = 0;
           #% Need to modify this bit so it goes recursively - faster
        
           covmat = (covmat+covmat.transpose())/2+1e-5#covmat = (covmat+covmat')/2+1e-5;
           #% Precaution to make sure values don't get negative
           Y = np.maximum((1-b)*ss.multivariate_normal(mean=X,cov=sd*covmat) + b*Y0,0) #Y = max((1-b)*mvnrnd(X,sd*covmat) + b*Y0,0);
           Y[inds] = vals#Y(inds) = vals;
        #end

        history[0:d,t] = Y#history(1:d,t) = Y;
        #% --- Decide whether to accept or not
        FY = F(Y)
        if np.random.uniform() < np.exp(FY-FX):
           #% Accept
           xsel = Y[:]#xsel = Y(:);
           FX = FY#;
           acc = acc+1
           history[d,t] = 1#history(end,t) = 1;
        else:
           #% Reject
           xsel = xsto[:,t-1]#xsel = xsto(:,t-1);
        #end
        xsto[:,t] = xsel#;xsto(:,t) = xsel;
        outsto[t] = FX#;outsto(t) = FX;
        xbar[:,t] = (xbar[:,t-1]*(t-1) + xsel)/t#xbar(:,t) = (xbar(:,t-1)*(t-1) + xsel)/t;
        #% Display options
        #if displ && (mod(t,round(n/25))==0); fprintf('%0.5g ', t/n*25); end
        #if displ && (mod(t,200)==0)
        #     plot(xsto(:,1:t-1)'); xlim([0 n]); drawnow;
        #end
    accept_rate = acc/n#;
    xsto = xsto.transpose()#xsto = xsto';
    history = history.transpose()#history = history';    

    return (xsto,outsto,history,accept_rate)


def MCMC_adaptive2(F, x0, n, sigma, fixinds, blockind, cov0, displ):
    xsto=None; outsto=None; history=None; accept_rate=None
    d = len(x0)
    b = 0.05
    sd = sigma*2.4*2.4/d
    if len(fixinds)!=0:
        inds=fixinds[0,:]
        vals = fixinds[1,:]
    else:
        inds=[]
        vals=[]

    #% Checks on the initial covariance matrix
    if len(cov0) ==0:
        cov0 = np.eye(d)
        cov0[inds,:] = 0
        cov0[:,inds] = 0
        if blockind!=[]:
            cov0[0:blockind,blockind:d] = 0
            cov0[blockind:d,0:blockind] = 0
    #end

    #% Initiate the output matrices
    xsto = np.zeros(d,n); outsto = np.zeros(1,n)#xsto = zeros(d,n); outsto = zeros(1,n);
    history = np.zeros(d+1,n)#history = zeros(d+1,n);                                                    % Rows: 1:d Proposed values 4. Accept or reject

    xsto[:,0] = x0[:] #xbar = xsto#xsto(:,1) = x0(:); xbar = xsto;
    FX = F(x0); outsto[1] = FX#FX = F(x0); outsto(1) = FX;
    acc = 0

    current_mean = x0
    current_cov = cov0
    #if displ; figure; end
    # % --- Start the MCMC loop -------------------------------------------------
    for t in range(1,n-1):
        X = xsto[:,t-1]
        #% --- Make a proposal from the distribution
        Y0 =  ss.multivariate_normal(mean=X, cov=0.1*0.1*cov0*sigma/d).rvs()#Y0 = mvnrnd(X,0.1^2*cov0*sigma/d) #??
        if t < 2*d:
           Y = np.maximum(Y0,0); Y[inds] = vals#Y = max(Y0,0); Y(inds) = vals;
        else:
           ind0 = 0; ind1 = t-1
           covmat = np.cov(xsto[:,ind0:ind1].transpose()) #covmat = cov(xsto(:,ind0:ind1)');
                        
           covmat[inds,:] = 0; covmat[:,inds] = 0 #covmat(inds,:) = 0; covmat(:,inds) = 0
           if blockind!=[]:
            covmat[0:(blockind-1),blockind:d] = 0; covmat[blockind:d,0:(blockind-1)] = 0#covmat(1:blockind,(blockind+1:end)) = 0; covmat((blockind+1):end,1:blockind) = 0;
           #% Need to modify this bit so it goes recursively - faster
        
           covmat = (covmat+covmat.transpose())/2+1e-5#covmat = (covmat+covmat')/2+1e-5;
           #% Precaution to make sure values don't get negative
           Y = np.maximum((1-b)*ss.multivariate_normal(mean=X,cov=sd*covmat) + b*Y0,0) #Y = max((1-b)*mvnrnd(X,sd*covmat) + b*Y0,0);
           Y[inds] = vals#Y(inds) = vals;
        #end

        history[0:d,t] = Y#history(1:d,t) = Y;
        #% --- Decide whether to accept or not
        FY = F(Y)
        if random.uniform() < np.exp(FY-FX):
           #% Accept
           xsel = Y[:]#xsel = Y(:);
           FX = FY#;
           acc = acc+1
           history[d,t] = 1#history(end,t) = 1;
        else:
           #% Reject
           xsel = xsto[:,t-1]#xsel = xsto(:,t-1);
        #end
        xsto[:,t] = xsel#;xsto(:,t) = xsel;
        outsto[t] = FX#;outsto(t) = FX;
        #xbar[:,t] = (xbar[:,t-1]*(t-1) + xsel)/t;#xbar(:,t) = (xbar(:,t-1)*(t-1) + xsel)/t;
        #delta_bf = xsel - current_mean
        #current_mean = (current_mean*(t-1)+xsel)/t
        #delta_af = xsel-current_mean
        #current_cov = current_cov + np.outer(delta_bf,delta_af)
        w0 = 1/(t+1); w1 = t/(t+1); w2= w1/(t+1)
        current_mean = w1*current_mean + w0*xsel
        delta_af = xsel-current_mean
        current_cov = w1*current_cov + w2*np.outer(delta_af,delta_af)
        #% Display options
        #if displ && (mod(t,round(n/25))==0); fprintf('%0.5g ', t/n*25); end
        #if displ && (mod(t,200)==0)
        #     plot(xsto(:,1:t-1)'); xlim([0 n]); drawnow;
        #end
    accept_rate = acc/n#;
    xsto = xsto.transpose()#xsto = xsto';
    history = history.transpose()#history = history';    

    return (xsto,outsto,history,accept_rate)