#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Tue Jul 26 12:00:40 EDT 2022
#README: get models that have available data across experiments(both historical and future) and variables (monthly ts/daily tas/daily tasmax), and no missing data
if __name__ == '__main__':
    import sys
    from misc.timer import Timer
    tt = Timer('start ' + ' '.join(sys.argv))
import sys, os.path, os, glob, datetime
#import xarray as xr, numpy as np, pandas as pd, matplotlib.pyplot as plt
#more imports
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
cashed = True if 'cashed' in sys.argv else False #only examine models saved in current directory
#amip = True if 'amip' in sys.argv else False
amip = True if os.path.basename(os.path.abspath(os.getcwd())) == 'highresSST' else False
def get_models(idir):
    models = [m for m in os.listdir(idir) if os.path.isdir(os.path.join(idir,m)) and not m.startswith('wy_')]
    models.sort()
    return models

idirs = []

#daily data
#present
daname = 'tasmax'
expname = 'highresSST-present' if amip else 'hist-1950'
idir = f'/tigress/wenchang/data/cmip6/daily/{daname}/{expname}'
idirs.append(idir)
models = get_models(idir)
models_com = models
print(daname, expname)
print(models)
print()
#future
expname = 'highresSST-future' if amip else 'highres-future'
idir = f'/tigress/wenchang/data/cmip6/daily/{daname}/{expname}'
idirs.append(idir)
models = get_models(idir)
models_com = sorted(list(set(models_com)&(set(models))))
print(daname, expname)
print(models)
print()

#daily data
daname = 'tas'
expname = 'highresSST-present' if amip else 'hist-1950'
idir = f'/tigress/wenchang/data/cmip6/daily/{daname}/{expname}'
idirs.append(idir)
models = get_models(idir)
models_com = sorted(list(set(models_com)&(set(models))))
print(daname, expname)
print(models)
print()
expname = 'highresSST-future' if amip else 'highres-future'
idir = f'/tigress/wenchang/data/cmip6/daily/{daname}/{expname}'
idirs.append(idir)
models = get_models(idir)
models_com = sorted(list(set(models_com)&(set(models))))
print(daname, expname)
print(models)
print()

#monthly ts
daname = 'ts'
expname = 'highresSST-present' if amip else 'hist-1950'
idir = f'/tigress/wenchang/data/cmip6/variables/{daname}/{expname}'
idirs.append(idir)
models = get_models(idir)
models_com = sorted(list(set(models_com)&(set(models))))
print(daname, expname)
print(models)
print()
expname = 'highresSST-future' if amip else 'highres-future'
idir = f'/tigress/wenchang/data/cmip6/variables/{daname}/{expname}'
idirs.append(idir)
models = get_models(idir)
models_com = sorted(list(set(models_com)&(set(models))))
print(daname, expname)
print(models)
print()

print('models in common:')
print(models_com)
print()
if cashed:
    models_com = [m for m in os.listdir(os.getcwd()) if os.path.isdir(m) and not m.endswith('_drop')]
    models_com.sort()
    print('models in common (saved early in current directory):')
    print(models_com)
    print()

print('model directories:')
for im,model in enumerate(models_com,start=1):
    print(f'{im:2d} of {len(models_com):2d}: {model}')
    missing_data = False
    members_com = []
    for jj,idir in enumerate(idirs):
        _, daname, expname = idir.split('/')[-3:]
        #print(_, daname, expname, model)
        dd = os.path.join(idir, model)
        #print(dd)
        #ensemble members
        members = [s for s in os.listdir(dd) if os.path.isdir(os.path.join(dd,s)) and s.startswith('r')]
        members.sort()
        if jj == 0:
            members_com = members
        else:
            members_com = sorted(list(set(members_com)&(set(members))))
    print(members_com)
            
    for member in members_com:
        missing_data = False
        for idir in idirs:
            _, daname, expname = idir.split('/')[-3:]
            print(_, daname, expname, model, member)
            #ncfiles
            dd = os.path.join(idir, model, member)
            print(dd)
            ncfiles = [s for s in os.listdir(dd) if s.endswith('.nc')]
            ncfiles.sort()
            #ncfile, e.g. ts_Amon_MPI-ESM1-2-XR_highres-future_r1i1p1f1_gn_201501-201512.nc
            year_start = int( ncfiles[0].split('_')[-1][0:4] )
            year_stop = int( ncfiles[-1].split('-')[-1][0:4] )
            #make sure it covers the standard year range, e.g. 1950-2014 for hist-1950/highresSST-present; 2015-2050 for highres-future/highresSST-future
            if expname.endswith('future'):
                if year_start > 2015 or year_stop < 2050:
                    missing_data = True
                    print('*'*10, 'warning: some years of data are missing...', '*'*10)
            else:
                if year_start > 1950 or year_stop < 2014:
                    missing_data = True
                    print('*'*10, 'warning: some years of data are missing...', '*'*10)
            print(f'    {ncfiles[0]}')
            if len(ncfiles)>2:
                print('    ...')
            print(f'    {ncfiles[-1]}')
            print(f'    {len(ncfiles)} nc files')
        
        if not missing_data:
            odir = os.path.join(model, member)
            if os.path.exists(odir):
                print('[dir exists]:', odir)
            else:
                os.makedirs(odir)
                print('[dir created]:', odir)
        else:
            print('*'*10, 'missing data for', model, member, '*'*10)
    
    print()
 
if __name__ == '__main__':
    #from wyconfig import * #my plot settings
    
    #savefig
    if len(sys.argv)>1 and 'savefig' in sys.argv[1:]:
        figname = __file__.replace('.py', f'.png')
        if 'overwritefig' in sys.argv[1:]:
            wysavefig(figname, overwritefig=True)
        else:
            wysavefig(figname)
    tt.check(f'**Done**')
    print()
    if 'notshowfig' in sys.argv:
        pass
    else:
        pass#plt.show()