def test_analog_single_tet():
    """This test tests uniform sampling within a single tetrahedron. This is
    done by dividing the tetrahedron in 4 smaller tetrahedrons and ensuring
    that each sub-tet is sampled equally.
    """
    seed(1953)
    mesh = iMesh.Mesh()
    v1 = [0, 0, 0]
    v2 = [1, 0, 0]
    v3 = [0, 1, 0]
    v4 = [0, 0, 1]
    verts = mesh.createVtx([v1, v2, v3, v4])
    mesh.createEnt(iMesh.Topology.tetrahedron, verts)
    m = Mesh(structured=False, mesh=mesh)
    m.src = IMeshTag(1, float)
    m.src[:] = np.array([1])
    m.mesh.save("tet.h5m")
    center = m.ve_center(list(m.iter_ve())[0])

    subtets = [[center, v1, v2, v3], [center, v1, v2, v4], [center, v1, v3, v4], [center, v2, v3, v4]]

    sampler = Sampler("tet.h5m", "src", np.array([0, 1]), False)
    num_samples = 5000
    score = 1.0 / num_samples
    tally = np.zeros(shape=(4))
    for i in range(num_samples):
        s = sampler.particle_birth([uniform(0, 1) for x in range(6)])
        assert_equal(s[4], 1.0)
        for i, tet in enumerate(subtets):
            if point_in_tet(tet, [s[0], s[1], s[2]]):
                tally[i] += score
                break

    for t in tally:
        assert abs(t - 0.25) / 0.25 < 0.2

def test_analog_single_hex():
    """This test tests that particles of sampled evenly within the phase-space 
    of a single mesh volume element with one energy group in an analog sampling
    scheme. This done by dividing each dimension (x, y, z, E) in half, then 
    sampling particles and tallying on the basis of which of the 2^4 = 8 regions
    of phase space the particle is born into. 
    """
    seed(1953)
    m = Mesh(structured=True, structured_coords=[[0, 1], [0, 1], [0, 1]], mats=None)
    m.src = IMeshTag(1, float)
    m.src[0] = 1.0
    m.mesh.save("sampling_mesh.h5m")
    sampler = Sampler("sampling_mesh.h5m", "src", np.array([0, 1]), False)

    num_samples = 5000
    score = 1.0 / num_samples
    num_divs = 2
    tally = np.zeros(shape=(num_divs, num_divs, num_divs, num_divs))

    for i in range(num_samples):
        s = sampler.particle_birth(np.array([uniform(0, 1) for x in range(6)]))
        assert_equal(s[4], 1.0)  # analog: all weights must be one
        tally[int(s[0] * num_divs), int(s[1] * num_divs), int(s[2] * num_divs), int(s[3] * num_divs)] += score

    # Test that each half-space of phase space (e.g. x > 0.5) is sampled about
    # half the time.
    for i in range(0, 4):
        for j in range(0, 2):
            assert abs(np.sum(np.rollaxis(tally, i)[j, :, :, :]) - 0.5) < 0.05

def test_analog_multiple_hex():
    """This test tests that particle are sampled uniformly from a uniform source
    defined on eight mesh volume elements in two energy groups. This is done
    using the exact same method ass test_analog_multiple_hex.
    """
    seed(1953)
    m = Mesh(structured=True, structured_coords=[[0, 0.5, 1], [0, 0.5, 1], [0, 0.5, 1]], mats=None)
    m.src = IMeshTag(2, float)
    m.src[:] = np.ones(shape=(8, 2))
    m.mesh.save("sampling_mesh.h5m")
    sampler = Sampler("sampling_mesh.h5m", "src", np.array([0, 0.5, 1]), False)

    num_samples = 5000
    score = 1.0 / num_samples
    num_divs = 2
    tally = np.zeros(shape=(num_divs, num_divs, num_divs, num_divs))
    for i in range(num_samples):
        s = sampler.particle_birth([uniform(0, 1) for x in range(6)])
        assert_equal(s[4], 1.0)
        tally[int(s[0] * num_divs), int(s[1] * num_divs), int(s[2] * num_divs), int(s[3] * num_divs)] += score

    for i in range(0, 4):
        for j in range(0, 2):
            halfspace_sum = np.sum(np.rollaxis(tally, i)[j, :, :, :])
            assert abs(halfspace_sum - 0.5) / 0.5 < 0.1

def test_vtx_iterator():
    #use vanilla izip as we"ll test using non-equal-length iterators
    izip = itertools.izip

    sm = Mesh(structured=True,
             structured_coords =[range(10,15), range(21,25), range(31,34)])
    it = sm.structured_set.iterate(iBase.Type.vertex, iMesh.Topology.point)

    # test the default order
    for (it_x, sm_x) in itertools.izip_longest(it, 
                                           sm.structured_iterate_vertex("zyx")):
        assert_equal(it_x,sm_x)

    # Do the same again, but use an arbitrary kwarg to structured_iterate_vertex 
    # to prevent optimization from kicking in
    it.reset()
    for (it_x, sm_x) in itertools.izip_longest(it,
                            sm.structured_iterate_vertex("zyx", no_opt=True)):
        assert_equal(it_x,sm_x)

    it.reset()
    for (it_x, sm_x) in izip(it, 
                               sm.structured_iterate_vertex("yx",z=sm.dims[2])):
        assert_equal(it_x,sm_x)

    it.reset()
    for (it_x, sm_x) in izip(it, sm.structured_iterate_vertex("x")):
        assert_equal(it_x,sm_x)

def test_photon_source_hdf5_to_mesh():
    """Tests the function photon source_h5_to_mesh."""

    if not HAVE_PYTAPS:
        raise SkipTest

    filename = os.path.join(thisdir, "files_test_alara", "phtn_src")
    photon_source_to_hdf5(filename, chunkshape=(10,))
    assert_true(os.path.exists(filename + ".h5"))

    mesh = Mesh(structured=True, structured_coords=[[0, 1, 2], [0, 1, 2], [0, 1]])

    tags = {("1001", "shutdown"): "tag1", ("TOTAL", "1 h"): "tag2"}
    photon_source_hdf5_to_mesh(mesh, filename + ".h5", tags)

    # create lists of lists of expected results
    tag1_answers = [[1] + [0] * 41, [2] + [0] * 41, [3] + [0] * 41, [4] + [0] * 41]
    tag2_answers = [[5] + [0] * 41, [6] + [0] * 41, [7] + [0] * 41, [8] + [0] * 41]

    ves = list(mesh.structured_iterate_hex("xyz"))
    for i, ve in enumerate(ves):
        assert_array_equal(mesh.mesh.getTagHandle("tag1")[ve], tag1_answers[i])
        assert_array_equal(mesh.mesh.getTagHandle("tag2")[ve], tag2_answers[i])

    if os.path.isfile(filename + ".h5"):
        os.remove(filename + ".h5")

def test_write_fluxin_single():
    """This function tests the flux_mesh_to_fluxin function for a single energy
    group case.
    """

    if not HAVE_PYTAPS:
        raise SkipTest

    output_name = "fluxin.out"
    forward_fluxin = os.path.join(thisdir, "files_test_alara", "fluxin_single_forward.txt")
    output = os.path.join(os.getcwd(), output_name)

    flux_mesh = Mesh(structured=True, structured_coords=[[0, 1, 2], [0, 1, 2], [0, 1]])
    tag_flux = flux_mesh.mesh.createTag("flux", 1, float)
    flux_data = [1, 2, 3, 4]
    ves = flux_mesh.structured_iterate_hex("xyz")
    for i, ve in enumerate(ves):
        tag_flux[ve] = flux_data[i]

    # test forward writting
    mesh_to_fluxin(flux_mesh, "flux", output_name, False)

    with open(output) as f:
        written = f.readlines()

    with open(forward_fluxin) as f:
        expected = f.readlines()

    assert_equal(written, expected)
    if os.path.isfile(output):
        os.remove(output)

def gen_mesh(mats=()):
    mesh_1 = Mesh(structured_coords=[[-1,0,1],[-1,0,1],[0,1]], structured=True, 
                  structured_ordering='zyx', mats=mats)
    volumes1 = list(mesh_1.structured_iterate_hex("xyz"))
    flux_tag = mesh_1.mesh.createTag("flux", 1, float)
    flux_data = [1.0, 2.0, 3.0, 4.0]
    flux_tag[volumes1] = flux_data
    return mesh_1

def test_get_divs():
    x = [1, 2.5, 4, 6.9]
    y = [-12, -10, -.5]
    z = [100, 200]

    sm = Mesh(structured_coords=[x, y, z], structured=True)

    assert_equal(sm.structured_get_divisions("x"), x)
    assert_equal(sm.structured_get_divisions("y"), y)
    assert_equal(sm.structured_get_divisions("z"), z)

def test_iterate_3d():        
    # use izip_longest in the lockstep iterations below; this will catch any
    # situations where one iterator turns out to be longer than expected.
    sm = Mesh(structured=True,
             structured_coords =[range(10,15), range(21,25), range(31,34)])
    I = range(0,4)
    J = range(0,3)
    K = range(0,2)
    izip = itertools.izip_longest

    it = sm.structured_set.iterate(iBase.Type.region, 
                                 iMesh.Topology.hexahedron)

    # Test the zyx order, which is default; it should be equivalent
    # to the standard imesh iterator
    for it_x, sm_x in izip(it, sm.structured_iterate_hex()):
        assert_equal(it_x, sm_x)

    #testing xyz

    all_indices_zyx = itertools.product(I, J, K)
    # Test the xyz order, the default from original mmGridGen
    for ijk_index, sm_x in izip(all_indices_zyx, 
                                 sm.structured_iterate_hex("xyz")):
        assert_equal(sm.structured_get_hex(*ijk_index), sm_x )

    def _tuple_sort(collection, indices ):
        # sorting function for order test
        def t(tup):
            # sort this 3-tuple according to the order of x, y, and z in 
            #indices
            return (tup["xyz".find(indices[0])]*100 +
                    tup["xyz".find(indices[1])]*10 +
                    tup["xyz".find(indices[2])])
        return sorted(collection, key = t)

    def test_order(order, *args,  **kw):
        all_indices = itertools.product(*args)
        for ijk_index, sm_x in izip(_tuple_sort(all_indices, order),
                                     sm.structured_iterate_hex(order,**kw)):
            assert_equal(sm.structured_get_hex(*ijk_index), sm_x)

    test_order("yxz", I, J, K)
    test_order("yzx", I, J, K)
    test_order("xzy", I, J, K)
    test_order("zxy", I, J, K)

    # Specify z=[1] to iterator
    test_order("xyz", I, J, [1], z=[1])
    # Specify y=2 to iterator
    test_order("zyx", I, [2], K, y=2)
    # specify x and y both to iterator
    test_order("yzx", [1,2,3],J[:-1], K, y=J[:-1], x=[1,2,3])

 def arithmetic_mesh_setup(self):
     self.mesh_1 = Mesh(structured_coords=[[-1,0,1],[-1,0,1],[0,1]], structured=True)
     volumes1 = list(self.mesh_1.structured_iterate_hex("xyz"))
     volumes2 = list(self.mesh_1.structured_iterate_hex("xyz"))
     flux_tag = self.mesh_1.mesh.createTag("flux", 1, float)
     flux_data = [1.0, 2.0, 3.0, 4.0]
     flux_tag[volumes1] = flux_data    
 
     self.mesh_2 = Mesh(structured_coords=[[-1,0,1],[-1,0,1],[0,1]], structured=True)
     volumes1 = list(self.mesh_2.structured_iterate_hex("xyz"))
     volumes2 = list(self.mesh_2.structured_iterate_hex("xyz"))
     flux_tag = self.mesh_2.mesh.createTag("flux", 1, float)
     flux_data = [1.1, 2.2, 3.3, 4.4]
     flux_tag[volumes1] = flux_data

def test_bias_spatial():
    """This test tests a user-specified biasing scheme for which the only 1
    bias group is supplied for a source distribution containing two energy 
    groups. This bias group is applied to both energy groups. In this test,
    the user-supplied bias distribution that was choosen, correspondes to 
    uniform sampling, so that results can be checked against Case 1 in the
    theory manual.
    """
    seed(1953)
    m = Mesh(structured=True, structured_coords=[[0, 3, 3.5], [0, 1], [0, 1]], mats=None)
    m.src = IMeshTag(2, float)
    m.src[:] = [[2.0, 1.0], [9.0, 3.0]]
    m.bias = IMeshTag(1, float)
    m.bias[:] = [1, 1]
    e_bounds = np.array([0, 0.5, 1.0])
    m.mesh.save("sampling_mesh.h5m")
    sampler = Sampler("sampling_mesh.h5m", "src", e_bounds, "bias")

    num_samples = 10000
    score = 1.0 / num_samples
    num_divs = 2
    num_e = 2
    spatial_tally = np.zeros(shape=(num_divs, num_divs, num_divs))
    e_tally = np.zeros(shape=(4))  # number of phase space groups
    for i in range(num_samples):
        s = sampler.particle_birth(np.array([uniform(0, 1) for x in range(6)]))
        if s[0] < 3.0:
            assert_almost_equal(s[4], 0.7)  # hand calcs
        else:
            assert_almost_equal(s[4], 2.8)  # hand calcs

        spatial_tally[int(s[0] * num_divs / 3.5), int(s[1] * num_divs / 1.0), int(s[2] * num_divs / 1.0)] += score

        if s[0] < 3 and s[3] < 0.5:
            e_tally[0] += score
        elif s[0] < 3 and s[3] > 0.5:
            e_tally[1] += score
        if s[0] > 3 and s[3] < 0.5:
            e_tally[2] += score
        if s[0] > 3 and s[3] > 0.5:
            e_tally[3] += score

    for i in range(0, 3):
        for j in range(0, 2):
            halfspace_sum = np.sum(np.rollaxis(spatial_tally, i)[j, :, :])
            assert abs(halfspace_sum - 0.5) / 0.5 < 0.1

    expected_e_tally = [4.0 / 7, 2.0 / 7, 3.0 / 28, 1.0 / 28]  # hand calcs
    for i in range(4):
        assert abs(e_tally[i] - expected_e_tally[i]) / expected_e_tally[i] < 0.1

def test_structured_get_hex():
    # mesh with valid i values 0-4, j values 0-3, k values 0-2
    sm = Mesh(structured_coords = [range(11,16), range(21,25), range(31,34)], 
              structured=True)
    def check(e):
        assert_true(isinstance(e, iBase.Entity))
    check(sm.structured_get_hex(0, 0, 0))
    check(sm.structured_get_hex(1, 1, 1))
    check(sm.structured_get_hex(3, 0, 0))
    check(sm.structured_get_hex(3, 2, 1))

    assert_raises(MeshError, sm.structured_get_hex,-1,-1,-1)
    assert_raises(MeshError, sm.structured_get_hex, 4, 0, 0)
    assert_raises(MeshError, sm.structured_get_hex, 0, 3, 0)
    assert_raises(MeshError, sm.structured_get_hex, 0, 0, 2)

def test_structured_get_vertex():
    # mesh with valid i values 0-4, j values 0-3, k values 0-2
    x_range = np.array(range(10,15),dtype=np.float64)
    y_range = np.array(range(21,24),dtype=np.float64)
    z_range = np.array(range(31,33),dtype=np.float64)

    sm = Mesh(structured_coords=[x_range, y_range, z_range], structured=True)

    for i,x in enumerate(x_range):
        for j,y in enumerate(y_range):
            for k,z in enumerate(z_range):
                print("{0} {1} {2}".format(i, j, k))
                vtx = sm.structured_get_vertex(i,j,k)
                vcoord = sm.mesh.getVtxCoords(vtx)
                assert_true(all(vcoord == [x,y,z]))

def test_uniform():
    """This test tests that the uniform biasing scheme:
    1. Samples space uniformly. This is checked using the same method
       described in test_analog_single_hex().
    2. Adjusts weights accordingly. Sample calculations are provided in Case 1
       in the Theory Manual.
    """
    seed(1953)
    m = Mesh(structured=True, structured_coords=[[0, 3, 3.5], [0, 1], [0, 1]], mats=None)
    m.src = IMeshTag(2, float)
    m.src[:] = [[2.0, 1.0], [9.0, 3.0]]
    e_bounds = np.array([0, 0.5, 1.0])
    m.mesh.save("sampling_mesh.h5m")
    sampler = Sampler("sampling_mesh.h5m", "src", e_bounds, True)

    num_samples = 10000
    score = 1.0 / num_samples
    num_divs = 2
    num_e = 2
    spatial_tally = np.zeros(shape=(num_divs, num_divs, num_divs))
    e_tally = np.zeros(shape=(4))  # number of phase space groups
    for i in range(num_samples):
        s = sampler.particle_birth(np.array([uniform(0, 1) for x in range(6)]))
        if s[0] < 3.0:
            assert_almost_equal(s[4], 0.7)  # hand calcs
        else:
            assert_almost_equal(s[4], 2.8)  # hand calcs

        spatial_tally[int(s[0] * num_divs / 3.5), int(s[1] * num_divs / 1.0), int(s[2] * num_divs / 1.0)] += score

        if s[0] < 3 and s[3] < 0.5:
            e_tally[0] += score
        elif s[0] < 3 and s[3] > 0.5:
            e_tally[1] += score
        if s[0] > 3 and s[3] < 0.5:
            e_tally[2] += score
        if s[0] > 3 and s[3] > 0.5:
            e_tally[3] += score

    for i in range(0, 3):
        for j in range(0, 2):
            halfspace_sum = np.sum(np.rollaxis(spatial_tally, i)[j, :, :])
            assert abs(halfspace_sum - 0.5) / 0.5 < 0.1

    expected_e_tally = [4.0 / 7, 2.0 / 7, 3.0 / 28, 1.0 / 28]  # hand calcs
    for i in range(4):
        assert abs(e_tally[i] - expected_e_tally[i]) / expected_e_tally[i] < 0.1

def test_photon_sampling_setup_structured():

    phtn_src = os.path.join(thisdir, "files_test_r2s", "phtn_src")
    coords = [[0, 1, 2], [0, 1, 2], [0, 1]]
    m = Mesh(structured=True, structured_coords=coords)
    tags = {(10010000, "1 h"): "tag1", ("TOTAL", "shutdown"): "tag2"}
    photon_sampling_setup(m, phtn_src, tags)

    exp_tag1 = [[1.1, 2.2], [3.3, 4.4], [5.5, 6.6], [7.7, 8.8]]
    exp_tag2 = [[11.1, 12.2], [13.3, 14.4], [15.5, 16.6], [17.7, 18.8]]

    m.tag1 = IMeshTag(2, float)
    m.tag2 = IMeshTag(2, float)

    for i, mat, ve in m:
        assert_array_equal(m.tag1[i], exp_tag1[i])
        assert_array_equal(m.tag2[i], exp_tag2[i])

def test_magic_single_e():
    """Test a single energy group MAGIC case"""
    
    # create mesh
    coords = [[0, 1, 2], [-1, 3, 4], [10, 12]]
    flux_data = [1.2, 3.3, 1.6, 1.7]
    flux_error = [0.11, 0.013, 0.14, 0.19]
    tally = Mesh(structured=True, structured_coords=coords)
    
    tally.particle = "neutron"
    tally.e_bounds = [0.0, 1.0]
    tally.n_flux = IMeshTag(1, float)
    tally.n_flux[:] = flux_data
    
    tally.n_rel_error = IMeshTag(1, float)
    tally.n_rel_error[:] = flux_error
    
    tolerance = 0.15
    null_value = 0.001
    
    magic(tally, "n_flux", "n_rel_error", tolerance=tolerance, null_value=null_value)
    
    expected_ww = [0.181818182, 0.5, 0.2424242, 0.001]
    
    assert_array_almost_equal(tally.ww_x[:], expected_ww[:])

def test_magic_multi_bins():
    """Test multiple energy bins MAGIC case"""

    # create a basic meshtally
    coords = [[0, 1, 2], [-1, 3, 4], [10, 12]]
    flux_data = [[1.2, 3.3], [1.6, 1.7], [1.5, 1.4], [2.6, 1.0]]
    flux_error = [[0.11, 0.013], [0.14, 0.19], [0.02, 0.16], [0.04, 0.09]]
    tally = Mesh(structured=True, structured_coords=coords)
    
    tally.particle = "neutron"
    tally.e_bounds = [0.0, 0.5, 1.0]
    tally.n_flux = IMeshTag(2, float)
    tally.n_flux[:] = flux_data
    
    tally.n_rel_error = IMeshTag(2, float)
    tally.n_rel_error[:] = flux_error
    
    tolerance = 0.15
    null_value = 0.001
    
    magic(tally, "n_flux", "n_rel_error", tolerance=tolerance, null_value=null_value)
    
    expected_ww = [[0.2307692308, 0.5], 
                   [0.3076923077, 0.001], 
                   [0.2884615385, 0.001], 
                   [0.5, 0.15151515]]
    
    assert_array_almost_equal(tally.ww_x[:], expected_ww[:])

def test_iterate_1d():
    sm = Mesh(structured=True,
             structured_coords =[range(10,15), range(21,25), range(31,34)])
    
    def test_equal(ijk_list, miter):
        for ijk, i in itertools.izip_longest(ijk_list, miter):
            assert_equal(sm.structured_get_hex(*ijk), i)

    test_equal([[0,0,0],[0,0,1]], 
                sm.structured_iterate_hex("z"))

    test_equal([[0,1,1],[0,2,1]],
                sm.structured_iterate_hex("y", y=[1,2], z=1))

    test_equal([[2,0,0],[2,1,0],[2,2,0]],
                sm.structured_iterate_hex("y", x=2)) 
    test_equal([[0,0,0],[1,0,0],[2,0,0]], 
        sm.structured_iterate_hex("x", x=[0,1,2]))

def discretize_geom_centers():
    from pyne import dagmc
    dagmc.load(path)

    coords = [0, 1]
    coords2 = [0, 2, 4]
    mesh = Mesh(structured=True, 
                structured_coords=[coords2, coords, coords])
    #  explicitly set to unstructured to trigger the ve centers sampling
    #  method
    mesh.structured = False
    res = dagmc.discretize_geom(mesh)
    assert_array_equal(res["idx"], [0, 1])
    assert_array_equal(res["cell"], [2, 3])
    assert_array_equal(res["vol_frac"], [1.0, 1.0])
    assert_array_equal(res["rel_error"], [1.0, 1.0])

    #  ensure kwargs are not accepted for unstructured mesh
    assert_raises(ValueError, dagmc.discretize_geom, mesh, num_rays=3, grid=True)

def test_imeshtag_expand():
    m = Mesh(structured=True, structured_coords=[[-1, 0, 1],[0, 1],[0, 1]])
    m.clam = IMeshTag(2, float)
    m.clam[:] = [[1.1, 2.2], [3.3, 4.4]]
    m.clam.expand()
    m.clam_000 = IMeshTag(1, float)
    assert_array_equal(m.clam_000[:], [1.1, 3.3])
    m.clam_001 = IMeshTag(1, float)
    assert_array_equal(m.clam_001[:], [2.2, 4.4])

    # corner case: mesh with a single volume element
    m = Mesh(structured=True, structured_coords=[[0, 1],[0, 1],[0, 1]])
    m.clam = IMeshTag(2, float)
    m.clam[:] = [[1.1, 2.2]]
    m.clam.expand()
    m.clam_000 = IMeshTag(1, float)
    assert_array_equal(m.clam_000[:], 1.1)
    m.clam_001 = IMeshTag(1, float)
    assert_array_equal(m.clam_001[:], 2.2)