shell:
    """
    prokka --force --cpus {threads} --prefix {wildcards.genome} --locustag {params.locustag} --kingdom Viruses --outdir {params.outdir}/{wildcards.genome} {input}
    cp {params.outdir}/{wildcards.genome}/{wildcards.genome}.gbk {output}
    """

shell:
    "seqkit grep -rp {params.search} {input} | seqkit replace -p '$' -r ' {wildcards.segment}' -o {output}"

suppressPackageStartupMessages(library(dplyr))
suppressPackageStartupMessages(library(tidyr))
suppressPackageStartupMessages(library(stringr))

stdin_fd <- file("stdin")
all.lines <- readLines(stdin_fd)
close(stdin_fd)
param.start <- 1
data.start  <- which(grepl("^ *No Hit", all.lines)) %>% first %>% `+`(1)

align.start <- which(grepl("^No 1", all.lines)) %>% first

param.end <- data.start - 2
data.end  <- align.start - 1
align.end <- length(all.lines)

if (is.na(align.start)) {
    data <- tibble(
        Query           = character(),
        No              =   integer(),
        Hit.ID          = character(),
        Hit.Description = character(),
        Q.ss_pred       = character(),
        Q.query         = character(),
        Q.consensus     = character(),
        Q.Start         =   integer(),
        Q.End           =   integer(),
        Q.Length        =   integer(),
        T.consensus     = character(),
        T.Start         =   integer(),
        T.End           =   integer(),
        T.Length        =   integer(),
        T.hit           = character(),
        T.ss_dssp       = character(),
        T.ss_pred       = character(),
        Aligned_cols    =   integer(),
        E.value         =   numeric(),
        Identities      =   numeric(),
        Probab          =   numeric(),
        Score           =   numeric(),
        Similarity      =   numeric(),
        Sum_probs       =   numeric(),
        Template_Neff   =   numeric()
    )
} else {
    metadata <- data.frame(key = all.lines[param.start:param.end]) %>%
        mutate(value = substr(key, 14, 10000) %>% trimws, key = substr(key, 1, 14) %>% trimws) %>%
        filter(key != "") %>%
        {setNames(.$value, .$key)} %>%
        as.list
    data <- data.frame(Query = sub(" .*", "", metadata$Query), line = all.lines[align.start:align.end], stringsAsFactors = F) %>%
        filter(line != "") %>%
        extract(line, into = c("name", "value"), regex = "([^ ]+) ?(.+)?", remove = F) %>%
        mutate(No = ifelse(name == "No", value, NA) %>% as.integer) %>%
        mutate(Hit.ID = ifelse(substr(name, 1, 1) == ">", substr(name, 2, nchar(.)), NA)) %>%
        mutate(Hit.Description = ifelse(substr(name, 1, 1) == ">", value, NA)) %>%
        mutate(Match = ifelse(grepl("=", name), line, NA)) %>%
        mutate(name = ifelse(grepl("Q Consensus", lag(line)) & grepl("T Consensus", lead(line)), "M", name)) %>%
        mutate(value = ifelse(name == "M", line, value)) %>%
        fill(No) %>%
        group_by(Query, No) %>%
        summarize(
            Hit.ID       = na.omit(Hit.ID) %>% first,
            Hit.Description = na.omit(Hit.Description) %>% first,
            Match        = na.omit(Match) %>% first,
            Q.ss_pred    = value[name == "Q" & grepl("^ss_pred ", value)]         %>% substr(., 16, nchar(.)) %>% paste(collapse = "") %>% gsub(" +", "", .),
            Q.query      = value[name == "Q" & grepl("^Consensus ", lead(value))] %>% substr(., 16, nchar(.)) %>% paste(collapse = " "),
            Q.consensus  = value[name == "Q" & grepl("^Consensus ", value)]       %>% substr(., 16, nchar(.)) %>% paste(collapse = " "),
            T.consensus  = value[name == "T" & grepl("^Consensus ", value)]       %>% substr(., 16, nchar(.)) %>% paste(collapse = " "),
            T.hit        = value[name == "T" & grepl("^Consensus ", lag(value))]  %>% substr(., 16, nchar(.)) %>% paste(collapse = " "),
            T.ss_dssp    = value[name == "T" & grepl("^ss_dssp ", value)]         %>% substr(., 16, nchar(.)) %>% paste(collapse = " ") %>% gsub(" +", "", .),
            T.ss_pred    = value[name == "T" & grepl("^ss_pred ", value)]         %>% substr(., 16, nchar(.)) %>% paste(collapse = "")  %>% gsub(" ", "", .),
            .groups = "drop"
        ) %>%
        extract(Q.consensus, into = c("Q.Start", "Q.End", "Q.Length"), regex = "^ *(\\d+) .+ (\\d+) +[(](\\d+)[)]$", remove = F, convert = T) %>%
        extract(T.consensus, into = c("T.Start", "T.End", "T.Length"), regex = "^ *(\\d+) .+ (\\d+) +[(](\\d+)[)]$", remove = F, convert = T) %>%
        mutate(
            Q.consensus  = gsub("[0-9() ]+", "", Q.consensus),
            Q.query      = gsub("[0-9() ]+", "", Q.query),
            T.consensus  = gsub("[0-9() ]+", "", T.consensus),
            T.hit        = gsub("[0-9() ]+", "", T.hit),
        ) %>%
        #extract(Hit.Description, into = "Hit.Organism",    regex = "[{]([^}]+)[}]",  remove = F) %>%
        #extract(Hit.Description, into = "Hit.Description", regex = "([^;]+)",        remove = F) %>%
        #extract(Hit.Description, into = "Hit.Keywords",    regex = "[^;]+; ([^;]+)", remove = F) %>%
        mutate(Match = str_split(Match, " +")) %>%
        unnest(cols = Match) %>%
        separate(Match, into = c("key", "value"), "=") %>%
        mutate(value = sub("%", "", value) %>% as.numeric) %>%
        spread(key, value) %>%
        rename(E.value = `E-value`) %>%
        mutate(Aligned_cols = as.integer(Aligned_cols))
}

write.table(data, quote = F, sep = "\t", row.names = F)

shell:
    "mafft --localpair --maxiterate 1000 --thread {threads} {input} > {output}"

shell:
    "trimal -in {input} -out {output} -gt {params.gt}"

shell:
    "hmmbuild {output} {input}"

shell:
    "hmmsearch -o /dev/null --tblout {output} {input.hmm} {input.faa}"

shell:
    "hmmalign --trim --outformat A2M {input.hmm} {input.fasta} > {output}"

shell:
    "raxml-ng --redo --evaluate --msa {input.fasta} --tree {input.iqtree} --model {params.model} --seed {params.seed} --prefix {params.prefix} &> {log}"

shell:
    "epa-ng --redo -s {input.fasta} -t {input.tree} --model {input.model} -q {input.short} -w {params.dir} &> {log}"

shell:
    "gappa examine graft --jplace-path {input} --fully-resolve --out-dir {params.out_dir}"

shell:
    "meme -minw {params.minw} -maxw {params.maxw} -nmotifs {params.nmotifs} -minsites {params.minsites} -maxsites Inf -dna -oc {params.outdir} {input}"

shell:
    "seqret -supper1 -filter -outseq {output} {input}"

shell:
    "fimo --oc {params.out_dir} {input.meme} {input.fasta}"

shell:
    "ffindex_apply -q -d {output.data} -i {output.index} {input.data} {input.index} -- hhblits -cpu 1 -v 0 -b 1 -z 1 -d {params.db} -i stdin -oa3m stdout -e {params.e} -n {params.n}"

shell:
    "mmseqs createdb {input} {output}"

shell:
    "mcl {input} -o {output} --abc -scheme {params.scheme} -te {threads} -I {params.I}"

import re
from Bio import SeqIO

input_gbk = str(snakemake.input)
output_gbk = str(snakemake.output)
flank = snakemake.params['flank']
segment = snakemake.params['segment']

coords_re = re.compile(r'(.+)_(\d+)-(\d+)[. ]*$')
def get_coords(description):
    search = coords_re.search(description)
    assert search, "Unexpected sequence name: %s" % description
    scaffold, start, end = search.groups()
    start = int(start)
    end = int(end) + flank
    if start > flank:
        start = start - flank
    else:
        start = 1
    return scaffold, start, end

found = False
with open(input_gbk) as fd:
    records = SeqIO.parse(fd, 'genbank')
    for record in records:
        scaffold, start, end = get_coords(record.description)
        if scaffold == segment['scaffold'] and start <= segment['scaffold_start'] and end >= segment['scaffold_end']:
            found = True
            break
assert found, "Segment %s not found" % segment['segment_id']

segment_start = segment['scaffold_start'] - start
segment_end   = segment['scaffold_end'] - start + 1

sub_record = record[segment_start:segment_end]
if segment['strand'] < 0:
    sub_record = sub_record.reverse_complement()

sub_record.name = segment['name']
sub_record.id = "%s_%d-%d" % (segment['scaffold'], segment['scaffold_start'], segment['scaffold_end'])
sub_record.description = sub_record.id

sub_record.annotations["molecule_type"] = "DNA"
SeqIO.write(sub_record, output_gbk, "genbank")

from Bio import SeqIO
from Bio.SeqFeature import SeqFeature, FeatureLocation, CompoundLocation
from collections import defaultdict

ref_annots = []

gff_file = str(snakemake.input["gff"])
fna_file  = str(snakemake.input["fna"])
blastn_file = str(snakemake.input["blastn"])
pfam_file = str(snakemake.input["pfam"])
hmm_dbs = snakemake.input["hmm_dbs"]
hmmscan_files = snakemake.input["hmmscan"]
markers_files = snakemake.input["markers"]

output_file = str(snakemake.output)

evalue_threshold = snakemake.params["evalue"]
coding_feature = snakemake.params["coding_feature"]

min_repeat_id = snakemake.params["min_repeat_id"]
min_repeat_gap = snakemake.params["min_repeat_gap"]
min_repeat_len = snakemake.params["min_repeat_len"]

aliases = {
    "MCP": "Major capsid protein (MCP)",
    "mCP": "Minor capsid protein (mCP)",
    "uncharacterized protein": "-",
    "DNA polymerase - helicase": "DNA polymerase-helicase"
}

features = defaultdict(list)

with open(blastn_file) as fh:
    for line in fh:
        qseqid, sseqid, pident, length, mismatch, gapopen, qstart, qend, sstart, send, evalue, bitscore = line.rstrip().split()
        qstart = int(qstart)
        qend   = int(qend)
        sstart = int(sstart)
        send   = int(send)
        if float(pident) > min_repeat_id and qseqid == sseqid and int(length) >= min_repeat_len and qstart < qend and sstart > send and sstart - qend >= min_repeat_gap:
            loc1 = FeatureLocation(qstart - 1, qend, strand = 1)
            loc2 = FeatureLocation(send - 1, sstart, strand = -1)
            feature = SeqFeature(CompoundLocation([loc1, loc2]), type = 'repeat_region')
            feature.qualifiers['rpt_type'] = 'inverted'
            feature.qualifiers['inference'] = 'ab initio prediction:blastn'
            feature.qualifiers['note'] = 'pident: %s%%' % pident
            features[qseqid].append(feature)

descs = {}
for fname in hmm_dbs:
    with open(fname) as fh:
        NAME = ACC = None
        for line in fh:
            if line.startswith('ACC'):
                key, ACC = line.rstrip().split()
            if line.startswith('NAME'):
                key, NAME = line.rstrip().split()
            if line.startswith('DESC'):
                key, DESC = line.rstrip().split(maxsplit = 1)
                if ACC: descs[ACC] = DESC
                if NAME: descs[NAME] = DESC
                NAME = ACC = None

hits = defaultdict(list)

with open(pfam_file) as fh:
    for line in fh:
        if not line.startswith('#') and line != '\n':
            seq_id, aln_start, aln_end, env_start, env_end, hmm_acc, hmm_name, type, hmm_start, hmm_end, hmm_len, bit_score, E_value, significance, clan = line.rstrip().split()
            hit = "%s: %s (%s)" % (hmm_acc, descs[hmm_acc], E_value)
            hits[seq_id].append(hit)

for fname in markers_files:
    with open(fname) as fh:
        for line in fh:
            if not line.startswith('#'):
                t_name, t_acc, q_name, q_acc, E_value, score, bias, best_E_value, best_score, best_bias, exp, reg, clu, ov, env, dom, rep, inc, desc = line.rstrip().split(maxsplit = 18)
                if float(E_value) <= evalue_threshold:
                    hit = "Marker %s: %s (%s)" % (q_name, descs[q_name], E_value)
                    hits[t_name].append(hit)

for fname in hmmscan_files:
    with open(fname) as fh:
        for line in fh:
            if not line.startswith('#'):
                t_name, t_acc, q_name, q_acc, E_value, score, bias, best_E_value, best_score, best_bias, exp, reg, clu, ov, env, dom, rep, inc, desc = line.rstrip().split(maxsplit = 18)
                if desc in aliases:
                    desc = aliases[desc]
                definition = ''
                if t_name != desc and desc != '-':
                    definition = ' ' + desc if t_name != desc and desc != '-' else ''
                if float(E_value) <= evalue_threshold:
                    hit = "%s:%s (%s)" % (t_name, definition, E_value)
                    hits[q_name].append(hit)

with open(gff_file) as fh:
    for line in fh:
        if not line.startswith('#') and line != '\n':
            seqname, source, feature_type, start, end, score, strand, frame, attributes = line.rstrip().split('\t')
            if feature_type == coding_feature:
                atts = {}
                for att in attributes.split(';'):
                    key, value = att.split('=')
                    atts[key] = value
                if 'Parent' in atts:
                    parent = atts['Parent']
                else:
                    parent = atts['ID']
                loc = FeatureLocation(int(start) - 1, int(end))
                feature = SeqFeature(loc, type = 'CDS', strand = int(strand + '1'))
                feature.qualifiers['locus_tag'] = parent
                feature.qualifiers['inference'] = "ab initio prediction:" + source.replace("_v", ":")
                if parent in hits:
                    feature.qualifiers['note'] = hits[parent]
                features[seqname].append(feature)

with open(output_file, 'w') as out_fh:
    with open(fna_file) as in_fh:
        rec_num = 1
        for record in SeqIO.parse(in_fh, 'fasta'):
            record.annotations['molecule_type'] = 'DNA'
            record.name = "locus_%d" % rec_num
            for feature in features[record.id]:
                if feature.type == 'CDS':
                    transl = feature.translate(record, cds = False)
                    feature.qualifiers['translation'] = transl.seq.rstrip('*')
                record.features.append(feature)
            SeqIO.write(record, out_fh, 'genbank')
            rec_num += 1

library(dplyr)
library(tidyr)
library(ape)
library(ggtree)
library(treeio)
library(phangorn)
library(stringr)
library(ggplot2)
library(phytools)

if (interactive()) {
    setClass("snake", slots = list(input = "list", output = "list"))
    snakemake <- new("snake", input  = list(
            tree = "analysis/phylogeny/MCP_NCLDV_epa/epa_result.newick",
            fasta = "analysis/phylogeny/MCP_NCLDV.fasta",
            outgroups    = "metadata/queries/MCP_NCLDV_outgroups.faa",
            synonyms = "metadata/organisms.txt",
            hmm = Sys.glob("hmm_algae/*.hmm")
    ), output = list(
        image = "test.svg",
        jtree = "output/MCP_NCLDV.jtree"
    ))
}

with(snakemake@input, {
    tree_file     <<- tree
    fasta_file    <<- fasta
    synonyms_file <<- synonyms
    outgroup_file <<- outgroups
})
with(snakemake@output, {
    out_image_file <<- image
    out_jtree_file <<- jtree
})
with(snakemake@params, {
    outgroup_rooting <<- outgroup_rooting
})

read.fasta.headers <- function(fnames) {
    file.info(fnames) %>%
        filter(size > 0) %>%
        rownames %>%
        lapply(treeio::read.fasta) %>%
        lapply(names) %>%
        unlist %>%
        data.frame(title = .)
}

synonyms <- read.table(synonyms_file, header = T, sep = "\t", fill = T, na.strings = "") %>%
    mutate(Collapse = ifelse(is.na(Collapse), Name, Collapse))

headers <- read.fasta.headers(fasta_file) %>%
    extract(title, into = c("label", "ID"), regex = "^([^ ]+) ([^ ]+)", remove = F) %>%
    left_join(synonyms, by = "ID")

no_name <- filter(headers, is.na(Name)) %>%
    pull(label) %>%
    paste(collapse = ", ")
if (no_name != "") {
    print(paste("No aliases found for: ", no_name))
    quit(status = 1)
}

tree <- read.tree(tree_file)
tree <- phangorn::midpoint(tree, node.labels = "support")
if (outgroup_rooting) {
    outgroup_df <- read.fasta.headers(outgroup_file)
    outgroups <- with(outgroup_df, sub(" .*", "", title))
    tree <- ape::root(tree, node = MRCA(tree, outgroups), edgelabel = T, resolve.root = T)
}
tree <- as_tibble(tree) %>%
    mutate(support = ifelse(node %in% parent & label != "", label, NA)) %>%
    separate(support, into = c("SH_aLRT", "UFboot"), sep = "/", convert = T) %>%
    left_join(headers, by = "label") %>%
    mutate(label.show = Name) %>%
    mutate(isInternal = node %in% parent) %>%
    `class<-`(c("tbl_tree", "tbl_df", "tbl", "data.frame"))
tree_data <- as.treedata(tree)
write.jtree(tree_data, file = out_jtree_file)

ntaxa <- filter(tree, ! node %in% parent) %>% nrow

colors <- list(
    Haptophyta = "orange",
    Chlorophyta = "green",
    Streptophyta = "darkgreen",
    MAG = "purple",
    Stramenopiles = "brown",
    Cryptophyta = "red",
    Amoebozoa = "gold4",
    Euglenozoa = "yellow",
    Choanoflagellata = "darkslateblue",
    Glaucophyta = "cyan",
    Animals = "blue",
    Dinoflagellata = "gray50",
    Rhizaria = "gray30"
)

scaleClades <- function(p, df) {
    with(df, Reduce(function(.p, .node) {
        offs <- offspring(.p$data, .node)
        scale <- 0.5 / (nrow(offs) - 1)
        scaleClade(.p, .node, scale)
    }, node, p))
}
collapseClades <- function(p, df) {
    with(df, Reduce(function(.p, .node) {
        fill <- unlist(colors[Host[node == .node]])
        .p$data[.p$data$node == .node, "label.show"] <- label.show[node == .node]
        collapse(.p, .node, "mixed", fill = fill)
    }, node, p))
}
#labelClades <- function(p) {
#    with(df, Reduce(function(.p, .node) {
#        .p + geom_cladelab(node = .node, label = label[node == .node], align = T, offset = .2, textcolor = 'blue')
#    }, node, p))
#}

multi_species <- allDescendants(tree_data@phylo) %>%
    lapply(function(x) filter(tree, node %in% x)) %>%
    bind_rows(.id = "ancestor") %>%
    group_by(ancestor) %>%
    filter(n_distinct(Collapse, na.rm = T) == 1, sum(!isInternal) > 1) %>% # , !any(Group == "Haptophyta")) %>%
    ungroup %>%
    mutate(ancestor = as.numeric(ancestor)) %>%
    filter(! ancestor %in% node) %>%
    filter(!is.na(Collapse)) %>%
    group_by(ancestor, Collapse) %>%
    summarize(num_tips = sum(!isInternal), Host = first(na.omit(Host))) %>%
    mutate(label.show = sprintf("%s (%d)", Collapse, num_tips)) %>%
    rename(node = ancestor)
p <- ggtree(tree_data) +
    geom_nodepoint(aes(x = branch, subset = !is.na(UFboot) & UFboot >= 90, size = UFboot)) +
    geom_tiplab(aes(label = label.show), size = 4, align = T, linesize = 0) +
    geom_text2(aes(subset = node %in% multi_species$node, x = max(x, na.rm = T), label = label.show), nudge_x = 0.01, size = 4, hjust = 0) +
    geom_tippoint(aes(color = Host), size = 3) +
    geom_treescale(width = 0.5) +
    scale_size_continuous(limits = c(90, 100), range = c(1, 3)) +
    scale_shape_manual(values = seq(0,15)) +
    scale_color_manual(values = colors)

p <- scaleClades(p, multi_species)
p <- collapseClades(p, multi_species)
# p <- facet_plot(p, mapping = aes(x = as.numeric(as.factor(query.name)), shape = DESC), data = genes, geom = geom_point, panel = 'Genes')

ggsave(out_image_file, p, height = ntaxa * 0.1, width = 7, limitsize = F)