//
//  Distributions.swift
//  Metrika
//
//  Author:  Simon-Pierre Boucher
//  Contact: contact@spboucher.ai
//  Copyright © 2026 Simon-Pierre Boucher. All rights reserved.
//

import Foundation

/// Special functions and distribution CDFs used for inference output.
/// Implementations follow the classic Lanczos / continued-fraction
/// formulations and are accurate to well beyond the 1e-10 project
/// tolerance for the parameter ranges used in regression output.
public enum ZQDistributions {

    /// Natural log of the gamma function (Lanczos approximation, g=7).
    public static func logGamma(_ x: Double) -> Double {
        precondition(x > 0, "logGamma requires x > 0")
        let coefficients: [Double] = [
            676.5203681218851, -1259.1392167224028, 771.32342877765313,
            -176.61502916214059, 12.507343278686905, -0.13857109526572012,
            9.9843695780195716e-6, 1.5056327351493116e-7,
        ]
        if x < 0.5 {
            // Reflection formula.
            return log(.pi / sin(.pi * x)) - logGamma(1 - x)
        }
        let z = x - 1
        var sum = 0.99999999999980993
        for (i, c) in coefficients.enumerated() {
            sum += c / (z + Double(i) + 1)
        }
        let t = z + 7.5
        return 0.5 * log(2 * .pi) + (z + 0.5) * log(t) - t + log(sum)
    }

    /// Regularized incomplete beta function I_x(a, b) via the Lentz
    /// continued-fraction algorithm.
    public static func incompleteBeta(a: Double, b: Double, x: Double) -> Double {
        precondition(a > 0 && b > 0, "incompleteBeta requires a, b > 0")
        if x <= 0 { return 0 }
        if x >= 1 { return 1 }

        let logBeta = logGamma(a + b) - logGamma(a) - logGamma(b)
        let front = exp(logBeta + a * log(x) + b * log(1 - x))

        // Use the symmetry relation to keep the continued fraction in its
        // rapidly converging region.
        if x < (a + 1) / (a + b + 2) {
            return front * betaContinuedFraction(a: a, b: b, x: x) / a
        } else {
            return 1 - front * betaContinuedFraction(a: b, b: a, x: 1 - x) / b
        }
    }

    private static func betaContinuedFraction(a: Double, b: Double, x: Double) -> Double {
        let tiny = 1e-300
        let epsilon = 1e-16
        let qab = a + b
        let qap = a + 1
        let qam = a - 1

        var c = 1.0
        var d = 1 - qab * x / qap
        if abs(d) < tiny { d = tiny }
        d = 1 / d
        var h = d

        for m in 1...300 {
            let dm = Double(m)
            // Even step.
            var numerator = dm * (b - dm) * x / ((qam + 2 * dm) * (a + 2 * dm))
            d = 1 + numerator * d
            if abs(d) < tiny { d = tiny }
            c = 1 + numerator / c
            if abs(c) < tiny { c = tiny }
            d = 1 / d
            h *= d * c
            // Odd step.
            numerator = -(a + dm) * (qab + dm) * x / ((a + 2 * dm) * (qap + 2 * dm))
            d = 1 + numerator * d
            if abs(d) < tiny { d = tiny }
            c = 1 + numerator / c
            if abs(c) < tiny { c = tiny }
            d = 1 / d
            let delta = d * c
            h *= delta
            if abs(delta - 1) < epsilon { break }
        }
        return h
    }

    /// Standard normal CDF.
    public static func normalCDF(_ z: Double) -> Double {
        0.5 * erfc(-z / 2.0.squareRoot())
    }

    /// Student t CDF with `df` degrees of freedom.
    public static func studentTCDF(_ t: Double, df: Double) -> Double {
        precondition(df > 0, "studentTCDF requires df > 0")
        let x = df / (df + t * t)
        let tail = 0.5 * incompleteBeta(a: df / 2, b: 0.5, x: x)
        return t > 0 ? 1 - tail : tail
    }

    /// Two-sided p-value for a t statistic, computed directly from the
    /// tail integral: 2·P(T > |t|) = I_{df/(df+t²)}(df/2, 1/2). Going
    /// through `2·(1 − CDF)` would lose all precision below ~1e-16.
    public static func tTestPValue(_ t: Double, df: Double) -> Double {
        incompleteBeta(a: df / 2, b: 0.5, x: df / (df + t * t))
    }

    /// F distribution CDF with (df1, df2) degrees of freedom.
    public static func fCDF(_ f: Double, df1: Double, df2: Double) -> Double {
        precondition(df1 > 0 && df2 > 0, "fCDF requires positive df")
        if f <= 0 { return 0 }
        let x = df1 * f / (df1 * f + df2)
        return incompleteBeta(a: df1 / 2, b: df2 / 2, x: x)
    }

    /// Upper-tail p-value for an F statistic, computed directly as
    /// I_{df2/(df2+df1·f)}(df2/2, df1/2) to preserve precision for large F.
    public static func fTestPValue(_ f: Double, df1: Double, df2: Double) -> Double {
        if f <= 0 { return 1 }
        return incompleteBeta(a: df2 / 2, b: df1 / 2, x: df2 / (df2 + df1 * f))
    }

    /// Regularized lower incomplete gamma P(a, x), via the series for
    /// x < a+1 and the Lentz continued fraction otherwise.
    public static func incompleteGamma(a: Double, x: Double) -> Double {
        precondition(a > 0, "incompleteGamma requires a > 0")
        if x <= 0 { return 0 }

        if x < a + 1 {
            // Series representation.
            var term = 1.0 / a
            var sum = term
            var ap = a
            for _ in 0..<500 {
                ap += 1
                term *= x / ap
                sum += term
                if abs(term) < abs(sum) * 1e-16 { break }
            }
            return sum * exp(-x + a * log(x) - logGamma(a))
        } else {
            // Continued fraction for Q(a, x); P = 1 − Q.
            let tiny = 1e-300
            var b = x + 1 - a
            var c = 1 / tiny
            var d = 1 / b
            var h = d
            for i in 1...500 {
                let an = -Double(i) * (Double(i) - a)
                b += 2
                d = an * d + b
                if abs(d) < tiny { d = tiny }
                c = b + an / c
                if abs(c) < tiny { c = tiny }
                d = 1 / d
                let delta = d * c
                h *= delta
                if abs(delta - 1) < 1e-16 { break }
            }
            let q = exp(-x + a * log(x) - logGamma(a)) * h
            return 1 - q
        }
    }

    /// Chi-square CDF with `df` degrees of freedom.
    public static func chiSquareCDF(_ x: Double, df: Double) -> Double {
        x <= 0 ? 0 : incompleteGamma(a: df / 2, x: x / 2)
    }

    /// Upper-tail p-value for a chi-square statistic.
    public static func chiSquarePValue(_ x: Double, df: Double) -> Double {
        1 - chiSquareCDF(x, df: df)
    }

    /// Standard normal quantile via bisection on the CDF (erfc-based, so
    /// accurate to ~1e-15) — used for z confidence intervals.
    public static func normalQuantile(_ p: Double) -> Double {
        precondition(p > 0 && p < 1, "quantile requires 0 < p < 1")
        if abs(p - 0.5) < 1e-15 { return 0 }
        var low = -40.0, high = 40.0
        for _ in 0..<200 {
            let mid = 0.5 * (low + high)
            if normalCDF(mid) < p { low = mid } else { high = mid }
            if high - low < 1e-15 * max(1, abs(mid)) { break }
        }
        return 0.5 * (low + high)
    }

    /// Student t quantile (inverse CDF) via bisection refined with Newton
    /// steps — used for confidence intervals.
    public static func studentTQuantile(_ p: Double, df: Double) -> Double {
        precondition(p > 0 && p < 1, "quantile requires 0 < p < 1")
        if abs(p - 0.5) < 1e-15 { return 0 }

        var low = -1000.0, high = 1000.0
        var mid = 0.0
        for _ in 0..<200 {
            mid = 0.5 * (low + high)
            if studentTCDF(mid, df: df) < p {
                low = mid
            } else {
                high = mid
            }
            if high - low < 1e-13 * max(1, abs(mid)) { break }
        }
        return 0.5 * (low + high)
    }
}