import logging
import nbformat
import sys
import yaml

def write_parameters(cell_source, params_yaml, outfp):
    with open(params_yaml, 'r') as ifp:
        y = yaml.safe_load(ifp)
        # print out all the lines in notebook
        write_code(cell_source, 'PARAMS from notebook', outfp)
        # print out YAML file; this will override definitions above
        formats = [
            '{} = {}', # for integers and floats
            '{} = "{}"', # for strings
        ]
        write_code(
            '\n'.join([
                formats[type(value) is str].format(key, value) for key, value in y.items()]),
            'PARAMS from YAML',
            outfp
        )

def write_code(cell_source, comment, outfp):
    lines = cell_source.split('\n')
    if len(lines) > 0 and lines[0].startswith('%%'):
        prefix = '#'
    else:
        prefix = ''
    
    print("### BEGIN {} ###".format(comment), file=outfp)
    for line in lines:
        line = prefix + line.replace('print(', 'logging.info(')
        if len(line) > 0 and (line[0] == '!' or line[0] == '%'):
            print('#' + line, file=outfp)
        else:
            print(line, file=outfp)
    print("### END {} ###\n".format(comment), file=outfp)
            
def convert_notebook(notebook_filename, params_yaml, outfp):
    write_code('import logging', 'code added by notebook conversion', outfp)
    with open(INPUT) as ifp:
        nb = nbformat.reads(ifp.read(), nbformat.NO_CONVERT)
        for cell in nb.cells:
            if cell.cell_type == 'code':
                if 'tags' in cell.metadata and 'display' in cell.metadata.tags:
                    logging.info('Ignoring cell # {} with display tag'.format(cell.execution_count))
                elif 'tags' in cell.metadata and 'parameters' in cell.metadata.tags:
                    logging.info('Writing params cell # {}'.format(cell.execution_count))
                    write_parameters(cell.source, PARAMS, outfp)
                else:
                    logging.info('Writing model cell # {}'.format(cell.execution_count))
                    write_code(cell.source, 'Cell #{}'.format(cell.execution_count), outfp)

import os

INPUT='../../06_feateng_keras/solution/taxifare_fc.ipynb'
PARAMS='./notebook_params.yaml'
OUTDIR='./container/trainer'

!mkdir -p $OUTDIR
OUTFILE=os.path.join(OUTDIR, 'model.py')
!touch $OUTDIR/__init__.py
with open(OUTFILE, 'w') as ofp:
    #convert_notebook(INPUT, PARAMS, sys.stdout)
    convert_notebook(INPUT, PARAMS, ofp)
#!cat $OUTFILE

"""
Explanation: Train ML model on Cloud AI Platform
This notebook shows how to:
* Export training code from a Keras notebook into a trainer file
* Create a Docker container based on a DLVM container
* Deploy training job to cluster
TODO: Export the data from BigQuery to GCS

Navigate to export_data.ipynb
Update 'your-gcs-project-here' to your GCP project name
Run all the notebook cells

TODO: Edit notebook parameters

Navigate to notebook_params.yaml
Replace the bucket name with your own bucket containing your model (likely gcp-project with -ml at the end)
Save the notebook
Return to this notebook and continue

Export code from notebook
This notebook extracts code from a notebook and creates a Python file suitable for use as model.py
End of explanation
"""


!python3 $OUTFILE

"""
Explanation: Try out model file
<b>Note</b> Once the training starts, Interrupt the Kernel (from the notebook ribbon bar above). Because it processes the entire dataset, this will take a long time on the relatively small machine on which you are running Notebooks.
End of explanation
"""


%%writefile container/Dockerfile
FROM gcr.io/deeplearning-platform-release/tf2-cpu

#RUN python3 -m pip install --upgrade --quiet tf-nightly-2.0-preview
RUN python3 -m pip install --upgrade --quiet cloudml-hypertune

COPY trainer /trainer
CMD ["python3", "/trainer/model.py"]

%%writefile container/push_docker.sh
export PROJECT_ID=$(gcloud config list project --format "value(core.project)")
export IMAGE_REPO_NAME=serverlessml_training_container
#export IMAGE_TAG=$(date +%Y%m%d_%H%M%S)
#export IMAGE_URI=gcr.io/$PROJECT_ID/$IMAGE_REPO_NAME:$IMAGE_TAG
export IMAGE_URI=gcr.io/$PROJECT_ID/$IMAGE_REPO_NAME

echo "Building  $IMAGE_URI"
docker build -f Dockerfile -t $IMAGE_URI ./
echo "Pushing $IMAGE_URI"
docker push $IMAGE_URI

!find container

"""
Explanation: Create Docker container
Package up the trainer file into a Docker container and submit the image.
End of explanation
"""


%%bash
cd container
bash push_docker.sh

"""
Explanation: <b>Note</b>: If you get a permissions error when running push_docker.sh from Notebooks, do it from CloudShell:
* Open CloudShell on the GCP Console
*  git clone https://github.com/GoogleCloudPlatform/training-data-analyst
*  cd training-data-analyst/quests/serverlessml/07_caip/solution/container
*  bash push_docker.sh
This next step takes 5 - 10 minutes to run
End of explanation
"""


%%bash
JOBID=serverlessml_$(date +%Y%m%d_%H%M%S)
REGION=us-central1
PROJECT_ID=$(gcloud config list project --format "value(core.project)")
BUCKET=$(gcloud config list project --format "value(core.project)")-ml

#IMAGE=gcr.io/deeplearning-platform-release/tf2-cpu
IMAGE=gcr.io/$PROJECT_ID/serverlessml_training_container

gcloud beta ai-platform jobs submit training $JOBID \
   --staging-bucket=gs://$BUCKET  --region=$REGION \
   --master-image-uri=$IMAGE \
   --master-machine-type=n1-standard-4 --scale-tier=CUSTOM

"""
Explanation: Deploy to AI Platform
Submit a training job using this custom container that we have just built. After you submit the job, monitor it here.
End of explanation
"""


