#ENVDataPostProcessor.py
'''Gathers RIE_value data from any parts into one file, for EPI model to run.

This will be needed if any more parts-based calculations of environmental 
suitability are run without the ews_plotting routine.

Note that gather_RIE_values_v1_1() is only needed for dates missing from 
/storage/app/EWS/Ethiopia/Plotting/output .

gather_RIE_values_v1_3() is used in operational.
'''
import argparse
from glob import glob
import os
from pathlib import Path

import numpy as np
import pandas as pd

# Create the parser
my_parser = argparse.ArgumentParser(description='Different optional arguments')

# Add the arguments
my_parser.add_argument(
        '-inpath','--inpath',
        metavar = 'path',
        type = str,
        required = True,
        default = '/home/jws52/projects/py-coordination/ENVIRONMENT_20200115/env_suit_wheatrust_20200115/StripeRustOutput/',
        help = 'the directory path to all of the separate part calcs')

my_parser.add_argument(
        '-outpath','--outpath',
        metavar = 'path',
        type = str,
        required = True,
        default = '/home/jws52/projects/py-coordination/ENVIRONMENT_20200115/EPI/ENV',
        help = 'the directory path to gather separate part calcs')

my_parser.add_argument(
        '-rusttype','--rusttype',
        choices = ['stripe','stem','leaf'],
        type = str,
        required = True,
        default = 'stripe',
        help = 'the directory path to gather separate part calcs')

args = my_parser.parse_args()

print(f"input arguments are:\n{args}")

assert os.path.exists(args.inpath)

# make output directory
outdir = f'{args.outpath}/{args.rusttype.title()}Rust'

Path(outdir).mkdir(parents=True, exist_ok=True)

def gather_RIE_values_v1_1(envpath=args.inpath,outpath=args.outpath,rusttype=args.rusttype):

    envparts = sorted(glob(envpath+'part_*/*/RIE_value.csv'))

    pdparts = [pd.read_csv(fn) for fn in envparts]

    pdall = pd.concat(
            pdparts,
            ignore_index=False,
            keys=['part1','part2','part3','part4','part5','part6'])

    pdall.rename({'Unnamed: 0':'X'},axis=1,inplace=True)

    pdall.to_csv(f'{outpath}/{rusttype}/test_RIE_value.csv',index=False)   

    #df2 = read_csv('test_RIE_value.csv')

    return pdall

def gather_RIE_values_v1_3(envpath=args.inpath,outpath=args.outpath,rusttype=args.rusttype):

    envparts = sorted(glob(envpath+'Region*/*0000/RIE_value.csv'))

    pdparts = [pd.read_csv(fn) for fn in envparts]

    pdall = pd.concat(
            pdparts,
            ignore_index=False,
            keys=['part1','part2','part3','part4','part5','part6'])

    pdall.set_index('Unnamed: 0',inplace=True)

    print(pdall)

    pdall.index.name = None

    pdall.to_csv(f'{outdir}/RIE_value.csv')
    
    #df2 = read_csv('test_RIE_value.csv')

    return pdall


def test_case():
    inpath_default = '/home/jws52/projects/py-coordination/ENVIRONMENT_20200115/env_suit_wheatrust_20200115/StripeRustOutput/'
    outpath_default = '/home/jws52/projects/py-coordination/ENVIRONMENT_20200115/EPI/ENV'
    rusttype_default = 'stripe'

    df2 = gather_RIE_values(inpath_default,outpath_default,rusttype_default)

    # example case to reproduce
    fn1 = '/storage/app/EWS/Ethiopia/Workspace/ENVIRONMENT_20200115/EPI/ENVIRONMENT/Stripe/RIE_value.csv'
    df1 = pd.read_csv(fn1)
    
    print('testing')
    assert np.allclose(df1,df2)

if __name__ == '__main__':

    df2 = gather_RIE_values_v1_3(args.inpath,args.outpath,args.rusttype)

print("Finished!")