#!/usr/bin/env python
import argparse
import subprocess
import os
import sys
from yolov3.utils.logger import get_logger, info_level, debug_level

logger = get_logger(debug_level)

prefix = '/opt/ml/'
input_path = os.path.join(prefix,'input/data')
output_path = os.path.join(prefix, 'output/data')
model_path = os.path.join(prefix, 'model')
config_path = os.path.join(prefix, 'input/config')
checkpoint_path = os.path.join(prefix, 'checkpoints')

#weights_path = os.path.join(input_path, 'model')
weights_path = ''
cfg_path = os.path.join(prefix, 'code/yolov3/cfg/yolov3-tiny.cfg')
data_path = os.path.join(input_path, 'training/obj.data')


training_channel = 'training'
training_path = os.path.join(input_path, training_channel)

validating_channel = 'validating'
validating_path = os.path.join(input_path, validating_channel)

training_script = 'code/yolov3/train.py'
default_params = ['--cfg',str(cfg_path),'--data',str(data_path),'--weights',str(weights_path),'--input',str(input_path), '--model',str(model_path),'--checkpoints',str(checkpoint_path),'--epochs',str(2)]

def _run(cmd):
    logger.info('Running training script')
    process = subprocess.Popen(cmd, env=os.environ)
    stdout, stderr = process.communicate()

    return_code = process.poll()
    if return_code:
        error_msg = 'Return code: {}  Err: {}'.format(return_code, stderr)
        raise Exception(error_msg)


def show_dirs():
    logger.info('Container structure')
    for root, dirs, files in os.walk(prefix):
        for name in dirs:
            logger.info(os.path.join(root,name))

if __name__=='__main__':
    show_dirs()

    train_cmd = [sys.executable, training_script] + default_params
    _run(train_cmd)

    logger.info('Training completed')