#!/usr/bin/env python3
#
# Copyright (C) 2026 Susi Lehtola
#
# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.
#
# Generates the reference data for xc-special-functions.c, computing each
# value with mpmath at high precision.  Regenerate the embedded header with:
#
#     python3 xc-special-functions-gen.py > xc-special-functions-data.h
#
# (run from the testsuite/ directory).  Requires mpmath (ships with sympy).

import sys
import mpmath as mp

mp.mp.dps = 40

# argument grids per function, chosen to span each domain and its edges
POINTS = {
    "lambertw": [-0.36787, -0.3, -0.1, 0.0, 0.5, 1.0, 2.718281828, 10.0, 1e3, 1e8],
    "dilog":    [-100.0, -10.0, -1.0, 0.25, 0.5, 0.9, 1.0, 1.5, 2.0, 10.0, 100.0],
    "e1":       [1e-6, 1e-3, 0.1, 1.0, 5.0, 20.0, 50.0],
    "e1s":      [1e-6, 1e-3, 0.1, 1.0, 5.0, 20.0, 50.0],
    "I0":       [1e-4, 0.5, 1.0, 5.0, 20.0, 100.0, 600.0],
    "I1":       [1e-4, 0.5, 1.0, 5.0, 20.0, 100.0, 600.0],
    "I0s":      [1e-4, 0.5, 1.0, 5.0, 20.0, 100.0, 1e6],
    "I1s":      [1e-4, 0.5, 1.0, 5.0, 20.0, 100.0, 1e6],
    "K0":       [1e-4, 0.5, 1.0, 5.0, 20.0, 100.0, 600.0],
    "K1":       [1e-4, 0.5, 1.0, 5.0, 20.0, 100.0, 600.0],
    "K0s":      [1e-4, 0.5, 1.0, 5.0, 20.0, 100.0, 1e6],
    "K1s":      [1e-4, 0.5, 1.0, 5.0, 20.0, 100.0, 1e6],
    "erfcx":    [-10.0, -1.0, 0.0, 0.5, 1.0, 10.0, 100.0, 1e4],
}

# reference definitions matching the libxc conventions:
#   *_scaled (I -> exp(-|x|) I, K -> exp(x) K, e1s -> exp(x) E1)
#   dilog is the real part of the dilogarithm Li2(x)
REF = {
    "lambertw": lambda x: mp.lambertw(x),
    "dilog":    lambda x: mp.re(mp.polylog(2, x)),
    "e1":       lambda x: mp.e1(x),
    "e1s":      lambda x: mp.e1(x) * mp.exp(x),
    "I0":       lambda x: mp.besseli(0, x),
    "I1":       lambda x: mp.besseli(1, x),
    "I0s":      lambda x: mp.besseli(0, x) * mp.exp(-abs(x)),
    "I1s":      lambda x: mp.besseli(1, x) * mp.exp(-abs(x)),
    "K0":       lambda x: mp.besselk(0, x),
    "K1":       lambda x: mp.besselk(1, x),
    "K0s":      lambda x: mp.besselk(0, x) * mp.exp(x),
    "K1s":      lambda x: mp.besselk(1, x) * mp.exp(x),
    "erfcx":    lambda x: mp.erfc(x) * mp.exp(x * x),
}


def main():
    out = sys.stdout
    out.write("/* Reference data for xc-special-functions.c -- DO NOT EDIT.\n")
    out.write("   Generated by xc-special-functions-gen.py (mpmath, %d digits). */\n\n"
              % mp.mp.dps)
    out.write("static const struct { const char *name; double x, ref; }"
              " xc_sf_cases[] = {\n")
    for name in POINTS:
        for x in POINTS[name]:
            r = float(REF[name](mp.mpf(x)))
            out.write('  {%-11s %.17g, %.17g},\n' % ('"%s",' % name, x, r))
    out.write("};\n")


if __name__ == "__main__":
    main()
