import http
import http.client
import json
import threading
import sys
import queue
from enum import Enum

from ..version import __version__
from ..proxy import Proxy

from .stringextension import StringExtension
from .httpendpoint import HttpEndpoint
from .injectionresponseparser import InjectionResponseParser
from .serialization.injectionrequest import InjectionRequest

class HttpRequest(object):
    """
    Wrapper of the Http Client to handle threading for async
    """

    class HttpRequestMethod(Enum):
        """
        Enumeration of HTTP Request Methods
        """
        GET = 0
        POST = 1
        PUT = 2
        DELETE = 3

    def __init__(self, method: HttpRequestMethod, endpoint: HttpEndpoint, timeout: int, authentication: str):
        """
        Creates a new instance of the HTTP Request class
        :param method: the HTTP request method
        :type method: HttpRequestMethod
        :param endpoint: the Http endpoint for the HTTP request
        :type endpoint: HttpEndpoint
        """
        self._request_method = method
        self._endpoint = endpoint
        self._http_proxy = None
        self._timeout = timeout
        self._authentication = authentication
        self._headers = {
            'User-Agent': self.__user_agent,
            'Content-Type': 'application/json; charset=utf-8',
            'Accept': 'application/json',
        }
        if not StringExtension.is_none_or_white_space(authentication):
            self._headers["Authorization"] = "Bearer " + authentication

    @property
    def __user_agent(self):
        """
        The User-Agent request header added to the Injection API Http request.
        Used to identify the source of the request.
        :return the SocketLabs User-Agent
        :rtype str
        """
        return "SocketLabs-python/{0};python({1})".format(__version__, sys.version.split(' ')[0])

    @property
    def __request_method(self):
        """
        The HTTP Request Method to use
        :return the HTTP request method
        :rtype HttpRequestMethod
        """
        return self._request_method

    @property
    def __endpoint(self):
        """
        The SocketLabs Injection API endpoint
        :return the Http endpoint for the HTTP request
        :rtype HttpEndpoint
        """
        return self._endpoint

    @property
    def __timeout(self):
        """
        The SocketLabs Injection API timeout
        :return the Http timeout for the HTTP request
        :rtype int
        """
        return self._timeout

    @property
    def proxy(self):
        """
        Get the Proxy to use when making the HTTP request
        :return the Proxy to use for the HTTP request
        :rtype Proxy
        """
        return self._http_proxy

    @proxy.setter
    def proxy(self, val: Proxy):
        """
        Set the Proxy to use when making the HTTP request
        :param val: the Proxy to use for the HTTP request
        :type val: Proxy
        """
        self._http_proxy = val

    def send_async_request(self, request: InjectionRequest, on_success_callback, on_error_callback):
        """
        Send an HTTP Request asynchronously
        :param request: the injection request to send
        :type request: InjectionRequest
        :param on_success_callback: the callback method for success
        :type on_success_callback: method
        :param on_error_callback: the callback method for error
        :type on_error_callback: method
        """

        try:
            print (request)
            th = threading.Thread(target=self.__queue_request,
                                  kwargs={
                                      "request": request,
                                      "on_success_callback": on_success_callback,
                                      "on_error_callback": on_error_callback
                                  })
            th.start()

        except Exception as e:
            on_error_callback(e)

    def __queue_request(self, request: InjectionRequest, on_success_callback, on_error_callback):
        """
        queue method for the threaded send request.
        :param request: the injection request to send
        :type request: InjectionRequest
        :param out_queue: the Queue object to handle the response
        :type out_queue: Queue
        """
        try:
            response = self.send_request(request)
            on_success_callback(response)

        except Exception:
            on_error_callback(sys.exc_info()[0])

    def send_request(self, request: InjectionRequest):
        """
        Send the HTTP Request
        :param request: the injection request to send
        :type request: InjectionRequest
        :return the injection response received from the request
        :rtype InjectionResponse
        """

        json_body = json.dumps(request.to_json())

        try:
            connection = self.__get_connection()
            connection.request("POST", self._endpoint.url, json_body, self._headers)
            response = connection.getresponse()

        except Exception as e:
            raise e

        return response

    def __get_connection(self):
        """
        Opens a socket connection to the server to set up an HTTP request.
        :return the HTTPS connection to use in the request
        :rtype HTTPSConnection
        """
        if self._http_proxy is not None:
            connection = http.client.HTTPSConnection(self._http_proxy.host, self._http_proxy.port, timeout=self.__timeout)
            connection.set_tunnel(self._endpoint.host, 443)
        else:
            connection = http.client.HTTPSConnection(self._endpoint.host, timeout=self.__timeout)

        return connection


#! venv/bin/python3
"""a little bot for send paom at One houre to Discord Chanell
"""
import logging
from json import loads
from typing import AnyStr, Dict

from discord.ext import commands
from discord import Intents
from requests import get

TOKEN="TOKEN"

intents = Intents.default()
intents.message_content = True
bot = commands.Bot(command_prefix="/", intents=intents)

logger = logging.getLogger("discord")

logging.basicConfig(level=logging.INFO)
handler = logging.FileHandler(filename="discord.log", encoding="utf-8", mode="w")
handler.setFormatter(
    logging.Formatter("%(asctime)s:%(levelname)s:%(name)s: %(message)s")
)
logger.addHandler(handler)

def get_random_poem_from_ganjoor() -> Dict:
    content = get("http://c.ganjoor.net/beyt-json.php").content
    json_poem = loads(content)

    m1 = json_poem["m1"]
    m2 = json_poem["m2"]
    return {"m1": m1, "m2": m2, "شاعر": json_poem["poet"]}

def send_poem_for_user() -> AnyStr:
    poem_poet = get_random_poem_from_ganjoor()

    return f"{poem_poet['m1']} // {poem_poet['m2']}\n" f"{poem_poet['شاعر']}\n"

@bot.hybrid_command()
async def poem(ctx):
    await ctx.send(send_poem_for_user())

if __name__ == "__main__":
    bot.run(TOKEN)


from __future__ import annotations

# TODO: RF and move all paths related functions from datalad.utils in here
import os
import os.path as op
# to not pollute API importing as _
from collections import defaultdict as _defaultdict
from collections.abc import (
    Iterable,
    Iterator,
)
from functools import wraps
from itertools import dropwhile
from pathlib import (
    Path,
    PurePosixPath,
)

from ..utils import (
    ensure_bytes,
    getpwd,
)

def _get_unicode_robust_version(f):

    @wraps(f)
    def wrapped(*args, **kwargs):
        try:
            return f(*args, **kwargs)
        except UnicodeEncodeError:
            return f(ensure_bytes(*args, **kwargs))
    doc = getattr(f, '__doc__', None)
    # adjust only if __doc__ is not completely absent (None)
    if doc is not None:
        wrapped.__doc__ = doc + \
            "\n\nThis wrapper around original function would encode forcefully " \
            "to utf-8 if initial invocation fails"
    return wrapped

abspath = op.abspath
basename = op.basename
curdir = op.curdir
dirname = op.dirname
exists = _get_unicode_robust_version(op.exists)
isdir = _get_unicode_robust_version(op.isdir)
isabs = _get_unicode_robust_version(op.isabs)
join = op.join
lexists = _get_unicode_robust_version(op.lexists)
normpath = op.normpath
pardir = op.pardir
pathsep = op.pathsep
relpath = op.relpath
realpath = _get_unicode_robust_version(op.realpath)
sep = op.sep

def robust_abspath(p: str | Path) -> str:
    """A helper which would not fail if p is relative and we are in non-existing directory

    It will rely on getpwd, which would rely on $PWD env variable to report
    the path.  Desired for improved resilience during e.g. reporting as in
    https://github.com/datalad/datalad/issues/2787
    """
    try:
        return abspath(p)
    except OSError:
        if not isabs(p):
            try:
                os.getcwd()
            except Exception:
                return normpath(join(getpwd(), p))
        # if no exception raised it was not the reason, raise original
        raise

def split_ext(filename: str) -> tuple[str, str]:
    """Use git-annex's splitShortExtensions rule for splitting extensions.

    Parameters
    ----------
    filename : str

    Returns
    -------
    A tuple with (root, extension)

    Examples
    --------
    >>> from datalad.local.addurls import split_ext
    >>> split_ext("filename.py")
    ('filename', '.py')

    >>> split_ext("filename.tar.gz")
    ('filename', '.tar.gz')

    >>> split_ext("filename.above4chars.ext")
    ('filename.above4chars', '.ext')
    """
    parts = filename.split(".")
    if len(parts) == 1:
        return filename, ""

    tail = list(dropwhile(lambda x: len(x) < 5,
                          reversed(parts[1:])))

    file_parts = parts[:1] + tail[::-1]
    ext_parts = parts[1+len(tail):]
    return ".".join(file_parts), "." + ".".join(ext_parts)

def get_parent_paths(paths: list[str], parents: list[str], only_with_parents: bool = False, *, sep: str = '/') -> list[str]:
    """Given a list of children paths, return their parent paths among parents
    or their own path if there is no known parent. A path is also considered its
    own parent (haven't you watched Predestination?) ;)

    All paths should be relative, not pointing outside (not starting
    with ../), and normalized (no // or dir/../dir and alike). Only minimal
    sanity checking of values is done.  By default paths are considered to be
    POSIX. Use 'sep' kwarg to set to `os.sep` to provide OS specific handling.

    Accent is made on performance to avoid O(len(paths) * len(parents))
    runtime.  ATM should be typically less than O(len(paths) * len(log(parents)))

    Initial intended use - for a list of paths in the repository
    to provide their paths as files/submodules known to that repository, to
    overcome difference in ls-tree and ls-files, where ls-files outputs nothing
    for paths within submodules.
    It is coded, so it could later be applied even whenever there are nested
    parents, e.g. parents = ['sub', 'sub/sub'] and then the "deepest" parent
    is selected

    Parameters
    ----------
    parents: list of str
    paths: list of str
    only_with_parents: bool, optional
      If set to True, return a list of only parent paths where that path had
      a parent
    sep: str, optional
      Path separator.  By default - '/' and thus treating paths as POSIX.
      If you are processing OS-specific paths (for both `parents` and `paths`),
      specify `sep=os.sep`.

    Returns
    -------
    A list of paths (without duplicates), where some entries replaced with
    their "parents" without duplicates.  So for 'a/b' and 'a/c' with a being
    among parents, there will be a single 'a'
    """
    # Let's do an early check even though then we would skip the checks on paths
    # being relative etc
    if not parents:
        return [] if only_with_parents else paths

    # We will create a lookup for known parent lengths
    parent_set = set(parents)  # O(log(len(parents))) lookup

    # Will be used in sanity checking that we got consistently used separators, i.e.
    # not mixing non-POSIX paths and POSIX parents
    asep = {'/': '\\', '\\': '/'}[sep]

    # rely on path[:n] be quick, and len(parent_lengths) << len(parent_set)
    # when len(parent_set) is large.  We will also bail checking any parent of
    # the length if at that length path has no directory boundary ('/').
    #
    # Create mapping for each length of
    # parent path to list of parents with that length
    parent_lengths_map: dict[int, set[str]] = _defaultdict(set)
    for parent in parent_set:
        _get_parent_paths_check(parent, sep, asep)
        parent_lengths_map[len(parent)].add(parent)

    # Make it ordered in the descending order so we select the deepest/longest parent
    # and store them as sets for faster lookup.
    # Could be an ordered dict but no need
    parent_lengths = [(l, parent_lengths_map[l]) for l in sorted(parent_lengths_map, reverse=True)]

    res = []
    seen = set()

    for path in paths:  # O(len(paths)) - unavoidable but could be parallelized!
        # Sanity check -- should not be too expensive
        _get_parent_paths_check(path, sep, asep)
        for parent_length, parents_ in parent_lengths:  # O(len(parent_lengths))
            if (len(path) < parent_length) or (len(path) > parent_length and path[parent_length] != sep):
                continue  # no directory deep enough
            candidate_parent = path[:parent_length]
            if candidate_parent in parents_:  # O(log(len(parent_set))) but expected one less due to per length handling
                if candidate_parent not in seen:
                    res.append(candidate_parent)
                    seen.add(candidate_parent)
                break  # it is!
        else:  # no hits
            if not only_with_parents:
                if path not in seen:
                    res.append(path)
                    seen.add(path)

    return res

def get_filtered_paths_(paths: Iterable[str|Path], filter_paths: Iterable[str | Path],
                        *, include_within_path: bool = False) \
        -> Iterator[str]:
    """Among paths (or Path objects) select the ones within filter_paths.

    All `paths` and `filter_paths` must be relative and POSIX.

    In case of `include_with_path=True`, if a `filter_path` points to some path
    under a `path` within `paths`, that path would be returned as well, e.g.
    `path` 'submod' would be returned if there is a `filter_path` 'submod/subsub/file'.

    Complexity is O(N*log(N)), where N is the largest of the lengths of `paths`
    or `filter_paths`.

    Yields
    ------
    paths, sorted (so order is not preserved), which reside under 'filter_paths' or
    path within 'filter_paths' is under that path.
    """
    # do conversion and sanity checks, O(N)
    def _harmonize_paths(l: Iterable[str | Path]) -> list[tuple[str, ...]]:
        ps = []
        for p in l:
            pp = PurePosixPath(p)
            if pp.is_absolute():
                raise ValueError(f"Got absolute path {p}, expected relative")
            if pp.parts and pp.parts[0] == '..':
                raise ValueError(f"Path {p} leads outside")
            ps.append(pp.parts)  # store parts
        return sorted(ps)  # O(N * log(N))

    paths_parts = _harmonize_paths(paths)
    filter_paths_parts = _harmonize_paths(filter_paths)

    # we will pretty much "scroll" through sorted paths and filter_paths at the same time
    for path_parts in paths_parts:
        while filter_paths_parts:
            filter_path_parts = filter_paths_parts[0]
            l = min(len(path_parts), len(filter_path_parts))
            # if common part is "greater" in the path -- we can go to the next "filter"
            if filter_path_parts[:l] < path_parts[:l]:
                # get to the next one
                filter_paths_parts = filter_paths_parts[1:]
            else:
                break  # otherwise -- consider this one!
        else:
            # no filter path left - the other paths cannot be the selected ones
            break
        if include_within_path:
            # if one identical or subpath of another one -- their parts match in the beginning
            # and we will just reuse that 'l'
            pass
        else:
            # if all components of the filter match, for that we also add len(path_parts) check below
            l = len(filter_path_parts)
        if len(path_parts) >= l and (path_parts[:l] == filter_path_parts[:l]):
            yield '/'.join(path_parts)

def _get_parent_paths_check(path: str, sep: str, asep: str) -> None:
    """A little helper for get_parent_paths"""
    if isabs(path) or path.startswith(pardir + sep) or path.startswith(curdir + sep):
        raise ValueError("Expected relative within directory paths, got %r" % path)
    if sep+sep in path:
        raise ValueError(f"Expected normalized paths, got {path} containing '{sep+sep}'")
    if asep in path:
        raise ValueError(f"Expected paths with {sep} as separator, got {path} containing '{asep}'")


import navis
import requests
import warnings

import cloudvolume as cv
import numpy as np
import pandas as pd
import trimesh as tm

from abc import ABC
from io import StringIO, BytesIO

from ..utils import make_iterable

use_pbars = True

class OnDemandDict(dict):
    """Initialized with a just a URL.

    Will fetch data from given URL on first request.

    """

    def __init__(self, url):
        """Initialize with just a URL."""
        self.url = url
        self.fetched = False
        super().__init__()

    def __getitem__(self, key):
        if not self.fetched:
            self.update_from_url()
        return super().__getitem__(key)

    def __contains__(self, key):
        if not self.fetched:
            self.update_from_url()
        return super().__contains__(key)

    def get(self, *args, **kwargs):
        if not self.fetched:
            self.update_from_url()
        return super().get(*args, **kwargs)

    def update_from_url(self):
        """Update content from URL."""
        resp = requests.get(self.url)
        resp.raise_for_status()
        self.update(resp.json())
        self.fetched = True

TRANSFORM_SERVICE_URL = 'https://services.itanna.io/app/transform-service'
SYNAPSE_SERVICE_URL = 'https://services.itanna.io/app/synapse-service'

TRANSFORM_SERVICE_INFO = OnDemandDict(f'{TRANSFORM_SERVICE_URL}/info')
SYNAPSE_SERVICE_INFO = OnDemandDict(f'{SYNAPSE_SERVICE_URL}/info')

def spine_transform_datasets():
    """Fetch datasets available on spine's transform service.

    Returns
    -------
    dict

    """
    resp = requests.get(f'{TRANSFORM_SERVICE_URL}/info')
    resp.raise_for_status()
    return resp.json()

class SpineService(ABC):
    """Abstract base class for spine services."""

    @property
    def info(self):
        """Return general info on given service."""
        if not hasattr(self, '_info'):
            url = self.makeurl('info')
            resp = self.session.get(url)
            resp.raise_for_status()
            self._info = resp.json()
        return self._info

    def urljoin(self, *args):
        """Join arguments into an url.

        Trailing but not leading slashes are stripped for each argument.

        """
        return "/".join(map(lambda x: str(x).rstrip('/'), args))

    def makeurl(self, *args):
        """Generate URL from base URL and args."""
        return self.urljoin(self.base_url, *args)

class FlyCacheService(SpineService):
    """Interface with cache services."""

    def __init__(self,
                 base_url='https://services.itanna.io/app/flycache-dev'):
        """Init class."""
        self.base_url = base_url

    def get_L2_centroids(self, ids, token, as_array=False, chunksize=50,
                         progress=True):
        """Fetch centroids of given l2 chunks.

        Coordinates are in nm.

        Parameter
        ---------
        ids :           iterable
                        Iterable with L2 IDs.
        token :         str
                        A ChunkedGraph/CAVE auth token.
        as_array :      bool
                        Determines output (see Returns).
        chunksize :     int
                        Query L2 IDs in chunks of this size.
        progress :      bool
                        Whether to show a progress bar or not.

        Returns
        -------
        array
                        If `as_array=True`: array in same order as queried `ids`.
                        Missing centroids are returns with coordinate (0, 0, 0).
        dict
                        If `as_array=False`: dictionary mapping ID to
                        centroid `{L2_ID: [x, y, z], ..}`.

        """
        url = self.makeurl('mesh/l2_centroid/flywire_fafb_production/')

        # Make sure we have an array of integers
        ids = make_iterable(ids, force_type=np.int64)

        with navis.config.tqdm(total=len(ids), desc='Fetching centroids',
                               leave=False, disable=not progress) as pbar:

            # First we will get everything that's cached in a single big query
            post = {
                      "token": token,
                      "query_ids": ids.tolist()
                    }
            resp = self.session.post(url + '?cache_only=1', json=post)
            resp.raise_for_status()

            # Parse response
            centroids = {np.int64(k): v for k, v in resp.json().items()}

            # Filter to remaining IDs
            to_fetch = ids[~np.isin(ids, list(centroids))]

            # Update progress bar
            pbar.update(len(centroids))

            # Now go over the remaining indices in chunks
            for i in range(0, len(to_fetch), int(chunksize)):
                this_chunk = to_fetch[i:i+chunksize]
                post = {
                          "token": token,
                          "query_ids": this_chunk.tolist()
                        }

                resp = self.session.post(url, json=post)
                resp.raise_for_status()

                # Parse response
                data = resp.json()

                # IDs will have been returned as strings
                centroids.update({np.int64(k): v for k, v in data.items()})

                pbar.update(len(this_chunk))

        if not as_array:
            return centroids

        centroids = np.array([centroids.get(i, [0, 0, 0]) for i in ids])

        return centroids

class SynapseService(SpineService):
    """Interface with synapse service on `services.itanna.io`.

    Kindly hosted by Eric Perlman and Davi Bock! Check out the online docs for
    available API endpoints:
    https://services.itanna.io/app/synapse-service/docs

    Parameters
    ----------
    base_url :      str
                    Base URL for the service on spine.

    """

    def __init__(self,
                 base_url='https://services.itanna.io/app/synapse-service'):
        """Init class."""
        self.base_url = base_url
        self.session = requests.Session()

    @property
    def alignments(self):
        """Return available alignments of synapse data."""
        if not hasattr(self, '_alignments'):
            url = self.makeurl('alignments')
            resp = self.session.get(url)
            resp.raise_for_status()
            self._alignments = resp.json()
        return self._alignments

    @property
    def collections(self):
        """Return available collections of synapse data."""
        if not hasattr(self, '_collections'):
            url = self.makeurl('collections')
            resp = self.session.get(url)
            resp.raise_for_status()
            self._collections = resp.json()
        return self._collections

    @property
    def segmentations(self):
        """Return available segmentations the synapse data is mapped to."""
        if not hasattr(self, '_segmentations'):
            url = self.makeurl('segmentations')
            resp = self.session.get(url)
            resp.raise_for_status()
            self._segmentations = resp.json()
        return self._segmentations

    def validate_alignment(self, alignment):
        """Check if alignment exists."""
        available = [c.get('name', 'NA') for c in self.alignments]
        if alignment not in available:
            raise ValueError(f'{alignment} not among available alignments: '
                             f'{",".join(available)}')

    def validate_collection(self, collection):
        """Check if collection exists."""
        available = [c.get('name', 'NA') for c in self.collections]
        if collection not in available:
            raise ValueError(f'{collection} not among available collections: '
                             f'{",".join(available)}')

    def validate_segmentation(self, segmentation):
        """Check if alignment exists."""
        available = [c.get('name', 'NA') for c in self.segmentations]
        if segmentation not in available:
            raise ValueError(f'{segmentation} not among available segmentations: '
                             f'{",".join(available)}')

    def get_synapse(self, synapse_id, collection):
        """Return all available info (pre/post IDs, location) for single synapse.

        Parameter
        ---------
        synapse_id :    int
                        ID of synapse to query.
        collection :    str
                        Collection to query for the synapse. See ``.collections``
                        for a list of available collections. Currently available:

                          - "buhmann2019": Buhmann synapses

        Returns
        -------
        dict

        """
        assert isinstance(synapse_id, (int, np.integer, str))

        self.validate_collection(collection)

        url = self.makeurl('collection', collection, 'synapse', synapse_id, 'info')
        resp = self.session.get(url)
        resp.raise_for_status()
        return resp.json()

    def get_connectivity(self, segmentation_ids, segmentation, locations=False,
                         nt_predictions=False):
        """Fetch all connections from/to given segmentation ID(s).

        Parameters
        ----------
        segmentation_ids :  int | list thereof
                            Segmentation/root IDs to query.
        segmentation :      str
                            Which segmentation to search. Currently available:

                              - "fafb-ffn1-20190805" (segment IDs)
                              - "fafb-ffn1-20200412" (segment IDs)
                              - "flywire_supervoxels" (supervoxel IDs)
        locations :         bool
                            Whether to also fetch x/y/z locations for each
                            synapse.
        nt_predictions :    bool
                            Whether to also fetch neurotransmitter predictions
                            for each synapses.

        Returns
        -------
        pandas.DataFrame

        """
        # This always returns a numpy array
        segmentation_ids = navis.utils.make_iterable(segmentation_ids)

        # Check if segmentation actually exists
        self.validate_segmentation(segmentation)

        url = self.makeurl('segmentation', segmentation, 'feather')

        param = []
        if locations:
            param.append("locations=true")
        if nt_predictions:
            param.append("nt=eckstein2020")

        if param:
            url += f'?{"&".join(param)}'

        post = {"query_ids": segmentation_ids.tolist()}

        resp = self.session.post(url, json=post)
        resp.raise_for_status()

        # Read into DataFrame
        return pd.read_feather(BytesIO(resp.content))

class TransformService(SpineService):
    """Interface with transform service on `services.itanna.io`.

    Kindly hosted by Eric Perlman and Davi Bock! Check out the online docs for
    available API endpoints:
    https://services.itanna.io/app/transform-service/docs

    Parameters
    ----------
    base_url :      str
                    Base URL for the service on spine.

    """

    # Hard-coded data types for now. From:
    # https://github.com/bocklab/transform_service/blob/master/app/config.py
    DTYPES = {
              'test': np.float32,
              'flywire_v1': np.float32,
              'flywire_v1_inverse': np.float32,
              'fanc_v4_to_v3': np.float32,
              'fafb-ffn1-20200412': np.uint64,
              'fafb-ffn1-20200412-gcs': np.uint64,
              'fafb-ffn1-20200412-zarr': np.uint64,
              'flywire_190410': np.uint64
              }

    def __init__(self,
                 base_url='https://services.itanna.io/app/transform-service'):
        """Init class."""
        self.base_url = base_url
        self.session = requests.Session()

    def validate_dataset(self, dataset):
        """Check if dataset exists for given service."""
        if dataset not in self.info:
            ds_str = ", ".join(self.info.keys())
            raise ValueError(f'"{dataset}" not among listed datasets: {ds_str}')

    def validate_mip(self, mip, dataset):
        """Validate mip for given dataset."""
        assert isinstance(mip, (int, np.integer))
        if mip < 0:
            available_scales = sorted(self.info[dataset]['scales'])
            mip = available_scales[-mip - 1]
        elif mip not in self.info[dataset]['scales']:
            raise ValueError(f'mip {mip} not available for dataset "{dataset}". '
                             f'Available scales: {self.info["scales"]}')
        return mip

    def validate_output(self, x, on_fail='ignore'):
        """Validate output."""
        # If mapping failed will contain NaNs
        if on_fail != 'ignore':
            is_nan = np.any(np.isnan(x),
                            axis=1 if x.ndim == 2 else 0)
            if np.any(is_nan):
                msg = f'{is_nan.sum()} points failed to transform.'
                if on_fail == 'warn':
                    warnings.warn(msg)
                elif on_fail == 'raise':
                    raise ValueError(msg)

    def to_voxels(self, x, dataset, coordinates='nm'):
        """Parse spatial data into voxels coordinates.

        Parameters
        ----------
        x :             np.ndarray (N, 3) | Neuron/List | mesh
                        Data to transform.
        coordinates :   "nm" | "voxel"
                        Units of the coordinates in ``x``.

        Returns
        -------
        vxl :           np.ndarray (N, 3)
                        x/y/z coordinates extracted from x and converted to
                        voxels.

        """
        assert coordinates in ['nm',
                               'nanometer', 'nanometers',
                               'nanometre', 'nanometres',
                               'vxl', 'voxel', 'voxels']

        if isinstance(x, (navis.NeuronList, navis.TreeNeuron)):
            x = x.nodes[['x', 'y', 'z']].values
        elif isinstance(x, (navis.Volume, navis.MeshNeuron, tm.Trimesh)):
            x = np.asarray(x.vertices)

        # At this point we are expecting a numpy array
        vxl = np.asarray(x)

        # Make sure data is now in the correct format
        if vxl.ndim != 2 or vxl.shape[1] != 3:
            raise TypeError(f'Expected (N, 3) array, got {x.shape}')

        # We need to convert to voxel coordinates
        # Note that we are rounding here to get to voxels.
        # This will have the most impact on the Z section.
        if coordinates not in ['vxl', 'voxel', 'voxels']:
            # Make sure we are working with numbers
            # -> if dtype is "object" we will get errors from np.round
            if not np.issubdtype(vxl.dtype, np.number):
                vxl = vxl.astype(np.float64)

            # Convert to voxels
            vxl_size = self.info[dataset]['voxel_size']
            vxl = np.round(vxl / vxl_size).astype(np.int64)

        return vxl

    def get_offsets(self, x, transform, coordinates='nm', mip=-1, limit_request=10e9,
                    on_fail='warn'):
        """Transform coordinates.

        Parameters
        ----------
        x :             np.ndarray (N, 3) | Neuron/List | mesh
                        Data to transform.
        transform :     str
                        Dataset to use for transform. Currently available:

                         - 'flywire_v1': flywire to FAFB
                         - 'flywire_v1_inverse': FAFB to flywire
                         - 'fanc_v4_to_v3': FANC v4 to v3

                        See also ``.info`` for up-to-date info.
        mip :           int
                        Resolution of mapping. Lower = more precise but much slower.
                        Negative values start counting from the highest possible
                        resolution: -1 = highest, -2 = second highest, etc.
        coordinates :   "nm" | "voxel"
                        Units of the coordinates in ``x``.
        on_fail :       "warn" | "ignore" | "raise"
                        What to do if points failed to xform.

        Returns
        -------
        response :      np.ndarray
                        (N, 2) of dx and dy offsets.

        """
        assert on_fail in ['warn', 'raise', 'ignore']

        # Check if dataset exists
        self.validate_dataset(transform)

        # Parse mip
        mip = self.validate_mip(mip, dataset=transform)

        # Extract voxels from x
        vxl = self.to_voxels(x, transform, coordinates=coordinates)

        # Generate URL
        url = self.makeurl('transform/dataset', transform, 's', mip,
                           'values_binary/format/array_float_Nx3')

        # Make sure we don't exceed the maximum size for each request
        stack = []
        limit_request = int(limit_request)
        for ix in np.arange(0, vxl.shape[0], limit_request):
            this_vxl = vxl[ix: ix + limit_request]

            # Make request
            resp = self.session.post(url,
                                 data=this_vxl.astype(np.single).tobytes(order='C'))

            # Check for errors
            resp.raise_for_status()

            # Extract data
            data = np.frombuffer(resp.content, dtype=self.DTYPES[transform])

            # Reshape to [[dx1, dy1], [dx2, dy2], ...]
            data = data.reshape(this_vxl.shape[0], 2)
            stack.append(data)

        stack = np.concatenate(stack, axis=0)

        # See if any points failed to xform, and raise/warn if requested
        self.validate_output(stack, on_fail=on_fail)

        return stack

    def get_segids(self, x, segmentation, coordinates='nm', mip=-1,
                   limit_request=10e9, on_fail='warn'):
        """Fetch segmentation/supervoxel IDs.

        Parameters
        ----------
        x :             np.ndarray (N, 3) | Neuron/List | mesh
                        Coordinates to query segmentation IDs for.
        segmentation :  str
                        Segmentation to query. Currently available:

                         - 'fafb-ffn1-20200412'
                         - 'fafb-ffn1-20200412-gcs'
                         - 'fafb-ffn1-20200412-zarr'
                         - 'flywire_190410' (supervoxels)

        mip :           int
                        Resolution of mapping. Lower = more precise but much slower.
                        Negative values start counting from the highest possible
                        resolution: -1 = highest, -2 = second highest, etc.
        coordinates :   "nm" | "voxel"
                        Units of the coordinates in ``x``.
        limit_request : int
                        Max number of locations to query per request.
        on_fail :       "warn" | "ignore" | "raise"
                        What to do if points fail to return segmentation IDs.

        Returns
        -------
        response :      np.ndarray
                        (N, 2) of dx and dy offsets.

        """
        assert on_fail in ['warn', 'raise', 'ignore']

        # Check if dataset exists
        self.validate_dataset(segmentation)

        # Parse mip
        mip = self.validate_mip(mip, dataset=segmentation)

        # Extract voxels from x
        vxl = self.to_voxels(x, segmentation, coordinates=coordinates)

        # Generate URL
        url = self.makeurl('query/dataset', segmentation, 's', mip,
                           'values_binary/format/array_float_Nx3')

        # Make sure we don't exceed the maximum size for each request
        stack = []
        limit_request = int(limit_request)
        for ix in np.arange(0, vxl.shape[0], limit_request):
            this_vxl = vxl[ix: ix + limit_request]

            # Make request
            resp = self.session.post(url,
                                 data=this_vxl.astype(np.single).tobytes(order='C'))

            # Check for errors
            resp.raise_for_status()

            # Extract data
            data = np.frombuffer(resp.content, dtype=self.DTYPES[segmentation])

            stack.append(data)

        stack = np.concatenate(stack, axis=0)

        # See if any points failed to xform, and raise/warn if requested
        self.validate_output(stack, on_fail=on_fail)

        return stack

def query_spine_transform(x, dataset, query, coordinates='nm', mip=2,
                          limit_request=10e9, on_fail='warn'):
    """Fetch data via the transform service on spine.

    DEPCREATED! Use `spine.transform` or `spine.synapses` instead.

    Kindly hosted by Eric Perlman and Davi Bock! Check out
    https://services.itanna.io/app/transform-service/docs for available
    API endpoints.

    Parameters
    ----------
    x :             np.ndarray (N, 3) | Neuron/List | mesh
                    Data to transform.
    dataset :       str
                    Dataset to use for transform. Currently available:

                     - 'flywire_v1' (transformation)
                     - 'flywire_v1_inverse' (transformation)
                     - 'fanc_v4_to_v3' (transformation)
                     - 'fafb-ffn1-20200412' (query)
                     - 'fafb-ffn1-20200412-gcs' (query)
                     - 'fafb-ffn1-20200412-zarr' (query)
                     - 'flywire_190410' (query)

                    See also ``spine_transform_datasets()`` for up-to-date info.
    query :         "query" | "transform"
                    Whether we "query" segmentation IDs at given location(s) or
                    we want to "transform" locations into a different space.
    mip :           int
                    Resolution of mapping. Lower = more precise but much slower.
                    Negative values start counting from the highest possible
                    resolution: -1 = highest, -2 = second highest, etc.
    coordinates :   "nm" | "voxel"
                    Units of the provided coordinates in ``x``.
    on_fail :       "warn" | "ignore" | "raise"
                    What to do if points failed to xform.

    Returns
    -------
    response :      np.ndarray
                    (N, 2) of dx and dy coordinates for "transform"
                    (N, ) of IDs for "query"

    """
    # Hard-coded data types for now. From:
    # https://github.com/bocklab/transform_service/blob/master/app/config.py
    dtypes = {
              'test': np.float32,
              'flywire_v1': np.float32,
              'flywire_v1_inverse': np.float32,
              'fanc_v4_to_v3': np.float32,
              'fafb-ffn1-20200412': np.uint64,
              'fafb-ffn1-20200412-gcs': np.uint64,
              'fafb-ffn1-20200412-zarr': np.uint64,
              'flywire_190410': np.uint64
              }

    assert on_fail in ['warn', 'raise', 'ignore']
    assert query in ['query', 'transform']
    assert coordinates in ['nm',
                           'nanometer', 'nanometers',
                           'nanometre', 'nanometres',
                           'vxl', 'voxel', 'voxels']
    assert isinstance(mip, (int, np.integer))

    if mip < 0:
        available_scales = sorted(TRANSFORM_SERVICE_INFO[dataset]['scales'])
        mip = available_scales[-mip - 1]

    if dataset not in TRANSFORM_SERVICE_INFO:
        ds_str = ", ".join(TRANSFORM_SERVICE_INFO.keys())
        raise ValueError(f'"{dataset}" not among listed datasets: {ds_str}')

    if mip not in TRANSFORM_SERVICE_INFO[dataset]['scales']:
        raise ValueError(f'mip {mip} not available for dataset "{dataset}". '
                         f'Available scales: {TRANSFORM_SERVICE_INFO[dataset]["scales"]}')

    if isinstance(x, (navis.NeuronList, navis.TreeNeuron)):
        x = x.nodes[['x', 'y', 'z']].values
    elif isinstance(x, (navis.Volume, navis.MeshNeuron, tm.Trimesh)):
        x = np.asarray(x.vertices)

    # At this point we are expecting a numpy array
    x = np.asarray(x)

    # Make sure data is now in the correct format
    if x.ndim != 2 or x.shape[1] != 3:
        raise TypeError('Expected (N, 3) array, got {}'.format(x.shape))

    # We need to convert to voxel coordinates
    # Note that we are rounding here to get to voxels.
    # This will have the most impact on the Z section.
    if coordinates not in ['vxl', 'voxel', 'voxels']:
        x = np.round(x / [4, 4, 40]).astype(np.int64)

    # Generate URL
    url = f'{TRANSFORM_SERVICE_URL}/{query}/dataset/{dataset}/s/{mip}/values_binary/format/array_float_Nx3'

    # Make sure we don't exceed the maximum size for each request
    stack = []
    limit_request = int(limit_request)
    for ix in np.arange(0, x.shape[0], limit_request):
        this_x = x[ix: ix + limit_request]

        # Make request
        resp = requests.post(url,
                             data=this_x.astype(np.single).tobytes(order='C'))

        # Check for errors
        resp.raise_for_status()

        # Extract data
        data = np.frombuffer(resp.content, dtype=dtypes[dataset])
        if query == 'transform':
            data = data.reshape(this_x.shape[0], 2)
        stack.append(data)

    stack = np.concatenate(stack, axis=0)

    # If mapping failed will contain NaNs
    if on_fail != 'ignore':
        is_nan = np.any(np.isnan(stack),
                        axis=1 if stack.ndim == 2 else 0)
        if np.any(is_nan):
            msg = f'{is_nan.sum()} points failed to transform.'
            if on_fail == 'warn':
                warnings.warn(msg)
            elif on_fail == 'raise':
                raise ValueError(msg)

    return stack

synapses = SynapseService()
transform = TransformService()
flycache = FlyCacheService()


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

from __future__ import absolute_import, division, print_function, unicode_literals
from warnings import warn
import math
import torch
from typing import Optional
from torchaudio import functional as F
from torchaudio.compliance import kaldi

__all__ = [
    'Spectrogram',
    'GriffinLim',
    'AmplitudeToDB',
    'MelScale',
    'InverseMelScale',
    'MelSpectrogram',
    'MFCC',
    'MuLawEncoding',
    'MuLawDecoding',
    'Resample',
    'ComplexNorm',
    'TimeStretch',
    'Fade',
    'FrequencyMasking',
    'TimeMasking',
]

class Spectrogram(torch.nn.Module):
    r"""Create a spectrogram from a audio signal.

    Args:
        n_fft (int, optional): Size of FFT, creates ``n_fft // 2 + 1`` bins. (Default: ``400``)
        win_length (int or None, optional): Window size. (Default: ``n_fft``)
        hop_length (int or None, optional): Length of hop between STFT windows. (Default: ``win_length // 2``)
        pad (int, optional): Two sided padding of signal. (Default: ``0``)
        window_fn (Callable[[...], torch.Tensor], optional): A function to create a window tensor
            that is applied/multiplied to each frame/window. (Default: ``torch.hann_window``)
        power (float or None, optional): Exponent for the magnitude spectrogram,
            (must be > 0) e.g., 1 for energy, 2 for power, etc.
            If None, then the complex spectrum is returned instead. (Default: ``2``)
        normalized (bool, optional): Whether to normalize by magnitude after stft. (Default: ``False``)
        wkwargs (Dict[..., ...] or None, optional): Arguments for window function. (Default: ``None``)
    """
    __constants__ = ['n_fft', 'win_length', 'hop_length', 'pad', 'power', 'normalized']

    def __init__(self, n_fft=400, win_length=None, hop_length=None,
                 pad=0, window_fn=torch.hann_window,
                 power=2., normalized=False, wkwargs=None):
        super(Spectrogram, self).__init__()
        self.n_fft = n_fft
        # number of FFT bins. the returned STFT result will have n_fft // 2 + 1
        # number of frequecies due to onesided=True in torch.stft
        self.win_length = win_length if win_length is not None else n_fft
        self.hop_length = hop_length if hop_length is not None else self.win_length // 2
        window = window_fn(self.win_length) if wkwargs is None else window_fn(self.win_length, **wkwargs)
        self.register_buffer('window', window)
        self.pad = pad
        self.power = power
        self.normalized = normalized

    def forward(self, waveform):
        r"""
        Args:
            waveform (torch.Tensor): Tensor of audio of dimension (..., time).

        Returns:
            torch.Tensor: Dimension (..., freq, time), where freq is
            ``n_fft // 2 + 1`` where ``n_fft`` is the number of
            Fourier bins, and time is the number of window hops (n_frame).
        """
        return F.spectrogram(waveform, self.pad, self.window, self.n_fft, self.hop_length,
                             self.win_length, self.power, self.normalized)

class GriffinLim(torch.nn.Module):
    r"""Compute waveform from a linear scale magnitude spectrogram using the Griffin-Lim transformation.
        Implementation ported from `librosa`.

    .. [1] McFee, Brian, Colin Raffel, Dawen Liang, Daniel PW Ellis, Matt McVicar, Eric Battenberg, and Oriol Nieto.
        "librosa: Audio and music signal analysis in python."
        In Proceedings of the 14th python in science conference, pp. 18-25. 2015.

    .. [2] Perraudin, N., Balazs, P., & Søndergaard, P. L.
        "A fast Griffin-Lim algorithm,"
        IEEE Workshop on Applications of Signal Processing to Audio and Acoustics (pp. 1-4),
        Oct. 2013.

    .. [3] D. W. Griffin and J. S. Lim,
        "Signal estimation from modified short-time Fourier transform,"
        IEEE Trans. ASSP, vol.32, no.2, pp.236–243, Apr. 1984.

    Args:
        n_fft (int, optional): Size of FFT, creates ``n_fft // 2 + 1`` bins. (Default: ``400``)
        n_iter (int, optional): Number of iteration for phase recovery process. (Default: ``32``)
        win_length (int or None, optional): Window size. (Default: ``n_fft``)
        hop_length (int or None, optional): Length of hop between STFT windows. (Default: ``win_length // 2``)
        window_fn (Callable[[...], torch.Tensor], optional): A function to create a window tensor
            that is applied/multiplied to each frame/window. (Default: ``torch.hann_window``)
        power (float, optional): Exponent for the magnitude spectrogram,
            (must be > 0) e.g., 1 for energy, 2 for power, etc. (Default: ``2``)
        normalized (bool, optional): Whether to normalize by magnitude after stft. (Default: ``False``)
        wkwargs (Dict[..., ...] or None, optional): Arguments for window function. (Default: ``None``)
        momentum (float, optional): The momentum parameter for fast Griffin-Lim.
            Setting this to 0 recovers the original Griffin-Lim method.
            Values near 1 can lead to faster convergence, but above 1 may not converge. (Default: ``0.99``)
        length (int, optional): Array length of the expected output. (Default: ``None``)
        rand_init (bool, optional): Initializes phase randomly if True and to zero otherwise. (Default: ``True``)
    """
    __constants__ = ['n_fft', 'n_iter', 'win_length', 'hop_length', 'power', 'normalized',
                     'length', 'momentum', 'rand_init']

    def __init__(self, n_fft=400, n_iter=32, win_length=None, hop_length=None,
                 window_fn=torch.hann_window, power=2., normalized=False, wkwargs=None,
                 momentum=0.99, length=None, rand_init=True):
        super(GriffinLim, self).__init__()

        assert momentum < 1, 'momentum=%s > 1 can be unstable' % momentum
        assert momentum > 0, 'momentum=%s < 0' % momentum

        self.n_fft = n_fft
        self.n_iter = n_iter
        self.win_length = win_length if win_length is not None else n_fft
        self.hop_length = hop_length if hop_length is not None else self.win_length // 2
        window = window_fn(self.win_length) if wkwargs is None else window_fn(self.win_length, **wkwargs)
        self.register_buffer('window', window)
        self.normalized = normalized
        self.length = length
        self.power = power
        self.momentum = momentum / (1 + momentum)
        self.rand_init = rand_init

    def forward(self, S):
        return F.griffinlim(S, self.window, self.n_fft, self.hop_length, self.win_length, self.power,
                            self.normalized, self.n_iter, self.momentum, self.length, self.rand_init)

class AmplitudeToDB(torch.nn.Module):
    r"""Turn a tensor from the power/amplitude scale to the decibel scale.

    This output depends on the maximum value in the input tensor, and so
    may return different values for an audio clip split into snippets vs. a
    a full clip.

    Args:
        stype (str, optional): scale of input tensor ('power' or 'magnitude'). The
            power being the elementwise square of the magnitude. (Default: ``'power'``)
        top_db (float, optional): minimum negative cut-off in decibels.  A reasonable number
            is 80. (Default: ``None``)
    """
    __constants__ = ['multiplier', 'amin', 'ref_value', 'db_multiplier']

    def __init__(self, stype='power', top_db=None):
        super(AmplitudeToDB, self).__init__()
        self.stype = stype
        if top_db is not None and top_db < 0:
            raise ValueError('top_db must be positive value')
        self.top_db = top_db
        self.multiplier = 10.0 if stype == 'power' else 20.0
        self.amin = 1e-10
        self.ref_value = 1.0
        self.db_multiplier = math.log10(max(self.amin, self.ref_value))

    def forward(self, x):
        r"""Numerically stable implementation from Librosa.
        https://librosa.github.io/librosa/_modules/librosa/core/spectrum.html

        Args:
            x (torch.Tensor): Input tensor before being converted to decibel scale.

        Returns:
            torch.Tensor: Output tensor in decibel scale.
        """
        return F.amplitude_to_DB(x, self.multiplier, self.amin, self.db_multiplier, self.top_db)

class MelScale(torch.nn.Module):
    r"""Turn a normal STFT into a mel frequency STFT, using a conversion
    matrix.  This uses triangular filter banks.

    User can control which device the filter bank (`fb`) is (e.g. fb.to(spec_f.device)).

    Args:
        n_mels (int, optional): Number of mel filterbanks. (Default: ``128``)
        sample_rate (int, optional): Sample rate of audio signal. (Default: ``16000``)
        f_min (float, optional): Minimum frequency. (Default: ``0.``)
        f_max (float or None, optional): Maximum frequency. (Default: ``sample_rate // 2``)
        n_stft (int, optional): Number of bins in STFT. Calculated from first input
            if None is given.  See ``n_fft`` in :class:`Spectrogram`. (Default: ``None``)
    """
    __constants__ = ['n_mels', 'sample_rate', 'f_min', 'f_max']

    def __init__(self, n_mels=128, sample_rate=16000, f_min=0., f_max=None, n_stft=None):
        super(MelScale, self).__init__()
        self.n_mels = n_mels
        self.sample_rate = sample_rate
        self.f_max = f_max if f_max is not None else float(sample_rate // 2)
        self.f_min = f_min

        assert f_min <= self.f_max, 'Require f_min: %f < f_max: %f' % (f_min, self.f_max)

        fb = torch.empty(0) if n_stft is None else F.create_fb_matrix(
            n_stft, self.f_min, self.f_max, self.n_mels, self.sample_rate)
        self.register_buffer('fb', fb)

    def forward(self, specgram):
        r"""
        Args:
            specgram (torch.Tensor): A spectrogram STFT of dimension (..., freq, time).

        Returns:
            torch.Tensor: Mel frequency spectrogram of size (..., ``n_mels``, time).
        """

        # pack batch
        shape = specgram.size()
        specgram = specgram.view(-1, shape[-2], shape[-1])

        if self.fb.numel() == 0:
            tmp_fb = F.create_fb_matrix(specgram.size(1), self.f_min, self.f_max, self.n_mels, self.sample_rate)
            # Attributes cannot be reassigned outside __init__ so workaround
            self.fb.resize_(tmp_fb.size())
            self.fb.copy_(tmp_fb)

        # (channel, frequency, time).transpose(...) dot (frequency, n_mels)
        # -> (channel, time, n_mels).transpose(...)
        mel_specgram = torch.matmul(specgram.transpose(1, 2), self.fb).transpose(1, 2)

        # unpack batch
        mel_specgram = mel_specgram.view(shape[:-2] + mel_specgram.shape[-2:])

        return mel_specgram

class InverseMelScale(torch.nn.Module):
    r"""Solve for a normal STFT from a mel frequency STFT, using a conversion
    matrix.  This uses triangular filter banks.

    It minimizes the euclidian norm between the input mel-spectrogram and the product between
    the estimated spectrogram and the filter banks using SGD.

    Args:
        n_stft (int): Number of bins in STFT. See ``n_fft`` in :class:`Spectrogram`.
        n_mels (int): Number of mel filterbanks. (Default: ``128``)
        sample_rate (int): Sample rate of audio signal. (Default: ``16000``)
        f_min (float): Minimum frequency. (Default: ``0.``)
        f_max (float, optional): Maximum frequency. (Default: ``sample_rate // 2``)
        max_iter (int): Maximum number of optimization iterations.
        tolerance_loss (float): Value of loss to stop optimization at.
        tolerance_change (float): Difference in losses to stop optimization at.
        sgdargs (dict): Arguments for the SGD optimizer.
    """
    __constants__ = ['n_stft', 'n_mels', 'sample_rate', 'f_min', 'f_max', 'max_iter', 'tolerance_loss',
                     'tolerance_change', 'sgdargs']

    def __init__(self, n_stft, n_mels=128, sample_rate=16000, f_min=0., f_max=None, max_iter=100000,
                 tolerance_loss=1e-5, tolerance_change=1e-8, sgdargs=None):
        super(InverseMelScale, self).__init__()
        self.n_mels = n_mels
        self.sample_rate = sample_rate
        self.f_max = f_max or float(sample_rate // 2)
        self.f_min = f_min
        self.max_iter = max_iter
        self.tolerance_loss = tolerance_loss
        self.tolerance_change = tolerance_change
        self.sgdargs = sgdargs or {'lr': 0.1, 'momentum': 0.9}

        assert f_min <= self.f_max, 'Require f_min: %f < f_max: %f' % (f_min, self.f_max)

        fb = F.create_fb_matrix(n_stft, self.f_min, self.f_max, self.n_mels, self.sample_rate)
        self.register_buffer('fb', fb)

    def forward(self, melspec):
        r"""
        Args:
            melspec (torch.Tensor): A Mel frequency spectrogram of dimension (..., ``n_mels``, time)

        Returns:
            torch.Tensor: Linear scale spectrogram of size (..., freq, time)
        """
        # pack batch
        shape = melspec.size()
        melspec = melspec.view(-1, shape[-2], shape[-1])

        n_mels, time = shape[-2], shape[-1]
        freq, _ = self.fb.size()  # (freq, n_mels)
        melspec = melspec.transpose(-1, -2)
        assert self.n_mels == n_mels

        specgram = torch.rand(melspec.size()[0], time, freq, requires_grad=True,
                              dtype=melspec.dtype, device=melspec.device)

        optim = torch.optim.SGD([specgram], **self.sgdargs)

        loss = float('inf')
        for _ in range(self.max_iter):
            optim.zero_grad()
            diff = melspec - specgram.matmul(self.fb)
            new_loss = diff.pow(2).sum(axis=-1).mean()
            # take sum over mel-frequency then average over other dimensions
            # so that loss threshold is applied par unit timeframe
            new_loss.backward()
            optim.step()
            specgram.data = specgram.data.clamp(min=0)

            new_loss = new_loss.item()
            if new_loss < self.tolerance_loss or abs(loss - new_loss) < self.tolerance_change:
                break
            loss = new_loss

        specgram.requires_grad_(False)
        specgram = specgram.clamp(min=0).transpose(-1, -2)

        # unpack batch
        specgram = specgram.view(shape[:-2] + (freq, time))
        return specgram

class MelSpectrogram(torch.nn.Module):
    r"""Create MelSpectrogram for a raw audio signal. This is a composition of Spectrogram
    and MelScale.

    Sources
        * https://gist.github.com/kastnerkyle/179d6e9a88202ab0a2fe
        * https://timsainb.github.io/spectrograms-mfccs-and-inversion-in-python.html
        * http://haythamfayek.com/2016/04/21/speech-processing-for-machine-learning.html

    Args:
        sample_rate (int, optional): Sample rate of audio signal. (Default: ``16000``)
        win_length (int or None, optional): Window size. (Default: ``n_fft``)
        hop_length (int or None, optional): Length of hop between STFT windows. (Default: ``win_length // 2``)
        n_fft (int, optional): Size of FFT, creates ``n_fft // 2 + 1`` bins. (Default: ``400``)
        f_min (float, optional): Minimum frequency. (Default: ``0.``)
        f_max (float or None, optional): Maximum frequency. (Default: ``None``)
        pad (int, optional): Two sided padding of signal. (Default: ``0``)
        n_mels (int, optional): Number of mel filterbanks. (Default: ``128``)
        window_fn (Callable[[...], torch.Tensor], optional): A function to create a window tensor
            that is applied/multiplied to each frame/window. (Default: ``torch.hann_window``)
        wkwargs (Dict[..., ...] or None, optional): Arguments for window function. (Default: ``None``)

    Example
        >>> waveform, sample_rate = torchaudio.load('test.wav', normalization=True)
        >>> mel_specgram = transforms.MelSpectrogram(sample_rate)(waveform)  # (channel, n_mels, time)
    """
    __constants__ = ['sample_rate', 'n_fft', 'win_length', 'hop_length', 'pad', 'n_mels', 'f_min']

    def __init__(self, sample_rate=16000, n_fft=400, win_length=None, hop_length=None, f_min=0., f_max=None,
                 pad=0, n_mels=128, window_fn=torch.hann_window, wkwargs=None):
        super(MelSpectrogram, self).__init__()
        self.sample_rate = sample_rate
        self.n_fft = n_fft
        self.win_length = win_length if win_length is not None else n_fft
        self.hop_length = hop_length if hop_length is not None else self.win_length // 2
        self.pad = pad
        self.n_mels = n_mels  # number of mel frequency bins
        self.f_max = f_max
        self.f_min = f_min
        self.spectrogram = Spectrogram(n_fft=self.n_fft, win_length=self.win_length,
                                       hop_length=self.hop_length,
                                       pad=self.pad, window_fn=window_fn, power=2.,
                                       normalized=False, wkwargs=wkwargs)
        self.mel_scale = MelScale(self.n_mels, self.sample_rate, self.f_min, self.f_max, self.n_fft // 2 + 1)

    def forward(self, waveform):
        r"""
        Args:
            waveform (torch.Tensor): Tensor of audio of dimension (..., time).

        Returns:
            torch.Tensor: Mel frequency spectrogram of size (..., ``n_mels``, time).
        """
        specgram = self.spectrogram(waveform)
        mel_specgram = self.mel_scale(specgram)
        return mel_specgram

class MFCC(torch.nn.Module):
    r"""Create the Mel-frequency cepstrum coefficients from an audio signal.

    By default, this calculates the MFCC on the DB-scaled Mel spectrogram.
    This is not the textbook implementation, but is implemented here to
    give consistency with librosa.

    This output depends on the maximum value in the input spectrogram, and so
    may return different values for an audio clip split into snippets vs. a
    a full clip.

    Args:
        sample_rate (int, optional): Sample rate of audio signal. (Default: ``16000``)
        n_mfcc (int, optional): Number of mfc coefficients to retain. (Default: ``40``)
        dct_type (int, optional): type of DCT (discrete cosine transform) to use. (Default: ``2``)
        norm (str, optional): norm to use. (Default: ``'ortho'``)
        log_mels (bool, optional): whether to use log-mel spectrograms instead of db-scaled. (Default: ``False``)
        melkwargs (dict, optional): arguments for MelSpectrogram. (Default: ``None``)
    """
    __constants__ = ['sample_rate', 'n_mfcc', 'dct_type', 'top_db', 'log_mels']

    def __init__(self, sample_rate=16000, n_mfcc=40, dct_type=2, norm='ortho', log_mels=False,
                 melkwargs=None):
        super(MFCC, self).__init__()
        supported_dct_types = [2]
        if dct_type not in supported_dct_types:
            raise ValueError('DCT type not supported'.format(dct_type))
        self.sample_rate = sample_rate
        self.n_mfcc = n_mfcc
        self.dct_type = dct_type
        self.norm = norm
        self.top_db = 80.0
        self.amplitude_to_DB = AmplitudeToDB('power', self.top_db)

        if melkwargs is not None:
            self.MelSpectrogram = MelSpectrogram(sample_rate=self.sample_rate, **melkwargs)
        else:
            self.MelSpectrogram = MelSpectrogram(sample_rate=self.sample_rate)

        if self.n_mfcc > self.MelSpectrogram.n_mels:
            raise ValueError('Cannot select more MFCC coefficients than # mel bins')
        dct_mat = F.create_dct(self.n_mfcc, self.MelSpectrogram.n_mels, self.norm)
        self.register_buffer('dct_mat', dct_mat)
        self.log_mels = log_mels

    def forward(self, waveform):
        r"""
        Args:
            waveform (torch.Tensor): Tensor of audio of dimension (..., time).

        Returns:
            torch.Tensor: specgram_mel_db of size (..., ``n_mfcc``, time).
        """

        # pack batch
        shape = waveform.size()
        waveform = waveform.view(-1, shape[-1])

        mel_specgram = self.MelSpectrogram(waveform)
        if self.log_mels:
            log_offset = 1e-6
            mel_specgram = torch.log(mel_specgram + log_offset)
        else:
            mel_specgram = self.amplitude_to_DB(mel_specgram)
        # (channel, n_mels, time).tranpose(...) dot (n_mels, n_mfcc)
        # -> (channel, time, n_mfcc).tranpose(...)
        mfcc = torch.matmul(mel_specgram.transpose(1, 2), self.dct_mat).transpose(1, 2)

        # unpack batch
        mfcc = mfcc.view(shape[:-1] + mfcc.shape[-2:])

        return mfcc

class MuLawEncoding(torch.nn.Module):
    r"""Encode signal based on mu-law companding.  For more info see the
    `Wikipedia Entry <https://en.wikipedia.org/wiki/%CE%9C-law_algorithm>`_

    This algorithm assumes the signal has been scaled to between -1 and 1 and
    returns a signal encoded with values from 0 to quantization_channels - 1

    Args:
        quantization_channels (int, optional): Number of channels. (Default: ``256``)
    """
    __constants__ = ['quantization_channels']

    def __init__(self, quantization_channels=256):
        super(MuLawEncoding, self).__init__()
        self.quantization_channels = quantization_channels

    def forward(self, x):
        r"""
        Args:
            x (torch.Tensor): A signal to be encoded.

        Returns:
            x_mu (torch.Tensor): An encoded signal.
        """
        return F.mu_law_encoding(x, self.quantization_channels)

class MuLawDecoding(torch.nn.Module):
    r"""Decode mu-law encoded signal.  For more info see the
    `Wikipedia Entry <https://en.wikipedia.org/wiki/%CE%9C-law_algorithm>`_

    This expects an input with values between 0 and quantization_channels - 1
    and returns a signal scaled between -1 and 1.

    Args:
        quantization_channels (int, optional): Number of channels. (Default: ``256``)
    """
    __constants__ = ['quantization_channels']

    def __init__(self, quantization_channels=256):
        super(MuLawDecoding, self).__init__()
        self.quantization_channels = quantization_channels

    def forward(self, x_mu):
        r"""
        Args:
            x_mu (torch.Tensor): A mu-law encoded signal which needs to be decoded.

        Returns:
            torch.Tensor: The signal decoded.
        """
        return F.mu_law_decoding(x_mu, self.quantization_channels)

class Resample(torch.nn.Module):
    r"""Resample a signal from one frequency to another. A resampling method can be given.

    Args:
        orig_freq (float, optional): The original frequency of the signal. (Default: ``16000``)
        new_freq (float, optional): The desired frequency. (Default: ``16000``)
        resampling_method (str, optional): The resampling method. (Default: ``'sinc_interpolation'``)
    """

    def __init__(self, orig_freq=16000, new_freq=16000, resampling_method='sinc_interpolation'):
        super(Resample, self).__init__()
        self.orig_freq = orig_freq
        self.new_freq = new_freq
        self.resampling_method = resampling_method

    def forward(self, waveform):
        r"""
        Args:
            waveform (torch.Tensor): Tensor of audio of dimension (..., time).

        Returns:
            torch.Tensor: Output signal of dimension (..., time).
        """
        if self.resampling_method == 'sinc_interpolation':

            # pack batch
            shape = waveform.size()
            waveform = waveform.view(-1, shape[-1])

            waveform = kaldi.resample_waveform(waveform, self.orig_freq, self.new_freq)

            # unpack batch
            waveform = waveform.view(shape[:-1] + waveform.shape[-1:])

            return waveform

        raise ValueError('Invalid resampling method: %s' % (self.resampling_method))

class ComplexNorm(torch.nn.Module):
    r"""Compute the norm of complex tensor input.

    Args:
        power (float, optional): Power of the norm. (Default: to ``1.0``)
    """
    __constants__ = ['power']

    def __init__(self, power=1.0):
        super(ComplexNorm, self).__init__()
        self.power = power

    def forward(self, complex_tensor):
        r"""
        Args:
            complex_tensor (Tensor): Tensor shape of `(..., complex=2)`.

        Returns:
            Tensor: norm of the input tensor, shape of `(..., )`.
        """
        return F.complex_norm(complex_tensor, self.power)

class ComputeDeltas(torch.nn.Module):
    r"""Compute delta coefficients of a tensor, usually a spectrogram.

    See `torchaudio.functional.compute_deltas` for more details.

    Args:
        win_length (int): The window length used for computing delta. (Default: ``5``)
        mode (str): Mode parameter passed to padding. (Default: ``'replicate'``)
    """
    __constants__ = ['win_length']

    def __init__(self, win_length=5, mode="replicate"):
        super(ComputeDeltas, self).__init__()
        self.win_length = win_length
        self.mode = mode

    def forward(self, specgram):
        r"""
        Args:
            specgram (torch.Tensor): Tensor of audio of dimension (..., freq, time).

        Returns:
            deltas (torch.Tensor): Tensor of audio of dimension (..., freq, time).
        """
        return F.compute_deltas(specgram, win_length=self.win_length, mode=self.mode)

class TimeStretch(torch.nn.Module):
    r"""Stretch stft in time without modifying pitch for a given rate.

    Args:
        hop_length (int or None, optional): Length of hop between STFT windows. (Default: ``win_length // 2``)
        n_freq (int, optional): number of filter banks from stft. (Default: ``201``)
        fixed_rate (float or None, optional): rate to speed up or slow down by.
            If None is provided, rate must be passed to the forward method. (Default: ``None``)
    """
    __constants__ = ['fixed_rate']

    def __init__(self, hop_length=None, n_freq=201, fixed_rate=None):
        super(TimeStretch, self).__init__()

        self.fixed_rate = fixed_rate

        n_fft = (n_freq - 1) * 2
        hop_length = hop_length if hop_length is not None else n_fft // 2
        self.register_buffer('phase_advance', torch.linspace(0, math.pi * hop_length, n_freq)[..., None])

    def forward(self, complex_specgrams, overriding_rate=None):
        # type: (Tensor, Optional[float]) -> Tensor
        r"""
        Args:
            complex_specgrams (Tensor): complex spectrogram (..., freq, time, complex=2).
            overriding_rate (float or None, optional): speed up to apply to this batch.
                If no rate is passed, use ``self.fixed_rate``. (Default: ``None``)

        Returns:
            (Tensor): Stretched complex spectrogram of dimension (..., freq, ceil(time/rate), complex=2).
        """
        assert complex_specgrams.size(-1) == 2, "complex_specgrams should be a complex tensor, shape (..., complex=2)"

        if overriding_rate is None:
            rate = self.fixed_rate
            if rate is None:
                raise ValueError("If no fixed_rate is specified"
                                 ", must pass a valid rate to the forward method.")
        else:
            rate = overriding_rate

        if rate == 1.0:
            return complex_specgrams

        return F.phase_vocoder(complex_specgrams, rate, self.phase_advance)

class Fade(torch.nn.Module):
    r"""Add a fade in and/or fade out to an waveform.

    Args:
        fade_in_len (int, optional): Length of fade-in (time frames). (Default: ``0``)
        fade_out_len (int, optional): Length of fade-out (time frames). (Default: ``0``)
        fade_shape (str, optional): Shape of fade. Must be one of: "quarter_sine",
            "half_sine", "linear", "logarithmic", "exponential". (Default: ``"linear"``)
    """
    def __init__(self, fade_in_len=0, fade_out_len=0, fade_shape="linear"):
        super(Fade, self).__init__()
        self.fade_in_len = fade_in_len
        self.fade_out_len = fade_out_len
        self.fade_shape = fade_shape

    def forward(self, waveform):
        # type: (Tensor) -> Tensor
        r"""
        Args:
            waveform (torch.Tensor): Tensor of audio of dimension (..., time).

        Returns:
            torch.Tensor: Tensor of audio of dimension (..., time).
        """
        waveform_length = waveform.size()[-1]

        return self._fade_in(waveform_length) * self._fade_out(waveform_length) * waveform

    def _fade_in(self, waveform_length):
        # type: (int) -> Tensor
        fade = torch.linspace(0, 1, self.fade_in_len)
        ones = torch.ones(waveform_length - self.fade_in_len)

        if self.fade_shape == "linear":
            fade = fade

        if self.fade_shape == "exponential":
            fade = torch.pow(2, (fade - 1)) * fade

        if self.fade_shape == "logarithmic":
            fade = torch.log10(.1 + fade) + 1

        if self.fade_shape == "quarter_sine":
            fade = torch.sin(fade * math.pi / 2)

        if self.fade_shape == "half_sine":
            fade = torch.sin(fade * math.pi - math.pi / 2) / 2 + 0.5

        return torch.cat((fade, ones)).clamp_(0, 1)

    def _fade_out(self, waveform_length):
        # type: (int) -> Tensor
        fade = torch.linspace(0, 1, self.fade_out_len)
        ones = torch.ones(waveform_length - self.fade_out_len)

        if self.fade_shape == "linear":
            fade = - fade + 1

        if self.fade_shape == "exponential":
            fade = torch.pow(2, - fade) * (1 - fade)

        if self.fade_shape == "logarithmic":
            fade = torch.log10(1.1 - fade) + 1

        if self.fade_shape == "quarter_sine":
            fade = torch.sin(fade * math.pi / 2 + math.pi / 2)

        if self.fade_shape == "half_sine":
            fade = torch.sin(fade * math.pi + math.pi / 2) / 2 + 0.5

        return torch.cat((ones, fade)).clamp_(0, 1)

class _AxisMasking(torch.nn.Module):
    r"""Apply masking to a spectrogram.

    Args:
        mask_param (int): Maximum possible length of the mask.
        axis (int): What dimension the mask is applied on.
        iid_masks (bool): Applies iid masks to each of the examples in the batch dimension.
    """
    __constants__ = ['mask_param', 'axis', 'iid_masks']

    def __init__(self, mask_param, axis, iid_masks):

        super(_AxisMasking, self).__init__()
        self.mask_param = mask_param
        self.axis = axis
        self.iid_masks = iid_masks

    def forward(self, specgram, mask_value=0.):
        # type: (Tensor, float) -> Tensor
        r"""
        Args:
            specgram (torch.Tensor): Tensor of dimension (..., freq, time).
            mask_value (float): Value to assign to the masked columns.

        Returns:
            torch.Tensor: Masked spectrogram of dimensions (..., freq, time).
        """

        # if iid_masks flag marked and specgram has a batch dimension
        if self.iid_masks and specgram.dim() == 4:
            return F.mask_along_axis_iid(specgram, self.mask_param, mask_value, self.axis + 1)
        else:
            return F.mask_along_axis(specgram, self.mask_param, mask_value, self.axis)

class FrequencyMasking(_AxisMasking):
    r"""Apply masking to a spectrogram in the frequency domain.

    Args:
        freq_mask_param (int): maximum possible length of the mask.
            Indices uniformly sampled from [0, freq_mask_param).
        iid_masks (bool, optional): whether to apply the same mask to all
            the examples/channels in the batch. (Default: ``False``)
    """

    def __init__(self, freq_mask_param, iid_masks=False):
        super(FrequencyMasking, self).__init__(freq_mask_param, 1, iid_masks)

class TimeMasking(_AxisMasking):
    r"""Apply masking to a spectrogram in the time domain.

    Args:
        time_mask_param (int): maximum possible length of the mask.
            Indices uniformly sampled from [0, time_mask_param).
        iid_masks (bool, optional): whether to apply the same mask to all
            the examples/channels in the batch. (Default: ``False``)
    """

    def __init__(self, time_mask_param, iid_masks=False):
        super(TimeMasking, self).__init__(time_mask_param, 2, iid_masks)

class Vol(torch.nn.Module):
    r"""Add a volume to an waveform.

    Args:
        gain (float): Interpreted according to the given gain_type:
            If `gain_type’ = ‘amplitude’, `gain’ is a positive amplitude ratio.
            If `gain_type’ = ‘power’, `gain’ is a power (voltage squared).
            If `gain_type’ = ‘db’, `gain’ is in decibels.
        gain_type (str, optional): Type of gain. One of: ‘amplitude’, ‘power’, ‘db’ (Default: ``"amplitude"``)
    """

    def __init__(self, gain, gain_type='amplitude'):
        super(Vol, self).__init__()
        self.gain = gain
        self.gain_type = gain_type

        if gain_type in ['amplitude', 'power'] and gain < 0:
            raise ValueError("If gain_type = amplitude or power, gain must be positive.")

    def forward(self, waveform):
        # type: (Tensor) -> Tensor
        r"""
        Args:
            waveform (torch.Tensor): Tensor of audio of dimension (..., time).

        Returns:
            torch.Tensor: Tensor of audio of dimension (..., time).
        """
        if self.gain_type == "amplitude":
            waveform = waveform * self.gain

        if self.gain_type == "db":
            waveform = F.gain(waveform, self.gain)

        if self.gain_type == "power":
            waveform = F.gain(waveform, 10 * math.log10(self.gain))

        return torch.clamp(waveform, -1, 1)