import gcp.bigquery as bq

"""
Explanation: Creating TCGA cohorts  (part 3)
This notebook will construct a cohort for a single tumor type based on data availability, while also taking into consideration annotations about the patients or samples.
As you've seen already, in order to work with BigQuery, the first thing we need to do is import the bigquery module:
End of explanation
"""


studyName = "BRCA"

"""
Explanation: Just so that this doesn't get buried in the code below, we are going to specify our tumor-type of interest here.  In TCGA each tumor-type is also a separate study within the TCGA project.  The studies are referred to based on the 2-4 letter tumor-type abbreviation.  A complete list of all study abbreviations, with the full study name can be found in this report.  For this particular exercise, we will look at the "Breast invasive carcinoma" study, abbreviated BRCA:
End of explanation
"""


d = bq.DataSet('isb-cgc:tcga_201607_beta')
for t in d.tables():
  print '%10d rows  %12d bytes   %s' \
      % (t.metadata.rows, t.metadata.size, t.name.table_id)

"""
Explanation: In this notebook, we are going to wind up making use of all of the available data types, so let's have a look at the entire tcga_201607_beta dataset:
End of explanation
"""


%%sql --module get_DNU_patients

SELECT
  ParticipantBarcode,
  annotationCategoryName AS categoryName,
  annotationClassification AS classificationName
FROM
  [isb-cgc:tcga_201607_beta.Annotations]
WHERE
  ( itemTypeName="Patient"
    AND (annotationCategoryName="History of unacceptable prior treatment related to a prior/other malignancy"
      OR annotationClassification="Redaction" ) )
GROUP BY
  ParticipantBarcode,
  categoryName,
  classificationName
ORDER BY
  ParticipantBarcode

"""
Explanation: In this next code cell, we define an SQL query called get_DNU_patients which finds all patients in the Annotations table which have either been 'redacted' or had 'unacceptable prior treatment'.
End of explanation
"""


DNUpatients = bq.Query(get_DNU_patients).results().to_dataframe()
DNUpatients.describe()

"""
Explanation: Now we'll use the query defined above to get the "Do Not Use" list of participants (aka patients):
End of explanation
"""


%%sql --module get_barcode_list

# this is a very general-purpose query which will return all unique barcodes of a specified type,
# from the specified study, and from the specified BigQuery table
SELECT
  $barcodeType,
FROM
  $tableName
WHERE
  ( Study=$studyName )
GROUP BY
  $barcodeType
ORDER BY
  $barcodeType

"""
Explanation: In this next cell, we'll define a general-purpose query that we can use on each of the tables to get the unique list of barcodes (either patient- or sample-barcodes) that appear in that particular table.
End of explanation
"""


def buildCleanBarcodeList ( studyName, bqDataset, DNUlist, barcodeType ):

  print " in buildCleanBarcodeList ... ", studyName

  # we're going to loop over the main input molecular data tables (choosing just
  # one of the per-chromosome Methylation tables to avoid scanning all of those
  # large tables unnecessarily, and excluding the RPPA data since only a subset
  # of all tumor samples were assayed on that platform) and assemble lists of 
  # unique barcodes -- at the end of the loop, Ulists will be a list of lists
  Ulists = []
  print "     --> looping over data tables: "
  for t in bqDataset.tables():
    if ( t.name.table_id in ["Copy_Number_segments", "DNA_Methylation_chr17", 
                            "miRNA_Expression", "mRNA_BCGSC_RPKM", "mRNA_UNC_RSEM",
                            "Somatic_Mutation_calls"] ):
      currTable = t
      try:
        barcodeField = t.schema[barcodeType]
        q = bq.Query(get_barcode_list,
                     barcodeType=barcodeField,
                     tableName=currTable,
                     studyName=studyName)
      except:
        print "          the %s table does not have a field called %s " % ( t.name.table_id, barcodeType )
        if ( barcodeType == "SampleBarcode" ):
          try:
            barcodeField = t.schema['Tumor_SampleBarcode']
            print "              using Tumor_SampleBarcode instead "
          except:
            continue
      
      aList = q.results().to_dataframe()
      #print t.name.table_id
      #print len(aList)
      if ( len(aList) > 0 ):
        print "         ", t.name.table_id, " --> ", len(aList[barcodeType]), " unique barcodes"
        Ulists += [ aList ]
      else:
        print "         ", t.name.table_id, " -->  no data returned"
    
  print "     --> we have %d lists to merge " % len(Ulists)
  
  # concatenate those separate lists into a single list of unique barcodes
  masterList = []
  for aList in Ulists:
    for aBarcode in aList[barcodeType]:
      if ( aBarcode not in masterList ):
        masterList += [ aBarcode ]
      
  print "     --> which results in a single list with %d barcodes " % len(masterList)
  
  print "     --> removing DNU barcodes: "
  # create a clean list by removing the DNU barcodes
  cleanList = []
  for aBarcode in masterList:
    if ( aBarcode not in DNUlist[barcodeType].tolist() ):
      cleanList += [ aBarcode ]
    else:
      print "         excluding this barcode: ", aBarcode
    
  print "     --> returning a clean list with %d barcodes " % len(cleanList)
  
  # and return it
  return (cleanList)

barcodeType = "ParticipantBarcode"
cleanPatientList = buildCleanBarcodeList ( studyName, d, DNUpatients, barcodeType )

"""
Explanation: Now we're gong to use the query defined previously in a function that builds a "clean" list of patients in the specified study, with available molecular data, and without any disqualifying annotations.
End of explanation
"""


%%sql --module get_DNU_samples

# there are many different types of annotations that are at the "sample" level
# in the Annotations table, and most of them seem like they should be "disqualifying"
# annotations, so for now we will just return all sample barcodes with sample-level
# annotations
DEFINE QUERY select_on_annotations
SELECT
  SampleBarcode,
FROM
  [isb-cgc:tcga_201607_beta.Annotations]
WHERE
  ( itemTypeName="Sample" )
GROUP BY
  SampleBarcode,
ORDER BY
  SampleBarcode

DNUsamples = bq.Query(get_DNU_samples).results().to_dataframe()
DNUsamples.describe()

"""
Explanation: Now we are going to repeat the same process, but at the sample barcode level.  Most patients will have provided two samples, a "primary tumor" sample, and a "normal blood" sample, but in some cases additional or different types of samples may have been provided, and sample-level annotations may exist that should result in samples being excluded from most downstream analyses.
End of explanation
"""


barcodeType = "SampleBarcode"
cleanSampleList = buildCleanBarcodeList ( studyName, d, DNUsamples, barcodeType )

"""
Explanation: And now we can re-use the previously defined function get a clean list of sample-level barcodes:
End of explanation
"""


finalSampleList = []
for aSample in cleanSampleList:
  aPatient = aSample[:12]
  if ( aPatient not in cleanPatientList ):
    print "     excluding this sample in the final pass: ", aSample
  else:
    finalSampleList += [aSample]
    
print " Length of final sample list: %d " % len(finalSampleList)

finalPatientList = []
for aSample in finalSampleList:
  aPatient = aSample[:12]
  if ( aPatient not in finalPatientList ):
    finalPatientList += [ aPatient ]
    
print " Lenth of final patient list: %d " % len(finalPatientList)

for aPatient in cleanPatientList:
  if ( aPatient not in finalPatientList ):
    print "     --> patient removed in final pass: ", aPatient

"""
Explanation: Now we're going to double-check first that we keep only sample-level barcodes that correspond to patients in the "clean" list of patients, and then we'll also filter the list of patients in case there are patients with no samples remaining in the "clean" list of samples.
End of explanation
"""


sampleCounts = {}
for aSample in finalSampleList:
  sType = str(aSample[13:15])
  if ( sType not in sampleCounts ): sampleCounts[sType] = 0
  sampleCounts[sType] += 1
  
for aKey in sorted(sampleCounts):
  print "     %5d samples of type %s " % ( sampleCounts[aKey], aKey )

"""
Explanation: We're also interested in knowing what types of samples we have.  The codes for the major types of samples are:
- 01 : primary solid tumor
- 02 : recurrent solid tumor
- 03 : primary blood derived cancer
- 06 : metastatic
- 10 : blood derived normal
- 11 : solid tissue normal
and a complete list of all sample type codes and their definitions can be found in the Sample type report.
End of explanation
"""


import pandas as pd

patientBarcodes = []
sampleBarcodes = []
for aSample in finalSampleList:
  sampleBarcodes += [aSample]
  patientBarcodes += [aSample[:12]]
df = pd.DataFrame ( { 'ParticipantBarcode': patientBarcodes,
                      'SampleBarcode': sampleBarcodes } )
df.describe()

# we can define the schema directly from the dataframe
schema = bq.Schema.from_dataframe(df)
# and then we create the table and load the data
d = bq.DataSet('isb-cgc:tcga_cohorts')
t = bq.Table('isb-cgc:tcga_cohorts.'+studyName).create(schema=schema, overwrite=True)
t.insert_data(df)
# create a descriptionString and update the table
dString = "Curated cohort table for TCGA %s study:  %d unique patients and %d unique samples." % \
    ( studyName, len(finalPatientList), len(finalSampleList) )
dString += "\nJOIN this table with the molecular data tables in the isb-cgc:tcga_201607_beta"
dString += " BigQuery dataset for additional analyses."
t.update(description=dString)

t

cohorts_bucket_object = 'gs://isb-cgc-open/cohorts/tcga/' + studyName + '.tsv'

## NOTE: the extract command should work to export this table to a CSV file, but seems to not be working
## because the table thinks it has 0 rows at this point.
## t.extract(destination=cohorts_bucket_object, format='csv', csv_delimiter='\t')
print " The cohort table just created has %d rows." % ( t.metadata.rows )
if ( t.metadata.rows == 0 ): print " ??? "

"""
Explanation: Now we are going to create a simple dataframe with all of the sample barcodes and the associated patient (participant) barcodes so that we can write this to a BigQuery "cohort" table.
End of explanation
"""


t.to_file(destination='/tmp/test.csv', format='csv', csv_delimiter='\t', csv_header=True)

lines = None
with open('/tmp/test.csv') as datafile:
  lines = datafile.readlines()
outFile = ''.join(lines)
print outFile

%storage  write -v outFile -o $cohorts_bucket_object

"""
Explanation: So instead we will write to a temporary CSV file, read it back in and write to cloud storage using the %storage magic.
End of explanation
"""


