# coding: utf-8
"""
Name: cycloen_deepeningrate.py

Find cyclone center location and
calucurate cyclone deepinrate.

example:
python3 cyclone_deepinrate.py -x 137 -y 40 -d <directory> -t yy-mm-dd-hh

Author: Ryosuke Tomita
Date: 2021/08/13
"""
import argparse
from datetime import datetime, timedelta
import math
import os
from os.path import abspath, dirname, join
import re
import sys
from typing import Union
from netCDF4 import Dataset
import numpy as np

def parse_args() -> dict:
    """parse_args.
    User have to determine ambiguous cyclen center location by weather map and so on.

    Args:

    Returns:
        dict:
    """
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "-y", "--lat", help="Center of low pressure latitude", type=float)
    parser.add_argument(
        "-x", "--lon", help="Center of low pressure longitude", type=float)
    parser.add_argument(
        "-d", "--dir", help="set directory name. This is the initial time.", type=str)
    parser.add_argument(
        "-t", "--time", help="set starttime. Format is yy-mm-dd-hh.", type=str)
    p = parser.parse_args()
    args = {"lat": p.lat, "lon": p.lon, "dir": p.dir, "time": p.time}
    return args

def mk_prmsl_file_list(data_root_dir: str, start_date: datetime) -> list:
    """mk_prmsl_file_list.
    split data_root_dir's into prmslfile or the others. These lists are sorted by datetime.

    Args:
        data_root_dir (str): data_root_dir
        start_date (datetime): start_date

    Returns:
        list:
    """
    def _get_abs_path(par_dir: str, child: str) -> str:
        """get_abs_path.

        Args:
            par_dir:
            child:
        """
        abs_dir = abspath(join(par_dir, child))
        return abs_dir

    prmsl_all_list = sorted([_get_abs_path(data_root_dir, i)
                            for i in os.listdir(data_root_dir) if "prmsl" in i])
    date_list = [
        datetime.strptime(re.search('[0-9]{4}-[0-9]{2}-[0-9]{2}_[0-9]{2}', i).group(),"%Y-%m-%d_%H")
        for i in prmsl_all_list
    ]
    prmsl_list = [prmsl_all_list[i]
                  for i in range(len(prmsl_all_list)) if date_list[i] >= start_date]
    #print(prmsl_list)
    return prmsl_list

def read_loc_prmsl(ncfile: str) -> Union[
            np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray
            ]:
    """read_loc_prmsl.
    read ncfile and extract near Japan data.

    Args:
        file_name (str): file_name

    Returns:
        Union[
                np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray
                ]:
    """
    read_ncfile = Dataset(ncfile)
    prmsl_raw = read_ncfile.variables['prmsl'][0]
    lat = [
        i
        if 20 <= i <= 60
        else -99
        for i in np.array(read_ncfile.variables['latitude'])
    ]
    lon =[
        i if 110 <= i <= 180
        else -99
        for i in np.array(read_ncfile.variables['longitude'])
    ]
    prmsl_1d = [
        prmsl_raw[i][j]
        for i in range(len(lat))
        if lat[i] != -99
        for j in range(len(lon))
        if lon[j] != -99
    ]

    near_jp_lat = np.array([i for i in lat if i != -99])
    near_jp_lon = np.array([i for i in lon if i != -99])
    prmsl = np.array(prmsl_1d).reshape(len(near_jp_lat), len(near_jp_lon))
    #read_ncfile.close()
    return lat, lon, near_jp_lat, near_jp_lon, prmsl

class CenterInfo:
    """CenterInfo.
    store cycloen center information.
    """

    def __init__(self, lat_index: np.ndarray, lon_index: np.ndarray,
                 prmsl: np.ndarray, lat: np.ndarray, lon: np.ndarray, date: str):
        """__init__.
        set varues.

        Args:
            lat_index (np.ndarray): lat_index
            lon_index (np.ndarray): lon_index
            prmsl (np.ndarray): prmsl
            lat (np.ndarray): lat
            lon (np.ndarray): lon
            date (str): date
        """
        self.lat_index = lat_index
        self.lon_index = lon_index
        self.prmsl = prmsl
        self.lat = lat
        self.lon = lon
        date_jts = datetime.strptime(date, "%Y-%m-%d_%H") + timedelta(hours=9)
        self.date = date_jts.strftime('%Y-%m-%d-%H')
        self.epsilon = -999  # epsilon is None.
        self.u_wind = None
        self.v_wind = None
        self.abs_wind = None
        self.wind_direction = None

    def __str__(self):
        """__str__.
        print format settings.
        """
        lat_data = f'latitude={self.lat},Index={self.lat_index}'
        lon_data = f'longitude={self.lon},Index={self.lon_index}'
        prmsl_data = f'prmslMin={self.prmsl}'
        date_data = f'date={self.date}'
        u_wind = f'u_wind={self.u_wind}'
        v_wind = f'v_wind={self.v_wind}'
        return f'{date_data},{lat_data},{lon_data},{prmsl_data},{u_wind},{v_wind}'

    def read_wind(self, ncfile: str, lat: np.ndarray, lon: np.ndarray,
            near_jp_lat: np.ndarray, near_jp_lon: np.ndarray, param: str) -> np.ndarray:
        read_ncfile = Dataset(ncfile.replace("-prmsl_hPa", ""))
        var_all = read_ncfile.variables[param][0]
        var_1d = [
                var_all[i][j]
            for i in range(len(lat))
            if lat[i] != -99
            for j in range(len(lon))
            if lon[j] != -99
    ]
        var = np.array(var_1d).reshape(len(near_jp_lat), len(near_jp_lon))
        return var[self.lat_index][self.lon_index]

def find_initial_prmslmin(lat: np.ndarray, lon: np.ndarray,
                          center_lat_index: float, center_lon_index: float,
                          prmsl: np.ndarray, date: str):# -> CenterInfo:
    """find center of cyclone using user's input lat,lon."""
    # if y,x are not varid value, then convert to int and recursion function.
    flag = True
    prmsl_min = 99999
    prmsl_min_lat, prmsl_min_lon = np.where(lat == center_lat_index), np.where(lon == center_lon_index)
    if lat[prmsl_min_lat].size == 0:
        center_lat_index = round(center_lat_index)
        flag = False
    if lon[prmsl_min_lon].size == 0:
        center_lon_index = round(center_lon_index)
        flag = False
    if not flag:
        return find_initial_prmslmin(lat, lon, center_lat_index, center_lon_index, prmsl, date)

    lat_search_area = np.arange(-2, 2.1, 0.5)
    lon_search_area = np.arange(-2, 2.1, 0.5)
    for la in [np.where(lat == i) for i in (lat_search_area + center_lat_index) if 20 <= i <= 60]:
        for lo in [np.where(lon == i) for i in (lon_search_area + center_lon_index) if 110 <= i <= 180]:
            la_index = (list(map(int, la))[0])
            lo_index = (list(map(int, lo))[0])

            if prmsl[la_index][lo_index] < prmsl_min:
                min_data = CenterInfo(
                    la_index, lo_index, prmsl[la_index][lo_index],
                    lat[la_index], lon[lo_index], date
                )
                prmsl_min = prmsl[la_index][lo_index]
    print(min_data)
    return min_data

def find_next_prmslmin(lat:np.ndarray, lon:np.ndarray, prmsl: np.ndarray,
                       prev_center_info, date: str) -> CenterInfo:
    """Using previous cyclone center location. Find next cyclone center location."""
    prmsl_min = 99999
    prev_min_lat = prev_center_info.lat
    prev_min_lon = prev_center_info.lon
    #lat_search_area = np.arange(-6, 6.1, 0.5)
    #lon_search_area = np.arange(-9, 9.1, 0.5)
    #lat_search_area = np.arange(-3, 3.1, 0.5)
    #lon_search_area = np.arange(-5, 5.1, 0.5)
    lat_search_area = np.arange(-3, 3.1, 0.5)
    lon_search_area = np.arange(-1, 7.1, 0.5)

    for la in [np.where(lat == i) for i in lat_search_area + prev_min_lat if 20 <= i <= 60]:
        for lo in [np.where(lon == i) for i in lon_search_area + prev_min_lon if 110 <= i <= 180]:
            la_index = (list(map(int, la))[0])
            lo_index = (list(map(int, lo))[0])

            if prmsl[la_index][lo_index] < prmsl_min:
                min_data = CenterInfo(
                    la_index, lo_index, prmsl[la_index][lo_index],
                    lat[la_index], lon[lo_index], date
                    )
                prmsl_min = prmsl[la_index][lo_index]
    print(min_data)
    return min_data

def cal_deeping_rate(prmsl_tracks: list):
    """calcurate cycloen deeping rate."""
    for i in range(1, len(prmsl_tracks)-1, 1):
        prmsl6h_ago = prmsl_tracks[i-1].prmsl
        prmsl6h_later = prmsl_tracks[i+1].prmsl
        nowlat_rad = math.radians(prmsl_tracks[i].lat)
        rad_45 = math.radians(45)

        prmsl_tracks[i].epsilon = (
            (prmsl6h_ago - prmsl6h_later)/12) * (math.sin(rad_45)/math.sin(nowlat_rad))

def write_data(prmsl_tracks, outname: str):
    """write_data.

    Args:
        prmsl_tracks:
        outname (str): outname
    """
    columns = ["year-month-day_Hour", "prmsl center (hPa)","latitude", "longitude",
               "deepingRate", "u_wind", "v_wind", "abs_wind", "direction of wind(rad)"]
    file_path = join(abspath(dirname(__file__)) + "/" + outname + ".csv")
    with open(file_path, mode='w') as f:
        f.write("{}\n".format(",".join(columns)))
        for p in prmsl_tracks:
            datas = list(map(str, [p.date, p.prmsl, p.lat, p.lon, p.epsilon, p.u_wind, p.v_wind, p.abs_wind, p.wind_direction]))
            f.write("{}\n".format(",".join(datas)))

def main():
    """main
    1. Set initial cyclone center info, initialtime,
        directory path from stdin.
    2. Make file list using directory path.
    3. Get lan lot data and extract near japan area.
    4. Find cyclone center location data and store in CenterInfo.
    5. Go to next files and find next cyclone center location data.
    6. Save to CSV file.
    """
    args = parse_args()

    data_root_dir = join(abspath(args["dir"]))
    data_root_dir_has_datetime = re.search('[0-9]{10}', data_root_dir)
    if data_root_dir_has_datetime:
        outname = data_root_dir_has_datetime.group()
    else:
        outname = "initialdata"
    start_date = datetime.strptime(args["time"], "%Y-%m-%d_%H")
    ncfiles = mk_prmsl_file_list(data_root_dir, start_date)

    # initial cyclone center data.
    prmsl_tracks = []
    lat, lon, near_jp_lat, near_jp_lon, prmsl = read_loc_prmsl(ncfiles[0])
    pressure_center = find_initial_prmslmin(
        near_jp_lat, near_jp_lon, args["lat"], args["lon"], prmsl, args["time"])
    pressure_center.u_wind = pressure_center.read_wind(
        ncfiles[0], lat, lon, near_jp_lat, near_jp_lon,
        'u10',
    )
    pressure_center.v_wind = pressure_center.read_wind(
        ncfiles[0], lat, lon, near_jp_lat, near_jp_lon,
        'v10',
    )
    prmsl_tracks.append(pressure_center)

    # next cyclone center data.
    for i in range(1, len(ncfiles), 1):
        lat, lon, near_jp_lat, near_jp_lon, prmsl = read_loc_prmsl(ncfiles[i])

        ncfiles_has_datetime = re.search('[0-9]{4}-[0-9]{2}-[0-9]{2}_[0-9]{2}', ncfiles[i])
        if ncfiles_has_datetime:
            date = ncfiles_has_datetime.group()
        else:
            raise Exception("ncfiles doesn't have datetime.")

        pressure_center = find_next_prmslmin(
                near_jp_lat, near_jp_lon, prmsl, pressure_center, date)
        pressure_center.u_wind = pressure_center.read_wind(
                ncfiles[i], lat, lon, near_jp_lat, near_jp_lon,
                'u10'
        )
        pressure_center.v_wind = pressure_center.read_wind(
                ncfiles[i], lat, lon, near_jp_lat, near_jp_lon,
                'v10'
        )
        pressure_center.abs_wind = (pressure_center.u_wind ** 2
                + pressure_center.v_wind ** 2) ** 0.5
        pressure_center.wind_direction = np.arctan2(pressure_center.u_wind, pressure_center.v_wind)

        prmsl_tracks.append(pressure_center)
        if pressure_center.lon == 180.0:
            break

    cal_deeping_rate(prmsl_tracks)
    write_data(prmsl_tracks, outname)

if __name__ == "__main__":
    main()


# pylint: disable=too-many-locals,too-many-lines
import os
import io
import math
import json
import shlex
import queue
import base64
import typing
import hashlib
import warnings
import logging
import fractions
import threading
import functools
import itertools
import subprocess
from urllib import request
from http import client

import numpy as np
import validators
import cv2

try:
    import PIL
    import PIL.Image
except ImportError:  # pragma: no cover
    PIL = None

LOGGER = logging.getLogger(__name__)

ImageInputType = typing.Union[str, np.ndarray, 'PIL.Image.Image', io.BytesIO]

SIZES = {'float32': 32, 'uint8': 8, 'bool': 1}

# Map codec names to the CUDA-accelerated version. Obtain
# from ffmpeg -codecs after building using CUDA.
CUDA_CODECS = {
    "h264": "h264_cuvid",
    "hevc": "hevc_cuvid",
    "mjpeg": "mjpeg_cuvid",
    "mpeg1video": "mpeg1_cuvid",
    "mpeg2video": "mpeg2_cuvid",
    "mpeg4": "mpeg4_cuvid",
    "vc1": "vc1_cuvid",
    "vp8": "vp8_cuvid",
    "vp9": "vp9_cuvid",
}

FramesWithIndexesAndTimestamps = typing.Generator[
    typing.Tuple[np.ndarray, typing.Optional[int], typing.
                 Optional[float]], None, None]

def get_ffprobe():
    return os.environ.get("PERCEPTION_FFPROBE_BINARY", "ffprobe")

def get_ffmpeg():
    return os.environ.get("PERCEPTION_FFMPEG_BINARY", "ffmpeg")

# pylint: disable=invalid-name
def compute_quality(image) -> int:
    """Compute a quality metric, using the calculation proposed by
    `Facebook <https://github.com/facebook/ThreatExchange/blob/master/hashing/hashing.pdf/>`_
    for their PDQ hash algorithm."""
    if len(image.shape) == 3:
        image = cv2.cvtColor(image, code=cv2.COLOR_RGB2GRAY)
    if image.shape[0] != 64 or image.shape[1] != 64:
        image = cv2.resize(src=image, dsize=(64, 64)).astype('float32')
    dx = 100 * np.abs(image[:, 1:] - image[:, :-1]) / 255
    dy = 100 * np.abs(image[1:] - image[:-1]) / 255
    dx = dx.astype('int').sum()
    dy = dy.astype('int').sum()
    return int(np.clip(a=int((dx + dy) / 90), a_min=0, a_max=100))

def compute_md5(filepath) -> str:
    """Compute the md5 hash for a file at `filepath`.

    Args:
        filepath: The path to the file
    """
    with open(filepath, 'rb') as f:  # pylint: disable=invalid-name
        hash_str = hashlib.md5(f.read()).hexdigest()
    return hash_str

def get_string_length(hash_length: int, dtype: str, hash_format='hex') -> int:
    """Compute the expected length of a hash string.

    Args:
        hash_length: The length of the hash vector
        dtype: The dtype of the vector
        hash_format: One of 'base64' or 'hex'

    Returns:
        The expected string length
    """
    hash_bytes = math.ceil(hash_length * SIZES[dtype] / 8)

    if hash_format == 'base64':
        return int((4 * hash_bytes / 3) + 3) & ~3
    if hash_format == 'hex':
        return 2 * hash_bytes
    raise NotImplementedError('Unknown hash format: ' + hash_format)

def vector_to_string(vector: np.ndarray, dtype: str,
                     hash_format: str) -> typing.Optional[str]:
    """Convert vector to hash.

    Args:
        vector: Input vector
    """
    # At times, a vector returned by a hasher is None (e.g., for hashes
    # that depend on the image not being featureless). In those cases,
    # we need to just return None, which is the least surprising outcome
    # because after all, the string representation of None is None.
    if vector is None:
        return None
    if hash_format == 'vector':
        # return vector.astype(dtype)  # old behavior
        raise DeprecationWarning("`hash_format` `vector` has been removed.")
    if dtype == 'uint8':
        vector_bytes = vector.astype('uint8')
    elif dtype == 'float32':
        vector_bytes = vector.astype('float32')
    elif dtype == 'bool':
        vector_bytes = np.packbits(vector.astype('bool'))
    else:
        raise NotImplementedError(f'Cannot convert hash of type {dtype}.')
    if hash_format == 'base64':
        return base64.b64encode(vector_bytes.tobytes()).decode('utf-8')
    if hash_format == 'hex':
        return vector_bytes.tobytes().hex()
    raise NotImplementedError(
        f'Cannot convert to string format: {hash_format}.')

def string_to_vector(hash_string: str,
                     dtype: str,
                     hash_length: int,
                     hash_format: str,
                     verify_length: bool = True) -> np.ndarray:
    """Convert hash back to vector.

    Args:
        hash_string: The input hash string
        dtype: The data type of the hash
        hash_length: The length of the hash vector
        hash_format: The input format of the hash (base64 or hex)
        verify_length: Whether to verify the string length
    """
    assert not verify_length or len(hash_string) == get_string_length(
        hash_length=hash_length, hash_format=hash_format,
        dtype=dtype), 'Incorrect string length for this hash format.'
    if hash_format == 'base64':
        vector_bytes = np.frombuffer(
            base64.b64decode(hash_string),
            dtype='uint8' if dtype in ['bool', 'uint8'] else dtype)
    elif hash_format == 'hex':
        vector_bytes = np.frombuffer(
            bytearray.fromhex(hash_string),
            dtype='uint8' if dtype in ['bool', 'uint8'] else dtype)
    else:
        raise NotImplementedError(
            f'Cannot convert to string format: {hash_format}')
    if dtype == 'uint8':
        return vector_bytes[:hash_length]
    if dtype == 'float32':
        return vector_bytes[:hash_length]
    if dtype == 'bool':
        return np.unpackbits(vector_bytes)[:hash_length].astype('bool')
    raise NotImplementedError(f'Cannot convert hash of type {dtype}.')

def hex_to_b64(hash_string: str,
               dtype: str,
               hash_length: int,
               verify_length: bool = True):
    """Convert a hex-encoded hash to base64.

    Args:
        hash_string: The input base64 hash string
        dtype: The data type of the hash
        hash_length: The length of the hash vector
        verify_length: Whether to verify the string length
    """
    return vector_to_string(
        string_to_vector(
            hash_string,
            hash_length=hash_length,
            hash_format='hex',
            dtype=dtype,
            verify_length=verify_length),
        dtype=dtype,
        hash_format='base64')

def b64_to_hex(hash_string: str,
               dtype: str,
               hash_length: int,
               verify_length: bool = True):
    """Convert a base64-encoded hash to hex.

    Args:
        hash_string: The input hex hash string
        dtype: The data type of the hash
        hash_length: The length of the hash vector
        verify_length: Whether to verify the string length
    """
    return vector_to_string(
        string_to_vector(
            hash_string,
            hash_length=hash_length,
            hash_format='base64',
            dtype=dtype,
            verify_length=verify_length),
        dtype=dtype,
        hash_format='hex')

def to_image_array(image: ImageInputType, require_color=True):
    if isinstance(image, np.ndarray):
        assert image.flags['C_CONTIGUOUS'], (
            'Provided arrays must be contiguous to avoid '
            'erroneous results when arrays are passed to '
            'underlying libraries. This can be achieved using'
            'np.ascontiguousarray(image)')
        assert not require_color or (len(image.shape) == 3
                                     and image.shape[-1] == 3), (
                                         'Provided images must be RGB images.')
        return image
    return read(image)

def get_common_framerates(id_rates: dict):
    """Compute an optimal set of framerates for a list
    of framerates. Optimal here means that reading the video
    at each of the framerates will allow one to collect all
    of the frames required with the smallest possible number of
    frames decoded.

    For example, consider if we need to read a video at
    3 fps, 5 fps, 1 fps and 0.5 fps. We could read the video
    4 times (once per framerate). But a more optimal approach
    is to read the video only twice, once at 3 frames per second
    and another time at 5 frames per second. For the 1 fps hasher,
    we simply pass every 3rd frame of the 3 fps pass. For the
    0.5 fps hasher, we pass every 6th frame of the 3 fps pass. So
    if you pass this function {A: 3, B: 5, C: 1, D: 0.5}, you will
    get back {3: [A, C, D], 5: C}.

    Args:
        id_rates: A dictionary with IDs as keys and frame rates as values.

    Returns:
        rate_ids: A dictionary with framerates as keys and a list of
            ids as values.
    """

    def partition(collection):
        """This function taken from
        https://stackoverflow.com/questions/19368375/set-partitions-in-python/30134039#30134039
        """
        if len(collection) == 1:
            yield [collection]
            return

        first = collection[0]
        for smaller in partition(collection[1:]):
            # insert `first` in each of the subpartition's subsets
            for n, subset in enumerate(smaller):
                yield smaller[:n] + [[first] + subset] + smaller[n + 1:]
            # put `first` in its own subset
            yield [[first]] + smaller

    framerates = list(id_rates.values())
    factor = 2 * 3 * 5 * 7 * 11 * 60 * 60
    assert min(framerates
               ) >= 1 / factor, 'Framerates must be at least 1 frame per hour.'
    best_frame_count = np.inf
    best_grouping: typing.Optional[typing.List] = None
    best_frame_rates: typing.Optional[typing.List] = None

    # We try every possible grouping of framerates to minimize the number
    # of frames we decode. There is likely a better way to do this,
    # but this seems to do the job for now.
    for grouping in partition(list(set(framerates))):
        current_frame_rates = [
            # pylint: disable=no-member
            functools.reduce(np.lcm,
                             (np.array(group) * factor).round().astype(int)) /
            factor for group in grouping
        ]
        current_frame_count = sum(current_frame_rates)
        if current_frame_count < best_frame_count:
            best_frame_count = current_frame_count
            best_frame_rates = current_frame_rates
            best_grouping = grouping

    assert best_frame_rates is not None
    assert best_grouping is not None
    return {
        framerate:
        tuple(name for name, rate in id_rates.items() if rate in group)
        for framerate, group in zip(best_frame_rates, best_grouping)
    }

def get_isometric_transforms(image: ImageInputType, require_color=True):
    image = to_image_array(image, require_color=require_color)
    return dict(
        r0=image,
        fv=np.ascontiguousarray(image[::-1, :]),
        fh=np.ascontiguousarray(image[:, ::-1]),
        r180=np.ascontiguousarray(image[::-1, ::-1]),
        r90=np.ascontiguousarray(image.transpose(1, 0, 2)[::-1, :, :]),
        r90fv=np.ascontiguousarray(image.transpose(1, 0, 2)),
        r90fh=np.ascontiguousarray(image.transpose(1, 0, 2)[::-1, ::-1]),
        r270=np.ascontiguousarray(image.transpose(1, 0, 2)[:, ::-1]))

def get_isometric_dct_transforms(dct: np.ndarray):
    # pylint: disable=invalid-name
    T1 = np.empty_like(dct)
    T1[::2] = 1
    T1[1::2] = -1

    # pylint: disable=invalid-name
    T2 = np.empty_like(dct)
    T2[::2, ::2] = 1
    T2[1::2, 1::2] = 1
    T2[::2, 1::2] = -1
    T2[1::2, ::2] = -1
    return dict(
        r0=dct,
        fv=dct * T1,
        fh=dct * T1.T,
        r180=dct * T2,
        r90=dct.T * T1,
        r90fv=dct.T,
        r90fh=dct.T * T2,
        r270=dct.T * T1.T)

def read(filepath_or_buffer: ImageInputType, timeout=None):
    """Read a file into an image object

    Args:
        filepath_or_buffer: The path to the file or any object
            with a `read` method (such as `io.BytesIO`)
        timeout: If filepath_or_buffer is a URL, the timeout to
            use for making the HTTP request.
    """
    if PIL is not None and isinstance(filepath_or_buffer, PIL.Image.Image):
        return np.array(filepath_or_buffer.convert("RGB"))
    if isinstance(filepath_or_buffer, (io.BytesIO, client.HTTPResponse)):
        image = np.asarray(
            bytearray(filepath_or_buffer.read()), dtype=np.uint8)
        image = cv2.imdecode(image, cv2.IMREAD_UNCHANGED)
    elif isinstance(filepath_or_buffer, str):
        if validators.url(filepath_or_buffer):
            return read(request.urlopen(filepath_or_buffer, timeout=timeout))
        if not os.path.isfile(filepath_or_buffer):
            raise FileNotFoundError('Could not find image at path: ' +
                                    filepath_or_buffer)
        image = cv2.imread(filepath_or_buffer)
    else:
        raise RuntimeError("Unhandled filepath_or_buffer type: " +
                           str(type(filepath_or_buffer)))
    if image is None:
        raise ValueError(f'An error occurred reading {filepath_or_buffer}.')
    # We use cvtColor here instead of just ret[..., ::-1]
    # in order to ensure that we provide a contiguous
    # array for later processing. Some hashers use ctypes
    # to pass the array and non-contiguous arrays can lead
    # to erroneous results.
    return cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

def _get_keyframes(filepath):
    """Get the keyframes for a video.

    Args:
        filepath: Path to the target file

    Returns:
        A list of frame indexes.
    """
    args = [
        get_ffprobe(), '-select_streams', 'v', '-i', f"'{filepath}'",
        '-print_format', 'json', '-show_entries',
        'frame=pict_type,coded_picture_number'
    ]
    with subprocess.Popen(
            args, stdout=subprocess.PIPE, stderr=subprocess.PIPE) as p:
        out, err = p.communicate()
        if p.returncode != 0:
            raise ValueError("{out}: {err}".format(out=str(out), err=str(err)))
        data = json.loads(out.decode('utf-8'))['frames']
        frames = [
            f['coded_picture_number'] for f in data if f['pict_type'] == 'I'
        ]
        frames = list(set(frames))
        frames.sort()
    return frames

def get_video_properties(filepath):
    cmd = f"""
    {get_ffprobe()} -select_streams v:0 -i '{filepath}'
    -print_format json -show_entries stream=width,height,avg_frame_rate,codec_name,start_time
    """
    with subprocess.Popen(
            shlex.split(cmd), stdout=subprocess.PIPE,
            stderr=subprocess.PIPE) as p:
        out, err = p.communicate()
        if p.returncode != 0:
            raise ValueError("{out}: {err}".format(out=str(out), err=str(err)))
        data = json.loads(out.decode("utf-8"))["streams"][0]
        numerator, denominator = tuple(
            map(int, data["avg_frame_rate"].split("/")[:2]))
        avg_frame_rate: typing.Optional[fractions.Fraction]
        if numerator > 0 and denominator > 0:
            avg_frame_rate = fractions.Fraction(
                numerator=numerator, denominator=denominator)
        else:
            avg_frame_rate = None
        return data["width"], data["height"], avg_frame_rate, data[
            "codec_name"], float(data.get("start_time", "0"))

# pylint: disable=too-many-branches,too-many-statements,too-many-arguments
def read_video_to_generator_ffmpeg(
        filepath,
        frames_per_second: typing.Optional[typing.Union[str, float]] = None,
        errors="raise",
        max_duration: float = None,
        max_size: int = None,
        interp: str = None,
        frame_rounding: str = "up",
        draw_timestamps=False,
        use_cuda=False) -> FramesWithIndexesAndTimestamps:
    """This is used by :code:`read_video` when :code:`use_ffmpeg` is True. It
    differs from :code:`read_video_to_generator` in that it uses FFMPEG instead of
    OpenCV and, optionally, allows for CUDA acceleration. CUDA acceleration
    can be faster for larger videos (>1080p) where downsampling is desired.
    For other videos, CUDA may be slower, but the decoding load will still be
    taken off the CPU, which may still be advantageous. You can specify which
    FFMPEG binary to use by setting PERCEPTION_FFMPEG_BINARY.

    Args:
        filepath: See read_video
        frames_per_second: See read_video
        errors: See read_video
        max_duration: See read_video
        max_size: See read_video
        interp: The interpolation method to use. When not using CUDA, you must choose one
            of the `interpolation options <https://ffmpeg.org/ffmpeg-scaler.html#sws_005fflags>`_
            (default: area). When using CUDA, you must choose from the
            `interp_algo options <http://underpop.online.fr/f/ffmpeg/help/scale_005fnpp.htm.gz>`_
            (default: super).
        frame_rounding: The frame rounding method.
        draw_timestamps: Draw original timestamps onto the frames (for debugging only)
        use_cuda: Whether to enable CUDA acceleration. Requires a
            CUDA-accelerated version of ffmpeg.

    To build FFMPEG with CUDA, do the following in a Docker
    container based on nvidia/cuda:10.1-cudnn7-devel-ubuntu18.04. The
    FFMPEG binary will be ffmpeg/ffmpeg.

    .. code-block:: bash

        git clone https://git.videolan.org/git/ffmpeg/nv-codec-headers.git
        cd nv-codec-headers
        make
        sudo make install
        cd ..
        git clone --branch release/4.3 https://git.ffmpeg.org/ffmpeg.git
        cd ffmpeg
        sudo apt-get update && sudo apt-get -y install yasm
        export PATH=$PATH:/usr/local/cuda/bin
        ./configure --enable-cuda-nvcc --enable-cuvid --enable-nvenc --enable-nvdec \
                    --enable-libnpp --enable-nonfree --extra-cflags=-I/usr/local/cuda/include \
                    --extra-ldflags=-L/usr/local/cuda/lib64
        make -j 10

    Returns:
        See :code:`read_video`
    """
    if interp is None:
        interp = "super" if use_cuda else "area"
    try:
        raw_width, raw_height, avg_frame_rate, codec_name, start_time = get_video_properties(
            filepath)
        start_time_offset = 0.0 if avg_frame_rate is None else float(
            (1 / (2 * avg_frame_rate)))
        LOGGER.debug(
            "raw_width: %s, raw_height: %s, avg_frame_rate: %s, codec_name: %s, start_time: %s",
            raw_width, raw_height, avg_frame_rate, codec_name, start_time)
        channels = 3
        scale = (min(max_size / raw_width, max_size / raw_height, 1)
                 if max_size is not None else 1)
        width, height = map(lambda d: int(round(scale * d)),
                            [raw_width, raw_height])
        # If there is no average frame rate, the offset tends to be unreliable.
        offset = max(start_time,
                     start_time_offset) if avg_frame_rate is not None else 0
        cmd = (f"{get_ffmpeg()} -hide_banner -an -vsync 0 -loglevel fatal "
               f"-itsoffset -{offset}")
        filters = []
        if draw_timestamps:
            pattern = "%{pts}-%{frame_num}"
            filters.append(f"drawtext=fontsize={int(raw_height * 0.1)}:"
                           f"fontcolor=yellow:text={pattern}"
                           ":x=(w-text_w):y=(h-text_h)")
        # Add frame rate filters.
        if frames_per_second is None:
            seconds_per_frame = float(
                1 / avg_frame_rate) if avg_frame_rate is not None else None
        elif frames_per_second == "keyframes":
            seconds_per_frame = None
            filters.append(r"select=eq(pict_type\,I)")
        else:
            assert isinstance(
                frames_per_second,
                (float, int)), f"Invalid framerate: {frames_per_second}"
            seconds_per_frame = 1 / frames_per_second
            filters.append(f"fps={frames_per_second}:"
                           f"round={frame_rounding}:"
                           f"start_time={offset}")
        # Add resizing filters.
        if use_cuda and codec_name in CUDA_CODECS:
            cuda_codec = CUDA_CODECS[codec_name]
            cmd += f" -hwaccel cuda -c:v {cuda_codec}"
            filters.append("hwupload_cuda")
            if scale != 1:
                filters.append(
                    f"scale_npp={width}:{height}:interp_algo={interp}")
            filters.extend([
                "hwdownload",
                "format=nv12",
            ])
        elif scale != 1:
            filters.append(f"scale={width}:{height}:flags={interp}")
        cmd += f" -i '{filepath}'"
        if filters:
            cmd += " -vf '{fstring}'".format(fstring=",".join(filters))
        cmd += " -pix_fmt rgb24 -f image2pipe -vcodec rawvideo -"
        LOGGER.debug("running ffmpeg with: %s", cmd)
        framebytes = width * height * channels
        bufsize = framebytes * int(
            os.environ.get("PERCEPTION_FFMPEG_BUFSIZE", "5"))
        with subprocess.Popen(
                shlex.split(cmd),
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                bufsize=bufsize) as p:
            assert p.stdout is not None, "Could not launch subprocess pipe."
            timestamp: typing.Optional[float] = 0
            frame_index: typing.Optional[int] = 0
            while True:
                batch = p.stdout.read(bufsize)
                if not batch:
                    break
                for image in np.frombuffer(
                        batch, dtype="uint8").reshape((-1, height, width,
                                                       channels)):
                    if frames_per_second != "keyframes":
                        yield (image, frame_index, timestamp)
                        if seconds_per_frame is not None:
                            assert timestamp is not None
                            timestamp += seconds_per_frame
                            frame_index = math.ceil(
                                avg_frame_rate * timestamp
                            ) if avg_frame_rate is not None else None
                        else:
                            timestamp = None
                            frame_index = None
                    else:
                        # Obtaining the keyframe indexes with ffprobe is very slow (slower
                        # than reading the video sometimes). We don't *have* to do it
                        # when using ffmpeg, so we don't. The OpenCV approach *does*
                        # get the keyframe indexes, but only because they're required
                        # in order to select them.
                        yield (image, None, None)
                    if (max_duration is not None and timestamp is not None
                            and timestamp > max_duration):
                        break
            stdout, stderr = p.communicate()
            if p.returncode != 0:
                raise ValueError(
                    f"Error parsing video: {stdout.decode('utf-8')} {stderr.decode('utf-8')}"
                )
    # pylint: disable=broad-except
    except Exception as e:
        if errors not in ["warn", "ignore"]:
            raise e
        if errors == "warn":
            warnings.warn(
                message=
                f"An error occurred while reading {filepath}. Processing may be truncated."
            )

# pylint: disable=too-many-branches,too-many-locals,too-many-statements
def read_video_to_generator(
        filepath,
        frames_per_second: typing.Optional[typing.Union[str, float]] = None,
        errors='raise',
        max_duration: float = None,
        max_size: int = None) -> FramesWithIndexesAndTimestamps:
    """This is used by :code:`read_video` when :code:`use_ffmpeg` is False (default).

    Args:
        filepath: See :code:`read_video`.
        frames_per_second: See :code:`read_video`.
        errors: See :code:`read_video`.
        max_duration: See :code:`read_video`.
        max_size: See :code:`read_video`.

    Returns:
        See :code:`read_video`.
    """
    # pylint: disable=no-member
    if cv2.__version__ < '4.1.1' and filepath.lower().endswith('gif'):
        message = 'Versions of OpenCV < 4.1.1 may read GIF files improperly. Upgrade recommended.'
        if errors == 'raise':
            raise ValueError(message)
        warnings.warn(message=message)

    if not os.path.isfile(filepath):
        raise FileNotFoundError(f'Could not find {filepath}.')
    cap = cv2.VideoCapture(filename=filepath, apiPreference=cv2.CAP_FFMPEG)
    try:
        # The purpose of the following block is largely to create a
        # frame_indexes (iterator or list) that indicates which
        # frames we should be returning to the user and then
        # yielding those frames as we come across them.
        file_frames_per_second = cap.get(cv2.CAP_PROP_FPS)
        if file_frames_per_second == 0:
            if errors == "raise":
                raise ValueError("Video file has framerate of 0fps.")
            # The known case where this occurs is for GIFs, where
            # 0 fps is typically inferred as 10 fps.
            file_frames_per_second = 10
            if errors == "warn":
                warnings.warn(
                    message=
                    "Video file has framerate of 0 fps. Guessing framerate of 10fps."
                )
        if frames_per_second is None:
            frames_per_second = file_frames_per_second
        seconds_between_desired_frames = None if (
            frames_per_second is not None
            and isinstance(frames_per_second,
                           str)) else 1 / frames_per_second  # type: ignore
        seconds_between_grabbed_frames = 1 / file_frames_per_second
        grabbed_frame_count = 0
        if frames_per_second == 'keyframes':
            frame_indexes: typing.Union[range, typing.List[int], typing.
                                        Iterator[int]] = _get_keyframes(
                                            filepath)
            # The repeat flag is used to handle the case where the
            # desired sampling rate is higher than the file's frame
            # rate. In this case, we will need to repeat frames in
            # order to provide the least-surprising behavior that
            # we can.
            repeat = False
        else:
            frame_indexes = itertools.count(
                0, max(1, file_frames_per_second / frames_per_second))
            repeat = file_frames_per_second < frames_per_second
        input_width = cap.get(cv2.CAP_PROP_FRAME_WIDTH)
        input_height = cap.get(cv2.CAP_PROP_FRAME_HEIGHT)
        if max_size is not None:
            scale = min(max_size / max(input_width, input_height), 1)
        else:
            scale = 1
        target_size: typing.Optional[typing.Tuple[int, int]]
        if scale < 1:
            target_size = (int(scale * input_width), int(scale * input_height))
        else:
            target_size = None
        for frame_index in frame_indexes:
            while grabbed_frame_count < frame_index:
                # We need to skip this frame.
                success = cap.grab()
                if not success:
                    break
                grabbed_frame_count += 1
            success, frame = cap.read()
            grabbed_frame_count += 1
            if not success:
                # The video is over or an error has occurred.
                break
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            if target_size is not None:
                frame = cv2.resize(
                    frame, target_size, interpolation=cv2.INTER_NEAREST)
            current_timestamp = frame_index / file_frames_per_second
            yield frame, grabbed_frame_count - 1, current_timestamp
            if max_duration is not None and current_timestamp > max_duration:
                break
            if repeat:
                next_desired_timestamp = current_timestamp + seconds_between_desired_frames
                next_timestamp = current_timestamp + seconds_between_grabbed_frames
                while next_desired_timestamp < next_timestamp:
                    yield (frame, grabbed_frame_count - 1,
                           next_desired_timestamp)
                    next_desired_timestamp += seconds_between_desired_frames
    # pylint: disable=broad-except
    except Exception as e:
        if errors not in ['warn', 'ignore']:
            raise e
        if errors == 'warn':
            warnings.warn(
                message=
                f'An error occurred while reading {filepath}. Processing may be truncated.'
            )
    finally:
        cap.release()

def read_video_into_queue(*args, video_queue, terminate, func, **kwargs):
    # We're inside a thread now and the queue is being read elsewhere.
    try:
        for frame, frame_index, timestamp in func(*args, **kwargs):
            if not terminate.isSet():
                video_queue.put((frame, frame_index, timestamp))
            else:
                break
    finally:
        video_queue.put((None, None, None))

# pylint: disable=too-many-arguments
def read_video(
        filepath,
        frames_per_second: typing.Optional[typing.Union[str, float]] = None,
        max_queue_size=128,
        use_queue=True,
        errors='raise',
        use_ffmpeg=False,
        **kwargs) -> FramesWithIndexesAndTimestamps:
    """Provides a generator of RGB frames, frame indexes, and timestamps from a
    video. This function requires you to have installed ffmpeg. All other
    arguments passed to read_video_to_generator.

    Args:
        filepath: Path to the video file
        frames_per_second: How many frames to provide for
            each second of video. If None, all frames
            are provided. If frames_per_second is "keyframes",
            we use ffmpeg to select I frames from the video.
        max_queue_size: The maximum number of frames to load in the queue
        use_queue: Whether to use a queue of frames during processing
        max_duration: The maximum length of the video to hash.
        max_size: The maximum size of frames to queue
        errors: Whether to 'raise', 'warn', or 'ignore' errors
        use_ffmpeg: Whether to use the FFMPEG CLI to read videos. If True, other
            kwargs (e.g., :code:`use_cuda`) are passed to
            :code:`read_video_to_generator_ffmpeg`.

    Yields:
        (frame, frame_index, timestamp) tuples
    """
    for ffmpeg_kwarg in [
            "interp", "frame_rounding", "draw_timestamps", "use_cuda"
    ]:
        if not use_ffmpeg and ffmpeg_kwarg in kwargs:
            warnings.warn(
                f"{ffmpeg_kwarg} is ignored when use_ffmpeg is False.",
                UserWarning)
            del kwargs[ffmpeg_kwarg]
    generator: typing.Callable[..., FramesWithIndexesAndTimestamps]
    if use_ffmpeg:
        generator = read_video_to_generator_ffmpeg
    else:
        generator = read_video_to_generator
    frame_index: typing.Optional[int]
    timestamp: typing.Optional[float]
    if use_queue:
        video_queue = queue.Queue(
            maxsize=max_queue_size
        )  # type: queue.Queue[typing.Tuple[np.ndarray, int, float]]
        terminate = threading.Event()
        thread = threading.Thread(
            target=read_video_into_queue,
            kwargs={
                'frames_per_second': frames_per_second,
                'func': generator,
                'video_queue': video_queue,
                'filepath': filepath,
                'errors': errors,
                'terminate': terminate,
                **kwargs
            })
        thread.start()
        try:
            while True:
                frame, frame_index, timestamp = video_queue.get()
                video_queue.task_done()
                if frame is None:
                    break
                yield (frame, frame_index, timestamp)
        finally:
            # Set the termination flag for the
            # background thread.
            terminate.set()
            try:
                # Unblock the thread, in the event
                # that it is waiting.
                video_queue.get_nowait()

                # Do it twice for the edge case
                # where the queue is completely
                # full and the end sentinel is
                # blocking.
                video_queue.get_nowait()
            except queue.Empty:
                # It doesn't matter if it's empty.
                pass
            # Wait for the background thread to terminate.
            thread.join()
    else:
        for frame, frame_index, timestamp in generator(
                filepath=filepath,
                frames_per_second=frames_per_second,
                errors=errors,
                **kwargs):
            yield (frame, frame_index, timestamp)

def compute_synchronized_video_hashes(filepath: str,
                                      hashers: dict,
                                      framerates=None,
                                      hash_format='base64',
                                      use_queue=True):
    """Compute the video hashes for a group of hashers with synchronized
    frame processing wherever possible.

    Args:
        filepath: Path to video file.
        hashers: A dictionary mapping hasher names to video hasher objects
        hash_format: The format in which to return the hashes
        use_queue: Whether to use queued video frames
    """
    if framerates is None:
        framerates = get_common_framerates({
            k: h.frames_per_second
            for k, h in hashers.items() if h.frames_per_second is not None
        })
    else:
        assert all(
            any(hasher_name in hasher_names
                for hasher_names in framerates.values())
            for hasher_name, hasher in hashers.items()
            if hasher.frames_per_second is not None
        ), 'Provided framerates do not have an entry for all required hashers.'

    results = {
        hasher_name: {
            'state':
            None,
            'hash':
            None,
            'relative_framerate':
            next(framerate / hasher.frames_per_second
                 for framerate, hasher_names in framerates.items()
                 if hasher_name in hasher_names)
        }
        for hasher_name, hasher in hashers.items()
        if hasher.frames_per_second is not None
    }
    for current_framerate, current_hasher_names in framerates.items():
        for frame_index, (frame, grabbed_frame_index,
                          frame_timestamp) in enumerate(
                              read_video(
                                  filepath=filepath,
                                  frames_per_second=current_framerate,
                                  use_queue=use_queue)):
            for hasher_name in current_hasher_names:
                config = results[hasher_name]
                hasher = hashers[hasher_name]
                assert config['relative_framerate'] is not None
                if frame_index % config['relative_framerate'] == 0:
                    config['state'] = hasher.process_frame(
                        frame=frame,
                        frame_index=grabbed_frame_index,
                        frame_timestamp=frame_timestamp,
                        state=config['state'])
        for hasher_name in current_hasher_names:
            config = results[hasher_name]
            hasher = hashers[hasher_name]
            current_hash = hasher.hash_from_final_state(state=config['state'])
            if hash_format == 'vector':
                config['hash'] = current_hash
            else:
                if not hasher.returns_multiple:
                    config['hash'] = hasher.vector_to_string(
                        current_hash, hash_format=hash_format)
                else:
                    config['hash'] = [
                        hasher.vector_to_string(h, hash_format=hash_format)
                        for h in current_hash
                    ]
            config['state'] = None
    hashes = {
        hasher_name: config['hash']
        for hasher_name, config in results.items()
    }
    for hasher_name, hasher in hashers.items():
        if hasher.frames_per_second is None:
            # This is a custom hasher that we just pass a video path to.
            hashes[hasher_name] = hasher.compute(filepath)
    return hashes

def unletterbox(image) -> typing.Optional[
        typing.Tuple[typing.Tuple[int, int], typing.Tuple[int, int]]]:
    """Return bounds of non-trivial region of image or None.

    Unletterboxing is cropping an image such that trivial edge regions
    are removed. Trivial in this context means that the majority of
    the values in that row or column are zero or very close to
    zero. This is why we don't use the terms "non-blank" or
    "non-empty."

    In order to do unletterboxing, this function returns bounds in the
    form (x1, x2), (y1, y2) where:

    - x1 is the index of the first column where over 10% of the pixels
      have means (average of R, G, B) > 2.
    - x2 is the index of the last column where over 10% of the pixels
      have means > 2.
    - y1 is the index of the first row where over 10% of the pixels
      have means > 2.
    - y2 is the index of the last row where over 10% of the pixels
      have means > 2.

    If there are zero columns or zero rows where over 10% of the
    pixels have means > 2, this function returns `None`.

    Note that in the case(s) of a single column and/or row of
    non-trivial pixels that it is possible for x1 = x2 and/or y1 = y2.

    Consider these examples to understand edge cases.  Given two
    images, `L` (entire left and bottom edges are 1, all other pixels
    0) and `U` (left, bottom and right edges 1, all other pixels 0),
    `unletterbox(L)` would return the bounds of the single bottom-left
    pixel and `unletterbox(U)` would return the bounds of the entire
    bottom row.

    Consider `U1` which is the same as `U` but with the bottom two
    rows all 1s. `unletterbox(U1)` returns the bounds of the bottom
    two rows.

    Args:
        image: The image from which to remove letterboxing.

    Returns:
        A pair of coordinates bounds of the form (x1, x2)
        and (y1, y2) representing the left, right, top, and
        bottom bounds.

    """
    # adj should be thought of as a boolean at each pixel indicating
    # whether or not that pixel is non-trivial (True) or not (False).
    adj = image.mean(axis=2) > 2

    if adj.all():
        return (0, image.shape[1] + 1), (0, image.shape[0] + 1)

    y = np.where(adj.sum(axis=1) > 0.1 * image.shape[1])[0]
    x = np.where(adj.sum(axis=0) > 0.1 * image.shape[0])[0]

    if len(y) == 0 or len(x) == 0:
        return None

    if len(y) == 1:
        y1 = y2 = y[0]
    else:
        y1, y2 = y[[0, -1]]
    if len(x) == 1:
        x1 = x2 = x[0]
    else:
        x1, x2 = x[[0, -1]]
    bounds = (x1, x2 + 1), (y1, y2 + 1)

    return bounds


'''
https://github.com/christiancosgrove/pytorch-spectral-normalization-gan

chainer: https://github.com/pfnet-research/sngan_projection
'''

# ResNet generator and discriminator
import torch
from torch import nn
import torch.nn.functional as F

# from spectral_normalization import SpectralNorm
import numpy as np
from torch.nn.utils import spectral_norm

channels = 3
bias = True
GEN_SIZE=64
DISC_SIZE=64

DIM_EMBED=128

class ConditionalBatchNorm2d(nn.Module):
    def __init__(self, num_features, dim_embed):
        super().__init__()
        self.num_features = num_features
        self.bn = nn.BatchNorm2d(num_features, affine=False)

        # self.embed = nn.Linear(dim_embed, num_features * 2, bias=False)
        # self.embed.weight.data[:, :num_features].normal_(1, 0.02)  # Initialise scale at N(1, 0.02)
        # self.embed.weight.data[:, num_features:].zero_()  # Initialise bias at 0
        # # self.embed = spectral_norm(self.embed) #seems not work

        self.embed_gamma = nn.Linear(dim_embed, num_features, bias=False)
        self.embed_beta = nn.Linear(dim_embed, num_features, bias=False)

    def forward(self, x, y):
        out = self.bn(x)

        # gamma, beta = self.embed(y).chunk(2, 1)
        # out = gamma.view(-1, self.num_features, 1, 1) * out + beta.view(-1, self.num_features, 1, 1)

        gamma = self.embed_gamma(y).view(-1, self.num_features, 1, 1)
        beta = self.embed_beta(y).view(-1, self.num_features, 1, 1)
        out = out + out*gamma + beta

        return out

class ResBlockGenerator(nn.Module):

    def __init__(self, in_channels, out_channels, dim_embed, bias=True):
        super(ResBlockGenerator, self).__init__()

        self.conv1 = nn.Conv2d(in_channels, out_channels, 3, 1, padding=1, bias=bias)
        self.conv2 = nn.Conv2d(out_channels, out_channels, 3, 1, padding=1, bias=bias)
        nn.init.xavier_uniform_(self.conv1.weight.data, np.sqrt(2))
        nn.init.xavier_uniform_(self.conv2.weight.data, np.sqrt(2))

        self.condbn1 = ConditionalBatchNorm2d(in_channels, dim_embed)
        self.condbn2 = ConditionalBatchNorm2d(out_channels, dim_embed)
        self.relu = nn.ReLU()
        self.upsample = nn.Upsample(scale_factor=2)

        # unconditional case
        self.model = nn.Sequential(
            nn.BatchNorm2d(in_channels),
            nn.ReLU(),
            nn.Upsample(scale_factor=2),
            self.conv1,
            nn.BatchNorm2d(out_channels),
            nn.ReLU(),
            self.conv2
            )

        self.bypass_conv = nn.Conv2d(in_channels,out_channels, 1, 1, padding=0, bias=bias) #h=h
        nn.init.xavier_uniform_(self.bypass_conv.weight.data, 1.0)
        self.bypass = nn.Sequential(
            nn.Upsample(scale_factor=2),
            self.bypass_conv,
        )

    def forward(self, x, y):
        if y is not None:
            out = self.condbn1(x, y)
            out = self.relu(out)
            out = self.upsample(out)
            out = self.conv1(out)
            out = self.condbn2(out, y)
            out = self.relu(out)
            out = self.conv2(out)
            out = out + self.bypass(x)
        else:
            out = self.model(x) + self.bypass(x)

        return out

class ResBlockDiscriminator(nn.Module):

    def __init__(self, in_channels, out_channels, stride=1):
        super(ResBlockDiscriminator, self).__init__()

        self.conv1 = nn.Conv2d(in_channels, out_channels, 3, 1, padding=1, bias=bias)
        self.conv2 = nn.Conv2d(out_channels, out_channels, 3, 1, padding=1, bias=bias)
        nn.init.xavier_uniform_(self.conv1.weight.data, np.sqrt(2))
        nn.init.xavier_uniform_(self.conv2.weight.data, np.sqrt(2))

        if stride == 1:
            self.model = nn.Sequential(
                nn.ReLU(),
                spectral_norm(self.conv1),
                nn.ReLU(),
                spectral_norm(self.conv2)
                )
        else:
            self.model = nn.Sequential(
                nn.ReLU(),
                spectral_norm(self.conv1),
                nn.ReLU(),
                spectral_norm(self.conv2),
                nn.AvgPool2d(2, stride=stride, padding=0)
                )

        self.bypass_conv = nn.Conv2d(in_channels,out_channels, 1, 1, padding=0, bias=bias)
        nn.init.xavier_uniform_(self.bypass_conv.weight.data, 1.0)
        if stride != 1:
            self.bypass = nn.Sequential(
                spectral_norm(self.bypass_conv),
                nn.AvgPool2d(2, stride=stride, padding=0)
            )
        else:
            self.bypass = nn.Sequential(
                spectral_norm(self.bypass_conv),
            )

    def forward(self, x):
        return self.model(x) + self.bypass(x)

# special ResBlock just for the first layer of the discriminator
class FirstResBlockDiscriminator(nn.Module):

    def __init__(self, in_channels, out_channels, stride=1):
        super(FirstResBlockDiscriminator, self).__init__()

        self.conv1 = nn.Conv2d(in_channels, out_channels, 3, 1, padding=1, bias=bias)
        self.conv2 = nn.Conv2d(out_channels, out_channels, 3, 1, padding=1, bias=bias)
        self.bypass_conv = nn.Conv2d(in_channels, out_channels, 1, 1, padding=0, bias=bias)
        nn.init.xavier_uniform_(self.conv1.weight.data, np.sqrt(2))
        nn.init.xavier_uniform_(self.conv2.weight.data, np.sqrt(2))
        nn.init.xavier_uniform_(self.bypass_conv.weight.data, 1.0)

        # we don't want to apply ReLU activation to raw image before convolution transformation.
        self.model = nn.Sequential(
            spectral_norm(self.conv1),
            nn.ReLU(),
            spectral_norm(self.conv2),
            nn.AvgPool2d(2)
            )
        self.bypass = nn.Sequential(
            nn.AvgPool2d(2),
            spectral_norm(self.bypass_conv),
        )

    def forward(self, x):
        return self.model(x) + self.bypass(x)

class cont_cond_cnn_generator(nn.Module):
    def __init__(self, nz=128, img_size=64, dim_embed=DIM_EMBED):
        super(cont_cond_cnn_generator, self).__init__()
        self.z_dim = nz
        self.dim_embed = dim_embed

        self.dense = nn.Linear(self.z_dim, 4 * 4 * GEN_SIZE*16, bias=True)
        self.final = nn.Conv2d(GEN_SIZE, channels, 3, stride=1, padding=1, bias=bias)
        nn.init.xavier_uniform_(self.dense.weight.data, 1.)
        nn.init.xavier_uniform_(self.final.weight.data, 1.)

        self.genblock0 = ResBlockGenerator(GEN_SIZE*16, GEN_SIZE*8, dim_embed=dim_embed) #4--->8
        self.genblock1 = ResBlockGenerator(GEN_SIZE*8, GEN_SIZE*4, dim_embed=dim_embed) #8--->16
        self.genblock2 = ResBlockGenerator(GEN_SIZE*4, GEN_SIZE*2, dim_embed=dim_embed) #16--->32
        self.genblock3 = ResBlockGenerator(GEN_SIZE*2, GEN_SIZE, dim_embed=dim_embed) #32--->64

        self.final = nn.Sequential(
            nn.BatchNorm2d(GEN_SIZE),
            nn.ReLU(),
            self.final,
            nn.Tanh()
        )

    def forward(self, z, y): #y is embedded in the feature space
        z = z.view(z.size(0), z.size(1))
        out = self.dense(z)
        out = out.view(-1, GEN_SIZE*16, 4, 4)

        out = self.genblock0(out, y)
        out = self.genblock1(out, y)
        out = self.genblock2(out, y)
        out = self.genblock3(out, y)
        out = self.final(out)

        return out

class cont_cond_cnn_discriminator(nn.Module):
    def __init__(self, img_size=64, dim_embed=DIM_EMBED):
        super(cont_cond_cnn_discriminator, self).__init__()
        self.dim_embed = dim_embed

        self.discblock1 = nn.Sequential(
            FirstResBlockDiscriminator(channels, DISC_SIZE, stride=2), #64--->32
            ResBlockDiscriminator(DISC_SIZE, DISC_SIZE*2, stride=2), #32--->16
            ResBlockDiscriminator(DISC_SIZE*2, DISC_SIZE*4, stride=2), #16--->8
        )
        self.discblock2 = ResBlockDiscriminator(DISC_SIZE*4, DISC_SIZE*8, stride=2) #8--->4
        self.discblock3 = nn.Sequential(
            ResBlockDiscriminator(DISC_SIZE*8, DISC_SIZE*16, stride=1), #4--->4;
            nn.ReLU(),
        )

        # self.linear1 = nn.Linear(DISC_SIZE*16, 1, bias=True)
        # nn.init.xavier_uniform_(self.linear1.weight.data, 1.)
        # self.linear1 = spectral_norm(self.linear1)
        # self.linear2 = nn.Linear(self.dim_embed, DISC_SIZE*16, bias=False)
        # nn.init.xavier_uniform_(self.linear2.weight.data, 1.)
        # self.linear2 = spectral_norm(self.linear2)

        self.linear1 = nn.Linear(DISC_SIZE*16*4*4, 1, bias=True)
        nn.init.xavier_uniform_(self.linear1.weight.data, 1.)
        self.linear1 = spectral_norm(self.linear1)
        self.linear2 = nn.Linear(self.dim_embed, DISC_SIZE*16*4*4, bias=False)
        nn.init.xavier_uniform_(self.linear2.weight.data, 1.)
        self.linear2 = spectral_norm(self.linear2)

    def forward(self, x, y):
        output = self.discblock1(x)
        output = self.discblock2(output)
        output = self.discblock3(output)

        # output = torch.sum(output, dim=(2,3))
        # output_y = torch.sum(output*self.linear2(y), 1, keepdim=True)
        # output = self.linear1(output) + output_y

        output = output.view(-1, DISC_SIZE*16*4*4)
        output_y = torch.sum(output*self.linear2(y), 1, keepdim=True)
        output = self.linear1(output) + output_y

        return output.view(-1, 1)


from __future__ import print_function
import numpy as np
import copy
import logging
import traceback
from simtk import openmm, unit
from perses.dispersed.feptasks import Particle, compute_reduced_potential
from perses.storage import NetCDFStorageView
from perses.annihilation.relative import HybridTopologyFactory
from perses.tests.utils import quantity_is_finite
from openmmtools.constants import kB
from openmmtools.cache import LRUCache, global_context_cache
from openmmtools.states import ThermodynamicState, SamplerState, CompoundThermodynamicState
from perses.annihilation.lambda_protocol import RelativeAlchemicalState, LambdaProtocol

default_temperature = 300.0*unit.kelvin
default_nsteps = 1
default_timestep = 1.0 * unit.femtoseconds
default_steps_per_propagation = 1
_logger = logging.getLogger("NCMCEngine")

class NaNException(Exception):
    def __init__(self, *args, **kwargs):
        super(NaNException,self).__init__(*args,**kwargs)

class NCMCEngine(object):
    """
    NCMC switching engine

    Examples
    --------

    Create a transformation for an alanine dipeptide test system where the N-methyl group is eliminated.

    >>> from openmmtools import testsystems
    >>> testsystem = testsystems.AlanineDipeptideVacuum()
    >>> from perses.rjmc.topology_proposal import TopologyProposal
    >>> new_to_old_atom_map = { index : index for index in range(testsystem.system.getNumParticles()) if (index > 3) } # all atoms but N-methyl
    >>> topology_proposal = TopologyProposal(old_system=testsystem.system, old_topology=testsystem.topology, old_chemical_state_key='AA', new_chemical_state_key='AA', new_system=testsystem.system, new_topology=testsystem.topology, logp_proposal=0.0, new_to_old_atom_map=new_to_old_atom_map, metadata=dict())
    >>> ncmc_engine = NCMCEngine(temperature=300.0*unit.kelvin, functions=default_functions, nsteps=50, timestep=1.0*unit.femtoseconds)
    >>> positions = testsystem.positions
    >>> [positions, logP_delete, potential_delete] = ncmc_engine.integrate(topology_proposal, positions, direction='delete')
    >>> [positions, logP_insert, potential_insert] = ncmc_engine.integrate(topology_proposal, positions, direction='insert')

    """

    def __init__(self, temperature=default_temperature, functions=None, nsteps=default_nsteps,
                 steps_per_propagation=default_steps_per_propagation, timestep=default_timestep,
                 constraint_tolerance=None, platform=None, write_ncmc_interval=1, measure_shadow_work=False,
                 integrator_splitting='V R O H R V', storage=None, verbose=False, LRUCapacity=10, pressure=None, bond_softening_constant=1.0, angle_softening_constant=1.0):
        """
        This is the base class for NCMC switching between two different systems.

        Arguments
        ---------
        temperature : simtk.unit.Quantity with units compatible with kelvin
            The temperature at which switching is to be run
        functions : dict of str:str, optional, default=default_functions
            functions[parameter] is the function (parameterized by 't' which switched from 0 to 1) that
            controls how alchemical context parameter 'parameter' is switched
        nsteps : int, optional, default=1
            The number of steps to use for switching.
        steps_per_propagation : int, optional, default=1
            The number of intermediate propagation steps taken at each switching step
        timestep : simtk.unit.Quantity with units compatible with femtoseconds, optional, default=1*femtosecond
            The timestep to use for integration of switching velocity Verlet steps.
        constraint_tolerance : float, optional, default=None
            If not None, this relative constraint tolerance is used for position and velocity constraints.
        platform : simtk.openmm.Platform, optional, default=None
            If specified, the platform to use for OpenMM simulations.
        write_ncmc_interval : int, optional, default=None
            If a positive integer is specified, a snapshot frame will be written to storage with the specified interval on NCMC switching.
            'storage' must also be specified.
        measure_shadow_work : bool, optional, default False
            Whether to measure shadow work
        integrator_splitting : str, optional, default='V R O H R V'
            NCMC internal integrator splitting based on OpenMMTools Langevin splittings
        storage : NetCDFStorageView, optional, default=None
            If specified, write data using this class.
        verbose : bool, optional, default=False
            If True, print debug information.
        LRUCapacity : int, default 10
            Capacity of LRU cache for hybrid systems
        pressure : float, default None
            The pressure to use for the simulation. If None, no barostat
        """
        # Handle some defaults.
        if functions == None:
            functions = LambdaProtocol.default_functions
        if nsteps == None:
            nsteps = default_nsteps
        if timestep == None:
            timestep = default_timestep
        if temperature == None:
            temperature = default_temperature

        self._temperature = temperature
        self._functions = copy.deepcopy(functions)
        self._nsteps = nsteps
        self._timestep = timestep
        self._constraint_tolerance = constraint_tolerance
        self._platform = platform
        self._integrator_splitting = integrator_splitting
        self._steps_per_propagation = steps_per_propagation
        self._verbose = verbose
        self._pressure = pressure
        self._bond_softening_constant = bond_softening_constant
        self._angle_softening_constant = angle_softening_constant
        self._disable_barostat = False
        self._hybrid_cache = LRUCache(capacity=LRUCapacity)
        self._measure_shadow_work = measure_shadow_work

        self._nattempted = 0

        self._storage = None
        if storage is not None:
            self._storage = NetCDFStorageView(storage, modname=self.__class__.__name__)
            self._save_configuration = True
        else:
            self._save_configuration = False
        if write_ncmc_interval is not None:
            self._write_ncmc_interval = write_ncmc_interval
        else:
            self._write_ncmc_interval = 1
        self._work_save_interval = write_ncmc_interval

    @property
    def beta(self):
        kT = kB * self._temperature
        beta = 1.0 / kT
        return beta

    def _compute_energy_contribution(self, hybrid_thermodynamic_state, initial_sampler_state, final_sampler_state):
        """
        Compute NCMC energy contribution to log probability.

        See Eqs. 62 and 63 (two-stage) and Eq. 45 (hybrid) of reference document.
        In both cases, the contribution is u(final_positions, final_lambda) - u(initial_positions, initial_lambda).

        Parameters
        ----------
        hybrid_thermodynamic_state : openmmtools.states.CompoundThermodynamicState
            The thermodynamic state of the hybrid sampler.
        initial_sampler_state : openmmtools.states.SamplerState
            The sampler state of the nonalchemical system at the start of the NCMC protocol with box vectors
        final_sampler_state : openmmtools.states.SamplerState
            The sampler state of the nonalchemical system at the end of the NCMC protocol

        Returns
        -------
        logP_energy : float
            The NCMC energy contribution to log probability.
        """
        hybrid_thermodynamic_state.set_alchemical_parameters(0.0)
        initial_reduced_potential = compute_reduced_potential(hybrid_thermodynamic_state, initial_sampler_state)

        hybrid_thermodynamic_state.set_alchemical_parameters(1.0)
        final_reduced_potential = compute_reduced_potential(hybrid_thermodynamic_state, final_sampler_state)

        return final_reduced_potential - initial_reduced_potential

    def _topology_proposal_to_thermodynamic_states(self, topology_proposal):
        """
        Convert a topology proposal to thermodynamic states for the end systems. This will be used to compute the
        "logP_energy" quantity.

        Arguments
        ---------
        topology_proposal : perses.rjmc.TopologyProposal
            topology proposal for whose endpoint systems we want ThermodynamicStates

        Returns
        -------
        old_thermodynamic_state : openmmtools.states.ThermodynamicState
            The old system (nonalchemical) thermodynamic state
        new_thermodynamic_state : openmmtools.states.ThermodynamicState
            The new system (nonalchemical) thermodynamic state
        """
        systems = [topology_proposal.old_system, topology_proposal.new_system]
        thermostates = []
        for system in systems:
            thermodynamic_state = ThermodynamicState(system, temperature=self._temperature, pressure=self._pressure)
            thermostates.append(thermodynamic_state)

        return thermostates[0], thermostates[1]

    def make_alchemical_system(self, topology_proposal, current_positions, new_positions):
        """
        Generate an alchemically-modified system at the correct atoms
        based on the topology proposal. This method generates a hybrid system using the new
        HybridTopologyFactory. It memoizes so that calling multiple times (within a recent time period)
        will immediately return a cached object.

        Arguments
        ---------
        topology_proposal : perses.rjmc.TopologyProposal
            Unmodified real system corresponding to appropriate leg of transformation.
        current_positions : np.ndarray of float
            Positions of "old" system
        new_positions : np.ndarray of float
            Positions of "new" system atoms

        Returns
        -------
        hybrid_factory : perses.annihilation.relative.HybridTopologyFactory
            a factory object containing the hybrid system
        """
        try:
            hybrid_factory = self._hybrid_cache[topology_proposal]

            #If we've retrieved the factory from the cache, update it to include the relevant positions
            hybrid_factory._old_positions = current_positions
            hybrid_factory._new_positions = new_positions
            hybrid_factory._compute_hybrid_positions()
        except KeyError:
            try:
                hybrid_factory = HybridTopologyFactory(topology_proposal, current_positions, new_positions, bond_softening_constant=self._bond_softening_constant, angle_softening_constant=self._angle_softening_constant)
                self._hybrid_cache[topology_proposal] = hybrid_factory
            except:
                hybrid_factory = None

        return hybrid_factory

    def integrate(self, topology_proposal, initial_sampler_state, proposed_sampler_state, iteration=None):
        """
        Performs NCMC switching to either delete or insert atoms according to the provided `topology_proposal`.

        For `delete`, the system is first modified from fully interacting to alchemically modified, and then NCMC switching is used to eliminate atoms.
        For `insert`, the system begins with eliminated atoms in an alchemically noninteracting form and NCMC switching is used to turn atoms on, followed by making system real.

        Parameters
        ----------
        topology_proposal : TopologyProposal
            Contains old/new Topology and System objects and atom mappings.
        initial_sampler_state : openmmtools.states.SamplerState representing the initial (old) system
            Configurational properties of the atoms at the beginning of the NCMC switching.
        proposed_sampler_state : openmmtools.states.SamplerState representing the proposed (post-geometry new) system
            Configurational properties new system atoms at beginning of NCMC switching
        iteration : int, optional, default=None
            Iteration number, for storage purposes.

        Returns
        -------
        final_old_sampler_state : openmmtools.State.SamplerState
            The final configurational properties of the old system after hybrid alchemical switching
        final_sampler_state : openmmtools.states.SamplerState
            The final configurational properties after `nsteps` steps of alchemical switching, and reversion to the nonalchemical system
        logP_work : float
            The NCMC work contribution to the log acceptance probability (Eqs. 62 and 63)
        logP_initial : float
            The initial logP of the hybrid configuration
        logP_final : float
            The final logP of the hybrid configuration
        """

        assert not initial_sampler_state.has_nan() and not proposed_sampler_state.has_nan()

        #generate or retrieve the hybrid topology factory:
        hybrid_factory = self.make_alchemical_system(topology_proposal, initial_sampler_state.positions, proposed_sampler_state.positions)

        if hybrid_factory is None:
            _logger.warning("Unable to construct hybrid system for {} -> {}".format(topology_proposal.old_chemical_state_key, topology_proposal.new_chemical_state_key))
            return initial_sampler_state, proposed_sampler_state, -np.inf, 0.0, 0.0

        topology = hybrid_factory.hybrid_topology

        #generate the corresponding thermodynamic and sampler states so that we can use the NonequilibriumSwitchingMove:

        #First generate the thermodynamic state:
        hybrid_system = hybrid_factory.hybrid_system
        hybrid_thermodynamic_state = ThermodynamicState(hybrid_system, temperature=self._temperature, pressure=self._pressure)

        #Now create an RelativeAlchemicalState from the hybrid system:
        alchemical_state = RelativeAlchemicalState.from_system(hybrid_system)
        alchemical_state.set_alchemical_parameters(0.0)

        #Now create a compound thermodynamic state that combines the hybrid thermodynamic state with the alchemical state:
        compound_thermodynamic_state = CompoundThermodynamicState(hybrid_thermodynamic_state, composable_states=[alchemical_state])

        #construct a sampler state from the hybrid positions and the box vectors of the initial sampler state:
        initial_hybrid_positions = hybrid_factory.hybrid_positions
        initial_hybrid_box_vectors = initial_sampler_state.box_vectors

        initial_hybrid_sampler_state = SamplerState(initial_hybrid_positions, box_vectors=initial_hybrid_box_vectors)
        final_hybrid_sampler_state = copy.deepcopy(initial_hybrid_sampler_state)

        #create the nonequilibrium move:
        #ne_move = NonequilibriumSwitchingMove(self._functions, self._integrator_splitting, self._temperature, self._nsteps, self._timestep,
        #                                      work_save_interval=self._write_ncmc_interval, top=topology,subset_atoms=None,
        #                                      save_configuration=self._save_configuration, measure_shadow_work=self._measure_shadow_work)

        ne_move = ExternalNonequilibriumSwitchingMove(self._functions, nsteps_neq=self._nsteps,
                                                      timestep=self._timestep, temperature=self._temperature,
                                                      work_configuration_save_interval=self._work_save_interval,
                                                      splitting="V R O R V")

        #run the NCMC protocol
        try:
            ne_move.apply(compound_thermodynamic_state, final_hybrid_sampler_state)
        except Exception as e:
            _logger.warn("NCMC failed because {}; rejecting.".format(str(e)))
            logP_work = -np.inf
            return [initial_sampler_state, proposed_sampler_state, -np.inf, 0.0, 0.0]

        #get the total work:
        logP_work = - ne_move.cumulative_work[-1]

        # Compute contribution of transforming to and from the hybrid system:
        context, integrator = global_context_cache.get_context(hybrid_thermodynamic_state)

        #set all alchemical parameters to zero:
        for parameter in self._functions.keys():
            context.setParameter(parameter, 0.0)

        initial_hybrid_sampler_state.apply_to_context(context, ignore_velocities=True)
        initial_reduced_potential = hybrid_thermodynamic_state.reduced_potential(context)

        #set all alchemical parameters to one:
        for parameter in self._functions.keys():
            context.setParameter(parameter, 1.0)

        final_hybrid_sampler_state.apply_to_context(context, ignore_velocities=True)
        final_reduced_potential = hybrid_thermodynamic_state.reduced_potential(context)

        #reset the parameters back to zero just in case
        for parameter in self._functions.keys():
            context.setParameter(parameter, 0.0)

        #compute the output SamplerState, which has the atoms only for the new system post-NCMC:
        new_positions = hybrid_factory.new_positions(final_hybrid_sampler_state.positions)
        new_box_vectors = final_hybrid_sampler_state.box_vectors
        final_sampler_state = SamplerState(new_positions, box_vectors=new_box_vectors)

        #compute the output SamplerState for the atoms only in the old system (required for geometry_logP_reverse)
        old_positions = hybrid_factory.old_positions(final_hybrid_sampler_state.positions)
        old_box_vectors = copy.deepcopy(new_box_vectors) #these are the same as the new system
        final_old_sampler_state = SamplerState(old_positions, box_vectors=old_box_vectors)

        #extract the trajectory and box vectors from the move:
        trajectory = ne_move.trajectory[::-self._write_ncmc_interval, :, :][::-1]
        topology = hybrid_factory.hybrid_topology
        position_varname = "ncmcpositions"
        nframes = np.shape(trajectory)[0]

        #extract box vectors:
        box_vec_varname = "ncmcboxvectors"
        box_lengths = ne_move.box_lengths[::-self._write_ncmc_interval, :][::-1]
        box_angles = ne_move.box_angles[::-self._write_ncmc_interval, :][::-1]
        box_lengths_and_angles = np.stack([box_lengths, box_angles])

        #write out the positions of the topology
        if self._storage:
            for frame in range(nframes):
                self._storage.write_configuration(position_varname, trajectory[frame, :, :], topology, iteration=iteration, frame=frame, nframes=nframes)

        #write out the periodict box vectors:
        if self._storage:
            self._storage.write_array(box_vec_varname, box_lengths_and_angles, iteration=iteration)

        #retrieve the protocol work and write that out too:
        protocol_work = ne_move.cumulative_work
        if self._storage:
            self._storage.write_array("protocolwork", protocol_work, iteration=iteration)

        # Return
        return [final_old_sampler_state, final_sampler_state, logP_work, -initial_reduced_potential, -final_reduced_potential]