r"""
## [ROCCO](https://github.com/nolan-h-hamilton/ROCCO/): chrom.py 
Run ROCCO on chromosome `--chrom` with samples' count/signal data in `tracks_chr[]`

This script is the workhorse for [`rocco gwide`](https://nolan-h-hamilton.github.io/ROCCO/rocco/gwide.html) which generates and combines
results for multiple chromosomes

Parameters:
    --chrom (str): (Required) Specifies the chromosome, e.g., `--chrom chr1`.
    --wig_path (str): (Required) Directory generated by `rocco prep` containing chromosome-specific
        .wig files for each sample, e.g., `--wig_path tracks_chr1`
    -b, --budget (float): Upper-bounds the fraction of loci selected as open.
        Defines the budget constraint (Default: 0.035)
    -g, --gamma (float): $\gamma$ fragmentation penalty in objective func. (Default: 1.0)
    -t, --tau (float): $\tau$ parameter (enrichment threshold) (Default: 0.0).
    --c1 (float): $c_1$ coefficient in the score function (enrichment reward) (Default: 1.0).
    --c2 (float): $c_2$ coefficient in the score function (dispersion penalty) (Default: 1.0).
    --c3 (float): $c_3$ coefficient in the score function (local shift) (Default: 1.0).
    --start (int): Beginning nucleotide position (Default: infer from wiggle files).
    --end (int): Ending nucleotide position (Default: infer from wiggle files).
    -N, --rr_iter (int): Number of RR iterations. $N$ in the paper. (Default: 50).
    --verbose (bool): invoke for verbose logging
    --solver (str): Solver used for the underlying optimization problem of ROCCO (Default: "CLARABEL").
        Other solvers:
            - `MOSEK`: fast, commercial grade solver. Users can instantly obtain a free academic\
                license or generous trial commericial license at
                https://www.mosek.com/products/academic-licenses

            - `PDLP`: first-order solver suitable for very large problems or limited computing environments.
                Uses a primal-dual hybrid gradient approach.

            - `ECOS`: Lightweight second-order conic solver. One of CVXPY's solvers included with the
                default installation, though there are plans to migrate to `CLARABEL`:
                (https://www.cvxpy.org/updates/index.html#ecos-deprecation)
    --bed_format (int): BED format, `3` for BED3 format and `6` for BED6 format (Default: 6).
    --identifiers (str): text file containing a subset of sample IDs, one on each line, for samples\
        to include in the experiment. If not invoked, defaults to `None`, and *all* samples with
        wig files in `--wig_path` are used.

Input:
    A `tracks_chr[]` directory containing samples wiggle tracks created by [`rocco prep`](https://nolan-h-hamilton.github.io/ROCCO/rocco/prep.html).
    For example: `--wig_path tracks_chr1`

Output:
    BED file containing predicted regions of open chromatin. Named with convention:
    ```ROCCO_out_[chrom]_[budget]_[gamma]_[tau]_[c1]_[c2]_[c3].bed```

Examples:
    - Run ROCCO on chromosome 21 with default parameters and using all samples for input:
        ```
        rocco chrom --wig_path tracks_chr21 --chrom chr21
        ```
    - Run ROCCO on chromosome 21 with default parameters on a subset of samples listed in `identifiers.txt`:
        ```
        rocco chrom --wig_path tracks_chr21 --chrom chr21 --identifiers identifiers.txt
        ```
        For instance, samples are named `sample1, sample2, sample3, ..., sampleK` and `identifiers.txt`
        contains:
        ```
        sample1
        sample2
        sample3
        ```
        Then the command above will ignore sample4,...,sampleK during peak calling.

    - Run ROCCO on chromosome 21. Select no more than $2\%$ of the $n$ loci
        ```
        rocco chrom --wig_path tracks_chr21 --chrom chr21 --budget 0.02
        ```
"""

import os
import sys
import gc
import argparse
import warnings
import subprocess
import numpy as np
import pandas as pd
from .locus import Locus
from .loci import Loci


def get_locus_size(wig_path) -> int:
    """
    Infers step size from replicates` wig files.

    Args:
        wig_path (str) : path to *directory* containing wig files


    Returns:
        int: an integer locus/step size used in the wig files
    """
    i = 0
    pos1 = 0
    pos2 = 0
    wig_file = [x for x in os.listdir(wig_path) if '.wig' in x][0]
    with open(wig_path + '/' + wig_file, mode='r', encoding='utf-8') as wig:
        for line in wig:
            line = line.strip()
            if line.replace('\t', '').replace(' ', '').isnumeric():
                if i == 0:
                    pos1 = int(line.split('\t')[0])
                if i == 1:
                    pos2 = int(line.split('\t')[0])
                if pos2 - pos1 > 0:
                    break
                i += 1
    return pos2 - pos1

def get_start_end(wig_path):
    """
    Infers starting and ending nucleotide position common to the
    multiple signal tracks

    To infer a reasonable genomic region over which to run ROCCO,
    this function returns the median starting point and endpoints
    among the multiple replicates' signal files in `wig_path`.

    Args:
        wig_path (str) : path to *directory* containing wig files

    Returns:
        start (int): median starting position
        end (int): median ending position
    """
    starts = []
    ends = []

    def get_start(track):
        with open(track, mode='r', encoding='utf-8') as wig:
            for line in wig:
                line = line.strip()

                if line.replace('\t', '').replace(' ', '').isnumeric():
                    start = int(line.split('\t')[0])
                    return start
        return None

    def get_end(track):
        with subprocess.Popen(["tail", "-n1", track],
                              stdout=subprocess.PIPE) as proc:
            for line in proc.stdout:
                line = line.decode()
                line = line.strip()
                return int(line.split('\t')[0])
        return None

    wig_files = [x for x in os.listdir(wig_path) if '.wig' in x]
    for wig_file in wig_files:
        starts.append(get_start(wig_path + '/' + wig_file))
        ends.append(get_end(wig_path + '/' + wig_file))
    return sorted(starts)[len(starts) // 2], sorted(ends)[len(ends) // 2]


def read_wig(wig_file, start=0, end=10**10, locus_size=50):
    r"""
    Processes a wig file for inclusion in the signal matrix $\mathbf{S}_{chr}$

    This function prepares signal data for to fit in a constant-sized
    $K \times n$ matrix. If a wig file begins at a position < `start`,
    it will be padded with zeros. Likewise if a wig file ends at a
    position > `end`, the extra values are ignored. Gaps between `start`
    and `end` are replaced with a zero-entry.

    Args:
        wig_file (str): path to wiggle formatted signal track
        start (int): inferred or manually-specified starting nucleotide
          position
        end (int): inferred or manually-specified ending nucleotide
          position
        locus_size (int): interval/locus size. Each wig file will have
          signal values at every `start + i*locus_size` nucleotide pos.
          for i = 0,1,2,...

    Returns:
        loci: list of starting nucleotide positions for each locus
        signal: enrichment signal value at each locus in `loci`
    """
    print("reading: {}".format(wig_file))
    loci = []
    signal = []
    with open(wig_file, mode='r', encoding='utf-8') as wig:
        for line in wig:
            line = line.strip()
            try:
                line = line.split('\t')
                line[0] = int(line[0])
                line[1] = float(line[1])
                if line[1] < 0:
                    line[1] = 0
                # case: wig begins after specified `start`
                if start < line[0] and len(loci) == 0:
                    for loc in range(start, line[0], locus_size):
                        loci.append(loc)
                        signal.append(0)
                    loci.append(line[0])
                    signal.append(line[1])
                    continue
                if line[0] < start:
                    continue
                if line[0] > end:
                    break
                # case: loci not contiguous
                if len(loci) > 0 and line[0] - loci[-1] > locus_size:
                    loci.append(loci[-1] + locus_size)
                    signal.append(0)
                    continue

                loci.append(line[0])
                signal.append(line[1])

            except ValueError:
                continue
            except IndexError:
                continue
    # case: wig ends prematurely
    if loci[-1] < end:
        for loc in range(loci[-1] + locus_size, end + locus_size, locus_size):
            loci.append(loc)
            signal.append(0)

    return loci, signal


def collect_wigs(wig_files, start=0, end=10**10, locus_size=50) -> pd.DataFrame:
    r"""
    creates a $K \times n$ dataframe of signals derived from wig files

    Args:
        wig_files (list): list of wig filepaths
        start (int): starting nucleotide position of wig files
        end (int): ending nucleotide position of wig files
        locus_size (int): initial size (in nucleotides) of Locus objects.
            Corresponds to $L$ in paper.

    Returns:
        pd.DataFrame: $K \times n$ dataframe of individuals' signals.
    """
    loci = [x for x in range(start, end + locus_size, locus_size)]
    all_wigs = pd.DataFrame(
        np.array([read_wig(x, start, end, locus_size)[1]
                  for x in wig_files],dtype=np.float32),
        columns=loci, index=None)
    return all_wigs


def emp_cdf(val, arr) -> float:
    """
    Get empirical cdf of `val` in `arr`

    Used to compute 0-1000 scaled score for BED6 files
    and visualization in genome browser.

    Note:
        Assumes `arr` is sorted!

    Args:
        val (float): a numeric element in `arr`
        arr (array): a SORTED array of peak values

    Returns:
        (float) estimate for percentile of `val` in `arr`
    """
    return np.searchsorted(arr, val) / len(arr)


def sig_selected(loci_obj: Loci) -> np.ndarray:
    """
    Get enrichment signals from all selected loci

    Args:
        loci_obj (Loci): Loci object after optimization

    Returns:
        np.ndarray: vector of enrichment signals from all selected
            (but still unmerged) loci
    """
    sig_ref = []
    for loc in loci_obj:
        if loc.accessible == 1:
            sig_ref.append(np.sum(loc.sig_data) / loc.size)
    return np.array(sig_ref)


def tags(fname) -> list:
    """
    Get sample IDs from file specified with command-line argument `--identifiers`

    Args:
        fname (str): path to file containing a subset of sample IDs--one line for each ID

    Returns:
        list: list of sample IDs in `fname`
    """
    tags_ = []
    for line in open(fname, mode='r', encoding='utf-8'):
        if line.strip().isspace() or len(line.strip()) == 0:
            continue
        tags_.append(line.strip())
    return tags_



def main(args):
    if args['verbose']:
        print(f'chrom: initial args\n{args}')
    if args['locus_size'] == -1:
        args['locus_size'] = get_locus_size(args['wig_path'])

    args['outdir'] = os.path.normpath(args['outdir'])
    if not os.path.exists(args['outdir']):
        os.mkdir(args['outdir'])

    if args['start'] == -1 or args['end'] == -1:
        args['start'], args['end'] = get_start_end(args['wig_path'])
    if args['bed_format'] not in [3, 6]:
        warnings.warn('Only BED3 and BED6 are supported--setting to BED6')
        args['bed_format'] = 6
    if args['verbose']:
        print(f'chrom: inferred args\n{args}')


    wig_files = [
        os.path.join(args['wig_path'], fname)
        for fname in os.listdir(args['wig_path'])
        if os.path.splitext(fname)[-1].lower() == '.wig'
    ]

    # select a subset of the wig files if `--identifiers` is defined
    if args['identifiers'] is not None:
        with open(os.path.join(args['outdir'], f"{args['chrom']}_sample_wig_match.log"), 'w') as samp_wig_log:
            new_wig_files = []
            for tag in tags(args['identifiers']):
                matching_wigs = []
                for wfname in wig_files:
                    if tag in wfname:
                        matching_wigs.append(wfname)
                tag_match = None
                if len(matching_wigs) > 1:
                    # pick the shortest match
                    tag_match = sorted(matching_wigs, key=len)[0]
                elif len(matching_wigs) == 1:
                    tag_match = matching_wigs[0]
                elif len(matching_wigs) == 0:
                    warnings.warn(f'chrom: no wig file found for sample identifier {tag}')
                if tag_match is not None:
                    print(f'chrom: sample/wig match: {tag}/{tag_match}')
                    samp_wig_log.write(f'{tag}\t{tag_match}\n')
                    new_wig_files.append(tag_match)
        wig_files = new_wig_files

    if len(wig_files) == 0:
        warnings.warn(f"chrom: did not any find wig files for the given samples + {args['chrom']}")
        sys.exit(0)

    # create a dataframe of signal values for `K` replicates and `n` loci
    signal_matrix = collect_wigs(wig_files,
                                 start=args['start'],
                                 end=args['end'],
                                 locus_size=args['locus_size'])

    InitLoci = Loci()
    for i, loc in enumerate(signal_matrix.columns):
        size_ = args['locus_size']
        NewLoc = Locus(position=args['start'] + i * size_,
                       size=size_,
                       sig_data=np.array(signal_matrix[loc],dtype=np.float32))
        InitLoci.append_locus(NewLoc)
    del signal_matrix
    gc.collect()

    InitLoci.rocco_lp(budget=args['budget'],
                        gam=args['gamma'],
                        tau=args['tau'],
                        c1=args['c1'],
                        c2=args['c2'],
                        c3=args['c3'],
                        verbose_=args['verbose'],
                        solver=args['solver'],
                        N=args['rr_iter'])

    InitLoci.combine_selected()

    fname = os.path.join(args['outdir'], "ROCCO_out_{}_{}_{}_{}_{}_{}_{}.bed".format(
        args['chrom'],
        args['budget'],
        args['gamma'],
        args['tau'],
        args['c1'],
        args['c2'],
        args['c3']))

    if os.path.exists(fname):
        warnings.warn(f'chrom: overwriting {fname}')
    print('chrom: writing output: {}'.format(fname))
    outfile = open(fname, mode='w', encoding='utf-8')
    output = ""
    if args['bed_format'] == 6:
        loc_count = 1
        sorted_sel_sig = np.sort(sig_selected(InitLoci))
        for loc in InitLoci:
            assert loc.size > 0
            if loc.accessible == 1:
                # the score can be used for coloring in ucsc genome browser
                loc_score = int(500 * emp_cdf(np.sum(loc.sig_data) / loc.size,
                                              sorted_sel_sig) + 500)
                loc_name = args['chrom'] + '_' + str(loc_count)
                output += "{}\t{}\t{}\t{}\t{}\t{}\n".format(
                    args['chrom'], loc.position, loc.position + loc.size,
                    loc_name, loc_score, '.')
                loc_count += 1

    if args['bed_format'] == 3:
        for loc in InitLoci:
            if loc.accessible == 1:
                output += "{}\t{}\t{}\n".format(
                    args['chrom'], loc.position, loc.position + loc.size)

    outfile.write(output)
    outfile.close()


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--chrom', type=str, help="Specifies the chromosome, e.g., `--chrom chr1`.")
    parser.add_argument('--start', type=int, default=-1, help='beginning nucleotide position (default: infer from wiggle files)')
    parser.add_argument('--end', type=int, default=-1, help='ending nucleotide position (default: infer from wiggle files)')
    parser.add_argument('-L','--locus_size', type=int, default=-1, help='fixed step size in the samples wiggle tracks (default: infer from wiggle files)')
    parser.add_argument('--wig_path', type=str, default=os.getcwd(), help='Directory generated by `rocco prep` containing chromosome-specific .wig files for each sample, e.g., `--wig_path tracks_chr1`')
    parser.add_argument('-N', '--rr_iter', type=int, default=50, help = 'number of RR iterations')
    parser.add_argument('--verbose', default=False, action="store_true", help='set to `True` for verbose logging')
    parser.add_argument('-b', '--budget', type=float, default=.035, help='Upper-bounds the fraction of loci selected as open. Defines the budget constraint (Default: 0.035)')
    parser.add_argument('-g', '--gamma', type=float, default=1.0, help='fragmentation penalty in objective func. (Default: 1.0)')
    parser.add_argument('-t', '--tau', type=float, default=0.0, help='tau parameter (enrichment threshold)')
    parser.add_argument('--c1', type=float, default=1.0, help='c_1 coefficient in score function (enrichment reward)')
    parser.add_argument('--c2', type=float, default=1.0, help='c_2 coefficient in score function (dispersion penalty) ')
    parser.add_argument('--c3', type=float, default=1.0, help='c_3 coefficient in score function (local shift)')
    parser.add_argument('--solver', default="CLARABEL", help='Optimization software used to solve the \
                        main LP. `CLARABEL` is used by default.')
    parser.add_argument('--bed_format', type=int, default=6,
                        help="`3` for BED3 format and `6` for BED6 format")
    parser.add_argument('--identifiers', default=None,
                        help="text file containing a subset of sample IDs, one on each line, for samples to include in the experiment. If not invoked, defaults to `None`, and *all* samples with wig files in `--wig_path` are used.")
    parser.add_argument('--outdir', type=str, default='.', help="directory in which to store output bed file")

    args = vars(parser.parse_args())
    main(args)
