try:
        parsed = urlparse(ioc)
        if parsed.netloc:
            parsed = parsed._replace(netloc=parsed.netloc.replace('.', '[.]'),
                                     scheme=parsed.scheme.replace('t', 'x'))
            return parsed.geturl()
    except ValueError:
        pass

"""Extract and optionally refang Indicators of Compromise (IOCs) from text.

All methods return iterator objects, not lists. If for some reason you need
a list, do e.g.: ``list(extract_iocs(my_data))``.

Otherwise, you can iterate over the objects (e.g. in a ``for`` loop) normally.
Each object yielded from the generators will by of type :class:`str`.
"""
import io
import sys
import itertools
import argparse
import binascii
import base64
import json

try:
    # python3
    from urllib.parse import urlparse, unquote
    unicode = str
except ImportError:
    from urlparse import urlparse
    from urllib import unquote

import ipaddress
import regex as re

# Reusable end punctuation regex.
END_PUNCTUATION = r"[\.\?>\"'\)!,}:;\u201d\u2019\uff1e\uff1c\]]*"

# Reusable regex for symbols commonly used to defang.
SEPARATOR_DEFANGS = r"[\(\)\[\]{}<>\\]"

# Split URLs on some characters that may be valid, but may also be garbage.
URL_SPLIT_STR = r"[>\"'\),};]"

# Get basic url format, including a few obfuscation techniques, main anchor is the uri scheme.
GENERIC_URL_RE = re.compile(r"""
        (
            [fhstu]\S\S?[px]s?

            # One of these delimiters/defangs.
            (?:
                :\/\/|
                :\\\\|
                \[:\]\/\/|
                :?__
            )

            # Any number of defang characters.
            (?:
                \x20|
                """ + SEPARATOR_DEFANGS + r"""
            )*

            # Domain/path characters.
            \w
            \S+?

            # CISCO ESA style defangs followed by domain/path characters.
            (?:\x20[\/\.][^\.\/\s]\S*?)*
        )
    """ + END_PUNCTUATION + r"""
        (?=\s|[^\x00-\x7F]|$)
    """, re.IGNORECASE | re.VERBOSE | re.UNICODE)

# Get some obfuscated urls, main anchor is brackets around the period.
BRACKET_URL_RE = re.compile(r"""
        \b
        (
            [\.\:\/\\\w\[\]\(\)-]+
            (?:
                \x20?
                [\(\[]
                \x20?
                \.
                \x20?
                [\]\)]
                \x20?
                \S*?
            )+
        )
    """ + END_PUNCTUATION + r"""
        (?=\s|[^\x00-\x7F]|$)
    """, re.VERBOSE | re.UNICODE)

# Get some obfuscated urls, main anchor is backslash before a period.
BACKSLASH_URL_RE = re.compile(r"""
        \b
        (
            [\:\/\\\w\[\]\(\)-]+
            (?:
                \x20?
                \\?\.
                \x20?
                \S*?
            )*?
            (?:
                \x20?
                \\\.
                \x20?
                \S*?
            )
            (?:
                \x20?
                \\?\.
                \x20?
                \S*?
            )*
        )
    """ + END_PUNCTUATION + r"""
        (?=\s|[^\x00-\x7F]|$)
    """, re.VERBOSE | re.UNICODE)

# Get hex-encoded urls.
HEXENCODED_URL_RE = re.compile(r"""
        (
            [46][86]
            (?:[57]4)?
            [57]4[57]0
            (?:[57]3)?
            3a2f2f
            (?:2[356def]|3[0-9adf]|[46][0-9a-f]|[57][0-9af])+
        )
        (?:[046]0|2[0-2489a-c]|3[bce]|[57][b-e]|[8-f][0-9a-f]|0a|0d|09|[
            \x5b-\x5d\x7b\x7d\x0a\x0d\x20
        ]|$)
    """, re.IGNORECASE | re.VERBOSE)

# Get urlencoded urls.
URLENCODED_URL_RE = re.compile(r"""
        (s?[hf]t?tps?%3A%2F%2F\w[\w%-]*?)(?:[^\w%-]|$)
    """, re.IGNORECASE | re.VERBOSE)

# Get base64-encoded urls.
B64ENCODED_URL_RE = re.compile(r"""
        (
            # Modified to ignore whitespace.
            (?:
                [\x2b\x2f-\x39A-Za-z]\s*[\x2b\x2f-\x39A-Za-z]\s*[\x31\x35\x39BFJNRVZdhlptx]\s*[Gm]\s*[Vd]\s*[FH]\s*[A]\s*\x36\s*L\s*y\s*[\x2b\x2f\x38-\x39]\s*|
                [\x2b\x2f-\x39A-Za-z]\s*[\x2b\x2f-\x39A-Za-z]\s*[\x31\x35\x39BFJNRVZdhlptx]\s*[Io]\s*[Vd]\s*[FH]\s*[R]\s*[Qw]\s*[O]\s*i\s*\x38\s*v\s*[\x2b\x2f-\x39A-Za-z]\s*|
                [\x2b\x2f-\x39A-Za-z]\s*[\x2b\x2f-\x39A-Za-z]\s*[\x31\x35\x39BFJNRVZdhlptx]\s*[Io]\s*[Vd]\s*[FH]\s*[R]\s*[Qw]\s*[Uc]\s*[z]\s*o\s*v\s*L\s*[\x2b\x2f-\x39w-z]\s*|
                [\x2b\x2f-\x39A-Za-z]\s*[\x30\x32EGUWkm]\s*[Z]\s*[\x30U]\s*[Uc]\s*[D]\s*o\s*v\s*L\s*[\x2b\x2f-\x39w-z]\s*|
                [\x2b\x2f-\x39A-Za-z]\s*[\x30\x32EGUWkm]\s*[h]\s*[\x30U]\s*[Vd]\s*[FH]\s*[A]\s*\x36\s*L\s*y\s*[\x2b\x2f\x38-\x39]\s*|
                [\x2b\x2f-\x39A-Za-z]\s*[\x30\x32EGUWkm]\s*[h]\s*[\x30U]\s*[Vd]\s*[FH]\s*[B]\s*[Tz]\s*[O]\s*i\s*\x38\s*v\s*[\x2b\x2f-\x39A-Za-z]\s*|
                [RZ]\s*[ln]\s*[R]\s*[Qw]\s*[O]\s*i\s*\x38\s*v\s*[\x2b\x2f-\x39A-Za-z]\s*|
                [Sa]\s*[FH]\s*[R]\s*[\x30U]\s*[Uc]\s*[D]\s*o\s*v\s*L\s*[\x2b\x2f-\x39w-z]\s*|
                [Sa]\s*[FH]\s*[R]\s*[\x30U]\s*[Uc]\s*[FH]\s*[M]\s*\x36\s*L\s*y\s*[\x2b\x2f\x38-\x39]\s*
            )
            # Up to 260 characters (pre-encoding, reasonable URL length).
            [A-Za-z0-9+/=\s]{1,357}
        )
        (?=[^A-Za-z0-9+/=\s]|$)
    """, re.VERBOSE)

# Get some valid obfuscated ip addresses.
IPV4_RE = re.compile(r"""
        (?:^|
            (?![^\d\.])
        )
        (?:
            (?:[1-9]?\d|1\d\d|2[0-4]\d|25[0-5])
            [\[\(\\]*?\.[\]\)]*?
        ){3}
        (?:[1-9]?\d|1\d\d|2[0-4]\d|25[0-5])
        (?:(?=[^\d\.])|$)
    """, re.VERBOSE)

# Experimental IPv6 regex, will not catch everything but should be sufficent for now.
IPV6_RE = re.compile(r"""
        \b(?:[a-f0-9]{1,4}:|:){2,7}(?:[a-f0-9]{1,4}|:)\b
    """, re.IGNORECASE | re.VERBOSE)

# Capture email addresses including common defangs.
EMAIL_RE = re.compile(r"""
        (
            [a-z0-9_.+-]+
            [\(\[{\x20]*
            (?:@|\Wat\W)
            [\)\]}\x20]*
            [a-z0-9-]+
            (?:
                (?:
                    (?:
                        \x20*
                        """ + SEPARATOR_DEFANGS + r"""
                        \x20*
                    )*
                    \.
                    (?:
                        \x20*
                        """ + SEPARATOR_DEFANGS + r"""
                        \x20*
                    )*
                    |
                    \W+dot\W+
                )
                [a-z0-9-]+?
            )+
        )
    """ + END_PUNCTUATION + r"""
        (?=\s|$)
    """, re.IGNORECASE | re.VERBOSE | re.UNICODE)

MD5_RE = re.compile(r"(?:[^a-fA-F\d]|\b)([a-fA-F\d]{32})(?:[^a-fA-F\d]|\b)")
SHA1_RE = re.compile(r"(?:[^a-fA-F\d]|\b)([a-fA-F\d]{40})(?:[^a-fA-F\d]|\b)")
SHA256_RE = re.compile(r"(?:[^a-fA-F\d]|\b)([a-fA-F\d]{64})(?:[^a-fA-F\d]|\b)")
SHA512_RE = re.compile(r"(?:[^a-fA-F\d]|\b)([a-fA-F\d]{128})(?:[^a-fA-F\d]|\b)")

# YARA regex.
YARA_PARSE_RE = re.compile(r"""
        (?:^|\s)
        (
            (?:
                \s*?import\s+?"[^\r\n]*?[\r\n]+|
                \s*?include\s+?"[^\r\n]*?[\r\n]+|
                \s*?//[^\r\n]*[\r\n]+|
                \s*?/\*.*?\*/\s*?
            )*
            (?:
                \s*?private\s+|
                \s*?global\s+
            )*
            rule\s*?
            \w+\s*?
            (?:
                :[\s\w]+
            )?
            \s+\{
            .*?
            condition\s*?:
            .*?
            \s*\}
        )
        (?:$|\s)
    """, re.MULTILINE | re.DOTALL | re.VERBOSE)

def extract_iocs(data, refang=False, strip=False):
    """Extract all IOCs.

    Results are returned as an itertools.chain iterable object which
    lazily provides the results of the other extract_* generators.

    :param data: Input text
    :param bool refang: Refang output?
    :param bool strip: Strip possible garbage from the end of URLs
    :rtype: :py:func:`itertools.chain`
    """
    return itertools.chain(
        extract_urls(data, refang=refang, strip=strip),
        extract_ips(data, refang=refang),
        extract_emails(data, refang=refang),
        extract_hashes(data),
        extract_yara_rules(data)
    )

def extract_urls(data, refang=False, strip=False):
    """Extract URLs.

    :param data: Input text
    :param bool refang: Refang output?
    :param bool strip: Strip possible garbage from the end of URLs
    :rtype: :py:func:`itertools.chain`
    """
    return itertools.chain(
        extract_unencoded_urls(data, refang=refang, strip=strip),
        extract_encoded_urls(data, refang=refang, strip=strip),
    )

def extract_unencoded_urls(data, refang=False, strip=False):
    """Extract only unencoded URLs.

    :param data: Input text
    :param bool refang: Refang output?
    :param bool strip: Strip possible garbage from the end of URLs
    :rtype: Iterator[:class:`str`]
    """
    unencoded_urls = itertools.chain(
        GENERIC_URL_RE.finditer(data),
        BRACKET_URL_RE.finditer(data),
        BACKSLASH_URL_RE.finditer(data),
    )
    for url in unencoded_urls:
        if refang:
            url = refang_url(url.group(1))
        else:
            url = url.group(1)

        if strip:
            url = re.split(URL_SPLIT_STR, url)[0]

        yield url

def extract_encoded_urls(data, refang=False, strip=False):
    """Extract only encoded URLs.

    :param data: Input text
    :param bool refang: Decode output?
    :param bool strip: Strip possible garbage from the end of URLs
    :rtype: Iterator[:class:`str`]
    """
    for url in HEXENCODED_URL_RE.finditer(data):
        if refang:
            yield binascii.unhexlify(url.group(1)).decode('utf-8')
        else:
            yield url.group(1)

    for url in URLENCODED_URL_RE.finditer(data):
        if refang:
            yield unquote(url.group(1))
        else:
            yield url.group(1)

    for url in B64ENCODED_URL_RE.finditer(data):
        # Strip whitespace.
        url = ''.join(url.group(1).split())

        # Truncate the string if it's not a multiple of 3 bytes long.
        # We don't care about the end of the string since it's probably garbage.
        if len(url) % 4:
            url = url[:-(len(url) % 4)]

        if refang:
            # Decode base64.
            url = base64.b64decode(url).decode('utf-8', 'replace')

            # Remove the first 1-2 bytes if we got back extra leading characters from the base64.
            # The only valid starts are "http" or "ftp", so look for h/f case insensitive.
            url = url[re.search('[hHfF]', url).start():]

            # Stop at the first whitespace or non-unicode character.
            url = url.split(u'\ufffd')[0].\
                      split()[0]

        if strip:
            url = re.split(URL_SPLIT_STR, url)[0]

        yield url

def extract_ips(data, refang=False):
    """Extract IP addresses.

    Includes both IPv4 and IPv6 addresses.

    :param data: Input text
    :param bool refang: Refang output?
    :rtype: :py:func:`itertools.chain`
    """
    return itertools.chain(
        extract_ipv4s(data, refang=refang),
        extract_ipv6s(data),
    )

def extract_ipv4s(data, refang=False):
    """Extract IPv4 addresses.

    :param data: Input text
    :param bool refang: Refang output?
    :rtype: Iterator[:class:`str`]
    """
    for ip_address in IPV4_RE.finditer(data):
        if refang:
            yield refang_ipv4(ip_address.group(0))
        else:
            yield ip_address.group(0)

def extract_ipv6s(data):
    """Extract IPv6 addresses.

    Not guaranteed to catch all valid IPv6 addresses.

    :param data: Input text
    :rtype: Iterator[:class:`str`]
    """
    for ip_address in IPV6_RE.finditer(data):
        yield ip_address.group(0)

def extract_emails(data, refang=False):
    """Extract email addresses.

    :param data: Input text
    :param bool refang: Refang output?
    :rtype: Iterator[:class:`str`]
    """
    for email in EMAIL_RE.finditer(data):
        if refang:
            email = refang_email(email.group(1))
        else:
            email = email.group(1)

        yield email

def extract_hashes(data):
    """Extract MD5/SHA hashes.

    Results are returned as an itertools.chain iterable object which
    lazily provides the results of the other extract_*_hashes generators.

    :param data: Input text
    :rtype: :py:func:`itertools.chain`
    """
    return itertools.chain(
        extract_md5_hashes(data),
        extract_sha1_hashes(data),
        extract_sha256_hashes(data),
        extract_sha512_hashes(data)
    )

def extract_md5_hashes(data):
    """Extract MD5 hashes.

    :param data: Input text
    :rtype: Iterator[:class:`str`]
    """
    for md5 in MD5_RE.finditer(data):
        yield md5.group(1)

def extract_sha1_hashes(data):
    """Extract SHA1 hashes.

    :param data: Input text
    :rtype: Iterator[:class:`str`]
    """
    for sha1 in SHA1_RE.finditer(data):
        yield sha1.group(1)

def extract_sha256_hashes(data):
    """Extract SHA256 hashes.

    :param data: Input text
    :rtype: Iterator[:class:`str`]
    """
    for sha256 in SHA256_RE.finditer(data):
        yield sha256.group(1)

def extract_sha512_hashes(data):
    """Extract SHA512 hashes.

    :param data: Input text
    :rtype: Iterator[:class:`str`]
    """
    for sha512 in SHA512_RE.finditer(data):
        yield sha512.group(1)

def extract_yara_rules(data):
    """Extract YARA rules.

    :param data: Input text
    :rtype: Iterator[:class:`str`]
    """
    for yara_rule in YARA_PARSE_RE.finditer(data):
        yield yara_rule.group(1).strip()

def extract_custom_iocs(data, regex_list):
    """Extract using custom regex strings.

    Will always yield only the first *group* match from each regex.

    Always use a single capture group! Do this::

        [
            r'(my regex)',  # This yields 'my regex' if the pattern matches.
            r'my (re)gex',  # This yields 're' if the pattern matches.
        ]

    NOT this::

        [
            r'my regex',  # BAD! This doesn't yield anything.
            r'(my) (re)gex',  # BAD! This yields 'my' if the pattern matches.
        ]

    For complicated regexes, you can combine capture and non-capture groups,
    like this::

        [
            r'(?:my|your) (re)gex',  # This yields 're' if the pattern matches.
        ]

    Note the (?: ) syntax for noncapture groups vs the ( ) syntax for the capture
    group.

    :param data: Input text
    :param regex_list: List of strings to treat as regex and match against data.
    :rtype: Iterator[:class:`str`]
    """
    # Compile all the regex strings first, so we can error out quickly.
    regex_objects = []
    for regex_string in regex_list:
        regex_objects.append(re.compile(regex_string))

    # Iterate over regex objects, running each against input data.
    for regex_object in regex_objects:
        for ioc in regex_object.finditer(data):
            yield ioc.group(1)

def _is_ipv6_url(url):
    """URL network location is an IPv6 address, not a domain.

    :param url: String URL
    :rtype: bool
    """
    # Fix urlparse exception.
    parsed = urlparse(url)

    # Handle RFC 2732 IPv6 URLs with and without port, as well as non-RFC IPv6 URLs.
    if ']:' in parsed.netloc:
        ipv6 = ':'.join(parsed.netloc.split(':')[:-1])
    else:
        ipv6 = parsed.netloc

    try:
        ipaddress.IPv6Address(unicode(ipv6.replace('[', '').replace(']', '')))
    except ValueError:
        return False

    return True

def _refang_common(ioc):
    """Remove artifacts from common defangs.

    :param ioc: String IP/Email Address or URL netloc.
    :rtype: str
    """
    return ioc.replace('[dot]', '.').\
               replace('(dot)', '.').\
               replace('[.]', '.').\
               replace('(', '').\
               replace(')', '').\
               replace(',', '.').\
               replace(' ', '').\
               replace(u'\u30fb', '.')

def refang_email(email):
    """Refang an email address.

    :param email: String email address.
    :rtype: str
    """
    email = re.sub('\W[aA][tT]\W', '@', email.lower())
    email = re.sub('\W*[dD][oO][tT]\W*', '.', email)

    # Then do other char replaces.
    return _refang_common(email).replace('[', '').\
                                 replace(']', '').\
                                 replace('{', '').\
                                 replace('}', '').\
                                 replace('{', '')

def refang_url(url):
    """Refang a URL.

    :param url: String URL
    :rtype: str
    """
    # First fix urlparse errors.
    # Fix ipv6 parsing exception.
    if '[.' in url and '[.]' not in url:
        url = url.replace('[.', '[.]')
    if '.]' in url and '[.]' not in url:
        url = url.replace('.]', '[.]')
    if '[dot' in url and '[dot]' not in url:
        url = url.replace('[dot', '[.]')
    if 'dot]' in url and '[dot]' not in url:
        url = url.replace('dot]', '[.]')
    if '[:]' in url:
        url = url.replace('[:]', ':')
    if '[/]' in url:
        url = url.replace('[/]', '/')

    # Since urlparse expects a scheme, make sure one exists.
    if '//' not in url:
        if '__' in url[:8]:
            # Support http__domain and http:__domain.
            if ':__' in url[:8]:
                url = url.replace(':__', '://', 1)
            else:
                url = url.replace('__', '://', 1)
        elif '\\\\' in url[:8]:
            # Support http:\\domain.
            url = url.replace('\\\\', '//', 1)
        else:
            # Support no-protocol.
            url = 'http://' + url

    url = url.replace('(/)', '/')

    # Refang some backslash-escaped characters.
    url = url.replace('\.', '.').\
              replace('\(', '(').\
              replace('\[', '[').\
              replace('\)', ')').\
              replace('\]', ']')

    try:
        _ = urlparse(url)
    except ValueError:
        # Last resort on ipv6 fail.
        url = url.replace('[', '').replace(']', '')

    # Now use urlparse and continue processing.
    parsed = urlparse(url)

    # Handle URLs with no scheme / obfuscated scheme.
    # Note: ParseResult._replace is a public member, this is safe.
    if parsed.scheme not in ['http', 'https', 'ftp']:
        if parsed.scheme.strip('s') in ['ftx', 'fxp']:
            parsed = parsed._replace(scheme='ftp')
            url = parsed.geturl().replace('ftp:///', 'ftp://')
        else:
            parsed = parsed._replace(scheme='http')
            url = parsed.geturl().replace('http:///', 'http://')

        try:
            _ = urlparse(url)
        except ValueError:
            # Last resort on ipv6 fail.
            url = url.replace('[', '').replace(']', '')

        parsed = urlparse(url)

    # Remove artifacts from common defangs.
    parsed = parsed._replace(netloc=_refang_common(parsed.netloc))
    parsed = parsed._replace(path=parsed.path.replace('[.]', '.'))

    # Fix example[.]com, but keep RFC 2732 URLs intact.
    if not _is_ipv6_url(url):
        parsed = parsed._replace(netloc=parsed.netloc.replace('[', '').replace(']', ''))

    return parsed.geturl()

def refang_ipv4(ip_address):
    """Refang an IPv4 address.

    :param ip_address: String IPv4 address.
    :rtype: str
    """
    return _refang_common(ip_address).replace('[', '').\
                                      replace(']', '').\
                                      replace('\\', '')

def defang(ioc):
    """Defang a URL, domain, or IPv4 address.

    :param ioc: String URL, domain, or IPv4 address.
    :rtype: str
    """
    # If it's a url, defang just the scheme and netloc.

    # If it's a domain or IP, defang up to the first slash.
    split_list = ioc.split('/')
    defanged = split_list[0].replace('.', '[.]')
    # Include everything after the first slash without modification.
    if len(split_list) > 1:
        defanged = '/'.join([defanged] + split_list[1:])

    return defanged

def main():
    """Run as a commandline utility."""
    parser = argparse.ArgumentParser(
        description="""Advanced Indicator of Compromise (IOC) extractor.
                       If no arguments are specified, the default behavior is
                       to extract all IOCs.""")
    parser.add_argument('--input', type=lambda x: io.open(x, 'r', encoding='utf-8', errors='ignore'),
                        default=io.open(0, 'r', encoding='utf-8', errors='ignore'), help="default: stdin")
    parser.add_argument('--output', type=lambda x: io.open(x, 'w', encoding='utf-8', errors='ignore'),
                        default=io.open(1, 'w', encoding='utf-8', errors='ignore'), help="default: stdout")
    parser.add_argument('--extract-emails', action='store_true')
    parser.add_argument('--extract-ips', action='store_true')
    parser.add_argument('--extract-ipv4s', action='store_true')
    parser.add_argument('--extract-ipv6s', action='store_true')
    parser.add_argument('--extract-urls', action='store_true')
    parser.add_argument('--extract-yara-rules', action='store_true')
    parser.add_argument('--extract-hashes', action='store_true')
    parser.add_argument('--custom-regex', type=lambda x: io.open(x, 'r', encoding='utf-8', errors='ignore'),
                        metavar='REGEX_FILE',
                        help="file with custom regex strings, one per line, with one capture group each")
    parser.add_argument('--refang', action='store_true', help="default: no")
    parser.add_argument('--strip-urls', action='store_true',
                        help="remove possible garbage from the end of urls. default: no")
    parser.add_argument('--wide', action='store_true',
                        help="preprocess input to allow wide-encoded character matches. default: no")
    parser.add_argument('--json', action='store_true')

    args = parser.parse_args()

    data = args.input.read()
    if args.wide:
        data = data.replace('\x00', '')

    # By default, extract all.
    extract_all = not (
        args.extract_ips or
        args.extract_urls or
        args.extract_yara_rules or
        args.extract_hashes or
        args.extract_ipv4s or
        args.extract_ipv6s or
        args.extract_emails or
        args.custom_regex
    )
    memo = {}

    if args.extract_emails or extract_all:
        memo["emails"] = list(extract_emails(data, refang=args.refang))
    if args.extract_ipv4s or args.extract_ips or extract_all:
        memo["ipv4s"] = list(extract_ipv4s(data, refang=args.refang))
    if args.extract_ipv6s or args.extract_ips or extract_all:
        memo["ipv6s"] = list(extract_ipv6s(data))
    if args.extract_urls or extract_all:
        memo["urls"] = list(extract_urls(data, refang=args.refang, strip=args.strip_urls))
    if args.extract_yara_rules or extract_all:
        memo["yara_rules"] = list(extract_yara_rules(data))
    if args.extract_hashes or extract_all:
        memo["hashes"] = list(extract_hashes(data))

    # Custom regex file, one per line.
    if args.custom_regex:
        regex_list = [l.strip() for l in args.custom_regex.readlines()]

        try:
            memo["custom_regex"] = list(extract_custom_iocs(data, regex_list))
        except (IndexError, re.error) as e:
            sys.stderr.write('Error in custom regex: {e}\n'.format(e=e))

    if args.json:
        ioc = json.dumps(memo, indent=4, sort_keys=True)
    else:
        ioc = "\n".join(sum(memo.values(), []))

    args.output.write(u"{ioc}\n".format(ioc=ioc))

if __name__ == "__main__":
    main()


import torch
import numpy as np
import os
import time
from sklearn.feature_extraction import image
from tqdm import tqdm
from glob import glob
from tifffile import imsave
from sklearn.cluster import MeanShift
from matplotlib import pyplot as plt
from IPython.display import clear_output

def normalize(img, mean, std):
    """Normalize an array of images with mean and standard deviation.
        Parameters
        ----------
        img: array
            An array of images.
        mean: float
            Mean of img array.
        std: float
            Standard deviation of img array.
        """
    return (img - mean)/std

def denormalize(img, mean, std):
    """Denormalize an array of images with mean and standard deviation.
    Parameters
    ----------
    img: array
        An array of images.
    mean: float
        Mean of img array.
    std: float
        Standard deviation of img array.
    """
    return (img * std) + mean

def convertToFloat32(train_images,val_images):
    """Converts the data to float 32 bit type.
    Parameters
    ----------
    train_images: array
        Training data.
    val_images: array
        Validation data.
        """
    x_train = train_images.astype('float32')
    x_val = val_images.astype('float32')
    return x_train, x_val

def getMeanStdData(train_images,val_images):
    """Compute mean and standrad deviation of data.
    Parameters
    ----------
    train_images: array
        Training data.
    val_images: array
        Validation data.
    """
    x_train_ = train_images.astype('float32')
    x_val_ = val_images.astype('float32')
    data = np.concatenate((x_train_,x_val_), axis=0)
    mean, std = np.mean(data), np.std(data)
    return mean, std

def convertNumpyToTensor(numpy_array):
    """Convert numpy array to PyTorch tensor.
    Parameters
    ----------
    numpy_array: numpy array
        Numpy array.
    """

def preprocess(train_patches,val_patches):
    data_mean, data_std = getMeanStdData(train_patches, val_patches)
    x_train, x_val = convertToFloat32(train_patches,val_patches)
    x_train_extra_axis = x_train[:,np.newaxis]
    x_val_extra_axis = x_val[:,np.newaxis]
    x_train_tensor = convertNumpyToTensor(x_train_extra_axis)
    x_val_tensor = convertNumpyToTensor(x_val_extra_axis)
    return x_train_tensor, x_val_tensor, data_mean, data_std

def get_trainval_patches(x,split_fraction=0.85,augment=True,patch_size=128,num_patches=None):
    np.random.shuffle(x)
    train_images = x[:int(0.85*x.shape[0])]
    val_images = x[int(0.85*x.shape[0]):]
    if(augment):
        train_images = augment_data(train_images)
    x_train_crops = extract_patches(train_images, patch_size, num_patches)
    x_val_crops = extract_patches(val_images, patch_size, num_patches)
    print("Shape of training patches:", x_train_crops.shape, "Shape of validation patches:", x_val_crops.shape)
    return x_train_crops, x_val_crops

def extract_patches(x,patch_size,num_patches):
    """Deterministically extract patches from array of images.
    Parameters
    ----------
    x: numpy array
        Array of images.
    patch_size: int
        Size of patches to be extracted from each image.
    num_patches: int
        Number of patches to be extracted from each image.
    """
    img_width = x.shape[2]
    img_height = x.shape[1]
    if(num_patches is None):
        num_patches = int(float(img_width*img_height)/float(patch_size**2)*2)
    patches = np.zeros(shape=(x.shape[0]*num_patches,patch_size,patch_size))

    for i in tqdm(range(x.shape[0])):
        patches[i*num_patches:(i+1)*num_patches] = image.extract_patches_2d(x[i],(patch_size,patch_size), num_patches,
                                                                           random_state=i)
    return patches

def augment_data(X_train):
    """Augment data by 8-fold with 90 degree rotations and flips.
    Parameters
    ----------
    X_train: numpy array
        Array of training images.
    """
    X_ = X_train.copy()
    X_train_aug = np.concatenate((X_train, np.rot90(X_, 1, (1, 2))))
    X_train_aug = np.concatenate((X_train_aug, np.rot90(X_, 2, (1, 2))))
    X_train_aug = np.concatenate((X_train_aug, np.rot90(X_, 3, (1, 2))))
    X_train_aug = np.concatenate((X_train_aug, np.flip(X_train_aug, axis=1)))
    return X_train_aug

def loadImages(path):
    """Load images from a given directory.
    Parameters
    ----------
    path: String
        Path of directory from where to load images from.
    """
    files = sorted(glob(path))
    data=[]
    print(path)
    for f in files:
        if '.png' in f:
            im_b = np.array(io.imread(f))
        if '.npy' in f:
            im_b = np.load(f)
        data.append(im_b)

    data = np.array(data).astype(np.float32)
    return data

def getSamples(vae, size=20,zSize=64, mu=None, logvar=None, samples=1, tq=False):
    """Generate synthetic samples from DivNoising network.
    Parameters
    ----------
    vae: VAE Object
        DivNoising model.
    size: int
        Size of generated image in the bottleneck.
    zSize: int
        Dimension of latent space for each pixel in bottleneck.
    mu: PyTorch tensor
        latent space mean tensor.
    logvar: PyTorch tensor
        latent space log variance tensor.
    samples: int
        Number of synthetic samples to generate.
    tq: boolean
        If tqdm should be active or not to indicate progress.
    """
    if mu is None:
        mu=torch.zeros(1,zSize,size,size).cuda()
    if logvar is None:
        logvar=torch.zeros(1,zSize,size,size).cuda()

    results=[]
    for i in tqdm(range(samples),disable= not tq):
        z = vae.reparameterize(mu, logvar)
        recon = vae.decode(z)
        recon_cpu = recon.cpu()
        recon_numpy = recon_cpu.detach().numpy()
        recon_numpy.shape=(recon_numpy.shape[-2],recon_numpy.shape[-1])
        results.append(recon_numpy)
    return np.array(results)

def interpolate(vae, z_start, z_end, steps, display, vmin=0,vmax=255,):
    results=[]
    for i in range(steps):
        alpha=(i/(steps-1.0))
        z=z_end*alpha + z_start*(1.0-alpha)
        recon = vae.decode(z)
        recon_cpu = recon.cpu()
        recon_numpy = recon_cpu.detach().numpy()
        recon_numpy.shape=(recon_numpy.shape[-2],recon_numpy.shape[-1])
        if display:
            clear_output(wait=True)
            plt.imshow(recon_numpy,vmin=vmin, vmax=vmax)
            plt.show()
        time.sleep(0.4)
        results.append(recon_numpy)
    return results

def tiledMode(im, ps, overlap, display=True, vmin=0,vmax=255,
              initBW=200, minBW=100, reduce=0.9):
    means=np.zeros(im.shape[1:])
    xmin=0
    ymin=0
    xmax=ps
    ymax=ps
    ovLeft=0
    while (xmin<im.shape[2]):
        ovTop=0
        while (ymin<im.shape[1]):

            inputPatch=im[:,ymin:ymax,xmin:xmax]
            a = findMode(inputPatch,
                        initBW, minBW, reduce)
            a=a[:a.shape[0], :a.shape[1]]
            means[ymin:ymax,xmin:xmax][ovTop:,ovLeft:] = a[ovTop:,ovLeft:]

            ymin=ymin-overlap+ps
            ymax=ymin+ps
            ovTop=overlap//2

        ymin=0
        ymax=ps
        xmin=xmin-overlap+ps
        xmax=xmin+ps
        ovLeft=overlap//2

        if display:
            plt.imshow(means,vmin=vmin, vmax=vmax)
            plt.show()
            clear_output(wait=True)

    return means

def findClosest(samples, q):
    """Find closest sample to a given sample.
    Parameters
    ----------
    samples: array
        Array of samples from which the closest image needs to be found.
    q: image(array)
        Image to which the closest image needs to be found.
    """
    dif=np.mean(np.mean((samples-q)**2, -1),-1)
    return samples[np.argmin(dif)]

def findMode(samples, initBW=200, minBW=100, reduce=0.9):
    """Find the modes of a distribution of images.
    Parameters
    ----------
    samples: array
        Array of samples from which the modes need to be found.
    initBW: int
        Initial bandwidth.
    minBW: int
        Minimum bandwidth.
    reduce: float
        Factor by which to reduce bandwith i n iterations.
    """
    imagesC=samples.copy()
    imagesC.shape=(samples.shape[0],samples.shape[1]*samples.shape[2])
    seed=np.mean(imagesC,axis=0)[np.newaxis,...]
    bw=initBW
    for i in range(15):

        clustering = MeanShift(bandwidth=bw, seeds=seed, cluster_all=True).fit(imagesC)
        centers=clustering.cluster_centers_.copy()
        seed=centers
        bw=bw*reduce
        if bw < minBW:
            break

    result=seed[0]
    result.shape=(samples.shape[1],samples.shape[2])
    return result

def plotProbabilityDistribution(signalBinIndex, histogram, gaussianMixtureNoiseModel, min_signal, max_signal, n_bin, device):
    """Plots probability distribution P(x|s) for a certain ground truth signal.
       Predictions from both Histogram and GMM-based Noise models are displayed for comparison.
        Parameters
        ----------
        signalBinIndex: int
            index of signal bin. Values go from 0 to number of bins (`n_bin`).
        histogram: numpy array
            A square numpy array of size `nbin` times `n_bin`.
        gaussianMixtureNoiseModel: GaussianMixtureNoiseModel
            Object containing trained parameters.
        min_signal: float
            Lowest pixel intensity present in the actual sample which needs to be denoised.
        max_signal: float
            Highest pixel intensity present in the actual sample which needs to be denoised.
        n_bin: int
            Number of Bins.
        device: GPU device
        """
    histBinSize=(max_signal-min_signal)/n_bin
    querySignal_numpy= (signalBinIndex/float(n_bin)*(max_signal-min_signal)+min_signal)
    querySignal_numpy +=histBinSize/2
    querySignal_torch = torch.from_numpy(np.array(querySignal_numpy)).float().to(device)

    queryObservations_numpy=np.arange(min_signal, max_signal, histBinSize)
    queryObservations_numpy+=histBinSize/2
    queryObservations = torch.from_numpy(queryObservations_numpy).float().to(device)
    pTorch=gaussianMixtureNoiseModel.likelihood(queryObservations, querySignal_torch)
    pNumpy=pTorch.cpu().detach().numpy()

    plt.figure(figsize=(12, 5))

    plt.subplot(1, 2, 1)
    plt.xlabel('Observation Bin')
    plt.ylabel('Signal Bin')
    plt.imshow(histogram**0.25, cmap='gray')
    plt.axhline(y=signalBinIndex+0.5, linewidth=5, color='blue', alpha=0.5)

    plt.subplot(1, 2, 2)
    plt.plot(queryObservations_numpy, histogram[signalBinIndex, :]/histBinSize, label='GT Hist: bin ='+str(signalBinIndex), color='blue', linewidth=2)
    plt.plot(queryObservations_numpy, pNumpy, label='GMM : '+' signal = '+str(np.round(querySignal_numpy,2)), color='red',linewidth=2)
    plt.xlabel('Observations (x) for signal s = ' + str(querySignal_numpy))
    plt.ylabel('Probability Density')
    plt.title("Probability Distribution P(x|s) at signal =" + str(querySignal_numpy))
    plt.legend()

def predict_mmse(vae, img, samples, device, returnSamples=False, tq=True):
    '''
    Predicts MMSE estimate.
    Parameters
    ----------
    vae: VAE object
        DivNoising model.
    img: array
        Image for which denoised MMSE estimate needs to be computed.
    samples: int
        Number of samples to average for computing MMSE estimate.
    returnSamples:
        Should the method also return the individual samples?
    tq:
        Should progress bar be shown.
    tta:
        Should test time augmentation be enabled.
    '''
    img_height,img_width=img.shape[0],img.shape[1]
    imgT=torch.Tensor(img.copy())
    image_sample = imgT.view(1,1,img_height,img_width).to(device)
    vae.num_samples = samples
    all_samples = np.array(vae(image_sample,tqdm_bar=tq))
    samples_array = all_samples[:,0,0,:,:]
    if returnSamples:
        return np.mean(samples_array,axis=0), samples_array
    else:
        return np.mean(samples_array,axis=0)

def normalize_minmse(x, target):
    """Affine rescaling of x, such that the mean squared error to target is minimal."""
    cov = np.cov(x.flatten(), target.flatten())
    alpha = cov[0, 1] / (cov[0, 0] + 1e-10)
    beta = target.mean() - alpha * x.mean()
    return alpha * x + beta

def PSNR(gt, img):
    '''
    Compute PSNR.
    Parameters
    ----------
    gt: array
        Ground truth image.
    img: array
        Predicted image.''
    '''
    mse = np.mean(np.square(gt - img))
    return 20 * np.log10(np.max(gt)-np.min(gt)) - 10 * np.log10(mse)

def tta_forward(x):
    """
    Augments x 8-fold: all 90 deg rotations plus lr flip of the four rotated versions.

    Parameters
    ----------
    x: data to augment

    Returns
    -------
    Stack of augmented x.
    """
    x_aug = [x, np.rot90(x, 1), np.rot90(x, 2), np.rot90(x, 3)]
    x_aug_flip = x_aug.copy()
    for x_ in x_aug:
        x_aug_flip.append(np.fliplr(x_))
    return x_aug_flip

def tta_backward(x_aug):
    """
    Inverts `tta_forward` and averages the 8 images.

    Parameters
    ----------
    x_aug: stack of 8-fold augmented images.

    Returns
    -------
    average of de-augmented x_aug.
    """
    x_deaug = [x_aug[0],
               np.rot90(x_aug[1], -1),
               np.rot90(x_aug[2], -2),
               np.rot90(x_aug[3], -3),
               np.fliplr(x_aug[4]),
               np.rot90(np.fliplr(x_aug[5]), -1),
               np.rot90(np.fliplr(x_aug[6]), -2),
               np.rot90(np.fliplr(x_aug[7]), -3)]
    return np.mean(x_deaug, 0)

def predict_and_save(img,vae,num_samples,device,
                     fraction_samples_to_export,export_mmse,export_results_path,tta):

    '''
    Predict denoised images and save results to disk.
    Parameters
    ----------
    img: array or list
        A stack of tif images.
    vae: DivNoising model
    num_samples: int
        Number of samples to generate and use for computing MMSE.
    device: cuda device or cpu
    fraction_samples_to_export: float between 0 (inclusive) and 1 (inclusive)
        Number of samples to save on disk for each noisy image.
    export_mmse: bool
        Should MMSE estimate also be exported?
    export_results_path: str
        path where all results will be exported.
    tta: bool
        Use test-time augmentation if set to True.

    '''
    mmse_results = []
    if isinstance(img,(list)):
        num_images = len(img)
    if isinstance(img,(np.ndarray)):
        num_images = img.shape[0]
    for i in range(num_images):
        print("Processing image:", i)
        if tta:
            aug_imgs = tta_forward(img[i])
            mmse_aug=[]
            for j in range(len(aug_imgs)):
                if(j==0):
                    mmse,samples = predict_mmse(vae, aug_imgs[j], num_samples, device=device, returnSamples=True)
                else:
                    mmse = predict_mmse(vae, aug_imgs[j], num_samples, device=device, returnSamples=False)
                mmse_aug.append(mmse)

            mmse_back_transformed = tta_backward(mmse_aug)
            mmse_results.append(mmse_back_transformed)
        else:
            mmse,samples = predict_mmse(vae, img[i], num_samples, device=device, returnSamples=True)
            mmse_results.append(mmse)
        if fraction_samples_to_export>0:
            subdir = export_results_path+"/"+str(i).zfill(3)+"/"
            if not os.path.exists(subdir):
                os.makedirs(subdir)
            imsave(subdir+"samples_for_image_"+str(i).zfill(3)+".tif",
                   np.array(samples[:int(num_samples*fraction_samples_to_export)]).astype("float32"))
    if (export_mmse):
        imsave(export_results_path+"/mmse_results.tif", np.array(mmse_results).astype("float32"))
    return mmse_results

def plot_qualitative_results(noisy_input,vae,device):
    '''
    Plot qualitative results on patches.
    Parameters
    ----------
    noisy_input: array or list
        A stack of tif images.
    '''

    for j in range(5):

        # we select a random crop
        size_uncropped=int(0.14*(np.minimum(noisy_input[0].shape[0],noisy_input[0].shape[1])))
        size = size_uncropped-(size_uncropped%(2**vae.n_depth))
        minx=np.random.randint(0,noisy_input[0].shape[0]-size)
        miny=np.random.randint(0,noisy_input[0].shape[1]-size)
        img=noisy_input[0][minx:minx+size,miny:miny+size]

        # generate samples and MMSE estimate
        imgMMSE, samps = predict_mmse(vae, img, samples=100, device=device, returnSamples=True)

        plt.figure(figsize=(20, 6.75))

        # We display the noisy input image
        ax=plt.subplot(1,6,1)
        ax.get_xaxis().set_visible(False)
        ax.get_yaxis().set_visible(False)
        plt.imshow(img,cmap='magma')
        plt.title('input')

        # We display the average of 100 predicted samples
        ax=plt.subplot(1,6,6)
        ax.get_xaxis().set_visible(False)
        ax.get_yaxis().set_visible(False)
        plt.imshow(imgMMSE,cmap='magma')
        plt.title('MMSE (100 samples)')

        #We also display the first 4 samples
        for i in range(4):
            ax=plt.subplot(1,6,i+2)
            ax.get_xaxis().set_visible(False)
            ax.get_yaxis().set_visible(False)
            plt.imshow(samps[i],cmap='magma')
            plt.title('prediction '+str(i+1))

        plt.show()


while (ymin<im.shape[1]):

            inputPatch=im[:,ymin:ymax,xmin:xmax]
            a = findMode(inputPatch,
                        initBW, minBW, reduce)
            a=a[:a.shape[0], :a.shape[1]]
            means[ymin:ymax,xmin:xmax][ovTop:,ovLeft:] = a[ovTop:,ovLeft:]

            ymin=ymin-overlap+ps
            ymax=ymin+ps
            ovTop=overlap//2

import torch
import numpy as np
import os
import time
from sklearn.feature_extraction import image
from tqdm import tqdm
from glob import glob
from tifffile import imsave
from sklearn.cluster import MeanShift
from matplotlib import pyplot as plt
from IPython.display import clear_output

def normalize(img, mean, std):
    """Normalize an array of images with mean and standard deviation.
        Parameters
        ----------
        img: array
            An array of images.
        mean: float
            Mean of img array.
        std: float
            Standard deviation of img array.
        """
    return (img - mean)/std

def denormalize(img, mean, std):
    """Denormalize an array of images with mean and standard deviation.
    Parameters
    ----------
    img: array
        An array of images.
    mean: float
        Mean of img array.
    std: float
        Standard deviation of img array.
    """
    return (img * std) + mean

def convertToFloat32(train_images,val_images):
    """Converts the data to float 32 bit type.
    Parameters
    ----------
    train_images: array
        Training data.
    val_images: array
        Validation data.
        """
    x_train = train_images.astype('float32')
    x_val = val_images.astype('float32')
    return x_train, x_val

def getMeanStdData(train_images,val_images):
    """Compute mean and standrad deviation of data.
    Parameters
    ----------
    train_images: array
        Training data.
    val_images: array
        Validation data.
    """
    x_train_ = train_images.astype('float32')
    x_val_ = val_images.astype('float32')
    data = np.concatenate((x_train_,x_val_), axis=0)
    mean, std = np.mean(data), np.std(data)
    return mean, std

def convertNumpyToTensor(numpy_array):
    """Convert numpy array to PyTorch tensor.
    Parameters
    ----------
    numpy_array: numpy array
        Numpy array.
    """
    return torch.from_numpy(numpy_array)

def preprocess(train_patches,val_patches):
    data_mean, data_std = getMeanStdData(train_patches, val_patches)
    x_train, x_val = convertToFloat32(train_patches,val_patches)
    x_train_extra_axis = x_train[:,np.newaxis]
    x_val_extra_axis = x_val[:,np.newaxis]
    x_train_tensor = convertNumpyToTensor(x_train_extra_axis)
    x_val_tensor = convertNumpyToTensor(x_val_extra_axis)
    return x_train_tensor, x_val_tensor, data_mean, data_std

def get_trainval_patches(x,split_fraction=0.85,augment=True,patch_size=128,num_patches=None):
    np.random.shuffle(x)
    train_images = x[:int(0.85*x.shape[0])]
    val_images = x[int(0.85*x.shape[0]):]
    if(augment):
        train_images = augment_data(train_images)
    x_train_crops = extract_patches(train_images, patch_size, num_patches)
    x_val_crops = extract_patches(val_images, patch_size, num_patches)
    print("Shape of training patches:", x_train_crops.shape, "Shape of validation patches:", x_val_crops.shape)
    return x_train_crops, x_val_crops

def extract_patches(x,patch_size,num_patches):
    """Deterministically extract patches from array of images.
    Parameters
    ----------
    x: numpy array
        Array of images.
    patch_size: int
        Size of patches to be extracted from each image.
    num_patches: int
        Number of patches to be extracted from each image.
    """
    img_width = x.shape[2]
    img_height = x.shape[1]
    if(num_patches is None):
        num_patches = int(float(img_width*img_height)/float(patch_size**2)*2)
    patches = np.zeros(shape=(x.shape[0]*num_patches,patch_size,patch_size))

    for i in tqdm(range(x.shape[0])):
        patches[i*num_patches:(i+1)*num_patches] = image.extract_patches_2d(x[i],(patch_size,patch_size), num_patches,
                                                                           random_state=i)
    return patches

def augment_data(X_train):
    """Augment data by 8-fold with 90 degree rotations and flips.
    Parameters
    ----------
    X_train: numpy array
        Array of training images.
    """
    X_ = X_train.copy()
    X_train_aug = np.concatenate((X_train, np.rot90(X_, 1, (1, 2))))
    X_train_aug = np.concatenate((X_train_aug, np.rot90(X_, 2, (1, 2))))
    X_train_aug = np.concatenate((X_train_aug, np.rot90(X_, 3, (1, 2))))
    X_train_aug = np.concatenate((X_train_aug, np.flip(X_train_aug, axis=1)))
    return X_train_aug

def loadImages(path):
    """Load images from a given directory.
    Parameters
    ----------
    path: String
        Path of directory from where to load images from.
    """
    files = sorted(glob(path))
    data=[]
    print(path)
    for f in files:
        if '.png' in f:
            im_b = np.array(io.imread(f))
        if '.npy' in f:
            im_b = np.load(f)
        data.append(im_b)

    data = np.array(data).astype(np.float32)
    return data

def getSamples(vae, size=20,zSize=64, mu=None, logvar=None, samples=1, tq=False):
    """Generate synthetic samples from DivNoising network.
    Parameters
    ----------
    vae: VAE Object
        DivNoising model.
    size: int
        Size of generated image in the bottleneck.
    zSize: int
        Dimension of latent space for each pixel in bottleneck.
    mu: PyTorch tensor
        latent space mean tensor.
    logvar: PyTorch tensor
        latent space log variance tensor.
    samples: int
        Number of synthetic samples to generate.
    tq: boolean
        If tqdm should be active or not to indicate progress.
    """
    if mu is None:
        mu=torch.zeros(1,zSize,size,size).cuda()
    if logvar is None:
        logvar=torch.zeros(1,zSize,size,size).cuda()

    results=[]
    for i in tqdm(range(samples),disable= not tq):
        z = vae.reparameterize(mu, logvar)
        recon = vae.decode(z)
        recon_cpu = recon.cpu()
        recon_numpy = recon_cpu.detach().numpy()
        recon_numpy.shape=(recon_numpy.shape[-2],recon_numpy.shape[-1])
        results.append(recon_numpy)
    return np.array(results)

def interpolate(vae, z_start, z_end, steps, display, vmin=0,vmax=255,):
    results=[]
    for i in range(steps):
        alpha=(i/(steps-1.0))
        z=z_end*alpha + z_start*(1.0-alpha)
        recon = vae.decode(z)
        recon_cpu = recon.cpu()
        recon_numpy = recon_cpu.detach().numpy()
        recon_numpy.shape=(recon_numpy.shape[-2],recon_numpy.shape[-1])
        if display:
            clear_output(wait=True)
            plt.imshow(recon_numpy,vmin=vmin, vmax=vmax)
            plt.show()
        time.sleep(0.4)
        results.append(recon_numpy)
    return results

def tiledMode(im, ps, overlap, display=True, vmin=0,vmax=255,
              initBW=200, minBW=100, reduce=0.9):
    means=np.zeros(im.shape[1:])
    xmin=0
    ymin=0
    xmax=ps
    ymax=ps
    ovLeft=0
    while (xmin<im.shape[2]):
        ovTop=0

        ymin=0
        ymax=ps
        xmin=xmin-overlap+ps
        xmax=xmin+ps
        ovLeft=overlap//2

        if display:
            plt.imshow(means,vmin=vmin, vmax=vmax)
            plt.show()
            clear_output(wait=True)

    return means

def findClosest(samples, q):
    """Find closest sample to a given sample.
    Parameters
    ----------
    samples: array
        Array of samples from which the closest image needs to be found.
    q: image(array)
        Image to which the closest image needs to be found.
    """
    dif=np.mean(np.mean((samples-q)**2, -1),-1)
    return samples[np.argmin(dif)]

def findMode(samples, initBW=200, minBW=100, reduce=0.9):
    """Find the modes of a distribution of images.
    Parameters
    ----------
    samples: array
        Array of samples from which the modes need to be found.
    initBW: int
        Initial bandwidth.
    minBW: int
        Minimum bandwidth.
    reduce: float
        Factor by which to reduce bandwith i n iterations.
    """
    imagesC=samples.copy()
    imagesC.shape=(samples.shape[0],samples.shape[1]*samples.shape[2])
    seed=np.mean(imagesC,axis=0)[np.newaxis,...]
    bw=initBW
    for i in range(15):

        clustering = MeanShift(bandwidth=bw, seeds=seed, cluster_all=True).fit(imagesC)
        centers=clustering.cluster_centers_.copy()
        seed=centers
        bw=bw*reduce
        if bw < minBW:
            break

    result=seed[0]
    result.shape=(samples.shape[1],samples.shape[2])
    return result

def plotProbabilityDistribution(signalBinIndex, histogram, gaussianMixtureNoiseModel, min_signal, max_signal, n_bin, device):
    """Plots probability distribution P(x|s) for a certain ground truth signal.
       Predictions from both Histogram and GMM-based Noise models are displayed for comparison.
        Parameters
        ----------
        signalBinIndex: int
            index of signal bin. Values go from 0 to number of bins (`n_bin`).
        histogram: numpy array
            A square numpy array of size `nbin` times `n_bin`.
        gaussianMixtureNoiseModel: GaussianMixtureNoiseModel
            Object containing trained parameters.
        min_signal: float
            Lowest pixel intensity present in the actual sample which needs to be denoised.
        max_signal: float
            Highest pixel intensity present in the actual sample which needs to be denoised.
        n_bin: int
            Number of Bins.
        device: GPU device
        """
    histBinSize=(max_signal-min_signal)/n_bin
    querySignal_numpy= (signalBinIndex/float(n_bin)*(max_signal-min_signal)+min_signal)
    querySignal_numpy +=histBinSize/2
    querySignal_torch = torch.from_numpy(np.array(querySignal_numpy)).float().to(device)

    queryObservations_numpy=np.arange(min_signal, max_signal, histBinSize)
    queryObservations_numpy+=histBinSize/2
    queryObservations = torch.from_numpy(queryObservations_numpy).float().to(device)
    pTorch=gaussianMixtureNoiseModel.likelihood(queryObservations, querySignal_torch)
    pNumpy=pTorch.cpu().detach().numpy()

    plt.figure(figsize=(12, 5))

    plt.subplot(1, 2, 1)
    plt.xlabel('Observation Bin')
    plt.ylabel('Signal Bin')
    plt.imshow(histogram**0.25, cmap='gray')
    plt.axhline(y=signalBinIndex+0.5, linewidth=5, color='blue', alpha=0.5)

    plt.subplot(1, 2, 2)
    plt.plot(queryObservations_numpy, histogram[signalBinIndex, :]/histBinSize, label='GT Hist: bin ='+str(signalBinIndex), color='blue', linewidth=2)
    plt.plot(queryObservations_numpy, pNumpy, label='GMM : '+' signal = '+str(np.round(querySignal_numpy,2)), color='red',linewidth=2)
    plt.xlabel('Observations (x) for signal s = ' + str(querySignal_numpy))
    plt.ylabel('Probability Density')
    plt.title("Probability Distribution P(x|s) at signal =" + str(querySignal_numpy))
    plt.legend()

def predict_mmse(vae, img, samples, device, returnSamples=False, tq=True):
    '''
    Predicts MMSE estimate.
    Parameters
    ----------
    vae: VAE object
        DivNoising model.
    img: array
        Image for which denoised MMSE estimate needs to be computed.
    samples: int
        Number of samples to average for computing MMSE estimate.
    returnSamples:
        Should the method also return the individual samples?
    tq:
        Should progress bar be shown.
    tta:
        Should test time augmentation be enabled.
    '''
    img_height,img_width=img.shape[0],img.shape[1]
    imgT=torch.Tensor(img.copy())
    image_sample = imgT.view(1,1,img_height,img_width).to(device)
    vae.num_samples = samples
    all_samples = np.array(vae(image_sample,tqdm_bar=tq))
    samples_array = all_samples[:,0,0,:,:]
    if returnSamples:
        return np.mean(samples_array,axis=0), samples_array
    else:
        return np.mean(samples_array,axis=0)

def normalize_minmse(x, target):
    """Affine rescaling of x, such that the mean squared error to target is minimal."""
    cov = np.cov(x.flatten(), target.flatten())
    alpha = cov[0, 1] / (cov[0, 0] + 1e-10)
    beta = target.mean() - alpha * x.mean()
    return alpha * x + beta

def PSNR(gt, img):
    '''
    Compute PSNR.
    Parameters
    ----------
    gt: array
        Ground truth image.
    img: array
        Predicted image.''
    '''
    mse = np.mean(np.square(gt - img))
    return 20 * np.log10(np.max(gt)-np.min(gt)) - 10 * np.log10(mse)

def tta_forward(x):
    """
    Augments x 8-fold: all 90 deg rotations plus lr flip of the four rotated versions.

    Parameters
    ----------
    x: data to augment

    Returns
    -------
    Stack of augmented x.
    """
    x_aug = [x, np.rot90(x, 1), np.rot90(x, 2), np.rot90(x, 3)]
    x_aug_flip = x_aug.copy()
    for x_ in x_aug:
        x_aug_flip.append(np.fliplr(x_))
    return x_aug_flip

def tta_backward(x_aug):
    """
    Inverts `tta_forward` and averages the 8 images.

    Parameters
    ----------
    x_aug: stack of 8-fold augmented images.

    Returns
    -------
    average of de-augmented x_aug.
    """
    x_deaug = [x_aug[0],
               np.rot90(x_aug[1], -1),
               np.rot90(x_aug[2], -2),
               np.rot90(x_aug[3], -3),
               np.fliplr(x_aug[4]),
               np.rot90(np.fliplr(x_aug[5]), -1),
               np.rot90(np.fliplr(x_aug[6]), -2),
               np.rot90(np.fliplr(x_aug[7]), -3)]
    return np.mean(x_deaug, 0)

def predict_and_save(img,vae,num_samples,device,
                     fraction_samples_to_export,export_mmse,export_results_path,tta):

    '''
    Predict denoised images and save results to disk.
    Parameters
    ----------
    img: array or list
        A stack of tif images.
    vae: DivNoising model
    num_samples: int
        Number of samples to generate and use for computing MMSE.
    device: cuda device or cpu
    fraction_samples_to_export: float between 0 (inclusive) and 1 (inclusive)
        Number of samples to save on disk for each noisy image.
    export_mmse: bool
        Should MMSE estimate also be exported?
    export_results_path: str
        path where all results will be exported.
    tta: bool
        Use test-time augmentation if set to True.

    '''
    mmse_results = []
    if isinstance(img,(list)):
        num_images = len(img)
    if isinstance(img,(np.ndarray)):
        num_images = img.shape[0]
    for i in range(num_images):
        print("Processing image:", i)
        if tta:
            aug_imgs = tta_forward(img[i])
            mmse_aug=[]
            for j in range(len(aug_imgs)):
                if(j==0):
                    mmse,samples = predict_mmse(vae, aug_imgs[j], num_samples, device=device, returnSamples=True)
                else:
                    mmse = predict_mmse(vae, aug_imgs[j], num_samples, device=device, returnSamples=False)
                mmse_aug.append(mmse)

            mmse_back_transformed = tta_backward(mmse_aug)
            mmse_results.append(mmse_back_transformed)
        else:
            mmse,samples = predict_mmse(vae, img[i], num_samples, device=device, returnSamples=True)
            mmse_results.append(mmse)
        if fraction_samples_to_export>0:
            subdir = export_results_path+"/"+str(i).zfill(3)+"/"
            if not os.path.exists(subdir):
                os.makedirs(subdir)
            imsave(subdir+"samples_for_image_"+str(i).zfill(3)+".tif",
                   np.array(samples[:int(num_samples*fraction_samples_to_export)]).astype("float32"))
    if (export_mmse):
        imsave(export_results_path+"/mmse_results.tif", np.array(mmse_results).astype("float32"))
    return mmse_results

def plot_qualitative_results(noisy_input,vae,device):
    '''
    Plot qualitative results on patches.
    Parameters
    ----------
    noisy_input: array or list
        A stack of tif images.
    '''

    for j in range(5):

        # we select a random crop
        size_uncropped=int(0.14*(np.minimum(noisy_input[0].shape[0],noisy_input[0].shape[1])))
        size = size_uncropped-(size_uncropped%(2**vae.n_depth))
        minx=np.random.randint(0,noisy_input[0].shape[0]-size)
        miny=np.random.randint(0,noisy_input[0].shape[1]-size)
        img=noisy_input[0][minx:minx+size,miny:miny+size]

        # generate samples and MMSE estimate
        imgMMSE, samps = predict_mmse(vae, img, samples=100, device=device, returnSamples=True)

        plt.figure(figsize=(20, 6.75))

        # We display the noisy input image
        ax=plt.subplot(1,6,1)
        ax.get_xaxis().set_visible(False)
        ax.get_yaxis().set_visible(False)
        plt.imshow(img,cmap='magma')
        plt.title('input')

        # We display the average of 100 predicted samples
        ax=plt.subplot(1,6,6)
        ax.get_xaxis().set_visible(False)
        ax.get_yaxis().set_visible(False)
        plt.imshow(imgMMSE,cmap='magma')
        plt.title('MMSE (100 samples)')

        #We also display the first 4 samples
        for i in range(4):
            ax=plt.subplot(1,6,i+2)
            ax.get_xaxis().set_visible(False)
            ax.get_yaxis().set_visible(False)
            plt.imshow(samps[i],cmap='magma')
            plt.title('prediction '+str(i+1))

        plt.show()