"""cuDSS uniform batch (CUDSS_CONFIG_UBATCH_SIZE > 128): NaN in the solution depending on the content of device memory
the library was never given.

A batch of identical, well-conditioned tridiagonal systems (n = 100, diagonal 5, off-diagonals 1) is solved with
cuDSS's uniform batch mode (ANALYSIS, FACTORIZATION, SOLVE; default configuration otherwise). Before creating the
cuDSS objects, the script allocates an unrelated device buffer and fills it with a byte pattern. That buffer is never
passed to cuDSS and stays allocated until the end. With 0xFF bytes (NaN as float64) a third of the solutions come
back NaN; with 0x00 bytes, or with UBATCH_SIZE <= 128, every solution is exact.

    python cudss_ubatch_repro.py [--batch 512] [--fill 255|0|none] [--mb 64]

Dependencies: cupy, scipy, nvmath-python (nvmath.bindings.cudss) and the cuDSS library.
"""

import argparse

import cupy as cp
import numpy as np
import scipy.sparse as sp
from nvmath.bindings import cudss

CUDA_R_64F, CUDA_R_32I = 1, 10


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--batch", type=int, default=512, help="CUDSS_CONFIG_UBATCH_SIZE")
    ap.add_argument("--fill", default="255", help="byte written to the unrelated buffer (255, 0) or 'none'")
    ap.add_argument("--mb", type=int, default=64, help="size of the unrelated buffer (MiB)")
    a = ap.parse_args()

    props = cp.cuda.runtime.getDeviceProperties(0)
    ver = [cudss.get_property(k) for k in (0, 1, 2)]  # MAJOR_VERSION, MINOR_VERSION, PATCH_LEVEL
    print(f"GPU {props['name'].decode()} | driver {cp.cuda.runtime.driverGetVersion()} | runtime "
          f"{cp.cuda.runtime.runtimeGetVersion()} | cuDSS {'.'.join(map(str, ver))} | CuPy {cp.__version__}")

    # 1. an unrelated device buffer, filled with a byte pattern, kept allocated, never given to cuDSS
    unrelated = None
    if a.fill != "none":
        unrelated = cp.cuda.runtime.malloc(a.mb << 20)
        cp.cuda.runtime.memset(unrelated, int(a.fill), a.mb << 20)
        cp.cuda.runtime.deviceSynchronize()

    # 2. batch of identical tridiagonal systems, exact solution x_true
    n, b = 100, a.batch
    m = sp.diags([np.ones(n - 1), np.full(n, 5.0), np.ones(n - 1)], [-1, 0, 1], format="csr")
    m.sort_indices()
    x_true = np.linspace(1.0, 2.0, n)
    rowptr = cp.asarray(m.indptr, dtype=cp.int32)
    colind = cp.asarray(m.indices, dtype=cp.int32)
    values = cp.asarray(np.tile(m.data, (b, 1)))  # b x nnz, contiguous
    rhs = cp.asarray(np.tile(m @ x_true, (b, 1)))  # b x n, contiguous
    x = cp.zeros((b, n))
    cp.cuda.Device().synchronize()

    # 3. cuDSS uniform batch: one matrix description, UBATCH_SIZE copies of values / rhs / x back to back
    handle = cudss.create()
    cudss.set_stream(handle, cp.cuda.get_current_stream().ptr)
    config = cudss.config_create()
    if b > 1:
        size = np.array([b], dtype=np.int32)
        cudss.config_set(config, cudss.ConfigParam.UBATCH_SIZE, size.ctypes.data, size.nbytes)
    data = cudss.data_create(handle)
    A = cudss.matrix_create_csr(n, n, m.nnz, rowptr.data.ptr, 0, colind.data.ptr, values.data.ptr, CUDA_R_32I,
                                CUDA_R_32I, CUDA_R_64F, cudss.MatrixType.GENERAL, cudss.MatrixViewType.FULL,
                                cudss.IndexBase.ZERO)
    B = cudss.matrix_create_dn(n, 1, n, rhs.data.ptr, CUDA_R_64F, cudss.Layout.COL_MAJOR)
    X = cudss.matrix_create_dn(n, 1, n, x.data.ptr, CUDA_R_64F, cudss.Layout.COL_MAJOR)
    for phase in (cudss.Phase.ANALYSIS, cudss.Phase.FACTORIZATION, cudss.Phase.SOLVE):
        cudss.execute(handle, phase, config, data, A, X, B)
    cp.cuda.Device().synchronize()

    sol = cp.asnumpy(x)
    bad = ~np.isfinite(sol).all(axis=1)
    err = np.abs(sol[~bad] - x_true).max() if (~bad).any() else float("nan")
    print(f"UBATCH_SIZE {b}, unrelated buffer {a.mb} MiB filled with {a.fill}: {int(bad.sum())} / {b} solutions "
          f"contain NaN; max error of the others {err:.1e}")

    for obj in (A, B, X):
        cudss.matrix_destroy(obj)
    cudss.data_destroy(handle, data)
    cudss.config_destroy(config)
    cudss.destroy(handle)
    if unrelated is not None:
        cp.cuda.runtime.free(unrelated)


if __name__ == "__main__":
    main()