#  Copyright 2019 Stian Soiland-Reyes
#
#  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.

from snakemake.utils import min_version
min_version("5.8")

workdir: "../../../../data/"

import os
import string
import itertools
from snakemake.remote.HTTP import RemoteProvider as HTTPRemoteProvider

HTTP = HTTPRemoteProvider()

TYPES = ["software", "dataset", "others"]

MAXSAMPLES=2000
def samples():
    """Predict GNU coreutil split's alphabetic filename enumeration."""
    # aa ab .. az ba .. yy yz  zaaa ..  zyzz zzaaaa .. zzyzzz zzzaaaaa
    LETTERS = string.ascii_lowercase
    MAX=len(LETTERS)-1 # "z"
    digits = [0, 0] # "aa"
    prefix = 0
    while True:
      yield "".join(LETTERS[c] for c in digits)
      # Increment the digits starting from the right,
      # but not beyond the prefix boundary
      for p in range(len(digits)-1, prefix-1, -1):
        if digits[p] < MAX:
          digits[p] += 1 ## increment this letter
          break # No increment further to the left
        else:
          digits[p] = 0
          continue # increment to the left
      if digits[prefix] == MAX:
        # Overflow; extend our array by growing our "zz" prefix one char
        prefix += 1
        # add another 0 at suffix (e.g. adding "z0" in middle)
        digits = [MAX,]*prefix + [0,]* (prefix+2)
        # +2 reflect digits is now 2 chars wider

SAMPLES = list(itertools.islice(samples(), MAXSAMPLES))

rule all:
    input:
        expand("{type}/listing/{sample}.txt", sample=SAMPLES, type=TYPES)

rule zipfiles:
    input: # https://doi.org/10.5281/zenodo.3531504
        HTTP.remote("zenodo.org/record/3531504/files/zenodo-records-json-2019-09-16-filtered.jsonseq.xz", keep_local=True)
    output:
        "zenodo-records/zipfiles.tsv"
    conda:
        "envs/jq.yaml"
    shell: # TODO: Split into .jq script file
        "xzcat {input} |"
        """jq -r '. | select(.metadata.access_right == "open") | 
            .metadata.resource_type.type as $rectype |
            . as $rec | 
            ( .files[]?  | select(.type == "zip") ) |
            [$rec.id, $rec.links.self, $rec.links.doi, .checksum, .links.self, .size, .type, .key, $rectype] | @tsv
            '"""
        "> zenodo-records/zipfiles.tsv"

rule seed:
    output:
        "seed"
    shell: 
        "dd if=/dev/urandom of=seed bs=1024 count=1024"

rule shuffled:
    input: 
        tsv=rules.zipfiles.output,
        seed=rules.seed.output
    output:
        "zenodo-records/zipfiles-shuffled.tsv"
    conda:
        "envs/coreutils.yaml"
    shell:
        "shuf --random-source={input.seed} {input.tsv} > {output}"

rule splitzipfiles:
    input: 
        rules.shuffled.output
    params:
        type="{type}"
    output:
        "zenodo-records/zipfiles-{type}.tsv"
    run:
        if wildcards.type=="others":
            # Note: This join will strangely also filter out our catch-all type
            # "others" which should anyway not appear in the TSV file
            pattern = r"(%s)$" % "|".join(TYPES)
            shell("egrep -v \"{pattern}\" {input} > {output}")
        else:
            shell("grep {wildcards.type}$ {input} > {output}")

rule samples:
    input:
        "zenodo-records/zipfiles-{type}.tsv"
    output:
        expand("{{type}}/sample/{sample}.tsv", sample=SAMPLES)
    params:
        maxsamples=MAXSAMPLES
    shell:
        "head --lines={params.maxsamples} {input} | "
        "split --additional-suffix .tsv --lines 1 - {wildcards.type}/sample/"

rule downloadzip:
    input:
        "{type}/sample/{sample}.tsv"
    output:
        zip=temp("{type}/zip/{sample}.zip"),
        listing="{type}/listing/{sample}.txt"
    threads: 4
    conda:
        "envs/curl-unzip.yaml"
    shell:
        """URL="`awk '{{print $5}}' {input}`";
        curl --fail --retry 5 --retry-connrefused --output {output.zip} $URL || touch {output.zip};
        (unzip -Z -1 {output.zip} || true) > {output.listing}
        """

