#!/usr/bin/env python
#
# This takes input about blast query sequences and the database and produces a variable called
# pipeline.WANT_EXEC_INSTANCES which can be fed into resize-cluster metric
#
# The input data currently needed is:
# misc.PROGRAM (which blast program)
# params.MAX_QUERY_SEQ_LEN
# params.MIN_QUERY_SEQ_LEN
# params.AVG_QUERY_SEQ_LEN
# params.NUM_QUERY_SEQ
# params.SCALE_FACTOR - How much to scale the total number of queries, can be a float
#                      defaults to 1 if not specified
import sys

#
# Some blast programs take the query and will translate it in multiple ways
# so this is the amount of work they add to the computation
WORK_FACTOR = {'blastn': 1,
               'blastp': 1,
               'blastx': 6,
               'tblastn': 6,
               'tblastx': 12
               }

# A very rough estimate of how many residues can be calculated per core per hour is 28k
# Since we are giving that a work factor of 6 we are going to estimate the blastn/blastp
# baseline for that is 28000 * WORK_FACTOR['blastx']
RESIDUES_PER_CORE_PER_HOUR = 28000 * WORK_FACTOR['blastx']

def calculateCPUHours(program, numSeqs, avgLen, scale):
    """
    + 1 because we always high-ball by an hour
    /scaleDown because we allow fo rthe total number of hours to be brought down by factor if the input DB
    is smaller than NR
    """
    return (WORK_FACTOR[program] * float(numSeqs * avgLen)/float(RESIDUES_PER_CORE_PER_HOUR)) * scale + 1


# Ensure we are dealing with kv pairs
if sys.stdin.readline().strip() != 'kv':
    raise Exception('Header needs to be kv')

# Write our header out
sys.stdout.write('kv\n')

kv = {'params.PROGRAM': None,
      'params.MAX_QUERY_SEQ_LEN': None,
      'params.MIN_QUERY_SEQ_LEN': None,
      'params.AVG_QUERY_SEQ_LEN': None,
      'params.NUM_QUERY_SEQ': None,
      'params.SCALE_FACTOR': 1,
      }

for line in sys.stdin:
    sys.stdout.write(line.strip() + '\n')
    k, v = line.strip().split('=', 1)
    if k in kv:
        kv[k] = v

if [v for v in kv.values() if not v]:
    raise Exception('Not all expected values present in config')

if kv['params.PROGRAM'] not in WORK_FACTOR:
    raise Exception('Unknown blast program ' + kv['params.PROGRAM'])

for i in ['params.MAX_QUERY_SEQ_LEN',
          'params.MIN_QUERY_SEQ_LEN',
          'params.NUM_QUERY_SEQ',]:
    kv[i] = int(kv[i])

for i in ['params.SCALE_FACTOR',
          'params.AVG_QUERY_SEQ_LEN',
          ]:
    kv[i] = float(kv[i])


#
# Insert magic calculuations here
totalHours = calculateCPUHours(kv['params.PROGRAM'],
                               kv['params.NUM_QUERY_SEQ'],
                               kv['params.AVG_QUERY_SEQ_LEN'],
                               kv['params.SCALE_FACTOR'])

sys.stdout.write('\n'.join(['pipeline.COMPUTED_CPU_HOURS=' + str(totalHours)]) + '\n')
