#!/opt/miniconda3/envs/gennaro/bin/python
"""
Read in the run-info.csv file for the given BioProject and: 

1. Check to see if the LibraryStrategy is RNA-Seq.
2. Check the LibraryLayout (SINGLE vs PAIRED)
3. Spawn a subprocess that downloads the raw fastq file via fasterq-dump
4. Use the library type to determine the mapping parameters
5. Map the raw fastq files
6. Remove the fastq files once the quants files are generated

"""
import argparse
import subprocess
import sys
from pathlib import Path

import pandas as pd 


def process_run_info(run_info_csv):
    """extract SRA accession, library type, and layout from run-info.csv"""
    df = pd.read_csv(run_info_csv)
    df = df[['Run', 'LibraryLayout', 'LibraryStrategy', 'TaxID']]
    records = df.to_records(index=False)

    return records


def download_fq(accession, layout_type, threads, out_dir):
    """Download fastq files from SRA using fasterq-dump"""
    if layout_type == "PAIRED":
        fq1 = Path(f"{out_dir}/{accession}_1.fastq")
        fq2 = Path(f"{out_dir}/{accession}_2.fastq")
        if not (fq1.exists() and fq2.exists()):
            subprocess.run(f"prefetch {accession} -O {Path(out_dir)}", shell=True)
            subprocess.run(f"fasterq-dump {accession} -e {threads} -p --outdir {out_dir}", shell=True)
        else:
            print(f"{fq1} and {fq2} already downloaded. skipping download...")
    else:
        fq = Path(f"{out_dir}/{accession}.fastq")
        if not fq.exists():
            subprocess.run(f"prefetch {accession} -O {Path(out_dir)}", shell=True)
            subprocess.run(f"fasterq-dump {accession} -e {threads} -p --outdir {out_dir}", shell=True)
        else:
            print(f"{fq} already downloaded. skipping download...")
    

def map_sample(accession, layout_type, salmon_idx, threads, out_dir):
    """Open a subprocess instance that can appropriately map paired-end or 
    single end samples and remove fastq files when finished"""
    quant_dir = Path(f"{out_dir}/quants/{accession}_quants")
    quant_file = Path(f"{quant_dir}/quant.sf")

    if layout_type == "PAIRED":
        fq1 = Path(f"{out_dir}/{accession}_1.fastq")
        fq2 = Path(f"{out_dir}/{accession}_2.fastq")
    
        if fq1.exists() and fq2.exists():
            subprocess.run(f"/usr/local/bin/salmon quant --seqBias --gcBias -i {salmon_idx} -l 'A' -p {threads} -1 {fq1} -2 {fq2} -o {quant_dir}", shell=True)
        else:
            sys.exit("Input fastq files do not exist")

        if quant_file.exists():
            subprocess.run(f"rm -f {fq1} {fq2}", shell=True)
            subprocess.run(f"rm -rf {Path(out_dir, accession)}", shell=True)
        else:
            sys.exit("quant.sf file not generated. Check Salmon logs.")
    else:
        fq = Path(f"{out_dir}/{accession}.fastq")

        if fq.exists():
            subprocess.run(f"/usr/local/bin/salmon quant --seqBias --gcBias -i {salmon_idx} -l 'A' -p {threads} -r {fq} -o {quant_dir}", shell=True)
        else:
            sys.exit("Input fastq files do not exist")
        
        if quant_file.exists():
            subprocess.run(f"rm -f {fq}", shell=True)
            subprocess.run(f"rm -rf {Path(out_dir, accession)}", shell=True)
        else:
            sys.exit("quant.sf file not generated. Check Salmon logs")


def main():
    parser = argparse.ArgumentParser(description='Download and map samples from SRA with REdiscoverTE pipeline')
    parser.add_argument('run_info', help='run-info.csv file generated by 02_get_run_info.sh')
    parser.add_argument('--out_dir', help='directory to output results. By default, this is the parent directory of the run-info.csv file')
    parser.add_argument('--salmon_idx', help='path to Salmon index directory', default="resources/REdiscoverTE_hg38/REdiscoverTE_hg38_GFP/REdiscoverTE_GFP_salmon_idx")
    parser.add_argument('--threads', help='number of threads to use in mapping', default=8)
    args = parser.parse_args()

    run_info = args.run_info
    out_dir = args.out_dir
    salmon_idx = args.salmon_idx
    threads = args.threads

    if not out_dir:
        out_dir = Path(run_info).parent
    out_dir = Path(out_dir)

    records = process_run_info(run_info)

    for record in records:
        sra_acc = record[0]
        layout = record[1]
        strategy = record[2]
        taxa = int(record[3])
        quant_file = Path(f"{out_dir}/quants/{sra_acc}_quants/quant.sf")
        quant_file_gz = Path(f"{out_dir}/quants/{sra_acc}_quants/quant.sf.gz")

        if strategy == "RNA-Seq" and taxa == 9606:
            if quant_file.exists() or quant_file_gz.exists():
                print(f"Mapping results exist for {sra_acc}. Skipping...")
            else:
                print(f"Downloading {sra_acc}...")
                download_fq(sra_acc, layout, threads, out_dir)
                print(f"Mapping {sra_acc}...")
                map_sample(sra_acc, layout, salmon_idx, threads, out_dir)


if __name__ == '__main__':
    main()