# This code is part of the Biopython distribution and governed by its
# license.  Please see the LICENSE file that should have been included
# as part of this package.
#

"""Example of using GenomeDiagram cross-links to mimic ACT."""

import os
import sys

from reportlab.lib import colors
from reportlab.lib.units import cm

from Bio import SeqIO
from Bio.Graphics.GenomeDiagram import CrossLink
from Bio.Graphics.GenomeDiagram import Diagram
from Bio.SeqFeature import SeqFeature
from Bio.SeqFeature import SimpleLocation

# Modify this line to point at the Artemis/ACT example data which is online at:
# https://github.com/sanger-pathogens/Artemis/tree/master/etc
input_folder = "/Applications/Artemis/Artemis.app/Contents/artemis/etc"

name = "af063097_v_b132222"
file_a = "af063097.embl"
file_b = "b132222.embl"
format_a = "embl"
format_b = "embl"
file_a_vs_b = "af063097_v_b132222.crunch"

for f in [file_a, file_b, file_a_vs_b]:
    if not os.path.isfile(os.path.join(input_folder, f)):
        print("Missing input file %s.fna" % f)
        sys.exit(1)

# Only doing a_vs_b here, could also have b_vs_c and c_vs_d etc
genomes = [
    (os.path.join(input_folder, file_a), format_a),
    (os.path.join(input_folder, file_b), format_b),
]
comparisons = [os.path.join(input_folder, file_a_vs_b)]

# Create diagram with tracks, each with a feature set
assert len(genomes) >= 2 and len(genomes) == len(comparisons) + 1
gd_diagram = Diagram(name, track_size=0.35, circular=False)
tracks = {}
feature_sets = {}
records = {}
for f, format in genomes:
    records[f] = SeqIO.read(f, format)
    tracks[f] = gd_diagram.new_track(
        1,
        name=f,
        start=0,
        end=len(records[f]),
        scale_smalltick_interval=1000,
        scale_largetick_interval=10000,
        greytrack=True,
        greytrack_labels=0,
    )
    feature_sets[f] = tracks[f].new_set()

print("Drawing matches...")
for i, crunch_file in enumerate(comparisons):
    q = genomes[i + 1][0]  # query file
    s = genomes[i][0]  # subject file
    q_set = feature_sets[q]
    s_set = feature_sets[s]
    with open(crunch_file) as handle:
        for line in handle:
            if line[0] == "#":
                continue
            parts = line.rstrip("\n").split(None, 7)
            # 0 = score
            # 1 = id
            # 2 = S1
            # 3 = E1
            # 4 = seq1
            # 5 = S2
            # 6 = E2
            # 7 = seq2
            try:
                q_start, q_end = int(parts[2]), int(parts[3])
                s_start, s_end = int(parts[5]), int(parts[6])
            except IndexError:
                sys.stderr.write(repr(line) + "\n")
                sys.stderr.write(repr(parts) + "\n")
                raise
            flip = False
            if q_start > q_end:
                flip = not flip
                q_start, q_end = q_end, q_start
            if s_start > s_end:
                flip = not flip
                s_start, s_end = s_end, s_start
            if flip:
                c = colors.Color(0, 0, 1, alpha=0.25)
                b = False
            else:
                c = colors.Color(1, 0, 0, alpha=0.25)
                b = False
            q_feature = q_set.add_feature(
                SeqFeature(SimpleLocation(q_start - 1, q_end)), color=c, border=b
            )
            s_feature = s_set.add_feature(
                SeqFeature(SimpleLocation(s_start - 1, s_end)), color=c, border=b
            )
            gd_diagram.cross_track_links.append(CrossLink(q_feature, s_feature, c, b))
            # NOTE: We are using the same colour for all the matches,
            # with transparency. This means overlaid matches will appear darker.
            # It also means the drawing order not very important.
            # Note ACT puts long hits at the back, and colours by hit score

print("Drawing CDS features...")
for f, format in genomes:
    record = records[f]
    feature_set = feature_sets[f]
    # Mark the CDS features
    for cds in record.features:
        if cds.type != "CDS":
            continue
        feature_set.add_feature(
            cds, sigil="ARROW", color=colors.lightblue, border=colors.blue
        )

gd_diagram.draw(
    format="linear", fragments=3, orientation="landscape", pagesize=(20 * cm, 10 * cm)
)
gd_diagram.write(name + ".pdf", "PDF")

gd_diagram.draw(format="circular", orientation="landscape", pagesize=(20 * cm, 20 * cm))
gd_diagram.write(name + "_c.pdf", "PDF")
