# Generated by Django 2.0.13 on 2019-12-18 09:34

from django.core.management.base import BaseCommand, CommandError

from django.db import migrations
from toxsign.genes.models import Gene, Pathway
from django.db.models import Q

import os, re, requests, shutil
from urllib.request import urlopen
from zipfile import ZipFile
import xml.etree.ElementTree as ET

from bs4 import BeautifulSoup


def cleanup(dirpath, specie, release):

    current_archive_files = [f for f in os.listdir(dirpath) if re.match(r'{}-[0-9]+\.zip'.format(specie), f)]
    # Missing latest release: cleanup other archives files, recreate dir
    if '{}-{}.zip'.format(specie, release) not in current_archive_file:
        for archive in current_archive_files:
            os.remove(dirpath + archive)
        if os.path.exists(dirpath + specie):
            shutil.rmtree(dirpath + specie)
        os.mkdir(dirpath + specie)


def download_specie_file(specie, url_dict):
    specie_name = specie.replace(' ', '_')
    release = "20191210"
    url = "http://data.wikipathways.org/20191210/gpml/wikipathways-20191210-gpml-{}.zip".format(specie_name)
    
    if specie_name in url_dict:
        release = url_dict[specie_name]['release']
        url = url_dict[specie_name]['url']

    dirpath = "/app/toxsign/media/jobs/admin/pathways/"
    archive_file = dirpath + "{}-{}.zip".format(specie_name, release)
    # If no file, download
    if not os.path.exists(archive_file):
        u = urlopen(url)
        f = open(archive_file, 'wb')
        meta = u.info()
        file_size = int(meta.get("Content-Length")[0])
        print("Downloading: %s"% (specie_name + ".zip"))
        file_size_dl = 0
        block_sz = 8192
        while True:
            buffer = u.read(block_sz)
            if not buffer:
                break

            file_size_dl += len(buffer)
            f.write(buffer)
        f.close()
        with ZipFile(archive_file, 'r') as zipObj:
            zipObj.extractall(path=dirpath + specie_name + "/")
    else:
        print("Skipping download for " + archive_file)
    return dirpath + specie_name + "/"

def process_pathway_file(path):

    it = ET.iterparse(path)
    for _, el in it:
        prefix, has_namespace, postfix = el.tag.partition('}')
        if has_namespace:
            el.tag = postfix  # strip all namespaces
    root = it.root

    content = {
        'id': re.search(r'WP\d+', path).group(),
        'name': root.attrib.get('Name', ''),
        'organism': root.attrib.get('Organism', ''),
        'entrez': set(),
        'ensembl': set()
    }

    for child in root:
        if child.tag == "DataNode" and child.attrib.get('Type') == 'GeneProduct':
            info = child.find('Xref')
            if info is not None and info.attrib.get('Database') and info.attrib.get('ID'):
                if info.attrib.get('Database') == "Entrez Gene":
                    content['entrez'].add(info.attrib.get('ID'))
                elif info.attrib.get('Database') == "Ensembl":
                    content['ensembl'].add(info.attrib.get('ID'))
    return content

def create_pathway(content):

    pathway = Pathway(pathway_id=content['id'], name=content['name'], organism=content['organism'])
    pathway.save()
    genes = Gene.objects.filter(Q(gene_id__in=content['entrez']) | Q(ensembl_id__in=content['ensembl']))
    pathway.genes.add(*genes)

def get_urls(url, base_url_layout):
    page = requests.get(url)
    soup = BeautifulSoup(page.text, 'html.parser')
    urls = set([link.get('href') for link in soup.find('table').findAll('a')])
    url_dict = {}
    for url in urls:
        search = re.search(r"wikipathways-(\d+)-gpml-(.+)\.zip", url)
        url_dict[search.group(2)] = {'url': base_url_layout.format(search.group(1), search.group(2)), 'release': search.group(1)}
    return url_dict

def process_pathways():

    # Check is genes are loaded:

    gene = Gene.objects.filter(id=1)
    if not gene:
        print('No genes are loaded. Did you run load_genes first?')
        return

    #fallback
    base_url_layout = "http://data.wikipathways.org/current/gpml/wikipathways-{}-gpml-{}.zip"
    url_list_url = "http://data.wikipathways.org/current/gpml/"

    url_dict = get_urls(url_list_url, base_url_layout)

    species = ['Bos taurus', 'Canis familiaris', 'Danio rerio', 'Equus caballus', 'Gallus gallus', 'Homo sapiens', 'Mus musculus', 'Pan troglodytes', 'Rattus norvegicus', 'Sus scrofa', 'Drosophila melanogaster']
    for specie in species:
        print("Processing " + specie)
        dir_path = download_specie_file(specie, url_dict)
        for pathway_file in os.listdir(dir_path):
            content = process_pathway_file(dir_path + pathway_file)
            create_pathway(content)

class Command(BaseCommand):
    help = 'Load pathways'

    def handle(self, *args, **options):
        process_pathways()
