# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .jl
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Julia 1.6.1
#     language: julia
#     name: julia-1.6
# ---

DATE = "2021-08-22"
TASK = "iterative-correction-via-resampling"
DIR = mkpath("$(homedir())/$(DATE)")

# +
import Pkg
pkgs = [
#     "BioAlignments",
    "BioSequences",
#     "Clustering",
#     "CodecZlib",
#     "Colors",
#     "Combinatorics",
#     "DataFrames",
#     "DataStructures",
    "Dates",
#     "DelimitedFiles",
#     "Distances",
#     "EzXML",
    "FASTX",
#     "GFF3",
#     "HTTP",
#     "Impute",
#     "JSON",
    "Graphs",
    "LSHFunctions",
#     "Measures",
    "MetaGraphs",
    "https://github.com/cjprybol/Mycelia.git",
#     "NumericIO",
#     "PlotlyJS",
#     "Plots",
    "Primes",
#     "Printf",
#     "ProgressMeter",
    "Random",
    "Revise",
    "SparseArrays",
#     "Statistics",
    "StatsBase",
#     "StatsPlots",
#     "StringDistances",
#     "uCSV",
#     "XLSX",
]

unregistered_packages = filter(pkg -> occursin(r"(^https|git$)", pkg), pkgs)
registered_packages = setdiff(pkgs, unregistered_packages)

for pkg in registered_packages
    try
        eval(Meta.parse("import $(pkg)"))
    catch
        Pkg.add(pkg)
        Pkg.build(pkg)
        eval(Meta.parse("import $(pkg)"))
    end
end

for pkg_url in unregistered_packages
    pkg_name = replace(basename(pkg_url), ".git" => "")
    try
        eval(Meta.parse("import $(pkg_name)"))
    catch
        Pkg.develop(url=pkg_url)
        Pkg.build(pkg_name)
        eval(Meta.parse("import $(pkg_name)"))
    end
end
# -

if Sys.isapple()
    BANDAGE = "/Applications/Bandage.app/Contents/MacOS/Bandage"
else
    BANDAGE = "Bandage"
end

function assess_kmer_sparsity_in_reads(k, observations)
#     @show k, observations
    canonical_kmer_counts = Mycelia.count_canonical_kmers(BioSequences.DNAMer{k}, observations)
#     canonical_kmer_counts = Mycelia.count_canonical_kmers(BioSequences.DNAMer{k}, first(observations))
# #     @show canonical_kmer_counts
#     for i in 2:length(observations)
#         canonical_kmer_counts_ = Mycelia.count_canonical_kmers(BioSequences.DNAMer{k}, observations[i])
#         canonical_kmer_counts = merge!(+, canonical_kmer_counts, canonical_kmer_counts_)
#     end
#     canonical_kmer_counts = map(o -> Mycelia.count_canonical_kmers(BioSequences.DNAMer{k}, o), observations)...)
    total_observed_canonical_kmers = length(keys(canonical_kmer_counts))
    total_possible_canonical_kmers = (4^k)/2
    sparsity = round(total_observed_canonical_kmers/total_possible_canonical_kmers*100, digits=2)
    @show sparsity
    return sparsity
end

function sequence_to_canonical_kmers(kmer_type, sequence)
    return (BioSequences.canonical(kmer.fw) for kmer in BioSequences.each(kmer_type, sequence))
end

function determine_edge_weights(graph)
    outgoing_edge_probabilities = determine_edge_weights(graph, true)
    incoming_edge_probabilities = determine_edge_weights(graph, false)
    return Dict(true => outgoing_edge_probabilities, false => incoming_edge_probabilities)
end

function determine_edge_weights(graph, strand)
    kmers = [graph.vprops[v][:kmer] for v in Graphs.vertices(graph)]
    outgoing_edge_weights = SparseArrays.spzeros(length(kmers), length(kmers))
    
    for (kmer_index, kmer) in enumerate(kmers)
        if !strand
            kmer = BioSequences.reverse_complement(kmer)
        end
        
        downstream_neighbors = []
        downstream_neighbor_weights = []
        for neighbor in BioSequences.neighbors(kmer)
            canonical_neighbor = BioSequences.canonical(neighbor)
            neighbor_orientation = neighbor == canonical_neighbor
            
            neighbor_index_range = searchsorted(kmers, canonical_neighbor)
            if !isempty(neighbor_index_range)
                @assert length(neighbor_index_range) == 1
                neighbor_index = first(neighbor_index_range)
                if Graphs.has_edge(graph, Graphs.Edge(kmer_index, neighbor_index))
                    edge_orientations = graph.eprops[Graphs.Edge(kmer_index, neighbor_index)][:orientations]
                    this_edge_orientation = (source_orientation = strand, destination_orientation = neighbor_orientation)
                    if this_edge_orientation in edge_orientations
                        edge_weight = graph.eprops[Graphs.Edge(kmer_index, neighbor_index)][:weight]
                        outgoing_edge_weights[kmer_index, neighbor_index] = edge_weight
                    end
                end
            end
        end
    end
    return outgoing_edge_weights
end

function determine_edge_probabilities(edge_weights)
    edge_probabilities = deepcopy(edge_weights)
    for orientation in [true, false]
        for row in 1:size(edge_probabilities[orientation], 1)
            edge_weights = edge_probabilities[orientation][row, :]
            edge_total = sum(edge_weights)
            for (col, edge_weight) in enumerate(edge_weights)
                if edge_total > 0
                    edge_probabilities[orientation][row, col] = edge_weight/edge_total
                else
                    edge_probabilities[orientation][row, col] = 0.0
                end
            end
        end
    end
    return edge_probabilities
end

function random_step(current_vertex, current_orientation, step_probabilities)
    outgoing_edge_likelihoods = step_probabilities[initial_orientation][current_vertex, :]
    chosen_step = StatsBase.sample(outgoing_edge_likelihoods.nzind, StatsBase.weights(outgoing_edge_likelihoods.nzval))
    possible_orientations = observation_graph.eprops[Graphs.Edge(current_vertex, chosen_step)][:orientations]
    possible_orientations = filter(o -> o.source_orientation == current_orientation, possible_orientations)
    chosen_orientation = rand(possible_orientations).destination_orientation
    chosen_oriented_step = (vertex = chosen_step, orientation = chosen_orientation)
    return chosen_oriented_step
end

function random_walk(observation_graph, initial_vertex, initial_orientation, step_probabilities; n_steps=size(step_probabilities[true], 1))
    step_count = 0
    outgoing_edge_likelihoods = step_probabilities[initial_orientation][initial_vertex, :]
    walk = Vector{NamedTuple{(:vertex, :orientation), Tuple{Int64, Bool}}}()
    done = (step_count > n_steps) || (sum(outgoing_edge_likelihoods) == 0)
    current_vertex = initial_vertex
    current_orientation = initial_orientation
    while !done
#         @show "choosing"
        chosen_step = StatsBase.sample(outgoing_edge_likelihoods.nzind, StatsBase.weights(outgoing_edge_likelihoods.nzval))
        chosen_step
        possible_orientations = observation_graph.eprops[Graphs.Edge(current_vertex, chosen_step)][:orientations]
        possible_orientations
        current_orientation
        possible_orientations = filter(o -> o.source_orientation == current_orientation, possible_orientations)
        possible_orientations
        chosen_orientation = rand(possible_orientations).destination_orientation
        chosen_oriented_step = (vertex = chosen_step, orientation = chosen_orientation)
        push!(walk, chosen_oriented_step)
        current_vertex = chosen_step
        current_orientation = chosen_orientation
        outgoing_edge_likelihoods = step_probabilities[last(walk).orientation][last(walk).vertex, :]
        step_count += 1
#         @show outgoing_edge_likelihoods
        done = (step_count > n_steps) || (sum(outgoing_edge_likelihoods) == 0)
    end
    return walk
end

# generate a genome
sequence_length = 10
sequence = BioSequences.randdnaseq(Random.seed!(sequence_length), sequence_length)
sequence_id = string(hash(sequence))
fasta_record = FASTX.FASTA.Record(sequence_id, sequence)

reference_fasta_file = "$(DIR)/LENGTH-$(sequence_length).fasta"
open(reference_fasta_file, "w") do io
    fastx_io = FASTX.FASTA.Writer(io)
    write(fastx_io, fasta_record)
    close(fastx_io)
end

# +
# randomly sample reads with errors from the genome
# 1%
error_rate = 0.01
n_reads = 100

# observations = [Mycelia.observe(FASTX.sequence(record), error_rate = error_rate) for i in 1:n_reads]
observations = [Mycelia.observe(fasta_record, error_rate = error_rate) for i in 1:n_reads];

# +
# build assembly graph
min_k = 7
max_k = min(minimum(length(observations)), 61)
k_options = Primes.primes(min_k, max_k)
k_index = findfirst(k -> assess_kmer_sparsity_in_reads(k, observations) < 100, k_options)
k = k_options[k_index]

observations_file = "$(DIR)/LENGTH-$(sequence_length)-DEPTH-$(n_reads).fastq"
open(observations_file, "w") do io
    fastq_io = FASTX.FASTQ.Writer(io)
    for record in observations
        write(fastq_io, record)
    end
    close(fastq_io)
end
# -

reference_graph = Mycelia.fastx_to_kmer_graph(BioSequences.DNAMer{k}, reference_fasta_file)
reference_kmers = [reference_graph.vprops[v][:kmer] for v in Graphs.vertices(reference_graph)]

observation_graph = Mycelia.fastx_to_kmer_graph(BioSequences.DNAMer{k}, observations_file)

# +
# visualize
gfa_file = observations_file * ".k-$k.gfa"
Mycelia.graph_to_gfa(observation_graph, gfa_file)

run(`$(BANDAGE) image $(gfa_file) $(gfa_file).svg --depwidth .9 --deppower .9`)
# --nodewidth <float> Average node width (0.5 to 1000, default: 5)
# --depwidth <float>  Depth effect on width (0 to 1, default: 0.5)
# --deppower <float>  Power of depth effect on width (0 to 1, default: 0.5)

html_path_to_svg = "./" * repeat("../", length(split(pwd(), '/')) - 3)
html_path_to_svg *= replace("$(gfa_file).svg", "$(homedir())/" => "")

x = "<img src=$(html_path_to_svg)>"
display("text/html", x)
# display("image/svg+xml", read("$(gfa_file).svg", String))
# -

kmer_counts = Dict(observation_graph.vprops[v][:kmer] => observation_graph.vprops[v][:weight] for v in Graphs.vertices(observation_graph))
total_observed_kmers = sum(values(kmer_counts))
kmer_probabilities = Dict(k => v/total_observed_kmers for (k,v) in kmer_counts)
scale = 250
Mycelia.plot_kmer_frequency_spectra(values(kmer_counts), size=(2scale,scale), title="kmer frequencies")

distance_to_reference_graph = 1 - LSHFunctions.jaccard(Set(reference_kmers), Set(keys(kmer_counts)))

# +
# step 1, hard-filter singletons
# -

filtered_vertices = findall(v -> observation_graph.vprops[v][:weight] >= 5, Graphs.vertices(observation_graph))

filtered_observation_graph, vertex_map = Graphs.induced_subgraph(observation_graph, filtered_vertices)

# +
# visualize
gfa_file = observations_file * ".k-$k.filtered.gfa"
Mycelia.graph_to_gfa(filtered_observation_graph, gfa_file)

run(`$(BANDAGE) image $(gfa_file) $(gfa_file).svg --depwidth .9 --deppower .9`)
# --nodewidth <float> Average node width (0.5 to 1000, default: 5)
# --depwidth <float>  Depth effect on width (0 to 1, default: 0.5)
# --deppower <float>  Power of depth effect on width (0 to 1, default: 0.5)

html_path_to_svg = "./" * repeat("../", length(split(pwd(), '/')) - 3)
html_path_to_svg *= replace("$(gfa_file).svg", "$(homedir())/" => "")

x = "<img src=$(html_path_to_svg)>"
display("text/html", x)
# display("image/svg+xml", read("$(gfa_file).svg", String))
# -

kmer_counts = Dict(filtered_observation_graph.vprops[v][:kmer] => filtered_observation_graph.vprops[v][:weight] for v in Graphs.vertices(filtered_observation_graph))
total_observed_kmers = sum(values(kmer_counts))
kmer_probabilities = Dict(k => v/total_observed_kmers for (k,v) in kmer_counts)
scale = 250
Mycelia.plot_kmer_frequency_spectra(values(kmer_counts), size=(2scale,scale), title="kmer frequencies")

distance_to_reference_graph = 1 - LSHFunctions.jaccard(Set(reference_kmers), Set(keys(kmer_counts)))

# +
# step 2, re-simulate reads as a means of error correction
# -

kmers = [filtered_observation_graph.vprops[v][:kmer] for v in Graphs.vertices(filtered_observation_graph)]
kmer_weights = [filtered_observation_graph.vprops[v][:weight] for v in Graphs.vertices(filtered_observation_graph)]
kmer_probabilities = kmer_weights ./ sum(kmer_weights)

edge_weights = determine_edge_weights(filtered_observation_graph)

edge_weights[true]

edge_weights = determine_edge_weights(filtered_observation_graph)
edge_probabilities = determine_edge_probabilities(edge_weights)

# +
kmers = [filtered_observation_graph.vprops[v][:kmer] for v in Graphs.vertices(filtered_observation_graph)]
kmer_weights = [filtered_observation_graph.vprops[v][:weight] for v in Graphs.vertices(filtered_observation_graph)]
kmer_probabilities = kmer_weights ./ sum(kmer_weights)

edge_weights = determine_edge_weights(filtered_observation_graph)
edge_probabilities = determine_edge_probabilities(edge_weights)

step_probabilities = deepcopy(edge_probabilities)

for orientation in [true, false]
#     @show orientation
    oriented_step_probabilities = step_probabilities[orientation] 
    for row in 1:size(step_probabilities[orientation], 1)
        step_probabilities_ = step_probabilities[orientation][row, :] .* kmer_probabilities
        if sum(step_probabilities_) > 0
            step_probabilities[orientation][row, :] .= step_probabilities_ ./ sum(step_probabilities_)
        end
    end
end
# -

new_records = FASTX.FASTQ.Record[]
for observation in observations
    read_length = length(FASTX.sequence(observation))
    max_steps = read_length - filtered_observation_graph.gprops[:k]
    # apply a squaring penalty to kmer_weights?
    initial_vertex = StatsBase.sample(1:length(kmers), StatsBase.weights(kmer_weights.^2))
#     @show "choosing initial vertex $initial_vertex with probability $(kmer_weights[initial_vertex])"
    initial_orientation = rand(Bool)
    forward_walk = random_walk(filtered_observation_graph, initial_vertex, initial_orientation, step_probabilities, n_steps=max_steps)
    remaining_steps = max_steps - length(forward_walk)
    reverse_walk = random_walk(filtered_observation_graph, initial_vertex, !initial_orientation, step_probabilities, n_steps = remaining_steps)
    full_walk = [[(vertex=x.vertex, orientation=!x.orientation) for x in reverse(reverse_walk)]..., [(vertex = initial_vertex, orientation = initial_orientation)]..., forward_walk...]
    oriented_path = [(x.vertex, x.orientation) for x in full_walk]
    new_seq = Mycelia.oriented_path_to_sequence(filtered_observation_graph, oriented_path)
    new_record = FASTX.FASTQ.Record(
        FASTX.identifier(observation)*"-k$k",
        FASTX.identifier(observation),
        new_seq,
        StatsBase.sample(FASTX.quality(observation), length(new_seq))
    )
#     @show "here"
    push!(new_records, new_record)
end
new_records;

corrected_fastq_file = replace(observations_file, r"\.fastq" => ".k$k.fastq")
open(corrected_fastq_file, "w") do io
    fastq_writer = FASTX.FASTQ.Writer(io)
    for record in new_records
        write(fastq_writer, record)
    end
    close(fastq_writer)
end

corrected_observation_graph = Mycelia.fastx_to_kmer_graph(BioSequences.DNAMer{k}, corrected_fastq_file)

# +
# visualize
gfa_file = corrected_fastq_file * ".k-$k.gfa"
Mycelia.graph_to_gfa(corrected_observation_graph, gfa_file)

run(`$(BANDAGE) image $(gfa_file) $(gfa_file).svg --depwidth .9 --deppower .9`)
# --nodewidth <float> Average node width (0.5 to 1000, default: 5)
# --depwidth <float>  Depth effect on width (0 to 1, default: 0.5)
# --deppower <float>  Power of depth effect on width (0 to 1, default: 0.5)

html_path_to_svg = "./" * repeat("../", length(split(pwd(), '/')) - 3)
html_path_to_svg *= replace("$(gfa_file).svg", "$(homedir())/" => "")

x = "<img src=$(html_path_to_svg)>"
display("text/html", x)
# display("image/svg+xml", read("$(gfa_file).svg", String))
# -

kmer_counts = Dict(corrected_observation_graph.vprops[v][:kmer] => corrected_observation_graph.vprops[v][:weight] for v in Graphs.vertices(corrected_observation_graph))
total_observed_kmers = sum(values(kmer_counts))
kmer_probabilities = Dict(k => v/total_observed_kmers for (k,v) in kmer_counts)
scale = 250
Mycelia.plot_kmer_frequency_spectra(values(kmer_counts), size=(2scale,scale), title="kmer frequencies")

# +
# compare the new graph to the error-free graph to the error-corrected graph
# -

distance_to_reference_graph = 1 - LSHFunctions.jaccard(Set(reference_kmers), Set(keys(kmer_counts)))

StarcoderdataJupyterScriptsDedupFiltered