# Jeremie Kalfon
# for BroadInsitute
# in 2019

from __future__ import print_function
from multiprocessing.sharedctypes import Value
import os
import signal
import re

import pandas as pd
import numpy as np

from genepy.google import gcp
from genepy.utils import helper as h
from tqdm import tqdm

size = {"GRCh37": 2864785220, "GRCh38": 2913022398}

cmaps = [
    "Greys",
    "Purples",
    "Blues",
    "Greens",
    "Oranges",
    "Reds",
    "YlOrBr",
    "YlOrRd",
    "OrRd",
    "PuRd",
    "RdPu",
    "BuPu",
    "GnBu",
    "PuBu",
    "YlGnBu",
    "PuBuGn",
    "BuGn",
    "YlGn",
]

chroms = {
    "chr1",
    "chr10",
    "chr11",
    "chr12",
    "chr13",
    "chr14",
    "chr15",
    "chr16",
    "chr17",
    "chr18",
    "chr19",
    "chr2",
    "chr20",
    "chr21",
    "chr22",
    "chr3",
    "chr4",
    "chr5",
    "chr6",
    "chr7",
    "chr8",
    "chr9",
    "chrX",
    "chrY",
    "1",
    "10",
    "11",
    "12",
    "13",
    "14",
    "15",
    "16",
    "17",
    "18",
    "19",
    "2",
    "20",
    "21",
    "22",
    "3",
    "4",
    "5",
    "6",
    "7",
    "8",
    "9",
    "X",
    "Y",
}


def fromGTF2BED(gtfname, bedname, gtftype="geneAnnot"):
    """
    transforms a  gtf file into a bed file

    Args:
    ----
      gtfname: filepath to gtf file
      bedname: filepath to beddfile
      gtftype: only geneAnnot for now

    Returns:
    --------
      newbed: the bedfile as a pandas.df

    """
    if gtftype == "geneAnnot":
        gtf = pd.read_csv(
            gtfname,
            sep="\t",
            header=0,
            names=[
                "chr",
                "val",
                "type",
                "start",
                "stop",
                "dot",
                "strand",
                "loc",
                "name",
            ],
        )
        gtf["name"] = [
            i.split('gene_id "')[-1].split('"; trans')[0] for i in gtf["name"]
        ]
        prevname = ""
        newbed = {"chr": [], "start": [], "end": [], "gene": []}
        for i, val in gtf.iterrows():
            h.showcount(i, len(gtf))
            if val["name"] == prevname:
                newbed["end"][-1] = val["stop"]
            else:
                newbed["chr"].append(val["chr"])
                newbed["start"].append(val["start"])
                newbed["end"].append(val["stop"])
                newbed["gene"].append(val["name"])
            prevname = val["name"]
        newbed = pd.DataFrame(newbed)
        newbed = newbed[~newbed.chr.str.contains("_fix")]
        newbed.to_csv(bedname + ".bed", sep="\t", index=None)
        newbed.to_csv(bedname + "_genes.bed", sep="\t", index=None)
        return newbed


def getBamDate(bams, split="-", order="des", unknown="U"):
    """
    from bam files (could be in a google bucket) returns their likely sequencing date if available in the header

    Args:
    -----
      bams: the bams file|bucket paths
      split: the splitter in the output date
      unknown: maybe the some dates can't be found the program will output unknown for them
      order: if 'asc', do d,m,y else do y,m,d

    Returns:
    -------
      a list of likely dates or [unknown]s
    """
    DTs = []
    for i, bam in enumerate(tqdm(bams)):
        data = os.popen(
            "export GCS_OAUTH_TOKEN=`gcloud auth application-default print-access-token`\
    && samtools view -H "
            + bam
            + ' | grep "^@RG"'
        )
        if data == signal.SIGINT:
            print("Awakened")
            break
        else:
            res = data.read()
            dt = re.findall("(?<=\tDT:).+?\t", res)
        if len(dt) > 1:
            arr = np.array(dt[0].split("T")[0].split(split)).astype(int)
            for val in dt[1:]:
                arr = np.vstack(
                    (arr, np.array(val.split("T")[0].split(split)).astype(int))
                )
            arr = arr.T
            i = (
                arr[0] * 365 + arr[1] * 31 + arr[2]
                if order == "asc"
                else arr[2] * 365 + arr[1] * 31 + arr[0]
            )
            DTs.append(dt[np.argsort(i)[0]].split("T")[0])
        elif len(dt) == 1:
            DTs.append(dt[0].split("T")[0])
        else:
            DTs.append(unknown)
    return DTs


async def indexBams(bams=None, bucketpath=None, cores=4):
    """
    given a bucket path, will index all .bam files without an associated index and return their paths
    """
    if bams is None:
        if bucketpath is None:
            raise ValueError("need one of bams or bucketpath")
        files = gcp.lsFiles([bucketpath])
        bams = [val for val in files if ".bam" in val[-4:]]
        unindexed = [
            val
            for val in bams
            if val[:-4] + ".bai" not in files and val[:4] + ".bam.bai" not in files
        ]
        print("found " + str(len(unindexed)) + " files to reindex")
    else:
        unindexed = bams
    h.parrun(
        [
            "export GCS_OAUTH_TOKEN=`gcloud auth application-default print-access-token` && samtools index "
            + val
            for val in unindexed
        ],
        cores,
    )
    return {val: val[:-4] + ".bam.bai" for val in unindexed}


def dropWeirdChromosomes(bedfile, keep=[], skip=0):
    """
    given a bedfile path, removes chromosomes that are not one of chroms

    Args:
    ----
      bedfile: str the filepath to the bedfile
      keep: list[str] of additional chromosomes to keep
    """
    if skip >= 20:
        raise ValueError("too many header lines!")
    try:
        bed = pd.read_csv(bedfile, sep="\t", header=None, skiprows=skip)
    except ParserError:
        dropWeirdChromosomes(bedfile, keep, skip + 1)
        return
    except EmptyDataError:
        print("empty bed")
        return
    initlen = len(bed)
    if initlen == 0:
        print("empty bed")
        return
    bed = bed[bed[0].isin(chroms | set(keep))]
    if len(bed) < skip and skip > 5:
        raise ValueError("too many header lines!")
    print("found " + str(skip) + " header line... removing")
    if len(bed) != initlen:
        print("removed " + str(initlen - len(bed)) + " lines")
    bed.to_csv(bedfile, sep="\t", header=None, index=None)


def extractPairedSingleEndFrom(folder, sep="-", namepos=2):
    """
    given a folder, find fastq files and sorts paired and single end based on the R1/R2 patterns

    Args:
    -----
      folder: the folder where the fastqs are
      sep: the separator in filename
      namepos: the location of the name in this separated list of name from filepath

    Returns:
    -------
      list of filepath to single end files
      df with R1 and R2 filepath
    """
    single = []
    paired = {}
    for val in os.listdir(folder):
        if ".fastq" in val or ".fq" in val:
            if "R1" in val:
                name = val.split(sep)[namepos]
                paired[name] = {"R1": val}
            elif "R2" in val:
                name = val.split(sep)[namepos]
                paired[name].update({"R2": val})
            else:
                single.append(val)
    return single, pd.DataFrame(paired)


def findReplicatesBams(folder, sep="-", namings="-r([0-9])", namepos=2):
    """
    creates a dict of name and replicate files given a regexp namging scheme
    """
    rep = {}
    for val in os.listdir(folder):
        if val[-4:] == ".bam":
            match = re.search(namings, val)
            if match:
                name = val.split(sep)[namepos]
                if name in rep:
                    rep[name].append(val)
                else:
                    rep[name] = [val]

    return rep


def singleEnd(
    singlend,
    folder="data/seqs/",
    numthreads=8,
    peaksFolder="peaks/",
    ismapped=False,
    mappedFolder="mapped/",
    refFolder="data/reference/index",
):
    """
    run the singleEnd pipeline
    for alignment etc, one can use pysam ready made implementation of samtools
    """
    print(
        "you need to have bowtie2 installed: http://bowtie-bio.sourceforge.net/bowtie2/index.shtml"
    )
    for val in singlend:
        out1 = folder + mappedFolder + val.split(".")[0] + ".mapped.sam"
        if not ismapped:
            in1 = folder + val
            os.system(
                "bowtie2 -x "
                + refFolder
                + " --threads "
                + str(numthreads)
                + " -t -k 1 --very-sensitive -U "
                + in1
                + " -S "
                + out1
            )
        out2 = folder + peaksFolder + val.split(".")[0]
        print(out1)
        os.system("macs2 callpeak -f SAM -t " + out1 + " --outdir " + out2)
        # it can take many TB so better delete


def pairedEnd(
    pairedend,
    folder="",
    numthreads=8,
    peaksFolder="peaks/",
    ismapped=False,
    mappedFolder="mapped/",
    refFolder="data/reference/index",
):
    """
    # run the paired end pipeline
    """
    print(
        "you need to have bowtie2 installed: http://bowtie-bio.sourceforge.net/bowtie2/index.shtml"
    )
    for _, val in pairedend.items():
        out1 = folder + mappedFolder + val[0].split(".")[0] + ".mapped.sam"
        in1 = folder + val[0]
        in2 = folder + val[1]
        os.system(
            "bowtie2 -x "
            + refFolder
            + " --threads "
            + str(numthreads)
            + " -t -k 1 \
    --very-sensitive -1 "
            + in1
            + " -2 "
            + in2
            + " - S "
            + out1
        )
        out2 = folder + peaksFolder + val[0].split(".")[0]
        print(out1)
        changefrom = out1
        changeto = out1[:-4] + ".bam"
        os.system("samtools view -b " + changefrom + " -o " + changeto)
        os.system(
            "macs2 callpeak --format 'BAMPE' --treatment "
            + changeto
            + " --outdir "
            + out2
        )
        # it can take many TB so better delete


async def mergeBams(rep):
    """
    uses samtools to merge a set of replicates considered into one file
    """
    in1 = ""
    for i, val in rep.items():
        out1 = i + ".merged.bam"
        for bam in val:
            in1 += " " + bam
        os.system("samtools merge " + out1 + in1)


def compare_gcloud_vcfs_overlap_methods(vcfs_met1_path, vcfs_met2_path):
    for i, j in zip(vcfs_met1_path, vcfs_met2_path):
        compare_gcloud_vcf_overlap(i, j)


def compare_gcloud_vcf_overlap(vcf1, vcf2, cols=["chr", "start", ".", "ref", "alt"]):
    import subprocess

    name1 = vcf1.split("/")[-1].split(".")[0] + "_1" + ".tsv"
    cmd1 = "gsutil cat " + vcf1 + " | gunzip | cut -f -5 > " + name1
    name2 = vcf2.split("/")[-1].split(".")[0] + "_2" + ".tsv"
    cmd2 = "gsutil cat " + vcf2 + " | gunzip | cut -f -5 > " + name2
    try:
        subprocess.run(
            cmd1,
            shell=True,
            check=True,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
        )
        subprocess.run(
            cmd2,
            shell=True,
            check=True,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
        )
    except subprocess.CalledProcessError as e:
        print(e.stderr)
        raise e
    val2 = pd.read_csv(name2, sep="\t", comment="#", names=cols)
    val1 = pd.read_csv(name1, sep="\t", comment="#", names=cols)
    val1["loc"] = (
        val1["chr"].astype(str)
        + ":"
        + val1["start"].astype(str)
        + ":"
        + val1["alt"].astype(str)
    )
    val2["loc"] = (
        val2["chr"].astype(str)
        + ":"
        + val2["start"].astype(str)
        + ":"
        + val2["alt"].astype(str)
    )
    print("length of vcf1:" + str(len(val1)))
    print("length of vcf2:" + str(len(val2)))
    print("overlap: " + str(len(set(val1["loc"]).intersection(val2["loc"]))))
    return val1, val2
