import argparse
import xarray as xr
import sys

############################################################################
# Script to extract data for each rbm id sites
############################################################################

## Helper function for formatting and saving data
def format_and_save(df, outpath):
    df = df[['temp']]
    df.to_csv(outpath)

def main():
    ## Parser
    parser = argparse.ArgumentParser()
    parser.add_argument('input_file', help="Sorted temp dataset")
    parser.add_argument('--filter', type=str,
                        help="Text with list of ids to extract. Default extracts all.")
    parser.add_argument('out_dir', nargs='?', default='./', help="Output file directory")
    args = parser.parse_args()
    
    ## Initialize values
    infile = args.input_file
    outdir = args.out_dir + "/"
    
    ds = xr.open_dataset(infile)
    flist = [ int(x) for x in args.filter.split(',') ]

    print("Segments to extract:", flist)
    for num in flist:
        print('--: ', num)
        seg = ds.sel(rbm=num).to_dataframe()
        out = "{}/seg{}.temp".format(outdir, num)
        format_and_save(seg, out)
    

if __name__ == "__main__":
    main()