from os.path import exists, basename, join, dirname

import nibabel as nib
import numpy as np
from scipy.signal import cubic
from scipy.sparse import kron, lil_array
from traitlets import (Instance,
                       Integer,
                       Unicode,
                       Bool,
                       Dict,
                       TraitError,
                       Enum,
                       Float)
from traitlets.config import Config

from mrHARDI.base.application import (mrHARDIBaseApplication,
                                           MultipleArguments, output_file_argument,
                                           output_prefix_argument,
                                           required_arg,
                                           required_file)
from mrHARDI.base.dwi import (AcquisitionType,
                                   Direction,
                                   DwiMetadata,
                                   DwiMismatchError,
                                   load_metadata_file,
                                   load_metadata,
                                   save_metadata)
from mrHARDI.compute.dwi import identify_shells, sh_order_from
from mrHARDI.compute.utils import resampling_affine
from mrHARDI.config.utils import DwiMetadataUtilsConfiguration

_mb_aliases = {
    'in': 'DwiMetadataUtils.dwis',
    'out': 'DwiMetadataUtils.output_folder',
    'suffix': 'DwiMetadataUtils.suffix',
    'json': 'DwiMetadataUtils.json_config'
}
_mb_flags = dict(
    overwrite=(
        {"DwiMetadataUtils": {'overwrite': True}},
        "Force overwriting of output images if present"
    ),
    update=(
        {"DwiMetadataUtils": {'update': True}},
        "Updates the already present metadata files"
    ),
    update_affine=(
        {"DwiMetadataUtils": {'update_affine': True}},
        "Only update affine file in metadata file"
    )
)

class DwiMetadataUtils(mrHARDIBaseApplication):
    name = u"DWI Metadata Utilities"
    description = (
        "Generates json metadata description files for the input datasets. "
        "This allow for easier management of multiband data, oscillating "
        "gradient and tensor-valued acquisitions, as well as other acquisition"
        "properties unavailable from the Nifti file alone."
    )

    configuration = Instance(DwiMetadataUtilsConfiguration).tag(config=True)

    dwis = required_arg(
        MultipleArguments, [],
        "Dwi volumes for which to create metadata files",
        traits_args=(Unicode(),)
    )

    suffix = Unicode(
        help="Suffix to append to metadata filenames"
    ).tag(config=True)

    output_folder = Unicode(
        help="Optional output folder for metadata files. Defaults to the same "
             "folder where the input dwis are."
    ).tag(config=True)

    json_config = Unicode(
        help="Json configuration file replacing the arguments"
    ).tag(config=True)

    metadata = Unicode(
        help="Force an update on a specific metadata file. Will "
             "output to a metadata file based on input image name."
    ).tag(config=True)

    update = Bool(
        False, help="If true, will update metadata files "
                    "based on given new parameters"
    ).tag(config=True)

    update_affine = Bool(
        False, help="If true, only update affine data"
    ).tag(config=True)

    overwrite = Bool(
        False, help="If True, overwrites the output file when writing"
    ).tag(config=True)

    aliases = Dict(default_value=_mb_aliases)
    flags = Dict(default_value=_mb_flags)

    def _validate_configuration(self):
        if self.update_affine:
            return

        if not self.update and self.configuration.acquisition is None:
            self.configuration.acquisition = AcquisitionType.Linear.name

        if self.json_config:
            assert exists(self.json_config)
            config = (
                self.json_config_loader_class(self.json_config).load_config()
            )
            self.update_config(self._split_config(dict(config)))

        if (
            self.configuration.multiband_factor and
            self.configuration.multiband_factor > 1
        ):
            self.configuration.traits()["slice_direction"].tag(
                required=True
            )

        super()._validate_configuration()

    def _split_config(self, config):
        traits = self.traits(config=True)
        conf_traits = self.configuration.traits(config=True)
        configuration = {
            self.__class__.__name__: {},
            self.configuration.__class__.__name__: {}
        }

        for k, v in config.items():
            if k in traits:
                configuration[self.__class__.__name__][k] = v
            elif k in conf_traits:
                configuration[self.configuration.__class__.__name__][k] = v
            else:
                raise TraitError(
                    "Trait {} not found in {} nor {}".format(
                        k, self.__class__.name,
                        self.configuration.__class__.name
                    )
                )

        return Config(configuration)

    def _get_slice_indexes(self, images):
        directions = self.get_slice_directions(images)
        idxs = np.array([
            np.pad(
                np.arange(img["data"].shape[np.absolute(d).argmax()]),
                (0, img["data"].shape[np.absolute(d).argmax()] %
                 self.configuration.multiband_factor
                 ),
                mode="constant",
                constant_values=-1
            ).reshape(
                (self.configuration.multiband_factor, -1)
            ).T.astype(int) for img, d in zip(images, directions)
        ])

        if self.configuration.interleaved:
            idxs = np.hstack((idxs[:, ::2, ...], idxs[:, 1::2, ...]))

        return [
            [list(filter(lambda i: not i == -1, k)) for k in ki]
            for ki in idxs
        ]

    def _shape_at_least_4D(self, image):
        return 1 if len(image.shape) < 4 else image.shape[-1]

    def _expand_to_slices(self, directions, images):
        return [
            np.repeat([direction], self._shape_at_least_4D(img["data"]), 0)
            for img, direction in zip(images, directions)
        ]

    def _unpack_directions(
        self, images, directions, val_extractor=lambda v: v.value
    ):
        if len(directions) == 1:
            d = directions[0]
            return self._expand_to_slices(
                np.repeat(
                    [val_extractor(Direction[d])],
                    len(images), 0
                ),
                images
            )
        elif len(directions) == len(images):
            return self._expand_to_slices(
                [
                    val_extractor(Direction[d])
                    for d in directions
                ],
                images
            )
        else:
            return []

    def get_phase_directions(self, images, val_extractor=lambda v: v.value):
        if len(self.configuration.direction) == 1:
            d = self.configuration.direction[0]
            return [{
                "dir": val_extractor(Direction[d]),
                "range": (0, self._shape_at_least_4D(i["data"]))
            } for i in images]
        elif len(self.configuration.direction) == len(images):
            return [{
                "dir": val_extractor(Direction[d]),
                "range": (0, self._shape_at_least_4D(i["data"]))
            } for d, i in zip(self.configuration.direction, images)]
        else:
            return []

    def get_slice_directions(self, images, val_extractor=lambda v: v.value):
        return self._unpack_directions(
            images, self.configuration.slice_direction, val_extractor
        )

    def preload_images(self):
        return [
            {"data": nib.load(name), "name": name} for name in self.dwis
        ]

    def _get_file_for(self, image_name):
        name = "{}_metadata".format(basename(image_name).split(".")[0])
        name = "{}_{}".format(name, self.suffix) if self.suffix else name
        return join(
            self.output_folder if self.output_folder else dirname(image_name),
            name
        )

    def _only_update_affine(self, images):
        base_meta = load_metadata_file(self.metadata) \
            if self.metadata else None

        for img in images:
            mt = base_meta.copy() if base_meta else load_metadata(img["name"])
            mt.affine = img["data"].affine.tolist()
            save_metadata(img["name"].split(".")[0], mt)

    def execute(self):
        images = self.preload_images()

        if self.update_affine:
            self._only_update_affine(images)

        directions = self.get_phase_directions(images)

        if self.configuration.multiband_factor in [None, 0]:
            self.configuration.multiband_factor = 1

        slice_indexes = self.configuration.slice_indexes
        if slice_indexes is None:
            slice_indexes = self._get_slice_indexes(images)
        else:
            slice_indexes = [slice_indexes for _ in range(len(images))]

        slice_dirs = self.configuration.slice_direction
        if len(self.configuration.slice_direction) == 1:
            slice_dirs = np.repeat(
                self.configuration.slice_direction, len(images)
            ).tolist()

        for name, img, d, ss, sd in zip(
            self.dwis, images, directions, slice_indexes, slice_dirs
        ):
            shape = img["data"].shape
            metadata = load_metadata(name) if self.metadata else DwiMetadata()
            metadata.n = shape[-1] if len(shape) > 3 else 1
            metadata.n_excitations = int(shape[
                np.argmax(np.absolute(Direction[sd].value))
            ] / self.configuration.multiband_factor)
            metadata.affine = img["data"].affine.tolist()
            metadata.directions = [d]
            metadata.readout = self.configuration.readout

            metadata.is_multiband = (
                self.configuration.multiband_factor and
                self.configuration.multiband_factor > 1
            )
            metadata.slice_order = ss if ss is not None else []
            metadata.multiband_corrected = \
                self.configuration.multiband_corrected or False

            metadata.is_tensor_valued = \
                self.configuration.tensor_valued or False
            metadata.acquisition_types = [
                AcquisitionType[self.configuration.acquisition].value
            ]
            metadata.acquisition_slices = [[0, None]]
            metadata.number_of_coils = self.configuration.number_of_coils

            metadata.generate_config_file(self._get_file_for(img["name"]))

_assert_aliases = {
    "in": "AssertDwiDimensions.dwi",
    "bvals": "AssertDwiDimensions.bvals",
    "bvecs": "AssertDwiDimensions.bvecs",
    "strat": "AssertDwiDimensions.strategy",
    "out": "AssertDwiDimensions.output",
    "ceil": "AssertDwiDimensions.b0_threshold"
}

class AssertDwiDimensions(mrHARDIBaseApplication):
    dwi = required_arg(
        Unicode,
        description="Input dwi file to which to compare bvals and bvecs"
    )

    bvals = required_arg(Unicode, description="Input b-values file")
    bvecs = required_arg(Unicode, description="Input b-vectors file")

    strategy = Enum(
        ["error", "fix"], "error",
        help="Strategy to employ when discrepancies are found, either "
             "output an error message or try to fix the problem"
    ).tag(config=True)

    b0_threshold = Integer(
        default_value=20, help="Upper threshold for b-values considered as b0"
    ).tag(config=True)

    output = output_prefix_argument(
        description="Output prefix for checked files. If supplied, will also "
                    "copy the input dwi to the new name. If none supplied, "
                    "will overwrite the input files.",
        required=False
    )

    aliases = Dict(default_value=_assert_aliases)

    def execute(self):
        img = nib.load(self.dwi)
        bvals, bvecs = np.loadtxt(self.bvals), np.loadtxt(self.bvecs).T

        if self.strategy == "error":
            if img.shape[-1] != len(bvals) != len(bvecs):
                raise DwiMismatchError(img.shape[-1], len(bvals), len(bvecs))
        else:
            if len(bvals) > img.shape[-1]:
                bvals = bvals[:img.shape[-1]]
            elif img.shape[-1] != len(bvals):
                if img.shape[-1] % len(bvals) == 0:
                    bvals = np.repeat(bvals, int(img.shape[-1] / len(bvals)))
                else:
                    raise DwiMismatchError(
                        img.shape[-1], len(bvals), len(bvecs),
                        "Could not fix b-values of the supplied dataset"
                    )

            if bvecs.shape[0] > len(bvals):
                bvecs = bvecs[:len(bvals), :]
            elif bvecs.shape[0] != len(bvals):
                if len(bvals) % bvecs.shape[0] == 0:
                    bvecs = np.repeat(
                        bvecs, int(len(bvals) / bvecs.shape[0]), axis=0
                    )
                else:
                    b0_mask = np.less_equal(bvals, self.b0_threshold)
                    zero_bvecs = np.isclose(np.linalg.norm(bvecs, axis=1), 0)
                    if np.sum(b0_mask) == (
                        len(bvals) - bvecs.shape[0] + np.sum(zero_bvecs)
                    ):
                        new_bvecs = np.zeros((len(bvals), 3))
                        new_bvecs[~b0_mask] = bvecs[~zero_bvecs]
                        bvecs = new_bvecs
                    else:
                        raise DwiMismatchError(
                            img.shape[-1], len(bvals), len(bvecs),
                            "Could not fix b-vectors the supplied dataset"
                        )

            np.savetxt(
                self.bvals if self.output is None else
                "{}.bval".format(self.output),
                bvals[None, :], fmt='%.2f'
            )
            np.savetxt(
                self.bvecs if self.output is None else
                "{}.bvec".format(self.output),
                bvecs.T, fmt='%.8f'
            )

            if self.output:
                nib.save(img, "{}.nii.gz".format(self.output))

                metadata = load_metadata(self.dwi)
                if metadata:
                    metadata.adapt_to_shape(len(bvals))
                    save_metadata(self.output, metadata)

_shells_aliases = {
    "in": "ExtractShells.dwi",
    "bvals": "ExtractShells.bvals",
    "bvecs": "ExtractShells.bvecs",
    "shells": "ExtractShells.shells",
    "count": "ExtractShells.count",
    "keep": "ExtractShells.keep",
    "out": "ExtractShells.output",
    "ceil": "ExtractShells.b0_threshold",
    "gap": "ExtractShells.shell_threshold"
}

_shells_flags = dict(
    with_b0=(
        {"ExtractShells": {"keep_b0": True}},
        'Keeps b0 volumes in the output'
    )
)

class ExtractShells(mrHARDIBaseApplication):
    dwi = required_arg(
        Unicode,
        description="Dwi file to extract shells from"
    )

    bvals = required_arg(Unicode, description="Input b-values file")
    bvecs = required_arg(Unicode, description="Input b-vectors file")

    shells = MultipleArguments(
        Float(), help="Shells group of threshold to work from. If None, "
                      "will execute on all shells, excepts b0 volumes."
    ).tag(config=True)

    count = Integer(
        default_value=0, help="Keep only shells which have a number "
                              "of directions greater than this parameter"
    ).tag(config=True)

    keep = Enum(
        ["all", "geq", "leq", "bigset", "smallset"], "all",
        help="Selection strategy to subset the initial group of shells."
    ).tag(config=True)

    keep_b0 = Bool(
        False, help="When False, removes the b0 volumes from the output volume"
    ).tag(config=True)

    b0_threshold = Integer(
        default_value=20, help="Upper threshold for b-values considered as b0"
    ).tag(config=True)
    shell_threshold = Integer(
        default_value=40, help="Threshold for gaps between shells"
    ).tag(config=True)

    output = output_prefix_argument()

    aliases = Dict(default_value=_shells_aliases)
    flags = Dict(default_value=_shells_flags)

    def execute(self):
        bvals = np.loadtxt(self.bvals)
        bvecs = np.loadtxt(self.bvecs)
        dwi = nib.load(self.dwi)

        b0_mask = np.less_equal(bvals, self.b0_threshold)
        shells, centroids = identify_shells(
            bvals[~b0_mask], self.shell_threshold
        )
        centroids = shells[centroids]
        _, cnt = np.unique(centroids, return_counts=True)

        cnt = np.array(cnt)
        cnt_mask = np.ones_like(centroids, bool)
        for shell, ct in zip(shells, cnt):
            if ct <= self.count:
                cnt_mask[centroids == shell] = False

        centroids = centroids[cnt_mask]
        shells = shells[cnt > self.count]

        if self.shells:
            shells = np.array(self.shells)

        mask = np.zeros_like(centroids, bool)
        if self.keep == "leq":
            mask |= centroids <= shells.max()
        elif self.keep == "geq":
            mask |= centroids >= shells.min()
        elif self.keep == "bigset":
            counts = [(centroids == s).sum() for s in shells]
            mask |= centroids == shells[counts.argmax()]
        elif self.keep == "smallset":
            counts = [(centroids == s).sum() for s in shells]
            mask |= centroids == shells[counts.argmin()]
        elif self.keep == "all":
            for shell in shells:
                mask |= centroids == shell

        extraction_mask = np.zeros_like(bvals, bool)
        cnt_mask[cnt_mask] = mask
        extraction_mask[~b0_mask] = cnt_mask
        if self.keep_b0:
            extraction_mask[b0_mask] = True

        np.savetxt(
            "{}.bval".format(self.output),
            bvals[extraction_mask],
            newline=" ",
            fmt="%d"
        )
        np.savetxt(
            "{}.bvec".format(self.output),
            bvecs[:, extraction_mask],
            fmt="%.8f"
        )
        nib.save(
            nib.Nifti1Image(
                dwi.get_fdata().astype(
                    dwi.get_data_dtype()
                )[..., extraction_mask],
                dwi.affine, dwi.header
            ),
            "{}.nii.gz".format(self.output)
        )

        metadata = load_metadata(self.dwi)
        if metadata:
            acq_types = np.array(
                metadata.acquisition_slices_to_list()
            )[extraction_mask]
            directions = np.array(metadata.directions)[extraction_mask, :]
            metadata.update_acquisition_from_list(acq_types.tolist())
            metadata.directions = directions.tolist()
            metadata.n = int(extraction_mask.sum())
            save_metadata(self.output, metadata)

_flip_aliases = {
    "in": "FlipGradientsOnReference.dwi",
    "bvecs": "FlipGradientsOnReference.bvecs",
    "out": "FlipGradientsOnReference.output"
}

class FlipGradientsOnReference(mrHARDIBaseApplication):
    dwi = required_file(description="Input dwi file used as reference")
    bvecs = required_file(description="Input b-vectors to flip")

    output = output_prefix_argument()

    aliases = Dict(default_value=_flip_aliases)

    def execute(self):
        affine = nib.load(self.dwi).affine[:3, :3]
        flips = np.sign(np.linalg.inv(affine) @ [1, 1, 1]) < 0
        bvecs = np.loadtxt(self.bvecs)
        bvecs[flips, :] *= -1.

        np.savetxt("{}.bvec".format(self.output), bvecs)

_duplicates_aliases = {
    "in": "CheckDuplicatedBvecsInShell.dwi",
    "bvals": "CheckDuplicatedBvecsInShell.bvals",
    "bvecs": "CheckDuplicatedBvecsInShell.bvecs",
    "out": "CheckDuplicatedBvecsInShell.output",
    "merge": "CheckDuplicatedBvecsInShell.merging",
    "abs-thr": "CheckDuplicatedBvecsInShell.abs_threshold",
    "ceil": "CheckDuplicatedBvecsInShell.b0_threshold"
}

class CheckDuplicatedBvecsInShell(mrHARDIBaseApplication):
    dwi = required_file(description="Input dwi file")
    bvals = required_file(description="Input b-values file")
    bvecs = required_file(description="Input b-vectors file")

    merging = Enum(
        ["first", "mean", "median"], "median",
        help="Merge strategy of duplicates"
    ).tag(config=True)

    abs_threshold = Float(
        1E-5, help="Absolute threshold on distance between directions"
    ).tag(config=True)

    b0_threshold = Integer(
        default_value=20, help="Upper threshold for b-values considered as b0"
    ).tag(config=True)

    output = output_prefix_argument()

    aliases = Dict(default_value=_duplicates_aliases)

    _mergers = {
        "first": lambda dt: dt[..., 0],
        "mean": lambda dt: np.mean(dt, -1),
        "median": lambda dt: np.median(dt, -1)
    }

    def execute(self):
        bvals = np.loadtxt(self.bvals)
        bvecs = np.loadtxt(self.bvecs).T
        dwi = nib.load(self.dwi)
        data = dwi.get_fdata().astype(dwi.get_data_dtype())

        merge_fn = self._mergers[self.merging]

        odwi = []
        obvals = []
        obvecs = []
        processed_vols = np.array([False] * data.shape[-1])

        b0_mask = np.less_equal(bvals, self.b0_threshold)
        meta_mask = np.zeros((len(bvals),),  dtype=bool)

        for i in range(data.shape[-1]):
            if not processed_vols[i]:
                processed_vols[i] = True
                meta_mask[i] = True
                obvals.append(bvals[i])
                obvecs.append(bvecs[i])
                if b0_mask[i]:
                    odwi.append(data[..., i, None])
                else:
                    shell_mask = np.isclose(bvals, bvals[i])
                    shell_mask[processed_vols] = False
                    distances = 1. - bvecs[shell_mask] @ bvecs[i]
                    close_idxs = np.where(shell_mask)[0][np.isclose(distances, 0.)]
                    processed_vols[close_idxs] = True
                    if len(close_idxs) == 1:
                        odwi.append(data[..., i, None])
                    else:
                        close_idxs = np.concatenate((close_idxs, [i]))
                        odwi.append(merge_fn(data[..., close_idxs])[..., None])

        np.savetxt("{}.bval".format(self.output), obvals, fmt="%d", newline=" ")
        np.savetxt(
            "{}.bvec".format(self.output), np.array(obvecs).T, fmt="%.6f"
        )
        nib.save(
            nib.Nifti1Image(
                np.concatenate(odwi, axis=-1), dwi.affine, dwi.header
            ),
            "{}.nii.gz".format(self.output)
        )

        metadata = load_metadata(self.dwi)
        if metadata:
            acq_types = np.array(metadata.acquisition_slices_to_list())[meta_mask]
            directions = np.array(metadata.directions)[meta_mask, :]
            metadata.update_acquisition_from_list(acq_types.tolist())
            metadata.directions = directions.tolist()
            metadata.n = int(meta_mask.sum())
            save_metadata(self.output, metadata)

_sh_order_aliases = {
    "bvals": "DetermineSHOrder.bvals",
    "bvecs": "DetermineSHOrder.bvecs",
    "out": "DetermineSHOrder.output",
    "ceil": "DetermineSHOrder.b0_threshold",
    "gap": "DetermineSHOrder.shell_threshold",
    "order": "DetermineSHOrder.sh_order",
}

_sh_order_flags = dict(
    strict=(
        {"DetermineSHOrder": {"strict": True}},
        'Images that do not meet the required SH order cause an error'
    ),
    msmt=(
        {"DetermineSHOrder": {"msmt": True}},
        'Compute the order per shell'
    ),
    full=(
        {"DetermineSHOrder": {"full_basis": True}},
        'Use full SH basis'
    )
)

class DetermineSHOrder(mrHARDIBaseApplication):
    bvals = required_file(description="b-values file")
    bvecs = required_file(description="b-vectors file")
    output = output_file_argument()
    sh_order = required_arg(Integer)

    b0_threshold = Integer(
        default_value=20, help="Upper threshold for b-values considered as b0"
    ).tag(config=True)
    shell_threshold = Integer(
        default_value=40, help="Threshold for gaps between shells"
    ).tag(config=True)
    strict = Bool(
        False, help="When true, images that do not meet the "
                    "required SH order cause an error"
    ).tag(config=True)
    msmt = Bool(False, help="Specify to compute the order per shell").tag(
        config=True
    )
    full_basis = Bool(False).tag(config=True)

    aliases = Dict(default_value=_sh_order_aliases)
    flags = Dict(default_value=_sh_order_flags)

    def execute(self):
        bvals = np.loadtxt(self.bvals)
        bvecs = np.loadtxt(self.bvecs)

        bvals_idxs = np.greater(bvals, self.b0_threshold)
        bvals = bvals[bvals_idxs]
        bvecs = bvecs[:, bvals_idxs]

        if self.msmt:
            num_ubvecs = []
            shells, centroids = identify_shells(bvals, self.shell_threshold)
            centroids = shells[centroids]
            for shell in shells:
                ubv = np.unique(bvecs[:, centroids == shell], axis=1)
                num_ubvecs.append(ubv.shape[1])

            num_ubvecs = min(num_ubvecs)
        else:
            num_ubvecs = np.unique(bvecs, axis=1).shape[1]

        sh_order = sh_order_from(num_ubvecs, self.full_basis)

        if self.strict and sh_order < self.sh_order:
            raise RuntimeError(
                "Insufficent number of volumes for SH order {}".format(
                    self.sh_order
                )
            )

        if sh_order == 0:
            if self.msmt:
                raise RuntimeError(
                    "Insufficent number of volumes in one or more shells "
                    "for SH reconstruction (calculated order = 0)"
                )

            raise RuntimeError(
                    "Insufficent number of volumes for SH "
                    "reconstruction (calculated order = 0)"
                )

        with open(self.output, "w+") as f:
            f.write("{}".format(sh_order))

_bsc_aliases = {
    "in": "ImageBSplineCoefficients.image",
    "spacing": "ImageBSplineCoefficients.knot_spacing",
    "out": "ImageBSplineCoefficients.output"
}

class ImageBSplineCoefficients(mrHARDIBaseApplication):
    image = required_file(
        description="Input image"
    )

    knot_spacing = required_arg(Float)

    output = output_file_argument()

    aliases = Dict(default_value=_bsc_aliases)

    def execute(self):
        img = nib.load(self.image)
        target_shape = img.shape[:3]
        target_zooms = img.header.get_zooms()[:3]

        bspline_zooms = np.repeat(self.knot_spacing, 3)
        bspline_affine = resampling_affine(
            img.affine,
            target_shape,
            target_zooms,
            bspline_zooms
        )
        bspline_shape = target_zooms / bspline_zooms * target_shape

        target_to_grid = np.linalg.inv(bspline_affine.affine) @ img.affine
        wd = []
        for axis in range(3):
            coords = np.zeros((3, target_shape[axis]))
            coords[axis] = np.arange(target_shape[axis])

            locs = nib.affines.apply_affine(target_to_grid, coords.T)[:, axis]
            knots = np.arange(bspline_shape[axis])

            distance = np.abs(locs[np.newaxis, ...] - knots[..., np.newaxis])
            within_support = distance < 2.0
            d_vals, d_idxs = np.unique(distance[within_support], return_inverse=True)
            bs_w = cubic(d_vals)

            colloc_ax = lil_array((bspline_shape[axis], bspline_shape[axis]))
            colloc_ax[within_support] = bs_w[d_idxs]

            wd.append(colloc_ax)

        nib.save(
            nib.Nifti1Image(
                kron(kron(wd[0], wd[1]), wd[2]),
                bspline_affine
            ),
            self.output
        )

_dftf_aliases = {
    "in": "DisplacementFieldToFieldmap.image",
    "readout": "DisplacementFieldToFieldmap.readout",
    "pe": "DisplacementFieldToFieldmap.pe_direction",
    "out": "DisplacementFieldToFieldmap.output"
}

class DisplacementFieldToFieldmap(mrHARDIBaseApplication):
    image = required_file(
        description="Input displacement field (4D with 3 "
                    "components in last dimension aka dx,dy,dz)"
    )

    readout = required_arg(Float)
    pe_direction = Enum(
        ["i", "j", "k", "i-", "j-", "k-"], None
    ).tag(config=True, required=True)

    output = output_file_argument()

    aliases = Dict(default_value=_dftf_aliases)

    def execute(self):
        disp = nib.load(self.image)
        dxdydz = disp.get_fdata(dtype="float32").reshape((-1, 3))
        dxdydz[:, (0, 1)] *= -1.

        inv_rotation = np.linalg.inv(disp.affine)
        inv_rotation[:3, 3] = 0

        didjdk = nib.affines.apply_affine(
            inv_rotation, dxdydz
        ).astype("float32")
        pe_axis = "ijk".index(self.pe_direction[0])
        vsm = didjdk[:, pe_axis].reshape(disp.shape[:3])

        scale = self.readout
        if self.pe_direction.endswith("-"):
            scale *= -1.

        fmap = nib.Nifti1Image(vsm / scale, disp.affine)
        fmap.header.set_intent("estimate", name="Delta_B0 [Hz]")
        fmap.header.set_xyzt_units("mm")
        fmap.header["cal_max"] = max((
            abs(np.asanyarray(fmap.dataobj).min()),
            np.asanyarray(fmap.dataobj).max(),
        ))
        fmap.header["cal_min"] = - fmap.header["cal_max"]

        nib.save(fmap, self.output)


import subprocess, os
from tqdm import tqdm

directory = '/datax/images/'

for profile in os.listdir(directory):
    print('PROCESSING:', profile)
    folder = directory + profile + '/'
    count = 1
    while count:
        if count > 10:
            print('SKIPPING:', profile)
            continue
        try:
            file = open('{}{}_detection.csv'.format(folder, profile), 'w+')
            count = 0
        except:
            print('TRY #{}:'.format(count), profile)
            count += 1
            continue
    for filename in tqdm(os.listdir(folder)):
        try:
            if filename.endswith('.jpg'):
                command = './darknet detect cfg/yolov3.cfg cfg/yolov3.weights {}{} -thresh 0.5'.format(folder, filename)
                subprocess.check_call(command.split())
                objects = open('prediction_details.txt', 'r').read().strip('\n').split('\n')
                file.write(filename)
                for obj in objects:
                    (label, conf, left, right, top, bot) = obj.split(',')
                    file.write(',"{},{},{},{},{},{}"'.format(label, conf, left, right, top, bot))
                file.write('\n')
        except:
            print('ERROR. SKIPPING:', profile, filename)
            continue


# For running in python 2.x
from __future__ import print_function, unicode_literals
from __future__ import absolute_import, division

import sys
import numpy as np
from sklearn.metrics import confusion_matrix
from scipy.spatial.distance import cdist
from keras.utils.np_utils import to_categorical

def relabel(labelTr, labelTs):
    unqLab, pInClass=np.unique(labelTr, return_counts=True)
    sortedUnqLab=np.argsort(pInClass, kind='mergesort')
    c=sortedUnqLab.shape[0]
    labelsNewTr=np.zeros((labelTr.shape[0],))-1
    labelsNewTs=np.zeros((labelTs.shape[0],))-1
    pInClass=np.sort(pInClass)
    classMap=list()
    for i in range(c):
        labelsNewTr[labelTr==unqLab[sortedUnqLab[i]]]=i
        labelsNewTs[labelTs==unqLab[sortedUnqLab[i]]]=i
        classMap.append(np.where(labelsNewTr==i)[0])
    return labelsNewTr, labelsNewTs, c, pInClass, classMap

def irFind(pInClass, c, irIgnore=1):
    ir=pInClass[-1]/pInClass
    imbalancedCls=np.arange(c)[ir>irIgnore]
    toBalance=np.subtract(pInClass[-1], pInClass[imbalancedCls])
    imbClsNum=toBalance.shape[0]
    if imbClsNum==0: sys.exit('No imbalanced classes found, exiting ...')
    return imbalancedCls, toBalance, imbClsNum, ir

def fileRead(fileName):
    dataTotal=np.loadtxt(fileName, delimiter=',')
    data=dataTotal[:, :-1]
    labels=dataTotal[:, -1]
    return data, labels

def indices(pLabel, tLabel):
    confMat=confusion_matrix(tLabel, pLabel)
    nc=np.sum(confMat, axis=1)
    tp=np.diagonal(confMat)
    tpr=tp/nc
    ppv = tp/(np.sum(confMat, axis = 0))

    acsa=np.mean(tpr)
    gm=np.prod(tpr)**(1/confMat.shape[0])
    acc=np.sum(tp)/np.sum(nc)
    return acsa, gm, tpr, confMat, acc, ppv

def smote(trainS, trainL, pInClass, toBalance, classMap, numNeighbor=5):
    n, d=trainS.shape[0], trainS.shape[1]
    trainSCopy, trainLCopy=np.copy(trainS), np.copy(trainL)
    checkImages=np.zeros((3*toBalance.shape[0], trainS.shape[1]))
    for i in range(toBalance.shape[0]):
        newPoints=np.zeros((toBalance[i], d))
        indices=np.random.randint(0, pInClass[i], (toBalance[i],))
        labelsNew=np.ones((toBalance[i],))*i
        for j in range(toBalance[i]):
            point=np.expand_dims(trainS[classMap[i][indices[j]], :], axis=0)
            distance=cdist(point, trainS[classMap[i], :])
            distInd=np.argsort(distance, axis=1)
            distInd=distInd[0, 1:1+numNeighbor]
            neighbor=np.random.randint(0, numNeighbor)
            neighborPoint=trainS[classMap[i][distInd[neighbor]], :]
            alpha=np.random.rand()
            newPoints[j, :]=(alpha*point)+((1-alpha)*neighborPoint)
        trainSCopy=np.vstack((trainSCopy, newPoints))
        trainLCopy=np.hstack((trainLCopy, labelsNew))
        checkImages[(i*3):((i+1)*3), :]=newPoints[0:3, :]
    return trainSCopy, trainLCopy, checkImages

def randomLabelGen(toBalance, batchSize, c):
    cumProb=np.cumsum(toBalance/np.sum(toBalance))
    bins=np.insert(cumProb, 0, 0)
    randomValue=np.random.rand(batchSize,)
    randLabel=np.digitize(randomValue, bins)-1
    randLabel_cat=to_categorical(randLabel)
    labelPadding=np.zeros((batchSize, c-randLabel_cat.shape[1]))
    randLabel_cat=np.hstack((randLabel_cat, labelPadding))
    return randLabel_cat

def batchDivision(n, batchSize):
    numBatches, residual=int(np.ceil(n/batchSize)), int(n%batchSize)
    if residual==0:
        residual=batchSize
    batchDiv=np.zeros((numBatches+1,1), dtype='int64')
    batchSizeStore=np.ones((numBatches, 1), dtype='int64')
    batchSizeStore[0:-1, 0]=batchSize
    batchSizeStore[-1, 0]=residual
    for i in range(numBatches):
        batchDiv[i]=i*batchSize
    batchDiv[numBatches]=batchDiv[numBatches-1]+residual
    return batchDiv, numBatches, batchSizeStore

def rearrange(labelsCat, numImbCls):
    labels=np.argmax(labelsCat, axis=1)
    arrangeMap=list()
    for i in range(numImbCls):
        arrangeMap.append(np.where(labels==i)[0])
    return arrangeMap


"""
===============================================
Cross-validation on diabetes Dataset Exercise
===============================================

A tutorial exercise which uses cross-validation with linear models.

This exercise is used in the :ref:`cv_estimators_tut` part of the
:ref:`model_selection_tut` section of the :ref:`stat_learn_tut_index`.

"""

# %%
# Load dataset and apply GridSearchCV
# -----------------------------------
import matplotlib.pyplot as plt
import numpy as np

from sklearn import datasets
from sklearn.linear_model import Lasso
from sklearn.model_selection import GridSearchCV

X, y = datasets.load_diabetes(return_X_y=True)
X = X[:150]
y = y[:150]

lasso = Lasso(random_state=0, max_iter=10000)
alphas = np.logspace(-4, -0.5, 30)

tuned_parameters = [{"alpha": alphas}]
n_folds = 5

clf = GridSearchCV(lasso, tuned_parameters, cv=n_folds, refit=False)
clf.fit(X, y)
scores = clf.cv_results_["mean_test_score"]
scores_std = clf.cv_results_["std_test_score"]

# %%
# Plot error lines showing +/- std. errors of the scores
# ------------------------------------------------------

plt.figure().set_size_inches(8, 6)
plt.semilogx(alphas, scores)

std_error = scores_std / np.sqrt(n_folds)

plt.semilogx(alphas, scores + std_error, "b--")
plt.semilogx(alphas, scores - std_error, "b--")

# alpha=0.2 controls the translucency of the fill color
plt.fill_between(alphas, scores + std_error, scores - std_error, alpha=0.2)

plt.ylabel("CV score +/- std error")
plt.xlabel("alpha")
plt.axhline(np.max(scores), linestyle="--", color=".5")
plt.xlim([alphas[0], alphas[-1]])

# %%
# Bonus: how much can you trust the selection of alpha?
# -----------------------------------------------------

# To answer this question we use the LassoCV object that sets its alpha
# parameter automatically from the data by internal cross-validation (i.e. it
# performs cross-validation on the training data it receives).
# We use external cross-validation to see how much the automatically obtained
# alphas differ across different cross-validation folds.

from sklearn.linear_model import LassoCV
from sklearn.model_selection import KFold

lasso_cv = LassoCV(alphas=alphas, random_state=0, max_iter=10000)
k_fold = KFold(3)

print("Answer to the bonus question:", "how much can you trust the selection of alpha?")
print()
print("Alpha parameters maximising the generalization score on different")
print("subsets of the data:")
for k, (train, test) in enumerate(k_fold.split(X, y)):
    lasso_cv.fit(X[train], y[train])
    print(
        "[fold {0}] alpha: {1:.5f}, score: {2:.5f}".format(
            k, lasso_cv.alpha_, lasso_cv.score(X[test], y[test])
        )
    )
print()
print("Answer: Not very much since we obtained different alphas for different")
print("subsets of the data and moreover, the scores for these alphas differ")
print("quite substantially.")

plt.show()


r"""Linear Ordinary Differential Equation Recurrent Neural Network."""

__all__ = [
    # Constants
    "__version__",
    "CONFIG",
    # Functions
    "info",
    # Sub-Modules
    "config",
    "initializations",
    "lib",
    "models",
    "projections",
    "regularizations",
    "utils",
    "testing",
]

from importlib import metadata

from linodenet import (
    config,
    initializations,
    lib,
    models,
    projections,
    regularizations,
    testing,
    utils,
)
from linodenet.config import CONFIG

# single-source version
try:
    __version__ = metadata.version(__package__ or __name__)
    r"""The version number of the `tsdm` package."""
except metadata.PackageNotFoundError:
    __version__ = "unknown"
    r"""The version number of the `tsdm` package."""
finally:
    del metadata

def info(obj: object | None = None) -> None:
    """Open the help page for the given object in a browser."""
    import inspect  # pylint: disable=import-outside-toplevel
    import webbrowser  # pylint: disable=import-outside-toplevel

    url = config.PROJECT.DOC_URL
    if obj is None:
        webbrowser.open(url)
        return

    pkg = inspect.getmodule(obj)
    if pkg is None:
        pkg = inspect.getmodule(type(obj))
    if pkg is None:
        raise ValueError("Could not determine package of object!")
    if not pkg.__name__.startswith(__package__):
        raise ValueError(f"Object does not belong to {__package__!r}!")
    webbrowser.open(f"{url}/apidoc/{pkg.__name__}.html")


import importlib.util  # needs backport to Python 2.7
# https://stackoverflow.com/questions/45350363/backport-of-importlib-for-python-2-7-from-3-6
import os
from flask import current_app, render_template
from flask_login import current_user
from flask import abort
ERROR_MSG = """
Wrong mapping format!
The syntax is either (<route>, <function>, <http_method>) or (<prefix>, <module>)
"""

class FlaskUrls(object):

    def __init__(self, app, abort_page=None):
        self._permissions = {}
        if app is not None:
            self.app = app
            self.init_app(app)
            self.abort_page = abort_page

    def init_app(self, app):
        app.url_value_preprocessors.setdefault(None, []).append(self._check_permissions)

    def _check_permissions(self, endpoint, values):
        """
        this function checks if a current user has the permission to 'use' the endpoint
        :param endpoint:
        :param values: not used but part of the signature
        :return:
        """
        with self.app.app_context():
            self.app.logger.info(current_app.name)
            self.app.logger.info("_check_permissions for: " + endpoint)
            if endpoint in self._permissions:
                self.app.logger.info(endpoint + " needs " + self._permissions[endpoint])
                if self._permissions[endpoint] in current_user.get_roles():
                    self.app.logger.info("access granted")
                else:
                    self.app.logger.info("access denied")
                    if self.abort_page:
                        abort(403)  # TODO: render page
                    else:
                        abort(403)

    def register_urls(self, urls, prefix=""):
        self.app.logger.info("Registering Urls")
        self._register_urls(urls, prefix)
        self.app.logger.info(self.app.url_map)

    def _register_urls(self, urls, prefix):
        """Main entry point of the module. Takes an array of url mappings and registers them as url_rule at the Flask app
        :param app: a Flask app object
        :param urls: an array of tuples of (<route>, <function>, <http_method>) or (<prefix>, <module>)
        :param prefix: a prefix to the route
        :return:
        """
        self.app.logger.info("_register_urls")
        for url in urls:
            number_of_args = len(url)

            if number_of_args == 2 and isinstance(url[0], str) and hasattr(url[1], '__call__'):
                self._register_endpoint(prefix, url[0], url[1], ["GET"])
            elif number_of_args == 2 and isinstance(url[0], str) and isinstance(url[1], str):
                self._register_component(url)
            elif number_of_args == 3 and isinstance(url[0], str) and hasattr(url[1], '__call__') and isinstance(url[2], list):
                self._register_endpoint(prefix, url[0], url[1], url[2])
            elif number_of_args == 4 and isinstance(url[0], str) and hasattr(url[1], '__call__') \
                    and isinstance(url[2], list) and isinstance(url[3], str):
                self._register_endpoint(prefix, url[0], url[1], url[2])
                self._permissions[os.path.join(prefix, url[1].__name__)] = url[3]
                self.app.logger.debug(self._permissions)
            else:
                raise Exception(ERROR_MSG)

    def _register_component(self, url):
        self.app.logger.info("Registering component: " + str(url))
        prefix = url[0]
        urls_file = url[1]
        self.app.logger.info("Importing: " + urls_file)
        urls_file = urls_file.split(".")
        package = urls_file[0]
        module_name = urls_file[1]
        module_path = os.path.join(self.app.root_path, package)
        urls_path = os.path.join(module_path, module_name + ".py")
        loaded_module = self._load_module(urls_path)
        self._register_urls(loaded_module.urls, prefix)
        self._append_template_path(module_path)

    def _append_template_path(self, module_path):
        template_path = os.path.join(module_path, "templates")
        if os.path.isdir(template_path):
            self.app.jinja_loader.searchpath.append(template_path)

    def _load_module(self, module_path):
        spec = importlib.util.spec_from_file_location("urls", module_path)
        foo = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(foo)
        return foo

    def _register_endpoint(self, prefix, route, view_func, methods):
        self.app.logger.info("Registering endpoint: " + str(route))
        if prefix:
            route = prefix + route
        self.app.add_url_rule(route, methods=methods, view_func=view_func)


import requests
from bs4 import BeautifulSoup

class JobScraper:

    def __init__(self, url):
        self.url = url

    # Scrapes this JobScraper's url and returns a list of all words
    # that are not part of html tags
    def scrape(self):
        r = requests.get(self.url)
        if(r.status_code != 200):
            print('Unable to access site')

        #soup = BeautifulSoup(r.text, 'html.parser')
        soup = BeautifulSoup(r.text, 'html.parser')
        plaintext = soup.encode("utf-8");

        allWords = []
        curWord = ""
        curTag = ""
        inWord = False
        inTag = False
        inScript = False
        for s in plaintext:
            char = chr(s)
            if(char == '<'):
                inTag = True
                #print('Tag'+ str(curTag))
                curTag = ""
            elif(char == '>'):
                if(inTag):
                    #if(curTag == "script"):
                    #    inScript = True
                    #elif(curTag == "/script"):
                    #    inScript = False
                    #print("Tag"+ str(curTag))
                    curTag = ""
                inTag = False
            elif(char == ' ' or char == '\n' or char == '\r' or char == '\t'):
                if(inTag):
                    #if(curTag == "script"):
                    #    inScript = True
                    #elif(curTag == "/script"):
                    #    inScript = False
                    #print("Tag"+ str(curTag))
                    curTag = ""

                if(inWord and (not inScript)):
                    exclude = set('!,.()?~|')
                    curWord = ''.join(char for char in curWord if char not in exclude)
                    #print("Word"+ str(curWord))
                    allWords.append(curWord)
                    curWord = ""
                inWord = False
            else:
                if(not inTag and not inScript):
                    inWord = True
                    curWord = curWord + str(char)
                else:
                    curTag = curTag + str(char)

        #should remove duplicates
        return allWords

#Arbitrary job posting for testing
#url = 'https://www.alarm.com/about/open-positions.aspx?gh_jid=4267530002'
#scraper = JobScraper(url)
#allWords = scraper.scrape()

#strip extras
#for word in allWords:
#    if word == '/r' or word == '/n' or word == '':
#        allWords.remove(word)
#
#for word in allWords:
#   print("RAYTHEON POST:" + str(word))
  #  pass


import logging
import os
import sys
from collections import Counter
from typing import Iterable, Optional, Union

import numpy as np
import tensorflow as tf
from tensorflow.keras import backend as K

class LRScheduler(tf.keras.optimizers.schedules.LearningRateSchedule):
    """Schedule learning rate linearly from max_learning_rate to min_learning_rate."""

    def __init__(
        self,
        total_steps: int,
        max_learning_rate: float,
        min_learning_rate: float,
        warmup_rate: Optional[float] = None,
        warmup_steps: Optional[int] = None,
    ):
        self.warmup_steps = max(int(total_steps * warmup_rate) if warmup_steps is None else warmup_steps, 1)
        self.increasing_delta = max_learning_rate / self.warmup_steps
        self.decreasing_delta = (max_learning_rate - min_learning_rate) / (total_steps - self.warmup_steps)
        self.max_learning_rate = tf.cast(max_learning_rate, tf.float32)
        self.min_learning_rate = tf.cast(min_learning_rate, tf.float32)

    def __call__(self, step: Union[int, tf.Tensor]) -> tf.Tensor:
        step = tf.cast(step, tf.float32)
        lr = tf.reduce_min([step * self.increasing_delta, self.max_learning_rate - step * self.decreasing_delta])
        return tf.reduce_max([lr, self.min_learning_rate])

def get_logger(name: str) -> logging.Logger:
    logger = logging.getLogger(name)
    logger.setLevel(logging.DEBUG)
    handler = logging.StreamHandler(sys.stdout)
    handler.setFormatter(logging.Formatter("[%(asctime)s] %(message)s"))
    logger.addHandler(handler)
    return logger

def path_join(*paths: Iterable[str]) -> str:
    """Join paths to string local paths and google storage paths also"""
    if paths[0].startswith("gs://"):
        return "/".join([path.strip("/") for path in paths])
    return os.path.join(*paths)

def set_mixed_precision(device: str):
    """Set mixed precision on"""
    mixed_type = "mixed_bfloat16" if device == "TPU" else "mixed_float16"
    policy = tf.keras.mixed_precision.experimental.Policy(mixed_type)
    tf.keras.mixed_precision.experimental.set_policy(policy)

def get_device_strategy(device) -> tf.distribute.Strategy:
    """Return tensorflow device strategy"""
    # Use TPU
    if device == "TPU":
        resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu=os.environ["TPU_NAME"])
        tf.config.experimental_connect_to_cluster(resolver)
        tf.tpu.experimental.initialize_tpu_system(resolver)
        strategy = tf.distribute.TPUStrategy(resolver)

        return strategy

    # Use GPU
    if device == "GPU":
        devices = tf.config.list_physical_devices("GPU")
        if len(devices) == 0:
            raise RuntimeError("Cannot find GPU!")
        for device in devices:
            tf.config.experimental.set_memory_growth(device, True)
        if len(devices) > 1:
            strategy = tf.distribute.MirroredStrategy()
        else:
            strategy = tf.distribute.OneDeviceStrategy("/gpu:0")

        return strategy

    # Use CPU
    return tf.distribute.OneDeviceStrategy("/cpu:0")

def sparse_categorical_crossentropy_nan_debug(y_true, y_pred):
    pred_nan = tf.math.is_nan(y_pred)
    if tf.math.reduce_any(pred_nan):
        tf.print(
            "\nWarning:",
            "The",
            tf.size(tf.where(pred_nan)),
            "number of output values are Nan!\n",
            output_stream=sys.stderr,
        )

    loss = tf.losses.sparse_categorical_crossentropy(y_true, y_pred, True)
    is_nan = tf.math.is_nan(loss)
    if tf.math.reduce_any(is_nan):
        tf.print(
            "\nWarning:", "The", tf.size(tf.where(is_nan)), "number of losses are Nan!\n", output_stream=sys.stderr
        )
        loss = tf.boolean_mask(loss, tf.logical_not(is_nan))
    return loss


from collections import defaultdict
from contextlib import contextmanager
from typing import IO, ContextManager, Generator, TextIO

@contextmanager
def file_read(filename: str) -> Generator[TextIO, None, None]:
    f: TextIO = open(filename)
    yield f
    f.close()

def line_processor(filename: str) -> Generator[str, None, None]:
    """
    This generator takes a file line by line
    and processes one by one.
    """
    cm: ContextManager[IO] = file_read(filename)
    with cm as input_file:
        for line in input_file:
            yield line.strip()

def create_graph_from_cave_connections(
    cave_connections: list[str],
) -> defaultdict[str, set[str]]:
    # create a graph using a defaultdict where the key is the starting vertex
    # and the value is a set of possible nodes to travel to
    graph: defaultdict[str, set[str]] = defaultdict(set)
    for connection in cave_connections:
        starting_vertex: str
        ending_vertex: str
        starting_vertex, ending_vertex = connection.split("-")
        graph[starting_vertex].add(ending_vertex)
        graph[ending_vertex].add(starting_vertex)
    return graph

def depth_first_search(
    starting_vertex: str,
    vertices_seen: set[str],
    graph: defaultdict[str, set[str]],
):
    if starting_vertex == "end":
        # if this new path starts at end, don't look for further ones
        return 1  # why?

    n_paths: int = 0
    # iterate through the possible vertices we can travel to
    for next_vertex in graph[starting_vertex]:  # type: str
        if next_vertex not in vertices_seen:
            # is it a big cave (UPPER) or small cave (lower)?
            # if big cave, don't add it to the vertices_seen since we
            # can visit it multiple times. only add the small caves
            # to the seen vertices
            # then repeat the process starting at the next_vertex
            if not next_vertex.isupper():
                n_paths += depth_first_search(
                    starting_vertex=next_vertex,
                    vertices_seen=vertices_seen.union({next_vertex}),
                    graph=graph,
                )
            else:
                n_paths += depth_first_search(
                    starting_vertex=next_vertex,
                    vertices_seen=vertices_seen,
                    graph=graph,
                )
    return n_paths

def depth_first_search_with_single_small_cave_once(
    starting_vertex: str,
    vertices_seen: set[str],
    graph: defaultdict[str, set[str]],
    visited_small_cave: bool,
):
    if starting_vertex == "end":
        # if this new path starts at end, don't look for further ones
        # as we care about the paths that end here
        # and it cannot be visited ore than once
        return 1  # why?

    n_paths: int = 0
    # iterate through the possible vertices we can travel to
    for next_vertex in graph[starting_vertex]:  # type: str
        if next_vertex not in vertices_seen:
            # is it a big cave (UPPER) or small cave (lower)?
            # if big cave, don't add it to the vertices_seen since we
            # can visit it multiple times. only add the small caves
            # to the seen vertices
            # then repeat the process starting at the next_vertex
            if not next_vertex.isupper():
                n_paths += depth_first_search_with_single_small_cave_once(
                    starting_vertex=next_vertex,
                    vertices_seen=vertices_seen.union({next_vertex}),
                    graph=graph,
                    visited_small_cave=visited_small_cave,
                )
            else:
                n_paths += depth_first_search_with_single_small_cave_once(
                    starting_vertex=next_vertex,
                    vertices_seen=vertices_seen,
                    graph=graph,
                    visited_small_cave=visited_small_cave,
                )
        # if we can visit caves more than once, then we will start at this cave
        # and call depth_first_search again but this time with visited_small_cave = False
        # since this small cave can only be visited once
        elif visited_small_cave and next_vertex != "start":
            n_paths += depth_first_search_with_single_small_cave_once(
                starting_vertex=next_vertex,
                vertices_seen=vertices_seen,
                graph=graph,
                visited_small_cave=False,
            )

    return n_paths

if __name__ == "__main__":
    cave_connections: list[str] = [line for line in line_processor("inputs/day12")]
    graph: defaultdict[str, set[str]] = create_graph_from_cave_connections(
        cave_connections
    )
    part_1: int = depth_first_search(
        starting_vertex="start", vertices_seen={"start"}, graph=graph
    )
    part_2: int = depth_first_search_with_single_small_cave_once(
        starting_vertex="start",
        vertices_seen={"start"},
        graph=graph,
        visited_small_cave=True,
    )
    print(part_1)
    print(part_2)


import time
import logging
from . import ExplorationTechnique

l = logging.getLogger(name=__name__)

class Timeout(ExplorationTechnique):
    """
    Timeout exploration technique that stops an active exploration if the run time exceeds
    a predefined timeout
    """

    def __init__(self, timeout=None):
        super().__init__()
        self.start_time = None
        self.timeout = timeout

    def setup(self, simgr):
        simgr.stashes["timeout"] = []

    def step(self, simgr, stash="active", **kwargs):
        if self.start_time is None:
            self.start_time = time.time()
        if self.timeout is not None and time.time() - self.start_time > self.timeout:
            self.start_time = None
            simgr.move(stash, "timeout")
            l.warning("exploration timeout in %s seconds!", self.timeout)
        else:
            simgr = simgr.step(stash=stash, **kwargs)
        return simgr


#!/usr/bin/env python3

import os
import sqlite3
from datetime import date

main_db = sqlite3.connect("dictionaries/wdweb/wikdict.sqlite3")
main_db.row_factory = sqlite3.Row
sitemap_tmpl = """
<?xml version="1.0" encoding="UTF-8"?>
<urlset xmlns="http://www.sitemaps.org/schemas/sitemap/0.9">
{}
</urlset>
""".strip()
sitemap_index_tmpl = """
<?xml version="1.0" encoding="UTF-8"?>
<sitemapindex xmlns="http://www.sitemaps.org/schemas/sitemap/0.9">
{}
</sitemapindex>
""".strip()

def make_sitemap(pair, lang):
    cur = sqlite3.connect("dictionaries/processed/%s.sqlite3" % lang).cursor()
    cur.execute("SELECT vocable FROM importance ORDER BY score DESC LIMIT 100")
    sorted_pair = "-".join(sorted(pair.split("-")))
    urls = "".join(
        """
        <url>
            <loc>https://www.wikdict.com/{}/{}</loc>
            <changefreq>monthly</changefreq>
        </url>
        """.format(
            sorted_pair, vocable.split("/")[1]
        )
        for (vocable,) in cur
    )
    filename = "sitemap/{}.xml".format(pair)
    with open(filename, "w") as f:
        f.write(sitemap_tmpl.format(urls))
    return filename

def make_sitemap_index(sitemaps):
    urls = "".join(
        """
        <sitemap>
            <loc>https://www.wikdict.com/static/{}</loc>
            <lastmod>{}</lastmod>
        </sitemap>
        """.format(
            url, date.today()
        )
        for url in sitemaps
    )
    with open("sitemap/index.xml", "w") as f:
        f.write(sitemap_index_tmpl.format(urls))

def main():
    os.makedirs("sitemap", exist_ok=True)
    cur = main_db.cursor()
    cur.execute("SELECT * FROM lang_pair")
    sitemaps = []
    for row in cur:
        pair = row["from_lang"] + "-" + row["to_lang"]
        print(pair)
        sitemaps.extend(
            [make_sitemap(pair, row["from_lang"]), make_sitemap(pair, row["to_lang"])]
        )
    make_sitemap_index(sitemaps)

main()


import logging

from homeassistant.components.media_player import (
    MediaPlayerEntity,
    MediaPlayerDeviceClass,
)
from homeassistant.components.media_player.const import (
    SUPPORT_PLAY,
    SUPPORT_SELECT_SOUND_MODE,
    SUPPORT_STOP,
    SUPPORT_VOLUME_SET,
    SUPPORT_VOLUME_STEP,
    MEDIA_TYPE_MUSIC,
)
from homeassistant.const import (
    STATE_IDLE,
    STATE_PLAYING,
)
from hatch_rest_api import RestIot
from .rest_entity import RestEntity

_LOGGER = logging.getLogger(__name__)

class RiotMediaEntity(RestEntity, MediaPlayerEntity):
    _attr_should_poll = False
    _attr_media_content_type = MEDIA_TYPE_MUSIC
    _attr_device_class = MediaPlayerDeviceClass.SPEAKER

    def __init__(self, rest_device: RestIot):
        super().__init__(rest_device, "Media Player")
        self._attr_sound_mode_list = self.rest_device.favorite_names()
        self._attr_supported_features = (
            SUPPORT_PLAY
            | SUPPORT_STOP
            | SUPPORT_SELECT_SOUND_MODE
            | SUPPORT_VOLUME_SET
            | SUPPORT_VOLUME_STEP
        )

    def _update_local_state(self):
        if self.platform is None:
            return
        _LOGGER.debug(f"updating state:{self.rest_device}")
        if self.rest_device.is_playing:
            self._attr_state = STATE_PLAYING
        else:
            self._attr_state = STATE_IDLE
        self._attr_sound_mode = self.rest_device.audio_track
        self._attr_volume_level = self.rest_device.volume / 100
        self._attr_device_info.update(sw_version=self.rest_device.firmware_version)
        self.async_write_ha_state()

    def set_volume_level(self, volume):
        self.rest_device.set_volume(volume * 100)

    def media_play(self):
        self.rest_device.set_favorite(self._attr_sound_mode_list[0])

    def select_sound_mode(self, sound_mode: str):
        self.rest_device.set_favorite(sound_mode)

    def media_stop(self):
        self.rest_device.turn_off()


import matplotlib.pyplot as plt # pip3 install matplotlib
import csv

from statistics import mean

def average_period(time_list):
    time_difference = []
    for i in range(1, len(time_list) - 1):
        time_difference.append(time_list[i] - time_list[i - 1])

    return mean(time_difference)

def average_reduction(data, window_size=20):
    averaged_data = []
    start = 0
    iteration = 1
    max_iterations = int(len(data) / window_size)
    max_iterations += 1 if len(data) % window_size > 0 else max_iterations
    while iteration <= max_iterations:
        data_sum = 0
        end = window_size * iteration if iteration < max_iterations else len(data)
        for i in range(start, end):
            data_sum += data[i]
        average = data_sum / abs(end - start)
        averaged_data.append(average)
        start += window_size
        iteration += 1
    return averaged_data

x = []
y = []

with open('sensors_log_15.1_2019-05-15T16-12-09.440914Z.csv','r') as csvfile:
    plots = csv.reader(csvfile, delimiter=',')
    for row in plots:
        x.append(int(row[1]))
        y.append(int(row[0]))
time_stamps = average_reduction(x, window_size=20)
data = average_reduction(y, window_size=20)

plt.plot(x, y, label='raw_data')
plt.plot(time_stamps, data, label='averaged_data')
plt.xlabel('milliseconds')
plt.ylabel('Galvanic response [0-1023]')
plt.title('Interesting Graph\nCheck it out')
plt.legend()
plt.show()


# -*- coding: utf-8 -*-

import sys

PY3 = False
if sys.version_info[0] >= 3: PY3 = True

import os
import codecs

from lib import serverhls

from core import httptools, scrapertools
from platformcode import logger, config

def get_video_url(page_url, url_referer=''):
    logger.info("(page_url='%s')" % page_url)

    ref = page_url.split('//', 1)

    resp = httptools.downloadpage(page_url, headers={"referer": ref[0] + "//" + ref[1]})

    if resp.code == 404:
        return "El archivo no existe o ha sido borrado"

    data = resp.data

    if PY3 and isinstance(data, bytes):
        data = "".join(chr(x) for x in bytes(data))

    v_type = "hls"

    url = page_url

    video_urls = list()

    if "const source = '" in data:
        url = scrapertools.find_single_match(data, "const source = '([^']+)")
        v_type = "hls"

    if '"file": ' in data:
        url, v_type = scrapertools.find_single_match(data, '"file": "([^"]+)",\s+"type": "([^"]+)"')

    headers = {"referer": page_url}

    if v_type == "mp4":
        url = httptools.downloadpage(url, headers=headers, follow_redirects=False).headers["location"]

        page_url = "%s|Referer=%s&User-Agent=%s" % (url, page_url, httptools.get_user_agent())

    elif v_type == "hls":
        if '.hls' in data: hls_data = data
        else:
            hls_data = httptools.downloadpage(url, headers=headers).data

            if PY3 and isinstance(hls_data, bytes):
                hls_data = "".join(chr(x) for x in bytes(hls_data))

        base_url = scrapertools.find_single_match(hls_data, "((?:https?:)?\/\/[^\/]+)")
        if base_url: hls_data = hls_data.replace(base_url, 'http://localhost:8781')

        if not base_url.startswith('http'): base_url = 'https:%s' % base_url

        m3u8 = os.path.join(config.get_data_path(), "m3u8hls.m3u8")

        outfile = open(m3u8, 'wb')
        outfile.write(codecs.encode(hls_data, "utf-8"))
        outfile.close()

        page_url = m3u8
        v_type = "m3u8"

        serverhls.start(base_url)
    else:
        return video_urls

    video_urls = [["%s" % v_type, page_url]]

    return video_urls


#Testcv4.py Looda for circles finds center

import sys
import cv2 as cv
import numpy as np
import time

def getTimeInMillis():
    return (int)(time.time()*1000.0)

def main():

    filename = 'Redlight13on.jpg'

    # Loads an image
    #src = cv.imread(filename, cv.IMREAD_COLOR)
    imageOrig = cv.imread("Redlight13on.jpg")
    scale = 0.25
    h = int(scale * imageOrig.shape[0])
    w = int(scale * imageOrig.shape[1])
    src = cv.resize(imageOrig, (w, h))
    print ("resized to: w = ", w, ", h = ",h)
    # Check if image is loaded fine
    if src is None:
        print ('Error opening image!')
        #print ('Usage: hough_circle.py [image_name -- default ' + default_file + '] \n')
        return -1
    cv.imshow("Scaled image",src)
    cv.waitKey(0)
    cv.destroyAllWindows()

    gray = cv.cvtColor(src, cv.COLOR_BGR2GRAY)
    cv.imshow("Grayed image",gray)
    cv.waitKey(0)
    cv.destroyAllWindows()

    gray = cv.medianBlur(gray, 5)
    cv.imshow("Blured image",gray)
    cv.waitKey(0)
    cv.destroyAllWindows()

    ret, binary = cv.threshold(gray, 100, 255,
      cv.THRESH_OTSU)
    binary = gray
    cv.imshow("Binary image",binary)
    cv.waitKey(0)
    cv.destroyAllWindows()

    rows = binary.shape[0]
    start = getTimeInMillis()
    circles = cv.HoughCircles(binary, cv.HOUGH_GRADIENT, 1, rows / 8,
                               param1=205, param2=8,
                               minRadius=20, maxRadius=324)
    computeTime = getTimeInMillis()- start

    print("Compute time = ",computeTime)
    if circles is not None:
        circles = np.uint16(np.around(circles))
        for i in circles[0, :]:
            center = (i[0], i[1])
            # circle center
            cv.circle(src, center, 1, (0, 100, 100), 3)
            # circle outline
            radius = i[2]
            cv.circle(src, center, radius, (255, 0, 255), 3)
    else:
        print ("No circles found")

    cv.imshow("detected circles", src)
    cv.waitKey(0)
    cv.destroyAllWindows()

    return 0
if __name__ == "__main__":
    main()


"""Model Wrapper for backward with activations checkpointing (for ZeRO)"""
#
#
#
# TODO:
#   1. ZeRO use activations checkpointing with pt-function style
#      Training Deep Nets with Sublinear Memory Cost
#      <https://github.com/Lyken17/pytorch-memonger>
#
# ----------------------------------------------------------------------------

import torch
import torch.nn as nn
import torch.nn.functional as F

from necklace.frmwrk.pytorch.mpnn import checkpointing as pt_ckpt
from necklace.frmwrk.pytorch.pttsutils import pt_is_tensor

# ------------------------------------------------------------------
# show more info when backward error
# ------------------------------------------------------------------
#torch.autograd.set_detect_anomaly(True)

# ================================================================
# TODO: use checkpointing with pt-function style
# ================================================================
class MdWrpWithCkpt(nn.Module):
    def __init__(self, submdls=[], md_device='cuda:0', mp_order=0):
        super(MdWrpWithCkpt, self).__init__()

        self.submdls = submdls
        self.md_device = md_device
        self.mp_order = mp_order

        #submdls = [m.to(md_device) for m in submdls]
        net = nn.Sequential(*submdls)  # TODO: nn.ModuleList
        self.net = net#.to(md_device)

    def forward(self, *args, **kwargs):
        out = pt_ckpt.checkpoint(self.net.forward, *args)
        return out

# NOTE: now we can NOT put functions like torch.xx and F.xx to nn.Sequential,
#       and the forward of the submodel may be very special or complex,
#       so in the most cases we must implement submodel mannually with
#       a whole forward function, and then wrap them with this class
class MdWrpWithSelfBkwd(nn.Module):
    def __init__(self, submdls=[], md_device='cuda:0', mp_order=0,
                 is_last_part=False, do_acts_ckpt=True):
        super(MdWrpWithSelfBkwd, self).__init__()

        self.submdls = submdls
        self.md_device = md_device
        self.mp_order = mp_order
        self.is_last_part = is_last_part
        self.do_acts_ckpt = do_acts_ckpt  # TODO:

        #submdls = [m.to(md_device) for m in submdls]
        net = nn.Sequential(*submdls)  # TODO: nn.ModuleList
        self.net = net#.to(md_device)

        self.inp_cache = None
        self.out_cache = None

    def cache_graded_input(self, *args, **kwargs):
        if self.mp_order > 0:  # not the first submodule which is not need to get output gradients
            inputs = []
            for a in args:
                if pt_is_tensor(a):
                    # NOTE: here the `a` may be a Tensor with requires_grad=True and so it is not a
                    # leaf of graph and we can NOT save its grads, so here we must detch it from the graph
                    a = a.detach()
                    a.requires_grad = True
                inputs.append(a)
        else:
            inputs = args

        self.inp_cache = inputs

        return inputs

    def forward(self, *args, **kwargs):
        inputs = self.cache_graded_input(*args)

        # ================================================================================
        # TODO: (do_acts_ckpt)
        #       if True (for now), we will do the forward without grad, and do the forward
        #       again with grad in the my_backward()
        # ================================================================================
        if self.is_last_part or (not self.do_acts_ckpt):
            with torch.enable_grad():
                out = self.net(*inputs)

        else:
            with torch.no_grad():
                out = self.net(*inputs)

        self.out_cache = out  # NOTE: cache the out for mp_backward

        return out

    def get_input_grads(self):
        if self.mp_order > 0:
            grads = [inp.grad for inp in self.inp_cache]
        else:
            grads = None

        return grads

    def mp_backward(self, grads=None, *args, **kwargs):

        # ================================================================================
        # TODO: (do_acts_ckpt)
        #       if True (for now), we will do the forward without grad, and do the forward
        #       again with grad in the my_backward()
        # SO, it is very slow !!!
        # and, this is NOT the real activations_checkpoint, see `pt_ckpt` for details
        # ================================================================================

        if grads is not None:  # not the last submodule which has already done the backward with loss

            #if not self.is_last_part:
            if self.do_acts_ckpt:
                with torch.enable_grad():
                    out = self.net(*self.inp_cache)
                self.out_cache = out
            else:
                pass

            if not isinstance(self.out_cache, (list, tuple, set)):
                outs = [self.out_cache]
            else:
                outs = self.out_cache

            for out, grad in zip(outs, grads):
                torch.autograd.backward(out, grad)
                # or,
                #out.backward(grad)

        inp_grads = self.get_input_grads()

        return inp_grads

class MdWrpWholeNet(nn.Module):
    def __init__(self, mdls=[], md_device='cpu'):
        super(MdWrpWholeNet, self).__init__()
        self.whole_md = nn.Sequential(*mdls)

    def forward(self, *args, **kwargs):
        out = self.whole_md(*args)
        return out


import sys
import contextlib
import time
import threading

from PySide6.QtWidgets import QApplication, QWidget, QToolButton, QSizePolicy
from PySide6.QtWidgets import QUndoView, QHBoxLayout, QVBoxLayout, QListView
from PySide6.QtWidgets import QMainWindow
from PySide6.QtGui import QPainter, QColor, QPixmap, QPen, QTabletEvent
from PySide6.QtGui import QPainterPath, QCursor, QBitmap, QIcon, QAction
from PySide6.QtGui import QUndoStack, QUndoCommand, QStandardItemModel
from PySide6.QtGui import QStandardItem, QUndoGroup, QPointingDevice
from PySide6.QtGui import QKeyEvent
from PySide6.QtCore import Qt, QEvent, QRect, QTimer, QFile, QObject, QSize
from PySide6.QtCore import Signal, QPointF, QRectF, QSizeF, QItemSelectionModel
from PySide6.QtUiTools import QUiLoader

import global_shortcuts

MAX_RADIUS = 100

COLORS = {
    'Red': (1, 0, 0),
    'Green': (0, 1, 0),
    'Blue': (0, 0, 1),
    'Yellow': (1, 1, 0),
    'Purple': (1, 0, 1),
    'Cyan': (0, 1, 1),
}

class Overlay():
    composition_mode = QPainter.CompositionMode_SourceOver
    opacity = 0.5
    _final = None
    def __init__(self, topleft=None, pixmap=None, prev=None):
        self.pixmap = pixmap
        self.prev = prev
        if topleft and pixmap:
            self.rect = QRect(
                topleft.x(), topleft.y(),
                pixmap.width(), pixmap.height()
            )
        else:
            self.rect = None

    def _opaque_copy(self):
        result = Overlay(self.rect.topLeft(), self.pixmap.copy())
        result.opacity = 1
        return result

    def __repr__(self):
        return f'<Overlay {self.rect}>'

    def paint(self, painter):
        if self.rect:
            painter.setOpacity(self.opacity)
            painter.setCompositionMode(self.composition_mode)
            painter.drawPixmap(self.rect.topLeft(), self.pixmap)

    def reserve(self, rect):
        if self.rect:
            new_rect = self.rect.united(rect)
        else:
            new_rect = rect
        new_pixmap = QPixmap(new_rect.width(), new_rect.height())
        new_pixmap.fill(QColor(0, 0, 0, 0))
        if self.pixmap:
            painter = QPainter(new_pixmap)
            painter.drawPixmap(
                self.rect.left() - new_rect.left(),
                self.rect.top() - new_rect.top(),
                self.pixmap
            )
            painter.end()
        self.pixmap = new_pixmap
        self.rect = new_rect

    def add(self, other_overlay):
        if self.rect:
            self.reserve(other_overlay.rect)
            painter = QPainter(self.pixmap)
            painter.drawPixmap(
                other_overlay.rect.left() - self.rect.left(),
                other_overlay.rect.top() - self.rect.top(),
                other_overlay.pixmap
            )
            painter.end()
        else:
            self.pixmap = other_overlay.pixmap.copy()
            self.rect = other_overlay.rect

    @contextlib.contextmanager
    def painter_context(self):
        self._final = None
        if self.pixmap:
            painter = QPainter(self.pixmap)
            painter.translate(-self.rect.topLeft())
            try:
                yield painter
            finally:
                painter.end()
        else:
            yield None

    @property
    def final(self):
        if self._final:
            return self._final
        if self.prev and self.prev.final.pixmap:
            self._final = self.prev.final._opaque_copy()
        else:
            self._final = Overlay()
            self._final.opacity = 1
        if self.pixmap:
            self._final.reserve(self.rect)
            with self._final.painter_context() as painter:
                self.paint(painter)
        return self._final

class DrawCommand(QUndoCommand):
    def __init__(self, widget, tool):
        super().__init__(f"Draw with {tool.name}")
        self.widget = widget
        self.scribbles = widget.scribbles
        if self.scribbles:
            self.scribble = Overlay(prev=self.scribbles[-1])
        else:
            self.scribble = Overlay()
        self.tool = tool

    def undo(self):
        popped = self.scribbles.pop()
        self.widget.current_wet = Overlay()
        if popped.rect:
            self.widget.update(popped.rect)
            assert popped == self.scribble

    def redo(self):
        self.scribbles.append(self.scribble)
        if self.scribble.rect:
            self.widget.update(self.scribble.rect)

class PictureItem(QStandardItem):
    def __init__(self, widget):
        super().__init__("Drawing")
        self.scribbles = []
        self.undo_stack = QUndoStack()
        self.undo_stack.indexChanged.connect(self.reset_props)
        self.widget = widget
        widget.undo_group.addStack(self.undo_stack)

    def start_scribble(self):
        cmd = DrawCommand(self.widget, self.widget.tool)
        self.undo_stack.push(cmd)

    def reset_props(self):
        if self.undo_stack.index() == self.undo_stack.count():
            self.setText(f"Drawing ({self.undo_stack.index()})")
        else:
            self.setText(f"Drawing ({self.undo_stack.index()}/{self.undo_stack.count()})")
        if self.scribbles:
            s = self.scribbles[-1]
            if s.final.pixmap:
                self.setIcon(s.final.pixmap)
            else:
                self.setIcon(QIcon())
        else:
            self.setIcon(QIcon())

class SignalingProperty:
    def __init__(self, signal, default=None, *, name=None):
        self.signal = signal
        self.default = default
        self.name = name

    def __set_name__(self, owner, name):
        if self.name is None:
            self.name = name

    def __get__(self, instance, owner):
        assert self.name
        if instance:
            return instance.__dict__.get(self.name, self.default)
        else:
            return super().__get__(instance, owner)

    def __set__(self, instance, value):
        assert self.name
        prev = instance.__dict__.get(self.name, self.default)
        if value != prev:
            instance.__dict__[self.name] = value
            self.signal.__get__(instance, type(instance)).emit(value)

class OverlayWidget(QWidget):
    grab_updated = Signal(bool)
    can_clear_changed = Signal(bool)
    can_undo_changed = Signal(bool)
    can_redo_changed = Signal(bool)
    can_clear = SignalingProperty(can_clear_changed)
    can_undo = SignalingProperty(can_undo_changed)
    can_redo = SignalingProperty(can_redo_changed)
    _last_cursor_pos = None
    _grabbing_mouse = False

    def __init__(self):
        super().__init__()
        self.setWindowTitle('pointout canvas')
        self.setWindowFlags(
            self.windowFlags()
            | Qt.Window
            | Qt.WindowTransparentForInput
            | Qt.WindowDoesNotAcceptFocus
            | Qt.FramelessWindowHint
            | Qt.WindowStaysOnTopHint
        )
        self.setAttribute(Qt.WA_TransparentForMouseEvents)
        self.setAttribute(Qt.WA_NoSystemBackground)
        self.setAttribute(Qt.WA_TranslucentBackground)
        self.setAttribute(Qt.WA_TabletTracking)

        self.current_wet = Overlay()

        self.undo_group = QUndoGroup()
        self.picture_model = QStandardItemModel()
        self.selection_model = QItemSelectionModel(self.picture_model)
        self.clear(force=True)
        self.selection_model.currentChanged.connect(self.picture_switched)
        self.picture_switched()
        self.undo_group.canUndoChanged.connect(self.update_action_availability)
        self.undo_group.canRedoChanged.connect(self.update_action_availability)

        cursor_bitmap = QBitmap.fromData(QSize(5, 5), bytes((
            0b00000, 0b00000, 0b00100, 0b00000, 0b00000,
        )))
        mask_bitmap = QBitmap.fromData(QSize(5, 5), bytes((
            0b00100, 0b00000, 0b10101, 0b00000, 0b00100,
        )))
        self.setCursor(QCursor(cursor_bitmap, mask_bitmap))

        self.anim_timer = QTimer()
        self.anim_timer.timeout.connect(self.anim_update)
        self.anim_timer.start(1000//30)
        self.anim_timer.setTimerType(Qt.CoarseTimer)
        self.wet_end = time.monotonic()

        self.eraser = Eraser()
        self.tool = Marker()
        self.last_point = 0

    def picture_switched(self):
        self.undo_group.setActiveStack(self.picture.undo_stack)
        self.update_action_availability()
        self.update()

    def update_action_availability(self):
        self.can_clear = bool(self.undo_stack.count())
        self.can_undo = (
            self.undo_stack.canUndo()
            or self.selection_model.currentIndex().row() > 0
        )
        rc = self.picture_model.rowCount()
        self.can_redo = (
            self.undo_stack.canRedo()
            or self.selection_model.currentIndex().row() < rc - 1
        )

    @property
    def picture(self):
        idx = self.selection_model.currentIndex()
        return self.picture_model.itemFromIndex(idx)

    @property
    def scribbles(self):
        return self.picture.scribbles

    @property
    def undo_stack(self):
        return self.picture.undo_stack

    @property
    def tool(self):
        return self._tool
    @tool.setter
    def tool(self, new_tool):
        self._tool = new_tool
        if new_tool is None:
            self.update_grab(False)

    def anim_update(self):
        if self.current_wet.rect is not None:
            with self.current_wet.painter_context() as painter:
                painter.setBrush(QColor(0, 0, 0, 255))
                painter.setPen(QPen(0))
                painter.setOpacity(0.1)
                painter.setCompositionMode(QPainter.CompositionMode_DestinationOut)
                painter.drawRect(self.current_wet.rect)
            self.update(self.current_wet.rect)
            if self.wet_end < time.monotonic():
                self.current_wet = Overlay()

    def paintEvent(self, e):
        painter = QPainter(self)
        if self.scribbles:
            self.scribbles[-1].final.paint(painter)
        painter.setOpacity(1)
        if self.current_wet:
            self.current_wet.paint(painter)
        painter.end()

    def tabletEvent(self, e):
        if e.type() == QEvent.TabletPress:
            self.start_line(e.posF())
        if e.type() in (QEvent.TabletMove, QEvent.TabletRelease):
            self.add_point(
                pos=e.posF(),
                pressure=e.pressure(),
                erase=e.pointerType() == QPointingDevice.PointerType.Eraser,
            )
        e.accept()

    def mousePressEvent(self, e):
        self.start_line(e.localPos())

    def mouseMoveEvent(self, e):
        self.add_point(e.localPos())

    def start_line(self, pos):
        self.last_point = pos
        self.picture.start_scribble()
        self.update_action_availability()

    def add_point(self, pos, *, pressure=0.5, erase=False):
        tool = self.tool
        if self.tool is None:
            return
        if erase:
            tool = self.eraser
        if not self.scribbles:
            self.start_line(pos)
        if self.last_point:
            tool.set_size(pressure)
            update_rect = QRect(self.last_point.toPoint(), pos.toPoint())
            update_rect = update_rect.normalized().adjusted(
                -tool.size-1, -tool.size-1, tool.size+1, tool.size+1,
            )
            tool.draw(
                self.last_point, pos,
                update_rect,
                self.scribbles, self.current_wet
            )
            self.update(update_rect)
        self.last_point = pos
        self.picture.reset_props()
        self.update_wet()

    def clear(self, *, force=False):
        if force or self.can_clear:
            pi = PictureItem(self)
            idx = self.selection_model.currentIndex()
            if idx:
                self.picture_model.insertRow(idx.row() + 1, pi)
            else:
                self.picture_model.appendRow(pi)
            self.selection_model.setCurrentIndex(
                self.picture_model.indexFromItem(pi),
                QItemSelectionModel.ClearAndSelect,
            )

    def undo(self):
        if self.undo_stack.canUndo():
            self.undo_stack.undo()
        elif self.can_undo:
            idx = self.selection_model.currentIndex()
            if idx.row() > 0:
                idx = self.picture_model.index(idx.row() - 1, 0)
                self.selection_model.setCurrentIndex(
                    idx,
                    QItemSelectionModel.ClearAndSelect,
                )

    def redo(self):
        if self.undo_stack.canRedo():
            self.undo_stack.redo()
        elif self.can_redo:
            idx = self.selection_model.currentIndex()
            if idx.row() < self.picture_model.rowCount() - 1:
                idx = self.picture_model.index(idx.row() + 1, 0)
                self.selection_model.setCurrentIndex(
                    idx,
                    QItemSelectionModel.ClearAndSelect,
                )

    def create_undo_action(self):
        act = QAction('Undo')
        act.triggered.connect(self.undo)
        act.setEnabled(False)
        self.can_undo_changed.connect(act.setEnabled)
        return act

    def create_redo_action(self):
        act = QAction('Redo')
        act.triggered.connect(self.redo)
        act.setEnabled(False)
        self.can_redo_changed.connect(act.setEnabled)
        return act

    def update_wet(self, seconds=1):
        self.wet_end = time.monotonic() + seconds

    def update_grab(self, grab):
        was_grabbing_mouse = self._grabbing_mouse
        if self.tool is None:
            grab = False
        if grab:
            if not self.tool:
                return False
            self._grabbing_mouse = True
            self.grabMouse()
        else:
            self.releaseMouse()
            self._grabbing_mouse = False
            if self._last_cursor_pos:
                QCursor.setPos(self._last_cursor_pos)
        self.grab_updated.emit(self._grabbing_mouse)

class Tool:
    name = 'tool'

    def __init__(self):
        self.pen = QPen(
            QColor(0, 0, 0, 0),
            0,
            Qt.SolidLine,
            Qt.RoundCap,
            Qt.BevelJoin,
        )
        self.set_size(1)

    def set_size(self, size):
        self.size = size
        self.alpha = 255
        if size < 1:
            self.alpha = int(255 * size)
            self.size = 1
        self.pen.setWidth(self.size)
        color = self.pen.color()
        color.setAlpha(self.alpha)
        self.pen.setColor(color)

    def draw(self, last, now, update_rect, scribbles, wet):
        for overlay in scribbles[-1], wet:
            if overlay:
                overlay.reserve(update_rect)
                with overlay.painter_context() as painter:
                    painter.setPen(self.pen)
                    painter.setRenderHint(QPainter.Antialiasing)
                    painter.drawLine(last, now)

class Marker(Tool):
    name = 'Marker'
    def set_size(self, size):
        super().set_size(size * MAX_RADIUS / 10)

class ColorMarker(Tool):
    name = 'Color Marker'
    def __init__(self, r, g, b, name=None):
        super().__init__()
        self.pen.setColor(QColor(int(r*255), int(g*255), int(b*255)))
        if name:
            self.name = f'{name} Marker'

    def set_size(self, size):
        super().set_size(size * MAX_RADIUS / 5)

class Highlighter(Tool):
    name = 'Highlighter'
    def set_size(self, size):
        self.pen.setColor(QColor(255, 250, 0, 255))
        super().set_size(size * MAX_RADIUS / 2)

class Eraser(Tool):
    name = 'Eraser'
    def set_size(self, size):
        super().set_size(size * MAX_RADIUS)
        self.pen.setColor(QColor(255, 255, 255, 255))

    def draw(self, last, now, update_rect, scribbles, wet):
        overlay = scribbles[-1]
        overlay.composition_mode = QPainter.CompositionMode_DestinationOut
        overlay.opacity = 1
        super().draw(last, now, update_rect, scribbles, wet)

class WidgetFinder:
    def __init__(self, obj):
        self.obj = obj

    def __getattr__(self, name):
        widget = self.obj.findChild(QObject, name)
        if widget is None:
            raise AttributeError(name)
        return widget

def make_tool_button(text, shortcut):
    btn = QToolButton()
    btn.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Fixed)
    btn.setText(text)
    btn.setShortcut(shortcut)
    btn.setCheckable(True)
    btn.setAutoExclusive(True)
    return btn

def make_toolbox_window(overlay_widget):
    window = QMainWindow()
    window.setWindowFlags(
        window.windowFlags()
        | Qt.FramelessWindowHint
        | Qt.WindowStaysOnTopHint
    )
    window.resize(0, 0)
    central = QWidget()
    window.setCentralWidget(central)
    window.shortcut_to_action = {}
    main_layout = QVBoxLayout()
    central.setLayout(main_layout)

    def add_layout():
        layout = QHBoxLayout(window)
        main_layout.addLayout(layout)
        return layout

    def tool_setter(tool):
        def func():
            overlay_widget.tool = tool
        return func

    layout = add_layout()

    for text, shortcut, tool, activate in (
        ("&Disable", "D", None, False),
        ("&Marker", "M", Marker(), True),
        ("&Hilite", "H", Highlighter(), False),
        ("&Eraser", "E", Eraser(), False),
    ):
        btn = make_tool_button(text, shortcut)
        layout.addWidget(btn)
        if activate:
            btn.setChecked(True)
        window.shortcut_to_action[shortcut] = btn.click
        btn.clicked.connect(tool_setter(tool))

    layout = add_layout()

    for i, (name, color) in enumerate(COLORS.items(), 1):
        tool = ColorMarker(*color, name)
        btn = make_tool_button(str(i), str(i))
        btn.setStyleSheet("background-color: rgb({}, {}, {});".format(
            *[c*255 for c in color])
        )
        layout.addWidget(btn)
        btn.setToolTip(name)
        window.shortcut_to_action[str(i)] = btn.click
        btn.clicked.connect(tool_setter(tool))

    toolbar = window.addToolBar("Main toolbar")

    act_draw = QAction('Draw', window)
    act_draw.setCheckable(True)
    act_draw.setShortcut('Esc')
    act_draw.toggled.connect(overlay_widget.update_grab)
    overlay_widget.grab_updated.connect(act_draw.setChecked)
    window.shortcut_to_action['Esc'] = act_draw.toggle
    window.shortcut_to_action['D'] = act_draw.toggle
    toolbar.addAction(act_draw)

    def add_action(text, func, icon, shortcut=None):
        act = QAction(QIcon.fromTheme(icon), text, window)
        act.triggered.connect(func)
        if shortcut:
            act.setShortcut(shortcut)
            window.shortcut_to_action[shortcut] = act.trigger
        toolbar.addAction(act)
        return act

    for action_factory, icon, shortcut in (
        (overlay_widget.create_undo_action, 'edit-undo-symbolic', 'Z'),
        (overlay_widget.create_redo_action, 'edit-redo-symbolic', 'Y'),
    ):
        act = action_factory()
        act.setIcon(QIcon.fromTheme(icon))
        act.setShortcut(shortcut)
        window.shortcut_to_action[shortcut] = act.trigger
        toolbar.addAction(act)

    clr = add_action('Clear', overlay_widget.clear, 'document-new-symbolic', 'Q')
    overlay_widget.can_clear_changed.connect(clr.setEnabled)
    clr.setEnabled(overlay_widget.can_clear)
    toolbar.addSeparator()
    add_action('Close', sys.exit, 'process-stop-symbolic')

    layout = add_layout()
    ilv = QListView()
    ilv.setModel(overlay_widget.picture_model)
    ilv.setSelectionModel(overlay_widget.selection_model)
    layout.addWidget(ilv)
    layout.addWidget(QUndoView(overlay_widget.undo_group))

    return window

def make_overlay_widget():
    w = OverlayWidget()

    for screen in reversed(app.screens()):
        print(screen.manufacturer())
        if screen.manufacturer().startswith(('Wacom', 'Chimei')):
            geom = screen.geometry()
            w.move(geom.left(), geom.top())
            w.resize(geom.width(), geom.height())

    return w

class Application(QApplication):
    def __init__(self, *args):
        super().__init__(*args)
        self._grabbing_mouse = False
        self._last_cursor_pos = QCursor.pos()

        def update_pos():
            if not self._grabbing_mouse:
                self._last_cursor_pos = QCursor.pos()

        self._timer = QTimer()
        self._timer.timeout.connect(update_pos)
        self._timer.start(100)

    def event(self, e):
        print(e)
        if e.type() == QEvent.TabletEnterProximity:
            print('enter', toolbox.geometry(), QCursor.pos())
            if toolbox.geometry().contains(QCursor.pos()):
                return False
            overlay_widget.update_grab(True)
            return True
        elif e.type() == QEvent.TabletLeaveProximity:
            print('leave')
            overlay_widget.update_grab(False)
            return True
        elif e.type() == QEvent.TabletTrackingChange:
            print('track')
            return True
        elif e.type() == QEvent.User:
            print('app handling', e.key)
            try:
                act = toolbox.shortcut_to_action[e.key]
            except KeyError:
                print(f'no {e.key}...')
            else:
                print(act)
                act()
        return False

def watch_shortcuts(app, widget):
    def post(key):
        app.postEvent(app, MyEvent(key))
    global_shortcuts.watch_shortcuts(post)

class MyEvent(QEvent):
    def __init__(self, key):
        super().__init__(QEvent.User)
        self.key = key

    def __repr__(self):
        return f'my event! {self.key}'

def main():
    global app, toolbox, overlay_widget
    app = Application(sys.argv)

    overlay_widget = make_overlay_widget()

    overlay_widget.showFullScreen()

    toolbox = make_toolbox_window(overlay_widget)
    toolbox.show()
    toolbox.move(overlay_widget.geometry().topLeft())

    app._toolbox = toolbox

    threading.Thread(
        target=watch_shortcuts,
        args=(app, overlay_widget, ),
        daemon=True
    ).start()

    sys.exit(app.exec())

if __name__ == '__main__':
    main()


import requests
from jsonpath_rw import parse

import cumulus
from cumulus.common import check_status

ssh_cluster = ['trad', 'ec2']

def get_assetstore_url_base(cluster):
    if cluster['type'] in ssh_cluster:
        return 'sftp_assetstores'
    elif cluster['type'] == 'newt':
        return 'newt_assetstores'
    else:
        raise Exception('Unsupported cluster type')

def get_assetstore_id(girder_token, cluster):
    if 'assetstoreId' in cluster:
        return cluster['assetstoreId']

    headers = {'Girder-Token':  girder_token}
    url_base = get_assetstore_url_base(cluster)
    create_url = '%s/%s' % (cumulus.config.girder.baseUrl, url_base)
    body = {
        'name': cluster['_id'],
        'host': cluster['config']['host'],
        'machine': cluster['config']['host'],
        'authKey': cluster['_id']
    }

    user = parse('config.ssh.user').find(cluster)
    if user:
        body['user'] = user[0].value

    r = requests.post(create_url, json=body, headers=headers)

    # If the assetstore has been created, patch the cluster to point to it
    if r.status_code == 200:
        assetstore_id = r.json()['_id']

        cluster_url = '%s/clusters/%s' % (cumulus.config.girder.baseUrl,
                                          cluster['_id'])
        body = {
            'assetstoreId': assetstore_id
        }
        r = requests.patch(cluster_url, json=body, headers=headers)
        check_status(r)

    # The assetstore may have already been created by another concurrently
    # running task flow. If thats the case, just fetch it and return its id.
    elif r.status_code == 400:
        body = r.json()
        if body.get('field') == 'name' and body.get('type') == 'validation':
            assetstores_url = '%s/assetstore/lookup' % (
                cumulus.config.girder.baseUrl)
            params = {'name': cluster['_id']}
            r = requests.get(assetstores_url, params=params, headers=headers)
            check_status(r)
            assetstores = r.json()
            if len(assetstores) == 0:
                raise Exception(
                    'Could not find assetstore with name "%s" even though '
                    'it should already exist.' % cluster['_id']
                )

            assetstore_id = assetstores[0]['_id']
        else:
            check_status(r)

    # Raise any other errors
    else:
        check_status(r)

    cluster['assetstoreId'] = assetstore_id

    return assetstore_id