from sagemaker.huggingface import HuggingFaceProcessor
from sagemaker.processing import ProcessingInput, ProcessingOutput
import os
import glob
import boto3
s3 = boto3.resource('s3')


#http://10.0.0.8:5000/company_a/-Users-bennicholl-Desktop-ml_apps_for_company-SM-data-data/0/[2,3,4,5]/5e-7/20/35/0.1/electra-base-discriminator


def main(company_name, s3_paths, input_indices, target_indices, learning_rate, batch_size, epochs, percent_of_testing_examples, model_size):  
    """
    s3_paths = []
    
    s3 = boto3.resource('s3')
    
    # use glob to get all the csv files 
    # in the folder
    csv_files = glob.glob(os.path.join(folder_path, "*.csv"))

    count = 0
    # loop over the list of csv files
    for f in csv_files:
        
        s3.meta.client.upload_file(f, 'training-stage', '{}/data{}.csv'.format(company_name, str(count)) )
        s3_path = "s3://training-stage/" + company_name + '/' + 'data' + str(count) + ".csv"
        s3_paths.append(s3_path)
        #print(f)
        count += 1
       
    #s3://training-stage/some_other_company/data0.csv  
    """
    
    
    
    role = 'arn:aws:iam::681897892690:role/service-role/AmazonSageMaker-ExecutionRole-20220605T131967'

    hfp = HuggingFaceProcessor(
        role=role, 
        instance_count=1,
        instance_type='ml.g4dn.2xlarge',
        transformers_version='4.17.0',
        pytorch_version='1.10.2', 
        base_job_name='frameworkprocessor-hf',
        py_version = 'py38'
    )

    

    hfp.run(
        code='train_process.py',
        source_dir='./train_transformer_source_dir',
        
        arguments=[
            '--s3_paths', str(s3_paths),
            '--learning_rate', learning_rate,
            '--batch_size', batch_size,
            '--epochs', epochs,
            #'--subtract_from_min_count', '1000',
            '--percent_of_testing_examples', percent_of_testing_examples,
            '--model_size', model_size,
            '--input_indices', input_indices,
            '--target_indices', target_indices]
        #inputs=[
        #    ProcessingInput(
         #       input_name='data',
         #       source=f's3://{BUCKET}/{S3_INPUT_PATH}',
         #       destination='/opt/ml/processing/input'
         #   )
        #],
        #outputs=[
        #    ProcessingOutput(output_name='data_structured', source='/opt/ml/processing/tmp/data_structured', destination=f's3://{BUCKET}/{S3_OUTPUT_PATH}'),
        #    ProcessingOutput(output_name='train', source='/opt/ml/processing/output/train', destination=f's3://{BUCKET}/{S3_OUTPUT_PATH}'),
        #    ProcessingOutput(output_name='validation', source='/opt/ml/processing/output/val', destination=f's3://{BUCKET}/{S3_OUTPUT_PATH}'),
        #    ProcessingOutput(output_name='test', source='/opt/ml/processing/output/test', destination=f's3://{BUCKET}/{S3_OUTPUT_PATH}'),
        #    ProcessingOutput(output_name='logs', source='/opt/ml/processing/logs', destination=f's3://{BUCKET}/{S3_OUTPUT_PATH}')
        #]
    )
        
    return company_name


main('company_a', ['s3://training-stage/company_a/data0.csv'], '0', '[2,3,4,5]', '5e-7', '20', '35', '0.1', 'electra-base-discriminator')