# Copyright (c) 2014. Mount Sinai School of Medicine
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import urllib2
import urllib
from StringIO import StringIO
import logging

import pandas as pd

from pipeline import PipelineElement

class IEDBMHCBinding(PipelineElement):

  def __init__(self, alleles=[], name="IEDB-MHC-Binding", method='recommended', lengths = [9,10,11], url='http://tools.iedb.org/tools_api/mhci/'):
    self.name = name
    self._method = method
    self._lengths = lengths
    self._url = url
    self._alleles = ",".join(alleles)

  def _get_iedb_request_params(self, sequence):
    params = {
        "method" : self._method,
        "length" : ",".join(str(l) for l in self._lengths),
        "sequence_text" : sequence,
        "allele" : self._alleles,
    }
    return params

  def query_iedb(self, sequence):
    request_values = self._get_iedb_request_params(sequence)
    logging.info("Calling iedb with {}, {}".format(sequence, self._alleles))
    try:
      data = urllib.urlencode(request_values)
      req = urllib2.Request(self._url, data)
      response = urllib2.urlopen(req).read()

      return pd.read_csv(StringIO(response), sep='\t', na_values=['-'])
    except:
      logging.error("Connection error: Failed on sequence {}".format(sequence))
      return pd.DataFrame()

  def apply(self,data):
    responses = {}
    for epitope in data.Peptide:
       responses[epitope] = self.query_iedb(epitope)
    responses = pd.concat(responses).reset_index(0)
    responses.rename(columns={'level_0':'Peptide'}, inplace=True)
    return data.merge(responses, on='Peptide')

class IEDBMHC1Binding(IEDBMHCBinding):
    def __init__(self,
            name = 'IEDB-MHC1-Binding',
            url='http://tools.iedb.org/tools_api/mhci/',
            alleles=[]):
        super(IEDBMHC1Binding, self).__init__(
            name = name,
            url = url,
            alleles = alleles)

class IEDBMHC2Binding(IEDBMHCBinding):
    def __init__(self,
            name = 'IEDB-MHC2-Binding',
            url='http://tools.iedb.org/tools_api/mhcii/',
            alleles=[]):
      super(IEDBMHC2Binding, self).__init__(
        name = name,
        url = url,
        alleles = alleles)

    def _get_iedb_request_params(self, sequence):
      params = {
        "method" : self._method,
        "sequence_text" : sequence,
        "allele" : self._alleles,
      }
      return params

@@ -1699,18 +1699,64 @@
   except
-:%0A
+ KeyboardInterrupt:%0A        raise%0A    except:%0A  
       lo
@@ -1767,16 +1767,29 @@
 g.error(
+%0A            
 %22Connect
@@ -1842,24 +1842,26 @@
 nce))%0A      
+ 
+ 
 return pd.Da
@@ -1871,16 +1871,17 @@
 rame()%0A%0A
+%0A
   def ap
@@ -1923,22 +1923,22 @@
     for 
+p
 ep
-itop
+tid
 e in dat
@@ -1948,16 +1948,57 @@
 eptide:%0A
+        if peptide not in response:%0A     
        r
@@ -2006,22 +2006,22 @@
 sponses%5B
+p
 ep
-itop
+tid
 e%5D = sel
@@ -2037,14 +2037,147 @@
 edb(
+p
 ep
-itop
+tide)%0A        else:%0A            logging.info(%0A                %22Skipping binding for peptide %25s, already queried%22,%0A                peptid
 e)%0A 

handle keyboard interrupt for iedb calls, don't repeat peptides