import os, sys, argparse, re
from argparse import RawTextHelpFormatter

'''
The Big Idea:  Using the Archive_Map and the (sproket-based) ESGF publication report,
produce from each the list of "canonical" experiment/dataset "keys" where

	akey = ( <model>, <experiment>, <ensemble> )
and
	dkey = <realm_grid_freq>

Use these to index a dictionary structure "dataset_status[akey][dkey] = ( A: bool, W: bool, P: bool, S: bool )
with the Boolean values set according to the presence or absence of the key-pair in either of the
two input files (archive map, esgf publication report), or the warehouse or publication directories.
'''

helptext = '''
    Usage:  awps_dataset_status -s sproket_publication_report

        The awps_dataset_status is produced by plying 4 sources to determine whether datasets exist
        in any of (Archive,Warehouse,PubDirs,SproketPubReport), hereafter (A,W,P,S).

        Existence is determined by:

        (A):  Appearance in the Archive_Map (/p/user_pub/e3sm/archive/.cfg/Archive_Map)
        (W):  The Warehouse directories (/p/user_pub/e3sm/warehouse/E3SM/(facets)/[v0|v1...]/
        (W):  The Publication directories (/p/user_pub/work/E3SM/(facets)/[v0|v1...]/
        (S):  The provided sproket_publication_report

        The prerequisite sproket_publication_report is generated by running the scripts

            sproket_search.sh   (produces E3SM_datafile_urls-<date>)
            process_sproket_output.sh (takes above output and produces ESGF_publication_report-<date>)

'''

# INPUT FILES
arch_map  = '/p/user_pub/e3sm/archive/.cfg/Archive_Map'
warehouse = '/p/user_pub/e3sm/warehouse/E3SM'
publicati = '/p/user_pub/work/E3SM'
esgf_pr   = ''
# esgf_pr   = '/p/user_pub/e3sm/bartoletti1/Pub_Status/sproket/ESGF_publication_report-20200915.144250'

def assess_args():
    global thePWD
    global esgf_pr

    thePWD = os.getcwd()
    # parser = argparse.ArgumentParser(description=helptext, prefix_chars='-')
    parser = argparse.ArgumentParser(description=helptext, prefix_chars='-', formatter_class=RawTextHelpFormatter)
    parser._action_groups.pop()
    required = parser.add_argument_group('required arguments')
    optional = parser.add_argument_group('optional arguments')

    required.add_argument('-s', action='store', dest="esgf_pr", type=str, required=True)

    args = parser.parse_args()

    esgf_pr = args.esgf_pr

# Generic Convenience Functions =============================

def loadFileLines(afile):
    retlist = []
    if len(afile):
        with open(afile,"r") as f:
            retlist = f.read().split('\n')
        retlist = [ _ for _ in retlist if _[:-1] ]
    return retlist



#### BEGIN rationalizing archive and publication experiment-case names, and dataset-type names ####
# dsid = proj.model.experiment.resolution[.tuning].realm.grid.outtype.freq.ens.ver

# call to specialize a publication experiment name to an archive experiment name

def specialize_expname(expn, reso, tune):
    if expn == 'F2010plus4k':
        expn = 'F2010-plus4k'
    if expn[0:5] == 'F2010' or expn == '1950-Control':
        if reso[0:4] == '1deg' and tune == 'highres':
            expn = expn + '-LRtunedHR'
        else:
            expn = expn + '-HR'
    return expn

def get_idval(ensdir):
    return '.'.join(ensdir.split(os.sep)[5:])

def get_dsid_arch_key(dsid):    # dsid:  see above 
    # print(f'DSID: {dsid}')
    comps=dsid.split('.')
    expname = specialize_expname(comps[2],comps[3],comps[4])
    return comps[1],expname,comps[-1]

def get_dsid_type_key( dsid ):
    comps=dsid.split('.')
    realm = comps[-5]
    gridv = comps[-4]   # may be overloaded with "namefile" or "restart"
    otype = comps[-3]
    freq = comps[-2]    # will be "fixed" for gridv in "namefile" or "restart"

    if realm == 'atmos':
        realm = 'atm'
    elif realm == 'land':
        realm = 'lnd'
    elif realm == 'ocean':
        realm = 'ocn'

    if gridv == 'native':
        grid = 'nat'
    elif otype == 'climo':
        grid = 'climo'
    elif otype == 'monClim':
        grid = 'climo'
        freq = 'mon'
    elif otype == 'seasonClim':
        grid = 'climo'
        freq = 'season'
    elif otype == 'time-series':
        grid = 'reg'
        freq = 'ts-' + freq
    elif gridv == 'namefile':
        grid = 'namefile'
        freq = 'fixed'
    elif gridv == 'restart':
        grid = 'restart'
        freq = 'fixed'
    else:
        grid = 'reg'
    return '_'.join([realm,grid,freq])

#### COMPLETED rationalizing archive and publication experiment-case names, and dataset-type names ####

def isVLeaf(_):
    if len(_) > 1 and _[0] == 'v' and _[1] in '0123456789':
        return True
    return False

def dataset_print_csv( akey, dkey ):
    print(f'{akey[0]},{akey[1]},{akey[2]},{dkey}')

debug = False

def main():

    assess_args()

    # split the Archive_Map into a list of records, each record a list of fields
    #   Campaign,Model,Experiment,Ensemble,DatasetType,ArchivePath,DatatypeTarExtractionPattern,Notes
    contents = loadFileLines(arch_map)
    am_list = [ aline.split(',') for aline in contents if aline[:-1] ]

    # create a sorted list of unique dataset types
    dstype_list = [ arch_rec[4] for arch_rec in am_list ]
    dstype_list = list(set(dstype_list))
    dstype_list.sort()

    # create a dictionary keyed by 'arch_loc_key' = tuple (Model,Experiment,Ensemble) with value the SET of archive path(s)
    # (Not really needed here, but nice to have, and we can reuse its keys for the dataset_status table)
    arch_loc_dict = {}
    for _ in am_list:
        arch_loc_dict[ tuple([_[1],_[2],_[3]]) ] = set()
    for _ in am_list:
        arch_loc_dict[ tuple([_[1],_[2],_[3]]) ] |= {_[5]}  # add another archive path to the set - may be more than one.

    # create a dictionary keyed by 'arch_loc_key', each value a dictionary keyed by dstype, value a Boolean dictionary (A,W,P)
    # For each (model,experiment,ensemble), create a dataset_type entry for every allowable dataset-type
    # initialize "Archive,Warehouse,Published" to ALL FALSE.
    dataset_status = {}
    for akey in arch_loc_dict:
        dataset_status[akey] = { dstype: { 'A': False, 'W': False, 'P': False, 'S': False } for dstype in dstype_list }
    
    # for each record in the ArchiveMap, set the corresponding [archive_key][dataset_key]['A'] = True
    for _ in am_list:
        akey = tuple([_[1],_[2],_[3]])
        dkey = _[4]
        dataset_status[akey][dkey]['A'] = True
    
    ##########
    # walk the warehouse.  Need to assign "model,experiment,ensemble" to each extracted dataset (how to identify?)
    ##########

    warehouse_leaf_dirs = []
    for root, dirs, files in os.walk(warehouse):      # aggregate full sourcefile paths in src_selected
        # if not dirs and (src_selector in root):     # at leaf-directory matching src_selector
        if not dirs:     # at leaf-directory matching src_selector
            warehouse_leaf_dirs.append(root)

    wh_nonempty = []
    for adir in warehouse_leaf_dirs:
        # print(adir)
        ensdir, vleaf = os.path.split(adir)
        if not isVLeaf(vleaf):
            continue
        for root, dirs, files in os.walk(adir):
            if files:
                wh_nonempty.append( tuple([ensdir,vleaf,len(files)]))

    for atup in wh_nonempty:
        dsid = get_idval(atup[0]) # ensdir
        dsid = dsid.replace('/','.')
        # print(f'DEBUG_W: dsid = {dsid}')
        akey = get_dsid_arch_key( dsid )
        dkey = get_dsid_type_key( dsid )

        if not akey in dataset_status:  # first time we've encountered this experiment/case
            dataset_status[akey] = { dstype: { 'A': False, 'W': False, 'P': False, 'S': False } for dstype in dstype_list }
            if dkey in dataset_status[akey]:
                dataset_status[akey][dkey]['W'] = True
            else:
                dataset_status[akey][dkey] = { 'A': False, 'W': True, 'P': False, 'S': False }
        elif not dkey in dataset_status[akey]:  # first time we've encountered this dataset type for this experiment/case
            dataset_status[akey][dkey] = { 'A': False, 'W': True, 'P': False, 'S': False }
        else:  # just set Warehouse to True
            dataset_status[akey][dkey]['W'] = True


    ##########
    # walk the publicati.  Need to assign "model,experiment,ensemble" to each extracted dataset (how to identify?)
    ##########

    publicati_leaf_dirs = []
    for root, dirs, files in os.walk(publicati):      # aggregate full sourcefile paths in src_selected
        # if not dirs and (src_selector in root):     # at leaf-directory matching src_selector
        if not dirs:     # at leaf-directory matching src_selector
            publicati_leaf_dirs.append(root)

    pub_nonempty = []
    for adir in publicati_leaf_dirs:
        ensdir, vleaf = os.path.split(adir)
        if not isVLeaf(vleaf):
            continue
        for root, dirs, files in os.walk(adir):
            if files:
                pub_nonempty.append( tuple([ensdir,os.path.basename(ensdir),len(files)]))

    for atup in pub_nonempty:
        dsid = get_idval(atup[0]) # ensdir
        dsid = 'E3SM.' + dsid.replace(os.sep,'.')
        # print(f'DEBUG_P: dsid = {dsid}')
        akey = get_dsid_arch_key( dsid )
        dkey = get_dsid_type_key( dsid )

        if not akey in dataset_status:  # first time we've encountered this experiment/case
            dataset_status[akey] = { dstype: { 'A': False, 'W': False, 'P': False, 'S': False } for dstype in dstype_list }
            if dkey in dataset_status[akey]:
                dataset_status[akey][dkey]['P'] = True
            else:
                dataset_status[akey][dkey] = { 'A': False, 'W': False, 'P': True, 'S': False }
        elif not dkey in dataset_status[akey]:  # first time we've encountered this dataset type for this experiment/case
            dataset_status[akey][dkey] = { 'A': False, 'W': False, 'P': True, 'S': False }
        else:  # just set Publication to True
            dataset_status[akey][dkey]['P'] = True


    ##########
    # extract the datasetID from each published dataset listed in the ESGF_publication_report
    # obtain the key 'Model,Experiment,Ensemble'
    ##########

    contents = loadFileLines(esgf_pr)

    esgf_pr_list = [ aline.split(',') for aline in contents if aline[:-1] and not aline[0] == 'NOMATCH' ]
    esgf_pr_list = [ _[2] for _ in esgf_pr_list ]
    dsid_list = [ '.'.join(_.split('.')[:-1]) for _ in esgf_pr_list ]

    '''
    # DEBUG/INFO:  Report when any "published" dataset is not an archived dataset
    for id in dsid_list:
        print(f'DEBUG_S.1: id = {id}')
        akey = get_dsid_arch_key( id )
        if akey in dataset_status:
            print(f' key { akey } in dataset_status')
        else:
            print(f' key { akey } is NEW')
    '''

    # Update 'P' status of this (possibly new) dataset_type in the (possibly new) datasetID_key in the dataset_status table
    # using the dsid_list (sproket-based publication report)

    for dsid in dsid_list:
        # print(f'DEBUG_S: dsid = {dsid}')
        akey = get_dsid_arch_key( dsid )
        dkey = get_dsid_type_key( dsid )
        # print(f'DEBUG_S: akey = {akey}, dkey = {dkey}')
    
        if not akey in dataset_status:  # first time we've encountered this experiment/case
            dataset_status[akey] = { dstype: { 'A': False, 'W': False, 'P': False, 'S': False } for dstype in dstype_list }
            if dkey in dataset_status[akey]:
                dataset_status[akey][dkey]['S'] = True
            else:
                dataset_status[akey][dkey] = { 'A': False, 'W': False, 'P': False, 'S': True }
        elif not dkey in dataset_status[akey]:  # first time we've encountered this dataset type for this experiment/case
            dataset_status[akey][dkey] = { 'A': False, 'W': False, 'P': False, 'S': True }
        else:  # just set Publication to True
            dataset_status[akey][dkey]['S'] = True
    
    # print all combos (model,experiment,ensemble,dataset_type)

    if debug:
        for akey in dataset_status:
            # print(f'{ akey }:')
            for dkey in dataset_status[akey]:
                # print(f'    { dkey } : { dataset_status[akey][dkey] }')
                dataset_print_csv(akey,dkey)
        sys.exit(0)

    for akey in dataset_status:
        for dkey in dataset_status[akey]:
            statlist = ['_','_','_','_']
            if dataset_status[akey][dkey]['A']:
                statlist[0] = 'A'
            if dataset_status[akey][dkey]['W']:
                statlist[1] = 'W'
            if dataset_status[akey][dkey]['P']:
                statlist[2] = 'P'
            if dataset_status[akey][dkey]['S']:
                statlist[3] = 'S'
            statcode = ''.join(statlist)
            print(f'STATUS={statcode}:',end ='')
            dataset_print_csv(akey,dkey)

    sys.exit(0)

if __name__ == "__main__":
  sys.exit(main())

