lambert_w 2.0.1

// Copyright 2024-2026 Johanna Sörngård
// SPDX-License-Identifier: MIT OR Apache-2.0

//! This module contains an implementation of the approximation of the principal
//! branch of the Lambert W function
//! with 50 bits of accuracy from Fukushima's paper.
//! It is based on the Fortran implementation of the name "dw0c" by Fukushima.

// The coefficients in these rational minimax functions all have excessive precision.
// By keeping the full precision in the source code we can ensure that there is no confusion
// when comparing with the paper.
#![allow(clippy::excessive_precision)]

use crate::{
    generic_math::{ln, rational_function, sqrt},
    NEG_INV_E,
};

/// The principal branch of the Lambert W function computed to 50 bits of accuracy on 64-bit floats with Fukushima's method[^1].
///
/// # Examples
///
/// #### Basic usage
///
/// ```
/// use lambert_w::lambert_w0;
/// use approx::assert_abs_diff_eq;
///
/// let Ω = lambert_w0(1.0);
///
/// assert_abs_diff_eq!(Ω, 0.567143290409783873);
/// ```
///
/// #### Special cases
///
/// For inputs of -1/e and 0 the function returns exactly -1 and 0 respectively,
/// while an infinite input gives [`INFINITY`](f64::INFINITY):
///
/// ```
/// use lambert_w::{lambert_w0, NEG_INV_E};
///
/// assert_eq!(lambert_w0(NEG_INV_E), -1.0);
/// assert_eq!(lambert_w0(0.0), 0.0);
/// assert_eq!(lambert_w0(f64::INFINITY), f64::INFINITY);
/// ```
///
/// Inputs smaller than -1/e, as well as inputs of [`NAN`](f64::NAN), result in [`NAN`](f64::NAN):
///
/// ```
/// use lambert_w::{lambert_w0, NEG_INV_E};
///
/// assert!(lambert_w0(NEG_INV_E.next_down()).is_nan());
/// assert!(lambert_w0(f64::NAN).is_nan());
/// ```
///
/// # Reference
///
/// [^1]: Toshio Fukushima. **Precise and fast computation of Lambert W function by piecewise minimax rational function approximation with variable transformation**. DOI: [10.13140/RG.2.2.30264.37128](https://doi.org/10.13140/RG.2.2.30264.37128). November 2020.
#[must_use = "this is a pure function that only returns a value and has no side effects"]
pub fn lambert_w0(z: f64) -> f64 {
    // The critical arguments used in the if statements are the numbers in table 3 of the paper, column two, with 1/e added.
    // The coefficients in the rational functions are the ones in tables 10 through 14 in the paper.

    let zc = z - NEG_INV_E;

    if z < NEG_INV_E {
        f64::NAN
    } else if z == NEG_INV_E {
        -1.0
    } else if z == 0.0 {
        0.0
    } else if zc <= 2.549_893_906_503_473_571_6 {
        // W <= 0.893, X_1

        rational_function(
            sqrt(zc),
            [
                -0.999_999_999_999_999_889_0,
                -2.739_966_866_820_365_930_4,
                0.026_164_207_726_990_399_347,
                6.370_916_807_894_900_917,
                7.101_328_651_785_402_668,
                2.980_082_678_300_685_257_3,
                0.488_195_968_137_898_650_43,
                0.023_753_035_787_333_611_915,
                0.000_077_365_760_093_772_431_00,
            ],
            [
                1.0,
                5.071_610_848_417_428_005,
                9.986_838_818_354_528_337,
                9.660_755_192_207_886_908,
                4.794_372_899_133_611_905_2,
                1.162_970_347_770_452_230_0,
                0.118_494_625_007_337_552_33,
                0.003_432_652_513_240_222_648_8,
            ],
        )
    } else if zc <= 43.613_924_462_669_367_895 {
        // W <= 2.754, X_2

        rational_function(
            sqrt(zc),
            [
                -0.999_978_018_005_789_167_49,
                -0.704_157_515_904_836_022_72,
                2.123_226_083_280_252_907_1,
                2.389_676_070_293_571_834_1,
                0.777_653_118_050_291_752_44,
                0.089_686_698_993_644_741_433,
                0.003_306_248_575_374_640_355_9,
                0.000_025_106_760_479_132_851_033,
            ],
            [
                1.0,
                3.035_602_682_808_541_088_4,
                3.143_453_015_128_677_705_7,
                1.372_315_656_659_244_727_5,
                0.258_446_974_157_442_111_42,
                0.019_551_162_251_819_044_265,
                0.000_487_759_332_445_301_231_01,
                2.316_511_684_107_315_271_7e-6,
            ],
        )
    } else if zc <= 598.453_533_718_782_769_46 {
        // W <= 4.821, X_3

        rational_function(
            sqrt(zc),
            [
                -0.989_674_203_372_735_063_93,
                0.595_876_806_063_943_827_48,
                1.422_508_301_815_194_314_8,
                0.448_828_891_683_238_097_98,
                0.044_504_943_332_390_033_511,
                0.001_521_879_483_541_957_855_4,
                0.000_016_072_263_556_502_220_023,
                3.372_337_302_030_651_084_3e-8,
            ],
            [
                1.0,
                1.695_940_239_462_619_805_2,
                0.809_685_734_155_009_008_96,
                0.140_020_349_998_170_219_55,
                0.009_357_187_849_379_016_448_0,
                0.000_232_514_875_933_897_734_64,
                1.806_017_075_150_298_864_5e-6,
                2.575_066_733_701_592_422_4e-9,
            ],
        )
    } else if zc <= 8_049.491_985_075_761_910_9 {
        // W <= 7.041, X_4

        rational_function(
            sqrt(zc),
            [
                -0.773_164_919_972_062_255_17,
                1.139_133_350_429_670_378_3,
                0.431_161_172_552_170_744_92,
                0.035_773_078_319_037_507_449,
                0.000_964_416_405_805_590_927_40,
                8.972_385_459_867_586_475_7e-6,
                2.562_350_314_411_772_321_7e-8,
                1.434_881_377_841_663_145_3e-11,
            ],
            [
                1.0,
                0.746_572_874_565_144_180_83,
                0.126_297_770_334_193_505_76,
                0.006_974_151_295_956_318_488_1,
                0.000_140_893_392_443_553_548_92,
                1.025_743_288_315_294_307_8e-6,
                2.290_268_719_011_923_094_0e-9,
                9.279_423_101_326_450_166_4e-13,
            ],
        )
    } else if zc <= 111_124.954_121_217_814_20 {
        // W <= 9.380, X_5

        rational_function(
            sqrt(zc),
            [
                0.120_071_016_715_536_884_30,
                0.833_526_408_299_128_228_96,
                0.070_142_775_916_948_337_582,
                0.001_484_635_798_547_512_484_9,
                0.000_010_478_757_366_110_155_290,
                2.571_589_298_707_103_852_7e-8,
                1.938_421_447_960_647_474_9e-11,
                2.844_704_903_913_940_965_2e-15,
            ],
            [
                1.0,
                0.253_967_388_456_191_266_30,
                0.012_839_238_907_330_317_393,
                0.000_202_753_756_325_109_973_71,
                1.148_295_607_344_914_138_4e-6,
                2.318_837_060_567_426_364_7e-9,
                1.427_199_416_574_256_341_9e-12,
                1.588_483_694_239_479_696_1e-16,
            ],
        )
    } else if zc <= 1.587_042_981_208_229_711_2e6 {
        // W <= 11.809, X_6

        rational_function(
            sqrt(zc),
            [
                1.722_110_443_993_771_011_2,
                0.399_195_942_864_842_756_05,
                0.007_988_554_014_068_502_893_7,
                0.000_042_889_742_253_257_920_541,
                7.814_682_818_052_986_498_1e-8,
                4.981_963_876_435_468_235_9e-11,
                9.765_088_971_426_529_460_6e-15,
                3.705_299_728_172_172_443_9e-19,
            ],
            [
                1.0,
                0.074_007_438_118_020_543_008,
                0.001_033_350_150_669_774_054_5,
                4.436_085_803_572_750_850_6e-6,
                6.782_291_231_637_104_157_0e-9,
                3.683_435_670_763_949_202_1e-12,
                6.083_615_956_026_604_116_8e-16,
                1.814_986_933_598_122_531_6e-20,
            ],
        )
    } else if zc <= 2.341_470_840_187_545_950_9e7 {
        // W <= 14.308, X_7

        rational_function(
            sqrt(zc),
            [
                3.752_931_402_343_454_425_6,
                0.154_913_426_903_578_065_25,
                0.000_756_631_406_759_007_845_05,
                1.027_160_923_596_997_905_9e-6,
                4.785_324_767_593_006_615_0e-10,
                7.832_804_077_027_547_441_0e-14,
                3.943_303_375_839_103_665_3e-18,
                3.823_286_220_566_028_397_8e-23,
            ],
            [
                1.0,
                0.020_112_985_338_854_443_555,
                0.000_074_712_286_154_830_141_768,
                8.480_059_800_369_383_746_9e-8,
                3.418_242_413_037_691_176_2e-11,
                4.886_625_913_969_095_789_9e-15,
                2.122_337_362_683_463_417_8e-19,
                1.664_298_567_126_058_251_5e-24,
            ],
        )
    } else if zc <= 3.557_647_430_800_996_522_5e8 {
        // W <= 16.865, X_8

        rational_function(
            sqrt(zc),
            [
                6.019_654_205_560_655_557_7,
                0.053_496_672_841_797_864_762,
                0.000_064_340_849_275_316_501_519,
                2.196_909_010_009_596_748_5e-8,
                2.592_798_893_703_306_107_0e-12,
                1.077_919_816_180_152_730_8e-16,
                1.378_042_409_101_789_830_1e-21,
                3.376_897_315_074_255_280_2e-27,
            ],
            [
                1.0,
                0.005_280_968_370_423_337_167_5,
                5.102_050_121_938_955_808_2e-6,
                1.501_831_229_227_083_210_3e-9,
                1.567_770_663_641_318_837_9e-13,
                5.799_204_123_891_187_836_1e-18,
                6.513_317_077_032_078_025_9e-23,
                1.320_508_013_921_340_607_1e-28,
            ],
        )
    } else if zc <= 5.550_171_629_616_362_785_4e9 {
        // W <= 19.468, X_9

        rational_function(
            sqrt(zc),
            [
                8.428_026_850_098_970_159_7,
                0.017_155_758_546_279_713_315,
                5.083_662_066_982_932_150_8e-6,
                4.335_490_369_183_258_180_2e-10,
                1.284_101_714_564_558_338_5e-14,
                1.341_910_676_974_588_592_7e-19,
                4.310_169_845_549_222_575_0e-25,
                2.642_243_342_208_818_754_9e-31,
            ],
            [
                1.0,
                0.001_357_200_675_459_530_031_5,
                3.353_524_348_142_620_369_4e-7,
                2.520_696_924_642_126_412_8e-11,
                6.713_622_627_306_053_049_6e-16,
                6.332_422_668_085_468_657_4e-21,
                1.812_816_740_001_377_419_4e-26,
                9.366_203_005_813_679_688_9e-33,
            ],
        )
    } else if zc <= 8.867_470_483_965_777_533_1e10 {
        // W <= 22.112, X_10

        rational_function(
            sqrt(zc),
            [
                10.931_063_230_472_498_189_,
                0.005_222_423_454_024_553_298_2,
                3.799_610_571_181_012_968_2e-7,
                8.030_579_353_341_035_582_4e-12,
                5.913_978_562_709_060_586_6e-17,
                1.538_202_035_953_302_872_4e-22,
                1.228_894_412_626_810_943_2e-28,
                1.866_508_927_066_012_239_8e-35,
            ],
            [
                1.,
                0.000_343_287_025_511_975_777_97,
                2.139_535_151_853_884_447_6e-8,
                4.052_417_018_663_159_415_9e-13,
                2.718_142_431_533_571_042_0e-18,
                6.453_898_663_835_549_089_4e-24,
                4.649_461_378_588_898_794_2e-30,
                6.044_202_436_729_938_761_6e-37,
            ],
        )
    } else if zc <= 1.447_779_186_527_290_281_6e12 {
        // W <= 24.791, X_11

        rational_function(
            sqrt(zc),
            [
                13.502_943_080_893_871_412,
                0.001_528_463_650_634_626_457_2,
                2.715_696_735_826_234_616_6e-8,
                1.411_039_405_124_216_177_2e-13,
                2.560_573_431_121_972_846_1e-19,
                1.642_129_372_442_533_746_3e-25,
                3.232_494_469_143_584_355_3e-32,
                1.205_466_264_125_178_315_5e-39,
            ],
            [
                1.0,
                0.000_085_701_512_879_089_462_255,
                1.331_124_443_575_269_156_3e-9,
                6.278_892_444_038_534_726_9e-15,
                1.048_378_815_225_220_482_4e-20,
                6.194_349_996_624_916_088_6e-27,
                1.110_156_786_034_091_729_4e-33,
                3.589_738_112_830_896_259_0e-41,
            ],
        )
    } else if zc <= 2.411_145_863_251_185_193_1e13 {
        // W <= 27.500, X_12

        rational_function(
            sqrt(zc),
            [
                16.128_076_167_439_014_775,
                0.000_433_603_851_764_670_691_31,
                1.869_640_387_182_091_646_6e-9,
                2.369_179_576_690_148_604_5e-15,
                1.050_319_182_696_315_489_3e-21,
                1.646_192_757_360_676_426_3e-28,
                7.913_827_608_347_452_293_1e-36,
                7.184_589_034_370_166_876_0e-44,
            ],
            [
                1.0,
                0.000_021_154_255_263_102_938_752,
                8.100_611_544_232_328_053_8e-11,
                9.415_598_602_216_990_573_8e-17,
                3.872_512_790_229_530_225_4e-23,
                5.634_465_111_557_056_506_6e-30,
                2.486_095_108_421_002_919_1e-37,
                1.978_830_473_742_778_740_5e-45,
            ],
        )
    } else if zc <= 4.089_703_644_260_084_556_4e14 {
        // W <= 30.236, X_13

        rational_function(
            sqrt(zc),
            [
                18.796_301_105_534_486_604,
                0.000_119_894_433_396_464_691_57,
                1.246_337_752_867_686_325_0e-10,
                3.821_945_685_801_036_817_2e-17,
                4.105_569_393_025_208_326_5e-24,
                1.559_523_145_604_846_424_6e-31,
                1.815_717_355_307_798_696_2e-39,
                3.980_799_776_432_616_624_5e-48,
            ],
            [
                1.0,
                5.169_103_198_835_992_232_9e-6,
                4.832_557_182_331_371_193_2e-12,
                1.370_788_874_691_692_810_7e-18,
                1.375_456_085_002_448_033_7e-25,
                4.881_188_297_566_180_518_4e-33,
                5.251_864_182_817_020_189_4e-41,
                1.019_211_959_313_475_644_0e-49,
            ],
        )
    } else if zc <= 7.055_590_147_678_997_240_2e15 {
        // W <= 32.996, X_14

        rational_function(
            sqrt(zc),
            [
                21.500_582_830_667_332_906,
                0.000_032_441_943_237_735_273_768,
                8.076_496_341_683_755_914_8e-12,
                5.948_844_550_612_288_352_3e-19,
                1.536_410_618_721_586_153_1e-26,
                1.403_323_129_700_238_699_5e-34,
                3.925_987_271_230_577_043_0e-43,
                2.062_908_638_225_773_751_7e-52,
            ],
            [
                1.0,
                1.251_531_764_243_385_019_7e-6,
                2.831_031_421_481_707_480_6e-13,
                1.942_366_641_612_363_799_8e-20,
                4.712_861_600_415_735_971_4e-28,
                4.043_334_739_183_994_596_0e-36,
                1.051_514_144_383_118_727_1e-44,
                4.931_649_093_543_692_730_7e-54,
            ],
        )
    } else if zc <= 1.236_660_755_797_672_728_7e17 {
        // W <= 35.779, X_15

        rational_function(
            sqrt(zc),
            [
                24.235_812_532_416_977_267,
                8.616_150_599_577_680_250_9e-6,
                5.103_343_156_186_827_369_2e-13,
                8.964_239_366_584_963_816_4e-21,
                5.525_436_418_109_742_077_7e-29,
                1.204_507_272_405_060_579_2e-37,
                8.037_299_717_652_684_018_4e-47,
                1.004_914_081_214_649_261_1e-56,
            ],
            [
                1.0,
                3.004_676_184_474_947_798_7e-7,
                1.630_910_427_085_546_322_3e-14,
                2.684_227_103_029_893_132_9e-22,
                1.561_967_263_245_888_119_5e-30,
                3.213_168_903_039_798_427_4e-39,
                2.003_239_624_530_768_413_4e-48,
                2.252_027_455_467_633_193_8e-58,
            ],
        )
    } else if zc <= 2.199_937_348_793_099_977_5e18 {
        // W <= 38.582, X_16

        rational_function(
            sqrt(zc),
            [
                26.998_134_347_987_436_511,
                2.251_225_776_757_228_586_6e-6,
                3.152_123_075_986_696_394_1e-14,
                1.311_403_571_979_063_154_1e-22,
                1.915_678_403_396_236_614_6e-31,
                9.896_700_305_344_479_916_3e-41,
                1.564_042_389_844_843_354_8e-50,
                4.621_619_304_066_487_260_6e-61,
            ],
            [
                1.0,
                7.157_267_637_090_757_389_8e-8,
                9.250_050_609_111_576_082_6e-16,
                3.623_981_958_278_757_303_1e-24,
                5.018_771_249_380_042_411_8e-33,
                2.456_586_198_821_806_903_9e-42,
                3.643_565_843_399_166_028_4e-52,
                9.743_249_064_015_534_600_4e-63,
            ],
        )
    } else if zc <= 3.968_539_219_834_401_615_5e19 {
        // W <= 41.404, X_17

        rational_function(
            sqrt(zc),
            [
                29.784_546_702_831_970_770,
                5.797_176_439_217_132_994_4e-7,
                1.906_987_279_260_195_080_8e-15,
                1.866_870_087_085_876_331_2e-24,
                6.420_051_095_337_094_007_5e-34,
                7.807_662_465_081_896_855_9e-44,
                2.902_963_869_695_631_565_4e-54,
                2.014_187_045_856_617_985_3e-65,
            ],
            [
                1.0,
                1.692_446_318_046_970_637_2e-8,
                5.170_393_431_125_454_011_1e-17,
                4.787_153_272_156_006_909_5e-26,
                1.566_440_583_254_514_936_8e-35,
                1.811_313_798_238_133_139_8e-45,
                6.345_415_028_949_541_952_9e-56,
                4.007_296_402_524_439_796_7e-67,
            ],
        )
    } else if zc <= 1.412_707_514_527_465_206_9e104 {
        // W <= 234.358, U_18

        rational_function(
            ln(zc),
            [
                0.744_134_994_601_267_761_43,
                0.414_032_436_180_059_111_60,
                0.260_125_641_667_734_161_70,
                0.021_450_457_095_960_295_520,
                0.000_518_723_772_647_059_075_77,
                4.357_469_356_831_997_599_6e-6,
                1.236_306_605_892_170_671_6e-8,
                9.019_414_776_630_995_753_7e-12,
            ],
            [
                1.0,
                0.334_878_110_674_670_109_07,
                0.023_756_834_394_570_626_395,
                0.000_542_256_330_089_077_351_60,
                4.437_898_005_257_962_303_7e-6,
                1.243_658_549_766_809_933_0e-8,
                9.022_582_586_763_185_221_5e-12,
                -4.205_783_627_010_971_665_4e-19,
            ],
        )
    } else if zc < f64::INFINITY {
        //   U_19

        rational_function(
            ln(zc),
            [
                -0.615_144_128_127_297_615_26,
                0.679_793_101_336_309_365_80,
                0.089_685_353_704_585_808_963,
                0.001_564_494_148_398_937_924_9,
                7.734_990_187_817_635_116_2e-6,
                1.289_164_754_669_943_522_9e-8,
                7.089_032_598_897_381_265_6e-12,
                9.841_979_033_427_971_145_3e-16,
            ],
            [
                1.0,
                0.097_300_263_710_401_439_315,
                0.001_610_367_274_844_205_865_1,
                7.824_774_100_307_700_001_2e-6,
                1.294_926_130_897_134_520_9e-8,
                7.098_691_121_934_282_713_0e-12,
                9.842_628_504_222_704_497_9e-16,
                -1.596_014_725_260_605_535_2e-24,
            ],
        )
    } else if zc.is_nan() {
        f64::NAN
    } else {
        f64::INFINITY
    }
}