from collections import OrderedDict
from datetime import datetime, timedelta
from glob import glob

import click
from click import style
from click_help_colors import HelpColorsGroup

from . import __version__
from .utils import (
    _get_filepath,
    _check_if_file_exists,
    _check_if_file_is_readable,
    echo,
)


class Timer:
    def __init__(self, label="run", fmt="%a %b %d %H:%M:%S %Y"):
        self.label = style(label, bold=True)
        self.fmt = fmt
        self.begin_time = datetime.now()
        echo(f"\nStarted {self.label} on {self.begin_time.strftime(self.fmt)}.")

    def stop(self):
        end_time = datetime.now()
        runtime = timedelta(seconds=(end_time - self.begin_time).seconds)
        echo(
            f"\nFinished {self.label} on {end_time.strftime(self.fmt)}. "
            f"Total runtime was {runtime} (hour:min:s)."
        )


class OrderedGroup(HelpColorsGroup):
    def __init__(self, commands=None, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.commands = commands or OrderedDict()

    def list_commands(self, ctx):
        return self.commands


def _print_kwargs(kwargs):
    echo("\nCommand line arguments:")
    for arg, value in kwargs.items():
        if value is not None:
            value_fmt = style(str(_get_filepath(value)), fg="blue")
            echo(f"{arg: >15}: {value_fmt}")
    echo()


def _get_input_file_paths(file_arg):
    if len(file_arg) == 1 and "*" in file_arg[0]:
        globbed = tuple(glob(file_arg[0]))
        if not globbed:
            raise Exception(f"Could not find a single file that matches {file_arg}")
        file_arg = globbed
    for file_path in file_arg:
        _check_if_file_exists(".", file_path, required=True)
        _check_if_file_is_readable(".", file_path, required=True)
    return file_arg


def _set_n_threads(ctx, param, value):
    """
    Set the number of CPU threads for numba.
    Arguments come straight from click option.
    """
    if value == -1:
        from numba import config

        return config.NUMBA_NUM_THREADS
    elif value == 0:
        return 1
    else:
        return value


# command group
@click.group(
    cls=OrderedGroup,
    help_headers_color="bright_white",
    help_options_color="green",
    help=f"""
        \b
                    {style("|", fg="blue")}
        {style(",---. ,---. |---. ,---.", fg="blue")} \
{style("version " + __version__, fg="green")}
        {style("`---. |     |   | `---.", fg="blue")}
        {style("`---' `---' `---' `---'", fg="blue")}

        Below you find a list of all available commands.
        To find out what they do and how to use them, check
        their help like this:

        {style("scbs [command] --help", fg="blue")}

        To use stdin or stdout, use the dash character
        {style("-", fg="blue")} instead of a file path.
        """,
)
@click.version_option(version=__version__)
def cli():
    pass


# prepare command
@cli.command(
    name="prepare",
    help=f"""
    Gathers single cell methylation data from multiple input files
    (one per cell) and creates a sparse matrix (position x cell) in CSR
    format for each chromosome. Methylated sites are represented by a 1,
    unmethylated sites are -1, missing values and other bases are 0.

    {style("INPUT_FILES", fg="green")} are single cell
    methylation files, for example '.cov'-files generated by Bismark.

    {style("DATA_DIR", fg="green")} is the output directory
    where the methylation data will be stored.

    Note: If you have many cells and encounter a "too many open files"-
    error, you need to increase the open file limit with e.g.
    'ulimit -n 99999'.
    """,
    short_help="Collect and store sc-methylation data for quick access",
    no_args_is_help=True,
)
@click.argument("input-files", nargs=-1)
@click.argument(
    "data-dir",
    type=click.Path(dir_okay=True, file_okay=False, writable=True),
)
@click.option(
    "--round-sites",
    is_flag=True,
    help="""
    Specify that you want to round methylation sites with ambigous status to
    0% or 100%. This means, for example, that a CpG site with 5 methylated reads and
    1 unmethylated read will be considered methylated in that cell. Otherwise, ambiguous
    sites will be discarded. Note that sites with the same number of methylated and
    unmethylated reads will always be discarded.
    """,
)
@click.option(
    "--chunksize",
    type=click.IntRange(min=1),
    metavar="INTEGER",
    default=1e7,
    help="""
    The data of each chromosome is read in chunks [default: 10 Mbp] to reduce memory
    requirements. If you are running out of RAM, decrease the chunk size (in bp).
    """,
)
@click.option(
    "--input-format",
    default="bismark",
    help="""
    Specify the format of the input files. Options: 'bismark' (default),
    'methylpy', 'allc', 'biscuit', 'biscuit_short' or custom (see below).

    \b
    You can specify a custom format by specifying the separator, whether the
    file has a header, and which information is stored in which columns. These
    values should be separated by ':' and enclosed by quotation marks, for
    example --input-format '1:2:3:4u:\\t:1'

    \b
    The six ':'-separated values denote the number of the columns that contain
    1. the chromosome name
    2. the genomic position
    3. the methylated counts
    4. either the total coverage (c) or the number of unmethylated counts (u), followed
    by either 'c' or 'u', e.g. '4c' to denote that the 4th column contains the coverage
    5. The separator, e.g. '\\t' or 'TAB' for tsv files or ',' for csv
    6. Either '1' if the file has a header or '0' if it does not have a header
    All column numbers are 1-indexed, i.e. to define the first column use '1' and not
    '0'.""",
)
def prepare_cli(**kwargs):
    from .prepare import prepare

    timer = Timer(label="prepare")
    kwargs["input_files"] = _get_input_file_paths(kwargs["input_files"])
    _print_kwargs(kwargs)
    prepare(**kwargs)
    timer.stop()


# filter command
@cli.command(
    name="filter",
    help=f"""
    Filters low-quality cells based on the number of observed methylation sites
    and/or the global methylation percentage.

    Alternatively, you may also provide a text file with the names of the cells you
    want to keep.

    {style("DATA_DIR", fg="green")} is the unfiltered directory containing the
    methylation matrices produced by running 'scbs prepare'.

    {style("FILTERED_DIR", fg="green")} is the output directory storing methylation
    data only for the cells that passed all filtering criteria.
    """,
    short_help="Filter low-quality cells based on coverage and mean methylation",
    no_args_is_help=True,
)
@click.argument(
    "data-dir",
    type=click.Path(exists=True, dir_okay=True, file_okay=False, readable=True),
)
@click.argument(
    "filtered-dir",
    type=click.Path(dir_okay=True, file_okay=False, writable=True),
)
@click.option(
    "--min-sites",
    type=click.IntRange(min=1),
    metavar="INTEGER",
    help="Minimum number of methylation sites required for a cell to pass filtering.",
)
@click.option(
    "--max-sites",
    type=click.IntRange(min=1),
    metavar="INTEGER",
    help="Maximum number of methylation sites required for a cell to pass filtering.",
)
@click.option(
    "--min-meth",
    type=click.FloatRange(min=0, max=100),
    metavar="PERCENT",
    help="Minimum average methylation percentage required for a cell to "
    "pass filtering.",
)
@click.option(
    "--max-meth",
    type=click.FloatRange(min=0, max=100),
    metavar="PERCENT",
    help="Maximum average methylation percentage required for a cell to "
    "pass filtering.",
)
@click.option(
    "--cell-names",
    type=click.File("r"),
    help="A text file with the names of the cells you want to keep (default) "
    "or remove. "
    "This is an alternative to the min/max filtering options. Each cell name "
    "must be on a new line.",
)
@click.option(
    "--keep/--discard",
    default=True,
    help="Specify whether the cells listed in your text file should be kept (default) "
    "or discarded from the data set. Only use together with --cell-names.",
)
def filter_cli(**kwargs):
    from .filter import filter_

    timer = Timer(label="filter")
    _print_kwargs(kwargs)
    filter_(**kwargs)
    timer.stop()


# smooth command
@cli.command(
    name="smooth",
    help=f"""
    This script will calculate the smoothed mean methylation over the
    whole genome.

    {style("DATA_DIR", fg="green")} is the directory containing the methylation matrices
    produced by running 'scbs prepare'.

    The smoothed methylation values will be written to
    {style("DATA_DIR/smoothed/", fg="green")}.
    """,
    short_help="Smooth the pseudobulk of single cell methylation data",
    no_args_is_help=True,
)
@click.argument(
    "data-dir",
    type=click.Path(
        exists=True, dir_okay=True, file_okay=False, readable=True, writable=True
    ),
)
@click.option(
    "-bw",
    "--bandwidth",
    default=1000,
    type=click.IntRange(min=1),
    metavar="INTEGER",
    show_default=True,
    help="Smoothing bandwidth in basepairs.",
)
@click.option(
    "--use-weights",
    is_flag=True,
    help="Use this to weigh each methylation site by log1p(coverage)."
    "  [default: off]",
)
def smooth_cli(**kwargs):
    from .smooth import smooth

    timer = Timer(label="smooth")
    _print_kwargs(kwargs)
    smooth(**kwargs)
    timer.stop()


# scan command
@cli.command(
    name="scan",
    help=f"""
    Scans the whole genome for variably methylated regions (VMRs). This works by
    sliding a window across the genome, calculating the variance of methylation per
    window, and selecting windows above a variance threshold.

    {style("DATA_DIR", fg="green")} is the directory containing the methylation
    matrices produced by running 'scbs prepare', as well as the smoothed methylation
    values produced by running 'scbs smooth'.

    {style("OUTPUT", fg="green")} is the path of the output file in '.bed' format,
    containing the VMRs that were found.
    """,
    short_help="Scan the genome to discover variably methylated regions",
    no_args_is_help=True,
)
@click.argument(
    "data-dir",
    type=click.Path(exists=True, dir_okay=True, file_okay=False, readable=True),
)
@click.argument("output", type=click.File("w"))
@click.option(
    "-bw",
    "--bandwidth",
    default=2000,
    type=click.IntRange(min=1),
    metavar="INTEGER",
    show_default=True,
    help="Bandwidth of the sliding window in basepairs. Increase this "
    "value to find larger VMRs.",
)
@click.option(
    "--stepsize",
    default=100,
    type=click.IntRange(min=1),
    metavar="INTEGER",
    show_default=True,
    help="Step size of the sliding window in basepairs. Increase "
    "this value to gain speed, at the cost of some accuracy.",
)
@click.option(
    "--var-threshold",
    default=0.02,
    show_default=True,
    type=click.FloatRange(min=0, max=1),
    metavar="FLOAT",
    help="The variance threshold, i.e. 0.02 means that the top 2% "
    "most variable genomic bins will be reported. Overlapping variable bins "
    "are merged.",
)
@click.option(
    "--min-cells",
    default=6,
    type=click.IntRange(min=1),
    metavar="INTEGER",
    show_default=True,
    help="The minimum number of cells required to report a VMR. For example, "
    "a value of 6 means that only VMRs with sequencing coverage in at least 6 "
    "cells are reported.",
)
@click.option(
    "--threads",
    default=-1,
    help="How many CPU threads to use in parallel.  [default: all available]",
    callback=_set_n_threads,
)
@click.option(
    "--write-header",
    is_flag=True,
    help="Write the column names of the output file.  [default: off]",
)
def scan_cli(**kwargs):
    from .scbs import scan

    timer = Timer(label="scan")
    _print_kwargs(kwargs)
    scan(**kwargs)
    timer.stop()


# diff command
@cli.command(
    name="diff",
    help=f"""
    Scans the whole genome for differentially methylated regions (DMRs)
    between two groups of cells. This works by sliding a window across the genome,
    performing a t-test for each window, and merging windows above a threshold.
    To control the false discovery rate, the same procedure is repeated on permutations
    of the data which are then used to calculate an adjusted p-value for each DMR.

    {style("DATA_DIR", fg="green")} is the directory containing the methylation
    matrices produced by running 'scbs prepare', as well as the smoothed methylation
    values produced by running 'scbs smooth'.

    {style("CELL_GROUPS", fg="green")} is a comma-separated text file that lists the
    group membership (e.g. cell type or treatment) of each cell. Each row contains two
    comma-separated values: The cell name and its group label. Cell names are denoted in
    {style("DATA_DIR/column_header.txt", fg="green")}. Only two cell groups can be
    compared, so there should only be two unique group labels (e.g. "neuron" and "glia"
    or "wildtype" and "KO"). To exclude cells that do not belong to either group from
    the analysis, you can assign them the group label '-' (dash character).

    {style("OUTPUT", fg="green")} is the path of the output file in '.bed' format,
    containing the DMR genome coordinates, their t-statistic, the cell group in which
    the DMR has lower methylation, and the adjusted p-value.
    """,
    short_help="Discover differentially methylated regions between groups of cells",
    no_args_is_help=True,
)
@click.argument(
    "data-dir",
    type=click.Path(exists=True, dir_okay=True, file_okay=False, readable=True),
)
@click.argument("cell_groups", type=click.File("r"))
@click.argument("output", type=click.File("w"))
@click.option(
    "-bw",
    "--bandwidth",
    default=2000,
    type=click.IntRange(min=1),
    metavar="INTEGER",
    show_default=True,
    help="Bandwidth of the sliding window in basepairs. Increase this "
    "value to find larger DMRs.",
)
@click.option(
    "--stepsize",
    default=1000,
    type=click.IntRange(min=1),
    metavar="INTEGER",
    show_default=True,
    help="Step size of the sliding window in basepairs. Increase "
    "this value to gain speed, at the cost of some accuracy.",
)
@click.option(
    "--threshold",
    default=0.02,
    show_default=True,
    type=click.FloatRange(min=0, max=1),
    metavar="FLOAT",
    help="The t-statistic threshold, i.e. 0.02 means that the top 2% "
    "most differentially methylated genomic bins will be reported. "
    "Overlapping bins are merged.",
)
@click.option(
    "--min-cells",
    default=6,
    type=click.IntRange(min=1),
    metavar="INTEGER",
    show_default=True,
    help="The minimum number of cells required to consider a genomic region for "
    "testing. For example, a value of 6 means that only regions with sequencing "
    "coverage in at least 6 cells per group are considered.",
)
@click.option(
    "--threads",
    default=-1,
    help="How many CPU threads to use in parallel.  [default: all available]",
    callback=_set_n_threads,
)
@click.option(
    "--write-header",
    is_flag=True,
    help="Write the column names of the output file.  [default: off]",
)
@click.option(
    "--debug",
    is_flag=True,
    help="Use this to also report DMRs that were identified in permutations.  "
    "[default: off]",
)
def diff_cli(**kwargs):
    from .diff import diff

    timer = Timer(label="diff")
    _print_kwargs(kwargs)
    diff(**kwargs)
    timer.stop()


# matrix command (makes a "count" matrix)
@cli.command(
    name="matrix",
    help=f"""
    From single cell methylation or NOMe-seq data, calculates the average methylation
    in genomic regions for every cell. The output is a long table that can be used e.g.
    for dimensionality reduction or clustering, analogous to a count matrix in
    scRNA-seq.

    {style("REGIONS", fg="green")} is a .bed file of regions for which methylation
    will be quantified in every cell.

    {style("DATA_DIR", fg="green")} is the directory containing the methylation
    matrices produced by running 'scbs prepare', as well as the smoothed methylation
    values produced by running 'scbs smooth'.

    \b
    {style("OUTPUT_DIR", fg="green")} is the output directory.
    It will contain four cell × region matrices ("count tables"):
    'methylated_sites.csv.gz': the number of sites that were methylated
    'total_sites.csv.gz': the total number of observed sites (sites with read coverage)
    'methylation_fractions.csv.gz': the average methylation, calculated as:
        # of methylated sites / # of total observed sites
    'mean_shrunken_residuals.csv.gz': the mean shrunken residuals, a more accurate
        measure of methylation in a genomic region.
    """,
    short_help="Make a methylation matrix, similar to a count matrix in scRNA-seq",
    no_args_is_help=True,
)
@click.argument("regions", type=click.File("r"))
@click.argument(
    "data-dir",
    type=click.Path(exists=True, dir_okay=True, file_okay=False, readable=True),
)
@click.argument("output_dir", type=click.Path(file_okay=False, writable=True))
@click.option(
    "--sparse",
    is_flag=True,
    help="Write the output as a sparse matrix, instead of the four .csv.gz files "
    "described above. This is faster and more space-efficient for large data sets. "
    "The output 'matrix.mtx.gz' contains four columns: row_index, col_index, shrunken "
    "residuals, and methylation fractions. Both indices are 1-indexed  [default: off]",
)
@click.option(
    "--threads",
    default=-1,
    help="How many CPU threads to use in parallel.  [default: all available]",
    callback=_set_n_threads,
)
def matrix_cli(**kwargs):
    from .matrix import matrix, matrix_sparse

    timer = Timer(label="matrix")
    _print_kwargs(kwargs)
    sparse_output = kwargs["sparse"]
    del kwargs["sparse"]
    if sparse_output:
        matrix_sparse(**kwargs)
    else:
        matrix(**kwargs)
    timer.stop()


# profile command
@cli.command(
    name="profile",
    help=f"""
    From single cell methylation or NOMe-seq data,
    calculates the average methylation profile of a set of
    genomic regions. Useful for plotting and visually comparing
    methylation between groups of regions or cells.

    {style("REGIONS", fg="green")} is an alphabetically sorted (!) .bed file of regions
    for which the methylation profile will be produced.

    {style("DATA_DIR", fg="green")} is the directory containing the methylation matrices
    produced by running 'scbs prepare'.

    {style("OUTPUT", fg="green")} is the file path where the methylation profile data
    will be written. Should end with '.csv'.
    """,
    short_help="Plot mean methylation around a group of genomic features",
    no_args_is_help=True,
)
@click.argument("regions", type=click.File("r"))
@click.argument(
    "data-dir",
    type=click.Path(exists=True, dir_okay=True, file_okay=False, readable=True),
)
@click.argument("output", type=click.File("w"))
@click.option(
    "--width",
    default=4000,
    show_default=True,
    type=click.IntRange(min=1),
    metavar="INTEGER",
    help="The total width of the profile plot in bp. "
    "The center of all bed regions will be "
    "extended in both directions by half of this amount. "
    "Shorter regions will be extended, longer regions "
    "will be shortened accordingly.",
)
@click.option(
    "--strand-column",
    type=click.IntRange(min=1),
    metavar="INTEGER",
    help="The bed column number (1-indexed) denoting "
    "the DNA strand of the region  [optional].",
)
@click.option(
    "--label",
    help="Specify a constant value to be added as a "
    "column to the output table. This can be "
    "useful to give each output a unique label when "
    "you want to concatenate multiple outputs  [optional].",
)
def profile_cli(**kwargs):
    from .profile import profile

    timer = Timer(label="profile")
    _print_kwargs(kwargs)
    profile(**kwargs)
    timer.stop()


# CLI template:
# @cli.command(
#     help=f"""
#     Blabla

#     {style("INPUT", fg="green")} blabla

#     {style("OUTPUT", fg="green")} blabla
#     """,
#     short_help="template for dev",
#     no_args_is_help=True,
# )
# @click.argument("input", type=click.File("r"))
# @click.argument("output", type=click.File("w"))
# @click.option("-o", "--option", type=int, default=4, show_default=True)
# @click.option("--flag", is_flag=True)
# def template(**kwargs):
#     timer = Timer(label="template")
#     _print_kwargs(kwargs)
#     timer.stop()
