#!/opt/conda/conda-bld/rbpbench_1726242255020/_h_env_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_plac/bin/python

import argparse
import os
import re
import gzip
from rbpbench import benchlib


###############################################################################

def setup_argument_parser():
    """Setup argparse parser."""
    help_description = """
    Extract most prominent transcript (MPT) regions from GTF file 
    and create BED file with MPT regions.

    """
    # Define argument parser.
    p = argparse.ArgumentParser(add_help=False,
                                prog="gtf_extract_mpt_region_bed.py",
                                description=help_description,
                                formatter_class=argparse.MetavarTypeHelpFormatter)

    # Required arguments.
    p.add_argument("-h", "--help",
                   action="help",
                   help="Print help message")
    p.add_argument("--gtf",
                   dest="in_gtf",
                   type=str,
                   metavar='str',
                   required = True,
                   help = "Input GTF file with genomic annotations to extract MPT regions from. Note that only genes on standard chromosomes (1,2,..,X,Y,MT) are currently used")
    p.add_argument("--out",
                   dest="out_folder",
                   type=str,
                   metavar='str',
                   required=True,
                   help="Output folder to store BED files")
    p.add_argument("--tr-list",
                   dest="tr_list",
                   type=str,
                   metavar='str',
                   default = False,
                   help = "Supply file with transcript IDs (one ID per row) to define which transcripts to extract from --gtf (overrides MPT selection)")
    p.add_argument("--mrna-only",
                   dest="only_mrna",
                   default = False,
                   action = "store_true",
                   help = "Set if only mRNAs should be extracted from --gtf file (default: False)")
    p.add_argument("--bed-col4-infos",
                   dest="bed_col4_infos",
                   type=int,
                   default=1,
                   choices=[1, 2],
                   help="Define what to store in BED column 4. 1: store transcript_id. 2: store transcript_id;gene_id (default: 1)")
    p.add_argument("--chr-id-style",
                   dest="chr_id_style",
                   type=int,
                   default=1,
                   choices=[1, 2, 3],
                   help="Define to which chromosome ID style to convert chromosome IDs to. 1: do not change chromosome IDs. 2: convert to chr1,chr2,...,chrM style. 3: convert to 1,2,...,MT style (default: 1)")
    return p


################################################################################

if __name__ == '__main__':

    parser = setup_argument_parser()
    args = parser.parse_args()

    assert os.path.exists(args.in_gtf), "--gtf file \"%s\" not found" % (args.in_gtf)

    if not os.path.exists(args.out_folder):
        os.makedirs(args.out_folder)

    mrna_regions_bed = args.out_folder + "/mrna_regions.bed"  # UTR CDS regions on mRNAs (i.e. transcript context).
    out_mpt_bed = args.out_folder + "/mpt_regions.bed"  # MPT regions (i.e. transcript context).


    print("Read in gene features from --gtf ... ")
    tr2gid_dic = {}
    tr_types_dic = {}  # Store transcript biotypes in GTF file.
    gid2gio_dic = benchlib.gtf_read_in_gene_infos(args.in_gtf,
                                                tr2gid_dic=tr2gid_dic,
                                                tr_types_dic=tr_types_dic,
                                                chr_style=args.chr_id_style,
                                                empty_check=False)
    assert gid2gio_dic, "no gene infos read in from --gtf. Please provide a valid/compatible GTF file (e.g. from Ensembl or ENCODE)"
    c_gene_infos = len(gid2gio_dic)
    print("# gene features read in from --gtf:", c_gene_infos)

    # Get most prominent transcripts or if --tr-list is set, read in transcript IDs.
    tr_ids_dic = {}
    if args.tr_list:
        tr_ids_dic = benchlib.read_ids_into_dic(args.tr_list,
                                                check_dic=False)
        assert tr_ids_dic, "no IDs read in from provided --tr-list file. Please provide a valid IDs file (one ID per row)"
        for tr_id in tr_ids_dic:
            assert tr_id in tr2gid_dic, "transcript ID \"%s\" from provided --tr-list file does not appear in --gtf file. Please provide compatible IDs + files" %(tr_id)
            tr_ids_dic[tr_id] = tr2gid_dic[tr_id]
        print("# of transcript IDs (read in from --tr-list): ", len(tr_ids_dic))
    else:
        # Get most prominent transcripts from gene infos.
        print("Select most prominent transcript (MPT) for each gene ... ")
        tr_ids_dic = benchlib.select_mpts_from_gene_infos(gid2gio_dic,
                                basic_tag=False,  # do not be strict (only_tsl=False too).
                                ensembl_canonical_tag=False,
                                prior_basic_tag=True,  # Prioritize basic tag transcript.
                                only_tsl=False)
        assert tr_ids_dic, "most prominent transcript selection from gene infos failed. Please contact developers"
        print("# of transcript IDs (most prominent transcripts): ", len(tr_ids_dic))

    # Check exon order (return True if minus strand exon 1 is most downstream, not most upstream, which is the correct way).
    print("Check minus-strand exon order in --gtf ... ")
    correct_min_ex_order = benchlib.gtf_check_exon_order(args.in_gtf)
    if correct_min_ex_order:
        print("Correct order encountered ... ")
    else:
        print("Reverse order encountered ... ")

    # Get transcript infos.
    print("Read in transcript infos from --gtf ... ")
    tid2tio_dic = benchlib.gtf_read_in_transcript_infos(args.in_gtf, 
                                                        tr_ids_dic=tr_ids_dic,
                                                        correct_min_ex_order=correct_min_ex_order,
                                                        chr_style=args.chr_id_style,
                                                        empty_check=False)

    assert tid2tio_dic, "no transcript infos read in from --gtf. Please provide a valid/compatible GTF file (e.g. from Ensembl or ENCODE)"

    # (in)sanity checks.
    for tr_id in tr_ids_dic:
        assert tr_id in tid2tio_dic, "transcript ID %s not in tid2tio_dic"
    for tr_id in tid2tio_dic:
        assert tr_id in tr_ids_dic, "transcript ID %s not in tr_ids_dic"

    c_tr_infos = len(tid2tio_dic)
    print("# transcript features read in from --gtf:", c_tr_infos)

    # If --only-mrna, only select mRNAs, which also triggers mRNA region occupancy plots generation.
    tid2regl_dic = {}
    if args.only_mrna:
        # Get mRNA region lengths (from transcript IDs with CDS feature).
        print("Get mRNA region lengths ... ")
        # tid2tio_dic contains infos for transcripts in tr_ids_dic (i.e. either MPT or --tr-list defined ones).
        tid2regl_dic = benchlib.get_mrna_region_lengths(tid2tio_dic)
        assert tid2regl_dic, "tid2regl_dic empty. If --tr-list was set, this means none of supplied transcript IDs contain a CDS. If --tr-list was not set, this means that none of the MPTs contain a CDS. In this case please provide a valid/compatible GTF file (e.g. from Ensembl or ENCODE) or contact developers"
        c_mrna_tids = len(tid2regl_dic)
        if args.tr_list:
            print("# mRNA transcripts (containing CDS) from --tr-list:", c_mrna_tids)
        else:
            print("# mRNA transcripts (containing CDS, out of MPT selected set):", c_mrna_tids)

        assert c_mrna_tids, "no mRNA transcripts (containing CDS) found in --gtf. Please provide a valid/compatible GTF file (e.g. from Ensembl or ENCODE). Alternatively, if --tr-list was given, make sure that the list contains mRNA transcripts, or do not set --only-mrna"

        # Output mRNA regions (5'UTR CDS 3'UTR) to BED.
        print("Output mRNA regions to BED ... ")
        benchlib.output_mrna_regions_to_bed(tid2regl_dic, mrna_regions_bed)
    
        # Assign mRNA transcripts as transcripts to extract sequences for.
        tr_ids_dic = tid2regl_dic
        
    """
    Output MPT regions to BED file.

    """

    print("Output MPT regions to BED ... ")

    c_out = 0
    OUTBED = open(out_mpt_bed, "w")

    for tid in tid2tio_dic:
        tio = tid2tio_dic[tid]
        chr_id = tio.chr_id
        tr_s = tio.tr_s - 1
        tr_e = tio.tr_e
        tr_pol = tio.tr_pol
        gene_id = tio.gene_id

        if tid not in tr_ids_dic:
            continue

        reg_id = tid
        if args.bed_col4_infos == 2:
            reg_id = tid + ";" + gene_id

        c_out += 1

        OUTBED.write("%s\t%i\t%i\t%s\t0\t%s\n" %(chr_id, tr_s, tr_e, reg_id, tr_pol))

    OUTBED.close()

    print("# MPT regions output to BED file:", c_out)

    print("Output BED file with MPT regions:\n%s" %(out_mpt_bed))
    if args.only_mrna:
        print("Output BED file with mRNA regions:\n%s" %(mrna_regions_bed))

    print("Done.")


################################################################################
