#
# Written by:
# -- 
# John L. Weatherwax                2009-04-21
# 
# email: wax@alum.mit.edu
# 
# Please send comments and especially bug reports to the
# above email address.
#
# Utility function for Chapter 16
#
#-----

get_a_sub_n = function(n){
    if( (n<3) | (n>8) ){
        return(NA)
    }else{
        ans = c( .886, .921, .940, .952, .959, .965 ) # from n = 3 - 8
        ans[n-2]
    }
}

get_b_sub_n = function(n){
    if( (n<3) | (n>8) ){
        return(NA)
    }else{
        bns = c( 1.693, 2.058, 2.325, 2.536, 2.706, 2.844 ) # from n = 3 - 8
        bns[n-2]
    }
}

get_c_sub_n = function(n){
    if( (n<3) | (n>8) ){
        return(NA)
    }else{
        cns = c( .888, .880, .864, .848, .833, .820 ) # from n = 3 - 8
        cns[n-2]
    }
}

get_k_sub_n = function(n){ 
    if( (n<4) | (n>8) ){
        return(NA)
    }else{
        kns = c( .596, .990, 1.282, 1.512, .942 ) # from n = 4 - 8
        kns[n-3]
    }
}

control_limit_utils = function(data, group){
    #
    # A utility function that computes quantities needed for producing control charts.
    #
    ugroup = unique( group )
    k = length( ugroup )

    # all groups must have the same sumber of samples (as the first one):
    n = sum( group==ugroup[1] )

    x_i_dot = rep( NA, k ) # the sum of the samples in the group
    xbar_i_dot = rep( NA, k ) # the mean for the samples in the group
    ss = rep( NA, k ) # the standard deviation of the samples in group
    srange = rep( NA, k ) # the range of the samples in group
    iqr = rep( NA, k ) # the interquartile range (IQR)
    for( gi in 1:k ){
        group_data = data[ group==ugroup[gi] ]
        stopifnot( length(group_data)==n )
        x_i_dot[gi] = sum( group_data )
        xbar_i_dot[gi] = x_i_dot[gi] / n # compute the group mean
        ss[gi] = sd(group_data)
        srange[gi] = diff( range(group_data) )
        iqr[gi] = IQR( group_data )
    }

    # Estimate the control limits:
    #
    x_bar_bar = mean( xbar_i_dot )
    sbar = mean( ss )
    rbar = mean( srange )
    sigma_est_from_std = mean( ss ) / get_a_sub_n(n)
    sigma_est_from_range = mean( srange ) / get_b_sub_n(n)
    sigma_est_from_IQR = mean( iqr ) / get_k_sub_n(n)

    an = get_a_sub_n( n )
    delta_s = 3 * sbar * ( sqrt( 1 - an^2 )/an )
    LCL = max( 0, sbar - delta_s )
    UCL = sbar + delta_s
    s_control_limits = c( LCL, UCL )

    rbar = mean( srange )
    bn = get_b_sub_n( n )
    cn = get_c_sub_n( n )
    delta_r = 3 * rbar * ( cn/bn )
    LCL = max( 0, rbar - delta_r )
    UCL = rbar + delta_r
    r_control_limits = c( LCL, UCL )
    
    # return everything as a list:
    list(
        k=k, n=n, # variables names used in this chapter
        I=k, J=n, # variable names that should be deprecated
        x_i_dot=x_i_dot, xbar_i_dot=xbar_i_dot, ss=ss, srange=srange,
        x_bar_bar=x_bar_bar,
        sigma_est_from_std=sigma_est_from_std,
        sigma_est_from_range=sigma_est_from_range,
        sigma_est_from_IQR=sigma_est_from_IQR,
        r_control_limits=r_control_limits,
        s_control_limits=s_control_limits
    )
}

CUMSUM = function( data, mu_0, k ){
    r = length(data)
    ds = rep( NA, r+1 ) # holds c( d_0, d_l, d_2, ... d_r )
    ds[1] = 0
    es = rep( NA, r+1 ) # holds c( e_0, e_1, e_2, ... e_r )
    es[1] = 0
    for( l in 1:r ){ # want to fill d_1, d_2, d_3, .. d_r
        ds[l+1] = max( 0, ds[l] + ( data[l] - (mu_0 + k) ) )
        es[l+1] = max( 0, es[l] - ( data[l] - (mu_0 - k) ) )
    }
    list( ds=ds[2:length(ds)], es=es[2:length(es)] ) # drop the first index (corresponds to 0)
}