#!/cvmfs/northgrid.gridpp.ac.uk/simonsobservatory/conda/so-conda-20231121/bin/python

import os,sys
import unittest

# Must do the thread setup before we import pixell tests, otherwise those
# dependent on the threadpool will fail.

import multiprocessing
max_threads = multiprocessing.cpu_count()
assert max_threads>=1

# On MacOS default values of `DUCC0_NUM_THREADS` (used for run-time thread number selection
# as ducc0 uses pthreads not OMP) are not correctly handled, we must set it to a valid 
# integer value.
if sys.platform == "darwin":
    if "DUCC0_NUM_THREADS" not in os.environ:
        os.environ["DUCC0_NUM_THREADS"] = str(max_threads)
        print(f"Setting DUCC0_NUM_THREADS (not present) to {max_threads}.")

    try:
        ducc0_num_threads = int(os.environ["DUCC0_NUM_THREADS"])
    except ValueError:
        # We have a non-integer value for this environment variable.
        # Ducc0 does not currently (2023-10-16) have error handling for this
        # (at least on MacOS 13), so we need to set it to a valid value.
        os.environ["DUCC0_NUM_THREADS"] = str(max_threads)

from pixell.tests import *

unittest.main(exit=False)

def run_alm_benchmark(nthreads):
    os.system(f"""OMP_NUM_THREADS={nthreads} python -c 'from pixell import curvedsky,enmap,utils ;
import time ;
import numpy as np ;

np.random.seed(100) ;
shape,wcs = enmap.fullsky_geometry(res=12.*utils.arcmin) ;
imap = enmap.enmap(np.random.random(shape),wcs) ;

nsims = 40 ; 
lmax = int(6000 * (2./16.)); 

t0 = time.time() ;
for i in range(nsims): alm = curvedsky.map2alm(imap,lmax=lmax) ; omap = curvedsky.alm2map(alm,enmap.empty(shape,wcs)) ;
t1 = time.time();

total = t1-t0;
print(f"{{total:.4f}} seconds.");
'
    """)

print("Single threaded alm test:")
run_alm_benchmark(1)

if max_threads==1:
    print("Multi-threading not detected.")
else:
    print(f"Multi-threaded alm test with {max_threads} threads:")
    run_alm_benchmark(max_threads)