#!/usr/bin/env python

"""
Step 9
======
In this step, we'll do everything from steps 1 through 8 - then,
when we have quality stats, we'll rename and annotate the HDAs containing the
forward and reverse reads.

The output will be in a new history named 'Step 9'.
"""
import os
import sys
import pprint
import time

import users_1
import histories_3
import tools_1
import workflows_1
import datasets_3

# we'll use the latest hdas module in order to get the update_hda function
#import hdas_2
import hdas_3


NEW_HISTORY_NAME = 'Step 9'
DATA_TO_UPLOAD = 'data/myIlluminaRun.solexa.fastq'

SOLEXA_QC_WORKFLOW_NAME = 'Joined Solexa QC'
SOLEXA_QC_WORKFLOW_INPUT_STEP = 6

STATISTICS_DATASET_NAME = 'statistics'
STATISTICS_DATASET_COLUMN = 5

# here's the name of the dataset containing the statistics - kinda silly in this case
FREAD_DATASET_NAME = 'forward reads'
RREAD_DATASET_NAME = 'reverse reads'

# ----------------------------------------------------------------------------- main
if __name__ == '__main__':
    try:
        # check the connection
        users = users_1.get_users()

        # create a new history
        new_history = histories_3.create_history( NEW_HISTORY_NAME )
        print 'created history!', new_history[ 'name' ]
        new_history_id = new_history[ 'id' ]
        new_history_details = histories_3.get_history( new_history_id )

        # upload a file
        tool_output_datasets = tools_1.upload_hda( new_history_id, DATA_TO_UPLOAD )
        uploaded_file_data = tool_output_datasets[ 'outputs' ][0]
        print 'uploaded hda!', uploaded_file_data[ 'name' ]
        uploaded_file_id = uploaded_file_data[ 'id' ]

        # and use get_hda to get details on the new, uploaded HDA
        uploaded_hda_details = hdas_3.get_hda( new_history_id, uploaded_file_id )
                               # ^^^ --------------------------------------- we're using a new version
                               #                                             - don't forget to change things like this
        uploaded_hda_state = uploaded_hda_details[ 'state' ]

        # wait for the upload to finish
        while uploaded_hda_state != 'ok':
            print '\t uploaded_hda_state:', uploaded_hda_state
            print '\t (waiting 4 seconds...)'
            time.sleep( 4.0 )

            # keep checking to get any new state the HDA might move into
            uploaded_hda_details = hdas_3.get_hda( new_history_id, uploaded_file_id )
            uploaded_hda_state = uploaded_hda_details[ 'state' ]

        # here's the new stuff...moving fast now

        # get the info of the all workflows available to us
        all_workflows = workflows_1.get_workflows()

        # let's search that info for the name of the one we want in the list of all workflows
        found_workflow = None
        for workflow in all_workflows:
            if workflow[ 'name' ] == SOLEXA_QC_WORKFLOW_NAME:
                found_workflow = workflow

        if not found_workflow:
            raise Exception( 'If you see this error, let one of the workhop presenters know' )

        print 'found workflow!', found_workflow[ 'name' ]
        target_workflow_id = found_workflow[ 'id' ]
        target_workflow_details = workflows_1.get_workflow( target_workflow_id )

        # now we'll run it:
        print 'running',  found_workflow[ 'name' ], 'workflow...'
        workflow_output = workflows_1.run_single_input_workflow_on_hda( target_workflow_id,
            new_history_id, uploaded_file_id, SOLEXA_QC_WORKFLOW_INPUT_STEP )
        print 'workflow started!'

        # the 'outputs' list of the workflow_output dictionary are the ids of the HDAs the workflow creates
        output_hda_ids = workflow_output[ 'outputs' ]

        # wait for them all to finish
        for hda_id in output_hda_ids:
            workflow_hda_details = hdas_3.get_hda( new_history_id, hda_id )
            workflow_hda_state = workflow_hda_details[ 'state' ]
            workflow_hda_name = workflow_hda_details[ 'name' ]
            print workflow_hda_name

            while workflow_hda_state != 'ok':
                print '\t state:', workflow_hda_state
                print '\t (waiting 4 seconds...)'
                time.sleep( 4.0 )

                # keep checking to get any new state the HDA might move into
                workflow_hda_details = hdas_3.get_hda( new_history_id, hda_id )
                workflow_hda_state = workflow_hda_details[ 'state' ]

            print '\t ok'
        print 'workflow complete!'

        # get some statistics using the datasets api
        found_statistics = None
        hda_summaries = hdas_3.get_hdas( new_history_id )
        for hda_summary in hda_summaries:
            if hda_summary[ 'name' ] == STATISTICS_DATASET_NAME:
                found_statistics = hda_summary
                break
        if not found_statistics:
            raise Exception( 'If you see this error, let one of the workhop presenters know' )
        statistics_dataset_id = found_statistics[ 'id' ]

        column_data = datasets_3.get_dataset_column( statistics_dataset_id, STATISTICS_DATASET_COLUMN )
        column_metadata = column_data[ 'meta' ][0]
        mean_perbase_quality = column_metadata[ 'mean' ]
        median_perbase_quality = column_metadata[ 'median' ]
        print 'mean:', mean_perbase_quality, 'median:', median_perbase_quality

        # now we'll rename and annotate the split forward and reverse reads
        # we'll need to find them first - familiar pattern
        forward_reads_hda = None
        reverse_reads_hda = None
        # we've already got a list from the statistics step above
        for hda_summary in hda_summaries:
            if   hda_summary[ 'name' ] == FREAD_DATASET_NAME:
                forward_reads_hda = hda_summary
            elif hda_summary[ 'name' ] == RREAD_DATASET_NAME:
                reverse_reads_hda = hda_summary
        if not forward_reads_hda or not reverse_reads_hda:
            raise Exception( 'If you see this error, let one of the workhop presenters know' )
        print 'found fwd/rev reads:', forward_reads_hda[ 'name' ], reverse_reads_hda[ 'name' ]

        # New stuff: let's update the names and annotations of each file using the update_hda API call

        # first, we'll use the first part of the file we uploaded for part of the new names
        # (this would be a good place to customize what you grab from the filename to make it the HDA name)
        uploaded_basename = os.path.basename( DATA_TO_UPLOAD )
        uploaded_prefix = uploaded_basename.split( '.' )[0]
        annotation_string  = 'mean perbase quality: ' + str( mean_perbase_quality ) + '; '
        annotation_string += 'median perbase quality: ' + str( median_perbase_quality ) + '; '

        forward_reads_hda_id = forward_reads_hda[ 'id' ]
        hdas_3.update_hda( new_history_id, forward_reads_hda_id, {
            'name'          : uploaded_prefix + '.fwd.fastqsanger',
            'annotation'    : annotation_string
        })

        reverse_reads_hda_id = reverse_reads_hda[ 'id' ]
        # the annotations are the same, but note the 'rev' in the name to indicate reverse reads
        hdas_3.update_hda( new_history_id, reverse_reads_hda_id, {
            'name'          : uploaded_prefix + '.rev.fastqsanger',
            'annotation'    : annotation_string
        })

        # let's make sure that worked by getting the details for these two HDAs
        forward_reads_hda_details = hdas_3.get_hda( new_history_id, forward_reads_hda_id )
        forward_reads_hda_name = forward_reads_hda_details[ 'name' ]
        reverse_reads_hda_details = hdas_3.get_hda( new_history_id, reverse_reads_hda_id )
        reverse_reads_hda_name = reverse_reads_hda_details[ 'name' ]
        # (currently the hdas.show API doesn't return the annotations: you'll have to double check
        #   that in the browser for now)
        
    except Exception, exc:
        print 'Error updating HDA names/annotations:', str( exc )
        sys.exit( 1 )

    print 'Forward:', forward_reads_hda_name
    print 'Reverse:', reverse_reads_hda_name