#!/usr/bin/env python
# -*- coding: utf-8 -*-

from Bio import SeqIO
from Bio import Entrez

import io
import zlib
import random
import logging


import requests


def ncbi(kingdom, n_genomes, output):
    """download random genomes sequences from ncbi genomes with entrez eutils
        and requests.

    Args:
        kingdom (string): the kingdom from which the sequences are from
        n_genomes (int): the number of genomes to download

    Returns:
        str: the output file
    """
    logger = logging.getLogger(__name__)
    Entrez.email = 'hadrien.gourle@slu.se'
    Entrez.tool = 'InSilicoSeq'
    Entrez.api_key = 'd784b36672ca73601f4a19c3865775a17207'
    full_id_list = Entrez.read(Entrez.esearch(
        'assembly',
        term='%s[Organism] AND "latest refseq"[filter] AND "complete genome"[filter]'
        % kingdom, retmax=100000))['IdList']
    genomes = []
    n = 0
    logger.info('Searching for %s to download' % kingdom)
    while n < n_genomes:
        ident = random.choice(full_id_list)
        genome_info = Entrez.read(
            Entrez.esummary(
                db='assembly',
                id=ident))["DocumentSummarySet"]["DocumentSummary"][0]
        if genome_info['FtpPath_RefSeq']:
            url = genome_info['FtpPath_RefSeq']
            url = "%s/%s_%s_genomic.fna.gz" \
                % (genome_info['FtpPath_RefSeq'],
                   genome_info['AssemblyAccession'],
                   genome_info['AssemblyName'])
            logger.info('Downloading %s' % genome_info['AssemblyAccession'])
            assembly_to_fasta(url, output)
            n += 1
    return output


def assembly_to_fasta(url, output, chunk_size=1024):
    """download an assembly from the ncbi ftp and append
    the chromosome sequence to a fasta file.

    This function discards the plsamid sequences!

    Args:
        url (string): an url to a fasta file
        output (string): the output file name

    Returns:
        str: the file name
    """
    logger = logging.getLogger(__name__)
    if url.startswith("ftp://"):  # requests doesnt support ftp
        url = url.replace("ftp://", "https://")
    if url:
        request = requests.get(url)
        request = zlib.decompress(
            request.content, zlib.MAX_WBITS | 32).decode()

    with io.StringIO(request) as fasta_io:
        seq_handle = SeqIO.parse(fasta_io, 'fasta')
        chromosome = filter_plasmids(seq_handle)

        try:
            f = open(output, 'a')
        except (IOError, OSError) as e:
            logger.error('Failed to open output file: %s' % e)
            sys.exit(1)
        else:
            logger.debug('Writing genome to %s' % output)
            with f:
                SeqIO.write(chromosome, f, 'fasta')
    return output


def filter_plasmids(handle):
    """returns the largest sequence from a sequence handle
    """
    n = 0
    for record in handle:
        if len(record) > n:
            n = len(record)
            largest = record
    return largest

@@ -96,16 +96,39 @@
 port io%0A
+import sys%0Aimport time%0A
 import z
@@ -161,17 +161,16 @@
 ogging%0A%0A
-%0A
 import r
@@ -179,16 +179,233 @@
 uests%0A%0A%0A
+class BadRequestError(Exception):%0A    %22%22%22Exception to raise when a http request does not return 200%0A    %22%22%22%0A%0A    def __init__(self, url, status_code):%0A        super().__init__('%25s returned %25d' %25 (url, status_code))%0A%0A%0A
 def ncbi
@@ -1105,25 +1105,8 @@
 t'%5D%0A
-    genomes = %5B%5D%0A
     
@@ -1764,24 +1764,45 @@
 ccession'%5D)%0A
+            try:%0A    
             
@@ -1848,14 +1848,333 @@
     
-n += 1
+except BadRequestError as e:%0A                logger.debug('Could not download %25s' %25%0A                             genome_info%5B'AssemblyAccession'%5D)%0A                logger.debug('Skipping and waiting two seconds')%0A                time.sleep(2)%0A            else:%0A                n += 1%0A                full_id_list.remove(ident)
 %0A   
@@ -2748,16 +2748,59 @@
 et(url)%0A
+        if request.status_code == 200:%0A    
         
@@ -2826,16 +2826,20 @@
 mpress(%0A
+    
         
@@ -2888,16 +2888,90 @@
 decode()
+%0A        else:%0A            raise BadRequestError(url, request.status_code)
 %0A%0A    wi

check for response status code when querying ncbi