import os
from timeit import default_timer as timer
import unittest
import pytest
from decorator import decorator

from hail.utils.java import Env
import hail as hl
from hail.backend.local_backend import LocalBackend

_initialized = False


def startTestHailContext():
    global _initialized
    if not _initialized:
        backend_name = os.environ.get('HAIL_QUERY_BACKEND', 'spark')
        if backend_name == 'spark':
            hl.init(master='local[1]', min_block_size=0, quiet=True)
        else:
            Env.hc()  # force initialization
        _initialized = True


def stopTestHailContext():
    pass

_test_dir = os.environ.get('HAIL_TEST_RESOURCES_DIR', '../src/test/resources')
_doctest_dir = os.environ.get('HAIL_DOCTEST_DATA_DIR', 'hail/docs/data')


def resource(filename):
    return os.path.join(_test_dir, filename)


def doctest_resource(filename):
    return os.path.join(_doctest_dir, filename)


def schema_eq(x, y):
    x_fds = dict(x)
    y_fds = dict(y)
    return x_fds == y_fds


def convert_struct_to_dict(x):
    if isinstance(x, hl.Struct):
        return {k: convert_struct_to_dict(v) for k, v in x._fields.items()}
    elif isinstance(x, list):
        return [convert_struct_to_dict(elt) for elt in x]
    elif isinstance(x, tuple):
        return tuple([convert_struct_to_dict(elt) for elt in x])
    elif isinstance(x, dict):
        return {k: convert_struct_to_dict(v) for k, v in x.items()}
    else:
        return x


_dataset = None


def get_dataset():
    global _dataset
    if _dataset is None:
        _dataset = hl.split_multi_hts(hl.import_vcf(resource('sample.vcf'))).cache()
    return _dataset

def assert_time(f, max_duration):
    start = timer()
    x = f()
    end = timer()
    assert (start - end) < max_duration
    print(f'took {end - start:.3f}')
    return x

def create_all_values():
    return hl.struct(
        f32=hl.float32(3.14),
        i64=hl.int64(-9),
        m=hl.null(hl.tfloat64),
        astruct=hl.struct(a=hl.null(hl.tint32), b=5.5),
        mstruct=hl.null(hl.tstruct(x=hl.tint32, y=hl.tstr)),
        aset=hl.set(['foo', 'bar', 'baz']),
        mset=hl.null(hl.tset(hl.tfloat64)),
        d=hl.dict({hl.array(['a', 'b']): 0.5, hl.array(['x', hl.null(hl.tstr), 'z']): 0.3}),
        md=hl.null(hl.tdict(hl.tint32, hl.tstr)),
        h38=hl.locus('chr22', 33878978, 'GRCh38'),
        ml=hl.null(hl.tlocus('GRCh37')),
        i=hl.interval(
            hl.locus('1', 999),
            hl.locus('1', 1001)),
        c=hl.call(0, 1),
        mc=hl.null(hl.tcall),
        t=hl.tuple([hl.call(1, 2, phased=True), 'foo', hl.null(hl.tstr)]),
        mt=hl.null(hl.ttuple(hl.tlocus('GRCh37'), hl.tbool)),
        nd=hl.nd.arange(0, 10).reshape((2, 5)),
    )

def prefix_struct(s, prefix):
    return hl.struct(**{prefix + k: s[k] for k in s})

def create_all_values_table():
    all_values = create_all_values()
    return (hl.utils.range_table(5, n_partitions=3)
            .annotate_globals(**prefix_struct(all_values, 'global_'))
            .annotate(**all_values)
            .cache())

def create_all_values_matrix_table():
    all_values = create_all_values()
    return (hl.utils.range_matrix_table(3, 2, n_partitions=2)
            .annotate_globals(**prefix_struct(all_values, 'global_'))
            .annotate_rows(**prefix_struct(all_values, 'row_'))
            .annotate_cols(**prefix_struct(all_values, 'col_'))
            .annotate_entries(**prefix_struct(all_values, 'entry_'))
            .cache())

def create_all_values_datasets():
    return (create_all_values_table(), create_all_values_matrix_table())

def skip_unless_spark_backend():
    from hail.backend.spark_backend import SparkBackend
    @decorator
    def wrapper(func, *args, **kwargs):
        if isinstance(hl.utils.java.Env.backend(), SparkBackend):
            return func(*args, **kwargs)
        else:
            raise unittest.SkipTest('requires Spark')

    return wrapper


fails_local_backend = pytest.mark.xfail(
    os.environ.get('HAIL_QUERY_BACKEND') == 'local',
    reason="doesn't yet work on local backend",
    strict=True)

def run_with_cxx_compile():
    @decorator
    def wrapper(func, *args, **kwargs):
        return

    return wrapper


def assert_evals_to(e, v):
    res = hl.eval(e)
    if res != v:
        raise ValueError(f'  actual: {res}\n  expected: {v}')


def assert_all_eval_to(*expr_and_expected):
    exprs, expecteds = zip(*expr_and_expected)
    assert_evals_to(hl.tuple(exprs), expecteds)


def lower_only():
    @decorator
    def wrapper(func, *args, **kwargs):
        flags = hl._get_flags()
        prev_lower = flags.get('lower')
        prev_lower_only = flags.get('lower_only')

        hl._set_flags(lower='1', lower_only='1')

        try:
            return func(*args, **kwargs)
        finally:
            hl._set_flags(lower=prev_lower, lower_only=prev_lower_only)
    return wrapper
[((28, 12, 28, 78), 'os.environ.get', 'os.environ.get', ({(28, 27, 28, 52): '"""HAIL_TEST_RESOURCES_DIR"""', (28, 54, 28, 77): '"""../src/test/resources"""'}, {}), "('HAIL_TEST_RESOURCES_DIR', '../src/test/resources')", False, 'import os\n'), ((29, 15, 29, 72), 'os.environ.get', 'os.environ.get', ({(29, 30, 29, 53): '"""HAIL_DOCTEST_DATA_DIR"""', (29, 55, 29, 71): '"""hail/docs/data"""'}, {}), "('HAIL_DOCTEST_DATA_DIR', 'hail/docs/data')", False, 'import os\n'), ((33, 11, 33, 44), 'os.path.join', 'os.path.join', ({(33, 24, 33, 33): '_test_dir', (33, 35, 33, 43): 'filename'}, {}), '(_test_dir, filename)', False, 'import os\n'), ((37, 11, 37, 47), 'os.path.join', 'os.path.join', ({(37, 24, 37, 36): '_doctest_dir', (37, 38, 37, 46): 'filename'}, {}), '(_doctest_dir, filename)', False, 'import os\n'), ((69, 12, 69, 19), 'timeit.default_timer', 'timer', ({}, {}), '()', True, 'from timeit import default_timer as timer\n'), ((71, 10, 71, 17), 'timeit.default_timer', 'timer', ({}, {}), '()', True, 'from timeit import default_timer as timer\n'), ((100, 11, 100, 53), 'hail.struct', 'hl.struct', ({}, {}), '(**{(prefix + k): s[k] for k in s})', True, 'import hail as hl\n'), ((147, 10, 147, 20), 'hail.eval', 'hl.eval', ({(147, 18, 147, 19): 'e'}, {}), '(e)', True, 'import hail as hl\n'), ((17, 23, 17, 68), 'os.environ.get', 'os.environ.get', ({(17, 38, 17, 58): '"""HAIL_QUERY_BACKEND"""', (17, 60, 17, 67): '"""spark"""'}, {}), "('HAIL_QUERY_BACKEND', 'spark')", False, 'import os\n'), ((134, 4, 134, 40), 'os.environ.get', 'os.environ.get', ({(134, 19, 134, 39): '"""HAIL_QUERY_BACKEND"""'}, {}), "('HAIL_QUERY_BACKEND')", False, 'import os\n'), ((154, 20, 154, 35), 'hail.tuple', 'hl.tuple', ({(154, 29, 154, 34): 'exprs'}, {}), '(exprs)', True, 'import hail as hl\n'), ((160, 16, 160, 31), 'hail._get_flags', 'hl._get_flags', ({}, {}), '()', True, 'import hail as hl\n'), ((164, 8, 164, 48), 'hail._set_flags', 'hl._set_flags', (), '', True, 'import hail as hl\n'), ((19, 12, 19, 68), 'hail.init', 'hl.init', (), '', True, 'import hail as hl\n'), ((21, 12, 21, 20), 'hail.utils.java.Env.hc', 'Env.hc', ({}, {}), '()', False, 'from hail.utils.java import Env\n'), ((78, 12, 78, 28), 'hail.float32', 'hl.float32', ({(78, 23, 78, 27): '(3.14)'}, {}), '(3.14)', True, 'import hail as hl\n'), ((79, 12, 79, 24), 'hail.int64', 'hl.int64', ({(79, 21, 79, 23): '(-9)'}, {}), '(-9)', True, 'import hail as hl\n'), ((80, 10, 80, 30), 'hail.null', 'hl.null', ({(80, 18, 80, 29): 'hl.tfloat64'}, {}), '(hl.tfloat64)', True, 'import hail as hl\n'), ((83, 13, 83, 42), 'hail.set', 'hl.set', ({(83, 20, 83, 41): "['foo', 'bar', 'baz']"}, {}), "(['foo', 'bar', 'baz'])", True, 'import hail as hl\n'), ((87, 12, 87, 49), 'hail.locus', 'hl.locus', ({(87, 21, 87, 28): '"""chr22"""', (87, 30, 87, 38): '(33878978)', (87, 40, 87, 48): '"""GRCh38"""'}, {}), "('chr22', 33878978, 'GRCh38')", True, 'import hail as hl\n'), ((92, 10, 92, 23), 'hail.call', 'hl.call', ({(92, 18, 92, 19): '(0)', (92, 21, 92, 22): '(1)'}, {}), '(0, 1)', True, 'import hail as hl\n'), ((93, 11, 93, 28), 'hail.null', 'hl.null', ({(93, 19, 93, 27): 'hl.tcall'}, {}), '(hl.tcall)', True, 'import hail as hl\n'), ((125, 22, 125, 49), 'hail.utils.java.Env.backend', 'hl.utils.java.Env.backend', ({}, {}), '()', True, 'import hail as hl\n'), ((128, 18, 128, 53), 'unittest.SkipTest', 'unittest.SkipTest', ({(128, 36, 128, 52): '"""requires Spark"""'}, {}), "('requires Spark')", False, 'import unittest\n'), ((169, 12, 169, 71), 'hail._set_flags', 'hl._set_flags', (), '', True, 'import hail as hl\n'), ((82, 24, 82, 58), 'hail.tstruct', 'hl.tstruct', (), '', True, 'import hail as hl\n'), ((84, 21, 84, 41), 'hail.tset', 'hl.tset', ({(84, 29, 84, 40): 'hl.tfloat64'}, {}), '(hl.tfloat64)', True, 'import hail as hl\n'), ((86, 19, 86, 47), 'hail.tdict', 'hl.tdict', ({(86, 28, 86, 37): 'hl.tint32', (86, 39, 86, 46): 'hl.tstr'}, {}), '(hl.tint32, hl.tstr)', True, 'import hail as hl\n'), ((88, 19, 88, 38), 'hail.tlocus', 'hl.tlocus', ({(88, 29, 88, 37): '"""GRCh37"""'}, {}), "('GRCh37')", True, 'import hail as hl\n'), ((90, 12, 90, 30), 'hail.locus', 'hl.locus', ({(90, 21, 90, 24): '"""1"""', (90, 26, 90, 29): '(999)'}, {}), "('1', 999)", True, 'import hail as hl\n'), ((91, 12, 91, 31), 'hail.locus', 'hl.locus', ({(91, 21, 91, 24): '"""1"""', (91, 26, 91, 30): '(1001)'}, {}), "('1', 1001)", True, 'import hail as hl\n'), ((81, 28, 81, 46), 'hail.null', 'hl.null', ({(81, 36, 81, 45): 'hl.tint32'}, {}), '(hl.tint32)', True, 'import hail as hl\n'), ((85, 19, 85, 39), 'hail.array', 'hl.array', ({(85, 28, 85, 38): "['a', 'b']"}, {}), "(['a', 'b'])", True, 'import hail as hl\n'), ((94, 20, 94, 46), 'hail.call', 'hl.call', (), '', True, 'import hail as hl\n'), ((94, 55, 94, 71), 'hail.null', 'hl.null', ({(94, 63, 94, 70): 'hl.tstr'}, {}), '(hl.tstr)', True, 'import hail as hl\n'), ((95, 29, 95, 48), 'hail.tlocus', 'hl.tlocus', ({(95, 39, 95, 47): '"""GRCh37"""'}, {}), "('GRCh37')", True, 'import hail as hl\n'), ((96, 11, 96, 30), 'hail.nd.arange', 'hl.nd.arange', ({(96, 24, 96, 25): '(0)', (96, 27, 96, 29): '(10)'}, {}), '(0, 10)', True, 'import hail as hl\n'), ((85, 61, 85, 77), 'hail.null', 'hl.null', ({(85, 69, 85, 76): 'hl.tstr'}, {}), '(hl.tstr)', True, 'import hail as hl\n'), ((104, 12, 104, 51), 'hail.utils.range_table', 'hl.utils.range_table', (), '', True, 'import hail as hl\n'), ((111, 12, 111, 61), 'hail.utils.range_matrix_table', 'hl.utils.range_matrix_table', (), '', True, 'import hail as hl\n')]