import sys
import os
import pickle
import json
import subprocess
from pathlib import Path
import numpy as np
import pandas as pd
import papermill as pm  # Render the ipython notebook
from kipoi_utils.external.flatten_json import flatten, unflatten
from weakref import WeakValueDictionary
import uuid
from threading import Lock
import logging
logger = logging.getLogger(__name__)
logger.addHandler(logging.NullHandler())


def jupyter_nbconvert(input_ipynb):
    # NOTE: cwd is used since the input_ipynb could contain some strange output
    # characters like '[' which would mess up the paths
    subprocess.call(["jupyter",
                     "nbconvert",
                     os.path.basename(input_ipynb),
                     "--to", "html"],
                    cwd=os.path.dirname(input_ipynb))


def render_ipynb(template_ipynb, rendered_ipynb, params=dict()):
    """Render the ipython notebook

    Args:
      template_ipynb: template ipython notebook where one cell defines the following metadata:
        {"tags": ["parameters"]}
      render_ipynb: output ipython notebook path
      params: parameters used to execute the ipython notebook
    """
    import jupyter_client

    os.makedirs(os.path.dirname(rendered_ipynb), exist_ok=True)
    kernel_name = os.environ.get("CONDA_DEFAULT_ENV", 'python3')
    if kernel_name not in jupyter_client.kernelspec.find_kernel_specs():
        logger.info(f"Installing the ipython kernel for the current conda environment: {kernel_name}")
        from ipykernel.kernelspec import install
        install(user=True, kernel_name=kernel_name)

    pm.execute_notebook(
        template_ipynb,  # input template
        rendered_ipynb,
        kernel_name=kernel_name,  # default kernel
        parameters=params
    )
    jupyter_nbconvert(rendered_ipynb)


def tqdm_restart():
    """Restart tqdm to not print to every line
    """
    from tqdm import tqdm as tqdm_cls
    inst = tqdm_cls._instances
    for i in range(len(inst)):
        inst.pop().close()


def touch_file(file, verbose=True):
    import subprocess
    if verbose:
        add = "v"
    else:
        add = ""
    subprocess.run(["vmtouch", f'-{add}tf', file])


def remove_exists(output_path, overwrite=False):
    if os.path.exists(output_path):
        if overwrite:
            os.remove(output_path)
        else:
            raise ValueError(f"File exists {str(output_path)}. Use overwrite=True to overwrite it")


def write_pkl(obj, fname, create_dirs=True, protocol=2):
    import cloudpickle
    if create_dirs:
        if os.path.dirname(fname):
            os.makedirs(os.path.dirname(fname), exist_ok=True)
    cloudpickle.dump(obj, open(fname, 'wb'), protocol=protocol)


def read_pkl(fname):
    import cloudpickle
    return cloudpickle.load(open(fname, 'rb'))


def read_json(fname):
    with open(fname) as f:
        return json.load(f)


class NumpyAwareJSONEncoder(json.JSONEncoder):

    def default(self, obj):
        if isinstance(obj, Path):
            return str(obj)
        elif isinstance(obj, np.ndarray):
            return obj.tolist()
        elif isinstance(obj, np.generic):
            return obj.item()
        return json.JSONEncoder.default(self, obj)


def write_json(obj, fname, **kwargs):
    with open(fname, "w") as f:
        return json.dump(obj, f, cls=NumpyAwareJSONEncoder, **kwargs)


dump = write_pkl
load = read_pkl


def _listify(arg):
    if hasattr(type(arg), '__len__'):
        return arg
    return [arg, ]


def reverse_complement(seq):
    alt_map = {'ins': '0'}
    complement = {'A': 'T', 'C': 'G', 'G': 'C', 'T': 'A'}
    for k, v in alt_map.items():
        seq = seq.replace(k, v)
    bases = list(seq)
    bases = reversed([complement.get(base, base) for base in bases])
    bases = ''.join(bases)
    for k, v in alt_map.items():
        bases = bases.replace(v, k)
    return bases


def related_dump_yaml(obj, path, verbose=False):
    import related
    generated_yaml = related.to_yaml(obj,
                                     suppress_empty_values=False,
                                     suppress_map_key_values=True)  # .strip()
    if verbose:
        print(generated_yaml)

    with open(path, "w") as f:
        f.write(generated_yaml)


def shuffle_list(l):
    return pd.Series(l).sample(frac=1).tolist()


def flatten_list(l):
    """Flattens a nested list
    """
    return [x for nl in l for x in nl]


class Logger(object):
    """tee functionality in python. If this object exists,
    then all of stdout gets logged to the file

    Adoped from:
    https://stackoverflow.com/questions/616645/how-do-i-duplicate-sys-stdout-to-a-log-file-in-python/3423392#3423392
    """

    def __init__(self, name, mode='a'):
        self.file = open(name, mode)
        self.stdout = sys.stdout
        sys.stdout = self

    def __del__(self):
        sys.stdout = self.stdout
        self.file.close()

    def write(self, data):
        self.file.write(data)
        self.stdout.write(data)
        # flush right away
        self.file.flush()
        self.stdout.flush()

    def flush(self):
        self.file.flush()
        self.stdout.flush()


def add_file_logging(output_dir, logger, name='stdout'):
    os.makedirs(os.path.join(output_dir, 'log'), exist_ok=True)
    log = Logger(os.path.join(output_dir, 'log', name + '.log'), 'a+')  # log to the file
    fh = logging.FileHandler(os.path.join(output_dir, 'log', name + '.log'), 'a+')
    fh.setFormatter(logging.Formatter('[%(asctime)s] - [%(levelname)s] - %(message)s'))
    fh.setLevel(logging.INFO)
    logger.addHandler(fh)
    return log


def halve(n):
    """Halve an integer"""
    return n // 2 + n % 2, n // 2


def expand_str_list(l, prefix="", suffix=""):
    """add strings to the beginning or to the end of the string
    """
    return [prefix + x + suffix for x in l]


def kv_string2dict(s):
    """Convert a key-value string: k=v,k2=v2,... into a dictionary
    """
    import yaml
    return yaml.load(s.replace(",", "\n").replace("=", ": "))


def dict_suffix_key(d, suffix):
    return {k + suffix: v for k, v in d.items()}


def dict_prefix_key(d, prefix):
    return {prefix + k: v for k, v in d.items()}


def kwargs_str2kwargs(hparams):
    """Converts a string to a dictionary of kwargs

    In[22]: params_str2kwargs("a=1;b=[1,2]")
    Out[22]: {'a': 1, 'b': [1, 2]}

    In[30]: hparams_str2kwargs("a=null")
    Out[30]: {'a': None}

    """
    import yaml
    return yaml.load(hparams.replace(";", "\n").replace("=", ": "))


def apply_parallel(df_grouped, func, n_jobs=-1, verbose=True):
    from joblib import Parallel, delayed
    import pandas as pd
    from tqdm import tqdm
    retLst = Parallel(n_jobs=n_jobs)(delayed(func)(group)
                                     for name, group in tqdm(df_grouped, disable=not verbose))
    return pd.concat(retLst)


def get_timestamp():
    """Get current time-stamp as a string: 2018-12-10_14:20:04
    """
    import datetime
    import time
    return datetime.datetime.fromtimestamp(time.time()).strftime('%Y-%m-%d_%H:%M:%S')


class ConditionalRun:
    """Simple class keeping track whether the command has already
    been ran or not
    """

    def __init__(self, main_cmd, cmd, output_dir, force=False):
        self.main_cmd = main_cmd
        self.cmd = cmd
        self.output_dir = output_dir
        self.force = force

    def set_cmd(self, cmd):
        self.cmd = cmd
        return self

    def get_done_file(self):
        return os.path.join(self.output_dir, f".{self.main_cmd}/{self.cmd}.done")

    def done(self):
        ret = os.path.exists(self.get_done_file())
        if self.force:
            # always run the command
            ret = False
        if ret:
            logger.info(f"Skipping {self.cmd}")
        else:
            logger.info(f"Running {self.cmd}")
        return ret

    def write(self):
        fname = self.get_done_file()
        os.makedirs(os.path.dirname(fname),
                    exist_ok=True)
        with open(fname, "w") as f:
            f.write(get_timestamp())


class Logger(object):
    """tee functionality in python. If this object exists,
    then all of stdout gets logged to the file
    Adoped from:
    https://stackoverflow.com/questions/616645/how-do-i-duplicate-sys-stdout-to-a-log-file-in-python/3423392#3423392
    """

    def __init__(self, name, mode='a'):
        self.file = open(name, mode)
        self.stdout = sys.stdout
        sys.stdout = self

    def __del__(self):
        sys.stdout = self.stdout
        self.file.close()

    def write(self, data):
        self.file.write(data)
        self.stdout.write(data)
        # flush right away
        self.file.flush()
        self.stdout.flush()

    def flush(self):
        self.file.flush()
        self.stdout.flush()


def fnmatch_any(s, patterns):
    from fnmatch import fnmatch  # unix-style pattern matching
    return any([fnmatch(s, p) for p in patterns])


def to_list(l):
    if isinstance(l, list):
        return l
    else:
        return [l]


def pd_first_cols(df: pd.DataFrame, cols):
    """Set `cols` to be the first columns in pd.DataFrame df
    """
    for c in cols:
        assert c in df
    other_cols = [c for c in df.columns if c not in cols]
    return df[cols + other_cols]


def pd_col_prepend(df: pd.DataFrame, column, prefix='', suffix=''):
    """Add a prefix or suffix to all the columns names in pd.DataFrame
    """
    if isinstance(column, list):
        for c in column:
            df[c] = prefix + df[c] + suffix
    else:
        df[column] = prefix + df[column] + suffix
    return df


def create_tf_session(visiblegpus, per_process_gpu_memory_fraction=0.45):
    import os
    import tensorflow as tf
    import keras.backend as K
    os.environ['CUDA_VISIBLE_DEVICES'] = str(visiblegpus)
    session_config = tf.ConfigProto()
    # session_config.gpu_options.deferred_deletion_bytes = DEFER_DELETE_SIZE
    session_config.gpu_options.per_process_gpu_memory_fraction = per_process_gpu_memory_fraction
    session = tf.Session(config=session_config)
    K.set_session(session)
    return session


class SerializableLock(object):
    _locks = WeakValueDictionary()
    """ A Serializable per-process Lock

    This wraps a normal ``threading.Lock`` object and satisfies the same
    interface.  However, this lock can also be serialized and sent to different
    processes.  It will not block concurrent operations between processes (for
    this you should look at ``multiprocessing.Lock`` or ``locket.lock_file``
    but will consistently deserialize into the same lock.

    So if we make a lock in one process::

        lock = SerializableLock()

    And then send it over to another process multiple times::

        bytes = pickle.dumps(lock)
        a = pickle.loads(bytes)
        b = pickle.loads(bytes)

    Then the deserialized objects will operate as though they were the same
    lock, and collide as appropriate.

    This is useful for consistently protecting resources on a per-process
    level.

    The creation of locks is itself not threadsafe.

    This class was taken from dask.utils.py

    ﻿Copyright (c) 2014-2018, Anaconda, Inc. and contributors
    All rights reserved.

    Redistribution and use in source and binary forms, with or without modification,
    are permitted provided that the following conditions are met:

    Redistributions of source code must retain the above copyright notice,
    this list of conditions and the following disclaimer.

    Redistributions in binary form must reproduce the above copyright notice,
    this list of conditions and the following disclaimer in the documentation
    and/or other materials provided with the distribution.

    Neither the name of Anaconda nor the names of any contributors may be used to
    endorse or promote products derived from this software without specific prior
    written permission.

    THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
    AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
    IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
    ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
    LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
    CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
    SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
    INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
    CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
    ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF
    THE POSSIBILITY OF SUCH DAMAGE.
    """

    def __init__(self, token=None):
        self.token = token or str(uuid.uuid4())
        if self.token in SerializableLock._locks:
            self.lock = SerializableLock._locks[self.token]
        else:
            self.lock = Lock()
            SerializableLock._locks[self.token] = self.lock

    def acquire(self, *args, **kwargs):
        return self.lock.acquire(*args, **kwargs)

    def release(self, *args, **kwargs):
        return self.lock.release(*args, **kwargs)

    def __enter__(self):
        self.lock.__enter__()

    def __exit__(self, *args):
        self.lock.__exit__(*args)

    def locked(self):
        return self.lock.locked()

    def __getstate__(self):
        return self.token

    def __setstate__(self, token):
        self.__init__(token)

    def __str__(self):
        return "<%s: %s>" % (self.__class__.__name__, self.token)

    __repr__ = __str__
