"""The medical image object.

This module defines :class:`MedicalVolume`, which is a wrapper for nD volumes.
"""
import warnings
from copy import deepcopy
from mmap import mmap
from numbers import Number
from typing import Sequence, Tuple, Union

import nibabel as nib
import numpy as np
import pydicom
from nibabel.spatialimages import SpatialFirstSlicer as _SpatialFirstSlicerNib
from numpy.lib.mixins import NDArrayOperatorsMixin
from packaging import version

from dosma.core import orientation as stdo
from dosma.core.device import Device, cpu_device, get_array_module, get_device, to_device
from dosma.core.io.format_io import ImageDataFormat
from dosma.defaults import SCANNER_ORIGIN_DECIMAL_PRECISION
from dosma.utils import env

if env.sitk_available():
    import SimpleITK as sitk
if env.cupy_available():
    import cupy as cp
if env.package_available("h5py"):
    import h5py

__all__ = ["MedicalVolume"]

# PyTorch version introducing complex tensor support.
_TORCH_COMPLEX_SUPPORT_VERSION = version.Version("1.5.0")

class MedicalVolume(NDArrayOperatorsMixin):
    """The class for medical images.

    Medical volumes use ndarrays to represent medical data. However, unlike standard ndarrays,
    these volumes have inherent spatial metadata, such as pixel/voxel spacing, global coordinates,
    rotation information, all of which can be characterized by an affine matrix following the
    RAS+ coordinate system. The code below creates a random 300x300x40 medical volume with
    scanner origin ``(0, 0, 0)`` and voxel spacing of ``(1,1,1)``:

    >>> mv = MedicalVolume(np.random.rand(300, 300, 40), np.eye(4))

    Medical volumes can also store header information that accompanies pixel data
    (e.g. DICOM headers). These headers are used to expose metadata, which can be fetched
    and set using :meth:`get_metadata()` and :meth:`set_metadata()`, respectively. Headers are
    also auto-aligned, which means that headers will be aligned with the slice(s) of data from
    which they originated, which makes Python slicing feasible. Currently, medical volumes
    support DICOM headers using ``pydicom`` when loaded with :class:`dosma.DicomReader`.

    >>> mv.get_metadata("EchoTime")  # Returns EchoTime
    >>> mv.set_metadata("EchoTime", 10.0)  # Sets EchoTime to 10.0

    Standard math and boolean operations are supported with other ``MedicalVolume`` objects,
    numpy arrays (following standard broadcasting), and scalars. Boolean operations are performed
    elementwise, resulting in a volume with shape as ``self.volume.shape``.
    If performing operations between ``MedicalVolume`` objects, both objects must have
    the same shape and affine matrix (spacing, direction, and origin). Header information
    is not deep copied when performing these operations to reduce computational and memory
    overhead. The affine matrix (``self.affine``) is copied as it is lightweight and
    often modified.

    2D images are also supported when viewed trivial 3D volumes with shape ``(H, W, 1)``:

    >>> mv = MedicalVolume(np.random.rand(10,20,1), np.eye(4))

    Many operations are in-place and modify the instance directly (e.g. `reformat(inplace=True)`).
    To allow chaining operations, operations that are in-place return ``self``.

    >>> mv2 = mv.reformat(ornt, inplace=True)
    >>> id(mv2) == id(mv)
    True

    Medical volumes can interface with the gpu using the :mod:`cupy` library.
    Volumes can be moved between devices (see :class:`Device`) using the ``.to()`` method.
    Only the volume data will be moved to the gpu. Headers and affine matrix will remain on
    the cpu. The following code moves a MedicalVolume to gpu 0 and back to the cpu:

    >>> from dosma import Device
    >>> mv = MedicalVolume(np.random.rand((10,20,30)), np.eye(4))
    >>> mv_gpu = mv.to(Device(0))
    >>> mv_cpu = mv.cpu()

    Note, moving data across devices results in a full copy. Above, ``mv_cpu.volume`` and
    ``mv.volume`` do not share memory. Saving volumes and converting to other images
    (e.g. ``SimpleITK.Image``) are only supported for cpu volumes. Volumes can also only
    be compared when on the same device. For example, both commands below will raise a
    RuntimeError:

    >>> mv_gpu == mv_cpu
    >>> mv_gpu.is_identical(mv_cpu)

    While CuPy requires the current device be set using ``cp.cuda.Device(X).use()`` or inside
    the ``with`` context, ``MedicalVolume`` automatically sets the appropriate context
    for performing operations. This means the CuPy current device need to be the same as the
    ``MedicalVolume`` object. For example, the following still works:

    >>> cp.cuda.Device(0).use()
    >>> mv_gpu = MedicalVolume(cp.ones((3,3,3)), np.eye(4))
    >>> cp.cuda.Device(1).use()
    >>> mv_gpu *= 2

    MedicalVolumes also have a limited NumPy/CuPy-compatible interface.
    Standard numpy/cupy functions that preserve array shapes can be performed
    on MedicalVolume objects:

    >>> log_arr = np.log(mv)
    >>> type(log_arr)
    <class 'dosma.io.MedicalVolume'>
    >>> exp_arr_gpu = cp.exp(mv_gpu)
    >>> type(exp_arr_gpu)
    <class 'dosma.io.MedicalVolume'>

    **ALPHA**: MedicalVolumes are also interoperable with popular image data structures
    with zero-copy, meaning array data will not be copied. Formats currently include the
    SimpleITK Image, Nibabel Nifti1Image, and PyTorch tensors:

    >>> sitk_img = mv.to_sitk()  # Convert to SimpleITK Image
    >>> mv_from_sitk = MedicalVolume.from_sitk(sitk_img)  # Convert back to MedicalVolume
    >>> nib_img = mv.to_nib()  # Convert to nibabel Nifti1Image
    >>> mv_from_nib = MedicalVolume.from_nib(nib_img)
    >>> torch_tensor = mv.to_torch()  # Convert to torch tensor
    >>> mv_from_tensor = MedicalVolume.from_torch(torch_tensor, affine)

    **ALPHA**: MedicalVolumes can also be used with memmapped arrays.
    This makes loading much faster and allows interaction with larger-than-memory
    arrays. Only when the volume is modified will the volume be loaded
    into memory and modified. If you take a slice of the memmaped array, the underlying
    array will also remain memmapped:

    >>> arr = np.load("/path/to/volume.npy", mmap_mode="r")
    >>> mv = MedicalVolume(arr, np.eye(4))
    >>> mv.is_mmap  # returns True

    We also preserve Nibabel's memmapping of certain file types (e.g. ``.nii``):

    >>> nib_img = nibabel.load("path/to/volume.nii")
    >>> mv = MedicalVolume.from_nib(nib_img, mmap=True)

    Args:
        volume (array-like): nD medical image.
        affine (array-like): 4x4 array corresponding to affine matrix transform in RAS+ coordinates.
            Must be on cpu (i.e. no ``cupy.ndarray``).
        headers (array-like[pydicom.FileDataset]): Headers for DICOM files.
    """

    def __init__(self, volume, affine, headers=None):
        if not isinstance(volume, np.memmap):
            xp = get_array_module(volume)
            volume = xp.asarray(volume)
        self._volume = volume
        self._affine = np.array(affine)
        self._headers = self._validate_and_format_headers(headers) if headers is not None else None

    def save_volume(self, file_path: str, data_format: ImageDataFormat = ImageDataFormat.nifti):
        """Write volumes in specified data format.

        Args:
            file_path (str): File path to save data. May be modified to follow convention
                given by the data format in which the volume will be saved.
            data_format (ImageDataFormat): Format to save data.
        """
        import dosma.core.io.format_io_utils

        device = self.device
        if device != cpu_device:
            raise RuntimeError(f"MedicalVolume must be on cpu, got {self.device}")

        writer = dosma.core.io.format_io_utils.get_writer(data_format)
        writer.save(self, file_path)

    def reformat(self, new_orientation: Sequence, inplace: bool = False) -> "MedicalVolume":
        """Reorients volume to a specified orientation.

        Flipping and transposing the volume array (``self.volume``) returns a view if possible.

        Reorientation method:
        ---------------------
        - Axis transpose and flipping are linear operations and therefore can be treated
        independently.
        - working example: ('AP', 'SI', 'LR') --> ('RL', 'PA', 'SI')
        1. Transpose volume and RAS orientation to appropriate column in matrix
            eg. ('AP', 'SI', 'LR') --> ('LR', 'AP', 'SI') - transpose_inds=[2, 0, 1]
        2. Flip volume across corresponding axes
            eg. ('LR', 'AP', 'SI') --> ('RL', 'PA', 'SI') - flip axes 0,1

        Reorientation method implementation:
        ------------------------------------
        1. Transpose: Switching (transposing) axes in volume is the same as switching columns
        in affine matrix

        2. Flipping: Negate each column corresponding to pixel axis to flip (i, j, k) and
        reestablish origins based on flipped axes

        Args:
            new_orientation (Sequence): New orientation.
            inplace (bool, optional): If `True`, do operation in-place and return ``self``.

        Returns:
            MedicalVolume: The reformatted volume. If ``inplace=True``, returns ``self``.
        """
        xp = self.device.xp
        device = self.device
        headers = self._headers

        new_orientation = tuple(new_orientation)
        if new_orientation == self.orientation:
            if inplace:
                return self
            return self._partial_clone(volume=self._volume)

        temp_orientation = self.orientation
        temp_affine = np.array(self._affine)

        transpose_inds = stdo.get_transpose_inds(temp_orientation, new_orientation)
        all_transpose_inds = transpose_inds + tuple(range(3, self._volume.ndim))

        with device:
            volume = xp.transpose(self.volume, all_transpose_inds)
        if headers is not None:
            headers = np.transpose(headers, all_transpose_inds)
        for i in range(len(transpose_inds)):
            temp_affine[..., i] = self._affine[..., transpose_inds[i]]

        temp_orientation = tuple([self.orientation[i] for i in transpose_inds])

        flip_axs_inds = list(stdo.get_flip_inds(temp_orientation, new_orientation))
        with device:
            volume = xp.flip(volume, axis=tuple(flip_axs_inds))
        if headers is not None:
            headers = np.flip(headers, axis=tuple(flip_axs_inds))
        a_vecs = temp_affine[:3, :3]
        a_origin = temp_affine[:3, 3]

        # phi is a vector of 1s and -1s, where 1 indicates no flip, and -1 indicates flip
        # phi is used to determine which columns in affine matrix to flip
        phi = np.ones([1, len(a_origin)]).flatten()
        phi[flip_axs_inds] *= -1

        b_vecs = np.array(a_vecs)
        for i in range(len(phi)):
            b_vecs[:, i] *= phi[i]

        # get number of pixels to shift by on each axis.
        # Should be 0 when not flipping - i.e. phi<0 mask
        vol_shape_vec = (
            (np.asarray(volume.shape[:3]) - 1) * (phi < 0).astype(np.float32)
        ).transpose()
        b_origin = np.round(
            a_origin.flatten() - np.matmul(b_vecs, vol_shape_vec).flatten(),
            SCANNER_ORIGIN_DECIMAL_PRECISION,
        )

        temp_affine = np.array(self.affine)
        temp_affine[:3, :3] = b_vecs
        temp_affine[:3, 3] = b_origin
        temp_affine[temp_affine == 0] = 0  # get rid of negative 0s

        if inplace:
            self._affine = temp_affine
            self._volume = volume
            self._headers = headers
            mv = self
        else:
            mv = self._partial_clone(volume=volume, affine=temp_affine, headers=headers)

        assert (
            mv.orientation == new_orientation
        ), f"Orientation mismatch: Expected: {self.orientation}. Got {new_orientation}"
        return mv

    def reformat_as(self, other, inplace: bool = False) -> "MedicalVolume":
        """Reformat this to the same orientation as ``other``.
        Equivalent to ``self.reformat(other.orientation, inplace)``.

        Args:
            other (MedicalVolume): The result volume has the same orientation as ``other``.
            inplace (bool, optional): If `True`, do operation in-place and return ``self``.

        Returns:
            MedicalVolume: The reformatted volume. If ``inplace=True``, returns ``self``.
        """
        return self.reformat(other.orientation, inplace=inplace)

    def is_identical(self, mv):
        """Check if another medical volume is identical.

        Two volumes are identical if they have the same pixel_spacing, orientation,
        scanner_origin, and volume.

        Args:
            mv (MedicalVolume): Volume to compare with.

        Returns:
            bool: `True` if identical, `False` otherwise.
        """
        if not isinstance(mv, MedicalVolume):
            raise TypeError("`mv` must be a MedicalVolume.")

        idevice = self.device
        odevice = mv.device
        if idevice != odevice:
            raise RuntimeError(f"Expected device {idevice}, got {odevice}.")

        with idevice:
            return self.is_same_dimensions(mv) and (mv.volume == self.volume).all()

    def _allclose_spacing(self, mv, precision: int = None, ignore_origin: bool = False):
        """Check if spacing between self and another medical volume is within tolerance.

        Tolerance is `10 ** (-precision)`.

        Args:
            mv (MedicalVolume): Volume to compare with.
            precision (`int`, optional): Number of significant figures after the decimal.
                If not specified, check that affine matrices between two volumes are identical.
                Defaults to `None`.
            ignore_origin (bool, optional): If ``True``, ignore matching origin in the affine
                matrix.

        Returns:
            bool: `True` if spacing between two volumes within tolerance, `False` otherwise.
        """
        if precision is not None:
            tol = 10 ** (-precision)
            return np.allclose(mv.affine[:3, :3], self.affine[:3, :3], atol=tol) and (
                ignore_origin or np.allclose(mv.scanner_origin, self.scanner_origin, rtol=tol)
            )
        else:
            return (mv.affine == self.affine).all() or (
                ignore_origin and (mv.affine[:, :3] == self.affine[:, :3]).all()
            )

    def is_same_dimensions(self, mv, precision: int = None, err: bool = False):
        """Check if two volumes have the same dimensions.

        Two volumes have the same dimensions if they have the same pixel_spacing,
        orientation, and scanner_origin.

        Args:
            mv (MedicalVolume): Volume to compare with.
            precision (`int`, optional): Number of significant figures after the decimal.
                If not specified, check that affine matrices between two volumes are identical.
                Defaults to `None`.
            err (bool, optional): If `True` and volumes do not have same dimensions,
                raise descriptive ValueError.

        Returns:
            bool: ``True`` if pixel spacing, orientation, and scanner origin
                between two volumes within tolerance, ``False`` otherwise.

        Raises:
            TypeError: If ``mv`` is not a MedicalVolume.
            ValueError: If ``err=True`` and two volumes do not have same dimensions.
        """
        if not isinstance(mv, MedicalVolume):
            raise TypeError("`mv` must be a MedicalVolume.")

        is_close_spacing = self._allclose_spacing(mv, precision)
        is_same_orientation = mv.orientation == self.orientation
        is_same_shape = mv.volume.shape == self.volume.shape
        out = is_close_spacing and is_same_orientation and is_same_shape

        if err and not out:
            tol_str = f" (tol: 1e-{precision})" if precision else ""
            if not is_close_spacing:
                raise ValueError(
                    "Affine matrices not equal{}:\n{}\n{}".format(tol_str, self._affine, mv._affine)
                )
            if not is_same_orientation:
                raise ValueError(
                    "Orientations not equal: {}, {}".format(self.orientation, mv.orientation)
                )
            if not is_same_shape:
                raise ValueError(
                    "Shapes not equal: {}, {}".format(self._volume.shape, mv._volume.shape)
                )
            assert False  # should not reach here

        return out

    def match_orientation(self, mv):
        """Reorient another MedicalVolume to orientation specified by self.orientation.

        Args:
            mv (MedicalVolume): Volume to reorient.
        """
        warnings.warn(
            "`match_orientation` is deprecated and will be removed in v0.1. "
            "Use `mv.reformat_as(self, inplace=True)` instead.",
            DeprecationWarning,
        )
        if not isinstance(mv, MedicalVolume):
            raise TypeError("`mv` must be a MedicalVolume.")

        mv.reformat(self.orientation, inplace=True)

    def match_orientation_batch(self, mvs):  # pragma: no cover
        """Reorient a collection of MedicalVolumes to orientation specified by self.orientation.

        Args:
            mvs (list[MedicalVolume]): Collection of MedicalVolumes.
        """
        warnings.warn(
            "`match_orientation_batch` is deprecated and will be removed in v0.1. "
            "Use `[x.reformat_as(self, inplace=True) for x in mvs]` instead.",
            DeprecationWarning,
        )
        for mv in mvs:
            self.match_orientation(mv)

    def clone(self, headers=True):
        """Clones the medical volume.

        Args:
            headers (bool, optional): If `True`, clone headers.
                If `False`, headers have shared memory.

        Returns:
            mv (MedicalVolume): A cloned MedicalVolume.
        """
        return MedicalVolume(
            self.volume.copy(),
            self.affine.copy(),
            headers=deepcopy(self._headers) if headers else self._headers,
        )

    def to(self, device):
        """Move to device.

        If on same device, no-op and returns ``self``.

        Args:
            device: The device to move to.

        Returns:
            MedicalVolume
        """
        device = Device(device)
        if self.device == device:
            return self

        return self._partial_clone(volume=to_device(self._volume, device))

    def cpu(self):
        """Move to cpu."""
        return self.to("cpu")

    def astype(self, dtype, **kwargs):
        """Modifies dtype of ``self._volume``.

        Note this operation is done in place. ``self._volume`` is modified, based
        on the ``astype`` implementation of the type associated with ``self._volume``.
        No new MedicalVolume is created - ``self`` is returned.

        Args:
            dtype (str or dtype): Typecode or data-type to which the array is cast.

        Returns:
            self
        """
        if (
            env.package_available("h5py")
            and isinstance(self._volume, h5py.Dataset)
            and version.parse(env.get_version(h5py)) < version.parse("3.0.0")
        ):
            raise ValueError("Cannot cast h5py.Dataset to dtype for h5py<3.0.0")

        self._volume = self._volume.astype(dtype, **kwargs)
        return self

    def to_nib(self):
        """Converts to nibabel Nifti1Image.

        Returns:
            nibabel.Nifti1Image: The nibabel image.

        Raises:
            RuntimeError: If medical volume is not on the cpu.

        Examples:
            >>> mv = MedicalVolume(np.ones((10,20,30)), np.eye(4))
            >>> mv.to_nib()
            <nibabel.nifti1.Nifti1Image>
        """
        device = self.device
        if device != cpu_device:
            raise RuntimeError(f"MedicalVolume must be on cpu, got {self.device}")

        return nib.Nifti1Image(self.A, self.affine.copy())

    def to_sitk(self, vdim: int = None, transpose_inplane: bool = False):
        """Converts to SimpleITK Image.

        SimpleITK Image objects support vector pixel types, which are represented
        as an extra dimension in numpy arrays. The vector dimension can be specified
        with ``vdim``.

        MedicalVolume must be on cpu. Use ``self.cpu()`` to move.

        SimpleITK loads DICOM files as individual slices that get stacked in ``(z, x, y)``
        order. Thus, ``sitk.GetArrayFromImage`` returns an array in ``(y, x, z)`` order.
        To return a SimpleITK Image that will follow this convention, set
        ``transpose_inplace=True``. If you have been using SimpleITK to load DICOM files,
        you will likely want to specify this parameter.

        Args:
            vdim (int, optional): The vector dimension.
            transpose_inplane (bool, optional): If ``True``, transpose inplane axes.
                Recommended to be ``True`` for users who are familiar with SimpleITK's
                DICOM loading convention.

        Returns:
            SimpleITK.Image

        Raises:
            ImportError: If `SimpleITK` is not installed.
            RuntimeError: If MedicalVolume is not on cpu.

        Note:
            Header information is not currently copied.
        """
        if not env.sitk_available():
            raise ImportError("SimpleITK is not installed. Install it with `pip install simpleitk`")
        device = self.device
        if device != cpu_device:
            raise RuntimeError(f"MedicalVolume must be on cpu, got {self.device}")

        arr = self.volume
        ndim = arr.ndim

        if vdim is not None:
            if vdim < 0:
                vdim = ndim + vdim
            axes = tuple(i for i in range(ndim) if i != vdim)[::-1] + (vdim,)
        else:
            axes = range(ndim)[::-1]
        arr = np.transpose(arr, axes)

        affine = self.affine.copy()
        affine[:2] = -affine[:2]  # RAS+ -> LPS+

        origin = tuple(affine[:3, 3])
        spacing = self.pixel_spacing
        direction = affine[:3, :3] / np.asarray(spacing)

        img = sitk.GetImageFromArray(arr, isVector=vdim is not None)
        img.SetOrigin(origin)
        img.SetSpacing(spacing)
        img.SetDirection(tuple(direction.flatten()))

        if transpose_inplane:
            pa = sitk.PermuteAxesImageFilter()
            pa.SetOrder([1, 0, 2])
            img = pa.Execute(img)

        return img

    def to_torch(
        self, requires_grad: bool = False, contiguous: bool = False, view_as_real: bool = False
    ):
        """Zero-copy conversion to torch tensor.

        If torch version supports complex tensors (i.e. torch>=1.5.0), complex MedicalVolume
        arrays will be converted into complex tensors (torch.complex64/torch.complex128).
        Otherwise, tensors will be returned as the real view, where the last dimension has
        two channels (`tensor.shape[-1]==2`). `[..., 0]` and `[..., 1]` correspond to the
        real/imaginary channels, respectively.

        Args:
            requires_grad (bool, optional): Set ``.requires_grad`` for output tensor.
            contiguous (bool, optional): Make output tensor contiguous before returning.
            view_as_real (bool, optional): If ``True`` and underlying array is complex,
                returns a real view of a complex tensor.

        Returns:
            torch.Tensor: The torch tensor.

        Raises:
            ImportError: If ``torch`` is not installed.

        Note:
            This method does not convert affine matrices and headers to tensor types.

        Examples:
            >>> mv = MedicalVolume(np.ones((2,2,2)), np.eye(4))  # zero-copy on CPU
            >>> mv.to_torch()
            tensor([[[1., 1.],
                     [1., 1.]],
                    [[1., 1.],
                     [1., 1.]]], dtype=torch.float64)
            >>> mv_gpu = MedicalVolume(cp.ones((2,2,2)), np.eye(4))  # zero-copy on GPU
            >>> mv.to_torch()
            tensor([[[1., 1.],
                     [1., 1.]],
                    [[1., 1.],
                     [1., 1.]]], device="cuda:0", dtype=torch.float64)
            >>> # view complex array as real tensor
            >>> mv = MedicalVolume(np.ones((3,4,5), dtype=np.complex), np.eye(4))
            >>> tensor = mv.to_torch(view_as_real)
            >>> tensor.shape
            (3, 4, 5, 2)
        """
        if not env.package_available("torch"):
            raise ImportError(  # pragma: no cover
                "torch is not installed. Install it with `pip install torch`. "
                "See https://pytorch.org/ for more information."
            )

        import torch
        from torch.utils.dlpack import from_dlpack

        device = self.device
        array = self.A

        if any(np.issubdtype(array.dtype, dtype) for dtype in (np.complex64, np.complex128)):
            torch_version = env.get_version(torch)
            supports_cplx = version.Version(torch_version) >= _TORCH_COMPLEX_SUPPORT_VERSION
            if not supports_cplx or view_as_real:
                with device:
                    shape = array.shape
                    array = array.view(dtype=array.real.dtype)
                    array = array.reshape(shape + (2,))

        if device == cpu_device:
            tensor = torch.from_numpy(array)
        else:
            tensor = from_dlpack(array.toDlpack())

        tensor.requires_grad = requires_grad
        if contiguous:
            tensor = tensor.contiguous()
        return tensor

    def headers(self, flatten=False):
        """Returns headers.

        If headers exist, they are currently stored as an array of
        pydicom dataset headers, though this is subject to change.

        Args:
            flatten (bool, optional): If ``True``, flattens header array
                before returning.

        Returns:
            Optional[ndarray[pydicom.dataset.FileDataset]]: Array of headers (if they exist).
        """
        if flatten and self._headers is not None:
            return self._headers.flatten()
        return self._headers

    def get_metadata(self, key, dtype=None, default=np._NoValue):
        """Get metadata value from first header.

        The first header is defined as the first header in ``np.flatten(self._headers)``.
        To extract header information for other headers, use ``self.headers()``.

        Args:
            key (``str`` or pydicom.BaseTag``): Metadata field to access.
            dtype (type, optional): If specified, data type to cast value to.
                By default for DICOM headers, data will be in the value
                representation format specified by pydicom. See
                ``pydicom.valuerep``.
            default (Any): Default value to return if `key`` not found in header.
                If not specified and ``key`` not found in header, raises a KeyError.

        Examples:
            >>> mv.get_metadata("EchoTime")
            '10.0'  # this is a number type ``pydicom.valuerep.DSDecimal``
            >>> mv.get_metadata("EchoTime", dtype=float)
            10.0
            >>> mv.get_metadata("foobar", default=0)
            0

        Raises:
            RuntimeError: If ``self._headers`` is ``None``.
            KeyError: If ``key`` not found and ``default`` not specified.

        Note:
            Currently header information is tied to the ``pydicom.FileDataset`` implementation.
            This function is synonymous to ``dataset.<key>`` in ``pydicom.FileDataset``.
        """
        if self._headers is None:
            raise RuntimeError("No headers found. MedicalVolume must be initialized with `headers`")
        headers = self.headers(flatten=True)

        if key not in headers[0] and default != np._NoValue:
            return default
        else:
            element = headers[0][key]
        val = element.value
        if dtype is not None:
            val = dtype(val)
        return val

    def set_metadata(self, key, value, force: bool = False):
        """Sets metadata for all headers.

        Args:
            key (str or pydicom.BaseTag): Metadata field to access.
            value (Any): The value.
            force (bool, optional): If ``True``, force the header to
                set key even if key does not exist in header.

        Raises:
            RuntimeError: If ``self._headers`` is ``None``.
        """
        if self._headers is None:
            if not force:
                raise ValueError(
                    "No headers found. To generate headers and write keys, `force` must be True."
                )
            self._headers = self._validate_and_format_headers([pydicom.Dataset()])
            warnings.warn(
                "Headers were generated and may not contain all attributes "
                "required to save the volume in DICOM format."
            )

        VR_registry = {float: "DS", int: "IS", str: "LS"}
        for h in self.headers(flatten=True):
            if force and key not in h:
                try:
                    setattr(h, key, value)
                except TypeError:
                    h.add_new(key, VR_registry[type(value)], value)
            else:
                h[key].value = value

    def materialize(self):
        if not self.is_mmap:
            return self

    def round(self, decimals=0, affine=False) -> "MedicalVolume":
        """Round array (and optionally affine matrix).

        Args:
            decimals (int, optional): Number of decimals to round to.
            affine (bool, optional): The rounded medical volume.

        Returns:
            MedicalVolume: MedicalVolume with rounded.
        """
        from dosma.core.numpy_routines import around

        return around(self, decimals, affine)

    def sum(
        self,
        axis=None,
        dtype=None,
        out=None,
        keepdims=False,
        initial=np._NoValue,
        where=np._NoValue,
    ) -> "MedicalVolume":
        """Compute the arithmetic sum along the specified axis. Identical to :meth:`sum_np`.

        See :meth:`sum_np` for more information.

        Args:
            axis: Same as :meth:`sum_np`.
            dtype: Same as :meth:`sum_np`.
            out: Same as :meth:`sum_np`.
            keepdims: Same as :meth:`sum_np`.
            initial: Same as :meth:`sum_np`.
            where: Same as :meth:`sum_np`.

        Returns:
            Union[Number, MedicalVolume]: If ``axis=None``, returns a number or a scalar type of
                the underlying ndarray. Otherwise, returns a medical volume containing sum
                values.
        """
        from dosma.core.numpy_routines import sum_np

        # `out` is required for cupy arrays because of how cupy calls array.
        if out is not None:
            raise ValueError("`out` must be None")
        return sum_np(self, axis=axis, dtype=dtype, keepdims=keepdims, initial=initial, where=where)

    def mean(
        self, axis=None, dtype=None, out=None, keepdims=False, where=np._NoValue
    ) -> Union[Number, "MedicalVolume"]:
        """Compute the arithmetic mean along the specified axis. Identical to :meth:`mean_np`.

        See :meth:`mean_np` for more information.

        Args:
            axis: Same as :meth:`mean_np`.
            dtype: Same as :meth:`mean_np`.
            out: Same as :meth:`mean_np`.
            keepdims: Same as :meth:`mean_np`.
            initial: Same as :meth:`mean_np`.
            where: Same as :meth:`mean_np`.

        Returns:
            Union[Number, MedicalVolume]: If ``axis=None``, returns a number or a scalar type of
            the underlying ndarray. Otherwise, returns a medical volume containing mean
            values.
        """
        from dosma.core.numpy_routines import mean_np

        # `out` is required for cupy arrays because of how cupy calls array.
        if out is not None:
            raise ValueError("`out` must be None")
        return mean_np(self, axis=axis, dtype=dtype, keepdims=keepdims, where=where)

    @property
    def A(self):
        """The pixel array. Same as ``self.volume``.

        Examples:
            >>> mv = MedicalVolume([[[1,2],[3,4]]], np.eye(4))
            >>> mv.A
            array([[[1, 2],
                    [3, 4]]])
        """
        return self.volume

    @property
    def volume(self):
        """ndarray: ndarray representing volume values."""
        return self._volume

    @volume.setter
    def volume(self, value):
        """
        If the volume is of a different shape, the headers are no longer valid,
        so delete all reorientations are done as part of MedicalVolume,
        so reorientations are permitted.

        However, external setting of the volume to a different shape array is not allowed.
        """
        if value.ndim != self._volume.ndim:
            raise ValueError("New volume must be same as current volume")

        if self._volume.shape != value.shape:
            self._headers = None

        self._volume = value
        self._device = get_device(self._volume)

    @property
    def pixel_spacing(self):
        """tuple[float]: Pixel spacing in order of current orientation."""
        vecs = self._affine[:3, :3]
        ps = tuple(np.sqrt(np.sum(vecs ** 2, axis=0)))

        assert len(ps) == 3, "Pixel spacing must have length of 3"
        return ps

    @property
    def orientation(self):
        """tuple[str]: Image orientation in standard orientation format.

        See orientation.py for more information on conventions.
        """
        nib_orientation = nib.aff2axcodes(self._affine)
        return stdo.orientation_nib_to_standard(nib_orientation)

    @property
    def scanner_origin(self):
        """tuple[float]: Scanner origin in global RAS+ x,y,z coordinates."""
        return tuple(self._affine[:3, 3])

    @property
    def affine(self):
        """np.ndarray: 4x4 affine matrix for volume in current orientation."""
        return self._affine

    @property
    def shape(self) -> Tuple[int, ...]:
        """The shape of the underlying ndarray."""
        return self._volume.shape

    @property
    def ndim(self) -> int:
        """int: The number of dimensions of the underlying ndarray."""
        return self._volume.ndim

    @property
    def device(self) -> Device:
        """The device the object is on."""
        return get_device(self._volume)

    @property
    def dtype(self):
        """The ``dtype`` of the ndarray. Same as ``self.volume.dtype``."""
        return self._volume.dtype

    @property
    def is_mmap(self) -> bool:
        """bool: Whether the volume is a memory-mapped array."""
        # important to check if .base is a python mmap object, since a view of a mmap
        # is also a memmap object, but should not be symlinked or copied
        return isinstance(self.A, np.memmap) and isinstance(self.A.base, mmap)

    @classmethod
    def from_nib(
        cls, image, affine_precision: int = None, origin_precision: int = None, mmap: bool = False
    ) -> "MedicalVolume":
        """Constructs MedicalVolume from nibabel images.

        Args:
            image (nibabel.Nifti1Image): The nibabel image to convert.
            affine_precision (int, optional): If specified, rounds the i/j/k coordinate
                vectors in the affine matrix to this decimal precision.
            origin_precision (int, optional): If specified, rounds the scanner origin
                in the affine matrix to this decimal precision.
            mmap (bool, optional): If True, memory map the image.

        Returns:
            MedicalVolume: The medical image.

        Examples:
            >>> import nibabel as nib
            >>> nib_img = nib.Nifti1Image(np.ones((10,20,30)), np.eye(4))
            >>> MedicalVolume.from_nib(nib_img)
            MedicalVolume(
                shape=(10, 20, 30),
                ornt=('LR', 'PA', 'IS')),
                spacing=(1.0, 1.0, 1.0),
                origin=(0.0, 0.0, 0.0),
                device=Device(type='cpu')
            )
        """
        affine = np.array(image.affine)  # Make a copy of the affine matrix.
        if affine_precision is not None:
            affine[:3, :3] = np.round(affine[:3, :3], affine_precision)
        if origin_precision:
            affine[:3, 3] = np.round(affine[:3, 3], origin_precision)

        data = image.dataobj.__array__() if mmap else image.get_fdata()
        mv = cls(data, affine)
        if mmap and not mv.is_mmap:
            raise ValueError(
                "Underlying array in the nibabel image is not mem-mapped. " "Please set mmap=False."
            )
        return mv

    @classmethod
    def from_sitk(cls, image, copy=False, transpose_inplane: bool = False) -> "MedicalVolume":
        """Constructs MedicalVolume from SimpleITK.Image.

        Use ``transpose_inplane=True`` if the SimpleITK image was loaded with SimpleITK's
        DICOM reader or if ``transpose_inplace=True`` was used to create the Image
        with :meth:`to_sitk`. See the discussion of SimpleITK's data ordering convention
        in :meth:`to_sitk` for more information.

        If you are getting a segmentation fault, try using ``copy=True``.

        Args:
            image (SimpleITK.Image): The image.
            copy (bool, optional): If ``True``, copies array.
            transpose_inplane (bool, optional): If ``True``, transposes the inplane axes.
                Set this to ``True`` if the SimpleITK image was loaded with SimpleITK's
                DICOM reader. May need to set ``copy=True`` to avoid segmentation fault.

        Returns:
            MedicalVolume

        Note:
            Metadata information is not copied.
        """
        if not env.sitk_available():
            raise ImportError("SimpleITK is not installed. Install it with `pip install simpleitk`")

        if len(image.GetSize()) < 3:
            raise ValueError("`image` must be 3D.")
        is_vector_image = image.GetNumberOfComponentsPerPixel() > 1

        if transpose_inplane:
            pa = sitk.PermuteAxesImageFilter()
            pa.SetOrder([1, 0, 2])
            image = pa.Execute(image)

        if copy:
            arr = sitk.GetArrayFromImage(image)
        else:
            arr = sitk.GetArrayViewFromImage(image)

        ndim = arr.ndim
        if is_vector_image:
            axes = tuple(range(ndim)[-2::-1]) + (ndim - 1,)
        else:
            axes = range(ndim)[::-1]
        arr = np.transpose(arr, axes)

        origin = image.GetOrigin()
        spacing = image.GetSpacing()
        direction = np.asarray(image.GetDirection()).reshape(-1, 3)

        affine = np.zeros((4, 4))
        affine[:3, :3] = direction * np.asarray(spacing)
        affine[:3, 3] = origin
        affine[:2] = -affine[:2]  # LPS+ -> RAS+
        affine[3, 3] = 1

        return cls(arr, affine)

    @classmethod
    def from_torch(cls, tensor, affine, headers=None, to_complex: bool = None) -> "MedicalVolume":
        """Zero-copy construction from PyTorch tensor.

        Args:
            tensor (torch.Tensor): A PyTorch tensor where first three dimensions correspond
                to spatial dimensions.
            affine (np.ndarray): See class parameters.
            headers (np.ndarray[pydicom.FileDataset], optional): See class parameters.
            to_complex (bool, optional): If ``True``, interprets tensor as real view of complex
                tensor and attempts to restructure it as a complex array.

        Returns:
            MedicalVolume: A medical image.

        Raises:
            RuntimeError: If ``affine`` is not on the cpu.
            ValueError: If ``tensor`` does not have at least three spatial dimensions.
            ValueError: If ``to_complex=True`` and shape is not size ``(..., 2)``.
            ImportError: If ``tensor`` on GPU and ``cupy`` not installed.

        Examples:
            >>> import torch
            >>> tensor = torch.ones((2,2,2))
            >>> MedicalVolume.from_torch(tensor, affine=np.eye(4))
            MedicalVolume(
                shape=(2, 2, 2),
                ornt=('LR', 'PA', 'IS')),
                spacing=(1.0, 1.0, 1.0),
                origin=(0.0, 0.0, 0.0),
                device=Device(type='cpu')
            )
            >>> tensor = torch.ones((2,2,2), device="cuda")  # zero-copy from GPU 0
            >>> MedicalVolume.from_torch(tensor, affine=np.eye(4))
            MedicalVolume(
                shape=(2, 2, 2),
                ornt=('LR', 'PA', 'IS')),
                spacing=(1.0, 1.0, 1.0),
                origin=(0.0, 0.0, 0.0),
                device=Device(type='cuda', index=0)
            )
            >>> tensor = torch.ones((3,4,5,2))  # treat this tensor as view of complex tensor
            >>> mv = MedicalVolume.from_torch(tensor, affine=np.eye(4), to_complex=True)
            >>> print(mv)
            MedicalVolume(
                shape=(3,4,5),
                ornt=('LR', 'PA', 'IS')),
                spacing=(1.0, 1.0, 1.0),
                origin=(0.0, 0.0, 0.0),
                device=Device(type='cuda', index=0)
            )
            >>> mv.dtype
            np.complex128
        """
        if not env.package_available("torch"):
            raise ImportError(  # pragma: no cover
                "torch is not installed. Install it with `pip install torch`. "
                "See https://pytorch.org/ for more information."
            )

        import torch
        from torch.utils.dlpack import to_dlpack

        torch_version = env.get_version(torch)
        supports_cplx = version.Version(torch_version) >= _TORCH_COMPLEX_SUPPORT_VERSION
        # Check if tensor needs to be converted to np.complex type.
        # If tensor is of torch.complex64 or torch.complex128 dtype, then from_numpy will take
        # care of conversion to appropriate numpy dtype, and we do not need to do the to_complex
        # logic.
        to_complex = to_complex and (
            not supports_cplx
            or (supports_cplx and tensor.dtype not in (torch.complex64, torch.complex128))
        )

        if isinstance(affine, torch.Tensor):
            if Device(affine.device) != cpu_device:
                raise RuntimeError("Affine matrix must be on the cpu")
            affine = affine.numpy()

        if (not to_complex and tensor.ndim < 3) or (to_complex and tensor.ndim < 4):
            raise ValueError(
                f"Tensor must have three spatial dimensions. Got shape {tensor.shape}."
            )
        if to_complex and tensor.shape[-1] != 2:
            raise ValueError(
                f"tensor.shape[-1] must have shape 2 when to_complex is specified. "
                f"Got shape {tensor.shape}."
            )

        device = Device(tensor.device)
        if device == cpu_device:
            array = tensor.detach().numpy()
        else:
            if env.cupy_available():
                array = cp.fromDlpack(to_dlpack(tensor))
            else:
                raise ImportError(  # pragma: no cover
                    "CuPy is required to convert a GPU torch.Tensor to array. "
                    "Follow instructions at https://docs.cupy.dev/en/stable/install.html to "
                    "install the correct binary."
                )

        if to_complex:
            with get_device(array):
                if array.dtype == np.float32:
                    array = array.view(np.complex64)
                elif array.dtype == np.float64:
                    array = array.view(np.complex128)

                array = array.reshape(array.shape[:-1])

        return cls(array, affine, headers=headers)

    def _partial_clone(self, **kwargs) -> "MedicalVolume":
        """Copies constructor information from ``self`` if not available in ``kwargs``."""
        if kwargs.get("volume", None) is False:
            # special use case to not clone volume
            kwargs["volume"] = self._volume
        for k in ("volume", "affine"):
            if k not in kwargs or (kwargs[k] is True):
                kwargs[k] = getattr(self, f"_{k}").copy()
        if "headers" not in kwargs:
            kwargs["headers"] = self._headers
        elif isinstance(kwargs["headers"], bool) and kwargs["headers"]:
            kwargs["headers"] = deepcopy(self._headers)
        return self.__class__(**kwargs)

    def _validate_and_format_headers(self, headers):
        """Validate headers are of appropriate shape and format into standardized shape.

        Headers are stored an ndarray of dictionary-like objects with explicit dimensions
        that match the dimensions of ``self._volume``. If header objects are not

        Assumes ``self._volume`` and ``self._affine`` have been set.
        """
        headers = np.asarray(headers)
        if headers.ndim > self._volume.ndim:
            raise ValueError(
                f"`headers` has too many dimensions. "
                f"Got headers.ndim={headers.ndim}, but volume.ndim={self._volume.ndim}"
            )
        for dim in range(-headers.ndim, 0)[::-1]:
            if headers.shape[dim] not in (1, self._volume.shape[dim]):
                raise ValueError(
                    f"`headers` must follow standard broadcasting shape. "
                    f"Got headers.shape={headers.shape}, but volume.shape={self._volume.shape}"
                )

        ndim = self._volume.ndim
        shape = (1,) * (ndim - len(headers.shape)) + headers.shape
        headers = np.reshape(headers, shape)
        return headers

    def _extract_input_array_ufunc(self, input, device=None):
        if device is None:
            device = self.device
        device_err = "Expected device {} but got device ".format(device) + "{}"
        if isinstance(input, Number):
            return input
        elif isinstance(input, np.ndarray):
            if device != cpu_device:
                raise RuntimeError(device_err.format(cpu_device))
            return input
        elif env.cupy_available() and isinstance(input, cp.ndarray):
            if device != input.device:
                raise RuntimeError(device_err.format(Device(input.device)))
            return input
        elif isinstance(input, MedicalVolume):
            if device != input.device:
                raise RuntimeError(device_err.format(Device(input.device)))
            assert self.is_same_dimensions(input, err=True)
            return input._volume
        else:
            return NotImplemented

    def _check_reduce_axis(self, axis: Union[int, Sequence[int]]) -> Tuple[int]:
        if axis is None:
            return None
        is_sequence = isinstance(axis, Sequence)
        if not is_sequence:
            axis = (axis,)
        axis = tuple(x if x >= 0 else self.volume.ndim + x for x in axis)
        assert all(x >= 0 for x in axis)
        if any(x < 3 for x in axis):
            raise ValueError("Cannot reduce MedicalVolume along spatial dimensions")
        if not is_sequence:
            axis = axis[0]
        return axis

    def _reduce_array(self, func, *inputs, **kwargs) -> "MedicalVolume":
        """
        Assumes inputs have been verified.
        """
        device = self.device
        xp = device.xp

        keepdims = kwargs.get("keepdims", False)
        reduce_axis = self._check_reduce_axis(kwargs["axis"])
        kwargs["axis"] = reduce_axis
        if not isinstance(reduce_axis, Sequence):
            reduce_axis = (reduce_axis,)
        with device:
            volume = func(*inputs, **kwargs)

        if xp.isscalar(volume) or volume.ndim == 0:
            return volume

        if self._headers is not None:
            headers_slices = tuple(
                slice(None) if x not in reduce_axis else slice(0, 1) if keepdims else 0
                for x in range(self._headers.ndim)
            )
            headers = self._headers[headers_slices]
        else:
            headers = None
        return self._partial_clone(volume=volume, headers=headers)

    def __getitem__(self, _slice):
        if isinstance(_slice, MedicalVolume):
            _slice = _slice.reformat_as(self).A

        slicer = _SpatialFirstSlicer(self)
        try:
            _slice = slicer.check_slicing(_slice)
        except ValueError as err:
            raise IndexError(*err.args)

        volume = self._volume[_slice]
        if any(dim == 0 for dim in volume.shape):
            raise IndexError("Empty slice requested")

        headers = self._headers
        if headers is not None:
            _slice_headers = []
            for idx, x in enumerate(_slice):
                if headers.shape[idx] == 1 and not isinstance(x, int):
                    _slice_headers.append(slice(None))
                elif headers.shape[idx] == 1 and isinstance(x, int):
                    _slice_headers.append(0)
                else:
                    _slice_headers.append(x)
            headers = headers[_slice_headers]

        affine = slicer.slice_affine(_slice)
        return self._partial_clone(volume=volume, affine=affine, headers=headers)

    def __setitem__(self, _slice, value):
        """
        Note:
            When ``value`` is a ``MedicalVolume``, the headers from that value
            are not copied over. This may be changed in the future.
        """
        if isinstance(value, MedicalVolume):
            image = self[_slice]
            assert value.is_same_dimensions(image, err=True)
            value = value._volume
        with self.device:
            self._volume[_slice] = value
        if self.is_mmap and self._volume.mode == "c":
            self._volume = np.asarray(self._volume)

    def __repr__(self) -> str:
        nl = "\n"
        nltb = "\n  "
        return (
            f"{self.__class__.__name__}({nltb}shape={self.shape},{nltb}"
            f"ornt={self.orientation}),{nltb}spacing={self.pixel_spacing},{nltb}"
            f"origin={self.scanner_origin},{nltb}device={self.device}{nl})"
        )

    def _iops(self, other, op):
        """Helper function for i-type ops (__iadd__, __isub__, etc.)"""
        if isinstance(other, MedicalVolume):
            assert self.is_same_dimensions(other, err=True)
            other = other.volume

        if isinstance(op, str):
            op = getattr(self._volume, op)
        op(other)

        if self.is_mmap and self._volume.mode == "c":
            self._volume = np.asarray(self._volume)
        return self

    def __iadd__(self, other):
        return self._iops(other, self._volume.__iadd__)

    def __ifloordiv__(self, other):
        return self._iops(other, self._volume.__ifloordiv__)

    def __imul__(self, other):
        return self._iops(other, self._volume.__imul__)

    def __ipow__(self, other):
        return self._iops(other, self._volume.__ipow__)

    def __isub__(self, other):
        return self._iops(other, self._volume.__isub__)

    def __itruediv__(self, other):
        return self._iops(other, self._volume.__itruediv__)

    def __array__(self):
        """Wrapper for performing numpy operations on MedicalVolume array.

        Examples:
            >>> a = np.asarray(mv)
            >>> type(a)
            <class 'numpy.ndarray'>

        Note:
            This is not valid when ``self.volume`` is a ``cupy.ndarray``.
            All CUDA ndarrays must first be moved to the cpu.
        """
        try:
            return np.asarray(self.volume)
        except TypeError:
            raise TypeError(
                "Implicit conversion to a NumPy array is not allowed. "
                "Please use `.cpu()` to move the array to the cpu explicitly "
                "before constructing a NumPy array."
            )

    def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
        def _extract_inputs(inputs, device):
            _inputs = []
            for input in inputs:
                input = self._extract_input_array_ufunc(input, device)
                if input is NotImplemented:
                    return input
                _inputs.append(input)
            return _inputs

        if method not in ["__call__", "reduce"]:
            return NotImplemented

        device = self.device
        _inputs = _extract_inputs(inputs, device)
        if _inputs is NotImplemented:
            return NotImplemented

        if method == "__call__":
            with device:
                volume = ufunc(*_inputs, **kwargs)
            if volume.shape != self._volume.shape:
                raise ValueError(
                    f"{self.__class__.__name__} does not support operations that change shape. "
                    f"Use operations on `self.volume` to modify array objects."
                )
            return self._partial_clone(volume=volume)
        elif method == "reduce":
            return self._reduce_array(ufunc.reduce, *_inputs, **kwargs)

    def __array_function__(self, func, types, args, kwargs):
        from dosma.core.numpy_routines import _HANDLED_NUMPY_FUNCTIONS

        if func not in _HANDLED_NUMPY_FUNCTIONS:
            return NotImplemented
        # Note: this allows subclasses that don't override
        # __array_function__ to handle MedicalVolume objects.
        if not all(issubclass(t, (MedicalVolume, self.__class__)) for t in types):
            return NotImplemented
        return _HANDLED_NUMPY_FUNCTIONS[func](*args, **kwargs)

    @property
    def __cuda_array_interface__(self):
        """Wrapper for performing cupy operations on MedicalVolume array."""
        if self.device == cpu_device:
            raise TypeError(
                "Implicit conversion to a CuPy array is not allowed. "
                "Please use `.to(device)` to move the array to the gpu explicitly "
                "before constructing a CuPy array."
            )
        return self.volume.__cuda_array_interface__

class _SpatialFirstSlicer(_SpatialFirstSlicerNib):
    def __init__(self, img):
        self.img = img

    def __getitem__(self, slicer):
        raise NotImplementedError("Slicing should be done by `MedicalVolume`")


from abc import ABCMeta, abstractmethod
from bisect import bisect_right
from typing import Any, Dict, Iterable, List, Optional, Tuple

from volatility3.framework import exceptions, interfaces
from volatility3.framework.configuration import requirements
from volatility3.framework.layers import linear

class NonLinearlySegmentedLayer(
    interfaces.layers.TranslationLayerInterface, metaclass=ABCMeta
):
    """A class to handle a single run-based layer-to-layer mapping.

    In the documentation "mapped address" or "mapped offset" refers to
    an offset once it has been mapped to the underlying layer
    """

    def __init__(
        self,
        context: interfaces.context.ContextInterface,
        config_path: str,
        name: str,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> None:
        super().__init__(
            context=context, config_path=config_path, name=name, metadata=metadata
        )

        self._base_layer = self.config["base_layer"]
        self._segments: List[Tuple[int, int, int, int]] = []
        self._minaddr: Optional[int] = None
        self._maxaddr: Optional[int] = None

        self._load_segments()

    @abstractmethod
    def _load_segments(self) -> None:
        """Populates the _segments variable.

        Segments must be (address, mapped address, length, mapped_length) and must be
        sorted by address when this method exits
        """

    def is_valid(self, offset: int, length: int = 1) -> bool:
        """Returns whether the address offset can be translated to a valid
        address."""
        try:
            base_layer = self._context.layers[self._base_layer]
            return all(
                [
                    base_layer.is_valid(mapped_offset)
                    for _i, _i, mapped_offset, _i, _s in self.mapping(offset, length)
                ]
            )
        except exceptions.InvalidAddressException:
            return False

    def _find_segment(
        self, offset: int, next: bool = False
    ) -> Tuple[int, int, int, int]:
        """Finds the segment containing a given offset.

        Returns the segment tuple (offset, mapped_offset, length, mapped_length)
        """

        if not self._segments:
            self._load_segments()

        # Find rightmost value less than or equal to x
        i = bisect_right(
            self._segments,
            (offset, self.context.layers[self._base_layer].maximum_address),
        )
        if i and not next:
            segment = self._segments[i - 1]
            if segment[0] <= offset < segment[0] + segment[2]:
                return segment
        if next:
            if i < len(self._segments):
                return self._segments[i]
        raise exceptions.InvalidAddressException(
            self.name, offset, f"Invalid address at {offset:0x}"
        )

    # Determines whether larger segments are in use and the offsets within them should be tracked linearly
    # When no decoding of the data occurs, this should be set to true
    _track_offset = False

    def mapping(
        self, offset: int, length: int, ignore_errors: bool = False
    ) -> Iterable[Tuple[int, int, int, int, str]]:
        """Returns a sorted iterable of (offset, length, mapped_offset, mapped_length, layer)
        mappings."""
        done = False
        current_offset = offset
        while not done:
            try:
                # Search for the appropriate segment that contains the current_offset
                logical_offset, mapped_offset, size, mapped_size = self._find_segment(
                    current_offset
                )
                # If it starts before the current_offset, bring the lower edge up to the right place
                if current_offset > logical_offset:
                    difference = current_offset - logical_offset
                    logical_offset += difference
                    if self._track_offset:
                        mapped_offset += difference
                    size -= difference
            except exceptions.InvalidAddressException:
                if not ignore_errors:
                    # If we're not ignoring errors, raise the invalid address exception
                    raise
                try:
                    # Find the next valid segment after our current_offset
                    (
                        logical_offset,
                        mapped_offset,
                        size,
                        mapped_size,
                    ) = self._find_segment(current_offset, next=True)
                    # We know that the logical_offset must be greater than current_offset so skip to that value
                    current_offset = logical_offset
                    # If it starts too late then we're done
                    if logical_offset > offset + length:
                        return
                except exceptions.InvalidAddressException:
                    return
            # Crop it to the amount we need left
            chunk_size = min(size, length + offset - logical_offset)
            yield logical_offset, chunk_size, mapped_offset, mapped_size, self._base_layer
            current_offset += chunk_size
            # Terminate if we've gone (or reached) our required limit
            if current_offset >= offset + length:
                done = True

    @property
    def minimum_address(self) -> int:
        if not self._segments:
            raise ValueError("SegmentedLayer must contain some segments")
        if self._minaddr is None:
            mapped, _, _, _ = self._segments[0]
            self._minaddr = mapped
        return self._minaddr

    @property
    def maximum_address(self) -> int:
        if not self._segments:
            raise ValueError("SegmentedLayer must contain some segments")
        if self._maxaddr is None:
            mapped, _, length, _ = self._segments[-1]
            self._maxaddr = mapped + length
        return self._maxaddr

    @property
    def dependencies(self) -> List[str]:
        """Returns a list of the lower layers that this layer is dependent
        upon."""
        return [self._base_layer]

    @classmethod
    def get_requirements(cls) -> List[interfaces.configuration.RequirementInterface]:
        return [
            requirements.TranslationLayerRequirement(name="base_layer", optional=False)
        ]

class SegmentedLayer(
    NonLinearlySegmentedLayer, linear.LinearlyMappedLayer, metaclass=ABCMeta
):
    _track_offset = True

    def mapping(
        self, offset: int, length: int, ignore_errors: bool = False
    ) -> Iterable[Tuple[int, int, int, int, str]]:
        # Linear mappings must return the same length of segment as that requested
        for offset, length, mapped_offset, mapped_length, layer in super().mapping(
            offset, length, ignore_errors
        ):
            yield offset, length, mapped_offset, length, layer


# Name: BoPclip.py
# Created on: 2017-03-06
# Author: Jessica Nephin

# Description:
#			Clips NCC bottom patches to correct boundary
#			Converts factor attributes
#			Removes fields

# Import modules
import os
import arcpy

# move up one directory
os.chdir('..')

# allow overwriting
arcpy.env.overwriteOutput = True

# Set environment settings
arcpy.env.workspace = os.getcwd()

# Bottom patches, get layer basename
g = "NCC_BoPs_v1.1.gdb"
inLayer = "BoP18_merged"

# Input feature
inFeature = "InputData/Polygons/NCC_BoPs_v1.1.gdb/BoP18_merged"
# Boundary layer
boundary = os.getcwd()+"/Boundary/NCC_Nearshore_Area_BoP.shp"
#Output feature
outFeature = "InputData/Polygons/NCC_BoPs_v1.1.gdb/BoP_clipped"

# Clip layer
arcpy.Clip_analysis(inFeature, boundary, outFeature)

#-----------------------------------------------------------#
# Combine Btype1 and Btype2

# # Add field
# fexists = arcpy.ListFields (outFeature, "BTypeComb")
# if len(fexists) != 1:
# 	arcpy.AddField_management (outFeature, "BTypeComb", "TEXT")
# #Calculate field
# arcpy.CalculateField_management (outFeature, "BTypeComb", "!BType1! + !BType2!", "PYTHON_9.3")

#-----------------------------------------------------------#
# Convert attributes to integer factors

# BType

# Rename BType1 to BType
arcpy.AlterField_management(outFeature, 'BType1', 'BType', 'BType')

# Add field
# fexists = arcpy.ListFields (outFeature, "BType")
# if len(fexists) != 1:
# 	arcpy.AddField_management (outFeature, "BType", "SHORT")
# #Calculate field
# arcpy.CalculateField_management (outFeature, "BType",
# 								'''{"1": "1", "1a": "2", "1b": "3",
# 								"2": "4", "2a": "5", "2b": "6",
# 								"3": "7", "3a": "8", "3b": "9"}.get(!BTypeComb!, "0")''',
# 								"PYTHON_9.3")

# DepthCode

# Add field
fexists = arcpy.ListFields (outFeature, "DepthBin")
if len(fexists) != 1:
	arcpy.AddField_management (outFeature, "DepthBin", "SHORT")
#Calculate field
arcpy.CalculateField_management (outFeature, "DepthBin",
								'''{"ITD": "1", "0to5": "2", "5to10": "3",
								"10to20": "4", "20to50": "5"}.get(!DepthCode!, "0")''',
								"PYTHON_9.3")

#-----------------------------------------------------------#
# Remove feilds

# add new feature
finalFeature = "InputData/Polygons/NCC_BoPs_v1.1.gdb/BoP"
arcpy.CopyFeatures_management(outFeature, finalFeature)

arcpy.DeleteField_management (finalFeature, ["SourceKey", "Confidence",
						"ObsBType1", "ObsBType2", "ObsArea","ObsDist", "GrabBType1",
						"GrabBType2", "GrabArea", "GrabDist", "SZBType1","SZBType2",
						"SZArea","SZITDDist", "SZ0to5Dist", "ObsDepthCode", "FragmentOrigFID",
						"GrabDepthCode","GrabSourceKey", "SZDepthCode","SZSourceKey",
						"ObsSourceKey", "BType1", "BType2", "BTypeComb", "DepthCode"])


from commands.Command import Command
import discord
import asyncio

class Cat(Command):
    def __init__(self):
        super().__init__(
            {
                'name': 'cat',
                'description': 'extracts the text content of your file',
                'argc': 1
            }
        )

    async def run(self, message: discord.Message, argc: int, argv: list):
        extension = argv[0]
        valid_extensions = ['cpp']
        if extension not in valid_extensions:
            emb = discord.Embed(
                title='Error', description="Invalid file extension.", colour=discord.Color.from_rgb(255, 0, 0)
            )
            await message.channel.send(embed=emb)
            return

        file = open(f'./files/{message.author}.{extension}', 'r')
        msg = '```cpp\n'
        for line in file.readlines():
            if len(line.strip()) != 0:
                msg += line
            else:
                msg += '\n'
        msg += '\n```'
        file.close()
        await message.channel.send(msg)


import random
secret = "anqomr"

def numOfMatch(w1, w2):
    ans = 0
    for i in range(len(w1)):
        if w1[i] == w2[i]:
            ans += 1
    return  ans

def guess(w):
    return numOfMatch(w, secret)

def findSecretWord(wordlist, guess):
    """
    :type wordlist: List[Str]
    :type master: Master
    :rtype: None
    """
    H = {}
    for i in range(len(wordlist)):
        if not wordlist[i] in H:
            H[wordlist[i]] = {}
        for j in range(len(wordlist)):
            wi = wordlist[i]
            wj = wordlist[j]
            m = numOfMatch(wi, wj)
            if not m in H[wi]:
                H[wi][m] = [wj]
            else:
                H[wi][m].append(wj)

    candidates = set(wordlist)
    g = wordlist[random.randint(0, len(wordlist)-1)]
    cnt = 10
    while cnt>0:
        m = guess(g)
        if m == 6:
            print('done')
            return
        candidates = candidates & set(H[g][m])
        g = list(candidates)[random.randint(0, len(candidates)-1)]

    return

print(findSecretWord(["pzrooh","aaakrw","vgvkxb","ilaumf","snzsrz","qymapx","hhjgwz","mymxyu","jglmrs","yycsvl","fuyzco","ivfyfx","hzlhqt","ansstc","ujkfnr","jrekmp","himbkv","tjztqw","jmcapu","gwwwmd","ffpond","ytzssw","afyjnp","ttrfzi","xkwmsz","oavtsf","krwjwb","bkgjcs","vsigmc","qhpxxt","apzrtt","posjnv","zlatkz","zynlqc","czajmi","smmbhm","rvlxav","wkytta","dzkfer","urajfh","lsroct","gihvuu","qtnlhu","ucjgio","xyycvd","fsssoo","kdtmux","bxnppe","usucra","hvsmau","gstvvg","ypueay","qdlvod","djfbgs","mcufib","prohkc","dsqgms","eoasya","kzplbv","rcuevr","iwapqf","ucqdac","anqomr","msztnf","tppefv","mplbgz","xnskvo","euhxrh","xrqxzw","wraxvn","zjhlou","xwdvvl","dkbiys","zbtnuv","gxrhjh","ctrczk","iwylwn","wefuhr","enlcrg","eimtep","xzvntq","zvygyf","tbzmzk","xjptby","uxyacb","mbalze","bjosah","ckojzr","ihboso","ogxylw","cfnatk","zijwnl","eczmmc","uazfyo","apywnl","jiraqa","yjksyd","mrpczo","qqmhnb","xxvsbx"], guess))


import os
import glob
import shutil
import traceback
from .context import KFJobContext

class DataPathCleaner(object):
    def __init__(self, base_path=None):
        self.base_path = (base_path or '').strip() or (KFJobContext.get_context().get_user_path() or '').strip()
        if self.base_path and not os.path.isabs(self.base_path):
            raise RuntimeError("base_path must be a absolute path, got '{}'".format(self.base_path))
        elif self.base_path:
            self.base_path = os.path.normpath(self.base_path)
        print("{}: set base_path='{}'".format(self, self.base_path))

    def clean(self, file_pattern):
        if not file_pattern:
            print("file_pattern not specified, will not do clean")
            return
        if not isinstance(file_pattern, list):
            file_pattern = [file_pattern]

        for fp in file_pattern:
            if os.path.isabs(fp):
                file_to_clean = os.path.abspath(fp)
            else:
                file_to_clean = os.path.abspath(os.path.join(KFJobContext.get_context().export_path, fp))

            if self.base_path and not file_to_clean.startswith(self.base_path):
                print("file_pattern '{}' go out of base path '{}', ignore it".format(fp, self.base_path))
                continue
            for ftc in glob.glob(file_to_clean):
                try:
                    if os.path.isdir(ftc):
                        shutil.rmtree(ftc, ignore_errors=True)
                        print("removed dir '{}' derived from '{}', pattern='{}'".format(ftc, file_to_clean, fp))
                    else:
                        os.remove(ftc)
                        print("removed file '{}' derived from '{}', pattern='{}'".format(ftc, file_to_clean, fp))
                except Exception as e:
                    print("clean '{}' dervied from '{}', pattern='{}' error: {}\n{}"
                          .format(ftc, file_to_clean, fp, e, traceback.format_exc()))


# -*- coding: utf-8 -*-
"""
Created on Wed Feb  9 12:39:55 2022

@author: pkinn
"""
def cvTrain(model, features, targets, nSplits, nEpochs, batchSz, initWts):
    from sklearn.model_selection import KFold
    import numpy as np
    kf = KFold(n_splits = nSplits, shuffle = True)
    fn = 1
    # Define per-fold score containers <-- these are new
    acc_per_fold = []
    loss_per_fold = []
    allHist = np.zeros((nSplits, nEpochs))
    for train, test in kf.split(features, targets):
        print('------------------------------------------------------------------------')
        print(f'Training for fold {fn} ...')
        model.set_weights(initWts)
        for kk in range(nEpochs):
            history = model.fit(features[train], targets[train],
                                batch_size = batchSz,
                                epochs = 1,
                                verbose = 0)
            scores = model.evaluate(features[test], targets[test], verbose = 0)
            allHist[fn-1, kk] = scores[1]
        print(f'Score for fold {fn}: {model.metrics_names[0]} of {scores[0]}; {model.metrics_names[1]} of {scores[1]*100}%')
        acc_per_fold.append(scores[1] * 100)
        loss_per_fold.append(scores[0])
        fn += 1

    return acc_per_fold, loss_per_fold, allHist


"""计算结果"""
from solve import solve

def quiz(show_style, is_int=True):
    """测验"""
    print("注意！")
    print("若结果为分数不需转换为小数。")
    if is_int:
        path = "problem_int.txt"
    else:
        path = "problem_fraction.txt"
    file = open(path, "r", encoding="utf8")
    problems = file.readlines()
    total = len(problems)
    correct = 0
    for problem in problems:
        problem = problem.strip('\n')
        if show_style == 1:
            print(problem.replace("^", "**"))
        else:
            print(problem)
        correct_answer = solve(problem.split(" "))
        while True:
            answer = input("请输入答案： ")
            if answer != "":
                break
        if answer == str(correct_answer):
            correct += 1
            print("此题正确")
        else:
            print("此题错误")
    file.close()
    print("测试结束！")
    print("你答对的题目数量为：", correct, "/", total)


import copy
import re

from knack.log import get_logger

from azdev.utilities import get_name_index

logger = get_logger(__name__)

_LOADER_CLS_RE = re.compile('.*azure/cli/command_modules/(?P<module>[^/]*)/__init__.*')

def filter_modules(command_loader, help_file_entries, modules=None, include_whl_extensions=False):
    """ Modify the command table and help entries to only include certain modules/extensions.

    : param command_loader: The CLICommandsLoader containing the command table to filter.
    : help_file_entries: The dict of HelpFile entries to filter.
    : modules: [str] list of module or extension names to retain.
    """
    return _filter_mods(command_loader, help_file_entries, modules=modules,
                        include_whl_extensions=include_whl_extensions)

def exclude_commands(command_loader, help_file_entries, module_exclusions, include_whl_extensions=False):
    """ Modify the command table and help entries to exclude certain modules/extensions.

    : param command_loader: The CLICommandsLoader containing the command table to filter.
    : help_file_entries: The dict of HelpFile entries to filter.
    : modules: [str] list of module or extension names to remove.
    """
    return _filter_mods(command_loader, help_file_entries, modules=module_exclusions, exclude=True,
                        include_whl_extensions=include_whl_extensions)

def _filter_mods(command_loader, help_file_entries, modules=None, exclude=False, include_whl_extensions=False):
    modules = modules or []

    # command tables and help entries must be copied to allow for seperate linter scope
    command_table = command_loader.command_table.copy()
    command_group_table = command_loader.command_group_table.copy()
    command_loader = copy.copy(command_loader)
    command_loader.command_table = command_table
    command_loader.command_group_table = command_group_table
    help_file_entries = help_file_entries.copy()
    name_index = get_name_index(include_whl_extensions=include_whl_extensions)

    for command_name in list(command_loader.command_table.keys()):
        try:
            source_name, _ = _get_command_source(command_name, command_loader.command_table)
        except LinterError as ex:
            # command is unrecognized
            logger.warning(ex)
            source_name = None

        try:
            long_name = name_index[source_name]
            is_specified = source_name in modules or long_name in modules
        except KeyError:
            is_specified = False
        if is_specified == exclude:
            # brute force method of ignoring commands from a module or extension
            command_loader.command_table.pop(command_name, None)
            help_file_entries.pop(command_name, None)

    # Remove unneeded command groups
    retained_command_groups = {' '.join(x.split(' ')[:-1]) for x in command_loader.command_table}
    excluded_command_groups = set(command_loader.command_group_table.keys()) - retained_command_groups

    for group_name in excluded_command_groups:
        command_loader.command_group_table.pop(group_name, None)
        help_file_entries.pop(group_name, None)

    return command_loader, help_file_entries

def share_element(first_iter, second_iter):
    return any(element in first_iter for element in second_iter)

def _get_command_source(command_name, command_table):
    from azure.cli.core.commands import ExtensionCommandSource  # pylint: disable=import-error
    command = command_table.get(command_name)
    # see if command is from an extension
    if isinstance(command.command_source, ExtensionCommandSource):
        return command.command_source.extension_name, True
    if command.command_source is None:
        raise LinterError('Command: `%s`, has no command source.' % command_name)
    # command is from module
    return command.command_source, False

# pylint: disable=line-too-long
def merge_exclusion(left_exclusion, right_exclusion):
    for command_name, value in right_exclusion.items():
        for rule_name in value.get('rule_exclusions', []):
            left_exclusion.setdefault(command_name, {}).setdefault('rule_exclusions', []).append(rule_name)
        for param_name in value.get('parameters', {}):
            for rule_name in value.get('parameters', {}).get(param_name, {}).get('rule_exclusions', []):
                left_exclusion.setdefault(command_name, {}).setdefault('parameters', {}).setdefault(param_name, {}).setdefault('rule_exclusions', []).append(rule_name)

class LinterError(Exception):
    """
    Exception thrown by linter for non rule violation reasons
    """
    pass  # pylint: disable=unnecessary-pass


import tqdm
import torch
import csv
import os
import os.path as osp
import random
import json
import h5py
import time

from collections import defaultdict

if __name__ == '__main__':
    from MiniImageNet import MiniImageNetDataset, TransformedImageLoader, h5load
    from base import MultiProcessImageLoader
else:
    from .base import MultiProcessImageLoader
    from .MiniImageNet import MiniImageNetDataset, TransformedImageLoader, h5load

def shuffle_slice(a, start, stop):
    assert stop <= len(a)
    i = start
    while (i < stop-1):
        idx = random.randrange(i, stop)
        a[i], a[idx] = a[idx], a[i]
        i += 1

class MiniImageNetMAMLDataset(MiniImageNetDataset):

    def __init__(self,
                 csv_dir,
                 split,
                 image_dir,
                 shuffle=True,
                 num_threads=2,
                 transform=None,
                 transformed_images=None,
                 imname_index_file=None,
                 separate_metasample={'train':True, 'test':False, 'val':False},
                 metasample_holdout=200,
                 ):

        super(MiniImageNetMAMLDataset, self).__init__(csv_dir, split, image_dir, shuffle, num_threads, transform,
                                                      transformed_images, imname_index_file)

        assert len(self.split_names) == len(separate_metasample)
        self.separate_metasample = separate_metasample
        #for i in range(len(separate_metasample)):
        #    self.separate_metasample.update({self.split_names[i]: separate_metasample[i]})
        self.metasample_holdout = metasample_holdout

        self.cat_ixs_meta = dict()      # holds current read index for the meta-update images

        for sp in self.split_names:
            if self.separate_metasample[sp]:
                self.cat_ixs_meta[sp] = dict()
            else:
                self.cat_ixs_meta[sp] = self.cat_ixs[sp]            # make the meta index refer to the normal index
            if self.separate_metasample[sp]:
                for cat in self.cat_to_files[sp].keys():
                    self.cat_ixs_meta[sp][cat] = len(self.cat_to_files[sp][cat]) - self.metasample_holdout

    def shuffle(self, split, cat):
        # shuffle all samples of given split and category
        if 'separate_metasample' in self.__dict__:
            if self.separate_metasample[split]:
                # shuffle separately
                self.shuffle_sq(split, cat)
                self.shuffle_meta(split, cat)
        else:
            random.shuffle(self.cat_to_files[split][cat])

    def shuffle_sq(self, split, cat):
        #print('sq shuffle', split, cat)
        num_files = len(self.cat_to_files[split][cat])
        shuffle_slice(self.cat_to_files[split][cat], 0, num_files - self.metasample_holdout)

    def shuffle_meta(self,split, cat):
        #print('meta shuffle', split, cat)
        num_files = len(self.cat_to_files[split][cat])
        shuffle_slice(self.cat_to_files[split][cat], num_files - self.metasample_holdout, num_files)

    def update_sq_iterator(self, split, cat):
        ix = self.cat_ixs[split][cat]
        ix += 1
        maxlen = len(self.cat_to_files[split][cat])
        if self.separate_metasample[split]:
            maxlen -= self.metasample_holdout
        if ix == maxlen:
            self.cat_ixs[split][cat] = 0
            if self.do_shuffle:
                if self.separate_metasample[split]:
                    self.shuffle_sq(split, cat)
                else:
                    self.shuffle(split, cat)
        else:
            self.cat_ixs[split][cat] = ix

    def update_meta_iterator(self, split, cat):
        ix = self.cat_ixs_meta[split][cat]
        ix += 1

        if ix == len(self.cat_to_files[split][cat]):
            if self.separate_metasample[split]:
                self.cat_ixs_meta[split][cat] = len(self.cat_to_files[split][cat]) - self.metasample_holdout
                if self.do_shuffle:
                    self.shuffle_meta(split, cat)
            else:
                self.cat_ixs_meta[split][cat] = 0
                if self.do_shuffle:
                    self.shuffle(split, cat)
        else:
            self.cat_ixs_meta[split][cat] = ix

    def fetch_images_MAML(self, split, cat, num, metanum):
        image_files = []
        for _ in range(num):
            ix = self.cat_ixs[split][cat]
            file = self.cat_to_files[split][cat][ix]
            image_files.append(file)
            self.update_sq_iterator(split, cat)
        for _ in range(metanum):
            ix = self.cat_ixs_meta[split][cat]
            file = self.cat_to_files[split][cat][ix]
            image_files.append(file)
            self.update_meta_iterator(split, cat)
        return image_files

    def get_episode_MAML(self, nway=5, nshot=4, nquery=1, nmeta=15, split='train', transform=None):
        categories = self.sample_categories(split, nway)
        images = dict()
        support_data = []
        support_lbls = []
        support_cats = []
        query_data = []
        query_lbls = []
        query_cats = []
        metasample_data = []
        metasample_lbls = []
        metasample_cats = []
        for lbl, cat in enumerate(categories):
            images[cat] = self.fetch_images_MAML(split, cat, nshot+nquery, nmeta)
            supports = images[cat][:nshot]
            queries = images[cat][nshot:nshot+nquery]
            metasamples = images[cat][nshot+nquery:]

            support_data.extend(supports)
            query_data.extend(queries)
            metasample_data.extend(metasamples)

            support_lbls.extend([lbl] * nshot)
            query_lbls.extend([lbl] * nquery)
            metasample_lbls.extend([lbl] * nmeta)

            support_cats.extend([cat] * nshot)
            query_cats.extend([cat] * nquery)
            metasample_cats.extend([cat] * nmeta)

        data_files = support_data + query_data + metasample_data
        data = self.load_images(data_files, transform)
        lbls = support_lbls + query_lbls + metasample_lbls
        cats = support_cats + query_cats + metasample_cats

        return {'data': data,
                'label': torch.LongTensor(lbls),
                'file': data_files,
                'category': cats,
                'nway': nway,
                'nshot': nshot,
                'nquery': nquery,
                'nmeta': nmeta,
                'split': split}

if __name__ == '__main__':
    import torchvision.transforms as transforms
    csv_dir = 'data/splits/mini_imagenet_split'
    split = 'Ravi'
    image_dir = 'data/raw/mini-imagenet'
    shuffle = True
    num_threads = 4
    transform = transforms.Compose([transforms.Resize((84, 84)),
                                    transforms.ToTensor()])
    separate_metasample = {'train': True, 'test': True, 'val': False}
    mini_loader = MiniImageNetMAMLDataset(csv_dir, split,
                                      image_dir, shuffle,
                                      num_threads, transform, separate_metasample=separate_metasample)
    for idx in tqdm.tqdm(range(100)):
        data = mini_loader.get_episode_MAML(20, 5, 15, 5, split='test')

        #data = mini_loader.get_episode_MAML(5, 4, 1, 15, split='test')
        #import ipdb
        #ipdb.set_trace()
    for idx in tqdm.tqdm(range(1000)):
        data = mini_loader.get_episode_MAML(5, 4, 1, 15, split='test')


from swift.common.swob import wsgify, HTTPInternalServerError, HTTPException
from swift.common.utils import get_logger
from zion.handlers import ProxyHandler
from zion.handlers import ComputeHandler
from zion.handlers import ObjectHandler
from zion.handlers.base import NotFunctionRequest
from distutils.util import strtobool
import redis

class FunctionHandlerMiddleware(object):

    def __init__(self, app, conf):
        self.app = app
        self.conf = conf
        self.exec_server = self.conf.get('execution_server')
        self.logger = get_logger(conf, name=self.exec_server +
                                 "-server Zion",
                                 log_route='function_handler')
        redis_host = self.conf.get('redis_host')
        redis_port = self.conf.get('redis_port')
        redis_db = self.conf.get('redis_db')
        self.redis_conn_pool = redis.ConnectionPool(host=redis_host,
                                                    port=redis_port,
                                                    db=redis_db)

        self.handler_class = self._get_handler(self.exec_server)

    def _get_handler(self, exec_server):
        """
        Generate Handler class based on execution_server parameter

        :param exec_server: Where this storlet_middleware is running.
                            This should value should be 'proxy' or 'compute'
        :raise ValueError: If exec_server is invalid
        """
        if exec_server == 'proxy':
            return ProxyHandler
        elif exec_server == 'compute':
            return ComputeHandler
        elif exec_server == 'object':
            return ObjectHandler
        else:
            raise ValueError('configuration error: execution_server must be '
                             'either proxy, object or compute but is %s' % exec_server)

    @wsgify
    def __call__(self, req):
        try:
            r = redis.Redis(connection_pool=self.redis_conn_pool)
            handler = self.handler_class(req, self.conf, self.app, self.logger, r)
            self.logger.debug('%s call in %s' % (req.method, req.path))

            return handler.handle_request()

        except NotFunctionRequest:
            self.logger.debug('No Zion Request, bypassing middleware')
            return req.get_response(self.app)

        except HTTPException:
            self.logger.exception('Zion execution failed')
            raise

        except Exception:
            self.logger.exception('Zion execution failed')
            raise HTTPInternalServerError(body='Unable to execute the requested operation: Zion execution failed\n')

def filter_factory(global_conf, **local_conf):
    """Standard filter factory to use the middleware with paste.deploy"""
    conf = global_conf.copy()
    conf.update(local_conf)
    # Common
    conf['execution_server'] = conf.get('execution_server')
    conf['functions_container'] = conf.get('functions_container', 'functions')
    conf['functions_visibility'] = strtobool(conf.get('functions_visibility', 'True'))
    # Paths
    conf['main_dir'] = conf.get('main_dir', '/opt/zion')
    # Worker paths
    conf['workers_dir'] = conf.get('workers_dir', 'workers')
    conf['java_runtime_dir'] = conf.get('java_runtime_dir', 'runtime/java')
    # Function Paths
    conf['functions_dir'] = conf.get('functions_dir', 'functions')
    conf['cache_dir'] = conf.get('cache_dir', 'cache')
    conf['log_dir'] = conf.get('log_dir', 'logs')
    conf['bin_dir'] = conf.get('bin_dir', 'bin')
    # Redis metastore
    conf['redis_host'] = conf.get('redis_host', 'localhost')
    conf['redis_port'] = int(conf.get('redis_port', 6379))
    conf['redis_db'] = int(conf.get('redis_db', 10))
    # Function defaults
    conf['default_function_timeout'] = int(conf.get('default_function_timeout', 10))
    conf['default_function_memory'] = int(conf.get('default_function_memory', 1024))
    conf['max_function_memory'] = int(conf.get('max_function_memory', 1024))
    # Compute Nodes
    conf['disaggregated_compute'] = strtobool(conf.get('disaggregated_compute', 'True'))
    conf['compute_nodes'] = conf.get('compute_nodes', 'localhost:8585')
    conf['docker_pool_dir'] = conf.get('docker_pool_dir', 'docker_pool')

    def swift_functions(app):
        return FunctionHandlerMiddleware(app, conf)

    return swift_functions


import numpy
screen = [[0]*120]*40
screen = numpy.array(screen)
objbrick = []
#component = [[]*120]*40
paddle = "^^^^^^"
brick1 = "▭"
ball = "◦"
exppaddle = "@"
shrpaddle = "*"
fastball = "&"
thruball = "%"
paddlegrab = "$"
ballmultiply = "#"
bullet = "."
shoot1 = "!"
bomb= "+"
ufo ="(__)"
objballmultiply = []
objpaddle = []
objball = []
objexppaddle = []
objshrpaddle = []
objthruball = []
objfastball = []
objpaddlegrab = []
objshoot = []
objbullet = []
objboss = []
objbombs = []
flag = 0
lives = 10
timelimit = 20
relball = 0
newlife = 0
health = 10
# yinc = -1
score  = 0
time = 0
paddlelength = 6
level = 1
timelimitlvl = 25
tilvl = 0
tatlvl = 0
decrease = 0
shoot = 0
timepowerup = 0


from graphviz import Digraph
import pandas as pd
import numpy as np
import glob
import os

#Associa cada cor na planilha a um par (cor de fundo,cor da fonte) do Graphviz
colors = {'Amarelo':('yellow','black'),'Azul':('blue','white'),'Branco':('white','black'),
         'Cinza':('grey','black'),'Marrom':('brown','white'),'Ouro':('gold3','black'),
         'Preto':('black','white'),'Roxo':('purple','white'),'Verde':('green','black'),
         'Vermelho':('red','black'),'Laranja':('orangered','black')}
colors_sub = { 'Amarelo':'#ffff005f','Azul':'cyan','Branco':'white','Cinza':'gray95',
             'Marrom':'#a52a2a5f','Ouro':'gold','Preto':'black','Roxo':'mediumpurple1',
             'Verde':'palegreen','Vermelho':'tomato','Laranja':'orange' }

#Associa os tipos de relação na planilha aos vértices no Graphviz
rel_types = {'Positivo':'blue', 'Negativo':'red', 'Neutro':'black'}
dirs = {'Sim':'both', 'Não':'forward'}

'''
Obtém o data frame do Pandas
Recebe o nome do arquivo Excel
Retorna as listas de atores e relacionamentos
'''
def get_data(fileloc, ext):
    #lendo o arquivo usando Pandas
    ext_engine={'.xls':'xlrd','.xlsx':'openpyxl','.ods':'odf'}
    try:
        df_at_rel = pd.read_excel(fileloc, ['atores','relacionamentos'], engine=ext_engine.get(ext,None))
    except:
        return None, None

    #Obtendo os atores e relacionamentos
    df_atores = df_at_rel['atores'].loc[:,['ator', 'cor', 'grupo']].dropna(subset=['ator']).replace({np.nan: None})
    df_relacionamentos = df_at_rel['relacionamentos'].loc[:,['de', 'relacionamento', 'para', 'tipo', 'bilateral']].dropna()
    return df_atores.to_numpy().tolist(), df_relacionamentos.to_numpy().tolist()

'''
Faz os nós dos atores
Recebe o objeto Digraph e a lista de atores
Retorna o Dicionário com os grupos e membros
'''
def make_actor_nodes(graph,atores):
    lst_grupos = set()
    for ator in atores:
        ator_nome, ator_cor, ator_grupo = ator
        if ator_grupo is not None: lst_grupos.add(ator_grupo)
    grupos = {}
    grupos_cores = {}
    for ator in atores:
        #expandindo os parâmetros da lista
        ator_nome, ator_cor, ator_grupo = ator
        if ator_cor is None: ator_cor = "Branco"
        if ator_nome in lst_grupos:
            grupos_cores[ator_nome] = ator_cor
        else:
            graph.node(ator_nome,shape='circle',fillcolor=colors[ator_cor][0],style='filled', fontcolor=colors[ator_cor][1], fixedsize='false',width='1')
        if ator_grupo is None: continue
        if ator_grupo not in grupos:
            grupos[ator_grupo] = list()
        grupos[ator_grupo].append(ator_nome)
    return grupos, grupos_cores

'''
Cria um subgrafo
Recebe o nome e cor do grupo
Retorna o subgrafo criado
'''
def makeGroup(grupo_nome, grupo_cor, nivel=1 ):
    sub_g = Digraph(name='cluster_'+grupo_nome)
    sub_g.attr(label=grupo_nome)
    sub_g.attr(style='rounded')
    sub_g.attr(bgcolor=colors_sub[grupo_cor])
    return sub_g

'''
Insere os nós e grupos em um subgrafo
'''
def fillGroup(sub_g, grupos, grupos_cores, membros, not_root, nivel=1 ):
    for membro in membros:
        if membro not in grupos.keys():
            sub_g.node(membro)
        else:
            new_g = makeGroup(membro, grupos_cores.get(membro,'Branco'), nivel+1)
            fillGroup(new_g, grupos, grupos_cores, grupos[membro], not_root, nivel+1)
            sub_g.subgraph(new_g)
            not_root.add(membro)

'''
Faz os agrupamentos
Recebe o objeto Digraph e Dicionário com os grupos e membros
'''
def makeGroups(graph, grupos, grupos_cores):
    sub_graphs = {}
    not_root = set()
    #Cria todos os grupos
    for grupo,membros in grupos.items():
        sub_g = makeGroup(grupo, grupos_cores.get(grupo,'Branco'), 1)
        fillGroup(sub_g, grupos, grupos_cores, membros, not_root, 1)
        sub_graphs[grupo]=sub_g

    #Insere tudo no grafo principal
    for k,v in sub_graphs.items():
        if k not in not_root:
            graph.subgraph(v)

def find_actor(grupos, grupo):
    for membro in grupos[grupo]:
        if membro not in grupos.keys():
            return membro
        else:
            return find_actor(grupos, membro)

'''
Cria os relacionamentos
Recebe o objeto Digraph, a lista de relacionamentos e o Dicionário com os grupos e membros
'''
def make_relationships(graph, relacionamentos, grupos):
    for relacionamento in relacionamentos:
        #expandindo os parâmetros da lista
        de, legenda, para, tipo, bilateral = relacionamento
        #Parâmetros básicos
        params = {'dir':dirs[bilateral], 'color':rel_types[tipo], 'fontcolor':rel_types[tipo], 'penwidth':'1.0', 'decorate':'false', 'minlen':'2'}

        #Ajustando caso o relacionamento envolva clusters
        if de in grupos.keys():
            #Criando a chave ltail para o relacionamento com o cluster
            params['ltail'] = 'cluster_'+de
            #Pegando um ator qualquer (o primeiro) dentro do cluster
            deN = find_actor(grupos, de)
        else:
            deN = de
        if para in grupos.keys():
            #Criando a chave lhead para o relacionamento com o cluster
            params['lhead'] = 'cluster_'+para
             #Pegando um ator qualquer (o primeiro) dentro do cluster
            paraN = find_actor(grupos, para)
        else:
            paraN = para
        graph.edge(deN, paraN, label=legenda, **params)

'''
Loop principal
'''
#Monta a lista arquivos .xls, .xlsx e .ods na pasta
files = {f:glob.glob('*'+f) for f in ['.ods','.xls','.xlsx']}

#Cria o gráfico para cada arquivo encontrado
for ext,lista_f in files.items():
    for fileloc in lista_f:
        #ignora arquivos temporários do Excel
        if fileloc.startswith('~$'):
            continue

        print("Processando arquivo: " + fileloc)

        k = fileloc.rfind(ext)
        gattr = dict()
        #allow edges between clusters
        gattr['compound'] = 'true'
        gattr['rankdir'] = 'LR'
        gattr['ranksep'] = '1.0'
        #gattr['dpi'] = '300'
        #gattr['ratio'] = '0.5625'
        gattr['newrank'] = 'true'
        gattr['overlap'] = 'false'
        gattr['fontsize'] = '20'

        g = Digraph(filename=fileloc[:k], engine='dot', format='png', graph_attr=gattr)

        #Constroi os dataframes do excel
        atores, relacionamentos = get_data(fileloc, ext)
        #Checa se get_data foi bem sucedida
        if atores is None and relacionamentos is None:
            continue
        #Faz os nós dos atores
        grupos, grupos_cores = make_actor_nodes(g, atores)
        #Cria os clusters
        makeGroups(g, grupos, grupos_cores)
        #Faz os relacionamentos
        make_relationships(g, relacionamentos, grupos)
        #Cria e abre o arquivo
        g.view()
        #Remove arquivo intermediário
        os.remove(fileloc[:k])


import random

def game_rules():
    print("Winning Rules of the Rock paper scissor game as follows: \n"
                                +"Rock vs paper->paper wins \n"
                                + "Rock vs scissor->Rock wins \n"
                                +"paper vs scissor->scissor wins \n")
game_rules()
def gaming_function() :

    #a list to store all possible options:
    choices = ["r", "p" , "s"]

    #ask the user to pick an option between "R", "P" or "S"
    print("Enter choice \n R for Rock, \n P for paper, and \n S for scissor \n ")

    choice = input("Your turn: ").lower()

    #search if choice is in available in the choices list
    def search(list, choice):
        for i in range(len(list)):
            if list[i] == choice:
                return True
        return False
    while search(choices, choice) == False:
        choice = input("enter valid input: ").lower()

    if choice == "r":
        choice_name = 'Rock'
    elif choice == "p":
        choice_name = 'Paper'
    else:
        choice_name = 'Scissor'
    print("Your choice is: " + choice_name)
    print("\nNow its computer turn.......")

    # Computer randomly chooses any option
    comp_choice = random.choice(choices)

    if comp_choice == "r":
        comp_choice_name = 'Rock'
    elif comp_choice == "p":
        comp_choice_name = 'Paper'
    else:
        comp_choice_name = 'Scissor'

    print("Computer choice is: " + comp_choice_name)

    print(choice_name + " V/s " + comp_choice_name)

    if choice == comp_choice:
        print("Draw=> ", end = "")
        result = "Draw"

    elif ((choice == "r" and comp_choice == 'p') or (choice == 'p' and comp_choice == 'r')):
        print("paper wins => ", end = "")
        result = "Paper"
    elif ((choice == "s" and comp_choice == 'r') or (choice == 'r' and comp_choice == 's')):
        print("rock wins => ", end = "")
        result = "Rock"
    else:
        print("scissor wins => ", end = "")
        result = "Scissor"

    def announce_winner(result) :
        if result == "Draw":
            print('It is a tie, so close!! try again' )
            gaming_function()
        elif result == choice_name:
            print("You win!! Iwolokan")
        else:
            print ("Opps!! Computer wins")

    announce_winner(result)

gaming_function()