source("HPE_functions.R")

# All routines written by Pantelis Z. Hadjipantelis and David Springate (30-Jul-2012) while still experimental this code is due to be appear under some GPL-related license

#HyperParameter Estimation
HPE <- function( tree, val, l = NULL, sf=NULL, sn=NULL, SNR= NULL){
	#tree: object of class ‘"phylo" as defined in ape(), see Paradis, E. (2012)
	#val : values of traits at the tips
	#l   : characteristic length scale
	#sf  : phylogeny related variational amplitude
	#sn  : non-phylogeny related  variational amplitude
	#SNR : sf^2/sn^2
	#Check tree and values are valid
	if ( dim(cophenetic(tree))[1] != length(val) ) { print ("Not all points at the tips are known / Insufficient number of known points. Aborting"); return();}

	z <-4; #Number of random initialization

	#if length scale is unknown 
	if ( is.null(l)){
		#and both sf and sn are known
		if( !(is.null(sf)) && !(is.null(sn)) ){
			#provide estimate for l
			return ( exp( GetHyperparameters_L (SN=sn, SF=sf,Y=val,REPS=z, derivative.function= fLogLik_OU1_L_Only_dF, likelihood.function= fLogLik_OU1_L_Only_F, X=cophenetic(tree))) );
			}
		#or say what went wrong
		else { print('Having no specific length-scale you need to specify both sf and sn. Aborting.'); return(NULL);}
		}
	#otherwise if length scale is known
	else{
		#and you know sf but not sn
		if ( !(is.null(sf)) && (is.null(sn)) ) {
			#provide estimate for sn
			return( exp( GetHyperparameters_SN(L=l, SF=sf, Y=val, REPS=z, derivative.function=fLogLik_OU1_SN_Only_dF, likelihood.function=fLogLik_OU1_SN_Only_F, X=cophenetic(tree))) );
		}
		#or if you know sn but not sf
		else if ( !(is.null(sn)) && (is.null(sf))  ) {
			#provide estimate for sf
			return(	exp( GetHyperparameters_SF(SN=sn, L=l, Y=val, REPS=z, derivative.function=fLogLik_OU1_SF_Only_dF, likelihood.function=fLogLik_OU1_SF_Only_F, X=cophenetic(tree))) );
		}
		#or if both sn and sf are unknown but you know the SNR (sf^2/sn^2)
		else if ((is.null(sn)) && (is.null(sf)) && !(is.null(SNR))  ){
			#provide estimate for sf (so user can later estimate sn)
			return(	exp( GetHyperparameters_knownL_and_SNR(SNR=SNR, L=l, Y=val, REPS=z, derivative.function=fLogLik_OU2_givenL_and_SNR_Only_dF, likelihood.function=fLogLik_OU2_givenL_and_SNR_Only_F,  X=cophenetic(tree))) );
		}
		#or if you just know l and nothing else, say what went wrong
		else{ print('Having a specific length-scale but no other information inference is not (currently) possible. Aborting.'); return(NULL); }
	}

}

#================================================================================================

GetHyperparameters_knownL_and_SNR <- function( Y, X, REPS, likelihood.function,derivative.function,  L,SNR){ 
    # Y : Trait values at tips
    # X : Distance matrix
    # REPS : number of random initializations.
    # returns vector with estimate of s_f
    num.thetas=1;
    Theta <- rnorm(num.thetas)+rnorm(1)   #Random Initialization
    #Next line: Optimize Log_Likelihood towards Theta given Y,X using BFGS
    W = optim(par = Theta, fn = likelihood.function, gr = derivative.function,method="BFGS", Y=Y, X=X,L=log(L),SNR=SNR )   
    Cost = W$value; Theta = W$par       #Optimal Theta and Cost at that Theta   
    for (i in 1:REPS){
	#print(i)
        New_Theta <- rnorm(num.thetas)+rnorm(1)    #New Random Initialization
        #Next line: Optimize Log_Likelihood towards Theta given Y,X using BFGS algorithm
        OptimizationObj = optim(par = New_Theta, fn =  likelihood.function, gr = derivative.function, method="BFGS", Y=Y, X=X,L=log(L),SNR=SNR ) 
        if (Cost >  OptimizationObj$value){ #If new solution is better
            Cost = OptimizationObj$value;   #Save new best logLikelihood
            if (Cost <0) {          #Negative variance / Burst in tears now.
                #print('Big problem amigo')
                Cost = 12^12;
            }
            Theta = OptimizationObj$par;     #Save new best Theta parameters
        }
    }
    return(Theta)
}

#================================================================================================

GetHyperparameters_SF <- function( Y, X, REPS, likelihood.function, 
                                    derivative.function,L,SN){ 
    # Y : Trait values at tips
    # X : Distance matrix
    # REPS : number of random initializations.
    # returns vector with estimate of s_f
    num.thetas=1;
    Theta <- rnorm(num.thetas)+rnorm(1)  #Random Initialization
    #Next line: Optimize Log_Likelihood towards Theta given Y,X using BFGS
    W = optim(par = Theta, fn = likelihood.function, gr = derivative.function,
                                                        method="BFGS", Y=Y, X=X,L=log(L),SN=log(SN) )   
    Cost = W$value; Theta = W$par       #Optimal Theta and Cost at that Theta   
    for (i in 1:REPS){
	#print(i)
        New_Theta <- rnorm(num.thetas)+rnorm(1)    #New Random Initialization
        #Next line: Optimize Log_Likelihood towards Theta given Y,X using BFGS algorithm
        OptimizationObj = optim(par = New_Theta, fn =  likelihood.function,
                                gr = derivative.function, method="BFGS", Y=Y, X=X,L=log(L),SN=log(SN) ) 
        if (Cost >  OptimizationObj$value){ #If new solution is better
            Cost = OptimizationObj$value;   #Save new best logLikelihood
            if (Cost <0) {          #Negative variance / Burst in tears now.
                #print('Big problem amigo')
                Cost = 12^12;
            }
            Theta = OptimizationObj$par;     #Save new best Theta parameters
        }
    }
    return(Theta)
}

#================================================================================================

GetHyperparameters_L <- function( Y, X, REPS, likelihood.function,derivative.function, SF,SN){ 
    # Y : Trait values at tips
    # X : Distance matrix
    # REPS : number of random initializations.
    # returns vector with  estimate of l
    num.thetas=1;
    Theta <- rnorm(num.thetas)+rnorm(1)   #Random Initialization
    #Next line: Optimize Log_Likelihood towards Theta given Y,X using BFGS
    W = optim(par = Theta, fn = likelihood.function, gr = derivative.function,
                                                        method="BFGS", Y=Y, X=X,SF=log(SF),SN=log(SN) )   
    Cost = W$value; Theta = W$par       #Optimal Theta and Cost at that Theta   
    for (i in 1:REPS){
	#print(i)
        New_Theta <- rnorm(num.thetas)+rnorm(1)    #New Random Initialization
        #Next line: Optimize Log_Likelihood towards Theta given Y,X using BFGS algorithm
        OptimizationObj = optim(par = New_Theta, fn =  likelihood.function,
                                gr = derivative.function, method="BFGS", Y=Y, X=X,SF=log(SF),SN=log(SN) ) 
        if (Cost >  OptimizationObj$value){ #If new solution is better
            Cost = OptimizationObj$value;   #Save new best logLikelihood
            if (Cost <0) {          #Negative variance / Burst in tears now.
                #print('Big problem amigo')
                Cost = 12^12;
            }
            Theta = OptimizationObj$par;     #Save new best Theta parameters
        }
    }
    return(Theta)
}

#================================================================================================

GetHyperparameters_SN <- function( Y, X, REPS, likelihood.function,derivative.function,  SF,L){ 
    # Y : Trait values at tips
    # X : Distance matrix
    # REPS : number of random initializations.
    # returns vector with estimate of s_n
    num.thetas=1;
    Theta <- rnorm(num.thetas)+rnorm(1)   #Random Initialization
    #Next line: Optimize Log_Likelihood towards Theta given Y,X using BFGS
    W = optim(par = Theta, fn = likelihood.function, gr = derivative.function,
                                                        method="BFGS", Y=Y, X=X,SF=log(SF),L=log(L) )   
    Cost = W$value; Theta = W$par       #Optimal Theta and Cost at that Theta   
    for (i in 1:REPS){
	#print(i)
        New_Theta <- rnorm(num.thetas)+rnorm(1)    #New Random Initialization
        #Next line: Optimize Log_Likelihood towards Theta given Y,X using BFGS algorithm
        OptimizationObj = optim(par = New_Theta, fn =  likelihood.function,
                                gr = derivative.function, method="BFGS", Y=Y, X=X,SF=log(SF),L=log(L) ) 
        if (Cost >  OptimizationObj$value){ #If new solution is better
            Cost = OptimizationObj$value;   #Save new best logLikelihood
            if (Cost <0) {          #Negative variance / Burst in tears now.
                #print('Big problem amigo')
                Cost = 12^12;
            }
            Theta = OptimizationObj$par;     #Save new best Theta parameters
        }
    }
    return(Theta)
}

#================================================================================================