import sagemaker
import logging
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.c5.18xlarge',    # 'ml.t3.medium',  #'ml.c5.4xlarge',  'ml.c5.18xlarge',
    command=['python3'],
    base_job_name="ProcessingJob"
)

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

#         today = datetime.today().strftime('%Y%m%d-%H%M%S')
run_id = '202406030927'
#run_id = '202401091205'
#run_id = 'test_run_to_debug'
CHUNK_COUNT=1
FOLDER = 'eimpara/pipelines/moments'

script1='/home/ec2-user/SageMaker/ml-collab/eimpara/Moments/ML_pipeline/Step1_savetable_S3.py'
script2='/home/ec2-user/SageMaker/ml-collab/eimpara/Moments/ML_pipeline/Step2_Fourier_ARIMA.py'
script3='/home/ec2-user/SageMaker/ml-collab/eimpara/Moments/ML_pipeline/Step3_merge_ARIMA_tables.py'
script4='/home/ec2-user/SageMaker/ml-collab/eimpara/Moments/ML_pipeline/Step3bis_merge_ARIMA_tables_timeseries.py'
script5='/home/ec2-user/SageMaker/ml-collab/eimpara/Moments/ML_pipeline/Step4_Manual_Clustering.py'

logging.info(f"Running: {script1}")
script_processor.run(
        code=script1,
        arguments = ["--run-id", str(run_id),
                    '--folder', FOLDER],
        inputs=[
            ProcessingInput(
                source='s3://dev-cucumbers/eimpara/requirements-2.txt',
                destination = '/opt/ml/processing/input/dependencies'
            )
        ],
        outputs=[
            ProcessingOutput(
                source='/opt/ml/processing/output',  
                destination=destination,
                output_name="Master_output",
            ),
        ]
    )
logging.info(f"Running: {script2}")
for chunk_id in range(CHUNK_COUNT):
    script_processor.run(
            code=script2,
            arguments = ['--dry-run', 'False', 
                         '--run-id', str(run_id), 
                         '--folder', FOLDER,
                         '--chunk-count', str(CHUNK_COUNT), 
                         '--chunk-id', str(chunk_id),
                         #'--parallelism', '1'
                        ],
            inputs=[
                ProcessingInput(
                    source='s3://dev-cucumbers/eimpara/requirements-2.txt',
                    destination = '/opt/ml/processing/input/dependencies'
                )
            ],
            outputs=[
                ProcessingOutput(
                    source='/opt/ml/processing/output', 
                    destination=destination,
                    output_name="Master_output",
                ),
            ]
        ) 
logging.info(f"Running: {script3}")
script_processor.run(
        code=script3,
        arguments = ["--run-id", str(run_id),
                    '--folder', FOLDER],
        inputs=[
            ProcessingInput(
                source='s3://dev-cucumbers/eimpara/requirements-2.txt',
                destination = '/opt/ml/processing/input/dependencies'
            )
        ],
        outputs=[
            ProcessingOutput(
                source='/opt/ml/processing/output',  
                destination=destination,
                output_name="Master_output",
            ),
        ]
    )
logging.info(f"Running: {script4}")
script_processor.run(
        code=script4,
        arguments = ["--run-id", str(run_id),
                    '--folder', FOLDER],
        inputs=[
            ProcessingInput(
                source='s3://dev-cucumbers/eimpara/requirements-2.txt',
                destination = '/opt/ml/processing/input/dependencies'
            )
        ],
        outputs=[
            ProcessingOutput(
                source='/opt/ml/processing/output',  
                destination=destination,
                output_name="Master_output",
            ),
        ]
    )
logging.info(f"Running: {script5}")
script_processor.run(
        code=script5,
        arguments = ["--run-id", str(run_id),
                    '--folder', FOLDER],
        inputs=[
            ProcessingInput(
                source='s3://dev-cucumbers/eimpara/requirements-2.txt',
                destination = '/opt/ml/processing/input/dependencies'
            )
        ],
        outputs=[
            ProcessingOutput(
                source='/opt/ml/processing/output',
                destination=destination,
                output_name="Master_output",
            ),
        ]
    )