###################################
# Copyright (c) <2021>, <J Lewis>
# All rights reserved.
#
# This source code is licensed under the GNU General public license found in the
# LICENSE file in the root directory of this source tree. 
#
# Author: Dr J Lewis
# Address: ScHARR, University of Sheffield
# Date: May 2021
# contact: jen.lewis@sheffield.ac.uk
###################################

# This function calculates sample size for a cRCT accounting for the 
# uncertainty in an ICC estimate calculated from a pilot trial.

# Returns estimates for the sample size for a main cRCT with a continuous outcome,
# equal sized clusters and two study arms. Estimates are calculated according to the 
# methods presented in Lewis & Julious.
# Results are returned using no adjustment, and accounting for uncertainty in the 
# ICC estimate using Swiger's variance, Searle's method and Fisher's normalising transformation. 
# 95% CI estimates are returned for the three methods of adjustment.

# LJ_cRCT_sampSize( rho, p_clust_size, p_n_clust, f_clust_size, eff_size, power = 0.9, alpha = 0.05 )

# Arguments:
# rho               estimate of the ICC
# p_clust_size      pilot trial cluster size
# p_n_clust         pilot trial number of clusters per arm
# f_clust_size      desired cluster size for main trial
# eff_size          standardised effect size for main trial
# power             desired power for main trial as proportion < 1 or percentage >=1. Default is 90% power
# alpha             Type 1 error rate as proportion < 1 or percentage >= 1. Defaults is 0.05


LJ_cRCT_sampSize <- function( rho, p_clust_size, p_n_clust, f_clust_size, eff_size, power = 0.9, alpha = 0.05 ){

    if( alpha >= 1 ){ alpha <- alpha/100 }
    z_alpha <- qnorm( 1-alpha/2)
    if( power >= 1 ){ power <- power/100 }
    z_beta <- qnorm( power )

    arms <- 2
    denominator <- eff_size^2
    
    step_size <- 0.001
    step_start <- 0.001
    step_end <- 0.999
    all_steps <- seq( step_start, step_end, by = step_size )
    
    # Subjects per arm (pilot)
    pN = p_n_clust*p_clust_size
    n_participants_pilot <- pN*arms
    
    
    ##########################
    # 0: No adjustment
    ##########################
    
    dEff <- max( 1, 1 + ( rho*( f_clust_size - 1 ) ))
    numerator <- 2*dEff*( z_alpha + z_beta )^2
    N_simple <- numerator/denominator
    nclust.unadjusted <- ceiling( sum( N_simple, na.rm = T  )/f_clust_size )
    n_participants.unadjusted <- nclust.unadjusted*f_clust_size
    
    
    ###########################
    # 1: Swiger's method
    ###########################
    
    pt.1 <- ( 1 + (( p_clust_size - 1 )* rho ))^2
    pt.2 <- ( 1 - rho )^2
    pt.3 <- 2*( ( pN*arms ) - 1 )
    pt.4 <- p_clust_size^2 * (( pN*arms ) - ( p_n_clust*arms )) * (( p_n_clust*arms ) - 1 )
    SwigerVar <- ( pt.3 * pt.2 * pt.1 ) / pt.4
    CI.swiger <- paste0( round( max( 0, rho-( sqrt( SwigerVar ) * 1.96 )), digits = 3 ),  ', ', 
                         round( min( 1, rho+( sqrt( SwigerVar ) * 1.96 )), digits = 3 ))
    
    # reset accumulators
    tot_N <- matrix( nrow = 1, ncol = length( all_steps ), data = NA )
    record_rho <- matrix( nrow = 1, ncol = length( all_steps ), data = NA )
    
    # integrate over rho for total N:
    for ( i in 1:( length( all_steps )-1 ) ){
        
        # calculate rho_i and rho_i+1
        # restrict so rho <= 1
        rho_i <- min( rho - ( qnorm( all_steps[ i ] ) * sqrt( SwigerVar )), 1 )
        rho_istep <- min( rho - ( qnorm( all_steps[ i + 1 ] ) * sqrt( SwigerVar )), 1 )
        
        # Design effect
        # restrict so dEff >=1
        dEff_a <- max( 1, 1 + ( rho_i*( f_clust_size - 1 )) )
        dEff_b <- max( 1, 1 + ( rho_istep*( f_clust_size - 1 )) )
        
        numer_a <- 2*dEff_a*( z_alpha + z_beta )^2
        numer_b <- 2*dEff_b*( z_alpha + z_beta )^2
        
        N_step <- 0.0005 * ( ( numer_a/denominator ) + ( numer_b/denominator ) )
        tot_N[ i ] <- N_step
        record_rho[ i ] <- rho_i
    }
    rm( i )
    
    # Store Swiger results
    nclust.swiger <- ceiling( sum( tot_N, na.rm = T  )/f_clust_size )
    n_participants.swiger <- nclust.swiger*f_clust_size
    
    
    #########################
    # 2: Searle's method
    #########################
    
    # First, find (n0-1), = cluster size-1, terminology = j
    j = p_clust_size-1
    
    # Calculate variance ratio from rho and j in absence of MSA and MSE
    f_num <- ( rho*j ) + 1
    f_denom <- 1 - rho
    var_ratio <- f_num/f_denom
    
    df1 <- ( p_n_clust*arms )-1
    df2 <- ( pN*arms )-1
    
    f_up <- qf( 0.975, df1, df2 )
    f_low <- qf( 0.025, df1, df2 )
    prop_up <- var_ratio/f_up
    prop_low <- var_ratio/f_low
    
    CI.searle <- paste0( round( max( 0, ( prop_up - 1 )/( p_clust_size + prop_up - 1 )), digits = 3 ), ', ',
                         round( min( 1, ( prop_low - 1 )/( p_clust_size + prop_low - 1 )), digits = 3 ))
    
    # reset accumulators etc
    tot_N <- matrix( nrow = 1, ncol = length( all_steps ), data = NA )
    record_rho <- matrix( nrow = 1, ncol = length( all_steps ), data = NA )
    
    # integrate over rho for total N:
    for ( i in 1:( length( all_steps )-1 ) ){
        
        # Get f quantile for step and step+1:
        f_quant <- qf( all_steps[ i ], df1, df2 )
        f_quant_step <- qf( all_steps[ i+1 ], df1, df2 )
        
        prop1 <- var_ratio/f_quant
        prop1_step <- var_ratio/f_quant_step
        
        # calculate rho_i and rho_i+1
        # restrict so rho <= 1
        rho_i <- min( ( prop1 - 1 )/( p_clust_size + prop1 - 1 ), 1 )
        rho_istep <- min( ( prop1_step - 1 )/( p_clust_size + prop1_step - 1 ), 1 )
        
        # Design effect
        # restrict so dEff >=1
        dEff_a <- max( 1, 1 + ( rho_i*( f_clust_size - 1 )) )
        dEff_b <- max( 1, 1 + ( rho_istep*( f_clust_size - 1 )) )
        
        numer_a <- 2*dEff_a*( z_alpha + z_beta )^2
        numer_b <- 2*dEff_b*( z_alpha + z_beta )^2
        
        N_step <- 0.0005 * ( ( numer_a/denominator ) + ( numer_b/denominator ) )
        tot_N[ i ] <- N_step
        record_rho[ i ] <- rho_i
    }
    rm( i )
    nclust.searle <- ceiling( sum( tot_N, na.rm = T  )/f_clust_size )
    n_participants.searle <- nclust.searle*f_clust_size
    
    
    ########################
    # 3: Fisher's method
    ########################
    
    j = p_clust_size - 1
    ZF <- 0.5*log(( 1 + ( j*rho ))/( 1 - rho))
    
    pt1 <- 1/( p_n_clust*arms - 1 )
    pt2 <- 1/( pN*arms - p_n_clust*arms )
    FisherVar <- 0.5*( pt1 + pt2 )
    
    Zlow <- ZF - ( 1.96 * sqrt( FisherVar ))
    Zup <- ZF + ( 1.96 * sqrt( FisherVar ))
    
    CI.fisher <- paste0( round( max( 0, ( exp( 2*Zlow ) - 1 )/( j + exp( 2*Zlow ))), digits = 3 ), ', ',
                         round( min( 1, ( exp( 2*Zup ) - 1 )/( j + exp( 2*Zup ))), digits = 3 ))
    
    # reset accumulators
    tot_N <- matrix( nrow = 1, ncol = length( all_steps ), data = NA )
    record_rho <- matrix( nrow = 1, ncol = length( all_steps ), data = NA )
    
    for( i in 1:( length( all_steps )-1 ) ){
        
        Zlim <- ZF - ( qnorm( all_steps[ i ] ) * sqrt( FisherVar ))
        lim <- ( exp( 2*Zlim ) - 1 )/( j + exp( 2* Zlim ))
        rho_i <- min( lim, 1 )
        
        Zlim_step <- ZF - ( qnorm( all_steps[ i+1 ] ) * sqrt( FisherVar ))
        lim_step <- ( exp( 2*Zlim_step ) - 1 )/( j + exp( 2* Zlim_step ))
        rho_istep <- min( lim_step, 1 )
        
        # Design effect
        # (restrict so dEff >=1)
        dEff_a <- max( 1, 1 + ( rho_i*( f_clust_size - 1 )) )
        dEff_b <- max( 1, 1 + ( rho_istep*( f_clust_size - 1 )) )
        
        numer_a <- 2*dEff_a*( z_alpha + z_beta )^2
        numer_b <- 2*dEff_b*( z_alpha + z_beta )^2
        
        N_step <- 0.0005 * ( ( numer_a/denominator ) + ( numer_b/denominator ) )
        tot_N[ i ] <- N_step
        record_rho[ i ] <- rho_i
    }
    rm( i )
    
    nclust.fisher <- ceiling( sum( tot_N, na.rm = T  )/f_clust_size )
    n_participants.fisher <- nclust.fisher*f_clust_size
    
    results <- data.frame( 'Unadjusted' = c( nclust.unadjusted, n_participants.unadjusted, NA ),
                           'Swiger' = c( nclust.swiger, n_participants.swiger, CI.swiger ), 
                           'Searle' = c( nclust.searle, n_participants.searle, CI.searle ), 
                           'Fisher' = c( nclust.fisher, n_participants.fisher, CI.fisher ))
    rownames( results ) <- c( 'Clusters per arm', 'Participants per arm', '95% CI' )
    return( results )

}