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
[((630, 696), 'os.environ.get', 'os.environ.get', (['"""HAIL_TEST_RESOURCES_DIR"""', '"""../src/test/resources"""'], {}), "('HAIL_TEST_RESOURCES_DIR', '../src/test/resources')\n", (644, 696), False, 'import os\n'), ((712, 769), 'os.environ.get', 'os.environ.get', (['"""HAIL_DOCTEST_DATA_DIR"""', '"""hail/docs/data"""'], {}), "('HAIL_DOCTEST_DATA_DIR', 'hail/docs/data')\n", (726, 769), False, 'import os\n'), ((807, 840), 'os.path.join', 'os.path.join', (['_test_dir', 'filename'], {}), '(_test_dir, filename)\n', (819, 840), False, 'import os\n'), ((886, 922), 'os.path.join', 'os.path.join', (['_doctest_dir', 'filename'], {}), '(_doctest_dir, filename)\n', (898, 922), False, 'import os\n'), ((1699, 1706), 'timeit.default_timer', 'timer', ([], {}), '()\n', (1704, 1706), True, 'from timeit import default_timer as timer\n'), ((1729, 1736), 'timeit.default_timer', 'timer', ([], {}), '()\n', (1734, 1736), True, 'from timeit import default_timer as timer\n'), ((2780, 2824), 'hail.struct', 'hl.struct', ([], {}), '(**{(prefix + k): s[k] for k in s})\n', (2789, 2824), True, 'import hail as hl\n'), ((4266, 4276), 'hail.eval', 'hl.eval', (['e'], {}), '(e)\n', (4273, 4276), True, 'import hail as hl\n'), ((341, 386), 'os.environ.get', 'os.environ.get', (['"""HAIL_QUERY_BACKEND"""', '"""spark"""'], {}), "('HAIL_QUERY_BACKEND', 'spark')\n", (355, 386), False, 'import os\n'), ((3994, 4030), 'os.environ.get', 'os.environ.get', (['"""HAIL_QUERY_BACKEND"""'], {}), "('HAIL_QUERY_BACKEND')\n", (4008, 4030), False, 'import os\n'), ((4469, 4484), 'hail.tuple', 'hl.tuple', (['exprs'], {}), '(exprs)\n', (4477, 4484), True, 'import hail as hl\n'), ((4588, 4603), 'hail._get_flags', 'hl._get_flags', ([], {}), '()\n', (4601, 4603), True, 'import hail as hl\n'), ((4703, 4743), 'hail._set_flags', 'hl._set_flags', ([], {'lower': '"""1"""', 'lower_only': '"""1"""'}), "(lower='1', lower_only='1')\n", (4716, 4743), True, 'import hail as hl\n'), ((435, 491), 'hail.init', 'hl.init', ([], {'master': '"""local[1]"""', 'min_block_size': '(0)', 'quiet': '(True)'}), "(master='local[1]', min_block_size=0, quiet=True)\n", (442, 491), True, 'import hail as hl\n'), ((518, 526), 'hail.utils.java.Env.hc', 'Env.hc', ([], {}), '()\n', (524, 526), False, 'from hail.utils.java import Env\n'), ((1887, 1903), 'hail.float32', 'hl.float32', (['(3.14)'], {}), '(3.14)\n', (1897, 1903), True, 'import hail as hl\n'), ((1917, 1929), 'hail.int64', 'hl.int64', (['(-9)'], {}), '(-9)\n', (1925, 1929), True, 'import hail as hl\n'), ((1941, 1961), 'hail.null', 'hl.null', (['hl.tfloat64'], {}), '(hl.tfloat64)\n', (1948, 1961), True, 'import hail as hl\n'), ((2093, 2122), 'hail.set', 'hl.set', (["['foo', 'bar', 'baz']"], {}), "(['foo', 'bar', 'baz'])\n", (2099, 2122), True, 'import hail as hl\n'), ((2323, 2360), 'hail.locus', 'hl.locus', (['"""chr22"""', '(33878978)', '"""GRCh38"""'], {}), "('chr22', 33878978, 'GRCh38')\n", (2331, 2360), True, 'import hail as hl\n'), ((2502, 2515), 'hail.call', 'hl.call', (['(0)', '(1)'], {}), '(0, 1)\n', (2509, 2515), True, 'import hail as hl\n'), ((2528, 2545), 'hail.null', 'hl.null', (['hl.tcall'], {}), '(hl.tcall)\n', (2535, 2545), True, 'import hail as hl\n'), ((3774, 3801), 'hail.utils.java.Env.backend', 'hl.utils.java.Env.backend', ([], {}), '()\n', (3799, 3801), True, 'import hail as hl\n'), ((3891, 3926), 'unittest.SkipTest', 'unittest.SkipTest', (['"""requires Spark"""'], {}), "('requires Spark')\n", (3908, 3926), False, 'import unittest\n'), ((4828, 4887), 'hail._set_flags', 'hl._set_flags', ([], {'lower': 'prev_lower', 'lower_only': 'prev_lower_only'}), '(lower=prev_lower, lower_only=prev_lower_only)\n', (4841, 4887), True, 'import hail as hl\n'), ((2043, 2077), 'hail.tstruct', 'hl.tstruct', ([], {'x': 'hl.tint32', 'y': 'hl.tstr'}), '(x=hl.tint32, y=hl.tstr)\n', (2053, 2077), True, 'import hail as hl\n'), ((2145, 2165), 'hail.tset', 'hl.tset', (['hl.tfloat64'], {}), '(hl.tfloat64)\n', (2152, 2165), True, 'import hail as hl\n'), ((2280, 2308), 'hail.tdict', 'hl.tdict', (['hl.tint32', 'hl.tstr'], {}), '(hl.tint32, hl.tstr)\n', (2288, 2308), True, 'import hail as hl\n'), ((2381, 2400), 'hail.tlocus', 'hl.tlocus', (['"""GRCh37"""'], {}), "('GRCh37')\n", (2390, 2400), True, 'import hail as hl\n'), ((2438, 2456), 'hail.locus', 'hl.locus', (['"""1"""', '(999)'], {}), "('1', 999)\n", (2446, 2456), True, 'import hail as hl\n'), ((2470, 2489), 'hail.locus', 'hl.locus', (['"""1"""', '(1001)'], {}), "('1', 1001)\n", (2478, 2489), True, 'import hail as hl\n'), ((1991, 2009), 'hail.null', 'hl.null', (['hl.tint32'], {}), '(hl.tint32)\n', (1998, 2009), True, 'import hail as hl\n'), ((2187, 2207), 'hail.array', 'hl.array', (["['a', 'b']"], {}), "(['a', 'b'])\n", (2195, 2207), True, 'import hail as hl\n'), ((2567, 2593), 'hail.call', 'hl.call', (['(1)', '(2)'], {'phased': '(True)'}), '(1, 2, phased=True)\n', (2574, 2593), True, 'import hail as hl\n'), ((2602, 2618), 'hail.null', 'hl.null', (['hl.tstr'], {}), '(hl.tstr)\n', (2609, 2618), True, 'import hail as hl\n'), ((2651, 2670), 'hail.tlocus', 'hl.tlocus', (['"""GRCh37"""'], {}), "('GRCh37')\n", (2660, 2670), True, 'import hail as hl\n'), ((2695, 2714), 'hail.nd.arange', 'hl.nd.arange', (['(0)', '(10)'], {}), '(0, 10)\n', (2707, 2714), True, 'import hail as hl\n'), ((2229, 2245), 'hail.null', 'hl.null', (['hl.tstr'], {}), '(hl.tstr)\n', (2236, 2245), True, 'import hail as hl\n'), ((2904, 2943), 'hail.utils.range_table', 'hl.utils.range_table', (['(5)'], {'n_partitions': '(3)'}), '(5, n_partitions=3)\n', (2924, 2943), True, 'import hail as hl\n'), ((3160, 3209), 'hail.utils.range_matrix_table', 'hl.utils.range_matrix_table', (['(3)', '(2)'], {'n_partitions': '(2)'}), '(3, 2, n_partitions=2)\n', (3187, 3209), True, 'import hail as hl\n')]