# Added by John O. Woods, SciRuby project.
# Derived from GSL-1.9 source files in the specfunc/ dir.

module Distribution
  module MathExtension
    module Beta
      class << self
        # Based on gsl_sf_lnbeta_e and gsl_sf_lnbeta_sgn_e
        # Returns result and sign in an array. If with_error is specified, also returns the error.
        def log_beta(x,y, with_error=false)
          sign = nil

          raise(ArgumentError, "x and y must be nonzero") if x == 0.0 || y == 0.0
          raise(ArgumentError, "not defined for negative integers") if [x,y].any? { |v| (v.is_a?(Fixnum) && v < 0) }

          # See if we can handle the positive case with min/max < 0.2
          if x > 0 && y > 0
            min, max = [x,y].minmax
            ratio    = min.quo(max)

            if ratio < 0.2
              gsx   = Gammastar.evaluate(x, with_error)
              gsy   = Gammastar.evaluate(y, with_error)
              gsxy  = Gammastar.evaluate(x+y, with_error)
              lnopr = Log::log_1plusx(ratio, with_error)

              gsx, gsx_err, gsy, gsy_err, gsxy, gsxy_err, lnopr, lnopr_err = [gsx,gsy,gsxy,lnopr].flatten if with_error

              lnpre = Math.log((gsx*gsy).quo(gsxy) * Math::SQRT2 * Math::SQRTPI)
              lnpre_err = gsx_err.quo(gsx) + gsy_err(gsy) + gsxy_err.quo(gsxy) if with_error

              t1    = min * Math.log(ratio)
              t2    = 0.5 * Math.log(min)
              t3    = (x+y-0.5)*lnopr

              lnpow       = t1 - t2 - t3
              lnpow_err   = Float::EPSILON * (t1.abs + t2.abs + t3.abs) + (x+y-0.5).abs * lnopr_err if with_error

              result      = lnpre + lnpow
              error       = lnpre_err + lnpow_err + 2.0*Float::EPSILON*result.abs if with_error

              return with_error ? [result, 1.0, error] : [result, 1.0]
            end
          end

          # General case: fallback
          lgx, sgx   = Math.lgamma(x)
          lgy, sgy   = Math.lgamma(y)
          lgxy, sgxy = Math.lgamma(x+y)
          sgn        = sgx * sgy * sgxy

          raise("Domain error: sign is -") if sgn == -1

          result = lgx + lgy - lgxy
          if with_error
            lgx_err, lgy_err, lgxy_err = begin
              STDERR.puts("Warning: Error is unknown for Math::lgamma, guessing.")
              [Math::EPSILON, Math::EPSILON, Math::EPSILON]
            end

            error  = lgx_err + lgy_err + lgxy_err + Float::EPSILON*(lgx.abs+lgy.abs+lgxy.abs) + 2.0*(Float::EPSILON)*result.abs
            return [result, sgn, error]
          else
            return [result, sgn]
          end

        end
      end
    end
    # Calculate regularized incomplete beta function
    module IncompleteBeta

      MAX_ITER = 512
      CUTOFF   = 2.0 * Float::MIN

      class << self

        # Evaluate aa * beta_inc(a,b,x) + yy
        #
        # No error mode available.
        #
        # From GSL-1.9: cdf/beta_inc.c, beta_inc_AXPY
        def axpy(aa,yy,a,b,x)
          return aa*0 + yy if x == 0.0
          return aa*1 + yy if x == 1.0

          ln_beta   = Math.logbeta(a, b)
          ln_pre    = -ln_beta   +   a * Math.log(x)   +   b * Math::Log.log1p(-x)
          prefactor = Math.exp(ln_pre)

          if x < (a+1).quo(a+b+2)
            # Apply continued fraction directly
            epsabs  = yy.quo((aa * prefactor).quo(a)).abs * Float::EPSILON
            cf      = continued_fraction(a, b, x, epsabs)
            return  aa * (prefactor * cf).quo(a) + yy
          else
            # Apply continued fraction after hypergeometric transformation
            epsabs = (aa + yy).quo( (aa*prefactor).quo(b) ) * Float::EPSILON
            cf     = continued_fraction(b, a, 1-x, epsabs)
            term   = (prefactor * cf).quo(b)
            return aa == -yy ? -aa*term : aa*(1-term)+yy
          end
        end


        # Evaluate the incomplete beta function
        # gsl_sf_beta_inc_e
        def evaluate(a,b,x,with_error=false)
          raise(ArgumentError, "Domain error: a(#{a}), b(#{b}) must be positive; x(#{x}) must be between 0 and 1, inclusive") if a <= 0 || b <= 0 || x < 0 || x > 1
          if x == 0
            return with_error ? [0.0,0.0] : 0.0
          elsif x == 1
            return with_error ? [1.0,0.0] : 1.0
          else

            ln_beta = Beta.log_beta(a,b, with_error)
            ln_1mx  = Log.log_1plusx(-x, with_error)
            ln_x    = Math.log(x)

            ln_beta, ln_beta_err, ln_1mx, ln_1mx_err, ln_x_err = begin
              #STDERR.puts("Warning: Error is unknown for Math::log, guessing.")
              [ln_beta,ln_1mx,Float::EPSILON].flatten
            end

            ln_pre      = -ln_beta + a*ln_x + b*ln_1mx
            ln_pre_err  = ln_beta_err + (a*ln_x_err).abs + (b*ln_1mx_err).abs if with_error

            prefactor, prefactor_err   = begin
              if with_error
                exp_err(ln_pre, ln_pre_err)
              else
                [Math.exp(ln_pre), nil]
              end
            end

            if x < (a+1).quo(a+b+2)
              # Apply continued fraction directly
              
              cf      = continued_fraction(a,b,x, nil, with_error)
              cf,cf_err = cf if with_error
              result  = (prefactor * cf).quo(a)
              return with_error ? [result, ((prefactor_err*cf).abs + (prefactor*cf_err).abs).quo(a)] : result
            else
              # Apply continued fraction after hypergeometric transformation

              cf      = continued_fraction(b, a, 1-x, nil)
              cf,cf_err = cf if with_error
              term    = (prefactor * cf).quo(b)
              result  = 1 - term

              return with_error ? [result, (prefactor_err*cf).quo(b) + (prefactor*cf_err).quo(b) + 2.0*Float::EPSILON*(1+term.abs)] : result
            end

          end
        end


        def continued_fraction_cutoff(epsabs)
          return CUTOFF if epsabs.nil?
          0.0/0 # NaN
        end

        # Continued fraction calculation of incomplete beta
        # beta_cont_frac from GSL-1.9
        #
        # If epsabs is set, will execute the version of the GSL function in the cdf folder. Otherwise, does the
        # basic one in specfunc.
        def continued_fraction(a,b,x,epsabs=nil,with_error=false)
          num_term = 1
          den_term = 1 - (a+b)*x.quo(a+1)
          k        = 0

          den_term = continued_fraction_cutoff(epsabs) if den_term.abs < CUTOFF
          den_term = 1.quo(den_term)
          cf       = den_term

          1.upto(MAX_ITER) do |k|
            coeff      = k *(b-k)*x.quo(((a-1)+2*k)*(a+2*k)) # coefficient for step 1
            delta_frac = nil
            2.times do

              den_term    = 1 + coeff*den_term
              num_term    = 1 + coeff.quo(num_term)

              den_term = continued_fraction_cutoff(epsabs) if den_term.abs < CUTOFF
              num_term = continued_fraction_cutoff(epsabs) if num_term.abs < CUTOFF

              den_term = 1.quo(den_term)

              delta_frac  = den_term * num_term
              cf         *= delta_frac

              coeff = -(a+k)*(a+b+k)*x.quo((a+2*k)*(a+2*k+1)) # coefficient for step 2
            end

            break if (delta_frac-1).abs < 2.0*Float::EPSILON
            break if !epsabs.nil? && (cf * (delta_frac-1).abs < epsabs)
          end

          if k > MAX_ITER
            raise("Exceeded maximum number of iterations") if epsabs.nil?
            return with_error ? [0.0/0, 0] : 0.0/0 # NaN if epsabs is set
          end

          with_error ? [cf, k * 4 * Float::EPSILON * cf.abs] : cf
        end


      end

    end
  end
end