"""
Runs multiple `rocco chrom` jobs, which execute ROCCO  on each respective
chromosome. Chromosome-specific parameters are collected in a CSV file specified
with argument `-p --param_file`. Consider using the `budgets` subcommand
to set chromosome-specific budgets based on read densities.

If any parameter entry in the CSV file `--param_file` is set to `NULL`, the script uses the
corresponding default value specified by the command-line arguments.

Supports simultaneous execution of multiple `rocco chrom` jobs via `--multi`.

The output files generated by `rocco chrom` are stored in the specified output
directory (`--outdir`, default is current working directory). If the `combine`
argument is provided not `None`, then the script combines the chromosome-specific
bed files into a single output file with the specified name.

A note on `--solver`: The default solver is ECOS, which is free, open-source, and
conveniently installed alongside cvxpy. ECOS should prove sufficient in general, but
using a commerical (e.g., MOSEK) solver can improve speed of execution. *Free* academic
MOSEK licenses can be obtained *instantly* with a valid .edu email address: https://www.mosek.com/products/academic-licenses/
Generous free trials are also available for commercial use.

After obtaining the license (`mosek.lic`, or something similar), create a folder `mosek` in
your home directory and move the license file to it.

```
mkdir ~/mosek
mv mosek.lic ~/mosek
```

Then, run:

```
pip install mosek
```

After these steps, you should be able to specify MOSEK with
with `--solver MOSEK`


"""

import os
import sys
import argparse
import subprocess
import tempfile
from . import rocco_aux

def get_params(param_file: str, budget: float, gamma: float, tau: float,
               c1: float, c2: float, c3: float) -> list:
    """
    Grabs parameters for each chromosome from `param_file`

    This function collects parameters from each chromosome-row of
    `params.csv`. If a `NULL` entry is encountered, it is replaced
    by the CLI-specified default (see `args`).

    Args:
        param_file (str) : a CSV file with a row containing parameter vals\
            to use for each chromosome.

    Returns:
        list: a list of lists, with each element containing
          the necessary parameters to run a `rocco chrom` job.
    """
    defaults = [None, None, budget, gamma, tau, c1, c2, c3]
    chrom_params = []
    with open(param_file, mode='r', encoding='utf-8') as par_file:
        header = True
        for line in par_file:
            if header is True:
                header = False
                continue
            if ',' not in line:
                continue
            line = line.strip()
            line = line.split(',')
            for i, entry in enumerate(line):
                if entry == 'NULL':
                    # replace NULL entries with the genome-wide default
                    line[i] = defaults[i]
            chrom_params.append(line)
    return chrom_params


def call_rocco(chrom, wig_path, budget, gamma, tau, c1, c2, c3, solver,
               bed_format, verbose=False, N=50, identifiers=None,
               outdir='.') -> str:
    r"""
    Formats a command to run `rocco chrom` for a given chromosome

    Args:
        chrom (str): chromosome name. Example: `chr1`
        wig_path (str): path to the *directory* containing samples' signal
            tracks for `chrom`
        budget (float): budget constraint
        tau (float): tau in $\mathcal{S}(i)$
        gam (float): gamma parameter in $f(\mathbf{\ell})$
        c1 (float): c1 value in $\mathcal{S}(i)$
        c2 (float): c2 value in $\mathcal{S}(i)$
        c3 (float): c3 value in $\mathcal{S}(i)$
        N (int): RR procedure iterations. If N <= 0,
            $\texttt{floor\_eps}$ procedure is applied...
            See `Loci.run_rr()`
        verbose_ (bool): Verbosity flag for the solver
        solver (str): the solver to use--either "ECOS" or "MOSEK"
    Returns:
        str: a formatted rocco chrom command with the appropriate parameters.
        substituted.
    """
    cli_args = ["rocco", "chrom",
                '--chrom', chrom,
                '--wig_path', wig_path,
                '--budget', budget,
                '--gamma', gamma,
                '--tau', tau,
                '--c1', c1,
                '--c2', c2,
                '--c3', c3,
                '--solver', solver,
                '--bed_format', bed_format,
                '--outdir', outdir,
                '--rr_iter', N,
                '--identifiers', identifiers,
                '--verbose']

    if not verbose:
        cli_args.remove('--verbose')
    if identifiers is None:
        cli_args.remove('--identifiers')
        cli_args.remove(identifiers)
    return ' '.join(cli_args)

def main(args):
    chrom_args = get_params(args['param_file'], args['budget'],
                            args['gamma'], args['tau'], args['c1'],
                            args['c2'], args['c3'])

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

    tmp = tempfile.NamedTemporaryFile(mode="w+")
    for i, arglist in enumerate(chrom_args):
        arglist = [str(x) for x in arglist]
        cmd = call_rocco(arglist[0], arglist[1], arglist[2], arglist[3],
                         arglist[4], arglist[5], arglist[6], arglist[7],
                         args['solver'], str(args['bed_format']),
                         args['verbose'], str(args['rr_iter']),
                         identifiers=args['identifiers'], outdir=args['outdir'])
        if not args['multi']:
            try:
                seq_process = subprocess.run(cmd.split(' '),
                                         capture_output=True, text=True, check=True)
                print(seq_process.stdout)
            except Exception as ex:
                print(ex)
                raise ex
        tmp.write(str(cmd + '\n'))

    tmp.flush()
    if args['multi']:
        rocco_aux.run_par(tmp.name, verbose=args['verbose'])
    tmp.close()

    if args['combine'] is not None:
        print('combining output files --> {}'.format(args['combine']))
        rocco_aux.sort_combine_bed(args['combine'], dir_=args['outdir'])


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('-p', '--param_file',
                        required=True)
    parser.add_argument('-b', '--budget', type=float, default=.035, help='budget parameter (largest allowed fraction of selected bp) used for each chromosome with a `NULL` entry observed in `--param_file`')
    parser.add_argument('-g', '--gamma', type=float, default=1.0, help='gamma parameter (discontig. penalty weight) used for each chromosome  with a `NULL` entry observed in `--param_file`')
    parser.add_argument('-t', '--tau', type=float, default=0.0, help='tau parameter (enrichment threshold) used for each chromosome  with a `NULL` entry observed in `--param_file`')
    parser.add_argument('--c1', type=float, default=1.0, help='g_1 coefficient in score function (enrichment reward) used for each chromosome  with a `NULL` entry observed in `--param_file`')
    parser.add_argument('--c2', type=float, default=1.0, help='g_2 coefficient in score function (dispersion penalty) used for each chromosome  with a `NULL` entry observed in `--param_file`')
    parser.add_argument('--c3', type=float, default=1.0, help='g_3 coefficient in score function (local shift) used for each chromosome  with a `NULL` entry observed in `--param_file`')
    parser.add_argument('-N', '--rr_iter', type=int, default=50, help = 'number of RR iterations')
    parser.add_argument('--solver', default="ECOS",
                        help="Optimization software used to solve the \
                        main LP. `ECOS` is used by default.")
    parser.add_argument('--bed_format', type=int, default=3,
                        help="Specifies BED3 or BED6 format.\
                        Default is BED6. Generate BED3 output with \
                        --bed_format 3")
    parser.add_argument('--identifiers', default=None,
                        help="(optional) a filename containing identifiers\
                          for samples to include in experiment. Each identi\
                          fier should be a uniquely-identifying substring of\
                          the respective `.wig` sample. If not specified, all\
                          samples are used.")
    parser.add_argument('--outdir', default='.',
                        help="directory in which to store output bed files from the calls to rocco chrom")
    parser.add_argument('--combine', default=None, help="if not None, combine\
                        output bed files and store in the file specified\
                        with this parameter. ex: `--combine bedname.bed` con-\
                        catenates the chromosome-specific bedfiles into `bedname.bed`.")
    parser.add_argument('--multi', default=False, action='store_true', help='run rocco chrom jobs\
        simultaneously to improve speed. May increase peak memory use.')
    parser.add_argument('--verbose', default=False, action="store_true")
    args = vars(parser.parse_args())
    main(args)
