# This software is open source software available under the BSD-3 license.
#
# Copyright (c) 2020 Triad National Security, LLC. All rights reserved.
# Copyright (c) 2020 Lawrence Livermore National Security, LLC. All rights
# reserved.
# Copyright (c) 2020 UT-Battelle, LLC. All rights reserved.
#
# Additional copyright and license information can be found in the LICENSE file
# distributed with this code, or at
# https://raw.githubusercontent.com/MPAS-Dev/MPAS-Analysis/master/LICENSE
"""
Functions for creating climatologies from monthly time series data
"""
# Authors
# -------
# Xylar Asay-Davis

from __future__ import absolute_import, division, print_function, \
    unicode_literals

import xarray as xr
import os
import numpy
from tempfile import TemporaryDirectory

from pyremap import Remapper, LatLonGridDescriptor, ProjectionGridDescriptor

from mpas_analysis.shared.constants import constants

from mpas_analysis.shared.timekeeping.utility import days_to_datetime

from mpas_analysis.shared.io.utility import build_config_full_path, \
    make_directories, fingerprint_generator
from mpas_analysis.shared.io import write_netcdf

from mpas_analysis.shared.climatology.comparison_descriptors import \
    get_comparison_descriptor

def get_remapper(config, sourceDescriptor, comparisonDescriptor,
                 mappingFilePrefix, method, logger=None):  # {{{
    """
    Given config options and descriptions of the source and comparison grids,
    returns a ``pyremap.Remapper`` object that can be used to remap from source
    files or data sets to corresponding data sets on the comparison grid.

    If necessary, creates the mapping file containing weights and indices
    needed to perform remapping.

    Parameters
    ----------
    config :  instance of ``MpasAnalysisConfigParser``
        Contains configuration options

    sourceDescriptor : ``MeshDescriptor`` subclass object
        A description of the source mesh or grid

    comparisonDescriptor : ``MeshDescriptor`` subclass object
        A description of the comparison grid

    mappingFilePrefix : str
        A prefix to be prepended to the mapping file name

    method : {'bilinear', 'neareststod', 'conserve'}
        The method of interpolation used.

    logger : ``logging.Logger``, optional
        A logger to which ncclimo output should be redirected

    Returns
    -------
    remapper : ``pyremap.Remapper`` object
        A remapper that can be used to remap files or data sets from the source
        grid or mesh to the comparison grid.
    """
    # Authors
    # -------
    # Xylar Asay-Davis

    mappingFileName = None

    if not _matches_comparison(sourceDescriptor, comparisonDescriptor):
        # we need to remap because the grids don't match

        mappingBaseName = '{}_{}_to_{}_{}.nc'.format(
            mappingFilePrefix,
            sourceDescriptor.meshName,
            comparisonDescriptor.meshName,
            method)

        tryCustom = config.get('diagnostics', 'customDirectory') != 'none'
        if tryCustom:
            # first see if mapping files are in the custom directory
            mappingSubdirectory = build_config_full_path(
                config, 'diagnostics', 'mappingSubdirectory',
                baseDirectoryOption='customDirectory')

            mappingFileName = '{}/{}'.format(mappingSubdirectory,
                                         mappingBaseName)
        if not tryCustom or not os.path.exists(mappingFileName):
            # second see if mapping files are in the base directory

            mappingSubdirectory = build_config_full_path(
                config, 'diagnostics', 'mappingSubdirectory',
                baseDirectoryOption='base_path')

            mappingFileName = '{}/{}'.format(mappingSubdirectory,
                                             mappingBaseName)

        if not os.path.exists(mappingFileName):
            # we don't have a mapping file yet, so get ready to create one
            # in the output subfolder if needed
            mappingSubdirectory = \
                build_config_full_path(config, 'output',
                                       'mappingSubdirectory')
            make_directories(mappingSubdirectory)
            mappingFileName = '{}/{}'.format(mappingSubdirectory,
                                             mappingBaseName)

    remapper = Remapper(sourceDescriptor, comparisonDescriptor,
                        mappingFileName)

    mpiTasks = config.getWithDefault('execute', 'mapMpiTasks', 1)
    esmf_parallel_exec = config.get('execute', 'mapParallelExec')
    if esmf_parallel_exec == 'None':
        esmf_parallel_exec = None

    mappingSubdirectory = \
        build_config_full_path(config, 'output',
                               'mappingSubdirectory')
    make_directories(mappingSubdirectory)
    with TemporaryDirectory(dir=mappingSubdirectory) as tempdir:
        remapper.build_mapping_file(method=method, logger=logger,
                                    mpiTasks=mpiTasks, tempdir=tempdir,
                                    esmf_parallel_exec=esmf_parallel_exec)

    return remapper  # }}}

def compute_monthly_climatology(ds, calendar=None, maskVaries=True):  # {{{
    """
    Compute monthly climatologies from a data set.  The mean is weighted but
    the number of days in each month of the data set, ignoring values masked
    out with NaNs.  If the month coordinate is not present, a data array
    ``month`` will be added based on ``Time`` and the provided calendar.

    Parameters
    ----------
    ds : xarray.Dataset or xarray.DataArray
        A data set with a ``Time`` coordinate expressed as days since
        0001-01-01 or ``month`` coordinate

    calendar : {'gregorian', 'gregorian_noleap'}, optional
        The name of one of the calendars supported by MPAS cores, used to
        determine ``month`` from ``Time`` coordinate, so must be supplied if
        ``ds`` does not already have a ``month`` coordinate or data array

    maskVaries : bool, optional
        If the mask (where variables in ``ds`` are ``NaN``) varies with time.
        If not, the weighted average does not need make extra effort to account
        for the mask.  Most MPAS fields will have masks that don't vary in
        time, whereas observations may sometimes be present only at some
        times and not at others, requiring ``maskVaries = True``.

    Returns
    -------
    climatology : object of same type as ``ds``
        A data set without the ``'Time'`` coordinate containing the mean
        of ds over all months in monthValues, weighted by the number of days
        in each month.
    """
    # Authors
    # -------
    # Xylar Asay-Davis

    def compute_one_month_climatology(ds):
        monthValues = list(ds.month.values)
        return compute_climatology(ds, monthValues, calendar, maskVaries)

    ds = add_years_months_days_in_month(ds, calendar)

    monthlyClimatology = \
        ds.groupby('month').map(compute_one_month_climatology)

    return monthlyClimatology  # }}}

def compute_climatology(ds, monthValues, calendar=None,
                        maskVaries=True):  # {{{
    """
    Compute a monthly, seasonal or annual climatology data set from a data
    set.  The mean is weighted but the number of days in each month of
    the data set, ignoring values masked out with NaNs.  If the month
    coordinate is not present, a data array ``month`` will be added based
    on ``Time`` and the provided calendar.

    Parameters
    ----------
    ds : xarray.Dataset or xarray.DataArray
        A data set with a ``Time`` coordinate expressed as days since
        0001-01-01 or ``month`` coordinate

    monthValues : int or array-like of ints
        A single month or an array of months to be averaged together

    calendar : {'gregorian', 'gregorian_noleap'}, optional
        The name of one of the calendars supported by MPAS cores, used to
        determine ``month`` from ``Time`` coordinate, so must be supplied if
        ``ds`` does not already have a ``month`` coordinate or data array

    maskVaries : bool, optional
        If the mask (where variables in ``ds`` are ``NaN``) varies with time.
        If not, the weighted average does not need make extra effort to account
        for the mask.  Most MPAS fields will have masks that don't vary in
        time, whereas observations may sometimes be present only at some
        times and not at others, requiring ``maskVaries = True``.

    Returns
    -------
    climatology : object of same type as ``ds``
        A data set without the ``'Time'`` coordinate containing the mean
        of ds over all months in monthValues, weighted by the number of days
        in each month.
    """
    # Authors
    # -------
    # Xylar Asay-Davis

    ds = add_years_months_days_in_month(ds, calendar)

    mask = xr.zeros_like(ds.month, bool)

    for month in monthValues:
        mask = numpy.logical_or(mask, ds.month == month)

    climatologyMonths = ds.where(mask, drop=True)

    climatology = _compute_masked_mean(climatologyMonths, maskVaries)

    return climatology  # }}}

def add_years_months_days_in_month(ds, calendar=None):  # {{{
    '''
    Add ``year``, ``month`` and ``daysInMonth`` as data arrays in ``ds``.
    The number of days in each month of ``ds`` is computed either using the
    ``startTime`` and ``endTime`` if available or assuming ``gregorian_noleap``
    calendar and ignoring leap years.  ``year`` and ``month`` are computed
    accounting correctly for the the calendar.

    Parameters
    ----------
    ds : ``xarray.Dataset`` or ``xarray.DataArray`` object
        A data set with a ``Time`` coordinate expressed as days since
        0001-01-01

    calendar : {'gregorian', 'gregorian_noleap'}, optional
        The name of one of the calendars supported by MPAS cores, used to
        determine ``year`` and ``month`` from ``Time`` coordinate

    Returns
    -------
    ds : object of same type as ``ds``
        The data set with ``year``, ``month`` and ``daysInMonth`` data arrays
        added (if not already present)
    '''
    # Authors
    # -------
    # Xylar Asay-Davis

    if ('year' in ds.coords and 'month' in ds.coords and
            'daysInMonth' in ds.coords):
        return ds

    ds = ds.copy()

    if 'year' not in ds.coords or 'month' not in ds.coords:
        if calendar is None:
            raise ValueError('calendar must be provided if month and year '
                             'coordinate is not in ds.')
        datetimes = days_to_datetime(ds.Time, calendar=calendar)

    if 'year' not in ds.coords:
        ds.coords['year'] = ('Time', [date.year for date in datetimes])

    if 'month' not in ds.coords:
        ds.coords['month'] = ('Time', [date.month for date in datetimes])

    if 'daysInMonth' not in ds.coords:
        if 'startTime' in ds.coords and 'endTime' in ds.coords:
            ds.coords['daysInMonth'] = ds.endTime - ds.startTime
        else:
            if calendar == 'gregorian':
                print('Warning: The MPAS run used the Gregorian calendar '
                      'but does not appear to have\n'
                      'supplied start and end times.  Climatologies '
                      'will be computed with\n'
                      'month durations ignoring leap years.')

            daysInMonth = numpy.array(
                [constants.daysInMonth[int(month) - 1] for
                 month in ds.month.values], float)
            ds.coords['daysInMonth'] = ('Time', daysInMonth)

    return ds  # }}}

def remap_and_write_climatology(config, climatologyDataSet,
                                climatologyFileName, remappedFileName,
                                remapper, logger=None):  # {{{
    """
    Given a field in a climatology data set, use the ``remapper`` to remap
    horizontal dimensions of all fields, write the results to an output file,
    and return the remapped data set.

    Note that ``climatologyFileName`` and ``remappedFileName`` will be
    overwritten if they exist, so if this behavior is not desired, the calling
    code should skip this call if the files exist and simply load the contents
    of ``remappedFileName``.

    Parameters
    ----------
    config :  instance of ``MpasAnalysisConfigParser``
        Contains configuration options

    climatologyDataSet : ``xarray.DataSet`` or ``xarray.DataArray`` object
        A data set containing a climatology

    fieldName : str
        A field within the climatology to be remapped

    climatologyFileName : str
        The name of the output file to which the data set should be written
        before remapping (if using ncremap).

    remappedFileName : str
        The name of the output file to which the remapped data set should
        be written.

    remapper : ``pyremap.Remapper`` object
        A remapper that can be used to remap files or data sets to a
        comparison grid.

    logger : ``logging.Logger``, optional
        A logger to which ncclimo output should be redirected

    Returns
    -------
    remappedClimatology : ``xarray.DataSet`` or ``xarray.DataArray`` object
        A data set containing the remapped climatology
    """
    # Authors
    # -------
    # Xylar Asay-Davis

    useNcremap = config.getboolean('climatology', 'useNcremap')

    if remapper.mappingFileName is None:
        # no remapping is needed
        remappedClimatology = climatologyDataSet
    else:
        renormalizationThreshold = config.getfloat(
            'climatology', 'renormalizationThreshold')
        parallel_exec = config.get(
            'execute', 'ncremapParallelExec')
        if parallel_exec == 'None':
            parallel_exec = None

        if useNcremap:
            if not os.path.exists(climatologyFileName):
                write_netcdf(climatologyDataSet, climatologyFileName)
            remapper.remap_file(inFileName=climatologyFileName,
                                outFileName=remappedFileName,
                                overwrite=True,
                                renormalize=renormalizationThreshold,
                                logger=logger,
                                parallel_exec=parallel_exec)
            remappedClimatology = xr.open_dataset(remappedFileName)
        else:

            remappedClimatology = remapper.remap(climatologyDataSet,
                                                 renormalizationThreshold)
            write_netcdf(remappedClimatology, remappedFileName)
    return remappedClimatology  # }}}

def get_unmasked_mpas_climatology_directory(config, op='avg'):  # {{{
    """
    Get the directory for an unmasked MPAS climatology produced by ncclimo,
    making the directory if it doesn't already exist

    Parameters
    ----------
    config :  ``MpasAnalysisConfigParser``
        configuration options

    op : {'avg', 'min', 'max'}
         operator for monthly stats
    """
    # Authors
    # -------
    # Xylar Asay-Davis

    climatologyOpDirectory = get_climatology_op_directory(config, op)

    mpasMeshName = config.get('input', 'mpasMeshName')

    directory = '{}/unmasked_{}'.format(climatologyOpDirectory,
                                        mpasMeshName)

    make_directories(directory)
    return directory  # }}}

def get_unmasked_mpas_climatology_file_name(config, season, componentName,
                                            op='avg'):
    # {{{
    """
    Get the file name for an unmasked MPAS climatology produced by ncclimo

    Parameters
    ----------
    config :  ``MpasAnalysisConfigParser``
        configuration options

    season : str
        One of the seasons in ``constants.monthDictionary``

    componentName : {'ocean', 'seaIce'}
        The MPAS component for which the climatology is being computed

    op : {'avg', 'min', 'max'}
         operator for monthly stats
    """
    # Authors
    # -------
    # Xylar Asay-Davis

    startYear = config.getint('climatology', 'startYear')
    endYear = config.getint('climatology', 'endYear')

    if componentName == 'ocean':
        ncclimoModel = 'mpaso'
    elif componentName == 'seaIce':
        ncclimoModel = 'mpascice'
    else:
        raise ValueError('component {} is not supported by ncclimo.\n'
                         'Check with Charlie Zender and Xylar Asay-Davis\n'
                         'about getting it added'.format(componentName))

    directory = get_unmasked_mpas_climatology_directory(config, op)

    make_directories(directory)
    monthValues = sorted(constants.monthDictionary[season])
    startMonth = monthValues[0]
    endMonth = monthValues[-1]

    suffix = '{:04d}{:02d}_{:04d}{:02d}_climo'.format(
        startYear, startMonth, endYear, endMonth)

    if season in constants.abrevMonthNames:
        season = '{:02d}'.format(monthValues[0])
    fileName = '{}/{}_{}_{}.nc'.format(directory, ncclimoModel,
                                       season, suffix)
    return fileName  # }}}

def get_masked_mpas_climatology_file_name(config, season, componentName,
                                          climatologyName, op='avg'):  # {{{
    """
    Get the file name for a masked MPAS climatology

    Parameters
    ----------
    config :  ``MpasAnalysisConfigParser``
        Configuration options

    season : str
        One of the seasons in ``constants.monthDictionary``

    componentName : {'ocean', 'seaIce'}
        The MPAS component for which the climatology is being computed

    climatologyName : str
        The name of the climatology (typically the name of a field to mask
        and later remap)

    op : {'avg', 'min', 'max'}
         operator for monthly stats
    """
    # Authors
    # -------
    # Xylar Asay-Davis

    startYear = config.getint('climatology', 'startYear')
    endYear = config.getint('climatology', 'endYear')
    mpasMeshName = config.get('input', 'mpasMeshName')

    if componentName == 'ocean':
        ncclimoModel = 'mpaso'
    elif componentName == 'seaIce':
        ncclimoModel = 'mpascice'
    else:
        raise ValueError('component {} is not supported by ncclimo.\n'
                         'Check with Charlie Zender and Xylar Asay-Davis\n'
                         'about getting it added'.format(componentName))

    climatologyOpDirectory = get_climatology_op_directory(config, op)

    stageDirectory = '{}/masked'.format(climatologyOpDirectory)

    directory = '{}/{}_{}'.format(
        stageDirectory, climatologyName,
        mpasMeshName)

    make_directories(directory)

    monthValues = sorted(constants.monthDictionary[season])
    startMonth = monthValues[0]
    endMonth = monthValues[-1]

    suffix = '{:04d}{:02d}_{:04d}{:02d}_climo'.format(
        startYear, startMonth, endYear, endMonth)

    if season in constants.abrevMonthNames:
        season = '{:02d}'.format(monthValues[0])
    fileName = '{}/{}_{}_{}.nc'.format(
        directory, ncclimoModel, season, suffix)

    return fileName  # }}}

def get_remapped_mpas_climatology_file_name(config, season, componentName,
                                            climatologyName,
                                            comparisonGridName,
                                            op='avg'):  # {{{
    """
    Get the file name for a masked MPAS climatology

    Parameters
    ----------
    config :  ``MpasAnalysisConfigParser``
        Configuration options

    season : str
        One of the seasons in ``constants.monthDictionary``

    componentName : {'ocean', 'seaIce'}
        The MPAS component for which the climatology is being computed

    climatologyName : str
        The name of the climatology (typically the name of a field to mask
        and later remap)

    comparisonGridName : str
        The name of the comparison grid to use for remapping.  If it is one
        of the default comparison grid names ``{'latlon', 'antarctic',
        'arctic'}``, the full grid name is looked up via
        get_comparison_descriptor

    op : {'avg', 'min', 'max'}
         operator for monthly stats
    """
    # Authors
    # -------
    # Xylar Asay-Davis

    startYear = config.getint('climatology', 'startYear')
    endYear = config.getint('climatology', 'endYear')
    mpasMeshName = config.get('input', 'mpasMeshName')

    if componentName == 'ocean':
        ncclimoModel = 'mpaso'
    elif componentName == 'seaIce':
        ncclimoModel = 'mpascice'
    else:
        raise ValueError('component {} is not supported by ncclimo.\n'
                         'Check with Charlie Zender and Xylar Asay-Davis\n'
                         'about getting it added'.format(componentName))

    climatologyOpDirectory = get_climatology_op_directory(config, op)

    if comparisonGridName in ['latlon', 'antarctic', 'arctic']:
        comparisonDescriptor = get_comparison_descriptor(config,
                                                         comparisonGridName)
        comparisonFullMeshName = comparisonDescriptor.meshName
    else:
        comparisonFullMeshName = comparisonGridName

    stageDirectory = '{}/remapped'.format(climatologyOpDirectory)

    directory = '{}/{}_{}_to_{}'.format(stageDirectory, climatologyName,
                                        mpasMeshName, comparisonFullMeshName)

    make_directories(directory)

    monthValues = sorted(constants.monthDictionary[season])
    startMonth = monthValues[0]
    endMonth = monthValues[-1]

    suffix = '{:04d}{:02d}_{:04d}{:02d}_climo'.format(
        startYear, startMonth, endYear, endMonth)

    if season in constants.abrevMonthNames:
        season = '{:02d}'.format(monthValues[0])
    fileName = '{}/{}_{}_{}.nc'.format(
        directory, ncclimoModel, season, suffix)

    return fileName  # }}}

def get_climatology_op_directory(config, op='avg'):
    '''
    Get the output directory for MPAS climatologies from output with the given
    monthly operator: avg, min or max
    '''
    climatologyBaseDirectory = build_config_full_path(
        config, 'output', 'mpasClimatologySubdirectory')

    return '{}/{}'.format(climatologyBaseDirectory, op)

def _compute_masked_mean(ds, maskVaries):  # {{{
    '''
    Compute the time average of data set, masked out where the variables in ds
    are NaN and, if ``maskVaries == True``, weighting by the number of days
    used to compute each monthly mean time in ds.
    '''
    # Authors
    # -------
    # Xylar Asay-Davis

    def ds_to_weights(ds):
        # make an identical data set to ds but replacing all data arrays with
        # nonnull applied to that data array
        weights = ds.copy(deep=True)
        if isinstance(ds, xr.core.dataarray.DataArray):
            weights = ds.notnull()
        elif isinstance(ds, xr.core.dataset.Dataset):
            for var in ds.data_vars:
                weights[var] = ds[var].notnull()
        else:
            raise TypeError('ds must be an instance of either xarray.Dataset '
                            'or xarray.DataArray.')

        return weights

    if maskVaries:
        dsWeightedSum = (ds * ds.daysInMonth).sum(dim='Time', keep_attrs=True)

        weights = ds_to_weights(ds)

        weightSum = (weights * ds.daysInMonth).sum(dim='Time')

        timeMean = dsWeightedSum / weightSum.where(weightSum > 0.)
    else:
        days = ds.daysInMonth.sum(dim='Time')

        dsWeightedSum = (ds * ds.daysInMonth).sum(dim='Time', keep_attrs=True)

        timeMean = dsWeightedSum / days.where(days > 0.)

    return timeMean  # }}}

def _matches_comparison(obsDescriptor, comparisonDescriptor):  # {{{
    '''
    Determine if the two meshes are the same
    '''
    # Authors
    # -------
    # Xylar Asay-Davis

    if isinstance(obsDescriptor, ProjectionGridDescriptor) and \
            isinstance(comparisonDescriptor, ProjectionGridDescriptor):
        # pretty hard to determine if projections are the same, so we'll rely
        # on the grid names
        match = obsDescriptor.meshName == comparisonDescriptor.meshName and \
            len(obsDescriptor.x) == len(comparisonDescriptor.x) and \
            len(obsDescriptor.y) == len(comparisonDescriptor.y) and \
            numpy.all(numpy.isclose(obsDescriptor.x,
                                    comparisonDescriptor.x)) and \
            numpy.all(numpy.isclose(obsDescriptor.y,
                                    comparisonDescriptor.y))
    elif isinstance(obsDescriptor, LatLonGridDescriptor) and \
            isinstance(comparisonDescriptor, LatLonGridDescriptor):
        match = ((('degree' in obsDescriptor.units and
                   'degree' in comparisonDescriptor.units) or
                  ('radian' in obsDescriptor.units and
                   'radian' in comparisonDescriptor.units)) and
                 len(obsDescriptor.lat) == len(comparisonDescriptor.lat) and
                 len(obsDescriptor.lon) == len(comparisonDescriptor.lon) and
                 numpy.all(numpy.isclose(obsDescriptor.lat,
                                         comparisonDescriptor.lat)) and
                 numpy.all(numpy.isclose(obsDescriptor.lon,
                                         comparisonDescriptor.lon)))
    else:
        match = False

    return match  # }}}

def _setup_climatology_caching(ds, startYearClimo, endYearClimo,
                               yearsPerCacheFile, cachePrefix,
                               monthValues):  # {{{
    '''
    Determine which cache files already exist, which are incomplete and which
    years are present in each cache file (whether existing or to be created).
    '''
    # Authors
    # -------
    # Xylar Asay-Davis

    cacheInfo = []

    cacheIndices = -1 * numpy.ones(ds.dims['Time'], int)
    monthsInDs = ds.month.values
    yearsInDs = ds.year.values

    # figure out which files to load and which years go in each file
    for firstYear in range(startYearClimo, endYearClimo + 1,
                           yearsPerCacheFile):
        years = range(firstYear, firstYear + yearsPerCacheFile)

        yearString, fileSuffix = _get_year_string(years[0], years[-1])
        outputFileClimo = '{}_{}.nc'.format(cachePrefix, fileSuffix)

        done = False
        if os.path.exists(outputFileClimo):
            # already cached
            dsCached = None
            try:
                dsCached = xr.open_dataset(outputFileClimo)
            except IOError:
                # assuming the cache file is corrupt, so deleting it.
                print('Warning: Deleting cache file {}, which appears to '
                      'have been corrupted.'.format(outputFileClimo))

                os.remove(outputFileClimo)

            monthsIfDone = len(monthValues) * len(years)
            if ((dsCached is not None) and
                    (dsCached.attrs['totalMonths'] == monthsIfDone)):
                # also complete, so we can move on
                done = True
            if dsCached is not None:
                dsCached.close()

        cacheIndex = len(cacheInfo)
        for year in years:
            for month in monthValues:
                mask = numpy.logical_and(yearsInDs == year,
                                         monthsInDs == month)
                cacheIndices[mask] = cacheIndex

        if numpy.count_nonzero(cacheIndices == cacheIndex) == 0:
            continue

        cacheInfo.append((outputFileClimo, done, yearString))

    ds = ds.copy()
    ds.coords['cacheIndices'] = ('Time', cacheIndices)

    return cacheInfo, cacheIndices  # }}}

def _cache_individual_climatologies(ds, cacheInfo, printProgress,
                                    yearsPerCacheFile, monthValues,
                                    calendar):  # {{{
    '''
    Cache individual climatologies for later aggregation.
    '''
    # Authors
    # -------
    # Xylar Asay-Davis

    for cacheIndex, info in enumerate(cacheInfo):
        outputFileClimo, done, yearString = info
        if done:
            continue
        dsYear = ds.where(ds.cacheIndices == cacheIndex, drop=True)

        if printProgress:
            print('     {}'.format(yearString))

        totalDays = dsYear.daysInMonth.sum(dim='Time').values

        monthCount = dsYear.dims['Time']

        climatology = compute_climatology(dsYear, monthValues, calendar,
                                          maskVaries=False)

        climatology.attrs['totalDays'] = totalDays
        climatology.attrs['totalMonths'] = monthCount
        climatology.attrs['fingerprintClimo'] = fingerprint_generator()

        write_netcdf(climatology, outputFileClimo)
        climatology.close()

    # }}}

def _cache_aggregated_climatology(startYearClimo, endYearClimo, cachePrefix,
                                  printProgress, monthValues,
                                  cacheInfo):  # {{{
    '''
    Cache aggregated climatology from individual climatologies.
    '''
    # Authors
    # -------
    # Xylar Asay-Davis

    yearString, fileSuffix = _get_year_string(startYearClimo, endYearClimo)
    outputFileClimo = '{}_{}.nc'.format(cachePrefix, fileSuffix)

    done = False
    if len(cacheInfo) == 0:
        climatology = None
        done = True

    if os.path.exists(outputFileClimo):
        # already cached
        climatology = None
        try:
            climatology = xr.open_dataset(outputFileClimo)

        except IOError:
            # assuming the cache file is corrupt, so deleting it.
            print('Warning: Deleting cache file {}, which appears to have '
                  'been corrupted.'.format(outputFileClimo))
            os.remove(outputFileClimo)

        if len(cacheInfo) == 1 and outputFileClimo == cacheInfo[0][0]:
            # theres only one cache file and it already has the same name
            # as the aggregated file so no need to aggregate
            done = True

        elif climatology is not None:
            monthsIfDone = (
                endYearClimo - startYearClimo + 1) * len(monthValues)
            if climatology.attrs['totalMonths'] == monthsIfDone:
                # also complete, so we can move on
                done = True
            else:
                climatology.close()

    if not done:
        if printProgress:
            print('   Computing aggregated climatology '
                  '{}...'.format(yearString))

        first = True
        for cacheIndex, info in enumerate(cacheInfo):
            inFileClimo = info[0]
            ds = xr.open_dataset(inFileClimo)
            days = ds.attrs['totalDays']
            months = ds.attrs['totalMonths']
            if first:
                totalDays = days
                totalMonths = months
                climatology = ds * days
                first = False
            else:
                totalDays += days
                totalMonths += months
                climatology = climatology + ds * days

            ds.close()
        climatology = climatology / totalDays

        climatology.attrs['totalDays'] = totalDays
        climatology.attrs['totalMonths'] = totalMonths
        climatology.attrs['fingerprintClimo'] = fingerprint_generator()

        write_netcdf(climatology, outputFileClimo)

    return climatology  # }}}

def _get_year_string(startYear, endYear):
    if startYear == endYear:
        yearString = '{:04d}'.format(startYear)
        fileSuffix = 'year{}'.format(yearString)
    else:
        yearString = '{:04d}-{:04d}'.format(startYear, endYear)
        fileSuffix = 'years{}'.format(yearString)

    return yearString, fileSuffix

# vim: foldmethod=marker ai ts=4 sts=4 et sw=4 ft=python


"""
Reference: https://github.com/openai/imitation
I follow the architecture from the official repository
"""
import gym
import tensorflow as tf
import numpy as np

from stable_baselines.common.mpi_running_mean_std import RunningMeanStd as MpiRunningMeanStd

from stable_baselines.common.running_mean_std import RunningMeanStd, RunningMinMax
from stable_baselines.common import tf_util as tf_util
from stable_baselines.common import zipsame

def logsigmoid(input_tensor):
    """
    Equivalent to tf.log(tf.sigmoid(a))

    :param input_tensor: (tf.Tensor)
    :return: (tf.Tensor)
    """
    return -tf.nn.softplus(-input_tensor)

def logit_bernoulli_entropy(logits):
    """
    Reference:
    https://github.com/openai/imitation/blob/99fbccf3e060b6e6c739bdf209758620fcdefd3c/policyopt/thutil.py#L48-L51

    :param logits: (tf.Tensor) the logits
    :return: (tf.Tensor) the Bernoulli entropy
    """
    ent = (1. - tf.nn.sigmoid(logits)) * logits - logsigmoid(logits)
    return ent
#
class TabularAdversary(object):
    def __init__(self, observation_space, action_space, hidden_size,
                 entcoeff=0.00, scope="adversary", normalize=True, expert_features=None,
                 exploration_bonus=False, bonus_coef=0.01, t_c=0.1):
        """
        Reward regression from observations and transitions

        :param observation_space: (gym.spaces)
        :param action_space: (gym.spaces)
        :param hidden_size: ([int]) the hidden dimension for the MLP
        :param entcoeff: (float) the entropy loss weight
        :param scope: (str) tensorflow variable scope
        :param normalize: (bool) Whether to normalize the reward or not
        """
        # TODO: support images properly (using a CNN)
        self.scope = scope
        self.observation_shape = observation_space.shape
        self.actions_shape = action_space.shape
        if isinstance(action_space, gym.spaces.Box):
            # Continuous action space
            self.discrete_actions = False
            self.n_actions = action_space.shape[0]
        elif isinstance(action_space, gym.spaces.Discrete):
            self.n_actions = action_space.n
            self.discrete_actions = True
        else:
            raise ValueError('Action space not supported: {}'.format(action_space))

        self.hidden_size = hidden_size
        self.normalize = normalize
        self.obs_rms = None
        self.expert_features = expert_features
        self.reward = expert_features
        normalization = np.linalg.norm(self.reward)
        self.norm_factor = np.sqrt(float(self.observation_shape[0]))
        if normalization > 1:
            self.reward = self.reward / (self.norm_factor * normalization)
        self.exploration_bonus = exploration_bonus
        self.t_c = t_c

        self.bonus_coef = bonus_coef
        if self.exploration_bonus:
            self.covariance_lambda = np.identity(self.observation_shape[0])
        else:
            self.covariance_lambda = None

    def update_reward(self, features):

        t_c = self.t_c
        # self.reward = (1-t_c) * self.reward + t_c * (self.expert_features - features)
        self.reward = self.reward + t_c * (self.expert_features - features) / self.norm_factor
        normalization = np.linalg.norm(self.reward)

        if normalization > 1:
            self.reward = self.reward / normalization

    def get_reward(self, observation):
        """
        Predict the reward using the observation and action

        :param obs: (tf.Tensor or np.ndarray) the observation
        :param actions: (tf.Tensor or np.ndarray) the action
        :return: (np.ndarray) the reward
        """
        if self.exploration_bonus:
            self.covariance_lambda = self.covariance_lambda \
                                + np.matmul(np.expand_dims(observation, axis=1), np.expand_dims(observation, axis=0))
            inverse_covariance = np.linalg.inv(self.covariance_lambda)
            reward = np.matmul(observation, np.reshape(self.reward, (self.reward.shape[0], 1))).squeeze()
            bonus = np.sqrt(np.matmul(np.matmul(observation,inverse_covariance), observation))
            return reward + self.bonus_coef * bonus
        else:
            reward = np.matmul(observation, np.reshape(self.reward, (self.reward.shape[0], 1))).squeeze()
            return reward

class TabularAdversaryTF(object):
    def __init__(self, sess, observation_space, action_space, hidden_size,
                 entcoeff=0.00, scope="adversary", normalize=True, expert_features=None,
                 exploration_bonus=False, is_action_features=True, bonus_coef=0.01, t_c=0.1):
        """
        Reward regression from observations and transitions

        :param observation_space: (gym.spaces)
        :param action_space: (gym.spaces)
        :param hidden_size: ([int]) the hidden dimension for the MLP
        :param entcoeff: (float) the entropy loss weight
        :param scope: (str) tensorflow variable scope
        :param normalize: (bool) Whether to normalize the reward or not
        """
        # TODO: support images properly (using a CNN)
        self.scope = scope
        self.sess = sess
        self.observation_shape = observation_space.shape
        self.actions_shape = action_space.shape
        self.is_action_features = is_action_features

        if isinstance(action_space, gym.spaces.Box):
            # Continuous action space
            self.discrete_actions = False
            self.n_actions = action_space.shape[0]
        elif isinstance(action_space, gym.spaces.Discrete):
            self.n_actions = action_space.n
            self.discrete_actions = True
        else:
            raise ValueError('Action space not supported: {}'.format(action_space))

        if self.is_action_features:
            self.n_features = self.observation_shape[0] + self.n_actions
        else:
            self.n_features = self.observation_shape[0]
        expert_features = expert_features[:self.n_features]

        self.hidden_size = hidden_size
        self.normalize = normalize
        self.obs_rms = None

        self.expert_features = tf.constant(expert_features, dtype=tf.float32)
        self.norm_factor = tf.sqrt(float(self.n_features))

        # self.normalization = tf.square(float(np.linalg.norm(expert_features)))
        self.normalization = np.linalg.norm(expert_features)

        expert_normalization = np.linalg.norm(expert_features)

        # if expert_normalization > 1:
        # self.reward_vec = tf.Variable(expert_features / (self.norm_factor * expert_normalization), dtype=tf.float32)
        self.reward_vec = tf.Variable(expert_features, dtype=tf.float32)

        # else:
        #     self.reward_vec = tf.Variable(expert_features / self.norm_factor)
            # self.reward_vec = tf.Variable(expert_features, dtype=tf.float32)

        # self.reward_vec = tf.Variable(expert_features / self.normalization, dtype=tf.float32)

        # if normalization > 1:
        #     self.reward_vec = tf.Variable(expert_features / normalization, dtype=tf.float32)
        # else:
        #     self.reward_vec = tf.Variable(expert_features, dtype=tf.float32)
        #

        self.exploration_bonus = exploration_bonus
        self.t_c = t_c

        self.bonus_coef = bonus_coef
        if self.exploration_bonus:
            self.covariance_lambda = tf.Variable(tf.eye(self.n_features), dtype=tf.float32)
            self.inverse_covariance = tf.eye(self.n_features)
        else:
            self.covariance_lambda = None
            self.inverse_covariance = None
        # Placeholders
        self.features_ph = tf.placeholder(tf.float32, (None,) + (self.n_features, ),
                                               name="observations_ph")
        self.successor_features_ph = tf.placeholder(tf.float32, (self.n_features, ),
                                               name="successor_features_ph")
        # Build graph
        with tf.variable_scope(self.scope, reuse=False):
            if self.normalize:
                with tf.variable_scope("obfilter"):
                    self.obs_rms = RunningMeanStd(shape=self.n_features)
                    # self.obs_rms = RunningMinMax(shape=self.observation_shape)
                obs_scaled = (tf.cast(self.features_ph, tf.float32) - self.obs_rms.mean)\
                      / tf.cast(tf.sqrt(self.obs_rms.var), tf.float32)
                reward_vec_scaled = (tf.cast(self.reward_vec, tf.float32) - self.obs_rms.mean)\
                             / (tf.cast(tf.sqrt(self.obs_rms.var), tf.float32))
                # obs_scaled = (tf.cast(self.features_ph, tf.float32)) / tf.cast(self.obs_rms.scale, tf.float32)
                obs = obs_scaled

                # reward_vec_scaled = (tf.cast(self.reward_vec, tf.float32)) / tf.cast(self.obs_rms.scale,
                #                                                                              tf.float32)
                reward_vec = reward_vec_scaled / tf.norm(reward_vec_scaled)

            else:
                obs = self.features_ph
                reward_vec = self.reward_vec

            if self.exploration_bonus:
                self.new_covariance_lambda = self.covariance_lambda \
                                         + tf.reduce_sum(
                                            tf.matmul(tf.expand_dims(tf.cast(self.features_ph, tf.float32), axis=2),
                                                         tf.expand_dims(tf.cast(self.features_ph, tf.float32), axis=1)), axis=0)
                self.update_covariance_op = tf.assign(self.covariance_lambda, self.new_covariance_lambda)

                bonus = tf.squeeze(tf.sqrt(tf.matmul(tf.matmul(tf.expand_dims(obs, axis=1), self.inverse_covariance),
                                          tf.expand_dims(obs, axis=2))))

                reward = tf.squeeze(tf.matmul(obs, tf.expand_dims(reward_vec, 1)))
                self.reward_op = reward + self.bonus_coef * bonus
            else:
                self.reward_op = tf.squeeze(tf.matmul(obs, tf.expand_dims(reward_vec, 1)))

            # Update reward
            self.new_reward_vec = self.reward_vec + self.t_c * (self.expert_features - self.successor_features_ph)
            # self.new_reward_vec = self.reward_vec\
            #                       + self.t_c * (self.expert_features - self.successor_features_ph)\
            #                       / (self.normalization * self.norm_factor)
            # normalization = tf.norm(self.new_reward_vec) * self.normalization
            # normalization = tf.norm(self.new_reward_vec)
            # self.new_reward_vec = tf.cond(normalization > 1.0,
            #                                 true_fn=lambda: self.new_reward_vec / normalization,
            #                                 false_fn=lambda: self.new_reward_vec)
            # self.new_reward_vec = self.new_reward_vec / normalization

            # reward_vec_unnormalized = self.reward_vec + self.t_c * (self.expert_features - self.successor_features_ph)
            # reward_vec_scaled = (tf.cast(reward_vec_unnormalized, tf.float32)) / tf.cast(self.obs_rms.scale, tf.float32)
            # self.new_reward_vec = reward_vec_scaled / tf.norm(reward_vec_scaled)
            self.update_reward_op = tf.assign(self.reward_vec, self.new_reward_vec)

    def update_reward(self, successor_features):
        #
        # sess = tf.get_default_session()
        # if len(features.shape) == 1:
            # features = np.expand_dims(features, 0)
        if not self.is_action_features:
            successor_features = successor_features[:self.observation_shape[0]]

        feed_dict = {self.successor_features_ph: successor_features}

        if self.exploration_bonus:
            self.inverse_covariance = tf.linalg.inv(self.covariance_lambda)

        self.sess.run(self.update_reward_op, feed_dict)

    def get_reward(self, obs, action=None):
        """
        Predict the reward using the observation and action

        :param obs: (tf.Tensor or np.ndarray) the observation
        :param actions: (tf.Tensor or np.ndarray) the action
        :return: (np.ndarray) the reward
        """
        # sess = tf.get_default_session()
        if len(obs.shape) == 1:
            obs = np.expand_dims(obs, 0)
        if len(action.shape) == 1:
            action = np.expand_dims(action, 0)

        if self.is_action_features:
            features = np.concatenate((obs, action), axis=1)
        else:
            features = obs

        feed_dict = {self.features_ph: features}

        if self.exploration_bonus:
            reward, _ = self.sess.run([self.reward_op, self.update_covariance_op], feed_dict)
        else:
            reward = self.sess.run(self.reward_op, feed_dict)
        return reward

class NeuralAdversary(object):
    def __init__(self, sess, observation_space, action_space, hidden_size=64, lipschitz_reg_coef=1.0, scope="adversary", normalize=True):
        """
        Reward regression from observations and transitions

        :param observation_space: (gym.spaces)
        :param action_space: (gym.spaces)
        :param hidden_size: ([int]) the hidden dimension for the MLP
        :param entcoeff: (float) the entropy loss weight
        :param scope: (str) tensorflow variable scope
        :param normalize: (bool) Whether to normalize the reward or not
        """
        # TODO: support images properly (using a CNN)
        self.sess = sess
        self.scope = scope
        self.observation_shape = observation_space.shape
        self.actions_shape = action_space.shape

        if isinstance(action_space, gym.spaces.Box):
            # Continuous action space
            self.discrete_actions = False
            self.n_actions = action_space.shape[0]
        elif isinstance(action_space, gym.spaces.Discrete):
            self.n_actions = action_space.n
            self.discrete_actions = True
        else:
            raise ValueError('Action space not supported: {}'.format(action_space))

        self.hidden_size = hidden_size
        self.normalize = normalize
        self.obs_rms = None

        # Placeholders
        self.policy_obs_ph = tf.placeholder(observation_space.dtype, (None,) + self.observation_shape,
                                               name="observations_ph")
        self.policy_acs_ph = tf.placeholder(action_space.dtype, (None,) + self.actions_shape,
                                               name="actions_ph")
        self.policy_gammas_ph = tf.placeholder(tf.float32, (None, 1), name="gammas_ph")
        self.expert_obs_ph = tf.placeholder(observation_space.dtype, (None,) + self.observation_shape,
                                            name="expert_observations_ph")
        self.expert_acs_ph = tf.placeholder(action_space.dtype, (None,) + self.actions_shape,
                                            name="expert_actions_ph")
        self.expert_gammas_ph = tf.placeholder(tf.float32, (None, 1), name="gammas_ph")
        self.mix_obs_ph = tf.placeholder(observation_space.dtype, (None,) + self.observation_shape,
                                            name="expert_observations_ph")
        self.mix_acs_ph = tf.placeholder(action_space.dtype, (None,) + self.actions_shape,
                                            name="expert_actions_ph")

        # Build graph
        policy_rewards = self.build_graph(self.policy_obs_ph, self.policy_acs_ph, reuse=False)
        expert_rewards = self.build_graph(self.expert_obs_ph, self.expert_acs_ph, reuse=True)
        # generator_rewards = tf.math.sigmoid(generator_logits)
        # expert_rewards = tf.math.sigmoid(expert_logits)
        # policy_scaled_rewards = tf.multiply(policy_rewards, self.policy_gammas_ph)
        policy_scaled_rewards = policy_rewards
        # policy_value = (1-0.99) * tf.reduce_sum(policy_scaled_rewards)
        policy_value = tf.reduce_mean(policy_scaled_rewards)
        # expert_scaled_rewards = tf.multiply(expert_rewards, self.expert_gammas_ph)
        expert_scaled_rewards = expert_rewards
        # expert_value = (1-0.99) * tf.reduce_sum(expert_scaled_rewards)
        expert_value = tf.reduce_mean(expert_scaled_rewards)

        # alpha = tf.random.uniform([], 0.0, 1.0, observation_space.dtype)
        # generator_obs_mix = tf.reduce_mean(self.generator_obs_ph, axis=0, keepdims=True)
        # generator_acs_mix = tf.reduce_mean(self.generator_acs_ph, axis=0, keepdims=True)
        # expert_obs_mix = tf.reduce_mean(self.expert_obs_ph, axis=0, keepdims=True)
        # expert_acs_mix = tf.reduce_mean(self.expert_acs_ph, axis=0, keepdims=True)
        # generator_obs_mix = (1-0.99) * tf.reduce_sum(tf.cast(self.generator_gammas_ph, observation_space.dtype) * self.generator_obs_ph, axis=0, keepdims=True)
        # generator_acs_mix = (1-0.99) * tf.reduce_sum(tf.cast(self.generator_gammas_ph, action_space.dtype) * self.generator_acs_ph, axis=0, keepdims=True)
        # expert_obs_mix = (1-0.99) * tf.reduce_sum(tf.cast(self.expert_gammas_ph, observation_space.dtype) * self.expert_obs_ph, axis=0, keepdims=True)
        # expert_acs_mix = (1-0.99) * tf.reduce_sum(tf.cast(self.expert_gammas_ph, action_space.dtype) * self.expert_acs_ph, axis=0, keepdims=True)
        # mixture_obs = alpha * generator_obs_mix + (1 - alpha) * tf.reduce_mean(expert_obs_mix)
        # mixture_acs = tf.cast(alpha, action_space.dtype) * generator_acs_mix\
        #               + tf.cast((1 - alpha), action_space.dtype) * expert_acs_mix
        mixture_rewards = self.build_graph(self.mix_obs_ph, self.mix_acs_ph, reuse=True)
        grads = tf.gradients(mixture_rewards, [self.mix_obs_ph, self.mix_acs_ph])[0]
        norm = tf.cast(tf.sqrt(tf.reduce_sum(tf.square(grads), axis=1)), tf.float32)
        lipschitz_reg = tf.reduce_mean(tf.square(norm - 1.0))
        lipschitz_reg_loss = lipschitz_reg_coef * lipschitz_reg

        rewards = tf.concat([policy_rewards, expert_rewards], 0)

        rewards_reg = - tf.reduce_mean(logit_bernoulli_entropy(rewards))
        rewards_reg_coef = 0.001

        # rewards_reg = tf.reduce_sum(tf.square(rewards))
        # rewards_reg_coef = 0.01

        rewards_reg_loss = rewards_reg_coef * rewards_reg
        policy_loss = policy_value - expert_value
        loss = policy_loss + lipschitz_reg_loss + rewards_reg_loss

        # Loss + Accuracy terms
        self.losses = [loss]
        self.loss_name = ["generator_loss", "expert_loss", "entropy", "entropy_loss", "generator_acc", "expert_acc"]
        # self.total_loss = loss
        # Build Reward for policy
        self.reward_op = tf.clip_by_value(policy_rewards, -10.0, 10.0)
        # self.reward_op = tf.stop_gradient(policy_rewards)
        # self.reward_op = generator_rewards

        var_list = self.get_trainable_variables()
        rewards_optimizer = tf.train.AdamOptimizer(learning_rate=3e-4)
        # rewards_optimizer = tf.train.AdamOptimizer(learning_rate=1e-5)

        # rewards_optimizer = tf.train.AdamOptimizer(learning_rate=3e-4, beta1=0)
        # rewards_optimizer = tf.train.AdamOptimizer(learning_rate=1e-3, beta1=0.5)

        # rewards_optimizer = tf.train.AdamOptimizer(learning_rate=1e-2)

        # grads, vars = zip(*rewards_optimizer.compute_gradients(loss, var_list=var_list))
        # accum_vars = [tf.Variable(tf.zeros_like(var.initialized_value()), trainable=False) for var in var_list]
        # accumulation_counter = tf.Variable(0.0, trainable=False)

        # zero_ops = [var.assign(tf.zeros_like(var)) for var in accum_vars]
        # zero_ops.append(accumulation_counter.assign(0.0))

        # gvs = rewards_optimizer.compute_gradients(loss, var_list)
        # accumulate_ops = [accum_vars[i].assign_add(gv[0]) for i, gv in enumerate(gvs)]
        # accumulate_ops.append(accumulation_counter.assign_add(1.0))

        # train_step = rewards_optimizer.apply_gradients([(accum_vars[i] / accumulation_counter, gv[1]) for i, gv in enumerate(gvs)])

        # grads, vars = list(zip(*grads_and_vars))
        # grads, norm = tf.clip_by_global_norm(grads, 300.0)
        # rewards_train_op = rewards_optimizer.apply_gradients(zip(grads, vars))
        # norm = tf.constant(0.)
        rewards_train_op = rewards_optimizer.minimize(loss, var_list=var_list)

        # rewards_train_op = [rewards_train_op, norm]

        # self.zero_grad = tf_util.function([], zero_ops)
        # self.compute_grads = tf_util.function(
        #     [self.generator_obs_ph, self.generator_acs_ph, self.generator_gammas_ph,
        #      self.expert_obs_ph, self.expert_acs_ph, self.expert_gammas_ph], accumulate_ops)
        # self.train = tf_util.function([], train_step)
        # print_op = tf.print("Value diff:", policy_value - expert_value, "Grad Regularizer:", lipschitz_reg)
        print_op = tf.no_op()
        self.train = tf_util.function(
            [self.policy_obs_ph, self.policy_acs_ph, self.policy_gammas_ph,
             self.expert_obs_ph, self.expert_acs_ph, self.expert_gammas_ph,
             self.mix_obs_ph, self.mix_acs_ph], [rewards_train_op, print_op])

    def build_graph(self, obs_ph, acs_ph, reuse=False):
        """
        build the graph

        :param obs_ph: (tf.Tensor) the observation placeholder
        :param acs_ph: (tf.Tensor) the action placeholder
        :param reuse: (bool)
        :return: (tf.Tensor) the graph output
        """
        with tf.variable_scope(self.scope):
            if reuse:
                tf.get_variable_scope().reuse_variables()

            if self.normalize:
                with tf.variable_scope("obfilter"):
                    self.obs_rms = RunningMeanStd(shape=self.observation_shape)
                obs = (tf.cast(obs_ph, tf.float32) - self.obs_rms.mean) / tf.cast(tf.sqrt(self.obs_rms.var), tf.float32)
            else:
                obs = tf.cast(obs_ph, tf.float32)

            if self.discrete_actions:
                one_hot_actions = tf.one_hot(acs_ph, self.n_actions)
                actions_ph = tf.cast(one_hot_actions, tf.float32)
            else:
                actions_ph = acs_ph

            _input = tf.concat([obs, actions_ph], axis=1)  # concatenate the two input -> form a transition
            p_h1 = tf.contrib.layers.fully_connected(_input, self.hidden_size, activation_fn=tf.nn.tanh)
            p_h2 = tf.contrib.layers.fully_connected(p_h1, self.hidden_size, activation_fn=tf.nn.tanh)
            # rewards = tf.contrib.layers.fully_connected(p_h2, 1, activation_fn=tf.nn.tanh)
            # rewards = tf.contrib.layers.fully_connected(p_h2, 1, activation_fn=tf.math.sigmoid)
            rewards = tf.contrib.layers.fully_connected(p_h2, 1, activation_fn=tf.identity)
        return rewards

    def get_trainable_variables(self):
        """
        Get all the trainable variables from the graph

        :return: ([tf.Tensor]) the variables
        """
        return tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, self.scope)

    def get_reward(self, obs, actions):
        """
        Predict the reward using the observation and action

        :param obs: (tf.Tensor or np.ndarray) the observation
        :param actions: (tf.Tensor or np.ndarray) the action
        :return: (np.ndarray) the reward
        """
        # sess = tf.get_default_session()
        if len(obs.shape) == 1:
            obs = np.expand_dims(obs, 0)
        if len(actions.shape) == 1:
            actions = np.expand_dims(actions, 0)
        elif len(actions.shape) == 0:
            # one discrete action
            actions = np.expand_dims(actions, 0)

        feed_dict = {self.policy_obs_ph: obs, self.policy_acs_ph: actions}
        reward = self.sess.run(self.reward_op, feed_dict)
        return reward

class NeuralAdversaryTRPO(object):
    def __init__(self, sess, observation_space, action_space, hidden_size=64, entcoeff=0.001,
                 scope="adversary", normalize=True):
        """
        Reward regression from observations and transitions

        :param observation_space: (gym.spaces)
        :param action_space: (gym.spaces)
        :param hidden_size: ([int]) the hidden dimension for the MLP
        :param entcoeff: (float) the entropy loss weight
        :param scope: (str) tensorflow variable scope
        :param normalize: (bool) Whether to normalize the reward or not
        """
        # TODO: support images properly (using a CNN)
        self.sess = sess
        self.scope = scope
        self.observation_shape = observation_space.shape
        self.actions_shape = action_space.shape

        if isinstance(action_space, gym.spaces.Box):
            # Continuous action space
            self.discrete_actions = False
            self.n_actions = action_space.shape[0]
        elif isinstance(action_space, gym.spaces.Discrete):
            self.n_actions = action_space.n
            self.discrete_actions = True
        else:
            raise ValueError('Action space not supported: {}'.format(action_space))

        self.hidden_size = hidden_size
        self.normalize = normalize
        self.obs_rms = None

        # Placeholders
        self.policy_obs_ph = tf.placeholder(observation_space.dtype, (None,) + self.observation_shape,
                                               name="observations_ph")
        self.policy_acs_ph = tf.placeholder(action_space.dtype, (None,) + self.actions_shape,
                                               name="actions_ph")
        self.policy_gammas_ph = tf.placeholder(tf.float32, (None, 1), name="gammas_ph")
        self.expert_obs_ph = tf.placeholder(observation_space.dtype, (None,) + self.observation_shape,
                                            name="expert_observations_ph")
        self.expert_acs_ph = tf.placeholder(action_space.dtype, (None,) + self.actions_shape,
                                            name="expert_actions_ph")
        self.expert_gammas_ph = tf.placeholder(tf.float32, (None, 1), name="gammas_ph")
        self.mix_obs_ph = tf.placeholder(observation_space.dtype, (None,) + self.observation_shape,
                                            name="expert_observations_ph")
        self.mix_acs_ph = tf.placeholder(action_space.dtype, (None,) + self.actions_shape,
                                            name="expert_actions_ph")

        # Build graph
        policy_rewards = self.build_graph(self.policy_obs_ph, self.policy_acs_ph, reuse=False)
        expert_rewards = self.build_graph(self.expert_obs_ph, self.expert_acs_ph, reuse=True)
        # policy_rewards = tf.math.sigmoid(policy_logits)
        # expert_rewards = tf.math.sigmoid(expert_logits)
        # policy_scaled_rewards = tf.multiply(policy_rewards, self.policy_gammas_ph)
        policy_scaled_rewards = policy_rewards
        # policy_value = tf.reduce_sum(policy_scaled_rewards)
        policy_value = tf.reduce_mean(policy_scaled_rewards)
        # expert_scaled_rewards = tf.multiply(expert_rewards, self.expert_gammas_ph)
        expert_scaled_rewards = expert_rewards
        # expert_value = tf.reduce_sum(expert_scaled_rewards)
        expert_value = tf.reduce_mean(expert_scaled_rewards)

        # alpha = tf.random.uniform([], 0.0, 1.0, observation_space.dtype)
        # generator_obs_mix = tf.reduce_mean(self.generator_obs_ph, axis=0, keepdims=True)
        # generator_acs_mix = tf.reduce_mean(self.generator_acs_ph, axis=0, keepdims=True)
        # expert_obs_mix = tf.reduce_mean(self.expert_obs_ph, axis=0, keepdims=True)
        # expert_acs_mix = tf.reduce_mean(self.expert_acs_ph, axis=0, keepdims=True)
        # generator_obs_mix = (1-0.99) * tf.reduce_sum(tf.cast(self.generator_gammas_ph, observation_space.dtype) * self.generator_obs_ph, axis=0, keepdims=True)
        # generator_acs_mix = (1-0.99) * tf.reduce_sum(tf.cast(self.generator_gammas_ph, action_space.dtype) * self.generator_acs_ph, axis=0, keepdims=True)
        # expert_obs_mix = (1-0.99) * tf.reduce_sum(tf.cast(self.expert_gammas_ph, observation_space.dtype) * self.expert_obs_ph, axis=0, keepdims=True)
        # expert_acs_mix = (1-0.99) * tf.reduce_sum(tf.cast(self.expert_gammas_ph, action_space.dtype) * self.expert_acs_ph, axis=0, keepdims=True)
        # mixture_obs = alpha * generator_obs_mix + (1 - alpha) * tf.reduce_mean(expert_obs_mix)
        # mixture_acs = tf.cast(alpha, action_space.dtype) * generator_acs_mix\
        #               + tf.cast((1 - alpha), action_space.dtype) * expert_acs_mix
        mixture_rewards = self.build_graph(self.mix_obs_ph, self.mix_acs_ph, reuse=True)
        grads = tf.gradients(mixture_rewards, [self.mix_obs_ph, self.mix_acs_ph])[0]
        norm = tf.cast(tf.sqrt(tf.reduce_sum(tf.square(grads), axis=1)), tf.float32)
        lipschitz_reg = tf.reduce_mean(tf.square(norm - 1.0))
        lipschitz_reg_coef = 0.0
        lipschitz_reg_loss = lipschitz_reg_coef * lipschitz_reg

        rewards = tf.concat([policy_rewards, expert_rewards], 0)

        rewards_reg = - tf.reduce_mean(logit_bernoulli_entropy(rewards))
        # rewards_reg = tf.reduce_sum(tf.square(rewards))

        rewards_reg_coef = 0.0

        rewards_reg_loss = rewards_reg_coef * rewards_reg

        policy_loss = policy_value - expert_value
        self.total_loss = policy_loss + lipschitz_reg_loss + rewards_reg_loss

        # Loss + Accuracy terms
        self.losses = []
        self.loss_name = ["generator_loss", "expert_loss", "entropy", "entropy_loss", "generator_acc", "expert_acc"]
        # self.total_loss = loss
        # Build Reward for policy
        # self.reward_op = tf.stop_gradient(policy_rewards)
        self.reward_op = tf.stop_gradient(tf.clip_by_value(policy_rewards, -1.0, 1.0))
        # self.reward_op = generator_rewards

        print_op = tf.print("Policy loss:", policy_loss,
                            "GradReg", lipschitz_reg,
                            "rewards_abs_mean", tf.reduce_mean(tf.abs(rewards)), "rewards_std", tf.math.reduce_std(rewards),
                            "abs_max", tf.math.reduce_max(tf.abs(rewards)))
        var_list = self.get_trainable_variables()
        self.lossandgrad = tf_util.function(
            [self.policy_obs_ph, self.policy_acs_ph, self.policy_gammas_ph,
             self.expert_obs_ph, self.expert_acs_ph, self.expert_gammas_ph,
             self.mix_obs_ph, self.mix_acs_ph],
            self.losses + [print_op] + [tf_util.flatgrad(self.total_loss, var_list)])

        # print_op = tf.print("Value diff:", policy_value - expert_value, "Grad Regularizer:", lipschitz_reg)
        # print_op = tf.no_op()
        # self.train = tf_util.function(
        #     [self.policy_obs_ph, self.policy_acs_ph, self.policy_gammas_ph,
        #      self.expert_obs_ph, self.expert_acs_ph, self.expert_gammas_ph,
        #      self.mix_obs_ph, self.mix_acs_ph], [rewards_train_op, print_op])

    def build_graph(self, obs_ph, acs_ph, reuse=False):
        """
        build the graph

        :param obs_ph: (tf.Tensor) the observation placeholder
        :param acs_ph: (tf.Tensor) the action placeholder
        :param reuse: (bool)
        :return: (tf.Tensor) the graph output
        """
        with tf.variable_scope(self.scope):
            if reuse:
                tf.get_variable_scope().reuse_variables()

            if self.normalize:
                with tf.variable_scope("obfilter"):
                    self.obs_rms = MpiRunningMeanStd(shape=self.observation_shape)
                obs = (tf.cast(obs_ph, tf.float32) - self.obs_rms.mean) / tf.cast(self.obs_rms.std, tf.float32)
            else:
                obs = tf.cast(obs_ph, tf.float32)

            if self.discrete_actions:
                one_hot_actions = tf.one_hot(acs_ph, self.n_actions)
                actions_ph = tf.cast(one_hot_actions, tf.float32)
            else:
                actions_ph = acs_ph

            _input = tf.concat([obs, actions_ph], axis=1)  # concatenate the two input -> form a transition
            p_h1 = tf.contrib.layers.fully_connected(_input, self.hidden_size, activation_fn=tf.nn.tanh)
            p_h2 = tf.contrib.layers.fully_connected(p_h1, self.hidden_size, activation_fn=tf.nn.tanh)
            # rewards = tf.contrib.layers.fully_connected(p_h2, 1, activation_fn=tf.nn.tanh)
            # rewards = tf.contrib.layers.fully_connected(p_h2, 1, activation_fn=tf.math.sigmoid)
            rewards = tf.contrib.layers.fully_connected(p_h2, 1, activation_fn=tf.identity)
            # last_layer_init = tf.contrib.layers.variance_scaling_initializer(factor=0.1, mode='FAN_AVG', uniform=True)
            # rewards = tf.contrib.layers.fully_connected(p_h2, 1, activation_fn=tf.identity, weights_initializer=last_layer_init)

        return rewards

    def get_trainable_variables(self):
        """
        Get all the trainable variables from the graph

        :return: ([tf.Tensor]) the variables
        """
        return tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, self.scope)

    def get_reward(self, obs, actions):
        """
        Predict the reward using the observation and action

        :param obs: (tf.Tensor or np.ndarray) the observation
        :param actions: (tf.Tensor or np.ndarray) the action
        :return: (np.ndarray) the reward
        """
        # sess = tf.get_default_session()
        if len(obs.shape) == 1:
            obs = np.expand_dims(obs, 0)
        if len(actions.shape) == 1:
            actions = np.expand_dims(actions, 0)
        elif len(actions.shape) == 0:
            # one discrete action
            actions = np.expand_dims(actions, 0)

        feed_dict = {self.policy_obs_ph: obs, self.policy_acs_ph: actions}
        reward = self.sess.run(self.reward_op, feed_dict)
        return reward

class NeuralAdversaryMDPO(object):
    def __init__(self, sess, observation_space, action_space, hidden_size=64, scope="adversary", normalize=True):
        """
        Reward regression from observations and transitions

        :param observation_space: (gym.spaces)
        :param action_space: (gym.spaces)
        :param hidden_size: ([int]) the hidden dimension for the MLP
        :param entcoeff: (float) the entropy loss weight
        :param scope: (str) tensorflow variable scope
        :param normalize: (bool) Whether to normalize the reward or not
        """
        # TODO: support images properly (using a CNN)
        self.sess = sess
        self.scope = scope
        self.observation_shape = observation_space.shape
        self.actions_shape = action_space.shape

        if isinstance(action_space, gym.spaces.Box):
            # Continuous action space
            self.discrete_actions = False
            self.n_actions = action_space.shape[0]
        elif isinstance(action_space, gym.spaces.Discrete):
            self.n_actions = action_space.n
            self.discrete_actions = True
        else:
            raise ValueError('Action space not supported: {}'.format(action_space))

        self.hidden_size = hidden_size
        self.normalize = normalize
        self.obs_rms = None

        # Placeholders
        self.policy_obs_ph = tf.placeholder(observation_space.dtype, (None,) + self.observation_shape,
                                               name="observations_ph")
        self.policy_acs_ph = tf.placeholder(action_space.dtype, (None,) + self.actions_shape,
                                               name="actions_ph")
        self.policy_gammas_ph = tf.placeholder(tf.float32, (None, 1), name="gammas_ph")
        self.expert_obs_ph = tf.placeholder(observation_space.dtype, (None,) + self.observation_shape,
                                            name="expert_observations_ph")
        self.expert_acs_ph = tf.placeholder(action_space.dtype, (None,) + self.actions_shape,
                                            name="expert_actions_ph")
        self.expert_gammas_ph = tf.placeholder(tf.float32, (None, 1), name="gammas_ph")
        self.mix_obs_ph = tf.placeholder(observation_space.dtype, (None,) + self.observation_shape,
                                            name="expert_observations_ph")
        self.mix_acs_ph = tf.placeholder(action_space.dtype, (None,) + self.actions_shape,
                                            name="expert_actions_ph")

        # Build graph
        policy_rewards = self.build_graph(self.policy_obs_ph, self.policy_acs_ph, reuse=False, scope=self.scope)
        expert_rewards = self.build_graph(self.expert_obs_ph, self.expert_acs_ph, reuse=True, scope=self.scope)
        old_policy_rewards = self.build_graph(self.policy_obs_ph, self.policy_acs_ph, reuse=False, scope="oldreward")
        old_expert_rewards = self.build_graph(self.expert_obs_ph, self.expert_acs_ph, reuse=True, scope="oldreward")

        # generator_rewards = tf.math.sigmoid(generator_logits)
        # expert_rewards = tf.math.sigmoid(expert_logits)
        # policy_scaled_rewards = tf.multiply(policy_rewards, self.policy_gammas_ph)
        policy_scaled_rewards = policy_rewards
        # policy_value = (1-0.99) * tf.reduce_sum(policy_scaled_rewards)
        policy_value = tf.reduce_mean(policy_scaled_rewards)
        # expert_scaled_rewards = tf.multiply(expert_rewards, self.expert_gammas_ph)
        expert_scaled_rewards = expert_rewards
        # expert_value = (1-0.99) * tf.reduce_sum(expert_scaled_rewards)
        expert_value = tf.reduce_mean(expert_scaled_rewards)

        # alpha = tf.random.uniform([], 0.0, 1.0, observation_space.dtype)
        # generator_obs_mix = tf.reduce_mean(self.generator_obs_ph, axis=0, keepdims=True)
        # generator_acs_mix = tf.reduce_mean(self.generator_acs_ph, axis=0, keepdims=True)
        # expert_obs_mix = tf.reduce_mean(self.expert_obs_ph, axis=0, keepdims=True)
        # expert_acs_mix = tf.reduce_mean(self.expert_acs_ph, axis=0, keepdims=True)
        # generator_obs_mix = (1-0.99) * tf.reduce_sum(tf.cast(self.generator_gammas_ph, observation_space.dtype) * self.generator_obs_ph, axis=0, keepdims=True)
        # generator_acs_mix = (1-0.99) * tf.reduce_sum(tf.cast(self.generator_gammas_ph, action_space.dtype) * self.generator_acs_ph, axis=0, keepdims=True)
        # expert_obs_mix = (1-0.99) * tf.reduce_sum(tf.cast(self.expert_gammas_ph, observation_space.dtype) * self.expert_obs_ph, axis=0, keepdims=True)
        # expert_acs_mix = (1-0.99) * tf.reduce_sum(tf.cast(self.expert_gammas_ph, action_space.dtype) * self.expert_acs_ph, axis=0, keepdims=True)
        # mixture_obs = alpha * generator_obs_mix + (1 - alpha) * tf.reduce_mean(expert_obs_mix)
        # mixture_acs = tf.cast(alpha, action_space.dtype) * generator_acs_mix\
        #               + tf.cast((1 - alpha), action_space.dtype) * expert_acs_mix
        mixture_rewards = self.build_graph(self.mix_obs_ph, self.mix_acs_ph, reuse=True, scope=self.scope)
        grads = tf.gradients(mixture_rewards, [self.mix_obs_ph, self.mix_acs_ph])[0]
        norm = tf.cast(tf.sqrt(tf.reduce_sum(tf.square(grads), axis=1)), tf.float32)
        lipschitz_reg = tf.reduce_mean(tf.square(norm - 1.0))
        lipschitz_reg_coef = 10
        lipschitz_reg_loss = lipschitz_reg_coef * lipschitz_reg
        #
        # rewards = tf.concat([policy_rewards, expert_rewards], 0)
        #
        # rewards_reg = - tf.reduce_mean(logit_bernoulli_entropy(rewards))
        # rewards_reg_coef = 0.001
        #

        rewards = tf.concat([policy_rewards, expert_rewards], 0)

        policy_clipped_rewards = tf.clip_by_value(policy_rewards, -10.0, 10.0)
        expert_clipped_rewards = tf.clip_by_value(expert_rewards, -10.0, 10.0)
        clipped_rewards = tf.concat([policy_clipped_rewards, expert_clipped_rewards], 0)

        # rewards_reg_loss = rewards_reg_coef * rewards_reg
        old_rewards = tf.concat([old_policy_rewards, old_expert_rewards], 0)

        old_policy_clipped_rewards = tf.clip_by_value(old_policy_rewards, -10.0, 10.0)
        old_expert_clipped_rewards = tf.clip_by_value(old_expert_rewards, -10.0, 10.0)
        old_clipped_rewards = tf.concat([old_policy_clipped_rewards, old_expert_clipped_rewards], 0)

        # rewards_reg_coef = 0.01
        # bregman = tf.reduce_mean(tf_util.huber_loss(old_rewards - rewards))
        bregman = tf.reduce_mean(tf.square(tf.stop_gradient(old_clipped_rewards) - rewards))

        bregman_coeff = 100
        bregman_loss = bregman_coeff * bregman
        #
        # stepsize = 0.001

        rewards_reg = - tf.reduce_mean(logit_bernoulli_entropy(rewards))
        rewards_reg_coeff = 0.00
        # rewards_reg = tf.reduce_mean(tf.square(rewards))
        # rewards_reg_coeff = 1
        rewards_reg_loss = rewards_reg_coeff * rewards_reg

        # rewards_reg_loss = rewards_reg_coef * rewards_reg

        # policy_loss = policy_value - expert_value
        # rewards_gradient = (tf.gradients(tf.reduce_mean(old_policy_clipped_rewards), [old_policy_rewards])[0]
        #                - tf.gradients(tf.reduce_mean(old_expert_clipped_rewards), [old_expert_rewards])[0])
        old_policy_loss = tf.reduce_mean(old_policy_rewards) - tf.reduce_mean(old_expert_rewards)
        old_rewards_gradient = tf.concat(tf.gradients(old_policy_loss, [old_policy_rewards, old_expert_rewards]), axis=0)

        policy_loss = tf.reduce_sum(tf.multiply(tf.stop_gradient(old_rewards_gradient), clipped_rewards))

        loss = policy_loss + bregman_loss + rewards_reg_loss + lipschitz_reg_loss

        # Loss + Accuracy terms
        self.losses = [loss]
        self.loss_name = ["generator_loss", "expert_loss", "entropy", "entropy_loss", "generator_acc", "expert_acc"]
        # self.total_loss = loss
        # Build Reward for policy
        self.reward_op = old_policy_clipped_rewards
        # self.reward_op = tf.stop_gradient(tf.clip_by_value(policy_rewards, -1.0, 1.0))
        # self.reward_op = generator_rewards

        self.update_old_rewards = \
            tf_util.function([], [], updates=[tf.assign(oldv, newv) for (oldv, newv) in
                                              zipsame(tf_util.get_globals_vars("oldreward"),
                                                      tf_util.get_globals_vars(self.scope))])

        var_list = self.get_trainable_variables()
        # rewards_optimizer = tf.train.AdamOptimizer(learning_rate=1e-5, epsilon=1e-5)
        # rewards_optimizer = tf.train.AdamOptimizer(learning_rate=3e-4, beta1=0.9895193, beta2=0.9999, epsilon=1e-5)

        # rewards_optimizer = tf.train.AdamOptimizer(learning_rate=3e-4, beta1=0, beta2=0.9)
        rewards_optimizer = tf.train.AdamOptimizer(learning_rate=3e-4)

        # rewards_optimizer = tf.train.AdamOptimizer(learning_rate=1e-3, beta1=0.5)

        # rewards_optimizer = tf.train.AdamOptimizer(learning_rate=1e-2)

        # accum_vars = [tf.Variable(tf.zeros_like(var.initialized_value()), trainable=False) for var in var_list]
        # accumulation_counter = tf.Variable(0.0, trainable=False)

        # zero_ops = [var.assign(tf.zeros_like(var)) for var in accum_vars]
        # zero_ops.append(accumulation_counter.assign(0.0))

        # gvs = rewards_optimizer.compute_gradients(loss, var_list)
        # accumulate_ops = [accum_vars[i].assign_add(gv[0]) for i, gv in enumerate(gvs)]
        # accumulate_ops.append(accumulation_counter.assign_add(1.0))

        # train_step = rewards_optimizer.apply_gradients([(accum_vars[i] / accumulation_counter, gv[1]) for i, gv in enumerate(gvs)])
        grads, vars = zip(*rewards_optimizer.compute_gradients(loss, var_list=var_list))
        grads, norm = tf.clip_by_global_norm(grads, 1e7)
        # grads, norm = tf.clip_by_global_norm(grads, 5.0)

        rewards_train_op = rewards_optimizer.apply_gradients(zip(grads, vars))
        # norm = tf.constant(0.)
        # rewards_train_op = rewards_optimizer.minimize(loss, var_list=var_list)

        # rewards_train_op = [rewards_train_op, norm]

        # self.zero_grad = tf_util.function([], zero_ops)
        # self.compute_grads = tf_util.function(
        #     [self.generator_obs_ph, self.generator_acs_ph, self.generator_gammas_ph,
        #      self.expert_obs_ph, self.expert_acs_ph, self.expert_gammas_ph], accumulate_ops)
        # self.train = tf_util.function([], train_step)
        linear_approximation = tf.stop_gradient(tf.reduce_sum(tf.multiply(old_rewards_gradient, rewards - old_rewards)))
        # print_op = tf.print("Value diff:", linear_approximation, "Bregman:", bregman_loss,
        #                     "MD objective",  linear_approximation + bregman_loss,
        #                     "Weight GradNorm:", norm, "Loss norm", tf.norm(old_rewards_gradient),
        #                     "rewards_mean", tf.reduce_mean(rewards), "rewards_std", tf.math.reduce_std(rewards),
        #                     "rewards_abs_max", tf.math.reduce_max(rewards))
        # print_op = tf.print("old_rewards_gradient", tf.shape(old_rewards_gradient), "rewards",tf.shape(rewards))
        print_op = tf.no_op()

        # var_list = self.get_trainable_variables()

        # grads, vars = zip(*rewards_optimizer.compute_gradients(loss, var_list=var_list))
        # grads, norm = tf.clip_by_global_norm(grads, 300.0)
        # rewards_train_op = rewards_optimizer.apply_gradients(zip(grads, vars))

        self.train = tf_util.function(
            [self.policy_obs_ph, self.policy_acs_ph, self.policy_gammas_ph,
             self.expert_obs_ph, self.expert_acs_ph, self.expert_gammas_ph,
             self.mix_obs_ph, self.mix_acs_ph], [rewards_train_op, print_op])

    def build_graph(self, obs_ph, acs_ph, reuse=False, scope=None):
        """
        build the graph

        :param obs_ph: (tf.Tensor) the observation placeholder
        :param acs_ph: (tf.Tensor) the action placeholder
        :param reuse: (bool)
        :return: (tf.Tensor) the graph output
        """
        with tf.variable_scope(scope):
            if reuse:
                tf.get_variable_scope().reuse_variables()

            if self.normalize:
                with tf.variable_scope("obfilter"):
                    self.obs_rms = RunningMeanStd(shape=self.observation_shape)
                obs = (tf.cast(obs_ph, tf.float32) - self.obs_rms.mean) / tf.cast(tf.sqrt(self.obs_rms.var), tf.float32)
            else:
                obs = tf.cast(obs_ph, tf.float32)

            if self.discrete_actions:
                one_hot_actions = tf.one_hot(acs_ph, self.n_actions)
                actions_ph = tf.cast(one_hot_actions, tf.float32)
            else:
                actions_ph = acs_ph

                _input = tf.concat([obs, actions_ph], axis=1)  # concatenate the two input -> form a transition
                p_h1 = tf.contrib.layers.fully_connected(_input, self.hidden_size, activation_fn=tf.nn.tanh)
                p_h2 = tf.contrib.layers.fully_connected(p_h1, self.hidden_size, activation_fn=tf.nn.tanh)
                rewards = tf.contrib.layers.fully_connected(p_h2, 1, activation_fn=tf.identity)

        return rewards

    def get_trainable_variables(self):
        """
        Get all the trainable variables from the graph

        :return: ([tf.Tensor]) the variables
        """
        return tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, self.scope)

    def get_reward(self, obs, actions):
        """
        Predict the reward using the observation and action

        :param obs: (tf.Tensor or np.ndarray) the observation
        :param actions: (tf.Tensor or np.ndarray) the action
        :return: (np.ndarray) the reward
        """
        # sess = tf.get_default_session()
        if len(obs.shape) == 1:
            obs = np.expand_dims(obs, 0)
        if len(actions.shape) == 1:
            actions = np.expand_dims(actions, 0)
        elif len(actions.shape) == 0:
            # one discrete action
            actions = np.expand_dims(actions, 0)

        feed_dict = {self.policy_obs_ph: obs, self.policy_acs_ph: actions}
        reward = self.sess.run(self.reward_op, feed_dict)
        return reward

class NeuralAdversaryMD(object):
    def __init__(self, sess, observation_space, action_space, hidden_size=64, entcoeff=0.001, lipschitz_reg_coef=0.0,
                 scope="adversary", normalize=True):
        """
        Reward regression from observations and transitions

        :param observation_space: (gym.spaces)
        :param action_space: (gym.spaces)
        :param hidden_size: ([int]) the hidden dimension for the MLP
        :param entcoeff: (float) the entropy loss weight
        :param scope: (str) tensorflow variable scope
        :param normalize: (bool) Whether to normalize the reward or not
        """
        # TODO: support images properly (using a CNN)
        self.sess = sess
        self.scope = scope
        self.observation_shape = observation_space.shape
        self.actions_shape = action_space.shape

        if isinstance(action_space, gym.spaces.Box):
            # Continuous action space
            self.discrete_actions = False
            self.n_actions = action_space.shape[0]
        elif isinstance(action_space, gym.spaces.Discrete):
            self.n_actions = action_space.n
            self.discrete_actions = True
        else:
            raise ValueError('Action space not supported: {}'.format(action_space))

        self.hidden_size = hidden_size
        self.normalize = normalize
        self.obs_rms = None

        # Placeholders
        self.policy_obs_ph = tf.placeholder(observation_space.dtype, (None,) + self.observation_shape,
                                               name="observations_ph")
        self.policy_acs_ph = tf.placeholder(action_space.dtype, (None,) + self.actions_shape,
                                               name="actions_ph")
        self.policy_gammas_ph = tf.placeholder(tf.float32, (None, 1), name="gammas_ph")
        self.expert_obs_ph = tf.placeholder(observation_space.dtype, (None,) + self.observation_shape,
                                            name="expert_observations_ph")
        self.expert_acs_ph = tf.placeholder(action_space.dtype, (None,) + self.actions_shape,
                                            name="expert_actions_ph")
        self.expert_gammas_ph = tf.placeholder(tf.float32, (None, 1), name="expert_gammas_ph")
        self.mix_obs_ph = tf.placeholder(observation_space.dtype, (None,) + self.observation_shape,
                                            name="expert_observations_ph")
        self.mix_acs_ph = tf.placeholder(action_space.dtype, (None,) + self.actions_shape,
                                            name="expert_actions_ph")

        if self.normalize:
            with tf.variable_scope("obfilter"):
                self.obs_rms = MpiRunningMeanStd(shape=self.observation_shape)

        # Build graph
        policy_rewards = self.build_graph(self.policy_obs_ph, self.policy_acs_ph, reuse=False, scope=self.scope)
        expert_rewards = self.build_graph(self.expert_obs_ph, self.expert_acs_ph, reuse=True, scope=self.scope)
        old_policy_rewards = self.build_graph(self.policy_obs_ph, self.policy_acs_ph, reuse=False, scope="oldreward")
        old_expert_rewards = self.build_graph(self.expert_obs_ph, self.expert_acs_ph, reuse=True, scope="oldreward")

        # policy_scaled_rewards = tf.multiply(policy_rewards, self.policy_gammas_ph)
        policy_scaled_rewards = policy_rewards
        # policy_value = tf.reduce_sum(policy_scaled_rewards)
        policy_value = tf.reduce_mean(policy_scaled_rewards)
        # expert_scaled_rewards = tf.multiply(expert_rewards, self.expert_gammas_ph)
        expert_scaled_rewards = expert_rewards
        # expert_value = tf.reduce_sum(expert_scaled_rewards)
        expert_value = tf.reduce_mean(expert_scaled_rewards)

        mixture_rewards = self.build_graph(self.mix_obs_ph, self.mix_acs_ph, reuse=True, scope=self.scope)
        grads = tf.gradients(mixture_rewards, [self.mix_obs_ph, self.mix_acs_ph])[0]
        norm = tf.cast(tf.sqrt(tf.reduce_sum(tf.square(grads), axis=1)), tf.float32)
        lipschitz_reg = tf.reduce_mean(tf.square(norm - 1.0))
        lipschitz_loss = lipschitz_reg_coef * lipschitz_reg
        #

        rewards = tf.concat([policy_rewards, expert_rewards], 0)

        policy_clipped_rewards = tf.clip_by_value(policy_rewards, -10.0, 10.0)
        expert_clipped_rewards = tf.clip_by_value(expert_rewards, -10.0, 10.0)
        clipped_rewards = tf.concat([policy_clipped_rewards, expert_clipped_rewards], 0)

        #
        # rewards_reg = - tf.reduce_mean(logit_bernoulli_entropy(rewards))
        # rewards_reg_coef = 0.001
        #
        # rewards_reg = tf.reduce_sum(tf.square(rewards))

        # rewards_reg_coef = 0.01

        old_rewards = tf.concat([old_policy_rewards, old_expert_rewards], 0)

        old_policy_clipped_rewards = tf.clip_by_value(old_policy_rewards, -10.0, 10.0)
        old_expert_clipped_rewards = tf.clip_by_value(old_expert_rewards, -10.0, 10.0)
        old_clipped_rewards = tf.concat([old_policy_clipped_rewards, old_expert_clipped_rewards], 0)

        # bregman = tf.reduce_mean(tf_util.huber_loss(old_clipped_rewards - rewards))
        bregman = tf.reduce_mean(tf.square(tf.stop_gradient(old_clipped_rewards) - rewards))

        bregman_coeff = 100
        bregman_loss = bregman_coeff * bregman

        #
        # stepsize = 0.001

        # old_policy_loss = tf.reduce_mean(tf.multiply(old_policy_rewards, self.policy_gammas_ph))\
        #                   - tf.reduce_mean(tf.multiply(old_expert_rewards, self.expert_gammas_ph))

        new_policy_loss = tf.reduce_mean(tf.multiply(policy_rewards, self.policy_gammas_ph))\
                          - tf.reduce_mean(tf.multiply(expert_rewards, self.expert_gammas_ph))

        old_policy_loss = tf.reduce_mean(old_policy_rewards) - tf.reduce_mean(old_expert_rewards)
        old_rewards_gradient = tf.concat(tf.gradients(old_policy_loss, [old_policy_rewards, old_expert_rewards]), axis=0)

        policy_loss = tf.reduce_sum(tf.multiply(tf.stop_gradient(old_rewards_gradient), clipped_rewards))

        rewards_reg = - tf.reduce_mean(logit_bernoulli_entropy(rewards))
        rewards_reg_coeff = 0.001
        # rewards_reg = tf.reduce_mean(tf.square(rewards))
        # rewards_reg = tf.reduce_mean(tf_util.huber_loss(rewards))
        # rewards_reg_coeff = 0
        rewards_reg_loss = rewards_reg_coeff * rewards_reg

        # rewards_reg_loss = rewards_reg_coef * rewards_reg
        # policy_loss = policy_value - exp  ert_value
        self.total_loss = policy_loss + bregman_loss + rewards_reg_loss + lipschitz_loss

        # Loss + Accuracy terms
        self.losses = []
        self.loss_name = ["generator_loss", "expert_loss", "entropy", "entropy_loss", "generator_acc", "expert_acc"]
        # Build Reward for policy
        self.reward_op = tf.stop_gradient(old_policy_clipped_rewards)
        # self.reward_op = old_policy_rewards

        # self.reward_op = tf.clip_by_value(policy_rewards, -1.0, 1.0)
        # self.reward_op = generator_rewards

        self.update_old_rewards = \
            tf_util.function([], [], updates=[tf.assign(oldv, newv) for (oldv, newv) in
                                              zipsame(tf_util.get_globals_vars("oldreward"),
                                                      tf_util.get_globals_vars(self.scope))])

        var_list = self.get_trainable_variables()

        # clip_weights = [tf.assign(var, tf.clip_by_value(var,  -0.5, 0.5)) for var in var_list]
        clip_weights = tf.no_op()
        self.clip_weights = tf_util.function([], [clip_weights])

        # rewards_optimizer = tf.train.AdamOptimizer(learning_rate=3e-5)
        # grads, vars = zip(*rewards_optimizer.compute_gradients(self.total_loss, var_list=var_list))
        # rewards_train_op = rewards_optimizer.apply_gradients(zip(grads, vars))

        linear_approximation = tf.stop_gradient(tf.reduce_sum(tf.multiply(old_rewards_gradient, rewards - old_rewards)))

        grads = tf.gradients(self.total_loss, var_list)
        grads, norm = tf.clip_by_global_norm(grads, 1e7)

        # self.train = tf_util.function(
        #     [self.policy_obs_ph, self.policy_acs_ph, self.policy_gammas_ph,
        #      self.expert_obs_ph, self.expert_acs_ph, self.expert_gammas_ph,
        #      self.mix_obs_ph, self.mix_acs_ph], [rewards_train_op, print_op])

        loss_op = tf.concat(axis=0, values=[tf.reshape(grad if grad is not None else tf.zeros_like(v), [tf_util.numel(v)])
                                            for (v, grad) in zip(var_list, grads)])

        print_op = tf.print("Policy Loss:", new_policy_loss, "Bregman:", bregman_loss,
                            "MD objective",  linear_approximation + bregman_loss,
                            "GradNorm", norm,
                            "Loss norm", tf.norm(old_rewards_gradient),
                            "averageEnt", rewards_reg_loss,
                            "mean", tf.reduce_mean(tf.abs(rewards)), "std", tf.math.reduce_std(rewards),
                            "max", tf.math.reduce_max(tf.abs(rewards)))

        self.lossandgrad = tf_util.function(
            [self.policy_obs_ph, self.policy_acs_ph, self.policy_gammas_ph,
             self.expert_obs_ph, self.expert_acs_ph, self.expert_gammas_ph,
             self.mix_obs_ph, self.mix_acs_ph],
             # self.losses + [tf_util.flatgrad(self.total_loss, var_list)]) #, clip_norm=0.5, clip_by_global_norm=True)])
             self.losses + [print_op] + [loss_op])

        # print_op = tf.print("Value diff:", policy_value - expert_value, "Grad Regularizer:", lipschitz_reg)
        # print_op = tf.no_op()
        # self.train = tf_util.function(
        #     [self.policy_obs_ph, self.policy_acs_ph, self.policy_gammas_ph,
        #      self.expert_obs_ph, self.expert_acs_ph, self.expert_gammas_ph,
        #      s, [rewards_train_op, print_op])

    def build_graph(self, obs_ph, acs_ph, reuse=False, scope=None):
        """
        build the graph

        :param obs_ph: (tf.Tensor) the observation placeholder
        :param acs_ph: (tf.Tensor) the action placeholder
        :param reuse: (bool)
        :return: (tf.Tensor) the graph output
        """
        with tf.variable_scope(scope):
            if reuse:
                tf.get_variable_scope().reuse_variables()

            if self.normalize:
                obs = (tf.cast(obs_ph, tf.float32) - self.obs_rms.mean) / tf.cast(self.obs_rms.std, tf.float32)
            else:
                obs = tf.cast(obs_ph, tf.float32)

            if self.discrete_actions:
                one_hot_actions = tf.one_hot(acs_ph, self.n_actions)
                actions_ph = tf.cast(one_hot_actions, tf.float32)
            else:
                actions_ph = acs_ph

                _input = tf.concat([obs, actions_ph], axis=1)  # concatenate the two input -> form a transition
                p_h1 = tf.contrib.layers.fully_connected(_input, self.hidden_size, activation_fn=tf.nn.tanh)
                p_h2 = tf.contrib.layers.fully_connected(p_h1, self.hidden_size, activation_fn=tf.nn.tanh)
                # rewards = tf.contrib.layers.fully_connected(p_h2, 1, activation_fn=tf.nn.tanh)
                # rewards = tf.contrib.layers.fully_connected(p_h2, 1, activation_fn=tf.math.sigmoid)
                rewards = tf.contrib.layers.fully_connected(p_h2, 1, activation_fn=tf.identity)
                # last_layer_init = tf.contrib.layers.variance_scaling_initializer(factor=0.01, mode='FAN_AVG', uniform=True)
                # rewards = tf.contrib.layers.fully_connected(p_h2, 1, activation_fn=tf.identity, weights_initializer=last_layer_init)

        return rewards

    def get_trainable_variables(self):
        """
        Get all the trainable variables from the graph

        :return: ([tf.Tensor]) the variables
        """
        return tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, self.scope)# + tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, "oldreward")

    def get_reward(self, obs, actions):
        """
        Predict the reward using the observation and action

        :param obs: (tf.Tensor or np.ndarray) the observation
        :param actions: (tf.Tensor or np.ndarray) the action
        :return: (np.ndarray) the reward
        """
        # sess = tf.get_default_session()
        if len(obs.shape) == 1:
            obs = np.expand_dims(obs, 0)
        if len(actions.shape) == 1:
            actions = np.expand_dims(actions, 0)
        elif len(actions.shape) == 0:
            # one discrete action
            actions = np.expand_dims(actions, 0)

        feed_dict = {self.policy_obs_ph: obs, self.policy_acs_ph: actions}
        reward = self.sess.run(self.reward_op, feed_dict)
        return reward