from __future__ import print_function

import os
import sys

import numpy as np
from tqdm import tqdm_notebook


def parse_enum(enum, v):
    if isinstance(v, enum):
        return v
    else:
        return enum[v]


def atoms_to_numpy(f, mol, dtype='float'):
    return np.fromiter(
        (f(a) for a in mol.GetAtoms()),
        dtype, mol.GetNumAtoms()
    )


def conformer_to_numpy(conf):
    return np.array(
        [list(conf.GetAtomPosition(i)) for i in range(conf.GetNumAtoms())]
    )


class Capture(object):
    def __init__(self, target='stderr'):
        self.target = target
        self.orig = getattr(sys, target)
        self.result = []

    def write(self, text):
        self.result.append(text)

    def flush(self):
        pass

    def __enter__(self):
        setattr(sys, self.target, self)
        return self

    def __exit__(self, *args):
        setattr(sys, self.target, self.orig)


class DummyBar(object):
    def __init__(self, *args, **kwargs):
        pass

    def __enter__(self):
        return self

    def __exit__(self, *args, **kws):
        pass

    def update(self, *args, **kws):
        pass

    @classmethod
    def write(cls, s, file=sys.stdout, end='\n'):
        print(s, file=file, end=end)


class NotebookWrapper(object):
    def __init__(self, **kwargs):
        self.bar = tqdm_notebook(**kwargs)

    def __enter__(self):
        return self

    def __exit__(self, *args):
        pass

    def update(self, *args, **kwargs):
        self.bar.update(*args, **kwargs)

    def write(self, *args, **kwargs):
        self.bar.update(*args, **kwargs)


def PathType(string):
    if not os.path.isfile(string):
        raise ValueError('file not exists: {}'.format(string))

    return string


def module_prog(pkg):
    return '{} -m {}'.format(os.path.basename(sys.executable), pkg)
