#!/opt/mambaforge/envs/bioconda/conda-bld/zol_1766863018713/_h_env_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_/bin/python

"""
Program: cgc
Author: Rauf Salamzade
Kalan Lab
UW Madison, Department of Medical Microbiology and Immunology
"""


# BSD 3-Clause License
#
# Copyright (c) 2023-2025, Kalan-Lab
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:   
#
# 1. Redistributions of source code must retain the above copyright notice, this
#    list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice, 
#    this list of conditions and the following disclaimer in the documentation
#    and/or other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its
#    contributors may be used to endorse or promote products derived from
#    this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, 
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

from collections import defaultdict
from operator import itemgetter
from rich_argparse import RawTextRichHelpFormatter
from zol import util
import argparse
import os
import shutil
import sys
import traceback

os.environ["OMP_NUM_THREADS"] = "1"
def create_parser(): 
    """ Parse arguments """
    parser = argparse.ArgumentParser(description = """
    Program: cgc
    Author: Rauf Salamzade
    Affiliation: Kalan Lab, UW Madison, Department of Medical Microbiology and
        Immunology

    collapsed gene clusters (cgc; pun: see(ಠಠ)-gc)

    cgc is a visualization tool that takes as input zol results and uses R libraries for
    plotting to generate a customizable figure that collapses information across hundreds
    to thousands of gene cluster instances in a "collapsed" figure.

    ---------------------------------------------------------------------------------------------
    Accepted tracks: 
    ---------------------------------------------------------------------------------------------
    - conservation, tajimas_d, entropy, upstream_entropy, median_beta_rd, max_beta_rd, median_gc, 
      bgc_score, viral_score, busted_pval.

    If comparative analysis was performed in zol, then addtional tracks include: 
    - focal_conservation, alternate_conservation, fst

    ---------------------------------------------------------------------------------------------
    Color scheme options: 
    ---------------------------------------------------------------------------------------------
    - white_to_black, red_white_blue, light_to_dark_gold, light_to_dark_blue, light_to_dark_green, 
    light_to_dark_purple, light_to_dark_red, grey, black, blue, red, purple, green, gold.

    Color scheme detailed customization can be performed by providing a tab delimited file with
    four columns: (1) track_name, 
        (2) low-value color (hex code), 
        (3) mid-value color (hex code), 

    (4) high-value color (hex-code).

    ---------------------------------------------------------------------------------------------
    Example command: 
    ---------------------------------------------------------------------------------------------
    cgc -i zol_Results_Directory/ -t conservation tajimas_d -c white_to_black blue_white_red

    """, formatter_class = RawTextRichHelpFormatter)

    parser.add_argument('-i', 
        '--zol-results-dir', 
        help = "Path to zol results directory.", 
        required = True)
    parser.add_argument('-o', 
        '--outdir', 
        help = "Output directory for cgc analysis.", 
        required = True)
    parser.add_argument('-t', 
        '--tracks', 
        nargs = '+', 
        help = "Tracks to use, note the first track is used to color the gene\n"
               "schematic directly. Should be provided in desired order (bottom\n"
               "to top) [Default is 'conservation tajimas_d']", 
        required = False, 
        default = ["conservation", 
        "tajimas_d"])
    parser.add_argument('-c', 
        '--color-schemes', 
        nargs = '+', 
        help = "Specify default color schemes avaibility by name in desired\n"
               "order matching --tracks [Default is 'white_to_black blue_white_red'].", 
        required = False, 
        default = ["white_to_black", 
        "red_white_blue"])
    parser.add_argument('-cc', 
        '--custom-color-scheme-file', 
        help = "Tab-delimited file with 4 columns: (1) track_name, (2) low-value\n"
               "color hexcode, (3) mid-value color hexcode, (4) high-value color hexcode.", 
        required = False, 
        default = None)
    parser.add_argument('-d', 
        '--annotation-db-faa', 
        help = "A FASTA file of proteins with headers to match to ortholog group\n"
               "consensus sequences. Will then use headers of matches (up until first\n"
               "space) as ortholog labels.", 
        required = False, 
        default = None)
    parser.add_argument('-ol', 
        '--og-label-file', 
        help = "Tab-delimited file with 2 columns: (1) ortholog group ID, (2) desired\n"
               "label.", 
        required = False, 
        default = None)
    parser.add_argument('-m', 
        '--min-conservation', 
        type = float, 
        help = "Minimum proportion of gene clusters an ortholog group needs to be found\n"
               "in to be shown in the figure [Default is 0.1].", 
        required = False, 
        default = 0.1)
    parser.add_argument('-rh', 
        '--relative-heights', 
        nargs = '+', 
        help = "Relative heights of individual tracks [Default is '1 1']", 
        required = False, 
        default = ["1", 
        "1"])
    parser.add_argument('-sl', 
        '--show-labels', 
        action = 'store_true', 
        help = "Show labels which by default are OG identifiers, if --og_label_file\n"
               "or --annotatIon_db_faa are not specified.", 
        required = False, 
        default = False)
    parser.add_argument('-ld', 
        '--label-distance', 
        type = float, 
        help = "The distance between genes and labels [Default is 0.025].", 
        required = False, 
        default = 0.025)
    parser.add_argument('-lts', 
        '--label-text-size', 
        type = float, 
        help = "Label text size [Default is 2.0].", 
        required = False, 
        default = 2.0)
    parser.add_argument('-sc', 
        '--show-copy-count-status', 
        action = 'store_true', 
        help = "Show whether genes are single-copy or not.", 
        required = False, 
        default = False)
    parser.add_argument('-q', 
        '--squish-factor', 
        type = float, 
        help = "A numeric value controlling vertical squishing of tracks in the plot - this\n"
               "value should be negative to squish or positive to give more space and\n"
               "proportional to the values provided for relative heights [Default is 0]", 
        required = False, 
        default = 0.0)
    parser.add_argument('-b', 
        '--bottom-spacing', 
        type = float, 
        help = "Extra space for bottom of the plot to make sure legend does not get cut off.\n"
               "Should be proportional to values for relative heights [Default is 0.1].", 
        required = False, 
        default = 0.1)
    parser.add_argument('-fl', 
        '--focal-set-title', 
        help = "Title to use for the focal set conservation track if requested.\n"
               "If spaces present, surround by quotes [Default is 'Focal Set Conservation'].", 
        required = False, 
        default = "Focal Set Conservation")
    parser.add_argument('-cl', 
        '--comparator-set-title', 
        help = "Title to use for the alternate set conservation track if requested. If spaces\n"
               "present, surround by quotes [Default is 'Comparator Set Conservation'].", 
        required = False, 
        default = "Comparator Set Conservation")
    parser.add_argument('-lgs', 
        '--legend-title-size', 
        type = float, 
        help = 'Change legend title size [Default is 10].', 
        required = False, 
        default = 10)
    parser.add_argument('-nl', 
        '--no-legend', 
        action = 'store_true', 
        help = "Do not show legend for coloring of the gene schematic.", 
        required = False, 
        default = False)
    parser.add_argument('-p', 
        '--png', 
        action = 'store_true', 
        help = "Create plot as PNG, default is PDF.", 
        required = False, 
        default = False)
    parser.add_argument('-l', 
        '--length', 
        type = int, 
        help = 'Specify the height/length of the heatmap plot in inches [Default is 7].', 
        required = False, 
        default = 7)
    parser.add_argument('-w', 
        '--width', 
        type = int, 
        help = 'Specify the width of the heatmap plot in inches [Default is 7].', 
        required = False, 
        default = 7)

    args = parser.parse_args()
    return args

INTEROG_DISTANCE = 100
default_color_schemes = {'white_to_black': ['#ffffff', '#878787', '#000000'], 
                         'red_white_blue': ['#bd3131', '#ffffff', '#548ce8'], 
                         'light_to_dark_gold': ['#fada89', '#c49723', '#705203'], 
                         'light_to_dark_blue': ['#a4c7fc', '#2c589c', '#0a2754'], 
                         'light_to_dark_green': ['#7ec47c', '#3d9e3a', '#106b0d'], 
                         'light_to_dark_purple': ['#b58df0', '#6a43a3', '#340e6b'], 
                         'light_to_dark_red': ['#f59795', '#a12e2b', '#630d0b'], 
                         'grey': ['#878787', '#878787', '#878787'], 
                         'black': ['#000000', '#000000', '#000000'], 
                         'red': ['#a12e2b', '#a12e2b', '#a12e2b'], 
                         'purple': ['#6a43a3', '#6a43a3', '#6a43a3'], 
                         'blue': ['#2c589c', '#2c589c', '#2c589c'], 
                         'green': ['#3d9e3a', '#3d9e3a', '#3d9e3a'], 
                         'gold': ['#c49723', '#c49723', '#c49723']}

track_name_to_full_label = {'conservation': 'Proportion of Total Gene Clusters with OG', 'tajimas_d': 'Tajima\'s D', 'entropy': 'Entropy', 
                            'upstream_entropy': 'Upstream Region Entropy', 'median_beta_rd': 'Median Beta-RD-gc', 'max_beta_rd': 'Max Beta-RD-gc', 
                            'focal_conservation': 'Proportion of Focal Gene Clusters with OG', 'fst': 'Fixation Index', 
                            'alterante_conservation': 'Proportion of Comparator Gene Clusters with OG', 
                            'bgc_score': 'BGC score (GECCO weights)', 'viral_score': 'Viral score (V-Score)', 
                            'busted_pval': 'P-value for gene-wide episodic selection by BUSTED'}

track_full_label_to_name = dict((v, k) for k, v in track_name_to_full_label.items())

def cgc(): 
    myargs = create_parser()

    zol_results_dir = os.path.abspath(myargs.zol_results_dir) + '/'
    outdir = os.path.abspath(myargs.outdir) + '/'
    tracks = myargs.tracks
    color_schemes = myargs.color_schemes
    custom_color_scheme_file = myargs.custom_color_scheme_file
    annotation_db_faa = myargs.annotation_db_faa
    og_label_file = myargs.og_label_file
    relative_heights = myargs.relative_heights
    show_labels = myargs.show_labels
    show_copy_count_status = myargs.show_copy_count_status
    plot_height = myargs.length
    plot_width = myargs.width
    min_conservation = myargs.min_conservation
    label_distance = myargs.label_distance
    label_text_size = myargs.label_text_size
    squish_factor = myargs.squish_factor
    bottom_spacing = myargs.bottom_spacing
    png_flag = myargs.png
    focal_set_title = myargs.focal_set_title
    comparator_set_title = myargs.comparator_set_title
    legend_title_size = myargs.legend_title_size
    no_legend = myargs.no_legend

    try: 
        assert (os.path.isdir(zol_results_dir))
    except Exception as e: 
        sys.stderr.write('The zol results directory required as input does not exist.\n')
        sys.exit(1)
    
    # create output directory if needed, or warn of over-writing
    if os.path.isdir(outdir):
        sys.stderr.write("Output directory exists. Files will be overwritten, but\n"
                       "checkpoints will be used to avoid redoing successfully\n"
                       "completed steps.\n"
                       "Do you wish to proceed? (yes/no): ")
        user_input = input().strip().lower()
        if user_input not in ['yes', 'y']:
            sys.stderr.write("Execution cancelled by user.\n")
            sys.exit(0)
    else:
        util.setup_ready_directory([outdir], delete_if_exist = False)

    # create logging object
    log_file = outdir + 'Progress.log'
    log_object = util.create_logger_object(log_file)
    version_string = util.get_version()

    sys.stdout.write(f'Running cgc version {version_string}\n')
    log_object.info(f'Running cgc version {version_string}')

    # log command used
    parameters_file = outdir + 'Command_Issued.txt'
    parameters_handle = open(parameters_file, 'a+')
    parameters_handle.write(' '.join(sys.argv) + '\n')
    parameters_handle.close()

    # Step 1: Parse relationships between gene clusters detected by fai and
    # prepTG genomes
    msg = '------------------Step 1------------------\nAssessing input provided.'
    log_object.info(msg)
    sys.stdout.write(msg + '\n')

    zol_spreadsheet_tsv = zol_results_dir + 'Final_Results/Consolidated_Report.tsv'

    try: 
        assert(os.path.isfile(zol_spreadsheet_tsv))
    except Exception as e: 
        msg = 'The "Final_Results/Consolidated_Report.tsv" table results from zol analysis is not available in the zol directory provided.'
        log_object.error(msg)
        sys.stderr.write(msg + '\n')
        sys.exit(1)

    try: 
        assert(type(tracks) is list and type(color_schemes) is list and type(relative_heights) is list)
        assert(len(tracks) == len(color_schemes) and len(tracks) == len(relative_heights))
    except Exception as e: 
        msg = 'The tracks, color_schemes, and relative_heights arguments values are not\nlists as expected or are not lists of the same length.'
        log_object.error(msg)
        sys.stderr.write(msg + '\n')
        sys.exit(1)

    try: 
        for t in tracks: 
            if not t in track_name_to_full_label: 
                msg = t + ' not an accepted track option!'
                log_object.error(msg)
                sys.stderr.write(msg + '\n')
                raise RuntimeError()
    except Exception as e: 
        msg = 'One or more of the tracks is not an accepted option, please check the --help/usage for cgc.'
        log_object.error(msg)
        sys.stderr.write(msg + '\n')
        sys.exit(1)

    try: 
        for c in color_schemes: 
            if not c in default_color_schemes: 
                msg = c + ' not an accepted default color scheme option! '
                log_object.error(msg)
                sys.stderr.write(msg + '\n')
                raise RuntimeError()
    except Exception as e: 
        msg = 'One or more of the default color scheme is not an accepted option, please check the --help/usage for cgc.'
        log_object.error(msg)
        sys.stderr.write(msg + '\n')
        sys.exit(1)

    color_list = []
    if custom_color_scheme_file != None: 
        msg = 'A custom color scheme file is provided, ignoring values provided to the argument --color_schemes if provided.'
        log_object.error(msg)
        sys.stderr.write(msg + '\n')

        try: 
            track_colors = {}
            with open(custom_color_scheme_file) as occsf: 
                for line in occsf: 
                    line = line.strip()
                    ls = line.split('\t')
                    track_colors[ls[0]] = [ls[1], ls[2], ls[3]]
            for t in tracks: 
                cs_info = track_colors[t]
                color_list.append(cs_info)
        except Exception as e: 
            sys.stderr.write(traceback.format_exc())
            msg = "Issue reading the custom color scheme file or perhaps colors for\nall tracks requested were not provided. This should be a\nheader-less tab-delmited file with four columns, with track\nname as first column andlow, mid, high value colors as second, third, and fourth columns."
            log_object.error(msg)
            sys.stderr.write(msg + '\n')
            sys.exit(1)
    else: 
        for cs in color_schemes: 
            cs_info = default_color_schemes[cs]
            color_list.append(cs_info)

    # Step 2: Sorting out labels
    og_labels = defaultdict(lambda: '')
    if show_labels or og_label_file != None or annotation_db_faa != None: 
        msg = '------------------Step 1.5------------------\nSorting out ortholog labels.'
        log_object.info(msg)   
        sys.stdout.write(msg + '\n')

    if og_label_file: 
        try: 
            with open(og_label_file) as olf: 
                for i, line in enumerate(olf): 
                    line = line.strip()
                    ls = line.split('\t')
                    og_labels[ls[0]] = ls[1]
        except Exception as e: 
            sys.stderr.write(traceback.format_exc())
            msg = "Issue reading the custom labels for ortholog groups from the provided file."
            log_object.error(msg)
            sys.stderr.write(msg + '\n')
            sys.exit(1)

        if annotation_db_faa:
            og_consensus_seqs_faa = zol_results_dir + 'OG_Consensus_Seqs.faa'

            try: 
                assert(util.is_fasta(annotation_db_faa))
            except Exception as e: 
                msg = "Issue validating the annotation_db_faa file exists or is in FASTA format"
                log_object.error(msg)
                sys.stderr.write(msg + '\n')
                sys.exit(1)

            try: 
                assert(util.is_fasta(og_consensus_seqs_faa))
            except Exception as e: 
                msg = "Issue validating the 'OG_Consensus_Seqs.faa' file exists in the zol results directory or is in FASTA format"
                log_object.error(msg)
                sys.stderr.write(msg + '\n')
                sys.exit(1)

            mapping_dir = outdir + 'Protein_Mapping/'
            proteomes_dir = mapping_dir + 'Inputs/'
            util.setup_ready_directory([mapping_dir, proteomes_dir], delete_if_exist = True)
            try: 
                shutil.copy(annotation_db_faa, proteomes_dir + '1_queries.faa')
                shutil.copy(og_consensus_seqs_faa, proteomes_dir + '2_references.faa')

                mapping_results_dir = mapping_dir + 'Results/'
                mapping_result_file = mapping_results_dir + 'Orthogroups.tsv'
                find_ogs_cmd = ['findOrthologs.py', '-p', proteomes_dir, '-o', mapping_results_dir, '-c', '1']
                util.run_cmd_via_subprocess(find_ogs_cmd, log_object=log_object, 
                                            check_files = [mapping_result_file])

                with open(mapping_result_file) as omrf: 
                    for i, line in enumerate(omrf): 
                        if i == 0: continue
                        line = line.strip('\n')
                        query, og = line.split('\t')[1: ]
                        # get one to one mappings only!
                        if og.strip() != '' and len(og.split(', ')) == 1 and query != '' and len(query.split()) == 1: 
                            og_labels[og.strip()] = query.strip()

            except Exception as e: 
                sys.stderr.write(traceback.format_exc())
                msg = "Issue reading the custom labels for ortholog groups from the provided file."
                log_object.error(msg)
                sys.stderr.write(msg + '\n')
                sys.exit(1)
        else: 
            try: 
                with open(zol_spreadsheet_tsv) as ozst: 
                    for i, line in enumerate(ozst): 
                        if i == 0: continue
                        line = line.strip('\n')
                        ls = line.split('\t')
                        og = ls[0]
                        og_labels[og] = og
            except Exception as e: 
                sys.stderr.write(traceback.format_exc())
                msg = f"Issue reading the zol consolidated spreadsheet: {zol_spreadsheet_tsv}"
                log_object.error(msg)
                sys.stderr.write(msg + '\n')
                sys.exit(1)

    # Step 3: Parsing out data from zol results
    msg = '------------------Step 2------------------\nParsing out data from zol results and creating '
    msg += 'plot input.'
    log_object.info(msg)
    sys.stdout.write(msg + '\n')

    track_set = set(tracks)
    dataframe = outdir + 'Plot_Input.txt'
    not_scc_count = 0
    try: 
        df_handle = open(dataframe, 'w')

        og_print_data = {}
        og_pos_data = []
        select_indices = set([])
        with open(zol_spreadsheet_tsv) as ozst: 
            for i, line in enumerate(ozst): 
                line = line.strip('\n')
                ls = line.split('\t')
                if i == 0: 
                    printlist = [ls[0]]
                    for j, val in enumerate(ls[1: ]): 
                        if val in track_full_label_to_name and track_full_label_to_name[val] in track_set: 
                            printlist.append(track_full_label_to_name[val])
                            select_indices.add(j)
                        elif val == 'OG is Single Copy?': 
                            printlist.append('single_copy')
                            select_indices.add(j)
                    printlist += ['x_start', 'x_end', 'x_midpoint', 'direction', 'y_start', 'label']
                    df_handle.write('\t'.join(printlist) + '\n')
                else: 
                    og = ls[0]
                    printlist = [og]
                    for j, val in enumerate(ls[1: ]): 
                        if j in select_indices: 
                            printlist.append(val)
                    og_print_data[og] = printlist
                    med_len = float(ls[3])
                    cons_order = int(ls[4])
                    cons_dir = 1
                    if ls[5] == '"-"': 
                        cons_dir = 0
                    conservation = float(ls[2])
                    if conservation > min_conservation: 
                        if ls[1] == 'False': 
                            not_scc_count += 1
                        og_pos_data.append([og, med_len, cons_order, cons_dir])

        previous_end = 1.0
        for op in sorted(og_pos_data, key = itemgetter(2)): 
            og, med_len, cons_order, cons_dir = op
            x_start = previous_end
            x_end = x_start + med_len
            x_midpoint = x_start + (x_end - x_start)/ 2.0
            opd = og_print_data[og] + [str(x) for x in [x_start, 
          x_end, 
          x_midpoint, 
          cons_dir, 
          0, 
          og_labels[og]]]
            df_handle.write('\t'.join(opd) + '\n')
            previous_end = x_end + INTEROG_DISTANCE
        df_handle.close()
    except Exception as e: 
        sys.stderr.write(traceback.format_exc())
        msg = f"Issue reading the zol consolidated spreadsheet: {zol_spreadsheet_tsv}"
        log_object.error(msg)
        sys.stderr.write(msg + '\n')
        sys.exit(1)

    # Step 4: Write R script to generate the final figure
    msg = '------------------Step 4------------------\nWriting R script for generating figure PDF/PNG.'
    log_object.info(msg)
    sys.stdout.write(msg + '\n')

    # flip as if you were writing top to bottom instead.
    track_list = tracks[::-1]
    color_list = color_list[::-1]
    relhei_list = relative_heights[::-1]

    if show_copy_count_status and not_scc_count > 0: 
        last_rh = float(relhei_list[-1])
        sc_rh = 0.1*last_rh
        gg_rh = 0.9*last_rh
        relhei_list = relhei_list[: -1] + [str(sc_rh), str(gg_rh)]

    relhei_list_updated = []
    for i, rh in enumerate(relhei_list): 
        if i == 0: 
            relhei_list_updated.append(rh)
        else: 
            relhei_list_updated.append(str(squish_factor))
            relhei_list_updated.append(rh)

    # new line character
    nl = '\n'

    rscript_path = outdir + 'cgc_script.R'
    result_pdf_file = outdir + 'cgc_plot.pdf'
    if png_flag: 
        result_pdf_file = outdir + 'cgc_plot.png'
    rscript_handle = open(rscript_path, 'w')

    rscript_handle.write('library(ggplot2)' + nl)
    rscript_handle.write('library(cowplot)' + nl)
    rscript_handle.write('library(gggenes)' + nl)
    rscript_handle.write(nl)
    rscript_handle.write('dat <- read.table("' + dataframe + '", header = T, sep = "\\t")' + nl)
    rscript_handle.write(nl)

    ggplot_objects = []
    for i, t in enumerate(track_list): 
        cs = color_list[i]
        gobj_id = 'g' +  str(i+1)
        if i < len(track_list)-1: 
            rscript_handle.write('# track for ' + t + nl)
            rscript_handle.write(gobj_id + ' <- ggplot(dat, aes(xmin = x_start, xmax = x_end, ymin = y_start, ymax = ' + t + ', fill = ' + t + ')) + theme_classic() + ' + nl)
            rscript_handle.write('geom_rect(color = "black", show.legend = F) + scale_fill_gradient2(low = "' + cs[0] + '", mid = "' + cs[1] + '", high = "' + cs[2] + '", na.value = "grey50") +' + nl)
            rscript_handle.write('theme(axis.line.x = element_line(colour = "white"), axis.title.x = element_blank(), axis.text.x = element_blank(), axis.ticks.x = element_blank()) +' + nl)
            if t == 'focal_conservation': 
                rscript_handle.write('ggtitle("' + focal_set_title + '")' + nl)
            elif t == 'alternate_conservation': 
                rscript_handle.write('ggtitle("' + comparator_set_title + '")' + nl)
            else: 
                rscript_handle.write('ggtitle("' + track_name_to_full_label[t] + '")' + nl)
            rscript_handle.write(nl)
        else: 
            if show_copy_count_status and not_scc_count > 0: 
                rscript_handle.write('# track for single-copy status' + nl)
                ggplot_objects.append('gsc')
                ggplot_objects.append("NULL")
                rscript_handle.write('xmin <- min(dat$x_start)' + nl)
                rscript_handle.write('xmax <- max(dat$x_end)' + nl)
                rscript_handle.write('dat.sc <- dat[dat$single_copy == "False", ]' + nl)
                rscript_handle.write('gsc <- ggplot(dat.sc, aes(x = x_midpoint, y = 0, color = single_copy)) + theme_void() + geom_point(color = "#a33f37", show.legend = F) + xlim(xmin, xmax)' + nl)
                rscript_handle.write(nl)
            rscript_handle.write('# track for ' + t + nl)
            rscript_handle.write(gobj_id + ' <- ggplot(dat, aes(xmin = x_start, xmax = x_end, y = 1, label = label, forward = direction)) +' + nl)
            rscript_handle.write('geom_gene_arrow(aes(fill = ' + t + '), color = "black") + theme_void() + scale_fill_gradient2(low = "' + cs[0] + '", mid = "' + cs[1] + '", high = "' + cs[2] + '", na.value = "grey50", guide = "colourbar", aesthetics = "fill") + ' + nl)
            rscript_handle.write(f'geom_text(aes(x = x_midpoint), angle = 45, y = {1-label_distance}, size = {label_text_size}) +' + nl)
            if no_legend: 
                rscript_handle.write('theme(legend.position = "none")' + nl)
            else: 
                rscript_handle.write(f'theme(legend.position = "bottom", legend.title = element_text(size = {legend_title_size}), legend.box = "horizontal")' + nl)
            rscript_handle.write(nl)
        ggplot_objects.append(gobj_id)
        ggplot_objects.append("NULL")

    rscript_handle.write(nl)
    if png_flag: 
        rscript_handle.write(f'png("{result_pdf_file}", height = {plot_height}, width = {plot_width}, res = 600, units = "in")' + nl)
    else: 
        rscript_handle.write(f'pdf("{result_pdf_file}", height = {plot_height}, width = {plot_width})' + nl)
    rscript_handle.write(f'print(plot_grid({", ".join(ggplot_objects)}, NULL, ncol = 1, axis = "lr", align = "v", rel_heights = c({", ".join(relhei_list_updated)}, {bottom_spacing})))' + nl)
    rscript_handle.write('dev.off()' + nl)

    rscript_handle.close()

    # Step 5: Running R script to generate the final figure
    msg = '------------------Step 5------------------\nRunning R script for generating figure PDF/PNG.'
    log_object.info(msg)
    sys.stdout.write(msg + '\n')

    plot_cmd = ['Rscript', rscript_path]
    util.run_cmd_via_subprocess(plot_cmd, log_object=log_object, 
                                check_files = [result_pdf_file])

    msg = f'cgc finished successfully!\nResulting PDF can be found at: {result_pdf_file}'
    log_object.info(msg)
    sys.stdout.write(msg + '\n')
    sys.exit(0)

if __name__ == '__main__': 
    cgc()
