#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Mon Jan 13 02:40:35 PM EST 2025
if __name__ == '__main__':
    import sys,os
    try:
        from misc.timer import Timer
        tt = Timer(f'[{os.getcwd()}] start ' + ' '.join(sys.argv))
    except:
        pass
import sys, os.path, os, glob, datetime
#import xarray as xr, numpy as np, pandas as pd, matplotlib.pyplot as plt
import pandas as pd, numpy as np
#more imports
wython = '/tigress/wenchang/wython'
if wython not in sys.path: sys.path.append(wython); print('added to python path:', wython)
#from misc import get_kws_from_argv
#
if __name__ == '__main__':
    try:
        tt.check('end import')
    except:
        pass
#
#start from here
ncpu_per_node = 112
#C360: 360 = 2**3 * 3**2 * 5**1
ivec = range(0, 6+1)
jvec = range(0, 4+1)
kvec = range(0, 2+1)

min_ncpu = 500
max_ncpu = 5000

iis = []
jjs = []
kks = []
ncpus = []
for kk in kvec:
    for jj in jvec:
        for ii in ivec:
            ncpu = 2**ii * 3**jj * 5**kk * 2 * 3
            if ncpu >= min_ncpu and ncpu<=max_ncpu:
                print(f'ii = {ii}; jj = {jj}; kk = {kk}; ncpu = {ncpu}')
                iis.append(ii)
                jjs.append(jj)
                kks.append(kk)
                ncpus.append(ncpu)

df = pd.DataFrame(dict(ii=iis, jj=jjs, kk=kks, ncpu=ncpus))
df = df.sort_values(by='ncpu')
df['nodes'] = df['ncpu']/ncpu_per_node
#df['reminder'] = df['ncpu']%ncpu_per_node
#df['reminderFrac'] = df['reminder']/112
df = df.set_index(np.array(range(df.ncpu.size)))
print(df)
ofile = __file__.replace('.py', '.csv')
df.to_csv(ofile, float_format='%.2f')
print('[saved]:', ofile)
 
 
if __name__ == '__main__':
    #from wyconfig import * #my plot settings
    
    #savefig
    if 'savefig' in sys.argv or 's' in sys.argv:
        figname = __file__.replace('.py', f'.png')
        if 'overwritefig' in sys.argv or 'o' in sys.argv:
            wysavefig(figname, overwritefig=True)
        else:
            wysavefig(figname)
 
    try:
        tt.check(f'**Done**')
    except:
        pass
    print()
    if 'notshowfig' in sys.argv or 'n' in sys.argv:
        pass
    else:
        if 'plt' in globals(): plt.show()