#!/usr/bin/env python3
"""
This file is part of MitoZ.

THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH
REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY
AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT,
INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM
LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE
OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
PERFORMANCE OF THIS SOFTWARE.

COPYRIGHT © 2019-2022 Guanliang Meng. ALL RIGHTS RESERVED.

"""
import sys
import os
import argparse
import subprocess
import time
from Bio import SeqIO
from pathlib import Path


def add_arguments(parser):
    parser.add_argument('--circos',
        metavar='<str>',
        default='circos',
        help='''absolute path of circos executable. otherwise, `circos` must be in your `PATH` variable [%(default)s]''')

    parser.add_argument('--gb',
        metavar='<file>',
        required=True,
        help='Your input Genbank file [required]')

    # whether draw GC content circle
    gc_track_group = parser.add_argument_group('''GC content track (the innermost one)
        In the resulting GC track (the innermost one), 
        the red line means 0.5, while each black lines mean 0.05. ''')

    gc_track_group.add_argument('--gc',
        choices=['yes', 'no'],
        default='no',
        help='whether draw GC content track. [%(default)s]')

    gc_track_group.add_argument('--win',
        metavar='<str>',
        default=50,
        type=int,
        help='windows size for calculating GC content [%(default)s]')

    gc_track_group.add_argument('--gc_fill',
        metavar='<str>',
        default='128,177,211',
        help='color for filling the GC track [%(default)s]')

    # whether draw depth abundance circle
    abundance_track_group = parser.add_argument_group('''Abundance track (the middle one)
        This reveals the depth distribution, with dark green for the outline.
        If the depth lower than the minimum value (default 20), it turns red, 
        whereas if the depth larger than upper quartile,
        it turns dark green as same with the outline''')

    abundance_track_group.add_argument('--depth_file',
        metavar='<file>',
        default='',
        help='''A file of
tab-separated table with three columns: reference name, position, and coverage
depth, which is generated by 'samtools depth' command. Or you can set
'--run_map yes' and '--fq1' and '--fq2' to get this information''')

    abundance_track_group.add_argument('--run_map',
        choices=['yes', 'no'],
        default='no',
        help='''If you do NOT have the '--depth_file', setting '--run_map yes' as well as '--fq1' and '--fq2' to map reads to references to get the '--depth_file' [%(default)s]''')

    abundance_track_group.add_argument('--fq1',
        metavar='<file>',
        default='',
        help="fastq 1 file. Do NOT forget to set '--run_map yes'")

    abundance_track_group.add_argument('--fq2',
        metavar='<file>',
        default='',
        help="fastq 2 file. Do NOT forget to set '--run_map yes'")

    abundance_track_group.add_argument('--bwa',
        metavar='<str>',
        default='bwa',
        help='''absolute path of bwa executable,
otherwise, `bwa` must be in your `PATH` variable [%(default)s]'''),

    abundance_track_group.add_argument('--thread',
        metavar='<int>',
        default=2,
        help="BWA thread number [%(default)s]")

    abundance_track_group.add_argument('--samtools',
        metavar='<str>',
        default='samtools',
        help='''absolute path of samtools executable,
otherwise, `samtools` must be in your `PATH` variable [%(default)s]''')

    abundance_track_group.add_argument('--opts_samtools',
        metavar='<str>',
        default="-a -a",
        help='optional arguments for samtools [%(default)s]')

    abundance_track_group.add_argument('--depth_fill',
        metavar='<str>',
        default='190,186,218',
        help='color for filling the abundance track [%(default)s]')

    # color setting
    color_group = parser.add_argument_group('Color setting (the outermost one)')

    color_group.add_argument('--cds_color',
        metavar='<str>',
        default='141,211,199',
        help='PCG color [%(default)s]')

    color_group.add_argument('--trna_color',
        metavar='<str>',
        default='251,128,114',
        help='tRNA color [%(default)s]')

    color_group.add_argument('--rrna_color',
        metavar='<str>',
        default='253,192,134',
        help='rRNA color [%(default)s]')

    color_group.add_argument('--label_color',
        metavar='<str>',
        default='black',
        help="gene name label color [%(default)s]")

    color_group.add_argument('--locus_color',
        metavar='<str>',
        default='black',
        help="locus name's clolor showed on center of circle [%(default)s]")

    # whether draw base around circle
    base_track_group = parser.add_argument_group('Base track')
    base_track_group.add_argument('--base',
        choices=['yes', 'no'],
        default='no',
        help='whether draw base track [%(default)s]')


    # outfile setting
    outfile_group = parser.add_argument_group('Outfile setting')

    outfile_group.add_argument('--bgc',
        metavar='<str>',
        default='white',
        help="background color or the path to a user-defined image file (as the background of the resulting file!) [%(default)s]")

    outfile_group.add_argument('--outdir',
        metavar='<str>',
        default='./outdir',
        help="output directory [%(default)s]")

    outfile_group.add_argument('--png',
        choices=['yes', 'no'],
        default='yes',
        help=argparse.SUPPRESS)
        # help='PNG format output [%(default)s]')

    outfile_group.add_argument('--svg',
        choices=['yes', 'no'],
        default='yes',
        help=argparse.SUPPRESS)
        # help='SVG format output [%(default)s]')

    return parser


def generate_config_file(config_file='visual_config.txt', args=None):
    config_content = '''
# where circos you install
circos_path     =   {circos}
bwa             =   {bwa}
samtools        =   {samtools}

threads         =   {thread}

# optional arguments for samtools
opts_samtools   =  {opts_samtools}

#color
cds             =   {cds_color}
rRNA            =   {rrna_color}
tRNA            =   {trna_color}

# whether draw GC content circle
gc              =   {gc}
win             =   {win}
gc_fill         =   {gc_fill}

# whether draw depth abundance circle
run_map           =   {run_map}
depth_file      =   {depth_file}
depth_fill      =   {depth_fill}
fq              =   {fq1},{fq2}

# whether draw base around circle
base            =   {base}

# locus name's clolor showed on center of circle
locus_color     =   {locus_color}

# gene name label color
label_color     =   {label_color}

# image
outdir          =   {outdir}
png             =   {png}
svg             =   {svg}

# color or file
background      =   {bgc}
#background     =   ./background.png

    '''.format(
            circos=args.circos, 
            bwa=args.bwa,
            samtools=args.samtools,
            thread=args.thread,
            opts_samtools=args.opts_samtools,
            cds_color=args.cds_color,
            rrna_color=args.rrna_color,
            trna_color=args.trna_color,
            gc=args.gc,
            win=args.win,
            gc_fill=args.gc_fill,
            run_map=args.run_map,
            depth_file=args.depth_file,
            depth_fill=args.depth_fill,
            fq1=args.fq1,
            fq2=args.fq2,
            base=args.base,
            locus_color=args.locus_color,
            label_color=args.label_color,
            outdir=args.outdir,
            png=args.png,
            svg=args.svg,
            bgc=args.bgc)

    with open(config_file, 'w') as fhout:
        print(config_content, file=fhout)

    # logger.info("circos using this config_file: " + os.path.abspath(config_file))
    return os.path.abspath(config_file)


def run_wrapped_circos(perl=None, script=None, gb_file=None, config_file=None, logger=None):
    cmd = '''{perl} {script} --gb {gb_file} --conf {config_file}
'''.format(perl=perl, script=script, gb_file=gb_file, config_file=config_file)

    runcmd(cmd, logger=logger)


def runcmd(command, logger=None):
    try:
        if logger:
            logger.info(command)
        subprocess.check_call(command, shell=True)
    except Exception as e:
        m = "Error occured when running command:\n{0}\n".format(command, e)
        if logger:
            logger.error(m)
        else:
            sys.exit(m)

def sort_gb_by_length(in_gb=None, out_gb=None, logger=None):
    logger.info("sort_gb_by_length():\nin_gb: {in_gb}\nout_gb: {out_gb}".format(in_gb=in_gb, out_gb=out_gb))
    seqid_len_rec = {}
    for rec in SeqIO.parse(in_gb, 'gb'):
        seqlen = len(rec)
        seqid_len_rec[str(rec.id)] = (seqlen, rec)

    fhout = open(out_gb, 'w')
    for seqid, seqlen_rec in sorted(seqid_len_rec.items(), key=lambda x: x[1][0]):
        rec = seqlen_rec[1]
        SeqIO.write(rec, fhout, 'gb')
    fhout.close()

    return out_gb


def main(args):
    logger = args.logger
    logger.info("visualize.main() got args:\n{}".format(args))

    args.gb = os.path.abspath(args.gb)
    args.outdir = os.path.abspath(args.outdir)

    if args.fq1 and Path(args.fq1).is_file():
        args.fq1 = os.path.abspath(args.fq1)
    else:
        args.fq1 = ''

    if args.fq2 and Path(args.fq2).is_file():
        args.fq2 = os.path.abspath(args.fq2)
    else:
        args.fq2 = ''

    if not args.fq1 and not args.fq2:
        args.run_map = 'no'

    logger.info("visualize.main() got updated args:\n{}".format(args))

    if not os.path.exists(args.outdir):
        logger.info("visualize.main() creating: {}".format(args.outdir))
        os.makedirs(args.outdir, exist_ok=True)

    logger.info("visualize.main() chdir to: {}".format(args.outdir))
    os.chdir(args.outdir)

    out_gb = os.path.basename(args.gb) + '.sortedByLength'
    sort_gb_by_length(in_gb=args.gb, out_gb=out_gb, logger=logger)
    args.gb = os.path.abspath(out_gb)
    logger.info("updated args.gb: {}".format(args.gb))

    config_file = generate_config_file(config_file='visual_config.txt', args=args)
    logger.info("circos using this config_file: " + config_file)

    script = os.path.join(os.path.dirname(__file__),
        'draw_circos_for_mitogenome_auto_depth.pl')
    run_wrapped_circos(
        perl='perl',
        script=script,
        gb_file=args.gb,
        config_file=config_file,
        logger=logger)

    # only to testing, because the within the Perl script,
    # circos does not draw.
    # not drawing is due to circos can't find GD module.
    # The solutin is just to run this 'circos -conf circos.conf' command again from user's commandline!!
    # circos_config_file = os.path.join(args.outdir, 'circos.conf')
    # command = 'circos -conf ' + circos_config_file
    # runcmd(command)
    logger.info("visualize.main() returns: {}".format(args.outdir))
    return args.outdir

















