shell:
    """
    ( tar -xzf {input} -C {params.untar_dir} &&
     cp -r /opt/RepeatMasker/Libraries/* {output.libdir} &&
     ln -sf {output.dfam} {output.rm} &&
     addRepBase.pl --libdir {output.libdir} ) &> {log}
    """

shell:
    """
    BuildDatabase -name {params.out} {input} &> {log}
    """

shell:
    """
    ( famdb.py -i {input.rm_db} families \
        --format fasta_name \
        --ancestors --descendants --include-class-in-name \
        'tetrapoda' > {params.rm_db_fasta} &&

    cat {params.rm_db_fasta} {input.tr_db} > {output} ) 2> {log}
    """

shell:
    """
    prefetch {wildcards.acc} -O {params.outdir} 2> {log}
    """

shell:
    """
    cd {params.outdir}
    fasterq-dump --split-3 \
        -e {threads} \
        --skip-technical \
        {wildcards.acc} 2> {log}
    """

shell:
    """
    gzip {input} 2> {log} 
    """

shell:
    """
    STAR --readFilesIn {input.R1} {input.R2} \
        --outSAMtype BAM SortedByCoordinate \
        --outSAMstrandField intronMotif \
        --twopassMode Basic \
        --genomeDir {input.star_build} \
        --runThreadN {threads} \
        --readFilesCommand zcat \
        --alignIntronMax 10000 \
        --outFileNamePrefix {params.out} &> {log} 
    """

shell:
    """
    ( wget --no-check-certificate https://bioinf.uni-greifswald.de/bioinf/partitioned_odb11/Vertebrata.fa.gz -P {params.outdir} && 
    gunzip {params.outdir}/Vertebrata.fa.gz ) 2> {log}
    """

shell:
    """
    curl -OJX GET "https://api.ncbi.nlm.nih.gov/datasets/v2alpha/genome/accession/{params.ref_seq}/download?include_annotation_type=PROT_FASTA&filename={wildcards.psp}.zip" -H "Accept: application/zip" &&
    unzip {wildcards.psp}.zip -d {params.outdir} &&
    cp  {params.outdir}/ncbi_dataset/data/{params.ref_seq}/protein.faa {output}
    rm -rf {params.outdir}/ncbi_dataset
    """

shell:
    """
    curl --output {output} '{params.url}' 2> {log}
    """

shell:
    """
    cat {input} > {output}
    """

shell:
    """
    braker.pl --genome {input.masked_genome} \
        --prot_seq {input.proteins} \
        --softmasking \
        --useexisting \
        --threads {threads} \
        --workingdir {params.outputdir} \
        --species "Podacris siculus prot" 2> {log}
    """ 

shell:
    """
    braker.pl --genome {input.masked_genome} \
        --bam {input.bam} \
        --softmasking \
        --useexisting \
        --threads {threads} \
        --workingdir {params.outputdir} \
        --species "Podacris siculus rna" 2> {log} 
    """

shell:
    """
    tsebra.py -g {input.prot_aug} \
        -k {input.rna_aug} \
        -c {params.config} \
        -e {input.hints_rna},{input.hints_prot} \
        -o {output} 2> {log} 
    """ 

shell:
    """
    rename_gtf.py --gtf {input} \
            --prefix Psic \
            --translation_tab {output.tab} \
            --out {output.gtf} 2> {log} 
    """

run:
    features = ['exon', 'intron', 'gene', 'transcript']
    with open(input[0], 'r') as fin:
        with open(output[0], 'w') as fout:
            lines = fin.readlines()
            for line in lines:
                sline = line.split('\t')
                if sline[2] in features:
                    pass
                else:
                    fout.write(line)

shell:
    """
    agat_convert_sp_gxf2gxf.pl --gtf {input} \
        --output {output} &> {log}
    """

shell:
    """
    sed 's/SEQ/0/g' {input} > {output}
    """

shell:
    """
    grep -v 'Psic_g38722' {input} > {output}
    """

shell:
    """
    gt gff3 -sort -tidy -retainids {input} > {output} 2> {log}
    """

run:
    with open(input[0], 'r') as fin:
        with open(output[0], 'w') as fout:
            lines = fin.readlines()
            for line in lines:
                if not line.startswith('#'):
                    sline = line.split('\t')
                    feature = sline[2]
                    if feature == 'gene':
                        #Remove transcript_id attribute from gene features
                        sline[8] = re.sub(r'(;transcript_id.*$)', '', sline[8])
                    elif feature == 'mRNA':
                        #Make sure transcript_id annotations for isoforms are correct (i.e., .t2, .t3, etc.)
                        #Use ID attribute since transcript_id attributes are incorrectly incremented and will be replaced
                        id_pattern = r"(?<=ID=)(.*)(?=;Parent)"
                        ID = re.search(id_pattern, sline[8]).group(1)
                        if ID.endswith('.t1'):
                            #First isoforms are fine
                            pass
                        else:
                            #Alternative isoforms need transcript_id replaced with ID
                            sline[8] = re.sub(r'(?<=;transcript_id=)(.*$)', ID, sline[8])
                    fout.write('\t'.join(sline))
                else:
                    fout.write(line) 

shell:
    """
    mkdir -p {output}
    download_eggnog_data.py -y --data_dir {output} 
    """

shell:
    """
    mkdir -p {params.outdir}
    funannotate annotate \
        --gff {input.gff} \
        --fasta {input.ref} \
        --species "Podacris siculus" \
        --out {params.outdir} \
        --iprscan {input.iprs} \
        --eggnog {input.enm} \
        --force \
        --database {input.db} \
        --busco_db tetrapoda \
        --cpus {threads} 2> {log}
    """

shell:
    """
    wget {params.url} --no-check-certificate -P {params.outdir}
    """

script:
    "../scripts/python/fixEC_incrementCDS_addLocusTags.py"

shell:
    """
    gff3_sort -g {input} -r -og {output} 2> {log}
    """

shell:
    """
    gffread -E -y {output} -g {input.ref} {input.gff} 2> {log}
    """

shell:
    """
    touch {output}
    """

shell:
    """
    busco -m genome \
            -i {input.fasta} \
            -o {params.out_name} \
            --out_path {params.out_path} \
            --lineage tetrapoda_odb10 \
            --force \
            --cpu {threads} &> {log}
    """

shell:
    """
    agat_sp_functional_statistics.pl --gff {input.gff} \
            -gs {input.fasta} \
            --output {output} 2> {log}
    """

shell:
    """
    busco -m protein \
            -i {input} \
            -o {params.out_name} \
            --out_path {params.out_path} \
            --lineage tetrapoda_odb10 \
            --force \
            --cpu {threads} &> {log}
    """

shell:
    """
    touch {output}
    """

shell:
    """
    ln -sf {input} {output.fasta} &&
    samtools faidx {output.fasta}
    """

shell:
    """
    cat {input.fai} | awk '$2 > 100000 {{print $1}}' > {output.scaffs}&&
    samtools faidx {input.fasta} -r {output.scaffs} > {output.fasta}
    """

import re
import gffutils

# Creat dictionary mapping products to EC numbers
EC_num_to_products = {}
with open(snakemake.input['ec'][0], 'r') as fin:
    lines = fin.readlines()
    for l in lines:
        if l.startswith('ID'):
            EC_number = l.split('   ')[1].strip()
        if l.startswith('DE'):
            prod = l.split('   ')[1].strip()
            if not EC_number in EC_num_to_products.keys():
                EC_num_to_products[EC_number] = prod
            else:
                EC_num_to_products[EC_number] += EC_num_to_products[EC_number] + prod
        else:
            pass

locus_tag_prefix = snakemake.params['locus_tag']
with open(snakemake.output[0], 'w') as fout:
    fout.write("##gff-version 3\n")
    gene_num = 1
    hyp_wEC_count = 0
    reassigned = 0
    for feat in gffutils.DataIterator(snakemake.input['gff']):
        if feat.featuretype == 'gene':
            # Set CDS counter
            cds_count = 1

            # Add locus tag to genes
            feat.attributes["locus_tag"] = f'{locus_tag_prefix}_' + str(int(gene_num)).zfill(5)
            gene_num += 1
        elif feat.featuretype == 'mRNA':
            if 'product' in feat.attributes.keys() and 'EC_number' in feat.attributes.keys():
                # Rename EC_number
                feat.attributes['ec_number'] = [feat.attributes['EC_number'][0]]
                feat.attributes.pop('EC_number')

                # Reassign product if EC Number to 4 digits, else delete EC Number and keep hypothetical
                EC_number = feat.attributes['ec_number'][0]
                prod = feat.attributes['product'][0]
                if prod == 'hypothetical protein':
                    EC_split = EC_number.split('.')
                    if len(EC_split) < 4:
                        hyp_wEC_count += 1
                        feat.attributes.pop('ec_number')
                    elif len(EC_split) == 4:
                        reassigned += 1
                        new_product = EC_num_to_products[EC_number]
                        # Handle cases where EC numbers have been reassigned
                        if 'Transferred' in new_product:
                            pattern = r'(?<=:\s)(\d\.\d+\.\d+\.\d+(?=[\.|,]))'
                            new_ec = re.search(pattern, new_product).group(1)
                            new_product = EC_num_to_products[new_ec]

                        new_product = new_product.replace('.', '')
                        feat.attributes['product'] = [new_product]

                        # Add note to GFF about origin of Product
                        feat.attributes['note'] = ["Funnanotate product changed from hypothetical protein based on EC_number"]
                else:
                    EC_split = EC_number.split('.')
                    if len(EC_split) < 4:
                        new_ec = '.'.join(EC_split + ['-'] * (4 - len(EC_split)))
                        feat.attributes['ec_number'] = [new_ec]
        elif feat.featuretype == 'CDS':
            # Assign increment CDS to ensure unique ID
            feat.attributes['ID'][0] = re.sub('cds', f'cds{cds_count}', feat.attributes['ID'][0])
            cds_count += 1
        else:
            pass
        fout.write(str(feat) + '\n')
print(f"{hyp_wEC_count} products were Hypothetical and contained EC Numbers")
print(f"{reassigned} hypothetical proteins have been reassigned based on their EC Numbers")