# -*- coding: utf-8 -*-
'''Chemical Engineering Design Library (ChEDL). Utilities for process modeling.
Copyright (C) 2016, 2017 Caleb Bell <Caleb.Andrew.Bell@gmail.com>

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.'''

from __future__ import division
try: # pragma: no cover
    from cStringIO import StringIO
except: # pragma: no cover
    from io import BytesIO as StringIO
import os
import gzip
import datetime
from calendar import isleap
from collections import namedtuple

import numpy as np
from fluids.core import F2K
from scipy.constants import mile, knot, inch
from scipy.spatial import KDTree, cKDTree
from scipy.stats import scoreatpercentile

try: # pragma: no cover
    from urllib.request import urlopen
    from urllib.error import HTTPError
except ImportError: # pragma: no cover
    from urllib2 import urlopen
    from urllib2 import HTTPError

try:  # pragma: no cover
    from appdirs import user_data_dir, user_config_dir
    data_dir = user_config_dir('fluids')
except ImportError:  # pragma: no cover
    data_dir = ''
    pass
# TODO: Import ephem and get hours/minutes of sunlight per day.

__all__ = ['get_clean_isd_history', 'IntegratedSurfaceDatabaseStation',
           'get_closest_station', 'get_station_year_text', 'gsod_day_parser',
           'StationDataGSOD', 'heating_degree_days', 'cooling_degree_days', 'stations']

folder = os.path.join(os.path.dirname(__file__), 'data')

def heating_degree_days(T, T_base=F2K(65), truncate=True):
    r'''Calculates the heating degree days for a period of time.

    .. math::
        \text{heating degree days} = max(T - T_{base}, 0)

    Parameters
    ----------
    T : float
        Measured temperature; sometimes an average over a length of time is used,
        other times the average of the lowest and highest temperature in a
        period are used, [K]
    T_base : float, optional
        Reference temperature for the degree day calculation, defaults
        to 65 °F (18.33 °C, 291.483 K), the value most used in the US, [K]
    truncate : bool
        If truncate is True, no negative values will be returned; if negative,
        the value is truncated to 0, [-]

    Returns
    -------
    heating_degree_days : float
        Degree above the base temperature multiplied by the length of time of
        the measurement, normally days [day*K]

    Notes
    -----
    Some common base temperatures are 18 °C (Canada), 15.5 °C (EU),
    17 °C (Denmark, Finland), 12 °C Switzerland. The base temperature
    should always be presented with the results.

    The time unit does not have to be days; it can be any time unit, and the
    calculation behaves the same.

    Examples
    --------
    >>> heating_degree_days(303.8)
    12.31666666666672

    >>> heating_degree_days(273)
    0.0

    >>> heating_degree_days(322, T_base=300)
    22

    References
    ----------
    .. [1] "Heating Degree Day." Wikipedia, January 24, 2018.
       https://en.wikipedia.org/w/index.php?title=Heating_degree_day&oldid=822187764.
    '''
    dd = T - T_base
    if truncate and dd < 0.0:
        dd = 0.0
    return dd

def cooling_degree_days(T, T_base=283.15, truncate=True):
    r'''Calculates the cooling degree days for a period of time.

    .. math::
        \text{cooling degree days} = max(T_{base} - T, 0)

    Parameters
    ----------
    T : float
        Measured temperature; sometimes an average over a length of time is used,
        other times the average of the lowest and highest temperature in a
        period are used, [K]
    T_base : float, optional
        Reference temperature for the degree day calculation, defaults
        to 10 °C, 283.15 K, a common value, [K]
    truncate : bool
        If truncate is True, no negative values will be returned; if negative,
        the value is truncated to 0, [-]

    Returns
    -------
    cooling_degree_days : float
        Degree below the base temperature multiplied by the length of time of
        the measurement, normally days [day*K]

    Notes
    -----
    The base temperature should always be presented with the results.

    The time unit does not have to be days; it can be time unit, and the
    calculation behaves the same.

    Examples
    --------
    >>> cooling_degree_days(250)
    33.14999999999998

    >>> cooling_degree_days(300)
    0.0

    >>> cooling_degree_days(250, T_base=300)
    50

    References
    ----------
    .. [1] "Heating Degree Day." Wikipedia, January 24, 2018.
       https://en.wikipedia.org/w/index.php?title=Heating_degree_day&oldid=822187764.
    '''
    dd = T_base - T
    if truncate and dd < 0.0:
        dd = 0.0
    return dd

def get_clean_isd_history(dest=os.path.join(folder, 'isd-history-cleaned.tsv'),
                          url="ftp://ftp.ncdc.noaa.gov/pub/data/noaa/isd-history.csv"): # pragma: no cover
    '''Basic method to update the isd-history file from the NOAA. This is
    useful as new weather stations are updated all the time.

    This function requires pandas to run. If fluids is installed for the
    superuser, this method must be called in an instance of Python running
    as the superuser (administrator).

    Retrieving the file from ftp typically takes several seconds.
    Pandas reads the file in ~30 ms and writes it in ~220 ms. Reading it with
    the code below takes ~220 ms but is necessary to prevent a pandas
    dependency.

    Parameters
    ----------
    dest : str, optional
        The file to store the data retrieved; leave as the default argument
        for it to be accessible by fluids.
    url : str, optional
        The location of the data file; this can be anywhere that can be read
        by pandas, including a local file as would be useful in an offline
        situation.
    '''
    import pandas as pd
    df = pd.read_csv(url)
    df.to_csv(dest, sep='\t', index=False, header=False)

class IntegratedSurfaceDatabaseStation(object):
    '''Class to hold data on a weather station in the Integrated Surface
    Database.

    License information for the database can be found at the following link:
    https://data.noaa.gov/dataset/global-surface-summary-of-the-day-gsod

    Parameters
    ----------
    USAF : int or None if unassigned
        Air Force station ID. May contain a letter in the first position.
    WBAN : int or None if unassigned
        NCDC WBAN number
    NAME : str
        Name of the station; ex. 'CENTRAL COLORADO REGIONAL AP'
    CTRY : str or None if unspecified
        FIPS country ID
    ST : str or None if not in the US
        State for US stations
    ICAO : str or None if not an airport
        ICAO airport code
    LAT : float
        Latitude with a precision of one thousandths of a decimal degree,
        [degrees]
    LON : float
        Longitude with a precision of one thousandths of a decimal degree,
        [degrees]
    ELEV : float
        Elevation of weather station, [m]
    BEGIN : float
        Beginning Period Of Record (YYYYMMDD). There may be reporting gaps
        within the P.O.R.
    END : Ending Period Of Record (YYYYMMDD). There may be reporting gaps
        within the P.O.R.
    '''
    __slots__ = ['USAF', 'WBAN', 'NAME', 'CTRY', 'ST', 'ICAO', 'LAT', 'LON',
                 'ELEV', 'BEGIN', 'END', 'raw_data', 'parsed_data']

    def __repr__(self):
        s = ('<Weather station registered in the Integrated Surface Database, '
            'name %s, country %s, USAF %s, WBAN %s, coords (%s, %s) '
            'Weather data from %s to %s>' )
        return s%(self.NAME, self.CTRY, self.USAF, self.WBAN, self.LAT, self.LON, str(self.BEGIN)[0:4], str(self.END)[0:4])

    def __init__(self, USAF, WBAN, NAME, CTRY, ST, ICAO, LAT, LON, ELEV, BEGIN,
                 END):
        try:
            self.USAF = int(USAF)
        except TypeError:
            self.USAF = USAF # Nones
        try:
            self.WBAN = int(WBAN)
        except TypeError:
            self.WBAN = WBAN
        self.NAME = NAME
        self.CTRY = CTRY
        self.ST = ST
        self.ICAO = ICAO
        self.LAT = LAT
        self.LON = LON
        self.ELEV = ELEV
        self.BEGIN = int(BEGIN)
        self.END = int(END)

class StationDataGSOD(object):
    # Holds data, caches and retrieves data
    def __init__(self, station):
        self.station = station
        self.begin = datetime.datetime.strptime(str(self.station.BEGIN), '%Y%m%d')
        self.end = datetime.datetime.strptime(str(self.station.END), '%Y%m%d')

        self.year_range = range(self.begin.year, self.end.year + 1)

#         Would be nice to create these later, when using a download_data method
        self.raw_text = {}
        self.raw_data = {}
        self.parsed_data = {}
        self.load_empty_vectors()
        self.download_data()
        self.parse_data()

    def load_empty_vectors(self):
        for year in self.year_range:
            days_in_year = 366 if isleap(year) else 365
            self.raw_data[year] = [None]*days_in_year
            self.parsed_data[year] = [None]*days_in_year
            self.raw_text[year] = None
#        days = [None]*days_in_year(y)

    def download_data(self):
        for year in self.year_range:
            if self.raw_text[year] is None:
                try:
                    year_data = get_station_year_text(self.station.USAF, self.station.WBAN, year)
                    self.raw_text[year] = year_data
                except:
                    pass

    def parse_data(self):
        for year, data in self.raw_text.items():
            if data is not None:
                days = self.parsed_data[year]
                for line in data.split('\n')[1:-1]:
                    parsed = gsod_day_parser(line)
                    doy = parsed.DATE.timetuple().tm_yday-1
                    days[doy] = parsed

    def coldest_month(self, older_year=None, newer_year=None, minimum_days=23):
        # Tested
        month_data = self.month_average_temperature(older_year=older_year,
                                                    newer_year=newer_year,
                                                    minimum_days=minimum_days)
        return month_data.index(min(month_data))

    def warmest_month(self, older_year=None, newer_year=None, minimum_days=23):
        # Tested
        month_data = self.month_average_temperature(older_year=older_year,
                                                    newer_year=newer_year,
                                                    minimum_days=minimum_days)
        return month_data.index(max(month_data))

    def month_average_temperature(self, older_year=None, newer_year=None,
                                  include_yearly=False, minimum_days=23):
        '''
        >> station = get_closest_station(38.8572, -77.0369)
        >> station_data = StationDataGSOD(station)
        >> station_data.month_average_temperature(1990, 2000, include_yearly=False)
        [276.1599380905833, 277.5375516246206, 281.1881231671554, 286.7367003367004, 291.8689638318671, 296.79545454545456, 299.51868686868687, 298.2097914630174, 294.4116161616162, 288.25883023786247, 282.3188552188553, 277.8282339524275]
        '''
        # Take years, make them inclusive; add minimum valid days.
        year_month_averages = {}
        year_month_counts = {}

        for year, data in self.parsed_data.items():
            if not (older_year <= year <= newer_year):
                continue # Ignore out-of-range years easily
            year_month_averages[year] = [0.0]*12
            year_month_counts[year] = [0]*12

            for i, day in enumerate(data):
                if day is None:
                    continue
                # Don't do these comparisions to make it fast
                if day.DATE.year < older_year or day.DATE.year > newer_year:
                    continue # Ignore out-of-range days as possible

                T = day.TEMP
                if T is None:
                    continue
                # Cache these lookups
                year_month_averages[year][day.DATE.month-1] += T
                year_month_counts[year][day.DATE.month-1] += 1

            for month in range(12):
                count = year_month_counts[year][month]
                if count < minimum_days:
                    ans = None
                else:
                    ans = year_month_averages[year][month]/count
                year_month_averages[year][month] = ans

        # Compute the average of the month
        actual_averages = [0.0]*12
        actual_averages_counts = [0]*12
        for year, average in year_month_averages.items():
            for month in range(12):
                if average is not None and average[month] is not None:
                    count = actual_averages_counts[month]
                    if count is None:
                        count = 1
                    else:
                        count += 1
                    actual_averages_counts[month] = count
                    month_average_sum = actual_averages[month]
                    if month_average_sum is None:
                        month_average_sum = average[month]
                    else:
                        month_average_sum += average[month]
                    actual_averages[month] = month_average_sum

        for month in range(12):
            actual_averages[month] = actual_averages[month]/actual_averages_counts[month]

        # Don't set anything as properties - too many variables used in calculating thems
        # Speed is not that important.
        if include_yearly:
            return actual_averages, year_month_averages
        else:
            return actual_averages

    # Copy and paste
    def month_average_windspeed(self, older_year=None, newer_year=None,
                                  include_yearly=False, minimum_days=23):
        # Take years, make them inclusive; add minimum valid days.
        year_month_averages = {}
        year_month_counts = {}

        for year, data in self.parsed_data.items():
            if not (older_year <= year <= newer_year):
                continue # Ignore out-of-range years easily
            year_month_averages[year] = [0.0]*12
            year_month_counts[year] = [0]*12

            for i, day in enumerate(data):
                if day is None:
                    continue
                # Don't do these comparisions to make it fast
                if day.DATE.year < older_year or day.DATE.year > newer_year:
                    continue # Ignore out-of-range days as possible

                wind_speed = day.WDSP
                if wind_speed is None:
                    continue
                # Cache these lookups
                year_month_averages[year][day.DATE.month-1] += wind_speed
                year_month_counts[year][day.DATE.month-1] += 1

            for month in range(12):
                count = year_month_counts[year][month]
                if count < minimum_days:
                    ans = None
                else:
                    ans = year_month_averages[year][month]/count
                year_month_averages[year][month] = ans

        # Compute the average of the month
        actual_averages = [0.0]*12
        actual_averages_counts = [0]*12
        for year, average in year_month_averages.items():
            for month in range(12):
                if average is not None and average[month] is not None:
                    count = actual_averages_counts[month]
                    if count is None:
                        count = 1
                    else:
                        count += 1
                    actual_averages_counts[month] = count
                    month_average_sum = actual_averages[month]
                    if month_average_sum is None:
                        month_average_sum = average[month]
                    else:
                        month_average_sum += average[month]
                    actual_averages[month] = month_average_sum

        for month in range(12):
            actual_averages[month] = actual_averages[month]/actual_averages_counts[month]

        # Don't set anything as properties - too many variables used in calculating thems
        # Speed is not that important.
        if include_yearly:
            return actual_averages, year_month_averages
        else:
            return actual_averages

    def percentile_extreme_condition(self, older_year=None, newer_year=None,
                                  include_yearly=False, minimum_days=23, attr='WDSP'):
        # Really need to normalize data with interpolation etc here.
        # Need to get the data, and process it and score interpolation regimes.
        # Or could just randomly drop data and try to fill it in.
        accepted_values = []
        for year, data in self.parsed_data.items():
            if not (older_year <= year <= newer_year):
                continue # Ignore out-of-range years easily

stations = []
_latlongs = []
'''Read in the parsed data into
1) a list of latitudes and longitudes, temporary, which will get converted to
a numpy array for use in KDTree
2) a list of IntegratedSurfaceDatabaseStation objects; the query will return
the index of the nearest weather stations.
'''
with open(os.path.join(folder, 'isd-history-cleaned.tsv')) as f:
    for line in f:
        values = line.split('\t')
        for i in range(0, 11):
            v = values[i]
            if not v:
                values[i] = None # '' case
            else:
                try:
                    values[i] = float(v)
                    if int(v) == 99999:
                        values[i] = None
                except:
                    continue
        lat, lon = values[6], values[7]
        if lat and lon:
            # Some stations have no lat-long; this isn't useful
            stations.append(IntegratedSurfaceDatabaseStation(*values))
            _latlongs.append((lat, lon))
_latlongs = np.array(_latlongs)
station_count = len(stations)

kd_tree = cKDTree(_latlongs) # _latlongs must be unchanged as data is not copied

def get_closest_station(latitude, longitude, minumum_recent_data=20140000,
                        match_max=100):
    '''Query function to find the nearest weather station to a particular
    set of coordinates. Optionally allows for a recent date by which the
    station is required to be still active at.

    Parameters
    ----------
    latitude : float
        Latitude to search for nearby weather stations at, [degrees]
    longitude : float
        Longitude to search for nearby weather stations at, [degrees]
    minumum_recent_data : int, optional
        Date that the weather station is required to have more recent
        weather data than; format YYYYMMDD; set this to 0 to not restrict data
        by date.
    match_max : int, optional
        The maximum number of results in the KDTree to search for before
        applying the filtering criteria; an internal parameter which is
        increased automatically if the default value is insufficient [-]

    Returns
    -------
    station : IntegratedSurfaceDatabaseStation
        Instance of IntegratedSurfaceDatabaseStation which was nearest
        to the requested coordinates and with sufficiently recent data
        available [-]

    Notes
    -----
    Searching for 100 stations is a reasonable choice as it takes, ~70
    microseconds vs 50 microsecond to find only 1 station. The search does get
    slower as more points are requested. Bad data is returned from a KDTree
    search if more points are requested than are available.

    Examples
    --------
    >>> get_closest_station(51.02532675, -114.049868485806, 20150000)
    <Weather station registered in the Integrated Surface Database, name CALGARY INTL CS, country CA, USAF 713930, WBAN None, coords (51.1, -114.0) Weather data from 2004 to 2017>
    '''
    # Both station strings may be important
    # Searching for 100 stations is fine, 70 microseconds vs 50 microsecond for 1
    # but there's little point for more points, it gets slower.
    # bad data is returned if k > station_count
    distances, indexes = kd_tree.query([latitude, longitude], k=min(match_max, station_count))
    #
    for i in indexes:
        latlon = _latlongs[i]
        enddate = stations[i].END
        # Iterate for all indexes until one is found whose date is current
        if enddate > minumum_recent_data:
            return stations[i]
    if match_max < station_count:
        return get_closest_station(latitude, longitude, minumum_recent_data=minumum_recent_data, match_max=match_max*10)
    raise Exception('Could not find a station with more recent data than '
                    'specified near the specified coordinates.')

# This should be agressively cached
def get_station_year_text(WMO, WBAN, year):
    '''Basic method to download data from the GSOD database, given a
    station idenfifier and year.

    Parameters
    ----------
    WMO : int or None
         World Meteorological Organization (WMO) identifiers, [-]
    WBAN : int or None
        Weather Bureau Army Navy (WBAN) weather station identifier, [-]
    year : int
        Year data should be retrieved from, [year]

    Returns
    -------
    data : str
        Downloaded data file
    '''
    if WMO is None:
        WMO = 999999
    if WBAN is None:
        WBAN = 99999
    station = str(int(WMO)) + '-' + str(int(WBAN))
    gsod_year_dir = os.path.join(data_dir, 'gsod', str(year))
    path = os.path.join(gsod_year_dir, station + '.op')
    if os.path.exists(path):
        data = open(path).read()
        if data and data != 'Exception':
            return data
        else:
            raise Exception(data)

    toget = ('ftp://ftp.ncdc.noaa.gov/pub/data/gsod/' + str(year) + '/'
             + station + '-' + str(year) +'.op.gz')
    try:
        data = urlopen(toget, timeout=5)
    except Exception as e:
        if not os.path.exists(gsod_year_dir):
            os.makedirs(gsod_year_dir)
        open(path, 'w').write('Exception')
        raise Exception('Could not obtain desired data; check '
                        'if the year has data published for the '
                        'specified station and the station was specified '
                        'in the correct form. The full error is %s' %(e))

    data = data.read()
    data_thing = StringIO(data)

    f = gzip.GzipFile(fileobj=data_thing, mode="r")
    year_station_data = f.read()
    try:
        year_station_data = year_station_data.decode('utf-8')
    except:
        pass

    # Cache the data for future use
    if not os.path.exists(gsod_year_dir):
        os.makedirs(gsod_year_dir)
    open(path, 'w').write(year_station_data)

    return year_station_data

gsod_fields = ['DATE', # 15-18 int year; 19-22 int month/day
               'TEMP', # 25-30 Real Mean temperature for the day in degrees Fahrenheit to tenths. Missing = 9999.9
               'TEMP_COUNT', # 32-33 Int. Number of observations used in calculating mean temperature
               'DEWP', # 36-41 Real Mean dew point for the day in degrees Fahrenheit to tenths.  Missing = 9999.9
               'DEWP_COUNT', # 43-44 Int. Number of observations used in calculating mean dew point
               'SLP', # 47-52 Real Mean sea level pressure for the day in millibars to tenths.  Missing = 9999.9
               'SLP_COUNT', # 54-55 Int. Number of observations used in calculating mean sea level pressure
               'STP', # 58-63 Real Mean station pressure for the day in millibars to tenths. Missing = 9999.9
               'STP_COUNT', # 65-66 Int. Number of observations used in calculating mean station pressure
               'VISIB', # 69-73 Real Mean visibility for the day in miles to tenths. Missing = 999.9
               'VISIB_COUNT', # 75-76 Int. Number of observations used in calculating mean visibility
               'WDSP', # 79-83 Real Mean wind speed for the day in knots to tenths. Missing = 999.9
               'WDSP_COUNT', # 85-86 Int. Number of observations used in calculating mean wind speed
               'MXSPD', # 89-93 Real Maximum sustained wind speed reported for the day in knots to tenths. Missing = 999.9
               'GUST', # 96-100 Real Maximum wind gust reported for the day in knots to tenths. Missing = 999.9
               'MAX', # 103-108 Real Maximum temperature reported during the
                      # day in Fahrenheit to tenths--time of max temp report varies by country and
                      # region, so this will sometimes not be the max for the calendar day.
                      # Missing = 9999.9; FLAG of '*' is present on 109-109!
               'MIN', # 111-116 Real Minimum temperature reported during the day in Fahrenheit to tenths--time of min
                      # temp report varies by country and region, so this will sometimes not be
                      # the min for the calendar day. Missing = 9999.9 FLAG of '*' is present on 117-117!
               'PRCP', # 119-123 Real Total precipitation (rain and/or melted snow) reported during the day in inches
                       # and hundredths; will usually not end with the midnight observation--i.e.,
                       # may include latter part of previous day. .00 indicates no measurable
                       # precipitation (includes a trace).
                       # Missing = 99.99

               'SNDP', # 126-130 Real Snow depth in inches to tenths--last report for the day if reported more than
                       # once.  Missing = 999.9 Note: Most stations do not report '0' on days with no snow on the
                       # ground--therefore, '999.9' will often appear on these days.
               'FRSHTT' # 133-138 Int. Indicators (1 = yes, 0 = no/not reported) for the occurrence during the day of:
                        # Fog ('F' - 1st digit).
                        # Rain or Drizzle ('R' - 2nd digit).
                        # Snow or Ice Pellets ('S' - 3rd digit).
                        # Hail ('H' - 4th digit).
                        # Thunder ('T' - 5th digit).
                        # Tornado or Funnel Cloud ('T' - 6th digit).
              ]
# Use TEMP and DEWP and STP to calculate wet bulb temperatures
# Values to be converted to floats always
gsod_float_fields = ('TEMP', 'DEWP', 'SLP', 'STP', 'VISIB', 'WDSP', 'MXSPD',
                     'GUST', 'MAX', 'MIN', 'PRCP', 'SNDP')
# Values to be converted to ints always
gsod_int_fields = ('TEMP_COUNT', 'DEWP_COUNT', 'SLP_COUNT', 'STP_COUNT',
                   'VISIB_COUNT', 'WDSP_COUNT')

# Values which signify flags
gsod_flag_chars = '*ABCDEFGHI'
# Values which should be converted to None, as normally there is no value
gsod_bad_values = set(['99.99', '999.9', '9999.9'])

gsod_indicator_names = ['fog', 'rain', 'snow_ice', 'hail', 'thunder',
                        'tornado']
five_ninths = 5.0/9.0

gsod_day = namedtuple('gsod_day', gsod_fields + gsod_indicator_names)

def gsod_day_parser(line, SI=True, to_datetime=True):
    '''One line (one file) parser of data in the format of the GSOD database.
    Returns all parsed results as a namedtuple for reduced memory consumption.
    Will convert all data to base SI units unless the `SI` flag is set to
    False. As the values are rounded to one or two decimal places in the
    GSOD database in Imperial units, it may be useful to look at the values
    directly.

    The names columns of the columns in the GSOD database are retained and used
    as the attributes of the namedtuple results.

    The day, month, and year are normally converted to a datetime instance in
    resulting namedtuple; this behavior can be disabled by setting the
    `datetime` flag to False; it will be a string in the format YYYYMMDD if so.
    This may be useful because datetime conversion roughly doubles the speed of
    this function.

    Parameters
    ----------
    line : str
        Line in format of GSOD documentation, [-]
    SI : bool
        Whether or not the results get converted to base SI units, [-]
    to_datetime : bool
        Whether or not the date gets converted to a datetime instance or stays
        as a string, [-]

    Returns
    -------
    gsod_day_instance : gsod_day
        namedtuple with fields described in the source (all values in SI units,
        if `SI` is True, i.e. meters, m/s, Kelvin, Pascal; otherwise the
        original unit set is used), [-]
    '''
    # Ignore STN--- and WBAN, 8-12 characters
    fields = line.strip().split()[2:]
    # For the case the field is blank, set it to None; strip it either way
    for i in range(len(fields)):
        field = fields[i].strip()
        if not field:
            field = None
        fields[i] = field

    obj = dict(zip(gsod_fields, fields))
    # Convert the date to a datetime object if specified
    if to_datetime and obj['DATE'] is not None:
        obj['DATE'] = datetime.datetime.strptime(obj['DATE'], '%Y%m%d')

    # Parse float values as floats
    for field in gsod_float_fields:
        value = obj[field].rstrip(gsod_flag_chars)
        if value in gsod_bad_values:
            value = None
        else:
            value = float(value)
        obj[field] = value

    if SI:
        # All temperatures are in deg F
        for field in ('TEMP', 'DEWP', 'MAX', 'MIN'):
            value = obj[field]
            if value is not None:
                # F2K inline for efficiency unfortunately
                obj[field] = (value + 459.67)*five_ninths

        # Convert visibility, wind speed, pressures
        # to si units of meters, Pascal, and meters/second.
        if obj['VISIB'] is not None:
            obj['VISIB'] = obj['VISIB']*mile
        if obj['PRCP'] is not None:
            obj['PRCP'] = obj['PRCP']*inch
        if obj['SNDP'] is not None:
            obj['SNDP'] = obj['SNDP']*inch
        if obj['WDSP'] is not None:
            obj['WDSP'] = obj['WDSP']*knot
        if obj['MXSPD'] is not None:
            obj['MXSPD'] = obj['MXSPD']*knot
        if obj['GUST'] is not None:
            obj['GUST'] = obj['GUST']*knot
        if obj['SLP'] is not None:
            obj['SLP'] = obj['SLP']*100.0
        if obj['STP'] is not None:
            obj['STP'] = obj['STP']*100.0

    # Parse int values as ints
    for field in gsod_int_fields:
        value = obj[field]
        if value is not None:
            obj[field] = int(value)

    indicator_values = [flag == '1' for flag in obj['FRSHTT']]
    obj.update(zip(gsod_indicator_names, indicator_values))
    return gsod_day(**obj)


from __future__ import print_function
#
# Computing the free energy difference of an organic crystal polymorph at different gamma values
#
# Copyright Eric Dybeck and Michael R. Shirts, University of Virginia, 2014
#

import numpy as np
import pymbar  # multistate Bennett acceptance ratio
from pymbar import timeseries  # timeseries analysis
from optparse import OptionParser # for parsing command-line options
import sys
import panedr
import os

def dA_MBAR(minimum=0, maximum=100, spacing=10, exponent=2, polymorphs='p1 p2', Molecules=72, Independent=4, Temp=200,
            bonds=False, primary_directory='.', added_directories=[]):
    # =============================================================================================
    # Setting up the values for gamma or lambda states
    # =============================================================================================
#    raw_value = minimum
#    values = []
    directory_names = np.arange(minimum, maximum + spacing, spacing)
    directory_names = np.sort(np.append(directory_names, added_directories))

#    while raw_value <= maximum:
#        if exponent >= 0:
#            value = int(100 * (float(raw_value) / float(maximum)) ** abs(exponent))
#        else:
#            value = int(100 * (1 - (float(maximum - raw_value) / float(maximum)) ** abs(exponent)))
#        values.append(value)
#        raw_value = raw_value + spacing
#    print(values)
#    print(directory_names)
#    exit()

    # POLYMORPH
    polymorphs = polymorphs.split()

    # =============================================================================================
    # READ IN RAW DATA
    # =============================================================================================
    # Constants.
    kB = 1.3806488e-23 * 6.0221413e23 / (1000.0 * 4.184)  # Boltzmann constant in kcal/mol

    # Parameters
    T_k = Temp * np.ones(len(directory_names), float)  # Convert temperatures to floats
    print(T_k)
  #  print(values)

    K = len(directory_names)  # How many states?

    # total number of states examined; 0 are unsampled if bonds are left on, 1 is unsampled if the bonds are removed
    Kbig = K

    # maximum number of snapshots/simulation (could make this automated) - doesn't matter, as long as it's long enough.
    N_max = 5000

    # beta factor for the different temperatures
    beta_k = 1.0 / (kB * T_k)
    dA = np.zeros([len(polymorphs), Kbig], float)
    ddA = np.zeros([len(polymorphs), Kbig], float)
    convert_units = 0.2390057 * np.ones(Kbig, float)  # Convert all energies to kcal/mol

    # Allocate storage for simulation data
    for i, poly in enumerate(polymorphs):
        # N_k[k] is the total number of snapshots from alchemical state k
        N_k = np.zeros([Kbig], np.int32)

        # N_k_s[k,s] is the total number of snapshots from alchemical state k from seed s
        N_k_s = np.zeros([Kbig], np.int32)

        # u_kln[k,l,n] is the adjusted energy of snapshot n from simulation k
        u_kln = np.zeros([K, Kbig, N_max], np.float64)

        # dhdl_kn[k,n] is the derivative of energy with respect to lambda of snapshot n from simulation k
        dhdl_kn = np.zeros([K, N_max], np.float64)

        #Load in the data for each run
        for k in range(K):
            n = 0

            # cycle through all the input total energy data
            if directory_names[k] == int(directory_names[k]):
                dirpath = polymorphs[i] + '/' + primary_directory + '/' + str(int(directory_names[k]))
            else:
                dirpath = polymorphs[i] + '/' + primary_directory + '/' + str(directory_names[k])
            if os.path.isdir(dirpath):
                fname = dirpath + '/PROD.edr'
                dhdlname = dirpath + '/dhdl_PROD.xvg'

                potential_energy = panedr.edr_to_df(fname)['Potential'].values
                print("loading " + fname)

                dhdl_energy = np.loadtxt(dhdlname, comments=['#', '$', '@', '!'])
                print("loading " + dhdlname)

                # Removing any non-equilibrated points of the simulation
                [start_production, _, _] = timeseries.detectEquilibration(potential_energy)
                potential_energy = potential_energy[start_production:]
                dhdl_energy = dhdl_energy[start_production:,:]

                # Cutting points if they exceed N_max
                if len(potential_energy) > N_max:
                    potential_energy = potential_energy[len(potential_energy) - N_max:]
                    dhdl_energy = dhdl_energy[len(dhdl_energy) - N_max:,:]

                # the energy of every configuration from each state evaluated at its sampled state
                n = len(potential_energy)
                dhdl_placement = len(dhdl_energy[0, :]) - K
                u_kln[k, :K, :n] = (potential_energy.reshape((n, 1)) + dhdl_energy[:, dhdl_placement:]).T * convert_units[k]
                dhdl_kn[k, :n] = (float(Independent) / Molecules) * \
                                 np.sum(dhdl_energy[:, 2:dhdl_placement], axis=1) * convert_units[k]

                N_k_s[k] = n
                N_k[k] = n

        # convert to nondimensional units from kcal/mol
        u_kln *= beta_k[0]

        #u_kln_save = u_kln.copy()
        u_kln_save = u_kln[:]
        g_k = np.zeros([K])

        print("Number of retained samples")
        print(N_k)
        print("Number of retained samples from each seed")
        print(N_k_s)

        # =============================================================================================
        # COMPUTE FREE ENERGY DIFFERENCE USING MBAR
        # =============================================================================================

        # Initialize MBAR.
        print("Running MBAR...")

        # generate the weights of each of the umbrella set
        mbar = pymbar.MBAR(u_kln, N_k, verbose=True, subsampling_protocol=[{'method': 'L-BFGS-B'}])

        print("MBAR Converged...")
        # testing

        for k in range(Kbig):
            w = np.exp(mbar.Log_W_nk[:, k])
            print("max weight in state %d is %12.7f" % (k, np.max(w)))
            neff = 1 / np.sum(w ** 2)
            print("Effective number of sample in state %d is %10.3f" % (k, neff))
            print("Efficiency for state %d is %d/%d = %10.4f" % (k, neff, len(w), neff / len(w)))

        # extract self-consistent weights and uncertainties
        (df_i, ddf_i, theta_i) = mbar.getFreeEnergyDifferences()

        print("Free Energies Optained...")

        # convert PMF to kcal/mol and normalize by the number of molecules
        df_i /= (beta_k[0] * float(Independent))
        ddf_i /= (beta_k[0] * float(Independent))

        dA[i, :] = df_i[-1]

        # =============================================================================================
        # COMPUTE UNCERTAINTY USING THE UNCORRELATED DATA
        # =============================================================================================

        for k in range(K):
            N_k[k] = 0
            n_old = 0

            g_k[k] = timeseries.statisticalInefficiency(dhdl_kn[k, n_old: (n_old + N_k_s[k])])
            print("Correlation time for sampled state %d is %10.3f" % (k, g_k[k]))
            # subsample the data to get statistically uncorrelated data
            indices = np.array(timeseries.subsampleCorrelatedData(u_kln[k, k, n_old:(n_old + N_k_s[k])],
                                                                          g=g_k[k]))  # subsample

            # not sure why we have to transpose
            if indices != []:
                u_kln[k, :, N_k[k]: (N_k[k] + len(indices))] = u_kln_save[k, :, (indices + n_old)].transpose()
                N_k[k] = N_k[k] + len(indices)
                n_old += N_k_s[k]

        print("Number of retained samples")
        print(N_k)
        print("Number of retained samples from each seed")
        print(N_k_s)

        # generate the weights of each of the umbrella set
        mbar = pymbar.MBAR(u_kln, N_k, verbose=True, subsampling_protocol=[{'method': 'L-BFGS-B'}])

        print("MBAR Converged...")

        # extract self-consistent weights and uncertainties
        try:
            (df_u, ddf_u, theta_i) = mbar.getFreeEnergyDifferences()
        except ValueError:
            pass

        print("Free Energies Optained...")

        # convert PMF to kcal/mol and normalize by the number of molecules
        df_u /= (beta_k[0] * float(Independent))
        ddf_u /= (beta_k[0] * float(Independent))

        ddA[i, :] = ddf_u[-1]
#        ddA[i, :] = ddf_i[-1]

        # Write out free energy differences
        print("Free Energy Difference (in units of kcal/mol)")
        print("  dA(Gamma) = A(Gamma) - A(Interactions Off)")
        for k in range(Kbig):
            print("%8.3f %8.3f" % (df_i[k, -1], ddf_u[k, -1]))

        del N_k
        del N_k_s
        del u_kln
        del dhdl_kn

    out_dA = np.zeros(len(polymorphs))
    out_ddA = np.zeros(len(polymorphs))
    for i, poly in enumerate(polymorphs):
        out_dA[i] = dA[i, 0]
        out_ddA[i] = ddA[i, 0]

    return out_dA, out_ddA


# This file is part of LayerModel_lib
#
#     A tool to compute the transmission behaviour of plane electromagnetic waves
#     through human tissue.
#
# Copyright (C) 2018 Jan-Christoph Brumm
#
# Licensed under MIT license.
#
import numpy as np
import os
from scipy import ndimage

from LayerModel_lib import VoxelModel, VoxelModelImporter, Coordinate, TissueProperties

current_directory = os.path.dirname(__file__)
base_path = os.path.join(current_directory, '..', '..', '..', '..', '..', 'Numerical Human Phantoms', 'Golem')

# path to the AVW File of this model
filename = os.path.join(base_path, 'segm_golem')
# path to the tissue_mapping file
tissue_file = os.path.join('ImportGolem_tissues.txt')

AVW_Data = VoxelModelImporter(filename, tissue_file, 'AVW')
model_orig = AVW_Data.data['image']
tissue_name_orig = AVW_Data.tissue_names
tissue_mapping = AVW_Data.tissue_mapping

Golem = VoxelModel()
Golem.show_progress_bar = True

# needs to be set manually from README.txt
Golem.set_scale(2.08, 2.08, 8)

Golem.name = 'Golem'
Golem.description = 'Golem model from the Helmholtz Zentrum München. ' \
                    'Resolution %.2fmm x %.2fmm x %.2fmm' % (Golem.scaling.x,
                                                             Golem.scaling.y, Golem.scaling.z)

# Golem is upside down and left right in wrong direction
model_orig = np.flip(model_orig, axis=0)
model_orig = np.flip(model_orig, axis=2)
# The space around the model is too much
model_orig = model_orig[0:166, :, :]

#  Calculate the outer_shape of the original and the complete model
outer_shape = AVW_Data.calculate_outer_shape(model_orig, tissue_mapping)

Golem.add_voxel_data(short_name='original',
                     name='Original data from AVW file',
                     model=model_orig,
                     outer_shape=outer_shape,
                     tissue_names=tissue_name_orig)

Golem.add_voxel_data(short_name='complete',
                     name='The \'original\' model converted to our TissueProperties.',
                     model=Golem.models['original'].data,
                     outer_shape=outer_shape,
                     tissue_mapping=tissue_mapping)

# Calculate the trunk model
start_slice = int(110)
end_slice = int(175)

(model_trunk, trunk_mask) = AVW_Data.calculate_trunk_model(Golem, 'complete', z_start=start_slice, z_end=end_slice)
outer_shape_trunk = AVW_Data.calculate_outer_shape(model_trunk)

# label the inner parts of the small intestine as GIcontents. For that erode the small intestine in
# 3D and replace the resulting voxels with GIcontents.
tp = TissueProperties()

trunk = model_trunk

# Golem has no labeled SmallIntestineContents. Therefore, we need to label that by ourself:
# select only the small intenstine
trunk_si = trunk == tp.get_tissue_id_for_name(['SmallIntestine'])
struct = ndimage.generate_binary_structure(3, 1)
# Erode the small intestine wall:
trunk_si_eroded = ndimage.binary_erosion(trunk_si, structure=struct, iterations=1)
trunk_si_out = np.copy(trunk)
# replace remaining inner small intestine with GIcontents
trunk_si_out[trunk_si_eroded] = tp.get_tissue_id_for_name(['GIcontents'])

model_trunk = trunk_si_out

Golem.add_voxel_data(short_name='trunk',
                     name="The trunk of the 'complete' model. Arms have been removed using "
                          "VoxelModel.remove_arms().",
                     outer_shape=outer_shape_trunk,
                     model=model_trunk,
                     mask=trunk_mask,
                     tissue_mapping=None)

surface = Golem.create_3d_model(model_type='trunk', patch_size=(30, 30))
Golem.models['trunk'].surface_3d = surface

Golem.models['trunk'].endpoints = []
for (i, s) in enumerate(surface):
    Golem.models['trunk'].endpoints.append(Coordinate(np.array(s['centroid'])))

Golem.save_model()


#!/usr/bin/env python

"""
short script to transform remodnav
(https://github.com/psychoinformatics-de/remodnav) outputs from naturalistic
stimulation with hollywood movie forrest gump into fixation vectors
"""

import numpy as np
import pandas as pd
from bisect import bisect_right
import os.path as op

def pursuits_to_fixations(npdata):
    """this function takes a numpy record array from Asims eye-event detection
    algorithm. Start and end points of pursuits are transformed into a fixation,
    the pursuit movement can then be simplified as a saccade. The function
    returns a recordarray"""
    # initialize empty rec array of the same shape
    newdata = np.recarray((0,), dtype=[('onset', '<f8'),
                                       ('duration', '<f8'),
                                       ('label', '<U10'),
                                       ('start_x', '<f8'),
                                       ('start_y', '<f8'),
                                       ('end_x', '<f8'),
                                       ('end_y', '<f8'),
                                       ('amp', '<f8'),
                                       ('peak_vel', '<f8'),
                                       ('med_vel', '<f8'),
                                       ('avg_vel', '<f8')])
    # reassemble rec array. split pursuits to use end and start as fixations later
    for i in range(0, len(npdata)):
        if npdata[i]['label'] == 'PURS':
            row_1 = npdata[i]
            row_1['duration'] = npdata[i]['duration'] / 2
            row_2 = row_1.copy()
            row_2['onset'] += row_2['duration']
            row_2['start_x'] = row_2['end_x']
            row_2['start_y'] = row_2['end_y']
            newdata = np.append(newdata, row_1)
            newdata = np.append(newdata, row_2)
        else:
            newdata = np.append(newdata, npdata[i])
    return newdata

def preprocess(data, sz=[1280, 720]):
    """
    data = n x 11 rec array
    sz = screen measurements
    preprocesses a recordarray with eye events (from pursuits_to_fixations()
    function. Assumes the datafile is sorted by time. Will filter to include
    only Fixations and Pursuits-start/end-points, will check for out-of-bound
    gazes.
    Returns clean Fixation data with onset, start_x, start_y and duration """
    # only fixations and pursuits
    Filterevents = data[np.logical_or(data['label'] == 'FIXA',
                                      data['label'] == 'PURS')]
    # within x coordinates?
    Filterxbounds = Filterevents[np.logical_and(Filterevents['start_x'] >= 0,
                                                Filterevents['start_x'] <= sz[0])]
    # within y coordinates?
    Filterybounds = Filterxbounds[np.logical_and(Filterxbounds['start_y'] >= 0,
                                                 Filterxbounds['end_y'] <= sz[1])]
    # give me onset times, start_x, start_y and duration
    fixations = Filterybounds[["onset", "start_x", "start_y",
                               "duration"]]
    return fixations

def takeClosestright(myList, myNumber):
    """return the integer closest to 'myNumber' in an ordered list.(shamelessly
    stolen from
    https://stackoverflow.com/questions/12141150/from-list-of-integers-get-number-
    closest-to-a-given-value)
    """
    pos = bisect_right(myList, myNumber)
    if pos == 0:
        return myList[0]
    if pos == len(myList):
        return myList[-1]
    after = myList[pos]
    return after

def createOnsets(data, dur):
    """create onset times of all shots of 'dur' seconds of length
    data = dataframe, should be location annotation
    dur = duration in seconds"""
    onsets = []
    for index, row in data.iterrows():
        if row['duration'] >= dur:
            onsets.append(row['onset'])
    return onsets

def createChunks(onsets, fixations, dur):
    """Create and return start and end indices to chunk eye movement data into
    segments to compare scanpaths across.
    onsets = output from CreateOnsets()
    fixations = output from preprocess(). n x 4 np record array
    dur = durations of segments in seconds"""
    # initialize empty lists
    startidx, endidx = [], []
    for shotonset in onsets:
        start = takeClosestright(fixations['onset'], shotonset)
        startidx.append(np.where(fixations['onset'] == start)[0].tolist())
        end = takeClosestright(fixations['onset'], shotonset + dur)
        endidx.append(np.where(fixations['onset'] == end)[0].tolist())
    # flatten the nested lists
    startidx = [element for sublist in startidx for element in sublist]
    endidx = [element for sublist in endidx for element in sublist]
    return startidx, endidx

def FixationsChunks(fixations, startid, endid):
    """Chunk eye movement data into segments of approximate length to compute
    scanpath similarities. Output is returned as a n x 3 fixation vector.
    startid, endid = output from createChunks
    fixations = output from preprocess"""
    fixation_vector = []
    # slice fixation data according to indices, take columns start_x, start_y and
    # duration
    for idx in range(0, len(startid)):
        ind = fixations[startid[idx]:endid[idx]][["start_x", "start_y", "duration"]]
        fixation_vector.append(ind)
    return fixation_vector

def longshot(shots, dur):
    """group movie shots without a cut together to obtain longer movie
    segments. This way, fewer but longer scanpaths are obtained. Example: use
    median shotlength of 4.92s.
    shots = dataframe, contains movie location annotation
    dur = length in seconds for movie shot
    """
    # turn pandas dataframe shots into record array
    structshots = shots.to_records()
    i = 0
    while i < len(structshots):
        # break before running into index error
        if structshots[i] == structshots[-1]:
            break
        else:
            if (structshots[i]['duration'] < dur) & \
                    (structshots[i + 1]['duration'] < dur) & \
                    (structshots[i]['locale'] == structshots[i + 1]['locale']):
                # add durations together and delete second row
                structshots[i]['duration'] += structshots[i + 1]['duration']
                structshots = np.delete(structshots, i + 1, 0)
            else:
                i += 1
    aggregated = pd.DataFrame({'onset': structshots['onset'].tolist(),
                               'duration': structshots['duration'].tolist()}, columns=['onset',
                                                                                'duration'])
    return aggregated

def savefile(fixation_vector, output, header):
    newheader = ''.join([w + '\t' for w in header]).strip()
    np.savetxt(output, fixation_vector, delimiter='\t', comments='', header=newheader)

def run(data1, shots, sz, dur, subname):
    newdata1 = pursuits_to_fixations(data1)
    fixations1 = preprocess(newdata1, sz)
    shots = longshot(shots, dur)
    onset = createOnsets(shots, dur)
    startid1, endid1 = createChunks(onset, fixations1, dur)
    fixation_vectors1 = FixationsChunks(fixations1, startid1, endid1)
    header = fixation_vectors1[0].dtype.names
    for i in range(0, len(onset)):
        output = args.output + '/fixvectors/segment_' + str(i) + '_' + subname + '.tsv'
        print('saving file', i, 'into', output)
        savefile(fixation_vectors1[i], output, header)

if __name__ == '__main__':
    import argparse

    parser = argparse.ArgumentParser()
    # define arguments
    parser.add_argument('-i', '--input1', nargs='+', help='Input: eyemovement data of one subject', metavar='PATH',
                        required=True)
    parser.add_argument('-o', '--output', help='Output: Specify path where output should be saved', metavar='PATH',
                        required=True)
    parser.add_argument('-sz', '--screensize',
                        help='Screensize: what are the dimensions of the screen the stimulus was displayed on in px, e.g. [1280, 720]',
                        default=[1280, 720])
    parser.add_argument('-s', '--shots', help='Input3: location annotation of the movie segment', metavar='PATH',
                        required=True)
    parser.add_argument('-d', '--duration',
                        help='approximate duration of video segments to derive fixation vectors from', default=5.0)
    args = parser.parse_args()

    data1 = np.recfromcsv(args.input1[0],
                          delimiter='\t',
                          dtype={'names': ('onset', 'duration', 'label', 'start_x', 'start_y',
                                           'end_x', 'end_y', 'amp', 'peak_vel', 'med_vel', 'avg_vel'),
                                 'formats': ('f8', 'f8', 'U10', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8',
                                             'f8', 'f8')})
    shots = pd.read_csv(args.shots, sep='\t')

    subname = op.basename(args.input1[0]).split('_')[0]

    if args.screensize:
        sz = args.screensize
    else:
        sz = [1280, 720]

    if args.duration:
        dur = float(args.duration)
    else:
        dur = 5.0

    # run everything
    run(data1, shots, sz, dur, subname)


"""Automatic differentiation filter-error method estimation models."""

import collections
import functools
import inspect
import itertools

import jax
from jax import numpy as jnp
from jax import scipy as jscipy

import numpy as onp

### Enable 64-bit in jax ###
jax.config.update("jax_enable_x64", True)

class BoundADFunction:
    def __init__(self, adfun, model):
        self.adfun = adfun
        self.model = model

    @property
    def __signature__(self):
        return bound_signature(self.adfun)

    @property
    def __name__(self):
        return self.adfun.__name__

    def __repr__(self):
        cls = type(self)
        name = getattr(self, '__name__', '*no name*')
        return f"<{cls.__module__}.{cls.__name__} '{name}'>"

    def __call__(self, *args, **kwargs):
        return self.adfun(self.model, *args, **kwargs)

    def hess_filter(self, dec_shapes=None, out_shape=None):
        filt_name = f'_{self.__name__}_hess_filter'
        try:
            return getattr(self.model, filt_name)
        except AttributeError:
            assert dec_shapes is not None
            assert out_shape is not None

            hess_filter = {}
            adfun = self.adfun
            wrt_seq = adfun.hessian
            ind_base = adfun._sparse_deriv_ind(wrt_seq, dec_shapes, out_shape)
            for wrt, ind in ind_base.items():
                if len(wrt) == 2 and wrt[0] == wrt[1]:
                    hess_filter[wrt] = ind[0] > ind[1]
            setattr(self.model, filt_name, hess_filter)
            return hess_filter

    def hess_nnz(self, dec_shapes, out_shape):
        adfun = self.adfun
        wrt_seq = adfun.hessian
        filter = self.hess_filter(dec_shapes, out_shape)
        return adfun._sparse_deriv_nnz(wrt_seq, dec_shapes, out_shape, filter)

    def hess_ind(self, dec_shapes, out_shape):
        adfun = self.adfun
        wrt_seq = adfun.hessian
        filter = self.hess_filter(dec_shapes, out_shape)
        return adfun._sparse_deriv_ind(wrt_seq, dec_shapes, out_shape, filter)

    def hess_val(self, *args, **kwargs):
        adfun = self.adfun
        model = self.model
        wrt_seq = adfun.hessian

        # The filter must be built by a previous call to hess_ind or hess_nnz
        filter = self.hess_filter()

        return adfun._sparse_deriv_val(wrt_seq, model, args, kwargs, filter)

class BoundADConstraint(BoundADFunction):
    def jac_nnz(self, dec_shapes, out_shape):
        adfun = self.adfun
        d1 = adfun.first_derivatives
        return adfun._sparse_deriv_nnz(d1, dec_shapes, out_shape)

    def jac_ind(self, dec_shapes, out_shape):
        adfun = self.adfun
        d1 = adfun.first_derivatives
        return adfun._sparse_deriv_ind(d1, dec_shapes, out_shape)

    def jac_val(self, *args, **kwargs):
        adfun = self.adfun
        d1 = adfun.first_derivatives
        return adfun._sparse_deriv_val(d1, args, kwargs)

class BoundADObjective(BoundADFunction):
    def grad(self, *args, **kwargs):
        adfun = self.adfun
        ret = collections.OrderedDict()
        for wrt in adfun.first_derivatives:
            # Calculate the gradient
            grad_fun = adfun.derivatives[wrt,]
            grad_val = grad_fun(self.model, *args, **kwargs)

            # skip empty gradients
            if not grad_val.size:
                continue

            # Get the shape of the wrt argument
            try:
                wrt_shape = onp.shape(kwargs[wrt])
            except KeyError:
                wrt_shape = onp.shape(args[adfun.argnum(wrt) - 1])

            # Accumulate so the gradient has the same shape as the variable
            ret[wrt] = grad_val.reshape(-1, *wrt_shape).sum(0)
        return ret

class ADFunction:
    """Helper for optimization function automatic differentiation."""

    BoundClass = BoundADFunction

    def __init__(self, fun, core_shape=''):
        self.fun = fun
        """Underlying function."""

        self.core_shape = core_shape
        """The shape of the output of the elementary function, if vectorized."""

        self.hessian = None
        """Hessian elements."""

        self.__signature__ = inspect.signature(self.fun)
        """The underlying function's signature."""

        self.args = list(self.__signature__.parameters)
        """The underlying function argument names."""

        self.derivatives = {}
        """Dictionary of function derivatives."""

        self.first_derivatives = []
        """Sequence of registered first derivatives."""

        self.isvectorized = False
        """Whether the function is vectorized"""

    def __call__(self, *args, **kwargs):
        return self.fun(*args, **kwargs)

    def __get__(self, instance, owner=None):
        if instance is None:
            return self
        else:
            return self.BoundClass(self, instance)

    def __set_name__(self, owner, name):
        self.__name__ = name

    def __repr__(self):
        cls = type(self)
        name = getattr(self, '__name__', '*no name*')
        return f"<{cls.__module__}.{cls.__name__} '{name}'>"

    def argnum(self, argname):
        return self.args.index(argname)

    def derivative(self, wrt):
        if isinstance(wrt, str):
            wrt_tuple = wrt,
            return self.derivative(wrt_tuple)

        # Trivial case, i.e., 0-th derivative
        if wrt == ():
            return self.fun

        # Return the registered derivative, if it exists
        try:
            return self.derivatives[wrt]
        except KeyError:
            pass

        # Compute the derivative
        assert len(wrt) >= 1
        fun = self.derivative(wrt[1:])
        argnum = self.argnum(wrt[0])
        deriv = jax.jacrev(fun, argnum)

        # Save it and return
        self.derivatives[wrt] = deriv
        return deriv

    def prepare_derivatives(self, decision):
        # Compute and save the first derivatives
        for d in self.args:
            if d in decision:
                self.derivative(d)
                self.first_derivatives.append(d)

        # Define default Hessian, if unset
        if self.hessian is None:
            first_deriv = self.first_derivatives
            hess_gen = itertools.combinations_with_replacement(first_deriv, 2)
            self.hessian = list(hess_gen)

        # Compute second derivatives
        for wrt_pair in self.hessian:
            self.derivative(wrt_pair)

    def vectorize(self, vectorized):
        vec_args = [a for a in self.args if a in vectorized]
        excluded = [i for i,a in enumerate(self.args) if a not in vectorized]

        if not vec_args:
            return

        arg_sig = ",".join(f'({vectorized[a]})' for a in vec_args)
        sig = f"{arg_sig}->({self.core_shape})"

        self.fun = jnp.vectorize(self.fun, excluded=excluded, signature=sig)
        for wrt, d in self.derivatives.items():
            wrtsig = (vectorized[var] for var in reversed(wrt))
            if self.core_shape:
                outsig = ','.join((self.core_shape, *wrtsig))
            else:
                outsig = ','.join(wrtsig)
            dsig = f"{arg_sig}->({outsig})"
            vecd = jnp.vectorize(d, excluded=excluded, signature=dsig)
            self.derivatives[wrt] = vecd

        core_shape = self.core_shape
        out_core_ndim = len(core_shape.split(',')) if core_shape else 0
        self.core_ndim = {a: len(vectorized[a].split(',')) for a in vec_args}
        self.core_ndim[None] = out_core_ndim
        self.isvectorized = True

    def _split_shape(self, shape, varname=None):
        """Split a variable's shape into extension and core."""
        try:
            core_ndim = self.core_ndim[varname]
        except KeyError:
            return (), shape #This variable is not vectorized

        # Test whether the core element is scalar (ndim==0)
        if core_ndim:
            return shape[:-core_ndim], shape[-core_ndim:]
        else:
            return shape, ()

    def _ext_shape(self, shape, varname=None):
        """Return a variable's shape extension."""
        return self._split_shape(shape, varname)[0]

    def _core_shape(self, shape, varname=None):
        """Return a variable's core shape."""
        return self._split_shape(shape, varname)[1]

    def _deriv_core_shape(self, wrt, dec_shapes, out_shape):
        out_ext, out_core = self._split_shape(out_shape)
        if len(wrt) == 0:
            return out_core
        else:
            wrt0, *wrt_rem = wrt
            wrt0_shape = dec_shapes[wrt0]
            wrt0_ext, wrt0_core = self._split_shape(wrt0_shape, wrt0)
            rem_core = self._deriv_core_shape(wrt_rem, dec_shapes, out_shape)
            return rem_core + wrt0_core

    def _deriv_core_ind(self, wrt, dec_shapes, out_shape):
        if len(wrt) == 0:
            out_ext, out_core = self._split_shape(out_shape)
            return [onp.arange(shape_size(out_core))]
        else:
            wrt0, *wrt_rem = wrt
            wrt0_shape = dec_shapes[wrt0]
            wrt0_ext, wrt0_core = self._split_shape(wrt0_shape, wrt0)
            rem_ind = self._deriv_core_ind(wrt_rem, dec_shapes, out_shape)

            wrt0_core_size = shape_size(wrt0_core)
            wrt0_tile = rem_ind[0].size
            wrt0_ind = onp.tile(onp.arange(wrt0_core_size), wrt0_tile)
            core_ind = [onp.repeat(i, wrt0_core_size) for i in rem_ind]
            core_ind.insert(0, wrt0_ind)
            return core_ind

    def _sparse_deriv_nnz(self, wrt_seq, dec_shapes, out_shape, filter={}):
        nnz = 0
        out_ext, out_core = self._split_shape(out_shape)
        ext_sz = shape_size(out_ext)
        for wrt in wrt_seq:
            if wrt in filter:
                nnz += filter[wrt].sum()
            else:
                deriv_core = self._deriv_core_shape(wrt, dec_shapes, out_shape)
                nnz += shape_size(deriv_core) * ext_sz
        return nnz

    def _sparse_deriv_ind(self, wrt_seq, dec_shapes, out_shape, filter={}):
        out_ext, out_core = self._split_shape(out_shape)
        core_out_sz = shape_size(out_core)

        ret = collections.OrderedDict()
        for wrt in wrt_seq:
            ind = []
            base_ind = self._deriv_core_ind(wrt, dec_shapes, out_shape)
            for wrt_name, wrt_ind in zip(wrt, base_ind):
                wrt_shape = dec_shapes[wrt_name]
                wrt_ext, wrt_core = self._split_shape(wrt_shape, wrt_name)
                wrt_core_sz = shape_size(wrt_core)

                wrt_offs = ndim_range(wrt_ext) * wrt_core_sz
                wrt_offs = onp.broadcast_to(wrt_offs, out_ext)
                ind.append((wrt_ind + wrt_offs[..., None]).ravel())

            # Extend the output indices
            out_ind = base_ind[-1]
            out_offs = ndim_range(out_ext) * core_out_sz
            ind.append((out_ind + out_offs[..., None]).ravel())

            # Save in dictionary
            selected = filter.get(wrt, slice(None))
            ret[wrt] = onp.array(ind)[:, selected]
        return ret

    def _sparse_deriv_val(self, wrt_seq, args, kwargs, filter={}):
        ret = collections.OrderedDict()
        for wrt in wrt_seq:
            deriv = self.derivatives[wrt]
            selected = filter.get(wrt, slice(None))
            ret[wrt] = deriv(*args, **kwargs).ravel()[selected]
        return ret

class ADConstraint(ADFunction):
    """Helper for constraint function automatic differentiation."""

    BoundClass = BoundADConstraint

class ADObjective(ADFunction):
    """Helper for objective function automatic differentiation."""

    BoundClass = BoundADObjective

def constraint(core_shape_or_fun=None):
    if callable(core_shape_or_fun):
        fun = core_shape_or_fun
        return constraint()(fun)
    else:
        core_shape = core_shape_or_fun
        def decorator(fun):
            return ADConstraint(fun, core_shape)
        return decorator

def objective(fun):
    return ADObjective(fun)

def hessian(*args):
    def decorator(obj):
        obj.hessian = args
        return obj
    return decorator

class ADModel:
    def __init_subclass__(cls):
        base_dec = getattr(super(), 'decision', set())
        cls_dec = getattr(cls, 'decision', set())
        cls.decision = set.union(base_dec, cls_dec)
        """The decision variables of this model."""

        base_vec = getattr(super(), 'vectorized', {})
        cls_vec = getattr(cls, 'vectorized', {})
        cls.vectorized = {**base_vec, **cls_vec}
        """The core shape of vectorized variables."""

        cls_items = cls.__dict__.items()
        adfuns = {k:v for k,v in cls_items if isinstance(v, ADFunction)}
        for name, adfun in adfuns.items():
            adfun.prepare_derivatives(cls.decision)
            adfun.vectorize(cls.vectorized)

class InnovationDTModel(ADModel):

    decision = {
        'x', 'en', 'xnext', 'xprev', 'enprev', 'ybias',
        'A', 'B', 'C', 'D', 'Ln', 'isRp_tril'
    }
    """Decision variables of the optimization problem."""

    vectorized = dict(
        x='nx', en='ny', xnext='nx', xprev='nx', enprev='ny',
        u='nu', y='ny', uprev='nu',
        A='nx,nx', B='nx,nu', C='ny,nx', D='ny,nu', Ln='nx,ny', ybias='ny',
        isRp_tril='nty',
    )
    """Decision variables of the optimization problem."""

    def __init__(self, nx, nu, ny):
        self.nx = nx
        """Number of states."""

        self.nu = nu
        """Number of inputs."""

        self.ny = ny
        """Number of outputs."""

    @hessian(('xprev', 'A'), ('enprev', 'Ln'))
    @constraint('nx')
    def dynamics(self, xnext, xprev, uprev, enprev, A, B, Ln):
        """Model dynamics defects."""
        xpred = A @ xprev + B @ uprev + Ln @ enprev
        return xnext - xpred

    @hessian(('x', 'C'), ('x', 'isRp_tril'), ('C', 'isRp_tril'),
             ('D', 'isRp_tril'), ('ybias', 'isRp_tril'))
    @constraint('ny')
    def innovation(self, y, en, x, u, C, D, ybias, isRp_tril):
        """Model normalized innovation constraint."""
        ymodel = C @ x + D @ u + ybias
        isRp = tril_mat(isRp_tril)
        return isRp @ (y - ymodel) - en

    @hessian(('en', 'en'), ('isRp_tril', 'isRp_tril'))
    @objective
    def Ln(self, en, isRp_tril):
        """Measurement log-likelihood."""
        isRp = tril_mat(isRp_tril)
        log_det_isRp = jnp.log(isRp.diagonal()).sum()
        return -0.5 * (en ** 2).sum() + log_det_isRp

def tril_mat(tril_elem):
    """Build a matrix from its lower-triangular elements."""
    ntril = len(tril_elem)
    n = int(round(0.5*(onp.sqrt(8*ntril + 1) - 1)))
    tril_ind = onp.tril_indices(n)
    M = jnp.zeros((n, n))
    return M.at[tril_ind].set(tril_elem)

def shape_size(shape):
    return onp.prod(shape, dtype=int)

def ndim_range(shape):
    assert isinstance(shape, tuple)
    return onp.arange(shape_size(shape)).reshape(shape)

def bound_signature(method):
    """Return the signature of a method when bound."""
    sig = inspect.signature(method)
    param = list(sig.parameters.values())[1:]
    return inspect.Signature(param, return_annotation=sig.return_annotation)

import numpy,import scipy,import jax,from jax


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Mon Sep 21 14:09:18 2020

@author: gao
"""

import numpy as np
import matplotlib.pyplot as plt
from matplotlib import colors
from mpl_toolkits.axes_grid1 import AxesGrid
import matplotlib as mpl
import os
from matplotlib.colors import LinearSegmentedColormap
#-------------------------------------------------------------------------------
# set display width
np.get_printoptions()['linewidth']
np.set_printoptions(linewidth=400)
#-------------------------------------------------------------------------------
"read size data at n=3, only size effects"

grid_chi=101                                           # grid size in figure, so check figure.py file first
grid_n=7                                               # 13[5], 23[6], 37[7]; 58[8]; 87[9]; 128[10]

chi_ratio_num=grid_chi
chi_ratio_log = np.linspace(-0.4,0.4, chi_ratio_num)    # log scale for chi/chi_neutral
chi_ratio_list=[np.power(10,i) for i in chi_ratio_log]       # chi/chi_neutral
ratio_index=72

with open("../simulation/LC.txt", "r") as file:
    lcs = eval(file.readline())                          # read lc list

#--------define a function to transform [1+2] to 1+2--------------------
def lc_add(x):                                           # x is the number for life cycles
	lcs[x].sort(reverse=True)
	off_num=len(lcs[x])
	ini=str(lcs[x][0])                                   #initial lcs
	for i in range(1,off_num):
		ini=ini+'+'+str(lcs[x][i])
	return ini

#----------------------------------------
N=3
lc_list=[5,56]

"add a list [[maturity size=3],[],[],;;;]"
data_list=[ ]

for chi_cluster in range(ratio_index,ratio_index+1):                        # how many figures or cell divisions
	for N_cluster in range(N-1,N):                        # y axes points = a
		all_list=[]                                          # algin all lcs' growth rates
		for i_th in lc_list:                         # read all lcs data = growth rate

			with open('../data/data_size_effect/%d_%d_%d.txt'%(N_cluster,chi_cluster,i_th), "r") as file:
				nan=float(np.nan)
				inf=np.inf
				grate = eval(file.readline())          # read growth rate
			data_list.append(np.array([chi_cluster,i_th,grate]))

data_size=np.array(data_list)

#-------------------------------------------------------------------------------
"read data of the threshold effects across k, only threshold effects"

grid_num=31                                   # grid size in figure, so check figure.py file first
grid_m_num=11                                        # 13[5], 23[6], 37[7]; 58[8]; 87[9]; 128[10]

b_range=np.array([1,16])
m_range=np.array([0,0.1])
k_range=np.array([1,2,3,4,5,6,7,8])          # first take 1,2,3,
k_cluster_list=np.array([0,1,2,3,4,5,6])

grid_b=np.linspace(b_range[0],b_range[1],num=grid_num,endpoint=True)
grid_m=np.linspace(m_range[0],m_range[1],num=grid_m_num,endpoint=True)

cell_number=[3,4,5,6,7,8]
end_lc_list=[3,7,13,23,37,58]                   # 13[5], 23[6], 37[7]; 58[8]; 87[9]; 128[10]     # how many lcs we care

"add a list [[maturity size=3],[],[],;;;]"
data_game=[]
max_size=8
for i in range(5,max_size-2):
	n=cell_number[i]
	end_lc=end_lc_list[i]
	#-------------------------------------------------------------------------------
	"read max and min dps from data into result []"
	num_k=len(k_cluster_list)
	for b_cluster in range(18,19):                             # b=10
		for m_cluster in range(1,2):                             # y axes points = a
			for k_cluster in k_cluster_list:
				for i_th in lc_list:                            # read all lcs data = growth rate
					with open('../data/data_threshold_effect/%d_%d_%d_%d.txt'%(b_cluster,m_cluster,k_cluster,i_th), "r") as file:
						nan=float(np.nan)
						inf=np.inf
						grate = eval(file.readline())          # read growth rate
					data_game.append([b_cluster,m_cluster,k_cluster,i_th,grate])
data_game=np.array(data_game)

#-------------------------------------------------------------------------------
"read data of both effects across k,"
grid_num=7                                   # grid size in figure, so check figure.py file first
t_ratio=1.5

result_list=[[],[]]                          # to store data of each k value
optimal_matrix_list=[np.zeros(shape=(grid_num,grid_num))*np.nan for i in range(2)]

"add a list [[maturity size=3],[],[],;;;]"
data_both=[]
for t_pterb_cluster in range(2,3):                        # perturbation at n=3
	for k_cluster in range(0,grid_num):                    # y axes points = a
		all_list=[]                                    # algin all lcs' growth rates
		for i_th in lc_list:                   # read all lcs data = growth rate

			with open('../data/data_size_threshold_chi_%s/%d_%d_%d.txt'%(t_ratio,t_pterb_cluster,k_cluster,i_th), "r") as file:
				nan=float(np.nan)
				inf=np.inf
				grate = eval(file.readline())          # read growth rate
			data_both.append(np.array([t_ratio,t_pterb_cluster,k_cluster,i_th,grate]))

data_both=np.array(data_both)

#================= draw figures====================================================
c = np.arange(1, 58 + 1)

norm = mpl.colors.Normalize(vmin=c.min(), vmax=c.max())
cmap = mpl.cm.ScalarMappable(norm=norm, cmap=mpl.cm.jet)
cmap.set_array([])

fig, ax = plt.subplots(1, 1, figsize=(4, 4.5))
fig.subplots_adjust(top=0.82)
ax.set_axisbelow(True)
k_list=[1,2,3,4,5,6,7]

"size figure"
s_size=1
for i in lc_list:
	y=data_size[np.where(data_size[:,1]==i)][0,2]
	ax.scatter(k_list, [y for i in range(7)], s=s_size,marker="^",c=cmap.to_rgba(i + 1),alpha=0.9)
	if s_size==1:
		ax.plot(k_list, [y for i in range(7)], ':',label="Size effects",linewidth=1.5,c=cmap.to_rgba(i + 1),alpha=0.95)
	else:
		ax.plot(k_list, [y for i in range(7)], ':',linewidth=1.5,c=cmap.to_rgba(i + 1),alpha=0.95)

	"game figure"
	y_game=data_game[np.where(data_game[:,3]==i)][:,4]
	ax.scatter(k_list, y_game, s=s_size,marker="^",c=cmap.to_rgba(i + 1),alpha=0.9)
	if s_size==1:
		ax.plot(k_list, y_game, '--',label="Threshold effets",linewidth=2,c=cmap.to_rgba(i + 1),alpha=0.95)
	else:
		ax.plot(k_list, y_game, '--',linewidth=2,c=cmap.to_rgba(i + 1),alpha=0.95)

	"size and game figure"
	y_both=data_both[np.where(data_both[:,3]==i)][:,4]
	ax.scatter(k_list, y_both, s=s_size,marker="^",c=cmap.to_rgba(i + 1),alpha=0.9)

	if  s_size==1:
		ax.plot(k_list, y_both, '-',label="Combined effects",linewidth=2,c=cmap.to_rgba(i + 1),alpha=0.75)
	else:
		ax.plot(k_list, y_both, '-',linewidth=2,c=cmap.to_rgba(i + 1),alpha=0.75)

	s_size+=1

leg=ax.legend(frameon=False,loc='upper center', bbox_to_anchor=(0.7, 0.98), shadow=True, ncol=1, fontsize=10)
for i in range(3):
	leg.legendHandles[i].set_color('k')

ax.annotate(lc_add(lc_list[0]), (2.3, 1.65),color=cmap.to_rgba(5 + 1))
ax.annotate(lc_add(lc_list[1]), (4.5, 1.85),color=cmap.to_rgba(56 + 1))

ax.set_xlabel(r"Contribution threshold $k$",fontsize=12)
ax.set_xticks([i for i in range(1,8)], minor=False)
ax.set_ylim([0.65, 3.5])
ax.set_ylabel(r"Population growth rate $\lambda$",fontsize=12)

ax.spines['right'].set_visible(False)
ax.spines['top'].set_visible(False)
ax.spines['left'].set_visible(False)
ax.yaxis.set_ticks_position('left')

plt.show();


# Main entrance of GAIL
import numpy as np
import torch
import torch.nn.functional as F
import gym
import safety_gym
import time
import os.path as osp

from torch.distributions.categorical import Categorical

from neural_nets import Discriminator, ActorCritic, ValorDiscriminator, count_vars

from utils import BufferS, BufferT, BufferActor, VALORBuffer
from utils import mpi_fork, proc_id, num_procs, EpochLogger,\
    average_gradients, sync_all_params

def ppo_penalized(env_fn,
            actor_critic=ActorCritic,
            ac_kwargs=dict(),
            seed=0,
            episodes_per_epoch=40,
            epochs=500,
            gamma=0.99,
            lam=0.97,
            pi_lr=3e-4,
            vf_lr=1e-3,
            train_v_iters=80,
            max_ep_len=1000,
            logger_kwargs=dict(),
            save_freq=10):

    logger = EpochLogger(**logger_kwargs)
    logger.save_config(locals())

    seed += 10000 * proc_id()
    torch.manual_seed(seed)
    np.random.seed(seed)

    env = env_fn()
    obs_dim = env.observation_space.shape
    act_dim = env.action_space.shape

    ac_kwargs['action_space'] = env.action_space

    # Models
    ac = actor_critic(input_dim=obs_dim[0], **ac_kwargs)

    # Set up model saving
    logger.setup_pytorch_saver(ac)

    # Buffers
    local_episodes_per_epoch = int(episodes_per_epoch / num_procs())
    buf = BufferActor(obs_dim[0], act_dim[0], local_episodes_per_epoch, max_ep_len)

    # Count variables
    var_counts = tuple(count_vars(module) for module in [ac.policy, ac.value_f])
    print("POLICY GRADIENT")
    logger.log('\nNumber of parameters: \t pi: %d, \t v: %d\n' % var_counts)

    # Optimizers
    train_pi = torch.optim.Adam(ac.policy.parameters(), lr=pi_lr)
    train_v = torch.optim.Adam(ac.value_f.parameters(), lr=vf_lr)

    # Parameters Sync
    sync_all_params(ac.parameters())

    def update(e):
        obs, act, adv, ret, lgp_old = [torch.Tensor(x) for x in buf.retrieve_all()]

        # Policy
        _, lgp, _ = ac.policy(obs, act)
        entropy = (-lgp).mean()

        # Policy loss # policy gradient term + entropy term
        pi_loss = -(lgp * adv).mean()

        # Train policy
        train_pi.zero_grad()
        pi_loss.backward()
        average_gradients(train_pi.param_groups)
        train_pi.step()

        # Value function
        v = ac.value_f(obs)
        v_l_old = F.mse_loss(v, ret)
        for _ in range(train_v_iters):
            v = ac.value_f(obs)
            v_loss = F.mse_loss(v, ret)

            # Value function train
            train_v.zero_grad()
            v_loss.backward()
            average_gradients(train_v.param_groups)
            train_v.step()

        # Log the changes
        _, lgp, _, v = ac(obs, act)
        entropy_new = (-lgp).mean()
        pi_loss_new = -(lgp * adv).mean()
        v_loss_new = F.mse_loss(v, ret)
        kl = (lgp_old - lgp).mean()
        logger.store(LossPi=pi_loss, LossV=v_l_old, DeltaLossPi=(pi_loss_new - pi_loss),
                     DeltaLossV=(v_loss_new - v_l_old), Entropy=entropy, KL=kl)

    start_time = time.time()
    o, r, d, ep_ret, ep_len = env.reset(), 0, False, 0, 0
    total_t = 0

    for epoch in range(epochs):
        ac.eval()
        # Policy rollout
        for _ in range(local_episodes_per_epoch):
            for _ in range(max_ep_len):
                obs = torch.Tensor(o.reshape(1, -1))
                a, _, lopg_t, v_t = ac(obs)

                buf.store(o, a.detach().numpy(), r, v_t.item(), lopg_t.detach().numpy())
                logger.store(VVals=v_t)

                o, r, d, _ = env.step(a.detach().numpy()[0])
                ep_ret += r
                ep_len += 1
                total_t += 1

                terminal = d or (ep_len == max_ep_len)
                if terminal:
                    buf.end_episode()
                    logger.store(EpRet=ep_ret, EpLen=ep_len)
                    o, r, d, ep_ret, ep_len = env.reset(), 0, False, 0, 0

        if (epoch % save_freq == 0) or (epoch == epochs - 1):
            # logger._torch_save(ac, fname="expert_torch_save.pt")
            # logger._torch_save(ac, fname="model.pt")
            logger.save_state({'env': env}, None, None)

        # Update
        ac.train()

        update(epoch)

        # Log
        logger.log_tabular('Epoch', epoch)
        logger.log_tabular('EpRet', with_min_and_max=True)
        logger.log_tabular('EpLen', average_only=True)
        logger.log_tabular('VVals', with_min_and_max=True)
        logger.log_tabular('TotalEnvInteracts', total_t)
        logger.log_tabular('LossPi', average_only=True)
        logger.log_tabular('DeltaLossPi', average_only=True)
        logger.log_tabular('LossV', average_only=True)
        logger.log_tabular('DeltaLossV', average_only=True)
        logger.log_tabular('Entropy', average_only=True)
        logger.log_tabular('KL', average_only=True)
        logger.log_tabular('Time', time.time() - start_time)
        logger.dump_tabular()

def gail_penalized(env_fn, actor_critic=ActorCritic, ac_kwargs=dict(),
         disc=Discriminator,
         dc_kwargs=dict(), seed=0,
         episodes_per_epoch=40,
         epochs=500,
         gamma=0.99, lam=0.97,
         pi_lr=3e-3, vf_lr=3e-3, dc_lr=5e-4, train_v_iters=80, train_dc_iters=80,
         max_ep_len=1000, logger_kwargs=dict(), save_freq=10):
    l_lam = 0  # balance two loss term

    print("starting now")

    logger = EpochLogger(**logger_kwargs)
    logger.save_config(locals())

    seed += 10000 * proc_id()
    torch.manual_seed(seed)
    np.random.seed(seed)

    env = env_fn()
    obs_dim = env.observation_space.shape
    act_dim = env.action_space.shape

    ac_kwargs['action_space'] = env.action_space

    # Models
    ac = actor_critic(input_dim=obs_dim[0], **ac_kwargs)
    disc = disc(input_dim=obs_dim[0], **dc_kwargs)

    # Set up model saving
    logger.setup_pytorch_saver([ac, disc])

    # TODO: Load expert policy here
    expert = actor_critic(input_dim=obs_dim[0], **ac_kwargs)
    # expert_name = "expert_torch_save.pt"
    expert_name = "model.pt"
    # expert = torch.load(osp.join(logger_kwargs['output_dir'],'pyt_save' , expert_name))
    expert = torch.load('/home/tyna/Documents/openai/research-project/data/anonymous-expert/anonymous-expert_s0/pyt_save/model.pt')

    print('RUNNING GAIL')

    # Buffers
    local_episodes_per_epoch = int(episodes_per_epoch / num_procs())
    buff_s = BufferS(obs_dim[0], act_dim[0], local_episodes_per_epoch, max_ep_len)
    buff_t = BufferT(obs_dim[0], act_dim[0], local_episodes_per_epoch, max_ep_len)

    # Count variables
    var_counts = tuple(count_vars(module) for module in [ac.policy, ac.value_f, disc.policy])
    print("GAIL")
    logger.log('\nNumber of parameters: \t pi: %d, \t v: %d, \t d: %d\n' % var_counts)

    # Optimizers
    train_pi = torch.optim.Adam(ac.policy.parameters(), lr=pi_lr)
    train_v = torch.optim.Adam(ac.value_f.parameters(), lr=vf_lr)
    train_dc = torch.optim.Adam(disc.policy.parameters(), lr=dc_lr)

    # Parameters Sync
    sync_all_params(ac.parameters())
    sync_all_params(disc.parameters())

    def update(e):
        obs_s, act, adv, ret, lgp_old = [torch.Tensor(x) for x in buff_s.retrieve_all()]
        obs_t, _ = [torch.Tensor(x) for x in buff_t.retrieve_all()]

        # Policy
        _, lgp, _ = ac.policy(obs_s, act)
        entropy = (-lgp).mean()

        # Policy loss
        # policy gradient term + entropy term
        pi_loss = -(lgp * adv).mean() - l_lam * entropy

        # Train policy
        if e > 10:
            train_pi.zero_grad()
            pi_loss.backward()
            average_gradients(train_pi.param_groups)
            train_pi.step()

        # Value function
        v = ac.value_f(obs_s)
        v_l_old = F.mse_loss(v, ret)

        for _ in range(train_v_iters):
            v = ac.value_f(obs_s)
            v_loss = F.mse_loss(v, ret)

            # Value function train
            train_v.zero_grad()
            v_loss.backward()
            average_gradients(train_v.param_groups)
            train_v.step()

        # Discriminator
        gt1 = torch.ones(obs_s.size()[0], dtype=torch.int)
        gt2 = torch.zeros(obs_t.size()[0], dtype=torch.int)
        _, lgp_s, _ = disc(obs_s, gt=gt1)
        _, lgp_t, _ = disc(obs_t, gt=gt2)
        dc_loss_old = - lgp_s.mean() - lgp_t.mean()

        for _ in range(train_dc_iters):
            _, lgp_s, _ = disc(obs_s, gt=gt1)
            _, lgp_t, _ = disc(obs_t, gt=gt2)
            dc_loss = - lgp_s.mean() - lgp_t.mean()

            # Discriminator train
            train_dc.zero_grad()
            dc_loss.backward()
            average_gradients(train_dc.param_groups)
            train_dc.step()

        _, lgp_s, _ = disc(obs_s, gt=gt1)
        _, lgp_t, _ = disc(obs_t, gt=gt2)
        dc_loss_new = - lgp_s.mean() - lgp_t.mean()

        # Log the changes
        _, lgp, _, v = ac(obs, act)
        entropy_new = (-lgp).mean()
        pi_loss_new = -(lgp * adv).mean() - l_lam * entropy
        v_loss_new = F.mse_loss(v, ret)
        kl = (lgp_old - lgp).mean()
        logger.store(LossPi=pi_loss, LossV=v_l_old, LossDC=dc_loss_old, DeltaLossPi=(pi_loss_new - pi_loss),
                     DeltaLossV=(v_loss_new - v_l_old), DeltaLossDC=(dc_loss_new - dc_loss_old),
                     DeltaEnt=(entropy_new - entropy),
                     Entropy=entropy, KL=kl)

    start_time = time.time()
    o, r, sdr, d, ep_ret, ep_sdr, ep_len = env.reset(), 0, 0, False, 0, 0, 0
    total_t = 0

    ep_len_t = 0
    for epoch in range(epochs):
        ac.eval()
        disc.eval()
        # We recognize the probability term of index [0] correspond to the teacher's policy

        # Student's policy rollout
        for _ in range(local_episodes_per_epoch):
            for _ in range(max_ep_len):
                obs = torch.Tensor(o.reshape(1, -1))
                a, _, lopg_t, v_t = ac(obs)

                buff_s.store(o, a.detach().numpy(), r, sdr, v_t.item(), lopg_t.detach().numpy())
                logger.store(VVals=v_t)

                o, r, d, _ = env.step(a.detach().numpy()[0])
                _, sdr, _ = disc(torch.Tensor(o.reshape(1, -1)), gt=torch.Tensor([0]))
                if sdr < -4:  # Truncate rewards
                    sdr = -4
                ep_ret += r
                ep_sdr += sdr
                ep_len += 1
                total_t += 1

                terminal = d or (ep_len == max_ep_len)
                if terminal:
                    buff_s.end_episode()
                    logger.store(EpRetS=ep_ret, EpLenS=ep_len, EpSdrS=ep_sdr)
                    print("Student Episode Return: \t", ep_ret)
                    o, r, sdr, d, ep_ret, ep_sdr, ep_len = env.reset(), 0, 0, False, 0, 0, 0

        # Teacher's policy rollout
        for _ in range(local_episodes_per_epoch):
            for _ in range(max_ep_len):
                obs = torch.Tensor(o.reshape(1, -1))
                a, _, _, _ = expert(obs)

                buff_t.store(o, a.detach().numpy(), r)

                o, r, d, _ = env.step(a.detach().numpy()[0])
                ep_ret += r
                ep_len += 1
                total_t += 1

                terminal = d or (ep_len == max_ep_len)
                if terminal:
                    buff_t.end_episode()
                    logger.store(EpRetT=ep_ret, EpLenT=ep_len)
                    print("Teacher Episode Return: \t", ep_ret)
                    o, r, d, ep_ret, ep_len = env.reset(), 0, False, 0, 0

        if (epoch % save_freq == 0) or (epoch == epochs - 1):
            logger.save_state({'env': env}, [ac, disc], None)

        # Update
        ac.train()
        disc.train()

        update(epoch)

        # Log
        logger.log_tabular('Epoch', epoch)
        logger.log_tabular('EpRetS', with_min_and_max=True)
        logger.log_tabular('EpSdrS', with_min_and_max=True)
        logger.log_tabular('EpLenS', average_only=True)
        logger.log_tabular('EpRetT', with_min_and_max=True)
        logger.log_tabular('EpLenT', average_only=True)
        logger.log_tabular('VVals', with_min_and_max=True)
        logger.log_tabular('TotalEnvInteracts', total_t)
        logger.log_tabular('LossPi', average_only=True)
        logger.log_tabular('DeltaLossPi', average_only=True)
        logger.log_tabular('LossV', average_only=True)
        logger.log_tabular('DeltaLossV', average_only=True)
        logger.log_tabular('LossDC', average_only=True)
        logger.log_tabular('DeltaLossDC', average_only=True)
        logger.log_tabular('Entropy', average_only=True)
        logger.log_tabular('DeltaEnt', average_only=True)
        logger.log_tabular('KL', average_only=True)
        logger.log_tabular('Time', time.time() - start_time)
        logger.dump_tabular()

def valor_penalized(env_fn, actor_critic=ActorCritic, ac_kwargs=dict(),
          disc=Discriminator, dc_kwargs=dict(), seed=0,
          episodes_per_epoch=40,
          epochs=50, gamma=0.99, pi_lr=3e-4, vf_lr=1e-3, dc_lr=5e-4,
          train_v_iters=80, train_dc_iters=10,
          train_dc_interv=10,
          lam=0.97, max_ep_len=1000, logger_kwargs=dict(), con_dim=5, save_freq=10, k=1):

    logger = EpochLogger(**logger_kwargs)
    logger.save_config(locals())

    seed += 10000 * proc_id()
    torch.manual_seed(seed)
    np.random.seed(seed)

    env = env_fn()
    obs_dim = env.observation_space.shape
    act_dim = env.action_space.shape

    ac_kwargs['action_space'] = env.action_space

    # Model
    ac = actor_critic(input_dim=obs_dim[0] + con_dim, **ac_kwargs)
    disc = disc(input_dim=obs_dim[0], context_dim=con_dim, **dc_kwargs)

    # Set up model saving
    logger.setup_pytorch_saver([ac, disc])

    # Buffer
    local_episodes_per_epoch = int(episodes_per_epoch / num_procs())
    buffer = VALORBuffer(con_dim, obs_dim[0], act_dim[0], local_episodes_per_epoch, max_ep_len, train_dc_interv)

    # Count variables
    var_counts = tuple(count_vars(module) for module in
                       [ac.policy, ac.value_f, disc.policy])
    logger.log('\nNumber of parameters: \t pi: %d, \t v: %d, \t d: %d\n' % var_counts)

    # Optimizers
    train_pi = torch.optim.Adam(ac.policy.parameters(), lr=pi_lr)
    train_v = torch.optim.Adam(ac.value_f.parameters(), lr=vf_lr)
    train_dc = torch.optim.Adam(disc.policy.parameters(), lr=dc_lr)

    # Parameters Sync
    sync_all_params(ac.parameters())
    sync_all_params(disc.parameters())

    def update(e):
        obs, act, adv, pos, ret, logp_old = [torch.Tensor(x) for x in buffer.retrieve_all()]

        # Policy
        _, logp, _ = ac.policy(obs, act)
        entropy = (-logp).mean()

        # Policy loss
        pi_loss = -(logp * (k * adv + pos)).mean()

        # Train policy
        train_pi.zero_grad()
        pi_loss.backward()
        average_gradients(train_pi.param_groups)
        train_pi.step()

        # Value function
        v = ac.value_f(obs)
        v_l_old = F.mse_loss(v, ret)
        for _ in range(train_v_iters):
            v = ac.value_f(obs)
            v_loss = F.mse_loss(v, ret)

            # Value function train
            train_v.zero_grad()
            v_loss.backward()
            average_gradients(train_v.param_groups)
            train_v.step()

        # Discriminator
        if (e + 1) % train_dc_interv == 0:
            print('Discriminator Update!')
            con, s_diff = [torch.Tensor(x) for x in buffer.retrieve_dc_buff()]
            _, logp_dc, _ = disc(s_diff, con)
            d_l_old = -logp_dc.mean()

            # Discriminator train
            for _ in range(train_dc_iters):
                _, logp_dc, _ = disc(s_diff, con)
                d_loss = -logp_dc.mean()
                train_dc.zero_grad()
                d_loss.backward()
                average_gradients(train_dc.param_groups)
                train_dc.step()

            _, logp_dc, _ = disc(s_diff, con)
            dc_l_new = -logp_dc.mean()
        else:
            d_l_old = 0
            dc_l_new = 0

        # Log the changes
        _, logp, _, v = ac(obs, act)
        pi_l_new = -(logp * (k * adv + pos)).mean()
        v_l_new = F.mse_loss(v, ret)
        kl = (logp_old - logp).mean()
        logger.store(LossPi=pi_loss, LossV=v_l_old, KL=kl, Entropy=entropy, DeltaLossPi=(pi_l_new - pi_loss),
                     DeltaLossV=(v_l_new - v_l_old), LossDC=d_l_old, DeltaLossDC=(dc_l_new - d_l_old))
        # logger.store(Adv=adv.reshape(-1).numpy().tolist(), Pos=pos.reshape(-1).numpy().tolist())

    start_time = time.time()
    o, r, d, ep_ret, ep_len = env.reset(), 0, False, 0, 0
    context_dist = Categorical(logits=torch.Tensor(np.ones(con_dim)))
    total_t = 0

    for epoch in range(epochs):
        ac.eval()
        disc.eval()
        for _ in range(local_episodes_per_epoch):
            c = context_dist.sample()
            c_onehot = F.one_hot(c, con_dim).squeeze().float()
            for _ in range(max_ep_len):
                concat_obs = torch.cat([torch.Tensor(o.reshape(1, -1)), c_onehot.reshape(1, -1)], 1)
                a, _, logp_t, v_t = ac(concat_obs)

                buffer.store(c, concat_obs.squeeze().detach().numpy(), a.detach().numpy(), r, v_t.item(),
                             logp_t.detach().numpy())
                logger.store(VVals=v_t)

                o, r, d, _ = env.step(a.detach().numpy()[0])
                ep_ret += r
                ep_len += 1
                total_t += 1

                terminal = d or (ep_len == max_ep_len)
                if terminal:
                    dc_diff = torch.Tensor(buffer.calc_diff()).unsqueeze(0)
                    con = torch.Tensor([float(c)]).unsqueeze(0)
                    _, _, log_p = disc(dc_diff, con)
                    buffer.end_episode(log_p.detach().numpy())
                    logger.store(EpRet=ep_ret, EpLen=ep_len)
                    o, r, d, ep_ret, ep_len = env.reset(), 0, False, 0, 0

        if (epoch % save_freq == 0) or (epoch == epochs - 1):
            logger.save_state({'env': env}, [ac, disc], None)

        # Update
        ac.train()
        disc.train()

        update(epoch)

        # Log
        logger.log_tabular('Epoch', epoch)
        logger.log_tabular('EpRet', with_min_and_max=True)
        logger.log_tabular('EpLen', average_only=True)
        logger.log_tabular('VVals', with_min_and_max=True)
        logger.log_tabular('TotalEnvInteracts', total_t)
        logger.log_tabular('LossPi', average_only=True)
        logger.log_tabular('DeltaLossPi', average_only=True)
        logger.log_tabular('LossV', average_only=True)
        logger.log_tabular('DeltaLossV', average_only=True)
        logger.log_tabular('LossDC', average_only=True)
        logger.log_tabular('DeltaLossDC', average_only=True)
        logger.log_tabular('Entropy', average_only=True)
        logger.log_tabular('KL', average_only=True)
        logger.log_tabular('Time', time.time() - start_time)
        logger.dump_tabular()

if __name__ == '__main__':
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument('--env', type=str, default='Safexp-PointGoal1-v0')
    parser.add_argument('--hid', type=int, default=64)
    parser.add_argument('--l', type=int, default=2)
    parser.add_argument('--gamma', type=float, default=0.99)
    parser.add_argument('--lam', type=float, default=0.97)
    parser.add_argument('--seed', '-s', type=int, default=0)
    parser.add_argument('--cpu', type=int, default=1)
    parser.add_argument('--episodes-per-epoch', type=int, default=5)
    # parser.add_argument('--episodes-per-epoch', type=int, default=40)
    parser.add_argument('--epochs', type=int, default=1000)
    parser.add_argument('--exp_name', type=str, default='valor-anonymous-expert')
    parser.add_argument('--con', type=int, default=5)
    args = parser.parse_args()

    mpi_fork(args.cpu)

    from utils import setup_logger_kwargs

    logger_kwargs = setup_logger_kwargs(args.exp_name, args.seed)

    # ppo_penalized(lambda: gym.make(args.env), actor_critic=ActorCritic, ac_kwargs=dict(hidden_dims=[args.hid]*args.l),
    #     gamma=args.gamma, lam=args.lam, seed=args.seed, episodes_per_epoch=args.episodes_per_epoch,
    #     epochs=args.epochs, logger_kwargs=logger_kwargs)

    # gail_penalized(lambda: gym.make(args.env), actor_critic=ActorCritic, ac_kwargs=dict(hidden_dims=[args.hid] * args.l),
    #      disc=Discriminator, dc_kwargs=dict(hidden_dims=[args.hid] * args.l), gamma=args.gamma, lam=args.lam,
    #      seed=args.seed, episodes_per_epoch=args.episodes_per_epoch, epochs=args.epochs, logger_kwargs=logger_kwargs)

    valor_penalized(lambda: gym.make(args.env), actor_critic=ActorCritic, ac_kwargs=dict(hidden_dims=[args.hid] * args.l),
          disc=ValorDiscriminator, dc_kwargs=dict(hidden_dims=args.hid),
          gamma=args.gamma, seed=args.seed, episodes_per_epoch=args.episodes_per_epoch, epochs=args.epochs,
          logger_kwargs=logger_kwargs, con_dim=args.con)