import sagemaker
import boto3
from sagemaker.processing import ScriptProcessor, ProcessingInput, ProcessingOutput, FrameworkProcessor
from sagemaker import get_execution_role

role = get_execution_role()
bucket = 'dev-cucumbers'
sagemaker_session = sagemaker.session.Session(default_bucket = bucket)

script_processor = ScriptProcessor(
    image_uri='683313688378.dkr.ecr.us-east-1.amazonaws.com/sagemaker-scikit-learn:0.20.0-cpu-py3',
    sagemaker_session=sagemaker_session,
    role=role,
    instance_count=1,
    instance_type='ml.t3.2xlarge',
    command=['python3'],
    base_job_name="ProcessingJob"
)

# script_processor = FrameworkProcessor(
#     image_uri='683313688378.dkr.ecr.us-east-1.amazonaws.com/sagemaker-scikit-learn:0.20.0-cpu-py3',
#     sagemaker_session=sagemaker_session,
#     role=role,
#     instance_count=1,
#     instance_type="ml.t3.2xlarge",
#     command=['python3'],
#     #source_dir = "/home/ec2-user/SageMaker/ml-collab/eimpara/Moments/2023_data/requirements.txt"
#     source_dir = "/home/ec2-user/SageMaker/ml-collab/eimpara/Moments/2023_data/"
# )


destination = '/opt/ml/processing/input'

#code='/home/ec2-user/SageMaker/ml-collab/eimpara/Moments/2023_data/Fourier_and_ARIMA.ipynb'
#code='/home/ec2-user/SageMaker/ml-collab/eimpara/Moments/2023_data/test.ipynb'
code='/home/ec2-user/SageMaker/ml-collab/eimpara/Moments/2023_data/V3_S3_prediction_table.py'
#code='/home/ec2-user/SageMaker/ml-collab/eimpara/Moments/2023_data/ANNOTATED_S3_Fourier_ARIMA.py'
#code='/home/ec2-user/SageMaker/ml-collab/eimpara/Moments/2023_data/ANNOTATES_Fourier_for_TESTING.py'


script_processor.run(
    code=code,
    inputs=[
        ProcessingInput(
            source='s3://dev-cucumbers/eimpara/Fourier/2023_data/', # also tried with empty string as the data source for Fourier_and_ARIMA.ipynb is SnowFlake and called in the jupyter notebook
            destination=destination
        ),
        ProcessingInput(
            source='s3://dev-cucumbers/eimpara/requirements.txt',
            destination = '/opt/ml/processing/input/dependencies'
        )
    ],
    outputs=[
        ProcessingOutput(
            source='/opt/ml/processing/output',  # also tried with empty string
            destination=destination,
            output_name="batch_moments_nov2023",
        ),
    ]
)