from __future__ import print_function

from functools import partial
from itertools import permutations
import numba.unittest_support as unittest

import numpy as np

from numba.compiler import compile_isolated, Flags
from numba import jit, types, from_dtype, errors, typeof
from numba.errors import TypingError
from .support import TestCase, MemoryLeakMixin, CompilationCache, tag

enable_pyobj_flags = Flags()
enable_pyobj_flags.set("enable_pyobject")

no_pyobj_flags = Flags()
no_pyobj_flags.set('nrt')

def from_generic(pyfuncs_to_use):
    """Decorator for generic check functions.
        Iterates over 'pyfuncs_to_use', calling 'func' with the iterated
        item as first argument. Example:

        @from_generic(numpy_array_reshape, array_reshape)
        def check_only_shape(pyfunc, arr, shape, expected_shape):
            # Only check Numba result to avoid Numpy bugs
            self.memory_leak_setup()
            got = generic_run(pyfunc, arr, shape)
            self.assertEqual(got.shape, expected_shape)
            self.assertEqual(got.size, arr.size)
            del got
            self.memory_leak_teardown()
    """
    def decorator(func):
        def result(*args, **kwargs):
            return (func(pyfunc, *args, **kwargs) for pyfunc in pyfuncs_to_use)
        return result
    return decorator

def array_reshape(arr, newshape):
    return arr.reshape(newshape)

def numpy_array_reshape(arr, newshape):
    return np.reshape(arr, newshape)

def flatten_array(a):
    return a.flatten()

def ravel_array(a):
    return a.ravel()

def ravel_array_size(a):
    return a.ravel().size

def numpy_ravel_array(a):
    return np.ravel(a)

def transpose_array(a):
    return a.transpose()

def numpy_transpose_array(a):
    return np.transpose(a)

def numpy_transpose_array_axes_kwarg(arr, axes):
    return np.transpose(arr, axes=axes)

def array_transpose_axes(arr, axes):
    return arr.transpose(axes)

def squeeze_array(a):
    return a.squeeze()

def expand_dims(a, axis):
    return np.expand_dims(a, axis)

def atleast_1d(*args):
    return np.atleast_1d(*args)

def atleast_2d(*args):
    return np.atleast_2d(*args)

def atleast_3d(*args):
    return np.atleast_3d(*args)

def as_strided1(a):
    # as_strided() with implicit shape
    strides = (a.strides[0] // 2,) + a.strides[1:]
    return np.lib.stride_tricks.as_strided(a, strides=strides)

def as_strided2(a):
    # Rolling window example as in https://github.com/numba/numba/issues/1884
    window = 3
    shape = a.shape[:-1] + (a.shape[-1] - window + 1, window)
    strides = a.strides + (a.strides[-1],)
    return np.lib.stride_tricks.as_strided(a, shape=shape, strides=strides)

def add_axis2(a):
    return a[np.newaxis, :]

def bad_index(arr, arr2d):
    x = arr.x,
    y = arr.y
    # note that `x` is a tuple, which causes a new axis to be created.
    arr2d[x, y] = 1.0

def bad_float_index(arr):
    # 2D index required for this function because 1D index
    # fails typing
    return arr[1, 2.0]

def numpy_fill_diagonal(arr, val, wrap=False):
    return np.fill_diagonal(arr, val, wrap)

class TestArrayManipulation(MemoryLeakMixin, TestCase):
    """
    Check shape-changing operations on arrays.
    """

    def setUp(self):
        super(TestArrayManipulation, self).setUp()
        self.ccache = CompilationCache()

    @tag('important')
    def test_array_reshape(self):
        pyfuncs_to_use = [array_reshape, numpy_array_reshape]

        def generic_run(pyfunc, arr, shape):
            cres = compile_isolated(pyfunc, (typeof(arr), typeof(shape)))
            return cres.entry_point(arr, shape)

        @from_generic(pyfuncs_to_use)
        def check(pyfunc, arr, shape):
            expected = pyfunc(arr, shape)
            self.memory_leak_setup()
            got = generic_run(pyfunc, arr, shape)
            self.assertPreciseEqual(got, expected)
            del got
            self.memory_leak_teardown()

        @from_generic(pyfuncs_to_use)
        def check_only_shape(pyfunc, arr, shape, expected_shape):
            # Only check Numba result to avoid Numpy bugs
            self.memory_leak_setup()
            got = generic_run(pyfunc, arr, shape)
            self.assertEqual(got.shape, expected_shape)
            self.assertEqual(got.size, arr.size)
            del got
            self.memory_leak_teardown()

        @from_generic(pyfuncs_to_use)
        def check_err_shape(pyfunc, arr, shape):
            with self.assertRaises(NotImplementedError) as raises:
                generic_run(pyfunc, arr, shape)
            self.assertEqual(str(raises.exception),
                             "incompatible shape for array")

        @from_generic(pyfuncs_to_use)
        def check_err_size(pyfunc, arr, shape):
            with self.assertRaises(ValueError) as raises:
                generic_run(pyfunc, arr, shape)
            self.assertEqual(str(raises.exception),
                             "total size of new array must be unchanged")

        @from_generic(pyfuncs_to_use)
        def check_err_multiple_negative(pyfunc, arr, shape):
            with self.assertRaises(ValueError) as raises:
                generic_run(pyfunc, arr, shape)
            self.assertEqual(str(raises.exception),
                             "multiple negative shape values")

        # C-contiguous
        arr = np.arange(24)
        check(arr, (24,))
        check(arr, (4, 6))
        check(arr, (8, 3))
        check(arr, (8, 1, 3))
        check(arr, (1, 8, 1, 1, 3, 1))
        arr = np.arange(24).reshape((2, 3, 4))
        check(arr, (24,))
        check(arr, (4, 6))
        check(arr, (8, 3))
        check(arr, (8, 1, 3))
        check(arr, (1, 8, 1, 1, 3, 1))
        check_err_size(arr, ())
        check_err_size(arr, (25,))
        check_err_size(arr, (8, 4))
        arr = np.arange(24).reshape((1, 8, 1, 1, 3, 1))
        check(arr, (24,))
        check(arr, (4, 6))
        check(arr, (8, 3))
        check(arr, (8, 1, 3))

        # F-contiguous
        arr = np.arange(24).reshape((2, 3, 4)).T
        check(arr, (4, 3, 2))
        check(arr, (1, 4, 1, 3, 1, 2, 1))
        check_err_shape(arr, (2, 3, 4))
        check_err_shape(arr, (6, 4))
        check_err_shape(arr, (2, 12))

        # Test negative shape value
        arr = np.arange(25).reshape(5,5)
        check(arr, -1)
        check(arr, (-1,))
        check(arr, (-1, 5))
        check(arr, (5, -1, 5))
        check(arr, (5, 5, -1))
        check_err_size(arr, (-1, 4))
        check_err_multiple_negative(arr, (-1, -2, 5, 5))
        check_err_multiple_negative(arr, (5, 5, -1, -1))

        # 0-sized arrays
        def check_empty(arr):
            check(arr, 0)
            check(arr, (0,))
            check(arr, (1, 0, 2))
            check(arr, (0, 55, 1, 0, 2))
            # -1 is buggy in Numpy with 0-sized arrays
            check_only_shape(arr, -1, (0,))
            check_only_shape(arr, (-1,), (0,))
            check_only_shape(arr, (0, -1), (0, 0))
            check_only_shape(arr, (4, -1), (4, 0))
            check_only_shape(arr, (-1, 0, 4), (0, 0, 4))
            check_err_size(arr, ())
            check_err_size(arr, 1)
            check_err_size(arr, (1, 2))

        arr = np.array([])
        check_empty(arr)
        check_empty(arr.reshape((3, 2, 0)))

        # Exceptions leak references
        self.disable_leak_check()

    def test_array_transpose_axes(self):
        pyfuncs_to_use = [numpy_transpose_array_axes_kwarg,
                          array_transpose_axes]

        def run(pyfunc, arr, axes):
            cres = self.ccache.compile(pyfunc, (typeof(arr), typeof(axes)))
            return cres.entry_point(arr, axes)

        @from_generic(pyfuncs_to_use)
        def check(pyfunc, arr, axes):
            expected = pyfunc(arr, axes)
            got = run(pyfunc, arr, axes)
            self.assertPreciseEqual(got, expected)
            self.assertEqual(got.flags.f_contiguous,
                             expected.flags.f_contiguous)
            self.assertEqual(got.flags.c_contiguous,
                             expected.flags.c_contiguous)

        @from_generic(pyfuncs_to_use)
        def check_err_axis_repeated(pyfunc, arr, axes):
            with self.assertRaises(ValueError) as raises:
                run(pyfunc, arr, axes)
            self.assertEqual(str(raises.exception),
                             "repeated axis in transpose")

        @from_generic(pyfuncs_to_use)
        def check_err_axis_oob(pyfunc, arr, axes):
            with self.assertRaises(ValueError) as raises:
                run(pyfunc, arr, axes)
            self.assertEqual(str(raises.exception),
                             "axis is out of bounds for array of given dimension")

        @from_generic(pyfuncs_to_use)
        def check_err_invalid_args(pyfunc, arr, axes):
            with self.assertRaises((TypeError, TypingError)):
                run(pyfunc, arr, axes)

        arrs = [np.arange(24),
                np.arange(24).reshape(4, 6),
                np.arange(24).reshape(2, 3, 4),
                np.arange(24).reshape(1, 2, 3, 4),
                np.arange(64).reshape(8, 4, 2)[::3,::2,:]]

        for i in range(len(arrs)):
            # First check `None`, the default, which is to reverse dims
            check(arrs[i], None)
            # Check supplied axis permutations
            for axes in permutations(tuple(range(arrs[i].ndim))):
                ndim = len(axes)
                neg_axes = tuple([x - ndim for x in axes])
                check(arrs[i], axes)
                check(arrs[i], neg_axes)

        # Exceptions leak references
        self.disable_leak_check()

        check_err_invalid_args(arrs[1], "foo")
        check_err_invalid_args(arrs[1], ("foo",))
        check_err_invalid_args(arrs[1], 5.3)
        check_err_invalid_args(arrs[2], (1.2, 5))

        check_err_axis_repeated(arrs[1], (0, 0))
        check_err_axis_repeated(arrs[2], (2, 0, 0))
        check_err_axis_repeated(arrs[3], (3, 2, 1, 1))

        check_err_axis_oob(arrs[0], (1,))
        check_err_axis_oob(arrs[0], (-2,))
        check_err_axis_oob(arrs[1], (0, 2))
        check_err_axis_oob(arrs[1], (-3, 2))
        check_err_axis_oob(arrs[1], (0, -3))
        check_err_axis_oob(arrs[2], (3, 1, 2))
        check_err_axis_oob(arrs[2], (-4, 1, 2))
        check_err_axis_oob(arrs[3], (3, 1, 2, 5))
        check_err_axis_oob(arrs[3], (3, 1, 2, -5))

    @tag('important')
    def test_expand_dims(self):
        pyfunc = expand_dims

        def run(arr, axis):
            cres = self.ccache.compile(pyfunc, (typeof(arr), typeof(axis)))
            return cres.entry_point(arr, axis)

        def check(arr, axis):
            expected = pyfunc(arr, axis)
            self.memory_leak_setup()
            got = run(arr, axis)
            self.assertPreciseEqual(got, expected)
            del got
            self.memory_leak_teardown()

        def check_all_axes(arr):
            for axis in range(-arr.ndim - 1, arr.ndim + 1):
                check(arr, axis)

        # 1d
        arr = np.arange(5)
        check_all_axes(arr)
        # 3d (C, F, A)
        arr = np.arange(24).reshape((2, 3, 4))
        check_all_axes(arr)
        check_all_axes(arr.T)
        check_all_axes(arr[::-1])
        # 0d
        arr = np.array(42)
        check_all_axes(arr)

    def check_atleast_nd(self, pyfunc, cfunc):
        def check_result(got, expected):
            # We would like to check the result has the same contiguity,
            # but we can't rely on the "flags" attribute when there are
            # 1-sized dimensions.
            self.assertStridesEqual(got, expected)
            self.assertPreciseEqual(got.flatten(), expected.flatten())

        def check_single(arg):
            check_result(cfunc(arg), pyfunc(arg))

        def check_tuple(*args):
            expected_tuple = pyfunc(*args)
            got_tuple = cfunc(*args)
            self.assertEqual(len(got_tuple), len(expected_tuple))
            for got, expected in zip(got_tuple, expected_tuple):
                check_result(got, expected)

        # 0d
        a1 = np.array(42)
        a2 = np.array(5j)
        check_single(a1)
        check_tuple(a1, a2)
        # 1d
        b1 = np.arange(5)
        b2 = np.arange(6) + 1j
        b3 = b1[::-1]
        check_single(b1)
        check_tuple(b1, b2, b3)
        # 2d
        c1 = np.arange(6).reshape((2, 3))
        c2 = c1.T
        c3 = c1[::-1]
        check_single(c1)
        check_tuple(c1, c2, c3)
        # 3d
        d1 = np.arange(24).reshape((2, 3, 4))
        d2 = d1.T
        d3 = d1[::-1]
        check_single(d1)
        check_tuple(d1, d2, d3)
        # 4d
        e = np.arange(16).reshape((2, 2, 2, 2))
        check_single(e)
        # mixed dimensions
        check_tuple(a1, b2, c3, d2)

    def test_atleast_1d(self):
        pyfunc = atleast_1d
        cfunc = jit(nopython=True)(pyfunc)
        self.check_atleast_nd(pyfunc, cfunc)

    def test_atleast_2d(self):
        pyfunc = atleast_2d
        cfunc = jit(nopython=True)(pyfunc)
        self.check_atleast_nd(pyfunc, cfunc)

    def test_atleast_3d(self):
        pyfunc = atleast_3d
        cfunc = jit(nopython=True)(pyfunc)
        self.check_atleast_nd(pyfunc, cfunc)

    def check_as_strided(self, pyfunc):
        def run(arr):
            cres = self.ccache.compile(pyfunc, (typeof(arr),))
            return cres.entry_point(arr)
        def check(arr):
            expected = pyfunc(arr)
            got = run(arr)
            self.assertPreciseEqual(got, expected)

        arr = np.arange(24)
        check(arr)
        check(arr.reshape((6, 4)))
        check(arr.reshape((4, 1, 6)))

    def test_as_strided(self):
        self.check_as_strided(as_strided1)
        self.check_as_strided(as_strided2)

    def test_flatten_array(self, flags=enable_pyobj_flags, layout='C'):
        a = np.arange(9).reshape(3, 3)
        if layout == 'F':
            a = a.T

        pyfunc = flatten_array
        arraytype1 = typeof(a)
        if layout == 'A':
            # Force A layout
            arraytype1 = arraytype1.copy(layout='A')

        self.assertEqual(arraytype1.layout, layout)
        cr = compile_isolated(pyfunc, (arraytype1,), flags=flags)
        cfunc = cr.entry_point

        expected = pyfunc(a)
        got = cfunc(a)
        np.testing.assert_equal(expected, got)

    def test_flatten_array_npm(self):
        self.test_flatten_array(flags=no_pyobj_flags)
        self.test_flatten_array(flags=no_pyobj_flags, layout='F')
        self.test_flatten_array(flags=no_pyobj_flags, layout='A')

    def test_ravel_array(self, flags=enable_pyobj_flags):
        def generic_check(pyfunc, a, assume_layout):
            # compile
            arraytype1 = typeof(a)
            self.assertEqual(arraytype1.layout, assume_layout)
            cr = compile_isolated(pyfunc, (arraytype1,), flags=flags)
            cfunc = cr.entry_point

            expected = pyfunc(a)
            got = cfunc(a)
            # Check result matches
            np.testing.assert_equal(expected, got)
            # Check copying behavior
            py_copied = (a.ctypes.data != expected.ctypes.data)
            nb_copied = (a.ctypes.data != got.ctypes.data)
            self.assertEqual(py_copied, assume_layout != 'C')
            self.assertEqual(py_copied, nb_copied)

        check_method = partial(generic_check, ravel_array)
        check_function = partial(generic_check, numpy_ravel_array)

        def check(*args, **kwargs):
            check_method(*args, **kwargs)
            check_function(*args, **kwargs)

        # Check 2D
        check(np.arange(9).reshape(3, 3), assume_layout='C')
        check(np.arange(9).reshape(3, 3, order='F'), assume_layout='F')
        check(np.arange(18).reshape(3, 3, 2)[:, :, 0], assume_layout='A')

        # Check 3D
        check(np.arange(18).reshape(2, 3, 3), assume_layout='C')
        check(np.arange(18).reshape(2, 3, 3, order='F'), assume_layout='F')
        check(np.arange(36).reshape(2, 3, 3, 2)[:, :, :, 0], assume_layout='A')

    def test_ravel_array_size(self, flags=enable_pyobj_flags):
        a = np.arange(9).reshape(3, 3)

        pyfunc = ravel_array_size
        arraytype1 = typeof(a)
        cr = compile_isolated(pyfunc, (arraytype1,), flags=flags)
        cfunc = cr.entry_point

        expected = pyfunc(a)
        got = cfunc(a)
        np.testing.assert_equal(expected, got)

    def test_ravel_array_npm(self):
        self.test_ravel_array(flags=no_pyobj_flags)

    def test_ravel_array_size_npm(self):
        self.test_ravel_array_size(flags=no_pyobj_flags)

    def test_transpose_array(self, flags=enable_pyobj_flags):
        @from_generic([transpose_array, numpy_transpose_array])
        def check(pyfunc):
            a = np.arange(9).reshape(3, 3)

            arraytype1 = typeof(a)
            cr = compile_isolated(pyfunc, (arraytype1,), flags=flags)
            cfunc = cr.entry_point

            expected = pyfunc(a)
            got = cfunc(a)
            np.testing.assert_equal(expected, got)

        check()

    def test_transpose_array_npm(self):
        self.test_transpose_array(flags=no_pyobj_flags)

    def test_squeeze_array(self, flags=enable_pyobj_flags):
        a = np.arange(2 * 1 * 3 * 1 * 4).reshape(2, 1, 3, 1, 4)

        pyfunc = squeeze_array
        arraytype1 = typeof(a)
        cr = compile_isolated(pyfunc, (arraytype1,), flags=flags)
        cfunc = cr.entry_point

        expected = pyfunc(a)
        got = cfunc(a)
        np.testing.assert_equal(expected, got)

    def test_squeeze_array_npm(self):
        with self.assertRaises(errors.TypingError) as raises:
            self.test_squeeze_array(flags=no_pyobj_flags)

        self.assertIn("squeeze", str(raises.exception))

    def test_add_axis2(self, flags=enable_pyobj_flags):
        a = np.arange(9).reshape(3, 3)

        pyfunc = add_axis2
        arraytype1 = typeof(a)
        cr = compile_isolated(pyfunc, (arraytype1,), flags=flags)
        cfunc = cr.entry_point

        expected = pyfunc(a)
        got = cfunc(a)
        np.testing.assert_equal(expected, got)

    def test_add_axis2_npm(self):
        with self.assertTypingError() as raises:
            self.test_add_axis2(flags=no_pyobj_flags)
        self.assertIn("unsupported array index type none in",
                      str(raises.exception))

    def test_bad_index_npm(self):
        with self.assertTypingError() as raises:
            arraytype1 = from_dtype(np.dtype([('x', np.int32),
                                              ('y', np.int32)]))
            arraytype2 = types.Array(types.int32, 2, 'C')
            compile_isolated(bad_index, (arraytype1, arraytype2),
                             flags=no_pyobj_flags)
        self.assertIn('unsupported array index type', str(raises.exception))

    def test_bad_float_index_npm(self):
        with self.assertTypingError() as raises:
            compile_isolated(bad_float_index,
                             (types.Array(types.float64, 2, 'C'),))
        self.assertIn('unsupported array index type float64',
                      str(raises.exception))

    def test_fill_diagonal_basic(self):
        pyfunc = numpy_fill_diagonal
        cfunc = jit(nopython=True)(pyfunc)

        def _shape_variations(n):
            # square
            yield (n, n)
            # tall and thin
            yield (2 * n, n)
            # short and fat
            yield (n, 2 * n)
            # a bit taller than wide; odd numbers of rows and cols
            yield ((2 * n + 1), (2 * n - 1))
            # 4d, all dimensions same
            yield (n, n, n, n)
            # weird edge case
            yield (1, 1, 1)

        def _val_variations():
            yield 1
            yield 3.142
            yield np.nan
            yield -np.inf
            yield True
            yield np.arange(4)
            yield (4,)
            yield [8, 9]
            yield np.arange(54).reshape(9, 3, 2, 1)  # contiguous C
            yield np.asfortranarray(np.arange(9).reshape(3, 3))  # contiguous F
            yield np.arange(9).reshape(3, 3)[::-1]  # non-contiguous

        # contiguous arrays
        def _multi_dimensional_array_variations(n):
            for shape in _shape_variations(n):
                yield np.zeros(shape, dtype=np.float64)
                yield np.asfortranarray(np.ones(shape, dtype=np.float64))

        # non-contiguous arrays
        def _multi_dimensional_array_variations_strided(n):
            for shape in _shape_variations(n):
                tmp = np.zeros(tuple([x * 2 for x in shape]), dtype=np.float64)
                slicer = tuple(slice(0, x * 2, 2) for x in shape)
                yield tmp[slicer]

        def _check_fill_diagonal(arr, val):
            for wrap in None, True, False:
                a = arr.copy()
                b = arr.copy()

                if wrap is None:
                    params = {}
                else:
                    params = {'wrap': wrap}

                pyfunc(a, val, **params)
                cfunc(b, val, **params)
                self.assertPreciseEqual(a, b)

        for arr in _multi_dimensional_array_variations(3):
            for val in _val_variations():
                _check_fill_diagonal(arr, val)

        for arr in _multi_dimensional_array_variations_strided(3):
            for val in _val_variations():
                _check_fill_diagonal(arr, val)

        # non-numeric input arrays
        arr = np.array([True] * 9).reshape(3, 3)
        _check_fill_diagonal(arr, False)
        _check_fill_diagonal(arr, [False, True, False])
        _check_fill_diagonal(arr, np.array([True, False, True]))

    def test_fill_diagonal_exception_cases(self):
        pyfunc = numpy_fill_diagonal
        cfunc = jit(nopython=True)(pyfunc)
        val = 1

        # Exceptions leak references
        self.disable_leak_check()

        # first argument unsupported number of dimensions
        for a in np.array([]), np.ones(5):
            with self.assertRaises(TypingError) as raises:
                cfunc(a, val)
            assert "The first argument must be at least 2-D" in str(raises.exception)

        # multi-dimensional input where dimensions are not all equal
        with self.assertRaises(ValueError) as raises:
            a = np.zeros((3, 3, 4))
            cfunc(a, val)
            self.assertEqual("All dimensions of input must be of equal length", str(raises.exception))

        # cases where val has incompatible type / value
        def _assert_raises(arr, val):
            with self.assertRaises(ValueError) as raises:
                cfunc(arr, val)
            self.assertEqual("Unable to safely conform val to a.dtype", str(raises.exception))

        arr = np.zeros((3, 3), dtype=np.int32)
        val = np.nan
        _assert_raises(arr, val)

        val = [3.3, np.inf]
        _assert_raises(arr, val)

        val = np.array([1, 2, 1e10], dtype=np.int64)
        _assert_raises(arr, val)

        arr = np.zeros((3, 3), dtype=np.float32)
        val = [1.4, 2.6, -1e100]
        _assert_raises(arr, val)

        val = 1.1e100
        _assert_raises(arr, val)

        val = np.array([-1e100])
        _assert_raises(arr, val)

if __name__ == '__main__':
    unittest.main()


import logging
import traceback
import os
from typing import Dict
from math import ceil
from itertools import repeat, takewhile

import multiprocessing
import numpy as np

from playground.main import get_args
from playground import ColorizedLogger, Configuration, timeit

# Create loggers with different colors to use in each problem
main_logger = ColorizedLogger('Main', 'yellow')
p1_logger = ColorizedLogger('Problem1', 'blue')
p2_logger = ColorizedLogger('Problem2', 'green')
p3_logger = ColorizedLogger('Problem3', 'magenta')
extra_ch_logger = ColorizedLogger('ExtraMain', 'yellow')
extra_sub_ch_logger = ColorizedLogger('ExtraSub', 'cyan')

# Global Vars (For the Extra Challenges)
# lock: multiprocessing.Lock
# multi_list: List = []

def my_pid(x: int) -> None:
    """ Problem 1 function to be called using pool.map

    Parameters:
        x: the id of the worker
    """

    pid = os.getpid()
    p1_logger.info(f"Hi, I’m worker {x} (with {pid})")

def problem1(conf: Dict) -> None:
    """ Problem 1 solution

    Parameters:
         conf: The config loaded from the yml file
            Example:
                properties:
                  pool_size: 4
                  chunk_size: 1
                  x_min: 0
                  x_max: 9
                conf_type: required
    """

    p1_logger.info("Starting Problem 1..")
    conf_props = conf['properties']
    # Generate iterable from `x_min` to `x_max`
    xs = range(conf_props["x_min"], conf_props["x_max"] + 1)
    p1_logger.info(f"Pool Size: {conf_props['pool_size']}: Will call `my_pid` for x in {tuple(xs)}")
    # Call my_pid() using pool.map()
    with multiprocessing.Pool(processes=conf_props['pool_size']) as pool:
        # https://docs.python.org/3/library/multiprocessing.html#multiprocessing.pool.Pool.map
        pool.map(func=my_pid,
                 iterable=xs,
                 chunksize=conf_props['chunk_size'])

# Call the timeit wrapper from main.py and pass it a custom string to print
# It already supports string formatting for func_name`, `args`, and `duration`
# To reference the first positional argument of the function I wrap, I can use {0}
@timeit(custom_print='N={0}: Calculation of pi took: {duration:2.5f} sec(s) total')
def py_pi(N: int, real_pi: float) -> None:
    """ Problem 2 function to be called using pool.starmap

    Parameters:
        N: Number of terms to be used to calculate pi
        real_pi: Correct value of pi
    """

    # Construct the equation
    first_term = (4 / N)
    list_to_be_summed = [1 / (1 + ((i - 0.5) / N) ** 2) for i in range(1, N + 1)]
    second_term = sum(list_to_be_summed)
    calced_pi = first_term * second_term
    # Calculate the absolute difference between the calculated pi and real value of pi
    pi_diff = abs(real_pi - calced_pi)
    p2_logger.info(f"N={N}: Pi({N}) = {calced_pi} (Real is {real_pi}, difference is {pi_diff})")

def problem2(conf: Dict) -> None:
    """ Problem 2 solution

    Parameters:
         conf: The config loaded from the yml file
            Example:
                properties:
                  pool_size: 4
                  chunk_size: 1
                  num_terms_min: 10
                  num_terms_step: 5
                  num_terms_max: 3906250
                conf_type: required

    """

    p2_logger.info("Starting Problem 2..")
    real_pi = np.pi
    conf_props = conf['properties']
    # Generate iterable with number of terms to be used
    # Lambda function that multiplies each element by 5 in the power of the element's index
    # e.g. [(10, 0), (10, 1), (10, 2)] -> [10*5^0, 10*5^1, 10*5^2]
    # the first ind_and_num would be (10, 0)
    multiplied_series = lambda ind_and_num: (ind_and_num[1] *
                                             (conf_props["num_terms_step"] ** (ind_and_num[0])))
    # Map an infinite enumerated iterable of (index, 10) to the lambda function
    # Stop when the series exceeds 3906250
    num_terms_range = takewhile(lambda series_el: series_el <= conf_props["num_terms_max"],
                                map(multiplied_series, enumerate(repeat(conf_props["num_terms_min"]))))
    num_terms_range = tuple(num_terms_range)  # I only do this because I want to print them
    # Create the iterable of arguments to be passed to py_pi using pool.starmap()
    # repeat() propagates the same `real_pi` value as many times as necessary
    # to zip it with num_terms_range
    args = zip(num_terms_range, repeat(real_pi))
    p2_logger.info(f"Will call `py_pi` for N in {tuple(num_terms_range)}")
    # Call py_pi() using pool.starmap() (starmap accepts iterable with multiple arguments)
    with multiprocessing.Pool(processes=conf_props['pool_size']) as pool:
        # https://docs.python.org/3/library/multiprocessing.html#multiprocessing.pool.Pool.map
        pool.starmap(func=py_pi,
                     iterable=args,
                     chunksize=conf_props['chunk_size'])

def py_pi_better(N: int, i_start: int, i_stop: int) -> float:
    """ Problem 3 function to be called using pool.starmap

    Parameters:
        N: Number of terms to be used to calculate pi
        i_start: The starting index of the sum used in pi's approximation
        i_stop: The ending index of the sum used in pi's approximation
    Returns:
        partial_calced_pi: The part of pi's approximation calculated
    """

    # Construct the equation
    first_term = (4 / N)
    list_to_be_summed = [1 / (1 + ((i - 0.5) / N) ** 2) for i in range(i_start, i_stop + 1)]
    second_term = sum(list_to_be_summed)
    cacled_pi = first_term * second_term
    return cacled_pi

def problem3(conf: Dict) -> None:
    """ Problem 3 solution

    Parameters:
         conf: The config loaded from the yml file
            Example:
                properties:
                  pool_size: 4
                  chunk_size: 1
                  num_terms:
                    - 100
                    - 500
                    - 1000
                    - 2000
                    - 10000
                    - 50000
                conf_type: required
    """

    p3_logger.info("Starting Problem 3..")
    conf_props = conf['properties']
    p3_logger.info(f"Will call `py_pi_better` for N in {tuple(conf_props['num_terms'])}")
    # Run the pi calculation once for each number of terms requested in the yml
    for num_term in conf_props["num_terms"]:
        # Split the work of `num_terms` into `pool_size` number of parts
        step = ceil(num_term / conf_props["pool_size"])
        i_start = range(1, num_term, step)
        i_stop = map(lambda el: el + step - 1 if el + step - 1 <= num_term else num_term, i_start)
        # Zip N with the i_start and i_stop iterables. Propagate the same N value using repeat
        args = zip(repeat(num_term), i_start, i_stop)
        # Call py_pi_better() using pool.starmap() (starmap accepts iterable with multiple arguments)
        with multiprocessing.Pool(processes=conf_props['pool_size']) as pool:
            # timeit can be used as a context manager too. Pass it a custom string and count the
            # total time to calculate pi
            custom_string = f'N={num_term}: Parallel calculation of pi took: ' + \
                            '{duration:2.5f} sec(s) total'
            with timeit(custom_print=custom_string):
                # https://docs.python.org/3/library/multiprocessing.html#multiprocessing.pool.Pool.map
                pi_chunks = pool.starmap(func=py_pi_better,
                                         iterable=args,
                                         chunksize=conf_props['chunk_size'])
                calced_pi = sum(pi_chunks)
        real_pi = np.pi
        pi_diff = abs(real_pi - calced_pi)
        p3_logger.info(f"N={num_term}: Pi({num_term}) = {calced_pi} (Real is {real_pi}, "
                       f"difference is {pi_diff})")

def extra_1(conf_props: Dict):
    """ Extra Challenge 1 solution

    Parameters:
         conf_props: The config loaded from the yml file
            Example:
                pool_sizes:
                  - 2
                  - 4
                  - 8
                  - 16
                  - 32
                chunk_size: 1
                num_term: 8000000
    """

    num_term = int(conf_props["num_term"])
    extra_ch_logger.info("Will call `py_pi_better` for pool_size in "
                         f"{tuple(conf_props['pool_sizes'])}")
    for pool_size in conf_props["pool_sizes"]:
        pool_size = int(pool_size)
        # Slightly modified code from problem 3
        extra_sub_ch_logger.info(f"Pool Size={pool_size}: Calling workers for N={num_term}")
        step = ceil(num_term / pool_size)
        i_start = range(1, num_term, step)
        i_stop = map(lambda el: el + step - 1 if el + step - 1 <= num_term else num_term, i_start)
        args = zip(repeat(num_term), i_start, i_stop)
        with multiprocessing.Pool(processes=pool_size) as pool:
            custom_string = f'Pool Size={pool_size}: Calculation of pi for N={num_term} took: ' + \
                            '{duration:2.5f} sec(s) total'
            with timeit(custom_print=custom_string):
                pi_chunks = pool.starmap(func=py_pi_better,
                                         iterable=args,
                                         chunksize=conf_props['chunk_size'])
                calced_pi = sum(pi_chunks)
        real_pi = np.pi
        pi_diff = abs(real_pi - calced_pi)
        extra_sub_ch_logger.info(f"Pool Size={pool_size}: Pi({num_term}) = {calced_pi} "
                                 f"(Real is {real_pi}, difference is {pi_diff})")

def py_pi_better_with_queue(N: int, i_start: int, i_stop: int,
                            m_queue: multiprocessing.Queue) -> None:
    """ Extra Challenge 2 function to be called using pool.call_async

    Parameters:
        N: Number of terms to be used to calculate pi
        i_start: The starting index of the sum used in pi's approximation
        i_stop: The ending index of the sum used in pi's approximation
        m_queue: A queue in which to append the results
    """

    # Construct the equation
    first_term = (4 / N)
    list_to_be_summed = [1 / (1 + ((i - 0.5) / N) ** 2) for i in range(i_start, i_stop + 1)]
    second_term = sum(list_to_be_summed)
    cacled_pi = first_term * second_term
    m_queue.put(cacled_pi)

def extra_2(conf_props: Dict):
    """ Extra Challenge 2 solution

    Parameters:
         conf_props: The config loaded from the yml file
            Example:
                pool_sizes:
                  - 2
                  - 4
                  - 8
                  - 16
                  - 32
                num_term: 8000000
    """

    num_term = int(conf_props["num_term"])

    extra_ch_logger.info("Will call `py_pi_better_with_queue` for pool_size in "
                         f"{tuple(conf_props['pool_sizes'])}")
    for pool_size in conf_props["pool_sizes"]:
        pool_size = int(pool_size)
        # Create a multiprocessing manager and a Queue
        multi_manager = multiprocessing.Manager()
        multi_queue = multi_manager.Queue()
        # Slightly modified code from problem 3
        extra_sub_ch_logger.info(f"Pool Size={pool_size}: Calling Async workers for N={num_term}")
        step = ceil(num_term / pool_size)
        i_start = range(1, num_term, step)
        i_stop = map(lambda el: el + step - 1 if el + step - 1 <= num_term else num_term, i_start)
        args = zip(repeat(num_term), i_start, i_stop, repeat(multi_queue))
        # Setup manually a Pool
        custom_string = f'Pool Size={pool_size}: Calculation of pi for N={num_term} and  ' + \
                        'took: {duration:2.5f} sec(s) total'
        with timeit(custom_print=custom_string):
            pool = multiprocessing.Pool(processes=pool_size)
            pool.starmap_async(func=py_pi_better_with_queue, iterable=args)
            pi_chunks = []
            while len(pi_chunks) < pool_size:
                pi_chunks.append(multi_queue.get())
            calced_pi = sum(pi_chunks)
        real_pi = np.pi
        pi_diff = abs(real_pi - calced_pi)
        extra_sub_ch_logger.info(f"Pool Size={pool_size}: Pi({num_term}) = {calced_pi}"
                                 f"(Real is {real_pi}, difference is {pi_diff})")
        pool.close()
        pool.join()

def init(_lock, _multi_list):
    global lock
    global multi_list
    lock = _lock
    multi_list = _multi_list

def py_pi_better_with_list(N: int, i_start: int, i_stop: int) -> None:
    """ Extra Challenge 2 function to be called using pool.call_async

    Parameters:
        N: Number of terms to be used to calculate pi
        i_start: The starting index of the sum used in pi's approximation
        i_stop: The ending index of the sum used in pi's approximation
    """

    # Construct the equation
    first_term = (4 / N)
    list_to_be_summed = [1 / (1 + ((i - 0.5) / N) ** 2) for i in range(i_start, i_stop + 1)]
    second_term = sum(list_to_be_summed)
    cacled_pi = first_term * second_term
    lock.acquire()
    multi_list.append(cacled_pi)
    lock.release()

def extra_3(conf_props: Dict):
    """ Extra Challenge 3 solution

    Parameters:
         conf_props: The config loaded from the yml file
    """

    num_term = int(conf_props["num_term"])

    extra_ch_logger.info("Will call `py_pi_better_with_list` for pool_size in "
                         f"{tuple(conf_props['pool_sizes'])}")
    for pool_size in conf_props["pool_sizes"]:
        pool_size = int(pool_size)
        # Create a multiprocessing manager and a Queue
        manager = multiprocessing.Manager()
        multi_list = manager.list()
        _lock = multiprocessing.Lock()
        # Slightly modified code from problem 3
        extra_sub_ch_logger.info(f"Pool Size={pool_size}: Calling Async workers for N={num_term}")
        step = ceil(num_term / pool_size)
        i_start = range(1, num_term, step)
        i_stop = map(lambda el: el + step - 1 if el + step - 1 <= num_term else num_term, i_start)
        args = zip(repeat(num_term), i_start, i_stop)
        # Setup manually a Pool
        custom_string = f'Pool Size={pool_size}: Calculation of pi for N={num_term} and  ' + \
                        'took: {duration:2.5f} sec(s) total'
        with timeit(custom_print=custom_string):
            pool = multiprocessing.Pool(initializer=init, initargs=(_lock, multi_list),
                                        processes=pool_size)
            pool.starmap_async(func=py_pi_better_with_queue, iterable=args)
            pi_chunks = []
            while len(multi_list) < pool_size:
                pi_chunks = multi_list
            calced_pi = sum(pi_chunks)
        real_pi = np.pi
        pi_diff = abs(real_pi - calced_pi)
        extra_sub_ch_logger.info(f"Pool Size={pool_size}: Pi({num_term}) = {calced_pi}"
                                 f"(Real is {real_pi}, difference is {pi_diff})")
        pool.close()
        pool.join()

def extra_challenges(conf: Dict) -> None:
    """ Extra Challenges solution

    Parameters:
         conf: The config loaded from the yml file
            Example: Examples of each sub-config in each problem's function
    """

    extra_ch_logger.info("Starting the Extra Challenges..")
    # Run all the extra challenge that have the property enabled: True
    if conf['properties']["ch1"]["enabled"]:
        extra_ch_logger.info("1: Test performance for different pool sizes.")
        extra_1(conf['properties']["ch1"])
    if conf['properties']["ch2"]["enabled"]:
        extra_ch_logger.info("2: Collect results from workers using multiprocessing list "
                             "(using starmap_async).")
        extra_2(conf['properties']["ch2"])
    if conf['properties']["ch3"]["enabled"]:
        extra_ch_logger.info("3: Collect results from workers using a pythonic list "
                             "(using starmap_async and Lock()).")
        extra_3(conf['properties']["ch3"])
    else:
        extra_ch_logger.info("3: Collect results from workers using a pythonic list "
                             "(using starmap_async and Lock()).")
        extra_ch_logger.info("Last progress report:")
        extra_ch_logger.info("For this and the 4th problem, I managed to create a Lock() and share it")
        extra_ch_logger.info("with the workers by creating an initializer function for the Pool().")
        extra_ch_logger.info("The problem was that when I did the same with the multiprocessing "
                             "List() or Array(), ")
        extra_ch_logger.info("I wasn't able to access them from the function that spawned them.")

@timeit()
def main():
    """ This is the main function of assignment.py

    Example:
        python assignment1/assignment.py \
            -c ../confs/assignment1.yml \
            -l ../logs/assignment.log
    """

    # Initialize
    args = get_args()
    ColorizedLogger.setup_logger(args.log, args.debug)
    main_logger.info("Starting Assignment 1")
    # Load the configuration
    conf = Configuration(config_src=args.config_file)
    # Start the problems defined in the configuration
    main_logger.info(f"{' Required Problems ':-^{100}}")
    check_required = lambda conf_type, tag: ((conf_type == 'required' or tag != 'required_only')
                                             and conf_type != 'disabled')
    # For each problem present in the config file, call the appropriate function
    if 'problem1' in conf.config_keys:
        for bench_conf in conf.get_config(config_name='problem1'):
            if check_required(bench_conf['type'], conf.tag):
                problem1(bench_conf)
    if 'problem2' in conf.config_keys:
        for bench_conf in conf.get_config(config_name='problem2'):
            if check_required(bench_conf['type'], conf.tag):
                problem2(bench_conf)
    if 'problem3' in conf.config_keys:
        for bench_conf in conf.get_config(config_name='problem3'):
            if check_required(bench_conf['type'], conf.tag):
                problem3(bench_conf)
    # Run the extra challenges if the tag of the conf is not set as "required_only"
    main_logger.info(f"{' Optional Problems ':-^{100}}")
    if 'extra_challenges' in conf.config_keys:
        for bench_conf in conf.get_config(config_name='extra_challenges'):
            if check_required(bench_conf['type'], conf.tag):
                extra_challenges(bench_conf)
    main_logger.info("Assignment 1 Finished")

if __name__ == '__main__':
    try:
        main()
    except Exception as e:
        logging.error(str(e) + '\n' + str(traceback.format_exc()))
        raise e


import qiskit, time
from numpy import pi
# define pi so that in string gates we can have pi as an angle.
# Because we use eval for string gates. For example, gate = "rz(pi/2, 1)".

name = "IBM"

simulators = simulator, unitary_simulator, state_simulator = (
    "qasm_simulator", "unitary_simulator",
    "statevector_simulator"
)

quantum_computer = "ibmqx4"

def apply_credentials():
    print("\nApplying credentials...\n")
#    with open("qSonify/qc/APItoken.txt") as f: APItoken = f.read().strip()

    try:
#        qiskit.IBMQ.enable_account(APItoken)
        qiskit.IBMQ.load_accounts()

        print('Available backends:')
        print(qiskit.IBMQ.backends())
        print(qiskit.Aer.backends())
        print("\nCredientials applied\n")
    except:
        print('Something went wrong.\nDid you enter a correct token?')

#### String algorithm methods ####

# With this, we can write an algorithm as a list with any of the keys in
# GATE_ARGUMENTS. So, for example,
#   alg = ["H(0)", "RX(pi/2, 1)", "CX(1, 2)", "u3(pi/2, pi/4, .2, 0)"]
# then apply it to a qiskit.QuantumCircuit and qiskit.QuantumRegister qc and r
# respectively by calling
#   apply_string_algorithm(alg, r, qc).

p = lambda x: ("reg[%d]",)*x
a = lambda x: ("%g",)*x
b = lambda x, y: "(" + ", ".join(a(x)+p(y)) + ")"

GATE_PARAMS = { ## The first number is the number of parameters,
                ## The second number is the number of qubit arguments.
    "ccx": (0, 3), "ch": (0, 2), "crz": (1, 2), "cswap": (0, 3), "cu1": (1, 2),
    "cu3": (3, 2), "cx": (0, 2), "cx_base": (0, 2), "cy": (0, 2), "cz": (0, 2),
    "h": (0, 1), "iden": (0, 1), "rx": (1, 1), "ry": (1, 1), "rz": (1, 1),
    "rzz": (1, 2), "s": (0, 1), "sdg": (0, 1), "swap": (0, 2), "t": (0, 1),
    "tdg": (0, 1), "u0": (1, 1), "u1": (1, 1), "u2": (2, 1), "u3": (3, 1),
    "u_base": (3, 1), "x": (0, 1), "y": (0, 1), "z": (0, 1),
}

GATE_ARGUMENTS = {gate: b(*args) for gate, args in GATE_PARAMS.items()}
GATE_ARGUMENTS["measure"] = "(reg[%d], c_reg[%d])"

def get_gate_info(gate):
    """
    gate: str, string gate. ie H(0), or "cx(1, 0)".
    returns: tuple, (gate_name (str), gate_args (tuple)).
    """
    gate = gate.strip().lower().replace("cnot", "cx")
    i = gate.index("(")
    gate_name, gate_args = gate[:i], eval(gate[i:])
    try: len(gate_args)
    except TypeError: gate_args = gate_args,
    return gate_name, gate_args

def get_num_qubits(algorithm):
    """
    Determine the max qubit value used in the algorithm.

    algorithm: iterable, each element must be a string gate, as in
                     apply_string_gate above.
                     ie, algorithm = ["h(0)", "cx(0, 1)", "rx(pi/4, 1)",..]

    returns: int, max qubit value in algorithm.
    """
    n = -1
    for gate in algorithm:
        gate_name, gate_args = get_gate_info(gate)
        if gate_name == "measure": m = gate_args[0]
#        elif sum(GATE_PARAMS[gate_name]) == 1: m = gate_args
        else: m = max(gate_args[GATE_PARAMS[gate_name][0]:])
        n = max(n, m)
    return n + 1

def apply_string_gate(gate, reg, cir, c_reg=None):
    """
    gate: str, one of the elements in GATE_ARGUMENTS.keys() + a tuple of
               arguments. ie, for a rx rotation by pi/2 radians on qubit 0,
               gate = "rx(pi/2, 0)".
    reg: qiskit.QuantumRegister, register to apply gate to.
    cir: qiskit.QuantumCircuit, circuit to add gate to.
    c_reg: qiskit.ClassicalRegister, must be supplied if gate is a measurement.
                                     Classical register to measure to.

    returns: int, if gate is a measure gate, then return the integer
                  corresponding to the classical register to measure to,
                  otherwise returns -1.
    """
    gate_name, gate_args = get_gate_info(gate)

    # apply gate
    eval("cir." + gate_name + GATE_ARGUMENTS[gate_name] % gate_args)

    # value of the classical register to measure to
    if "measure" in gate: return gate_args[-1]
    else: return -1

def apply_string_algorithm(algorithm, reg, cir, c_reg=None):
    """
    algorithm: iterable, each element must be a string gate, as in
                         apply_string_gate above.
                         ie, algorithm = ["h(0)", "cx(0, 1)", "rx(pi/4, 1)",..]
    reg: qiskit.QuantumRegister, register to apply algorithm to.
    cir: qiskit.QuantumCircuit, circuit to add gates in algorithm to.
    c_reg: qiskit.ClassicalRegister, must be supplied if gate is a measurement.
                                     Classical register to measure to.

    returns: int, if the algorithm has any measure gates, then returns the
                  integer corresponding to the largest index of the classical
                  register that is measured to, otherwise returns -1.
    """
    if not algorithm: return -1
    return max(apply_string_gate(gate, reg, cir, c_reg) for gate in algorithm)

def _make_job(qc, backend, num_samples):
    """
    Begin the execution of the circuit qc on the backend with shots=num_samples

    qc: qiskit.QuantumCircuit or list of circuits, circuits to run.
    backend: str, IBM backend to run circuit on. Can be 'ibmqx4', 'ibmqx5',
                  'local_qasm_simulator', 'local_unitary_simulator', etc.
    num_samples: int, number of samples to take from the quantum computer in
                      in order to determine the probabilities for each state.

    returns: qiskit Job object from qiskit.backends.
    """
    if backend in simulators: f = qiskit.Aer
    else: f = qiskit.IBMQ
    try:
        return qiskit.execute(qc, backend=f.get_backend(backend),
                                  shots=num_samples, max_credits=3)
    except LookupError:
        apply_credentials()
        return qiskit.execute(qc, backend=f.get_backend(backend),
                                  shots=num_samples, max_credits=3)

class Result(dict):
    """ Just a dictionary that automatically gives default values = 0.0 """
    def __getitem__(self, key):
        """ Return 0.0 if key not in result dictionary """
        return self.get(key, 0.0)

def run(algorithm, num_qubits=None, num_samples=8000, backend=simulator):
    """
    Create a quantum circuit, run the algorithm, return the resulting
    probability distribution.

    algorithm: algorithm (list of strings) or list of algorithms,
               each string is a gate in GATE_ARGUMENTS.keys() with whatever
               arguments required to define the gate.
    num_qubits: int, number of qubits to run each algorithm on. Can be None,
                     in which case the algorithm will be run on the minimum
                     number of qubits required.
    num_samples: int, number of samples to take from the quantum computer in
                      in order to determine the probabilities for each state.
    backend: str, IBM backend to run the algorithm on. If backend is not
                  a local simulator then credentials must have already
                  been applied.

    returns: dict (common.Result), keys are states, values are probabilities
                                   found to be in that state.
    """
    multiple = bool(algorithm and isinstance(algorithm[0], list))
    if not multiple: algorithm = [algorithm]

    n = len(algorithm)
    if num_qubits is None:
        num_qubits = max(get_num_qubits(a) for a in algorithm)
    q = qiskit.QuantumRegister(num_qubits)
    c = [qiskit.ClassicalRegister(num_qubits) for _ in range(n)]
    qc = [qiskit.QuantumCircuit(q, c[j]) for j in range(n)]
    for j in range(n):
        i = apply_string_algorithm(algorithm[j], q, qc[j], c[j])
        if i == -1: qc[j].measure(q, c[j])
        else: c[j].size = i + 1

    job_exp = _make_job(qc, backend, num_samples)

    # Often there are random queue errors that have happened to
    # me that cause the job to never complete. Two things I have
    # encountered: I lose connection or something, and I get an
    # error, or for some reason their server tells me that the
    # job is running indefinitely, ie it just get stuck running.
    # So if either of those things happen, we reset and
    # reinitialize our job(s) into the queue.
    if backend not in simulators:
        lapse, interval = 0, 30
        done = False
        while not done:
            str_status = str(job_exp.status())
            queue_position = job_exp.queue_position()
            error = job_exp.error_message()
            print('\nStatus @ %d seconds' % (interval * lapse))
            print("queue position =", queue_position)
            print(str_status)
            done = queue_position is not None and queue_position < 1

            if error:
                print("\nEncountered an error")
                print(error)
                print("reentering job into queue\n")
                job_exp.cancel()
                job_exp = _make_job(qc, backend, num_samples)
                lapse = 0

            lapse += 1
            time.sleep(interval)

    res = job_exp.result()

    ## qiskit orders their bits opposite to Cirq nad ProjectQ, and in my
    ## opinion in a much less intuitive way. So I flip the order of the bits
    ## here.
    if multiple:
        return [
            Result(
                {k[::-1]: v/num_samples
                 for k, v in res.get_counts(cir).items()}
            ) for cir in qc
        ]
    else:
        return Result(
            {k[::-1]: v/num_samples for k, v in res.get_counts(qc[0]).items()}
        )

def algorithm_unitary(algorithm, num_qubits=None):
    """
    Find the unitary corresponding to the algorithm.

    algorithm: list of strings, each string is a gate in GATE_ARGUMENTS.keys()
                                with whatever arguments required to define the
                                gate.
    num_qubits: int, number of qubits to run the algorithm on.

    returns: np.array, unitary matrix corresponding to the algorithm.
    """
    if num_qubits is None: num_qubits = get_num_qubits(algorithm)
    if not algorithm: algorithm = ["iden(0)"]

    ## qiskit orders their bits opposite to Cirq nad ProjectQ, and in my
    ## opinion in a much less intuitive way. So I flip the order of the bits
    ## here.
    a = []
    for gate in algorithm:
        gate_name, gate_args = get_gate_info(gate)
        i = GATE_PARAMS[gate_name][0]
        params = gate_args[:i]
        qubits = gate_args[i:]
        qubits = tuple(num_qubits-q-1 for q in qubits)
        a.append(gate_name + str(params + qubits))

    q = qiskit.QuantumRegister(num_qubits)
    qc = qiskit.QuantumCircuit(q)
    apply_string_algorithm(a, q, qc)
    return qiskit.execute(
        qc, backend=qiskit.Aer.get_backend(unitary_simulator)
    ).result().get_data(qc)["unitary"]

def prepare_state(state):
    """
    state: string, string of 0's and 1's.
    returns: algorithm (list of strings), algorithm to prepare the state.
    """
    return ["x(%d)" % i for i in range(len(state)) if state[i] == "1"]

def sample(algorithm, num_qubits=None,
           num_samples=1, backend=simulator):
    """
    Get a list of all the outputs from an algorithm. Differs from `run` because
    `run` returns the determined probabilities of each state, but `sample`
    returns a list of the outputs.

    algorithm: algorithm (list of strings), NOT a list of algorithms,
               each string is a gate in GATE_ARGUMENTS.keys() with whatever
               arguments required to define the gate.
    num_qubits: int, number of qubits to run each algorithm on.
    num_samples: int, number of samples to take from the quantum computer.
    backend: str, IBM backend to run the algorithm on. If backend is not
                  a local simulator then credentials must have already
                  been applied.

    returns: list, each element is the measured state.
    """
    d = run([algorithm]*num_samples,
            num_qubits=num_qubits,
            num_samples=1, backend=backend)
    return [list(x.keys())[0] for x in d]

def single_sample(algorithm, num_qubits=None, backend=simulator):
    """
    Same as `sample` with one sample, but returns a state instead of a list of
    one state.
    """
    return sample(algorithm, num_qubits, 1, backend)[0]

def markovian_sample(algorithm, num_qubits=None,
                     num_samples=1, backend=simulator):
    """
    Get a list of all the outputs from an algorithm, where the previous output
    is prepared as the starting point for the next algorithm; ie the
    measurement of the algorithm is input to run the algorithm again.

    algorithm: algorithm (list of strings), NOT a list of algorithms,
               each string is a gate in GATE_ARGUMENTS.keys() with whatever
               arguments required to define the gate.
    num_qubits: int, number of qubits to run each algorithm on.
    num_samples: int, number of samples to take from the quantum computer.
    backend: str, IBM backend to run the algorithm on. If backend is not
                  a local simulator then credentials must have already
                  been applied.

    returns: list, each element is the measured state.
    """
    if num_samples < 1: raise ValueError("Must have >= 1 sample")
    if num_qubits is None: num_qubits = get_num_qubits(algorithm)
    args = num_qubits, backend
    res = [single_sample(algorithm, *args)]
    for _ in range(num_samples-1):
        res.append(single_sample(prepare_state(res[-1])+algorithm, *args))
    return res

if __name__ == "__main__":
    ## Examples

    # `run` returns a dictionary mapping states to probabilities, ie
    # run(["h(0)", "cx(0, 1)"]) should return {"00":0.5, "11": 0.5}.

    # if no "measure" is included in alg, then by default everything will
    # be measured.
    alg = ["H(0)", "CX(0, 1)", "u3(pi, pi/2, pi/4, 0)"]
    print(run(alg, 3, num_samples=10000))

    # since a measure is included, only that register will be measured.
    alg = ["H(0)", "CX(0, 1)", "u3(pi, pi/2, pi/4, 0)", "measure(0, 0)"]
#    print(run(alg, 3, num_samples=1000, backend="ibmqx4"))

    # run multiple circuits at once
    alg0 = ["h(0)", "cx(0, 1)", "measure(0, 0)", "measure(1, 1)"]
    alg1 = ["x(0)", "H(1)", "ccx(0, 1, 2)"]
    print(run([alg0, alg1]))

    # convert alg to its unitary respresentation.
    alg = ["h(0)", "cx(0, 1)"]
    print(algorithm_unitary(alg, 2))


"""
Unit testing with Python for Session1 function

@brief get started with python testing by looking at https://docs.pytest.org/en/latest/getting-started.html#getstarted. Note that any script to be ran within the python testing framework (pytest) should follow the standard test discovery rules (https://docs.pytest.org/en/latest/goodpractices.html#test-discovery)
"""

import os
import numpy as np

print('Starting test script from working directory : ' + os.getcwd())

#testing session 1 functions
def load_s1_script():
    """
        utility function that tris to load the script written along the first lesson
        @throws an ImportError exception if the script file does not exist
        @return the script as a loaded module
    """
    s1_script_filename = 'assignments/Session1/S1_algotools.py'
    print('Trying to load target scripts : ' + s1_script_filename)
    import imp
    s1_algotools=imp.load_source('session_1_script', s1_script_filename)
    return  s1_algotools

#load the scripts to check
def test_session1script_exists():
    try:
        load_s1_script()
        assert True
    except  ImportError:
        print('Expected script not found, carrefuly check the assignement instructions ')
        assert False

# ----------- TEST FOR average_above_zero FUNCTION -----------

def test_s1_selective_average_non_zeros_values():
    ##
    # @test validates average_above_zero works fine with integer values > 0
    assert load_s1_script().average_above_zero([1, 2, 3, 4, -7]) == 2.5

def test_s1_selective_average_with_zeros_values():
    ##
    # @test validates average_above_zero works fine with integer values >= 0
    assert load_s1_script().average_above_zero([0, 1, 2, 3, 4, -7]) == 2.0

def test_s1_selective_average_with_negative_values():
    ##
    # @test validates average_above_zero works fine with integer values <= 0
    assert load_s1_script().average_above_zero([0, -7]) == 0

def test_s1_selective_average_with_string_values():
    ##
    # @test validates average_above_zero works fine with integer values <= 0
    try:
        load_s1_script().average_above_zero(['ad', 'c'])
        assert False
    except ZeroDivisionError:
        assert True

def test_s1_selective_average_with_empty_values():
    ##
    # @test validates average_above_zero works fine with an empty list
    try:
        load_s1_script().average_above_zero([])
        assert False
    except ValueError:
        assert True

def test_s1_selective_average_with_string():
    ##
    # @test validates average_above_zero works fine with an string
    try:
        load_s1_script().average_above_zero('string')
        assert False
    except TypeError:
        assert True

# ----------- TEST FOR max_value FUNCTION -----------

def test_s1_max_value():
    ##
    # @test validates max_value works fine with integer values > 0
    assert load_s1_script().max_value([1, 2, 3, 4, 10]) == (10, 4)

def test_s1_max_value_with_negative_value():
    ##
    # @test validates max_value works fine with integer values < 0
    assert load_s1_script().max_value([-1, -2, -3, -4, -10]) == (-1, 0)

def test_s1_max_value_with_max_at_index_0():
    ##
    # @test validates max_value works fine with integer values > 0
    assert load_s1_script().max_value([77, 2, 3, 4, 10]) == (77, 0)

def test_s1_max_value_with_max_at_index_2():
    ##
    # @test validates max_value works fine with integer values > 0
    assert load_s1_script().max_value([1, 2, 42, 4, 10]) == (42, 2)

def test_s1_max_value_with_empty_array():
  ##
  # @test validates max_value works fine with empty array
  try:
    load_s1_script().max_value([])
    assert False
  except ValueError:
    assert True

def test_s1_max_value_with_no_array_value():
  ##
  # @test validates max_value works fine with an no array value
  try:
    load_s1_script().max_value('array')
    assert False
  except TypeError:
    assert True

# ----------- TEST FOR reverse_table FUNCTION -----------

def test_reverse_table_value():
    ##
    # @test validates reverse_table works fine with correct array
    assert load_s1_script().reverse_table([1, 2, 3, 4]) == [4, 3, 2, 1]

def test_reverse_table_value_negative():
    ##
    # @test validates reverse_table works fine with correct array, with negative value
    assert load_s1_script().reverse_table([1, -2, -3, 4]) == [4, -3, -2, 1]

def test_reverse_table_with_no_int_values():
    ##
    # @test validates reverse_table works fine with correct array but with not only number inside
    assert load_s1_script().reverse_table([1, -2, 'v', [4, 3]]) == [[4, 3], 'v', -2, 1]

def test_reverse_table_with_string():
    ##
    # @test validates reverse_table works fine with string instead of array
    try:
        load_s1_script().reverse_table('array')
        assert False
    except TypeError:
        assert True

# ----------- TEST FOR roi_bbox FUNCTION -----------

def test_roi_bbox_no_np_array():
    ##
    # @test validates roi_bbox works fine with no np.array given.
    try:
        load_s1_script().roi_bbox(42)
        assert False
    except TypeError:
        assert True

def test_roi_bbox_with_7_7_array():
    ##
    # @test validates roi_bbox works fine with 7*7 array.
    assert load_s1_script().roi_bbox(np.array([
            [0, 0, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 1, 0],
            [0, 0, 0, 1, 0, 0, 0],
            [0, 0, 0, 1, 1, 0, 0],
            [0, 0, 1, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 0, 0]
        ])).tolist() == [[2, 1], [5, 1], [2, 4], [5, 4]]

# ----------- TEST FOR random_fill_sparse FUNCTION -----------

def test_random_fill_sparse_value_with_no_int_k():
    ##
    # @test validates random_fill_sparse works fine with k as not Int
    try:
        load_s1_script().random_fill_sparse(np.array([['', '', ''], ['', '', '']]), 'dd')
        assert False
    except TypeError:
        assert True

def test_random_fill_sparse_value_with_no_np_array_table():
    ##
    # @test validates random_fill_sparse works fine with table as not np.array
    try:
        load_s1_script().random_fill_sparse([['', '', ''], ['', '', '']], 2)
        assert False
    except TypeError:
        assert True

# ----------- TEST FOR remove_whitespace FUNCTION -----------

def test_remove_whitespace_with_one_whitespace():
    ##
    # @test validates remove_whitespace works fine with string with 1 whitespace.
    assert load_s1_script().remove_whitespace('str ing') == 'string'

def test_remove_whitespace_with_many_whitespace():
    ##
    # @test validates remove_whitespace works fine with string with many whitespace.
    assert load_s1_script().remove_whitespace('   s    t   r i   n   g   ') == 'string'

def test_remove_whitespace_with_no_string_value():
    ##
    # @test validates remove_whitespace works fine with no string given.
    try:
        load_s1_script().remove_whitespace(42)
        assert False
    except TypeError:
        assert True

# ----------- TEST FOR shuffle FUNCTION -----------

def test_shuffle_with_no_array_value():
    ##
    # @test validates shuffle works fine with no array given.
    try:
        load_s1_script().shuffle(42)
        assert False
    except TypeError:
        assert True

def test_shuffle_with_array():
    ##
    # @test validates shuffle works fine with array, test if all number are there.
    array_test = [i for i in range(10)]
    init_sum = sum(array_test)
    result = load_s1_script().shuffle(array_test)

    assert init_sum == sum(result)

# ----------- TEST FOR sort_selective FUNCTION -----------

def test_sort_selective_with_no_array_value():
    ##
    # @test validates sort_selective works fine with no array given.
    try:
        load_s1_script().sort_selective(42)
        assert False
    except TypeError:
        assert True

def test_sort_selective_with_array_of_ten():
    ##
    # @test validates sort_selective works fine with array of 10 values.
    assert load_s1_script().sort_selective([2, 5, 3, 7, 9, 6, 7, 7, 1, 0]) == [0, 1, 2, 3, 5, 6, 7, 7, 7, 9]

# ----------- TEST FOR sort_bubble FUNCTION -----------

def test_sort_bubble_with_no_array_value():
    ##
    # @test validates sort_bubble works fine with no array given.
    try:
        load_s1_script().sort_bubble(42)
        assert False
    except TypeError:
        assert True

def test_sort_bubble_with_array_of_ten():
    ##
    # @test validates sort_bubble works fine with array of 10 values.
    assert load_s1_script().sort_bubble([2, 5, 3, 7, 9, 6, 7, 7, 1, 0]) == [0, 1, 2, 3, 5, 6, 7, 7, 7, 9]


"""

Reads a file containing columns of data and creates a dictionary according to header information
using the following format:

1. comments start with ``#`` in the **first column** : ``# My Comment``

2. header information starts with ``#!``, where ``#`` is also in the **first column**. It must precede the data.

3. each columns is represented by ``name[dtype, col.nr.]/``,  *dtype* (optional) is the data type and *col.nr*
   (starting at 0) is the column number

4. data are separated by white space **NOT commas!**

5. data types:

      - s : string
      - f : float
      - i : integer

6. blank lines are ignored

 **NOTE**: if data formats are entered they should be specified for all columns

Example data::

   #! p_miss[0]/ siglt[2]/ s01[3]/ alt[4]/
   200. 1.35e-4 -1.e-3    0.1
   220. 2.56e-4 -2.e-4    -0.1
   230. 3.47e-6 -3.e-5    1.1

The header can also contain data type information::

   #! p_miss[f,0]/ siglt[f,2]/ s01[f,3]/ alt[f,4]/

   200. 1.35e-4 -1.e-3    0.1
   220. 2.56e-4 -2.e-4    -0.1
   230. 3.47e-6 -3.e-5    1.1

Example that opens a file and create a dfile object::

   >>> f0 = dfile('sig_LT.dat')

Loop over content::

   >>> for l in f0:
   >>> ...pm = l['p_miss']
   >>> ...sig_lt = l['siglt']*10000.
   >>> ...print pm, sig_lt, l['alt']
   >>> # end for l

Variables can also be accessed directly as numpy arrays (if installed):

   >>> pm = f0['p_miss']
   >>> sig_lt = f0['siglt']*10000.

They can also be converted to attributes:

   >>> f0.make_attr()
   >>> print f0.p_miss
   >>> print f0.sig_lt

Variable names are also called keys.

Save the data file as a csv file, this is useful for exporting and formatting for
documents (other than latex)

   >>> f0.write_csv(filename)

---------------------------------------------------------

"""

import sys
import os
import string
import math
import re
import pdb

numpy_ok = True
try:
    import numpy as np
except:
    numpy_ok = False

# function to create a dictionary from 2 lists: key, values
format_dict = {'i':'{:d}', 'f':'{:f}', 's':'{:s}'}
# create a dictionary from  an array of keys and values
#----------------------------------------------------------------------
def find_duplicates(array):
    # array with no duplicates
    na = []
    # counter array
    count = []
    # arrays with duplucates
    da = []
    for a in array:
        try:
            na.index(a)
        except:
            na.append(a)
            count.append(0.)
    # na has no duplicates
    for i,a in enumerate(na):
        for aa in array:
            if aa == a:
                count[i]+= 1
    for i,c in enumerate(count):
        if c>1:
            da.append(na[i])
    return da
#----------------------------------------------------------------------

class dfile:
    """
    Open a file, read and interpret the contents and return a dfile object:

    >>> df = dfile('my_datafile')

    keywords:

        debug = False (default) : print additional iformation

        fast = False (default) : load data in a fast way (using np.loadtxt) this has
                                 some limits for string entries. Best used for large purely
                                 numerical data.

        use_numpy = True (if numpy is installed) : set automatically but can be overridden
                                                   some attributes of datafile are not availables when
                                                   numpy is not present

    """
    def __init__(self, filename, debug = False, new=False, skip = True, use_numpy = numpy_ok, fast = False):
        self.H=re.compile("^#\!") # pattern for header
        self.C=re.compile("^#")   # pattern for comment
        # pattern for splitting header
        # currently:
        # ^#\! : matches line beginning with #!
        # \[[0-9]+\]\/: matches [0]/ ... [8796]/
        # \[ *\w+ *, *\w *\]\/ : matches [   f  , blabla   ]/
        #
        self.S=re.compile("^#\!|\[[0-9]+\]\/|\[ *\w *, *\w+ *\]\/")
        # pattern to find format information in the header
        self.F=re.compile("\w+ *, *\w+")
        # pattern to find a character
        self.Fc=re.compile("[a-zA-Z]")
        # pattern to find a number
        self.Fn=re.compile("[0-9]+ *\]")
        # supported formats: the output format dictionary key must match the supported formats
        self.fmt_characters = ['f','i','s']
        self.dt_formats = {'f':'f8', 'i':'i4','s':'S32'}
        self.output_format= {'f':'%r ','i':'%d ','s':'%s '}
        # empty dictionary for the formats
        self.formats = {}
        self.formats_dt = {}
        self.cols = {}
        # flag for printing debugging information
        self.debug = debug
        # flat to ignore bad line in input
        self.skip = skip
        self.use_numpy = use_numpy
        if not self.use_numpy:
            print('numpy will not be used !')
        if fast and (not use_numpy):
            print ('fast mode not possible as numpy cannot be used!')
            fast = False
        self.fast = fast
        #
        self.header = None
        self.data = []
        self.adata = []
        self.fdata=[]
        self.keys=[]
        self.new = False
        if new:
            new=True
            self.filename = filename
            self.keys.append('indx')
            self.headindex=0
            self.headerline = ('#! ')
            self.adata.append('#! ')
            return
        # pdb.set_trace()
        # open file
        self.filename = filename
        i=open(filename,"r")      # open file
        # remove leading and trailing spaces
        self.adata=[ x.strip() for x in i.readlines()]   # read all data
        if self.debug :
            print("datafile --> data read !")
        if self.debug :
            print("datafile --> spaces removed read !")
        self.remove_blanks()
        if self.debug :
            print("datafile --> blank lines removed read !")
        i.close()
        if (self.find_header() != 0): # find the header
            print("cannot interpret data, header probably missing or wrong")
            return
        if self.debug :
            print("datafile --> create arrays !")
        self.make_array()  # read data
        if self.debug :
            print("datafile --> arrays created !")
    def __getitem__(self,i):
        # allows 'direct' access to the data
        if type(i) is int:
            # it'sm index
            return self.data[i]
        elif type(i) is str:
            # its a lkey return the list of data
            return self.get_data(i)
    def __len__(self):
        return len(self.data)
    def remove_blanks(self):
        # work through the list in reverse order
        indices = list(range( len( self.adata ))) # list of indices
        indices.reverse() # reverse order
        for i in indices: # loop through from the top
            if len(self.adata[i].strip()) == 0:
                del self.adata[i]
    def find_header(self):
        # find header information : look for a line starting with #!
        for l in self.adata:
            if (self.H.match(l) != None):
                self.headindex=self.adata.index(l)
                self.headerline=l
                self.header=re.split(self.S,l)[1:-1] # remove first and last element
                self.keys=[ x.strip() for x in self.header]
                self.keys.append('indx') # store the index into the original
                # handle data format
                # check for duplicate keys
                duplicate_keys = find_duplicates(self.keys)
                if duplicate_keys != []:
                    print("duplicate keys found : ", duplicate_keys)
                    return -1
                # pdb.set_trace()
                fmt = re.findall(self.F,l)
                col_t = re.findall(self.Fn,l)
                col = [cc.replace(']','') for cc in col_t]
                # no format specification, set default format
                if fmt == []:
                    for i,k in enumerate(self.keys[:-1]) : # skip index
                        self.formats[k] = 'f'
                        self.formats_dt[k] = self.dt_formats['f']
                # use format specifications
                else:
                    for i,f in enumerate(fmt):
                        fc = re.findall(self.Fc,f)[0] # use 1st char. for format
                        self.formats[self.keys[i]] = fc
                        self.formats_dt[self.keys[i]] = self.dt_formats[fc]
                # addformat for the index
                self.formats['indx'] ='i'
                self.formats_dt['indx'] = self.dt_formats['i']
                # use format specifications
                for i,cc in enumerate(col):
                    f = self.formats[self.keys[i]]
                    if f == 's':
                        self.cols[self.keys[i]] = ''
                    else:
                        self.cols[self.keys[i]] = col[i]
                self.cols['indx'] ='-1'
                # addformat for the index
                # adata array
                # do some checks
                if (len(self.keys)  != len(self.formats)):
                    print(self.filename,": problem in data formats !")
                    print(self.adata)
                    return -1
                return 0
            #
        print(self.filename,": no header information !")
        print('data dump : ')
        print(self.adata)
        return -1

    def parse_line(self, l):
        if (self.C.match(l) != None) : # not a comment
            return None
        ll=l.split()
        # handle file format problems
        if len(ll) > (len(self.keys)-1) :
            n_data = (len(self.keys)-1)
            print(self.filename, ' error: more data than variable names !')
            print(' ---> check header : ', self.headerline)
            print('current line : ', self.line_n+1, ' ', ll)
            if self.skip:
                print('skipping !')
                return None
            else:
                sys.exit(-1)
        if len(ll) < (len(self.keys)-1) :
            print(self.filename, ' error: too few data values !')
            print(' ---> check header : ', self.headerline)
            print('current line : ', self.line_n+1, ' ', ll)
            if self.skip:
                print('skipping !')
                return None
            else:
                sys.exit(-1)
        f = []
        # now convert to values
        for i,k in enumerate(ll):
            # check if there are format specifications
            f_err = False
            try:
                fmt = self.formats[self.keys[i]]
                fc = fmt
                # check that the format is supported
                fi = self.fmt_characters.index(fmt)
            except:
                f_err = True
                print('unknown format : ', fc, ' set to f')
                fc = 'f'
                self.formats[self.keys[i]] = 'f'
            # now try to convert according to the format
            try:
                if fc == 'f':
                    f.append(float(k)) # create a float
                if fc == 'i':
                    f.append(int(k))   # create an integer
                if fc == 's':
                    f.append(k)        # leave as string
            except:
                f.append(k) # if not possible leave it,
                # but make a comment
                if len(self.formats) == 0 :
                    print('cannot convert : ', k, ' to float, line : ', self.line_n +1)
                else:
                    if (f_err ):
                        print('cannot convert : ', k, ' line : ', self.line_n+1,\
                            ', problem with unknown format : ', fmt)
                    else:
                        print('cannot convert : ', k, ' line : ', self.line_n+1,\
                            ', with format :', fmt)
        f.append(self.adata.index(l))
        # pdb.set_trace()
        return f

    def make_array(self):
        if self.fast:
            # use loadtxt to get the data
            data_types = [self.formats_dt[k] for k in self.keys[:-1]]
            self.data = np.loadtxt(self.filename, dtype = {'names':self.keys[:-1], 'formats':data_types})
        else:
               # create an array of dictionaries
            self.data = []
            n_data = len(self.adata)
            for self.line_n, l in enumerate(self.adata):
                f = self.parse_line(l)
                if f == None:
                    continue
                self.data.append(self.make_dict(self.keys,f) )
                if self.debug:
                    print("processed data : ", self.line_n, " out of ", n_data)

    def make_dict(self,keys,values):
        # create a dictionary from a list of keys and values
        p = []
        for j,key in enumerate(keys): # loop over keys
            try:
                p.append( (key.strip(),values[j]) ) # an array of tuples
                d=dict(p) # create a dictionary
            except:
                print('problem in data : ', key, values, j)
            #
        return d

    def make_attr(self):
        # make attributes out of the keys
        for k in self.keys:
            ka = ''
            try:
                self.__dict__[k]
                # attrtibute exists append _d
                ka = k + '_d'
                print(("{:s} exists using : {:s}".format(k, ka)))
            except:
                ka = k
            # replace periods with underscores
            if self.use_numpy:
                setattr(self, ka.replace('.','_'), np.array(self[k]))
            else:
                setattr(self, ka.replace('.','_'), self[k])

    def scale(self, key, factor):
        """
        multiply all values of key with a factor
        """
        for i, d in enumerate(self.data):
            self.data[i][key] = factor*d[key]
        #

    def sort(self, key, **kwargs):
        """
        sort the data according to the values in key
        """
        index_array = []
        for d in self.data:
            index_array.append((d[key], self.data.index(d))) # store index and values pair
        sorted_array = sorted(index_array,**kwargs) # create a sorted array
        temp = []
        for ip in sorted_array:
            temp.append(self.data[ip[1]])
        self.data = temp
        temp=[]

    def name(self):
        """
        print the filename associate with this instance

        """
        print("Input file name : ", self.filename)
    def show_keys(self):
        """
        print a list of variable names in the dictionary

        """
        print(self.keys)
    def get_keys(self):
        """

        return a list of keys

        """
        return self.keys
    def get_header(self):
        """

        return the header lines

        """
        return self.adata[self.headindex]

    def get_full_header(self):
        """

        return all line up to the header line

        """
        return self.adata[:self.headindex+1]

    def show_data(self,keylist):
        """

        print all the data corresponding to the key list::

           >>> df.show_data('key1:key2:key3')

        """
        akey=keylist.split(":")
        for l in self.data:
            ll=[]
            for i in akey:
                format = self.output_format[ self.formats[i] ]
                ll.append( format%( l[i] ) )
            print(ll)

    def check_data(self, func, data, key, *args):
        """

        function used to check if data fulfill a condition provided
        by the user. The function is assumed to return True of False.

        """
        return func( data, key, *args)

    def select_data(self, sel_func = None, sel_args = None):
        """
        returns an iterator for the data. As in get_data() a selector function and
        its arguments can be supplied:

        conditions can be applied:

        1) sel_func: a user provided function returning True or False

        2) sel_args: a list of arguments used in the sel_func function

        The iterator returned can be used as follows:

        >>> for d in df.select_data():
        >>> ...print d

        or the result can be converted to a list:

        >>> list(df.select_data())

        Using a selector function:

        Example: assume the data contain a variable (key) called 'name'
           you want to select only those data where name contains a certain substring 'sub'

        Define this function:

        >>> def myfind(data, key , what):
        >>> ...where = data[key]
        >>> ...return (str.find(where, what) >= 0)

        now you can select the data using:

        >>> list( df.select_data( myfind, 'name', 'sub') )

        """
        if sel_func == None:
            for i,data in enumerate(self.data):
                yield data
        else:
            for i,data in enumerate(self.data):
                if self.check_data(sel_func, data, *sel_args):
                    yield data
        # that's all

    def select_data_eval(self, eval_str):
        """
        returns an iterator for data selected with an *eval* expression
        stored in the string `eval_str` each dataset item is accessed using the name data::

        >>> df.select_data_eval("data['x'] >= 0.")

        returns only those data items where the value of the x-column in the file
        is larger than or equal to 0.

        """
        for i,data in enumerate(self.data):
            if eval( eval_str):
                yield data
        # that's all

    def eval_data(self, eval_str):
        """
        iterator over all data evaluating an expression contained in eval_str
        """
        for i,data in enumerate(self.data):
            yield eval(eval_str)
        # that's all

    def get_data(self,key, sel_func = None, sel_args = None):
        """

        return all data for `key` subject to the results of a selector function.

        one can define a selector function (`sel_func`) using the arguments stored in `sel_args`.
        This function is evaluated for each data record and only
        those data are returned for which the condition is fulfilled.

        1) sel_func: a user provided function returning True or False

        2) sel_args: a list of arguments used in the sel_func function

        Example::
           assume the data contain a variable (key) called 'name'
           you want to select only those data where name contains a certain substring 'Jo'

        Define this function::

        >>> def myfind(data, key , what):
        >>> ...where = data[key]
        >>> ...return (str.find(where, what) >= 0)

        now you can select the data using::

        >>> df.get_data('name',myfind, ['name','Jo'] )

        This should return a list of names containing the substring 'Jo'

        """
        array=[]
        for l in self.data:
            if sel_func == None:
                array.append(l[key])
            else :
                if self.check_data( sel_func, l, *sel_args):
                    array.append(l[key])
        if self.use_numpy:
            return np.array(array)
        else:
            return array

    def get_data_eval(self,key,eval_str): # return data for a  key
        """

        return all data for the `key` under the condition that the expression in `eval_str`
        is True.

        """
        array=[]
        for l in self.select_data_eval(eval_str):
            array.append(l[key])
        if self.use_numpy:
            return np.array(array)
        else:
            return array

    def get_data_list(self,keylist, sel_func = None, sel_args = None):

        """

        return all the data corresponding to the key list
        as follows::

           >>> a = df.get_data_list('key1:key2:key3')

           >>> a = df.get_data_list('key1:key2:key3', myfind, ['name','J'])

        return those data where the `name` values contain the character `J`

        *a* contains the list of data

        """

        akey=keylist.split(":")
        dd = []
        if sel_func == None:
            for l in self.data:
                ll=[]
                for i in akey:
                    ll.append( l[i] )
                dd.append(ll)
            if self.use_numpy:
                return np.array(dd)
            else:
                return dd
        else:
            for l in self.data:
                if self.check_data(sel_func, l, *sel_args):
                    ll=[]
                    for i in akey:
                        ll.append( l[i] )
                    dd.append(ll)
            if self.use_numpy:
                return np.array(dd)
            else:
                return dd

    def get_data_list_eval(self,keylist, eval_str):
        """

        similar function as select_data_eval but it only returns the
        values for the keys defined in `keylist`.

        """
        akey=keylist.split(":")
        dd = []

        for l in self.select_data_eval(eval_str):
            ll=[]
            for i in akey:
                ll.append( l[i] )
            dd.append(ll)
        if self.use_numpy:
            return np.array(dd)
        else:
            return dd

    def show_all_data(self):
        """

        print all data and keys stored

        """
        for l in self.data:
            print(l)
    def get_all_data(self):
        """

        return a list of all data stored

        """
        return self.data
    def write_complete_header(self,fp):
        """
        write entire header including comments and internal parameters to file with
        handle fp. Example:

        fp = open('mydile.data','w')

        datafile.write_comlete_header(fp)

        """
        for i,l in enumerate(self.adata[0:self.headindex]):
            fp.write(l+'\n')
        # write header line
        h = '#! '
        n = 0
        for k in self.keys[:-1]:
            h = h + k+'['+'%s,%d'%(self.formats[k],n)+']/ '
            n = n +1
        fp.write(h+'\n')
    def update_header(self):
        """
        update header line of this data file including format. This
        makes sure that the header line is in sync with the dictionary keys
        """
        # write header line
        h = '#! '
        n = 0
        for k in self.keys[:-1]:
            h = h + k+'['+'%s,%d'%(self.formats[k],n)+']/ '
            n = n +1
        # update data arrays
        self.adata[self.headindex]=h
        self.headerline = h
    def write_header(self,fp):
        """
        write only header line of this data file into file fp, including format
        """
        # write header line
        h = '#! '
        n = 0
        for k in self.keys[:-1]:
            h = h + k+'['+'%s,%d'%(self.formats[k],n)+']/ '
            n = n +1
        fp.write(h+'\n')
    def write_line(self,fp, i):
        """
        write data line i into file fp
        """
        l = ''
        # check array boundary
        ndata = len(self.data)
        if (i<0) :
            print('index below lower bound, set to 0')
            i = 0
        if (i>=ndata):
            print('index above upper bound, set to : ', ndata-1)
        for k in self.keys[:-1]:
            format = self.output_format[ self.formats[k] ]
            l = l + format%(self.data[i][k])
        fp.write(l+'\n')
    def write_all(self,fp, complete_header = False):
        """
        write all data to a file associated to fp

        if complete_header = True include the complete header including all comments
        """
        if complete_header:
            self.write_complete_header(fp)
        else:
            self.write_header(fp)
        for i in range(len(self.data)):
            self.write_line(fp, i)
        fp.close()
    #   new function to add a key
    def write_selected(self,fp,index_list, complete_header = False):
        """
        write a new datafile with an identical header but enter only
        those data with an index given in index_list. If the index does not
        exist print a message and skip it.

        if complete_header = True include the complete header including all comments

        """
        if complete_header:
            self.write_complete_header(fp)
        else:
            self.write_header(fp)
        for i in index_list:
            try:
                self.write_line(fp, i)
            except:
                print('cannot write index : ', i, ' does it exist ?')
        fp.close()

    def add_key(self,key,format='f'):
        """
        add a new key, this is useful if you want to add new data
        assign new values using a loop like::

        >>> df.add_key('newkey',format='i')

        The new data need to be stored as follows::

        >>> for d in df.data:
        >>>     d['newkey']= some_new_value

        where df is the datafile and 'newkey' the new key. The new data set
        can be saved using ``fp = open('new_file','w')`` and ``df.write_all(fp)``
        """
        # check if key already exists
        try:
            self.keys.index(key)
            print(key, " exists already !")
            return
        except:
            for d in self.data:
                # set initial value to none
                d[key] = None
        # add the new key at the end of the keys list but before indx
        self.keys.insert(len(self.keys)-1,key)
        # add the format information
        try:
            fi = self.fmt_characters.index(format)
        except:
            print('unknown format : ', format, ' use f instead')
            format = 'f'
        self.formats[key] =  format
        self.update_header()
    # end
    def delete_key(self,key):
        """
        remove a key and all values associated with it
        """
        # remove the data for the key
        for d in self.data:
            del d[key]
        # remove the format for this key
        fmt = self.formats.pop(key)
        # now remove the key from the key list
        self.keys.remove(key)
        print('key : ', key , ' and format : ', fmt, ' removed')
        self.update_header()
    # end

    def add_data(self, keys,line):
        """

        add data to the data file

        d.add_data('x:y:z','xval yval zval')

        xval, yval, zval are the values as they would be netered in a data file line

        both arguments are strings !

        IMPORTANT: add the values to all current lines of the data file !

        """

        key_l = keys.split(':')
        # check the key list
        key_index=[]
        for k in self.keys[:-1]:  # skip the indx
            try:
                ki = key_l.index(k)
                key_index.append(ki)
            except:
                print(k, ' does not exist !, nothing added')
                return
        # sort the argumens according to the sequence of the keys
        lsp = line.split()
        fline = [lsp[i] for i in key_index]
        line = ''
        for fl in fline: line += fl+' '
        self.adata.append(line.strip())
        print(line)
        """
        f = self.parse_line(line.strip())
        self.data.append(self.make_dict(self.keys,f) )
        """
        self.make_array()

    def add_header_comment(self, text):
        """
        add a comment line to the header. The # at the beginning is automatically added !

        * for a normal comment start with a space
        * to add a parameter start the comment with a \
        """
        self.adata.insert(self.headindex, '#' + text)
        self.headindex += 1

    def add_parameter(self, name, value):
        """
        add a parameter to the comment section
        """
        cmd = '\ {} = {} '.format(name, value)
        self.add_header_comment(cmd)

    def save(self, file = None):
        """

        Save the current datafile.

        With the keyword: *file* = 'new_file'

        the datafile will be written into the new_file name

        """
        if file == None:
            o = open(self.filename,'w')
        else:
            o = open( file, 'w')
        self.write_all(o, complete_header = True)

    def make_fmt(self):
        fmt = ''
        for k in self.keys[:-1]:
            fmt += format_dict[self.formats[k]] + ', '
        return fmt.strip()[:-1]

    def make_csv(self):
        keys = self.keys[:-1]
        csv = [ ''.join([k + ',' for k in keys]) ]
        for l in self.data:
            fmt = self.make_fmt()
            data = [l[k] for k in keys]
            ll = fmt.format(*data)
            csv.append(ll)
        return csv

    def write_csv(self, f):
        """

        save the current file as a csv file

        f : file name  to be used

        """
        o = open(f,'w')
        csv = self.make_csv()
        for l in csv:
            o.write(l + '\n')
        o.close()

#


import numpy as np      #the numpy library
import matplotlib.pyplot as plt     #this is matplotlib's pyplot

import sys  #gives access to C-like sys libraries
import os   #gives access to operating systems

print(sys.argv)     #command line arguments
print(os.getcwd)    #this gives the current working directory