#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""

Created on Sun Aug 23 12:11:15 2020

Modified from the cornstover biorefinery constructed in Cortes-Peña et al., 2020,
with modification of fermentation system for 2,3-Butanediol instead of the original ethanol

[1] Cortes-Peña et al., BioSTEAM: A Fast and Flexible Platform for the Design,
    Simulation, and Techno-Economic Analysis of Biorefineries under Uncertainty.
    ACS Sustainable Chem. Eng. 2020, 8 (8), 3302–3310.
    https://doi.org/10.1021/acssuschemeng.9b07040.

All units are explicitly defined here for transparency and easy reference

[1] Cortes-Peña et al., BioSTEAM: A Fast and Flexible Platform for the Design,
    Simulation, and Techno-Economic Analysis of Biorefineries under Uncertainty.
    ACS Sustainable Chem. Eng. 2020, 8 (8), 3302–3310.
    https://doi.org/10.1021/acssuschemeng.9b07040.

Naming conventions:
    D = Distillation column
    F = Flash tank
    H = Heat exchange
    M = Mixer
    P = Pump (including conveying belt)
    R = Reactor
    S = Splitter (including solid/liquid separator)
    T = Tank or bin for storage
    U = Other units
    PS = Process specificiation, not physical units, but for adjusting streams

Processes:
    100: Feedstock preprocessing
    200: Pretreatment
    300: Conversion
    400: Separation
    500: Wastewater treatment
    600: Facilities

@author: sarangbhagwat
"""

# %% Setup

import biosteam as bst
import thermosteam as tmo
import flexsolve as flx
import numpy as np
from biosteam import main_flowsheet as F
from biorefineries import BST222
from biorefineries.cornstover import CellulosicEthanolTEA
from biosteam import System
from thermosteam import Stream
from biorefineries import BDO as bdo
from biorefineries.BDO import units, facilities
from biorefineries.BDO._process_specification import ProcessSpecification
from biorefineries.BDO.process_settings import price, CFs
from biorefineries.BDO.utils import find_split, splits_df, baseline_feedflow
from biorefineries.BDO.chemicals_data import BDO_chemicals, chemical_groups, \
                                soluble_organics, combustibles
from biorefineries.BDO.tea import BDOTEA
from biorefineries.HP.lca import LCA

bst.speed_up()
flowsheet = bst.Flowsheet('BDO')
bst.main_flowsheet.set_flowsheet(flowsheet)
bst.System.default_relative_molar_tolerance = 1e-4
bst.System.default_maxiter = 100
bst.System.default_molar_tolerance = 0.1

# Speeds up ShortcutDistillation
bst.units.ShortcutColumn.minimum_guess_distillate_recovery = 0

# Baseline cost year is 2016
bst.CE = 541.7
# _labor_2007to2016 = 22.71 / 19.55

# Set default thermo object for the system
tmo.settings.set_thermo(BDO_chemicals)

BDO_sys = bdo.create_system_broth()
u = flowsheet.unit
feedstock = BDO_sys.ins[0]
# feedstock.mol[:] = baseline_feedflow[:]
conc_aqueous_broth, = BDO_sys.outs
get_flow_tpd = lambda: (feedstock.F_mass-feedstock.imass['H2O'])*24/907.185
u = F.unit
s = F.stream
# globals().update(F.to_dict())

# %%
# =============================================================================
# TEA
# =============================================================================

# BDO_tea = bst.CombinedTEA([BDO_no_BT_tea, BT_tea], IRR=0.10)
BDO_tea = CellulosicEthanolTEA(system=BDO_sys, IRR=0.10, duration=(2016, 2046),
        depreciation='MACRS7', income_tax=0.21, operating_days=0.9*365,
        lang_factor=None, construction_schedule=(0.08, 0.60, 0.32),
        startup_months=3, startup_FOCfrac=1, startup_salesfrac=0.5,
        startup_VOCfrac=0.75, WC_over_FCI=0.05,
        finance_interest=0.08, finance_years=10, finance_fraction=0.4,
        # biosteam Splitters and Mixers have no cost,
        # cost of all wastewater treatment units are included in WWT_cost,
        # BT is not included in this TEA
        OSBL_units=bst.get_OSBL(BDO_sys.units),
        warehouse=0.04, site_development=0.09, additional_piping=0.045,
        proratable_costs=0.10, field_expenses=0.10, construction=0.20,
        contingency=0.10, other_indirect_costs=0.10,
        labor_cost=3212962*get_flow_tpd()/2205,
        labor_burden=0.90, property_insurance=0.007, maintenance=0.03,
        steam_power_depreciation='MACRS20', boiler_turbogenerator=u.BT701)

# sub_units = BDO_sys.units[:BDO_sys.units.index(u.R401)]
# sub_sys = bst.System(sub_units)
# sub_tea = CellulosicEthanolTEA(system=sub_sys, IRR=0.10, duration=(2016, 2046),
#         depreciation='MACRS7', income_tax=0.21, operating_days=0.9*365,
#         lang_factor=None, construction_schedule=(0.08, 0.60, 0.32),
#         startup_months=3, startup_FOCfrac=1, startup_salesfrac=0.5,
#         startup_VOCfrac=0.75, WC_over_FCI=0.05,
#         finance_interest=0.08, finance_years=10, finance_fraction=0.4,
#         # biosteam Splitters and Mixers have no cost,
#         # cost of all wastewater treatment units are included in WWT_cost,
#         # BT is not included in this TEA
#         OSBL_units=bst.get_OSBL(sub_sys.units),
#         warehouse=0.04, site_development=0.09, additional_piping=0.045,
#         proratable_costs=0.10, field_expenses=0.10, construction=0.20,
#         contingency=0.10, other_indirect_costs=0.10,
#         labor_cost=3212962*get_flow_tpd()/2205,
#         labor_burden=0.90, property_insurance=0.007, maintenance=0.03,
#         steam_power_depreciation='MACRS20', boiler_turbogenerator=u.BT)

area_names = [
    'feedstock',
    'pretreatment',
    'conversion',
    'separation',
    'wastewater',
    'storage',
    'co-heat and power',
    'cooling tower and chilled water package',
    'other facilities',
    'heat exchanger network',
]
# u.CWP901.ID = 'CWP802'
for ui in u:
    if type(ui) == bst.ChilledWaterPackage:
        ui.ID = 'CWP802' # group with CT for system cooling demand
        break
unit_groups = bst.UnitGroup.group_by_area(BDO_sys.units)
for i, j in zip(unit_groups, area_names): i.name = j
for i in unit_groups: i.autofill_metrics(shorthand=True,
                                         electricity_production=True,
                                         material_cost=True)
for i in unit_groups:
    if i.name == 'storage' or i.name=='other facilities' or i.name == 'cooling tower and chilled water package':
        i.metrics[-1].getter = lambda: 0. # Material cost
    if i.name == 'cooling tower and chilled water package':
        i.metrics[1].getter = lambda: 0. # Cooling duty
HXN = None
for HXN_group in unit_groups:
    if HXN_group.name == 'heat exchanger network':
        HXN_group.filter_savings = False
        HXN = HXN_group.units[0]
        assert isinstance(HXN, bst.HeatExchangerNetwork)

unit_groups_dict = {}
for i in unit_groups:
    unit_groups_dict[i.name] = i
# HXN.force_ideal_thermo = True
CT = u.CT801
BT = u.BT701
CWP = u.CWP802

# %%
# =============================================================================
# Simulate system and get results
# =============================================================================

TEA_feeds = set([i for i in BDO_sys.feeds if i.price])

TEA_products = set([i for i in BDO_sys.products if i.price])
conc_aqueous_broth.price = 0.
def get_conc_aqueous_broth_MPSP():
    BDO_sys.simulate()

    for i in range(3):
        conc_aqueous_broth.price = BDO_tea.solve_price(conc_aqueous_broth)
    return conc_aqueous_broth.price

M401, F401 = flowsheet('M401'), flowsheet('F401')
def set_target_BDO_x(target_BDO_x):
    F401.target_BDO_x = target_BDO_x
    M401.target_BDO_x = target_BDO_x

def get_x(chem_ID, stream):
    return stream.imol[chem_ID]/sum(stream.imol[tuple([i.ID for i in stream.vle_chemicals])])

# get_conc_aqueous_broth_MPSP()

# R301 = F('R301') # Fermentor
seed_train_system = bst.System('seed_train_system', path=(u.S302, u.R303, u.T301))

# yearly_production = 125000 # ton/yr
spec = ProcessSpecification(
    evaporator = u.F301,
    evaporator_pump = u.F301_P,
    pump = u.M304_H_P,
    mixer = u.M304,
    heat_exchanger = u.M304_H,
    seed_train_system = seed_train_system,
    reactor= u.R302,
    reaction_name='fermentation_reaction',
    substrates=('Xylose', 'Glucose'),
    products=('BDO',),
    spec_1=0.36,
    spec_2=109.9,
    spec_3=1.,
    xylose_utilization_fraction = 0.80,
    feedstock = feedstock,
    dehydration_reactor = None,
    byproduct_streams = [],
    HXN = u.HXN1001,
    tolerable_HXN_energy_balance_percent_error=5,
    maximum_inhibitor_concentration = 1.,
    # pre_conversion_units = process_groups_dict['feedstock_group'].units + process_groups_dict['pretreatment_group'].units + [u.H301], # if the line below does not work (depends on BioSTEAM version)
    # pre_conversion_units = BDO_sys.split(u.F301.ins[0])[0],
    baseline_titer = 54.8,
    feedstock_mass = feedstock.F_mass,
    pretreatment_reactor = u.R201)

spec.load_spec_1 = spec.load_yield
spec.load_spec_2 = spec.load_titer
spec.load_spec_3 = spec.load_productivity

# spec = ProcessSpecification(
#     evaporator = u.F301,
#     mixer = u.M304,
#     reactor=u.R302,
#     reaction_name='fermentation_reaction',
#     substrates=('Xylose', 'Glucose'),
#     products=('BDO',),
#     spec_1=0.8,
#     spec_2=109.9,
#     spec_3=1.,
#     path = (u.M304_H, u.M304_H_P),
#     xylose_utilization_fraction = 0.80,
#     feedstock = feedstock,
#     dehydration_reactor = u.R401,
#     byproduct_streams = [s.isobutanol],
#     evaporator_pump = u.F301_P)

path = (u.F301, u.R302)
@np.vectorize
def calculate_titer(V):
    u.F301.V = V
    for i in path: i._run()
    return spec._calculate_titer()

@np.vectorize
def calculate_MPSP(V):
    u.F301.V = V
    BDO_sys.simulate()
    MPSP = conc_aqueous_broth.price = BDO_tea.solve_price(conc_aqueous_broth)
    return MPSP

# vapor_fractions = np.linspace(0.20, 0.80)
# titers = calculate_titer(vapor_fractions)
# MPSPs = calculate_MPSP(vapor_fractions)
# import matplotlib.pyplot as plt
# plt.plot(vapor_fractions, titers)
# plt.show()

# plt.plot(titers, MPSPs)
# plt.show()

# %%

# =============================================================================
# Life cycle analysis (LCA), waste disposal emission not included
# =============================================================================

# 100-year global warming potential (GWP) from material flows
LCA_streams = TEA_feeds.copy()
LCA_stream = Stream('LCA_stream', units='kg/hr')

get_Isobutanol_GWP = lambda: 0.
get_Isobutanol_FEC = lambda: 0.

def get_material_GWP():
    LCA_stream.mass = sum(i.mass for i in LCA_streams)
    chemical_GWP = LCA_stream.mass*CFs['GWP_CF_stream'].mass
    # feedstock_GWP = feedstock.F_mass*CFs['GWP_CFs']['Corn stover']
    return chemical_GWP.sum()/conc_aqueous_broth.F_mass

# GWP from combustion of non-biogenic carbons
get_non_bio_GWP = lambda: (s.natural_gas.get_atomic_flow('C')
                           + s.oleyl_alcohol.get_atomic_flow('C')) \
                           * BDO_chemicals.CO2.MW / conc_aqueous_broth.F_mass

# GWP from electricity
get_electricity_use = lambda: sum(i.power_utility.rate for i in BDO_sys.units)
get_electricity_GWP = lambda: get_electricity_use()*CFs['GWP_CFs']['Electricity'] \
    / conc_aqueous_broth.F_mass

# CO2 fixed in lactic acid product
get_fixed_GWP = lambda: \
    conc_aqueous_broth.get_atomic_flow('C')*BDO_chemicals.CO2.MW/conc_aqueous_broth.F_mass

get_GWP = lambda: get_material_GWP()+get_non_bio_GWP()+get_electricity_GWP() - get_Isobutanol_GWP()

# Fossil energy consumption (FEC) from materials
def get_material_FEC():
    LCA_stream.mass = sum(i.mass for i in LCA_streams)
    chemical_FEC = LCA_stream.mass*CFs['FEC_CF_stream'].mass
    # feedstock_FEC = feedstock.F_mass*CFs['FEC_CFs']['Corn stover']
    return chemical_FEC.sum()/conc_aqueous_broth.F_mass

# FEC from electricity
get_electricity_FEC = lambda: \
    get_electricity_use()*CFs['FEC_CFs']['Electricity']/conc_aqueous_broth.F_mass

# Total FEC
get_FEC = lambda: get_material_FEC() + get_electricity_FEC() - get_Isobutanol_FEC()

get_SPED = lambda: u.BT.system_heating_demand*0.001/conc_aqueous_broth.F_mass
conc_aqueous_broth_LHV = 31.45 # MJ/kg conc_aqueous_broth

BDO_lca = LCA(BDO_sys, BDO_chemicals, CFs, feedstock, 'Corn stover', conc_aqueous_broth, [CT, CWP], BT=BT, CT=CT)

#%% TEA breakdown
def TEA_breakdown():
    metric_breakdowns = {i.name: {} for i in unit_groups[0].metrics}
    for ug in unit_groups:
        for metric in ug.metrics:
            # storage_metric_val = None
            if not ug.name=='storage':
                if ug.name=='other facilities':
                    metric_breakdowns[metric.name]['storage and ' + ug.name] = metric() + unit_groups[5].metrics[ug.metrics.index(metric)]()
                else:
                    metric_breakdowns[metric.name][ug.name] = metric()
            # else:
            #     storage_metric_val = metric()

    # return metric_breakdowns
    for i in unit_groups[0].metrics:
        print(f"\n\n----- {i.name} ({i.units}) -----")
        metric_breakdowns_i = metric_breakdowns[i.name]
        for j in metric_breakdowns_i.keys():
            print(f"{j}: {format(metric_breakdowns_i[j], '.3f')}")

# %% LCA breakdown

def FEC_breakdown():
    return {
        'feedstock': BDO_lca.feedstock_FEC,
        'heating demand': BDO_lca.heating_demand_FEC,
        'cooling demand': BDO_lca.cooling_demand_FEC,
        'electricity demand (non-cooling)': BDO_lca.electricity_demand_non_cooling_FEC,
        'other materials': BDO_lca.material_FEC,
    }

def GWP_breakdown():
    return {
        'feedstock': BDO_lca.FGHTP_GWP,
        'heating demand': BDO_lca.heating_demand_GWP,
        'cooling demand': BDO_lca.cooling_demand_GWP,
        'electricity demand (non-cooling)': BDO_lca.electricity_demand_non_cooling_GWP,
        'other direct non-biogenic emissions': BDO_lca.non_BT_direct_emissions_GWP,
        'other materials': BDO_lca.material_GWP,
    }

# %% Full analysis
get_conc_aqueous_broth_MPSP()

R302 = flowsheet('R302')
set_target_BDO_x(250e-6)
# Overall yield assuming yield on xylose is 80% of yield on glucose:
# (0.64*0.36+0.36*0.8*0.36) * 2 = 0.33408 * 2
# saccharified stream mass ratio of glucose:xylose is 0.64:0.36

# to achieve $1.98/kg: spec_1=0.353*2
spec.load_specifications(0.33408*2, 109.9, 1.0)

def print_recycles():
    sys = BDO_sys.copy()
    sys.flatten()
    for i in sys.recycle:
        print(i, i.imass['OleylAlcohol'] / 1e3)

def print_capital_cost():
    print('TCI', BDO_tea.TCI / 1e6, 'MM$')
    print('Fermentation installed cost', u.R302.installed_cost / 1e6, 'MM$')
    print(u.R302.results())

def print_specs():
    print(spec.spec_1, spec.spec_2, spec.spec_3)

def simulate_and_print():
    get_conc_aqueous_broth_MPSP()
    print('\n---------- Simulation Results ----------')
    # print_recycles()
    # print_specs()
    # print_capital_cost()
    print(f'MPSP is ${get_conc_aqueous_broth_MPSP():.3f}/kg')
    print(f'GWP is {BDO_lca.GWP:.3f} kg CO2-eq/kg conc_aqueous_broth')
    # print(f'Non-bio GWP is {():.3f} kg CO2-eq/kg conc_aqueous_broth')
    print(f'FEC is {BDO_lca.FEC:.2f} MJ/kg conc_aqueous_broth or {get_FEC()/conc_aqueous_broth_LHV:.2f} MJ/MJ conc_aqueous_broth')
    # print(f'SPED is {get_SPED():.2f} MJ/kg conc_aqueous_broth or {get_SPED()/conc_aqueous_broth_LHV:.2f} MJ/MJ conc_aqueous_broth')
    print('--------------------\n')

simulate_and_print()
# spec.load_specifications(0.36, 13.7, 1.0)
# simulate_and_print()
# spec.load_specifications(0.36, 109.9, 1.0)
# simulate_and_print()

# %% Maximum BDO titer (given by maximum sugar concentration)

# %%% conc_aqueous_broth environmental impacts

# GREET 2020:
#     GHG100 = 1.04 kgCO2-eq/kg
#     FEC = 43 MJ/kg
# Ecoinvent 3.7.1:
#     IPCC 2013 GWP100 = 1.9166 kgCO2-eq/kg
#     CED-f = 58.707 MJ/kg

# %%

# =============================================================================
# For Monte Carlo and analyses
# =============================================================================

# BDO_sub_sys = {
#     'feedstock_sys': (U101,),
#     'pretreatment_sys': (T201, M201, M202, M203,
#                          R201, R201_H, T202, T203,
#                          F201, F201_H,
#                          M204, T204, T204_P,
#                          M205, M205_P),
#     'conversion_sys': (H301, M301, M302, R301, R302, T301),
    # 'separation_sys': (u.S401, u.M401, u.M401_P,
    #                     u.S402,
    #                     # F401, F401_H, F401_P,
    #                     u.D401, u.D401_H, u.D401_P, u.S403,
    #                     u.M402_P, u.S403,
    #                     u.D403, u.D403_H, u.D403_P,
    #                     u.M501,
    #                     u.T606, u.T606_P)
                        # F402, F402_H, F402_P,
                        # D405, D405_H1, D405_H2, D405_P,
                        # M401, M401_P)
#     'wastewater_sys': (M501, WWT_cost, R501,
#                        M502, R502, S501, S502, M503,
#                        M504, S503, S504, M505),
#     'HXN': (HXN,),
#     'BT': (BT,),
#     'CT': (CT,),
#     'other_facilities': (T601, S601,
#                          T602, T603,
#                          T604, T604_P,
#                          T605, T605_P,
#                          T606, T606_P,
#                          PWC, CIP, ADP, FWT)
    # }


# -*- coding: utf-8 -*-
# vim: tabstop=4 shiftwidth=4 softtabstop=4

#
# LICENSE
#
# Copyright (C) 2010-2018 GEM Foundation, G. Weatherill, M. Pagani,
# D. Monelli.
#
# The Hazard Modeller's Toolkit is free software: you can redistribute
# it and/or modify it under the terms of the GNU Affero General Public
# License as published by the Free Software Foundation, either version
# 3 of the License, or (at your option) any later version.
#
# You should have received a copy of the GNU Affero General Public License
# along with OpenQuake. If not, see <http://www.gnu.org/licenses/>
#
# DISCLAIMER
#
# The software Hazard Modeller's Toolkit (openquake.hmtk) provided herein
# is released as a prototype implementation on behalf of
# scientists and engineers working within the GEM Foundation (Global
# Earthquake Model).
#
# It is distributed for the purpose of open collaboration and in the
# hope that it will be useful to the scientific, engineering, disaster
# risk and software design communities.
#
# The software is NOT distributed as part of GEM’s OpenQuake suite
# (https://www.globalquakemodel.org/tools-products) and must be considered as a
# separate entity. The software provided herein is designed and implemented
# by scientific staff. It is not developed to the design standards, nor
# subject to same level of critical review by professional software
# developers, as GEM’s OpenQuake software suite.
#
# Feedback and contribution to the software is welcome, and can be
# directed to the hazard scientific staff of the GEM Model Facility
# (hazard@globalquakemodel.org).
#
# The Hazard Modeller's Toolkit (openquake.hmtk) is therefore distributed WITHOUT
# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
# FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License
# for more details.
#
# The GEM Foundation, and the authors of the software, assume no
# liability for use of the software.

import numpy as np

from openquake.hmtk.seismicity.declusterer.base import (
    BaseCatalogueDecluster, DECLUSTERER_METHODS)
from openquake.hmtk.seismicity.utils import decimal_year, haversine
from openquake.hmtk.seismicity.declusterer.distance_time_windows import (
    TIME_DISTANCE_WINDOW_FUNCTIONS)

@DECLUSTERER_METHODS.add(
    "decluster",
    time_distance_window=TIME_DISTANCE_WINDOW_FUNCTIONS,
    time_window=np.float)
class Afteran(BaseCatalogueDecluster):
    """
    This implements the Afteran algorithm as described in this paper:
    Musson, R. (1999), Probabilistic seismic hazard maps for the North
    Balkan Region, Annali Di Geofisica, 42(6), 1109 - 1124
    """

    def decluster(self, catalogue, config):
        """
        catalogue_matrix, window_opt=TDW_GARDNERKNOPOFF, time_window=60.):

        :param catalogue: a catalogue object
        :type catalogue: Instance of the openquake.hmtk.seismicity.catalogue.Catalogue()
                         class
        :keyword window_opt: method used in calculating distance and time
            windows
        :type window_opt: string
        :keyword time_window: Length (in days) of moving time window
        :type time_window: positive float
        :returns: **vcl vector** indicating cluster number,
                  **flagvector** indicating which earthquakes belong to a
                  cluster
        :rtype: numpy.ndarray
        """
        # Convert time window from days to decimal years
        time_window = config['time_window'] / 365.
        # Pre-processing steps are the same as for Gardner & Knopoff
        # Get relevent parameters
        mag = catalogue.data['magnitude']
        neq = np.shape(mag)[0]  # Number of earthquakes
        # Get decimal year (needed for time windows)
        year_dec = decimal_year(catalogue.data['year'],
                                catalogue.data['month'],
                                catalogue.data['day'])
        # Get space windows corresponding to each event
        sw_space, _ = (
            config['time_distance_window'].calc(catalogue.data['magnitude']))

        # Pre-allocate cluster index vectors
        vcl = np.zeros(neq, dtype=int)
        flagvector = np.zeros(neq, dtype=int)
        # Rank magnitudes into descending order
        id0 = np.flipud(np.argsort(mag, kind='heapsort'))

        clust_index = 0
        for imarker in id0:
            # Earthquake not allocated to cluster - perform calculation
            if vcl[imarker] == 0:
                # Perform distance calculation
                mdist = haversine(
                    catalogue.data['longitude'],
                    catalogue.data['latitude'],
                    catalogue.data['longitude'][imarker],
                    catalogue.data['latitude'][imarker]).flatten()

                # Select earthquakes inside distance window, later than
                # mainshock and not already assigned to a cluster
                vsel1 = np.where(
                    np.logical_and(vcl == 0,
                                   np.logical_and(
                                       mdist <= sw_space[imarker],
                                       year_dec > year_dec[imarker])))[0]
                has_aftershocks = False
                if len(vsel1) > 0:
                    # Earthquakes after event inside distance window
                    temp_vsel1, has_aftershocks = self._find_aftershocks(
                        vsel1,
                        year_dec,
                        time_window,
                        imarker,
                        neq)
                    if has_aftershocks:
                        flagvector[temp_vsel1] = 1
                        vcl[temp_vsel1] = clust_index + 1

                # Select earthquakes inside distance window, earlier than
                # mainshock and not already assigned to a cluster
                has_foreshocks = False
                vsel2 = np.where(
                    np.logical_and(
                        vcl == 0,
                        np.logical_and(mdist <= sw_space[imarker],
                                       year_dec < year_dec[imarker])))[0]
                if len(vsel2) > 0:
                    # Earthquakes before event inside distance window
                    temp_vsel2, has_foreshocks = self._find_foreshocks(
                        vsel2,
                        year_dec,
                        time_window,
                        imarker,
                        neq)
                    if has_foreshocks:
                        flagvector[temp_vsel2] = -1
                        vcl[temp_vsel2] = clust_index + 1

                if has_aftershocks or has_foreshocks:
                    # Assign mainshock to cluster
                    vcl[imarker] = clust_index + 1
                    clust_index += 1

        return vcl, flagvector

    def _find_aftershocks(self, vsel, year_dec, time_window, imarker, neq):
        '''
        Function to identify aftershocks from a set of potential
        events inside the distance window of an earthquake.
        :param vsel: Pointer vector to the location of the events in distance
                     window
        :type vsel: numpy.ndarray
        :param year_dec: Vector of decimal catalogue event times
        :type year_dec: numpy.ndarray
        :param time_window: Moving time window for selection of time clusters
        :type time_window: float
        :param imarker: Index of the mainshock in the catalogue vector
        :type imarker: Integer
        :param neq: Number of events in distance window of mainshock
        :type neq: Integer
        '''
        temp_vsel1 = np.zeros(neq, dtype=bool)
        has_aftershocks = False

        # Finds the time difference between events
        delta_time = np.diff(
            np.hstack([year_dec[imarker], year_dec[vsel]]))
        for iloc in range(0, len(vsel)):
            # If time difference between event is smaller than
            # time window - is an aftershock -> continue

            if delta_time[iloc] < time_window:
                temp_vsel1[vsel[iloc]] = True
                has_aftershocks = True
            else:
                # Time difference between events is larger than
                # window -> no more aftershocks -> return
                return temp_vsel1, has_aftershocks

        return temp_vsel1, has_aftershocks

    def _find_foreshocks(self, vsel, year_dec, time_window, imarker, neq):
        '''
        Finds foreshocks from a set of potential events within
        the distance window of a mainshock.
        :param vsel: Pointer vector to the location of the events in distance
                     window
        :type vsel: numpy.ndarray
        :param year_dec: Vector of decimal catalogue event times
        :type year_dec: numpy.ndarray
        :param time_window: Moving time window for selection of time clusters
        :type time_window: float
        :param imarker: Index of the mainshock in the catalogue vector
        :type imarker: Integer
        :param neq: Number of events in distance window of mainshock
        :type neq: Integer
        '''

        temp_vsel2 = np.zeros(neq, dtype=bool)
        has_foreshocks = False

        # The initial time is the time of the mainshock
        initial_time = year_dec[imarker]
        year_dec = year_dec[vsel]
        for jloc in range(len(vsel) - 1, -1, -1):
            # If the time between the mainshock and the preceeding
            # event is smaller than the time_window then event
            # is a foreshock

            if (initial_time - year_dec[jloc]) < time_window:
                temp_vsel2[vsel[jloc]] = True
                has_foreshocks = True
                # Update target time to consider current foreshock
                # Then continue
                initial_time = year_dec[jloc]
            else:
                # No events inside time window
                # end of foreshock sequence - return
                return temp_vsel2, has_foreshocks

        return temp_vsel2, has_foreshocks


# Helper functions for running autoencoder

import os
import argparse
import pandas as pd
import matplotlib.pyplot as plt
import tensorflow as tf

# To ensure reproducibility using Keras during development
# https://keras.io/getting-started/faq/#how-can-i-obtain-reproducible-results-using-keras-during-development
import numpy as np
import random as rn

# The below is necessary in Python 3.2.3 onwards to
# have reproducible behavior for certain hash-based operations.
# See these references for further details:
# https://docs.python.org/3.4/using/cmdline.html#envvar-PYTHONHASHSEED
# https://github.com/keras-team/keras/issues/2280#issuecomment-306959926
randomState = 123
import os
os.environ['PYTHONHASHSEED'] = '0'

# The below is necessary for starting Numpy generated random numbers
# in a well-defined initial state.

np.random.seed(42)

# The below is necessary for starting core Python generated random numbers
# in a well-defined state.

rn.seed(12345)

# Force TensorFlow to use single thread.
# Multiple threads are a potential source of
# non-reproducible results.
# For further details, see: https://stackoverflow.com/questions/42022950/which-seeds-have-to-be-set-where-to-realize-100-reproducibility-of-training-res

session_conf = tf.ConfigProto(
    intra_op_parallelism_threads=1, inter_op_parallelism_threads=1)

from keras import backend as K

# The below tf.set_random_seed() will make random number generation
# in the TensorFlow backend have a well-defined initial state.
# For further details, see: https://www.tensorflow.org/api_docs/python/tf/set_random_seed

tf.set_random_seed(1234)

sess = tf.Session(graph=tf.get_default_graph(), config=session_conf)
K.set_session(sess)

from keras.layers import Input, Dense, Lambda, Layer, Activation
from keras.layers.normalization import BatchNormalization
from keras.models import Model, Sequential
from keras import metrics, optimizers
from keras.callbacks import Callback

# Functions
#
# Based on publication by Way et. al. (https://www.ncbi.nlm.nih.gov/pmc/articles/PMC5728678/)
# Github repo (https://github.com/greenelab/tybalt/blob/master/scripts/vae_pancancer.py)

# Function for reparameterization trick to make model differentiable

def sampling_maker(epsilon_std):
    def sampling(args):
        # Function with args required for Keras Lambda function
        z_mean, z_log_var = args

        # Draw epsilon of the same shape from a standard normal distribution
        epsilon = K.random_normal(shape=tf.shape(z_mean), mean=0.,
                                  stddev=epsilon_std)

        # The latent vector is non-deterministic and differentiable
        # in respect to z_mean and z_log_var
        z = z_mean + K.exp(z_log_var / 2) * epsilon
        return z
    return sampling

class CustomVariationalLayer(Layer):
    """
    Define a custom layer that learns and performs the training
    """

    def __init__(self, original_dim, z_log_var_encoded,
                 z_mean_encoded, beta, **kwargs):
        # https://keras.io/layers/writing-your-own-keras-layers/
        self.is_placeholder = True
        self.original_dim = original_dim
        self.z_log_var_encoded = z_log_var_encoded
        self.z_mean_encoded = z_mean_encoded
        self.beta = beta

        super(CustomVariationalLayer, self).__init__(**kwargs)

    def vae_loss(self, x_input, x_decoded):
        reconstruction_loss = self.original_dim * \
            metrics.binary_crossentropy(x_input, x_decoded)

        kl_loss = - 0.5 * K.sum(1 + self.z_log_var_encoded
                                - K.square(self.z_mean_encoded)
                                - K.exp(self.z_log_var_encoded), axis=-1)

        return K.mean(reconstruction_loss + (K.get_value(self.beta) * kl_loss))

    def call(self, inputs):
        x = inputs[0]
        x_decoded = inputs[1]
        loss = self.vae_loss(x, x_decoded)
        self.add_loss(loss, inputs=inputs)
        # We won't actually use the output.
        return x

class WarmUpCallback(Callback):
    def __init__(self, beta, kappa):
        self.beta = beta
        self.kappa = kappa

    # Behavior on each epoch
    def on_epoch_end(self, epoch, logs={}):
        if K.get_value(self.beta) <= 1:
            K.set_value(self.beta, K.get_value(self.beta) + self.kappa)


# Copyright (c) Pymatgen Development Team.
# Distributed under the terms of the MIT License.

"""
This module implements a FloatWithUnit, which is a subclass of float. It
also defines supported units for some commonly used units for energy, length,
temperature, time and charge. FloatWithUnit also support conversion to one
another, and additions and subtractions perform automatic conversion if
units are detected. An ArrayWithUnit is also implemented, which is a subclass
of numpy's ndarray with similar unit features.
"""

import collections
import numbers
from functools import partial

import numpy as np
import scipy.constants as const

__author__ = "Shyue Ping Ong, Matteo Giantomassi"
__copyright__ = "Copyright 2011, The Materials Project"
__version__ = "1.0"
__maintainer__ = "Shyue Ping Ong, Matteo Giantomassi"
__status__ = "Production"
__date__ = "Aug 30, 2013"

"""
Some conversion factors
"""
Ha_to_eV = 1 / const.physical_constants["electron volt-hartree relationship"][0]
eV_to_Ha = 1 / Ha_to_eV
Ry_to_eV = Ha_to_eV / 2
amu_to_kg = const.physical_constants["atomic mass unit-kilogram relationship"][0]
mile_to_meters = const.mile
bohr_to_angstrom = const.physical_constants["Bohr radius"][0] * 1e10
bohr_to_ang = bohr_to_angstrom
ang_to_bohr = 1 / bohr_to_ang
kCal_to_kJ = const.calorie
kb = const.physical_constants["Boltzmann constant in eV/K"][0]

"""
Definitions of supported units. Values below are essentially scaling and
conversion factors. What matters is the relative values, not the absolute.
The SI units must have factor 1.
"""
BASE_UNITS = {
    "length": {
        "m": 1,
        "km": 1000,
        "mile": mile_to_meters,
        "ang": 1e-10,
        "cm": 1e-2,
        "pm": 1e-12,
        "bohr": bohr_to_angstrom * 1e-10,
    },
    "mass": {
        "kg": 1,
        "g": 1e-3,
        "amu": amu_to_kg,
    },
    "time": {
        "s": 1,
        "min": 60,
        "h": 3600,
        "d": 3600 * 24,
    },
    "current": {"A": 1},
    "temperature": {
        "K": 1,
    },
    "amount": {"mol": 1, "atom": 1 / const.N_A},
    "intensity": {"cd": 1},
    "memory": {
        "byte": 1,
        "Kb": 1024,
        "Mb": 1024**2,
        "Gb": 1024**3,
        "Tb": 1024**4,
    },
}

# Accept kb, mb, gb ... as well.
BASE_UNITS["memory"].update({k.lower(): v for k, v in BASE_UNITS["memory"].items()})

# This current list are supported derived units defined in terms of powers of
# SI base units and constants.
DERIVED_UNITS = {
    "energy": {
        "eV": {"kg": 1, "m": 2, "s": -2, const.e: 1},
        "meV": {"kg": 1, "m": 2, "s": -2, const.e * 1e-3: 1},
        "Ha": {"kg": 1, "m": 2, "s": -2, const.e * Ha_to_eV: 1},
        "Ry": {"kg": 1, "m": 2, "s": -2, const.e * Ry_to_eV: 1},
        "J": {"kg": 1, "m": 2, "s": -2},
        "kJ": {"kg": 1, "m": 2, "s": -2, 1000: 1},
        "kCal": {"kg": 1, "m": 2, "s": -2, 1000: 1, kCal_to_kJ: 1},
    },
    "charge": {
        "C": {"A": 1, "s": 1},
        "e": {"A": 1, "s": 1, const.e: 1},
    },
    "force": {
        "N": {"kg": 1, "m": 1, "s": -2},
        "KN": {"kg": 1, "m": 1, "s": -2, 1000: 1},
        "MN": {"kg": 1, "m": 1, "s": -2, 1e6: 1},
        "GN": {"kg": 1, "m": 1, "s": -2, 1e9: 1},
    },
    "frequency": {
        "Hz": {"s": -1},
        "KHz": {"s": -1, 1000: 1},
        "MHz": {"s": -1, 1e6: 1},
        "GHz": {"s": -1, 1e9: 1},
        "THz": {"s": -1, 1e12: 1},
    },
    "pressure": {
        "Pa": {"kg": 1, "m": -1, "s": -2},
        "KPa": {"kg": 1, "m": -1, "s": -2, 1000: 1},
        "MPa": {"kg": 1, "m": -1, "s": -2, 1e6: 1},
        "GPa": {"kg": 1, "m": -1, "s": -2, 1e9: 1},
    },
    "power": {
        "W": {"m": 2, "kg": 1, "s": -3},
        "KW": {"m": 2, "kg": 1, "s": -3, 1000: 1},
        "MW": {"m": 2, "kg": 1, "s": -3, 1e6: 1},
        "GW": {"m": 2, "kg": 1, "s": -3, 1e9: 1},
    },
    "emf": {"V": {"m": 2, "kg": 1, "s": -3, "A": -1}},
    "capacitance": {"F": {"m": -2, "kg": -1, "s": 4, "A": 2}},
    "resistance": {"ohm": {"m": 2, "kg": 1, "s": -3, "A": -2}},
    "conductance": {"S": {"m": -2, "kg": -1, "s": 3, "A": 2}},
    "magnetic_flux": {"Wb": {"m": 2, "kg": 1, "s": -2, "A": -1}},
    "cross_section": {"barn": {"m": 2, 1e-28: 1}, "mbarn": {"m": 2, 1e-31: 1}},
}

ALL_UNITS = dict(list(BASE_UNITS.items()) + list(DERIVED_UNITS.items()))  # type: ignore
SUPPORTED_UNIT_NAMES = tuple(i for d in ALL_UNITS.values() for i in d.keys())

# Mapping unit name --> unit type (unit names must be unique).
_UNAME2UTYPE = {}  # type: ignore
for utype, d in ALL_UNITS.items():
    assert not set(d.keys()).intersection(_UNAME2UTYPE.keys())
    _UNAME2UTYPE.update({uname: utype for uname in d})
del utype, d

def _get_si_unit(unit):
    unit_type = _UNAME2UTYPE[unit]
    si_unit = filter(lambda k: BASE_UNITS[unit_type][k] == 1, BASE_UNITS[unit_type].keys())
    return list(si_unit)[0], BASE_UNITS[unit_type][unit]

class UnitError(BaseException):
    """
    Exception class for unit errors.
    """

def _check_mappings(u):
    for v in DERIVED_UNITS.values():
        for k2, v2 in v.items():
            if all(v2.get(ku, 0) == vu for ku, vu in u.items()) and all(
                u.get(kv2, 0) == vv2 for kv2, vv2 in v2.items()
            ):
                return {k2: 1}
    return u

class Unit(collections.abc.Mapping):
    """
    Represents a unit, e.g., "m" for meters, etc. Supports compound units.
    Only integer powers are supported for units.
    """

    Error = UnitError

    def __init__(self, unit_def):
        """
        Constructs a unit.

        Args:
            unit_def: A definition for the unit. Either a mapping of unit to
                powers, e.g., {"m": 2, "s": -1} represents "m^2 s^-1",
                or simply as a string "kg m^2 s^-1". Note that the supported
                format uses "^" as the power operator and all units must be
                space-separated.
        """

        if isinstance(unit_def, str):
            unit = collections.defaultdict(int)
            import re

            for m in re.finditer(r"([A-Za-z]+)\s*\^*\s*([\-0-9]*)", unit_def):
                p = m.group(2)
                p = 1 if not p else int(p)
                k = m.group(1)
                unit[k] += p
        else:
            unit = {k: v for k, v in dict(unit_def).items() if v != 0}
        self._unit = _check_mappings(unit)

    def __mul__(self, other):
        new_units = collections.defaultdict(int)
        for k, v in self.items():
            new_units[k] += v
        for k, v in other.items():
            new_units[k] += v
        return Unit(new_units)

    def __rmul__(self, other):
        return self.__mul__(other)

    def __div__(self, other):
        new_units = collections.defaultdict(int)
        for k, v in self.items():
            new_units[k] += v
        for k, v in other.items():
            new_units[k] -= v
        return Unit(new_units)

    def __truediv__(self, other):
        return self.__div__(other)

    def __pow__(self, i):
        return Unit({k: v * i for k, v in self.items()})

    def __iter__(self):
        return self._unit.__iter__()

    def __getitem__(self, i):
        return self._unit[i]

    def __len__(self):
        return len(self._unit)

    def __repr__(self):
        sorted_keys = sorted(self._unit.keys(), key=lambda k: (-self._unit[k], k))
        return " ".join(
            [f"{k}^{self._unit[k]}" if self._unit[k] != 1 else k for k in sorted_keys if self._unit[k] != 0]
        )

    def __str__(self):
        return self.__repr__()

    @property
    def as_base_units(self):
        """
        Converts all units to base SI units, including derived units.

        Returns:
            (base_units_dict, scaling factor). base_units_dict will not
            contain any constants, which are gathered in the scaling factor.
        """
        b = collections.defaultdict(int)
        factor = 1
        for k, v in self.items():
            derived = False
            for d in DERIVED_UNITS.values():
                if k in d:
                    for k2, v2 in d[k].items():
                        if isinstance(k2, numbers.Number):
                            factor *= k2 ** (v2 * v)
                        else:
                            b[k2] += v2 * v
                    derived = True
                    break
            if not derived:
                si, f = _get_si_unit(k)
                b[si] += v
                factor *= f**v
        return {k: v for k, v in b.items() if v != 0}, factor

    def get_conversion_factor(self, new_unit):
        """
        Returns a conversion factor between this unit and a new unit.
        Compound units are supported, but must have the same powers in each
        unit type.

        Args:
            new_unit: The new unit.
        """
        uo_base, ofactor = self.as_base_units
        un_base, nfactor = Unit(new_unit).as_base_units
        units_new = sorted(un_base.items(), key=lambda d: _UNAME2UTYPE[d[0]])
        units_old = sorted(uo_base.items(), key=lambda d: _UNAME2UTYPE[d[0]])
        factor = ofactor / nfactor
        for uo, un in zip(units_old, units_new):
            if uo[1] != un[1]:
                raise UnitError(f"Units {uo} and {un} are not compatible!")
            c = ALL_UNITS[_UNAME2UTYPE[uo[0]]]
            factor *= (c[uo[0]] / c[un[0]]) ** uo[1]
        return factor

class FloatWithUnit(float):
    """
    Subclasses float to attach a unit type. Typically, you should use the
    pre-defined unit type subclasses such as Energy, Length, etc. instead of
    using FloatWithUnit directly.

    Supports conversion, addition and subtraction of the same unit type. E.g.,
    1 m + 20 cm will be automatically converted to 1.2 m (units follow the
    leftmost quantity). Note that FloatWithUnit does not override the eq
    method for float, i.e., units are not checked when testing for equality.
    The reason is to allow this class to be used transparently wherever floats
    are expected.

    >>> e = Energy(1.1, "Ha")
    >>> a = Energy(1.1, "Ha")
    >>> b = Energy(3, "eV")
    >>> c = a + b
    >>> print(c)
    1.2102479761938871 Ha
    >>> c.to("eV")
    32.932522246000005 eV
    """

    Error = UnitError

    @classmethod
    def from_string(cls, s):
        """
        Initialize a FloatWithUnit from a string. Example Memory.from_string("1. Mb")
        """
        # Extract num and unit string.
        s = s.strip()
        for i, char in enumerate(s):
            if char.isalpha() or char.isspace():
                break
        else:
            raise Exception(f"Unit is missing in string {s}")
        num, unit = float(s[:i]), s[i:]

        # Find unit type (set it to None if it cannot be detected)
        for unit_type, d in BASE_UNITS.items():
            if unit in d:
                break
        else:
            unit_type = None

        return cls(num, unit, unit_type=unit_type)

    def __new__(cls, val, unit, unit_type=None):
        """Overrides __new__ since we are subclassing a Python primitive/"""
        new = float.__new__(cls, val)
        new._unit = Unit(unit)
        new._unit_type = unit_type
        return new

    def __init__(self, val, unit, unit_type=None):
        """
        Initializes a float with unit.

        Args:
            val (float): Value
            unit (Unit): A unit. E.g., "C".
            unit_type (str): A type of unit. E.g., "charge"
        """
        if unit_type is not None and str(unit) not in ALL_UNITS[unit_type]:
            raise UnitError(f"{unit} is not a supported unit for {unit_type}")
        self._unit = Unit(unit)
        self._unit_type = unit_type

    def __repr__(self):
        return super().__repr__()

    def __str__(self):
        s = super().__str__()
        return f"{s} {self._unit}"

    def __add__(self, other):
        if not hasattr(other, "unit_type"):
            return super().__add__(other)
        if other.unit_type != self._unit_type:
            raise UnitError("Adding different types of units is not allowed")
        val = other
        if other.unit != self._unit:
            val = other.to(self._unit)
        return FloatWithUnit(float(self) + val, unit_type=self._unit_type, unit=self._unit)

    def __sub__(self, other):
        if not hasattr(other, "unit_type"):
            return super().__sub__(other)
        if other.unit_type != self._unit_type:
            raise UnitError("Subtracting different units is not allowed")
        val = other
        if other.unit != self._unit:
            val = other.to(self._unit)
        return FloatWithUnit(float(self) - val, unit_type=self._unit_type, unit=self._unit)

    def __mul__(self, other):
        if not isinstance(other, FloatWithUnit):
            return FloatWithUnit(float(self) * other, unit_type=self._unit_type, unit=self._unit)
        return FloatWithUnit(float(self) * other, unit_type=None, unit=self._unit * other._unit)

    def __rmul__(self, other):
        if not isinstance(other, FloatWithUnit):
            return FloatWithUnit(float(self) * other, unit_type=self._unit_type, unit=self._unit)
        return FloatWithUnit(float(self) * other, unit_type=None, unit=self._unit * other._unit)

    def __pow__(self, i):
        return FloatWithUnit(float(self) ** i, unit_type=None, unit=self._unit**i)

    def __truediv__(self, other):
        val = super().__truediv__(other)
        if not isinstance(other, FloatWithUnit):
            return FloatWithUnit(val, unit_type=self._unit_type, unit=self._unit)
        return FloatWithUnit(val, unit_type=None, unit=self._unit / other._unit)

    def __neg__(self):
        return FloatWithUnit(super().__neg__(), unit_type=self._unit_type, unit=self._unit)

    def __getnewargs__(self):
        """Function used by pickle to recreate object."""
        # print(self.__dict__)
        # FIXME
        # There's a problem with _unit_type if we try to unpickle objects from file.
        # since self._unit_type might not be defined. I think this is due to
        # the use of decorators (property and unitized). In particular I have problems with "amu"
        # likely due to weight in core.composition
        if hasattr(self, "_unit_type"):
            args = float(self), self._unit, self._unit_type
        else:
            args = float(self), self._unit, None

        return args

    def __getstate__(self):
        state = self.__dict__.copy()
        state["val"] = float(self)
        return state

    def __setstate__(self, state):
        self._unit = state["_unit"]

    @property
    def unit_type(self) -> str:
        """
        :return: The type of unit. Energy, Charge, etc.
        """
        return self._unit_type

    @property
    def unit(self) -> str:
        """
        :return: The unit, e.g., "eV".
        """
        return self._unit

    def to(self, new_unit):
        """
        Conversion to a new_unit. Right now, only supports 1 to 1 mapping of
        units of each type.

        Args:
            new_unit: New unit type.

        Returns:
            A FloatWithUnit object in the new units.

        Example usage:
        >>> e = Energy(1.1, "eV")
        >>> e = Energy(1.1, "Ha")
        >>> e.to("eV")
        29.932522246 eV
        """
        return FloatWithUnit(
            self * self.unit.get_conversion_factor(new_unit),
            unit_type=self._unit_type,
            unit=new_unit,
        )

    @property
    def as_base_units(self):
        """
        Returns this FloatWithUnit in base SI units, including derived units.

        Returns:
            A FloatWithUnit object in base SI units
        """
        return self.to(self.unit.as_base_units[0])

    @property
    def supported_units(self):
        """
        Supported units for specific unit type.
        """
        return tuple(ALL_UNITS[self._unit_type].keys())

class ArrayWithUnit(np.ndarray):
    """
    Subclasses `numpy.ndarray` to attach a unit type. Typically, you should
    use the pre-defined unit type subclasses such as EnergyArray,
    LengthArray, etc. instead of using ArrayWithFloatWithUnit directly.

    Supports conversion, addition and subtraction of the same unit type. E.g.,
    1 m + 20 cm will be automatically converted to 1.2 m (units follow the
    leftmost quantity).

    >>> a = EnergyArray([1, 2], "Ha")
    >>> b = EnergyArray([1, 2], "eV")
    >>> c = a + b
    >>> print(c)
    [ 1.03674933  2.07349865] Ha
    >>> c.to("eV")
    array([ 28.21138386,  56.42276772]) eV
    """

    Error = UnitError

    def __new__(cls, input_array, unit, unit_type=None):
        """
        Override __new__.
        """
        # Input array is an already formed ndarray instance
        # We first cast to be our class type
        obj = np.asarray(input_array).view(cls)
        # add the new attributes to the created instance
        obj._unit = Unit(unit)
        obj._unit_type = unit_type
        return obj

    def __array_finalize__(self, obj):
        """
        See http://docs.scipy.org/doc/numpy/user/basics.subclassing.html for
        comments.
        """
        if obj is None:
            return
        self._unit = getattr(obj, "_unit", None)
        self._unit_type = getattr(obj, "_unit_type", None)

    @property
    def unit_type(self) -> str:
        """
        :return: The type of unit. Energy, Charge, etc.
        """
        return self._unit_type

    @property
    def unit(self) -> str:
        """
        :return: The unit, e.g., "eV".
        """
        return self._unit

    def __reduce__(self):
        # print("in reduce")
        reduce = list(super().__reduce__())
        # print("unit",self._unit)
        # print(reduce[2])
        reduce[2] = {"np_state": reduce[2], "_unit": self._unit}
        return tuple(reduce)

    def __setstate__(self, state):
        # pylint: disable=E1101
        super().__setstate__(state["np_state"])
        self._unit = state["_unit"]

    def __repr__(self):
        return f"{np.array(self).__repr__()} {self.unit}"

    def __str__(self):
        return f"{np.array(self).__str__()} {self.unit}"

    def __add__(self, other):
        if hasattr(other, "unit_type"):
            if other.unit_type != self.unit_type:
                raise UnitError("Adding different types of units is not allowed")

            if other.unit != self.unit:
                other = other.to(self.unit)

        return self.__class__(np.array(self) + np.array(other), unit_type=self.unit_type, unit=self.unit)

    def __sub__(self, other):
        if hasattr(other, "unit_type"):
            if other.unit_type != self.unit_type:
                raise UnitError("Subtracting different units is not allowed")

            if other.unit != self.unit:
                other = other.to(self.unit)

        return self.__class__(np.array(self) - np.array(other), unit_type=self.unit_type, unit=self.unit)

    def __mul__(self, other):
        # FIXME
        # Here we have the most important difference between FloatWithUnit and
        # ArrayWithFloatWithUnit:
        # If other does not have units, I return an object with the same units
        # as self.
        # if other *has* units, I return an object *without* units since
        # taking into account all the possible derived quantities would be
        # too difficult.
        # Moreover Energy(1.0) * Time(1.0, "s") returns 1.0 Ha that is a
        # bit misleading.
        # Same protocol for __div__
        if not hasattr(other, "unit_type"):
            return self.__class__(
                np.array(self).__mul__(np.array(other)),
                unit_type=self._unit_type,
                unit=self._unit,
            )
        # Cannot use super since it returns an instance of self.__class__
        # while here we want a bare numpy array.
        return self.__class__(np.array(self).__mul__(np.array(other)), unit=self.unit * other.unit)

    def __rmul__(self, other):
        # pylint: disable=E1101
        if not hasattr(other, "unit_type"):
            return self.__class__(
                np.array(self).__rmul__(np.array(other)),
                unit_type=self._unit_type,
                unit=self._unit,
            )
        return self.__class__(np.array(self).__rmul__(np.array(other)), unit=self.unit * other.unit)

    def __div__(self, other):
        # pylint: disable=E1101
        if not hasattr(other, "unit_type"):
            return self.__class__(
                np.array(self).__div__(np.array(other)),
                unit_type=self._unit_type,
                unit=self._unit,
            )
        return self.__class__(np.array(self).__div__(np.array(other)), unit=self.unit / other.unit)

    def __truediv__(self, other):
        # pylint: disable=E1101
        if not hasattr(other, "unit_type"):
            return self.__class__(
                np.array(self).__truediv__(np.array(other)),
                unit_type=self._unit_type,
                unit=self._unit,
            )
        return self.__class__(np.array(self).__truediv__(np.array(other)), unit=self.unit / other.unit)

    def __neg__(self):
        return self.__class__(np.array(self).__neg__(), unit_type=self.unit_type, unit=self.unit)

    def to(self, new_unit):
        """
        Conversion to a new_unit.

        Args:
            new_unit:
                New unit type.

        Returns:
            A ArrayWithFloatWithUnit object in the new units.

        Example usage:
        >>> e = EnergyArray([1, 1.1], "Ha")
        >>> e.to("eV")
        array([ 27.21138386,  29.93252225]) eV
        """
        return self.__class__(
            np.array(self) * self.unit.get_conversion_factor(new_unit),
            unit_type=self.unit_type,
            unit=new_unit,
        )

    @property
    def as_base_units(self):
        """
        Returns this ArrayWithUnit in base SI units, including derived units.

        Returns:
            An ArrayWithUnit object in base SI units
        """
        return self.to(self.unit.as_base_units[0])

    # TODO abstract base class property?
    @property
    def supported_units(self):
        """
        Supported units for specific unit type.
        """
        return ALL_UNITS[self.unit_type]

    # TODO abstract base class method?
    def conversions(self):
        """
        Returns a string showing the available conversions.
        Useful tool in interactive mode.
        """
        return "\n".join(str(self.to(unit)) for unit in self.supported_units)

def _my_partial(func, *args, **kwargs):
    """
    Partial returns a partial object and therefore we cannot inherit class
    methods defined in FloatWithUnit. This function calls partial and patches
    the new class before returning.
    """
    newobj = partial(func, *args, **kwargs)
    # monkey patch
    newobj.from_string = FloatWithUnit.from_string
    return newobj

Energy = partial(FloatWithUnit, unit_type="energy")
"""
A float with an energy unit.

Args:
    val (float): Value
    unit (Unit): E.g., eV, kJ, etc. Must be valid unit or UnitError is raised.
"""
EnergyArray = partial(ArrayWithUnit, unit_type="energy")

Length = partial(FloatWithUnit, unit_type="length")
"""
A float with a length unit.

Args:
    val (float): Value
    unit (Unit): E.g., m, ang, bohr, etc. Must be valid unit or UnitError is
        raised.
"""
LengthArray = partial(ArrayWithUnit, unit_type="length")

Mass = partial(FloatWithUnit, unit_type="mass")
"""
A float with a mass unit.

Args:
    val (float): Value
    unit (Unit): E.g., amu, kg, etc. Must be valid unit or UnitError is
        raised.
"""
MassArray = partial(ArrayWithUnit, unit_type="mass")

Temp = partial(FloatWithUnit, unit_type="temperature")
"""
A float with a temperature unit.

Args:
    val (float): Value
    unit (Unit): E.g., K. Only K (kelvin) is supported.
"""
TempArray = partial(ArrayWithUnit, unit_type="temperature")

Time = partial(FloatWithUnit, unit_type="time")
"""
A float with a time unit.

Args:
    val (float): Value
    unit (Unit): E.g., s, min, h. Must be valid unit or UnitError is
        raised.
"""
TimeArray = partial(ArrayWithUnit, unit_type="time")

Charge = partial(FloatWithUnit, unit_type="charge")
"""
A float with a charge unit.

Args:
    val (float): Value
    unit (Unit): E.g., C, e (electron charge). Must be valid unit or UnitError
        is raised.
"""
ChargeArray = partial(ArrayWithUnit, unit_type="charge")

Memory = _my_partial(FloatWithUnit, unit_type="memory")
"""
A float with a memory unit.

Args:
    val (float): Value
    unit (Unit): E.g., Kb, Mb, Gb, Tb. Must be valid unit or UnitError
        is raised.
"""

def obj_with_unit(obj, unit):
    """
    Returns a `FloatWithUnit` instance if obj is scalar, a dictionary of
    objects with units if obj is a dict, else an instance of
    `ArrayWithFloatWithUnit`.

    Args:
        unit: Specific units (eV, Ha, m, ang, etc.).
    """
    unit_type = _UNAME2UTYPE[unit]

    if isinstance(obj, numbers.Number):
        return FloatWithUnit(obj, unit=unit, unit_type=unit_type)
    if isinstance(obj, collections.abc.Mapping):
        return {k: obj_with_unit(v, unit) for k, v in obj.items()}
    return ArrayWithUnit(obj, unit=unit, unit_type=unit_type)

def unitized(unit):
    """
    Useful decorator to assign units to the output of a function. You can also
    use it to standardize the output units of a function that already returns
    a FloatWithUnit or ArrayWithUnit. For sequences, all values in the sequences
    are assigned the same unit. It works with Python sequences only. The creation
    of numpy arrays loses all unit information. For mapping types, the values
    are assigned units.

    Args:
        unit: Specific unit (eV, Ha, m, ang, etc.).

    Example usage::

        @unitized(unit="kg")
        def get_mass():
            return 123.45

    """

    def wrap(f):
        def wrapped_f(*args, **kwargs):
            val = f(*args, **kwargs)
            unit_type = _UNAME2UTYPE[unit]

            if isinstance(val, (FloatWithUnit, ArrayWithUnit)):
                return val.to(unit)

            if isinstance(val, collections.abc.Sequence):
                # TODO: why don't we return a ArrayWithUnit?
                # This complicated way is to ensure the sequence type is
                # preserved (list or tuple).
                return val.__class__([FloatWithUnit(i, unit_type=unit_type, unit=unit) for i in val])
            if isinstance(val, collections.abc.Mapping):
                for k, v in val.items():
                    val[k] = FloatWithUnit(v, unit_type=unit_type, unit=unit)
            elif isinstance(val, numbers.Number):
                return FloatWithUnit(val, unit_type=unit_type, unit=unit)
            elif val is None:
                pass
            else:
                raise TypeError(f"Don't know how to assign units to {str(val)}")
            return val

        return wrapped_f

    return wrap

if __name__ == "__main__":
    import doctest

    doctest.testmod()


import torch
import torch.utils.data
import math
import numpy as np
from collections import defaultdict

def pad(species):
    """Put different species together into single tensor.

    If the species are from molecules of different number of total atoms, then
    ghost atoms with atom type -1 will be added to make it fit into the same
    shape.

    Arguments:
        species (:class:`collections.abc.Sequence`): sequence of species.
            Species must be of shape ``(N, A)``, where ``N`` is the number of
            3D structures, ``A`` is the number of atoms.

    Returns:
        :class:`torch.Tensor`: species batched together.
    """
    max_atoms = max([s.shape[1] for s in species])
    padded_species = []
    for s in species:
        natoms = s.shape[1]
        if natoms < max_atoms:
            padding = torch.full((s.shape[0], max_atoms - natoms), -1,
                                 dtype=torch.long, device=s.device)
            s = torch.cat([s, padding], dim=1)
        padded_species.append(s)
    return torch.cat(padded_species)

def pad_atomic_properties(atomic_properties, padding_values=defaultdict(lambda: 0.0, species=-1)):
    """Put a sequence of atomic properties together into single tensor.

    Inputs are `[{'species': ..., ...}, {'species': ..., ...}, ...]` and the outputs
    are `{'species': padded_tensor, ...}`

    Arguments:
        species_coordinates (:class:`collections.abc.Sequence`): sequence of
             atomic properties.
        padding_values (dict): the value to fill to pad tensors to same size
    """
    keys = list(atomic_properties[0])
    anykey = keys[0]
    max_atoms = max(x[anykey].shape[1] for x in atomic_properties)
    padded = {k: [] for k in keys}
    for p in atomic_properties:
        num_molecules = 1
        for v in p.values():
            assert num_molecules in {1, v.shape[0]}, 'Number of molecules in different atomic properties mismatch'
            if v.shape[0] != 1:
                num_molecules = v.shape[0]
        for k, v in p.items():
            shape = list(v.shape)
            padatoms = max_atoms - shape[1]
            shape[1] = padatoms
            padding = v.new_full(shape, padding_values[k])
            v = torch.cat([v, padding], dim=1)
            shape = list(v.shape)
            shape[0] = num_molecules
            v = v.expand(*shape)
            padded[k].append(v)
    return {k: torch.cat(v) for k, v in padded.items()}

# @torch.jit.script
def present_species(species):
    """Given a vector of species of atoms, compute the unique species present.

    Arguments:
        species (:class:`torch.Tensor`): 1D vector of shape ``(atoms,)``

    Returns:
        :class:`torch.Tensor`: 1D vector storing present atom types sorted.
    """
    # present_species, _ = species.flatten()._unique(sorted=True)
    present_species = species.flatten().unique(sorted=True)
    if present_species[0].item() == -1:
        present_species = present_species[1:]
    return present_species

def strip_redundant_padding(atomic_properties):
    """Strip trailing padding atoms.

    Arguments:
        atomic_properties (dict): properties to strip

    Returns:
        dict: same set of properties with redundant padding atoms stripped.
    """
    species = atomic_properties['species']
    non_padding = (species >= 0).any(dim=0).nonzero().squeeze()
    for k in atomic_properties:
        atomic_properties[k] = atomic_properties[k].index_select(1, non_padding)
    return atomic_properties

def map2central(cell, coordinates, pbc):
    """Map atoms outside the unit cell into the cell using PBC.

    Arguments:
        cell (:class:`torch.Tensor`): tensor of shape (3, 3) of the three
            vectors defining unit cell:

            .. code-block:: python

                tensor([[x1, y1, z1],
                        [x2, y2, z2],
                        [x3, y3, z3]])

        coordinates (:class:`torch.Tensor`): Tensor of shape
            ``(molecules, atoms, 3)``.

        pbc (:class:`torch.Tensor`): boolean vector of size 3 storing
            if pbc is enabled for that direction.

    Returns:
        :class:`torch.Tensor`: coordinates of atoms mapped back to unit cell.
    """
    # Step 1: convert coordinates from standard cartesian coordinate to unit
    # cell coordinates
    inv_cell = torch.inverse(cell)
    coordinates_cell = torch.matmul(coordinates, inv_cell)
    # Step 2: wrap cell coordinates into [0, 1)
    coordinates_cell -= coordinates_cell.floor() * pbc.to(coordinates_cell.dtype)
    # Step 3: convert from cell coordinates back to standard cartesian
    # coordinate
    return torch.matmul(coordinates_cell, cell)

class EnergyShifter(torch.nn.Module):
    """Helper class for adding and subtracting self atomic energies

    This is a subclass of :class:`torch.nn.Module`, so it can be used directly
    in a pipeline as ``[input->AEVComputer->ANIModel->EnergyShifter->output]``.

    Arguments:
        self_energies (:class:`collections.abc.Sequence`): Sequence of floating
            numbers for the self energy of each atom type. The numbers should
            be in order, i.e. ``self_energies[i]`` should be atom type ``i``.
        fit_intercept (bool): Whether to calculate the intercept during the LSTSQ
            fit. The intercept will also be taken into account to shift energies.
    """

    def __init__(self, self_energies, fit_intercept=False):
        super(EnergyShifter, self).__init__()

        self.fit_intercept = fit_intercept
        if self_energies is not None:
            self_energies = torch.tensor(self_energies, dtype=torch.double)

        self.register_buffer('self_energies', self_energies)

    def sae_from_dataset(self, atomic_properties, properties):
        """Compute atomic self energies from dataset.

        Least-squares solution to a linear equation is calculated to output
        ``self_energies`` when ``self_energies = None`` is passed to
        :class:`torchani.EnergyShifter`
        """
        species = atomic_properties['species']
        energies = properties['energies']
        present_species_ = present_species(species)
        X = (species.unsqueeze(-1) == present_species_).sum(dim=1).to(torch.double)
        # Concatenate a vector of ones to find fit intercept
        if self.fit_intercept:
            X = torch.cat((X, torch.ones(X.shape[0], 1).to(torch.double)), dim=-1)
        y = energies.unsqueeze(dim=-1)
        coeff_, _, _, _ = np.linalg.lstsq(X, y, rcond=None)
        return coeff_.squeeze()

    def sae(self, species):
        """Compute self energies for molecules.

        Padding atoms will be automatically excluded.

        Arguments:
            species (:class:`torch.Tensor`): Long tensor in shape
                ``(conformations, atoms)``.

        Returns:
            :class:`torch.Tensor`: 1D vector in shape ``(conformations,)``
            for molecular self energies.
        """
        intercept = 0.0
        if self.fit_intercept:
            intercept = self.self_energies[-1]

        self_energies = self.self_energies[species]
        self_energies[species == -1] = 0
        return self_energies.sum(dim=1) + intercept

    def subtract_from_dataset(self, atomic_properties, properties):
        """Transformer for :class:`torchani.data.BatchedANIDataset` that
        subtract self energies.
        """
        if self.self_energies is None:
            self_energies = self.sae_from_dataset(atomic_properties, properties)
            self.self_energies = torch.tensor(self_energies, dtype=torch.double)

        species = atomic_properties['species']
        energies = properties['energies']
        device = energies.device
        energies = energies.to(torch.double) - self.sae(species).to(device)
        properties['energies'] = energies
        return atomic_properties, properties

    def forward(self, species_energies):
        """(species, molecular energies)->(species, molecular energies + sae)
        """
        species, energies = species_energies
        sae = self.sae(species).to(energies.dtype).to(energies.device)
        return species, energies + sae

class ChemicalSymbolsToInts:
    """Helper that can be called to convert chemical symbol string to integers

    Arguments:
        all_species (:class:`collections.abc.Sequence` of :class:`str`):
            sequence of all supported species, in order.
    """

    def __init__(self, all_species):
        self.rev_species = {}
        for i, s in enumerate(all_species):
            self.rev_species[s] = i

    def __call__(self, species):
        """Convert species from squence of strings to 1D tensor"""
        rev = [self.rev_species[s] for s in species]
        return torch.tensor(rev, dtype=torch.long)

def hessian(coordinates, energies=None, forces=None):
    """Compute analytical hessian from the energy graph or force graph.

    Arguments:
        coordinates (:class:`torch.Tensor`): Tensor of shape `(molecules, atoms, 3)`
        energies (:class:`torch.Tensor`): Tensor of shape `(molecules,)`, if specified,
            then `forces` must be `None`. This energies must be computed from
            `coordinates` in a graph.
        forces (:class:`torch.Tensor`): Tensor of shape `(molecules, atoms, 3)`, if specified,
            then `energies` must be `None`. This forces must be computed from
            `coordinates` in a graph.

    Returns:
        :class:`torch.Tensor`: Tensor of shape `(molecules, 3A, 3A)` where A is the number of
        atoms in each molecule
    """
    if energies is None and forces is None:
        raise ValueError('Energies or forces must be specified')
    if energies is not None and forces is not None:
        raise ValueError('Energies or forces can not be specified at the same time')
    if forces is None:
        forces = -torch.autograd.grad(energies.sum(), coordinates, create_graph=True)[0]
    flattened_force = forces.flatten(start_dim=1)
    force_components = flattened_force.unbind(dim=1)
    return -torch.stack([
        torch.autograd.grad(f.sum(), coordinates, retain_graph=True)[0].flatten(start_dim=1)
        for f in force_components
    ], dim=1)

def vibrational_analysis(masses, hessian, unit='cm^-1'):
    """Computing the vibrational wavenumbers from hessian."""
    if unit != 'cm^-1':
        raise ValueError('Only cm^-1 are supported right now')
    assert hessian.shape[0] == 1, 'Currently only supporting computing one molecule a time'
    # Solving the eigenvalue problem: Hq = w^2 * T q
    # where H is the Hessian matrix, q is the normal coordinates,
    # T = diag(m1, m1, m1, m2, m2, m2, ....) is the mass
    # We solve this eigenvalue problem through Lowdin diagnolization:
    # Hq = w^2 * Tq ==> Hq = w^2 * T^(1/2) T^(1/2) q
    # Letting q' = T^(1/2) q, we then have
    # T^(-1/2) H T^(-1/2) q' = w^2 * q'
    inv_sqrt_mass = (1 / masses.sqrt()).repeat_interleave(3, dim=1)  # shape (molecule, 3 * atoms)
    mass_scaled_hessian = hessian * inv_sqrt_mass.unsqueeze(1) * inv_sqrt_mass.unsqueeze(2)
    if mass_scaled_hessian.shape[0] != 1:
        raise ValueError('The input should contain only one molecule')
    mass_scaled_hessian = mass_scaled_hessian.squeeze(0)
    eigenvalues, eigenvectors = torch.symeig(mass_scaled_hessian, eigenvectors=True)
    angular_frequencies = eigenvalues.sqrt()
    frequencies = angular_frequencies / (2 * math.pi)
    # converting from sqrt(hartree / (amu * angstrom^2)) to cm^-1
    wavenumbers = frequencies * 17092
    modes = (eigenvectors.t() * inv_sqrt_mass).reshape(frequencies.numel(), -1, 3)
    return wavenumbers, modes

__all__ = ['pad', 'pad_atomic_properties', 'present_species', 'hessian',
           'vibrational_analysis', 'strip_redundant_padding',
           'ChemicalSymbolsToInts']


"""
File: pylinex/util/RectangularBinner.py
Author: Keith Tauscher
Date: 20 Sep 2018

Description: File containing a class and function which bins using a
             rectangular window function defined by bin edges.
"""
from __future__ import division
import numpy as np
from distpy import Savable, Loadable, sequence_types, create_hdf5_dataset,\
    get_hdf5_value

class RectangularBinner(Savable, Loadable):
    """
    Class which bins using a rectangular window function defined by bin edges.
    """
    def __init__(self, unbinned_x_values, bin_edges):
        """
        Initializes a RectangularBinner with unbinned x values and edges of
        bins with which to define new points.

        unbinned_x_values: x_values associated with points to bin
        bin_edges: edges of bins with which to create new x values
        """
        self.unbinned_x_values = unbinned_x_values
        self.bin_edges = bin_edges
        self.digitize()

    def digitize(self):
        """
        Digitizes the bins so that binning can proceed quickly once data values
        are given.
        """
        bin_indices = np.digitize(self.unbinned_x_values, self.bin_edges)
        (bins_to_keep, unique_indices, unique_counts) =\
            np.unique(bin_indices, return_index=True, return_counts=True)
        if bins_to_keep[0] == 0:
            bins_to_keep = bins_to_keep[1:]
            unique_indices = unique_indices[1:]
            unique_counts = unique_counts[1:]
        if bins_to_keep[-1] == len(self.bin_edges):
            bins_to_keep = bins_to_keep[:-1]
            unique_indices = unique_indices[:-1]
            unique_counts = unique_counts[:-1]
        self._bins_to_keep = bins_to_keep - 1
        self._unique_indices = unique_indices
        self._unique_counts = unique_counts

    @property
    def unique_indices(self):
        """
        Property storing indices in the space of unbinned_x_values which share
        bins.
        """
        if not hasattr(self, '_unique_indices'):
            raise AttributeError("unique_indices was referenced before " +\
                "bins were digitized.")
        return self._unique_indices

    @property
    def unique_counts(self):
        """
        """
        if not hasattr(self, '_unique_counts'):
            raise AttributeError("unique_counts was referenced before bins " +\
                "were digitized.")
        return self._unique_counts

    @property
    def unbinned_x_values(self):
        """
        Property storing the x values of the unbinned points.
        """
        if not hasattr(self, '_unbinned_x_values'):
            raise AttributeError("unbinned_x_values was referenced before " +\
                "it was set.")
        return self._unbinned_x_values

    @unbinned_x_values.setter
    def unbinned_x_values(self, value):
        """
        Setter for the independent variable of the underlying model.
        """
        if type(value) in sequence_types:
            value = np.array(value)
            if value.ndim == 1:
                self._unbinned_x_values = value
            else:
                raise ValueError("unbinned_x_values was set to something " +\
                    "other than a 1D numpy.ndarray.")
        else:
            raise TypeError("unbinned_x_values was set to something other " +\
                "than a 1D numpy.ndarray.")

    @property
    def bin_edges(self):
        """
        Property storing the bin edges.
        """
        if not hasattr(self, '_bin_edges'):
            raise AttributeError("bin_edges was referenced before it was set.")
        return self._bin_edges

    @bin_edges.setter
    def bin_edges(self, value):
        """
        Setter for the bin edges.

        value: 1D array of bin edges (must be of length num_channels+1)
        """
        if type(value) in sequence_types:
            value = np.array(value)
            if value.ndim == 1:
                self._bin_edges = value
            else:
                raise ValueError("bin_edges was set to something other " +\
                    "than a 1D numpy.ndarray.")
        else:
            raise TypeError("bin_edges was set to something other than a " +\
                "1D numpy.ndarray.")

    @property
    def nbins(self):
        """
        Property storing the number of bins implied by the bin edges.
        """
        if not hasattr(self, '_nbins'):
            self._nbins = len(self.bin_edges) - 1
        return self._nbins

    @property
    def bins_to_keep(self):
        """
        Property storing the indices of bin with nonzero weight in order.
        """
        if not hasattr(self, '_bins_to_keep'):
            raise AttributeError("bins_to_keep referenced before it was set.")
        return self._bins_to_keep

    @property
    def nbins_to_keep(self):
        """
        Property storing the number of bins in the results of this binner.
        """
        if not hasattr(self, '_nbins_to_keep'):
            self._nbins_to_keep = len(self.bins_to_keep)
        return self._nbins_to_keep

    @property
    def binned_x_values(self):
        """
        Property storing the binned x values.
        """
        if not hasattr(self, '_binned_x_values'):
            self._binned_x_values = (self.bin_edges[self.bins_to_keep] +\
                self.bin_edges[self.bins_to_keep+1]) / 2
        return self._binned_x_values

    @property
    def x_samples(self):
        """
        Property storing the x values in each bin. It is a list of arrays.
        """
        if not hasattr(self, '_x_samples'):
            self._x_samples = [np.array([])] * self.nbins
            for (final_bin_index, original_bin_index) in\
                enumerate(self.bins_to_keep):
                index = self.unique_indices[final_bin_index]
                count = self.unique_counts[final_bin_index]
                self._x_samples[original_bin_index] =\
                    self.unbinned_x_values[index:index+count]
        return self._x_samples

    def bin(self, old_y_values, weights=None, return_weights=False):
        """
        Bins the given data.

        old_y_values: data to bin
        weights: weights to associate with each unbinned y value (should be of
                 same shape as old_y_values)
        return_weights: if True, weights are returned alongside binned data

        returns: if return_weights is False, new_y_values
                 if return_weights is True, (new_y_values, new_weights)
        """
        shape = old_y_values.shape[:-1] + (len(self.bins_to_keep),)
        new_y_values = np.zeros(shape)
        if type(weights) is type(None):
            weights = np.ones_like(old_y_values)
        if return_weights:
            new_weights = np.zeros(shape)
        for final_bin_index in range(len(self.bins_to_keep)):
            unique_index = self.unique_indices[final_bin_index]
            unique_count = self.unique_counts[final_bin_index]
            where = slice(unique_index, unique_index + unique_count)
            old_y_slice = old_y_values[...,where]
            weight_slice = weights[...,where]
            new_weight = np.sum(weight_slice, axis=-1)
            new_y_values[...,final_bin_index] =\
                np.sum(old_y_slice * weight_slice, axis=-1) / new_weight
            if return_weights:
                new_weights[...,final_bin_index] = new_weight
        if return_weights:
            return (new_y_values, new_weights)
        else:
            return new_y_values

    def bin_error(self, old_error, weights=None, return_weights=False):
        """
        Bins the given error vector(s).

        old_error: error to bin containing positive numbers with the last axis
                   being the binning axis.
        weights: weights to associate with each unbinned y value (should be of
                 same shape as old_y_values)
        return_weights: if True, weights are returned alongside binned data

        returns: if return_weights is False, new_y_values
                 if return_weights is True, (new_y_values, new_weights)
        """
        shape = old_error.shape[:-1] + (self.nbins_to_keep,)
        new_error = np.zeros(shape)
        if type(weights) is type(None):
            weights = np.ones_like(old_error)
        if return_weights:
            new_weights = np.zeros(shape)
        for final_bin_index in range(self.nbins_to_keep):
            unique_index = self.unique_indices[final_bin_index]
            unique_count = self.unique_counts[final_bin_index]
            where = slice(unique_index, unique_index + unique_count)
            old_error_slice = old_error[...,where]
            weight_slice = weights[...,where]
            new_weight = np.sum(weight_slice, axis=-1)
            new_error[...,final_bin_index] = np.sqrt(np.sum(np.power(\
                weight_slice * old_error_slice, 2), axis=-1)) / new_weight
            if return_weights:
                new_weights[...,final_bin_index] = new_weight
        if return_weights:
            return (new_error, new_weights)
        else:
            return new_error

    def __call__(self, old_y_values, weights=None, return_weights=False):
        """
        Bins the given data.

        old_y_values: data to bin
        weights: weights to associate with each unbinned y value (should be of
                 same shape as old_y_values)
        return_weights: if True, weights are returned alongside binned data

        returns: if return_weights is False, new_y_values
                 if return_weights is True, (new_y_values, new_weights)
        """
        return self.bin(old_y_values, weights=weights,\
            return_weights=return_weights)

    def fill_hdf5_group(self, group):
        """
        Fills the given hdf5 file group with data about this RectangularBinner.
        """
        create_hdf5_dataset(group, 'unbinned_x_values',\
            data=self.unbinned_x_values)
        create_hdf5_dataset(group, 'bin_edges', data=self.bin_edges)

    @staticmethod
    def load_from_hdf5_group(group):
        """
        Loads a RectangularBinner from the given hdf5 file group.

        group: group where RectangularBinner was once saved

        returns: RectangularBinner whose info was saved in the given group
        """
        unbinned_x_values = get_hdf5_value(group['unbinned_x_values'])
        bin_edges = get_hdf5_value(group['bin_edges'])
        return RectangularBinner(unbinned_x_values, bin_edges)

    def __eq__(self, other):
        """
        Checks if other is the same binner as this.

        other: object to check for equality

        returns: True only if other is a RectangularBinner with the same
                 unbinned x values and bin edges. False otherwise
        """
        if not isinstance(other, RectangularBinner):
            return False
        if (self.unbinned_x_values.shape != other.unbinned_x_values.shape) or\
            np.any(self.unbinned_x_values != other.unbinned_x_values):
            return False
        if (self.bin_edges.shape != other.bin_edges.shape) or\
            np.any(self.bin_edges != other.bin_edges):
            return False
        return True

    def __ne__(self, other):
        """
        Checks if other is not the same binner as this.

        other: object to check for inequality

        returns: False only if other is a RectangularBinner with the same
                 unbinned x values and bin edges. True otherwise
        """
        return (not self.__eq__(other))

def rect_bin(bin_edges, unbinned_x_values, old_y_values, weights=None,\
    return_weights=False):
    """
    Bins on final axis.

    bin_edges: 1D numpy.ndarray of length Nafter+1 containing bin edges
    unbinned_x_values: 1D numpy.ndarray of length Nbefore containing x values
                       of points
    old_y_values: ND numpy.ndarray whose final axis length is Nbefore
    return_weights: if True, extra array of new weights is returned
                    (Default False)

    returns: (binned_x_values, new_y_values[, new_weights])
             binned_x_values: 1D numpy.ndarray of length Nafter
             new_y_values: ND numpy.ndarray of the same shape as old_y_values
                           with Nafter instead of Nbefore as the final axis
                           length
             [new_weights]: new weights to associate with the binned data
                            points, given by the sum of the weights that went
                            into each bin (returned only if return_weights is
                            True)
    """
    binner = RectangularBinner(unbinned_x_values, bin_edges)
    bin_results =\
        binner(old_y_values, weights=weights, return_weights=return_weights)
    if return_weights:
        return (binner.binned_x_values,) + bin_results
    else:
        return (binner.binned_x_values, binner.new_y_values)