content: # -*- coding: utf-8 -*-
  
  import pytest
  import collections
  import numpy as np
  import qutip
  from qutip.core.dimensions import (
      flatten, unflatten, enumerate_flat, deep_remove, deep_map,
      dims_idxs_to_tensor_idxs, dims_to_tensor_shape, dims_to_tensor_perm,
      einsum, to_tensor_rep, from_tensor_rep,
      Dimensions, Field, Space, SuperSpace
  )
  from qutip.core.energy_restricted import EnrSpace
  
  
  @pytest.mark.parametrize(["base", "flat"], [
      pytest.param([[[0], 1], 2], [0, 1, 2], id="standard"),
      pytest.param(SuperSpace(Dimensions(Space(Space(2), Space(3)), Space(1))),
                   [1, 2, 3], id="space"),
      pytest.param([1, 2, [3, [4]], [5, 6], [7, [[[[[[[8]]]]]]]]],
                   [1, 2, 3, 4, 5, 6, 7, 8], id="deep nested"),
      pytest.param([1, 2, 3], [1, 2, 3], id="already flat"),
      pytest.param([], [], id="empty list"),
      pytest.param([[], [], [[[], [], []]]], [], id="nested empty lists"),
  ])
  class TestFlattenUnflatten:
      def test_flatten(self, base, flat):
          assert flatten(base) == flat
  
      def test_unflatten(self, base, flat):
          labels = enumerate_flat(base)
          expected = base.as_list() if isinstance(base, Space) else base
          assert unflatten(flat, labels) == expected
  
  
  @pytest.mark.parametrize(["base", "expected"], [
      pytest.param([1, 2, 3], [0, 1, 2], id="flat"),
      pytest.param([[1], [2], [3]], [[0], [1], [2]], id="nested"),
      pytest.param([[[1], [2, 3]], 4], [[[0], [1, 2]], 3], id="nested"),
      pytest.param([], [], id="empty"),
      pytest.param([[], [], [[[], [], []]]],
                   [[], [], [[[], [], []]]], id="nested empty lists"),
  ])
  def test_enumerate_flat(base, expected):
      assert enumerate_flat(base) == expected
  
  
  @pytest.mark.parametrize(["base", "to_remove", "expected"], [
      pytest.param([[[0], 1], 2], (1,), [[[0]], 2], id="simple"),
      pytest.param([[[[0, 1, 2]], [3, 4], [5], [6, 7]]], (0, 5),
                   [[[[1, 2]], [3, 4], [], [6, 7]]], id="harder"),
      pytest.param([1, 2, 3], (), [1, 2, 3], id="no-op"),
      pytest.param([], (), [], id="empty"),
  ])
  def test_deep_remove(base, to_remove, expected):
      assert deep_remove(base, *to_remove) == expected
  
  
  @pytest.mark.parametrize("mapping", [
      pytest.param(lambda x: x*2, id="(x -> 2x)"),
      pytest.param(lambda x: [x], id="(x -> [x])"),
  ])
  @pytest.mark.parametrize("base", [
      pytest.param([[[0], 1], 2], id="standard"),
      pytest.param([1, 2, [3, [4]], [7, [[[[[[[8]]]]]]]]], id="standard"),
      pytest.param([1, 2, 3], id="flat"),
      pytest.param([], id="empty list"),
      pytest.param([[], [], [[[], [], []]]], id="nested empty lists"),
  ])
  def test_deep_map(base, mapping):
      """
      Test the deep mapping.  To simplify generation of edge-cases, this tests
      against an equivalent (but slower) operation of flattening and unflattening
      the list.  We can get false negatives if the `flatten` or `unflatten`
      functions are broken, but other tests should catch those.
      """
      # This function might not need to be public, and consequently might not
      # need to be tested here.
      labels = enumerate_flat(base)
      expected = unflatten([mapping(x) for x in flatten(base)], labels)
      assert deep_map(mapping, base) == expected
  
  
  _Indices = collections.namedtuple('_Indices', ['base', 'permutation', 'shape'])
  
  
  @pytest.mark.parametrize("indices", [
      pytest.param(_Indices([[2], [1]], [0, 1], (2, 1)), id="ket preserved"),
      pytest.param(_Indices([[2, 3], [1]], [0, 1, 2], (2, 3, 1)),
                   id="tensor-ket preserved"),
      pytest.param(_Indices([[1], [2]], [0, 1], (1, 2)), id="bra preserved"),
      pytest.param(_Indices([[1], [2, 2]], [0, 1, 2], (1, 2, 2)),
                   id="tensor-bra preserved"),
      pytest.param(_Indices([[2], [3]], [0, 1], (2, 3)), id="oper preserved"),
      pytest.param(_Indices([[2, 3], [1, 2]], [0, 1, 2, 3], (2, 3, 1, 2)),
                   id="tensor-oper preserved"),
      pytest.param(_Indices([[[2, 4], [6, 8]], [[1, 3], [5, 7]]],
                            [2, 3, 0, 1, 6, 7, 4, 5],
                            (6, 8, 2, 4, 5, 7, 1, 3)),
                   id="super-oper"),
      pytest.param(_Indices([[[1], [2, 3]], [[4, 5], [6]]],
                            [2, 0, 1, 4, 5, 3],
                            (2, 3, 1, 6, 4, 5)),
                   id="super-oper-uneven"),
      pytest.param(_Indices([[[2, 4], [6, 8]], [1]],
                            [2, 3, 0, 1, 4], (6, 8, 2, 4, 1)),
                   id="operator-ket"),
      pytest.param(_Indices([[1], [[2, 4], [6, 8]]],
                            [0, 3, 4, 1, 2], (1, 6, 8, 2, 4)),
                   id="operator-bra"),
  ])
  class TestSuperOperatorDimsModification:
      def test_dims_to_tensor_perm(self, indices):
          # This function might not need to be public, and consequently might not
          # need to be tested here.
          assert dims_to_tensor_perm(indices.base) == indices.permutation
  
      def test_dims_idxs_to_tensor_idxs(self, indices):
          test_indices = list(range(len(flatten(indices.base))))
          assert (dims_idxs_to_tensor_idxs(indices.base, test_indices)
                  == indices.permutation)
  
      def test_dims_to_tensor_shape(self, indices):
          assert dims_to_tensor_shape(indices.base) == indices.shape
  
  
  @pytest.mark.parametrize("dims", [
      pytest.param([[2, 3], [1]], id="ket"),
      pytest.param([[1], [2, 3]], id="bra"),
      pytest.param([[2, 3], [5, 7]], id="oper"),
      pytest.param([[[2], [3]], [[5], [7]]], id="super-oper"),
      pytest.param([[[2], [3]], [1]], id="oper-ket"),
      pytest.param([[1], [[2, 3], [5, 7]]], id="oper-bra"),
      pytest.param([[[2], [3, 5]], [[7], [11]]], id="super-oper-uneven"),
  ])
  class TestToTensorRep:
      class mylist(list):
          def at(self, i, j):
              out = self.copy()
              out[i] = j
              return tuple(out)
  
      def _build_qobj(self, dims, idx=-1):
          """
          Create an Qobj with given dims which values only increment for the
          `idx` dims.
          """
          if isinstance(dims[0], int):
              ket = qutip.tensor([
                  (
                      qutip.Qobj(np.arange(N))
                      if i == idx else (qutip.Qobj(np.ones(N)))
                  ) for i, N in enumerate(dims)
              ])
              ket.dims = [ket.dims[0], [1]]
              return ket
  
          left = self._build_qobj(dims[0], idx)
          right = self._build_qobj(dims[1], idx - len(flatten(dims[0])))
          if left.isoper and left.type != 'scalar':
              left = qutip.operator_to_vector(left)
              left.dims = [left.dims[0], [1]]
          if right.isoper and right.type != 'scalar':
              right = qutip.operator_to_vector(right)
              left.dims = [left.dims[0], [1]]
          return left @ right.dag()
  
      def test_tensor_shape(self, dims):
          qobj = self._build_qobj(dims)
          array = to_tensor_rep(qobj)
          assert array.shape == tuple(flatten(dims))
  
      def test_tensor_rep(self, dims):
          N = len(flatten(dims))
          for i in range(N):
              qobj = self._build_qobj(dims, i)
              array = to_tensor_rep(qobj)
              slices = self.mylist([slice(j) for j in array.shape])
              assert all(
                  np.all(array.__getitem__(slices.at(i, j)) == j)
                  for j in range(array.shape[i])
              )
  
      def from_tensor_rep(self, dims):
          qobj = self._build_qobj(dims)
          array = to_tensor_rep(qobj)
          back = from_tensor_rep(array, dims)
          assert qobj == back
  
  
  class TestTypeFromDims:
      @pytest.mark.parametrize(["base", "expected"], [
          pytest.param([[2], [2]], 'oper'),
          pytest.param([[2, 3], [2, 3]], 'oper'),
          pytest.param([[2], [3]], 'oper'),
          pytest.param([[2], [1]], 'ket'),
          pytest.param([[1], [2]], 'bra'),
          pytest.param([[[2, 3], [2, 3]], [1]], 'operator-ket'),
          pytest.param([[1], [[2, 3], [2, 3]]], 'operator-bra'),
          pytest.param([[[3], [3]], [[2, 3], [2, 3]]], 'super'),
      ])
      def test_Dimensions_type(self, base, expected):
          assert Dimensions(base).type == expected
  
  
  class TestCollapseDims:
      @pytest.mark.parametrize(["base", "expected"], [
          pytest.param([[1], [3]], [[1], [3]], id="ket trivial"),
          pytest.param([[1, 1], [2, 3]], [[1], [6]], id="ket tensor"),
          pytest.param(Dimensions(Field(), EnrSpace([2, 2], 1)),
                       Dimensions(Field(), Space(3)), id="ket ENR"),
          pytest.param([[2], [1]], [[2], [1]], id="bra trivial"),
          pytest.param([[2, 3], [1, 1]], [[6], [1]], id="bra tensor"),
          pytest.param([[5], [5]], [[5], [5]], id="oper trivial"),
          pytest.param([[2, 3], [2, 3]], [[6], [6]], id="oper tensor"),
      ])
      def test_oper(self, base, expected):
          assert Dimensions(base).collapse() == Dimensions(expected)
  
      @pytest.mark.parametrize(["base", "expected"], [
          pytest.param([[[1]], [[2, 3], [2, 3]]],
                       [[[1]], [[6], [6]]], id="operator-ket"),
          pytest.param([[[2, 3], [2, 3]], [[1]]],
                       [[[6], [6]], [[1]]], id="operator-bra"),
          pytest.param([[[2, 3], [2, 3]], [[2, 3], [2, 3]]],
                       [[[6], [6]], [[6], [6]]], id="super"),
      ])
      def test_super(self, base, expected):
          assert Dimensions(base).collapse() == Dimensions(expected)
  
  
  @pytest.mark.parametrize("dims_list", [
      pytest.param([0], id="zero"),
      pytest.param([], id="empty"),
      pytest.param([1, [2]], id="mixed depth"),
      pytest.param([[2], [3], [4]], id="bay type"),
  ])
  def test_bad_dims(dims_list):
      with pytest.raises(ValueError):
          Dimensions([dims_list, [1]])
  
  
  @pytest.mark.parametrize("space_l", [[1], [2], [2, 3]])
  @pytest.mark.parametrize("space_m", [[1], [2], [2, 3]])
  @pytest.mark.parametrize("space_r", [[1], [2], [2, 3]])
  def test_dims_matmul(space_l, space_m, space_r):
      dims_l = Dimensions([space_l, space_m])
      dims_r = Dimensions([space_m, space_r])
      assert dims_l @ dims_r == Dimensions([space_l, space_r])
  
  
  def test_dims_matmul_bad():
      dims_l = Dimensions([[1], [3]])
      dims_r = Dimensions([[2], [2]])
      with pytest.raises(TypeError):
          dims_l @ dims_r
  
  
  def test_dims_comparison():
      assert Dimensions([[1], [2]]) == Dimensions([[1], [2]])
      assert not Dimensions([[1], [2]]) != Dimensions([[1], [2]])
      assert Dimensions([[1], [2]]) != Dimensions([[2], [1]])
      assert not Dimensions([[1], [2]]) == Dimensions([[2], [1]])
      assert Dimensions([[1], [2]])[1] == Dimensions([[1], [2]])[1]
      assert Dimensions([[1], [2]])[0] != Dimensions([[1], [2]])[1]
      assert not Dimensions([[1], [2]])[1] != Dimensions([[1], [2]])[1]
      assert not Dimensions([[1], [2]])[0] != Dimensions([[1], [2]])[0]
  
  
  @pytest.mark.parametrize(["subscripts", "operands", "expected"], [
      pytest.param("ii", [qutip.sigmaz()], 0),
      pytest.param("ij", [qutip.sigmax()], qutip.sigmax()),
      pytest.param("ij->ji", [qutip.sigmay()], qutip.sigmay().trans()),
      pytest.param("ij,ji", [qutip.sigmaz(), qutip.sigmaz()], 2),
      pytest.param("ijij", [qutip.tensor(qutip.thermal_dm(2,1), qutip.thermal_dm(2,1))], 1),
      pytest.param("ikjl,jm->ikml", [qutip.tensor(qutip.sigmaz(), qutip.sigmaz()),
                               qutip.sigmaz()], qutip.tensor(qutip.qeye(2), qutip.sigmaz())),
      pytest.param("ijkl->kjil", [qutip.tensor(qutip.sigmam(), qutip.sigmaz())], qutip.tensor(qutip.sigmap(), qutip.sigmaz()))
  ])
  def test_einsum(subscripts, operands, expected):
      assert einsum(subscripts, *operands) == expected
  
  
  @pytest.mark.parametrize(["list_dims", "expected"], [
      pytest.param([[4, 4], [1, 1, 1]], [[4, 4], [1]]),
      pytest.param([[1, 1], [1, 1, 1]], [[1], [1]]),
      pytest.param([[1, 1, 1, 1], [5, 5, 2, 6]], [[1], [5, 5, 2, 6]]),
      pytest.param(
          [[[2, 3, 5], [2, 3, 5]], [[1, 1, 1], [1, 1, 1]]],
          [[[2, 3, 5], [2, 3, 5]], [1]]
      )
  ])
  def test_scalar_dims(list_dims, expected):
      with qutip.CoreOptions(auto_tidyup_dims=True):
          assert Dimensions(list_dims).as_list() == expected
      with qutip.CoreOptions(auto_tidyup_dims=False):
          assert Dimensions(list_dims).as_list() == list_dims
  