#
#    Licensed under the Apache License, Version 2.0 (the "License");
#    you may not use this file except in compliance with the License.
#    You may obtain a copy of the License at
#
#        http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS,
#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#    See the License for the specific language governing permissions and
#    limitations under the License.

"""
==================================================================================
  Compute Interference caused by a grant for all the incumbent types
  APIs in this file are used by IAP and Aggregate Interference Reference Models

  The main routines are:

    computeInterference
    computeInterferencePpaGwpzPoint
    computeInterferenceEsc
    computeInterferenceFssCochannel
    computeInterferenceFssBlocking
    getEffectiveSystemEirp

  The common utility APIs are:

    findOverlappingGrantsInsideNeighborhood
    getProtectedChannels

  The routines return a interference caused by a grant in the neighborhood of
  FSS/GWPZ/PPA/ESC incumbent types
==================================================================================
"""
from collections import namedtuple
import numpy as np

from reference_models.common import data
from reference_models.antenna import antenna
from reference_models.geo import vincenty
from reference_models.propagation import wf_itm
from reference_models.propagation import wf_hybrid

# Set constant parameters based on requirements in the WINNF-TS-0112
# [R2-SGN-16]
GWPZ_NEIGHBORHOOD_DIST = 40  # neighborhood distance from a CBSD to a given protection
# point (in km) in GWPZ protection area
PPA_NEIGHBORHOOD_DIST = 40  # neighborhood distance from a CBSD to a given protection
# point (in km) in PPA protection area

FSS_CO_CHANNEL_NEIGHBORHOOD_DIST = 150  # neighborhood distance from a CBSD to FSS for
# co-channel protection

FSS_BLOCKING_NEIGHBORHOOD_DIST = 40  # neighborhood distance from a CBSD to FSS
# blocking protection

ESC_NEIGHBORHOOD_DIST_A = 40  # neighborhood distance from a ESC to category A CBSD

ESC_NEIGHBORHOOD_DIST_B = 80  # neighborhood distance from a ESC to category B CBSD

# Frequency used in propagation model (in MHz) [R2-SGN-04]
FREQ_PROP_MODEL_MHZ = 3625.0

# CBRS Band Frequency Range (Hz)
CBRS_LOW_FREQ_HZ = 3550.e6
CBRS_HIGH_FREQ_HZ = 3700.e6

# FSS Passband low frequency range  (Hz)
FSS_LOW_FREQ_HZ = 3600.e6

# FSS Passband for TT&C (Hz)
FSS_TTC_LOW_FREQ_HZ = 3700.e6
FSS_TTC_HIGH_FREQ_HZ = 4200.e6

# ESC IAP for Out-of-Band Category A CBSDs in Frequency Range (Hz)
ESC_CAT_A_LOW_FREQ_HZ = 3550.e6
ESC_CAT_A_HIGH_FREQ_HZ = 3660.e6

# ESC Passband Frequency Range (Hz)
ESC_LOW_FREQ_HZ = 3550.e6
ESC_HIGH_FREQ_HZ = 3680.e6

# ESC Channel 21 Centre Frequency
ESC_CH21_CF_HZ = 3652.5e6

# One Mega Hertz
MHZ = 1.e6

# Channel bandwidth over which SASs execute the aggregate interference and
# IAP process
RBW_HZ = 5.e6

# GWPZ and PPA height (m)
GWPZ_PPA_HEIGHT = 1.5

# In-band insertion loss
IN_BAND_INSERTION_LOSS = 0.5

# Global container to store neighborhood distance type of all the protection
_DISTANCE_PER_PROTECTION_TYPE = {
    data.ProtectedEntityType.GWPZ_AREA :  (GWPZ_NEIGHBORHOOD_DIST, GWPZ_NEIGHBORHOOD_DIST),
    data.ProtectedEntityType.PPA_AREA : ( PPA_NEIGHBORHOOD_DIST,  PPA_NEIGHBORHOOD_DIST),
    data.ProtectedEntityType.FSS_CO_CHANNEL : ( FSS_CO_CHANNEL_NEIGHBORHOOD_DIST,  FSS_CO_CHANNEL_NEIGHBORHOOD_DIST),
    data.ProtectedEntityType.FSS_BLOCKING : ( FSS_BLOCKING_NEIGHBORHOOD_DIST,  FSS_BLOCKING_NEIGHBORHOOD_DIST),
    data.ProtectedEntityType.ESC: (ESC_NEIGHBORHOOD_DIST_A, ESC_NEIGHBORHOOD_DIST_B)
}

def dbToLinear(x):
  """This function returns dBm to mW converted value"""
  return 10**(x / 10.)

def linearToDb(x):
  """This function returns mW to dBm converted value"""
  return 10 * np.log10(x)

def getProtectedChannels(low_freq_hz, high_freq_hz):
  """Gets protected channels list.

  Performs 5MHz IAP channelization and returns a list of tuple containing
  (low_freq,high_freq)

  Args:
    low_freq_hz: Low frequency of the protected entity(Hz).
    high_freq_hz: High frequency of the protected entity(Hz)
  Returns:
    An array of protected channel frequency range tuple
    (low_freq_hz,high_freq_hz).
  """
  if low_freq_hz >= high_freq_hz:
    raise ValueError('Low frequency is greater than high frequency')
  # Align the low_freq to multiple of 5MHZ
  low_freq_hz = int(low_freq_hz / (5*MHZ)) * (5*MHZ)
  channels = np.arange( max(low_freq_hz, 3550*MHZ), min(high_freq_hz, 3700*MHZ), 5*MHZ)
  return [(low, high) for low, high in zip(channels, channels+5*MHZ)]

def findGrantsInsideNeighborhood(grants, protection_point, entity_type):
  """Finds grants inside protection entity neighborhood.

  Args:
    grants: An iterable of CBSD grants of type |data.CbsdGrantInfo|.
    protection_point: The location of a protected entity as (longitude, latitude) tuple.
    entity_type: The entity type (|data.ProtectedEntityType|).
  Returns:
    grants_inside: a list of grants, each one being a namedtuple of type
                   |data.CbsdGrantInfo|, of all CBSDs inside the neighborhood
                   of the protection constraint.
  """
  # Initialize an empty list
  grants_inside = []

  # Loop over each CBSD grant
  for grant in grants:
    # Compute distance from CBSD location to protection constraint location
    dist_km, _, _ = vincenty.GeodesicDistanceBearing(grant.latitude,
                      grant.longitude, protection_point[1], protection_point[0])

    # Check if CBSD is inside the neighborhood of protection constraint
    if dist_km <= _DISTANCE_PER_PROTECTION_TYPE[entity_type][grant.cbsd_category == 'B']:
      grants_inside.append(grant)

  return grants_inside

def grantFrequencyOverlapCheck(grant, ch_low_freq, ch_high_freq, protection_ent_type):
  """Checks if grant frequency overlaps with protection constraint frequency range.

  Args:
    grant: A CBSD grant of type |data.CbsdGrantInfo|.
    ch_low_freq: The low frequency of protection channel.
    ch_high_freq: The high frequency of protection channel.
    protection_ent_type: An enum of type |data.ProtectedEntityType|.
  Returns:
    True if grant frequency overlaps with protection constraint frequency range,
    False otherwise.
  """
  # Special case of ESC: ESC Passband is 3550-3680MHz
  # Category A CBSD grants are considered in the neighborhood only for
  # constraint frequency range 3550-3660MHz
  if (protection_ent_type == data.ProtectedEntityType.ESC and
      grant.cbsd_category == 'A' and
      (grant.low_frequency >= ESC_CAT_A_HIGH_FREQ_HZ or
       ch_low_freq >= ESC_CAT_A_HIGH_FREQ_HZ)):
    return False

  # Check frequency range overlap
  overlapping_bw = (min(grant.high_frequency, ch_high_freq)
                    - max(grant.low_frequency, ch_low_freq))
  return (overlapping_bw > 0)

def findOverlappingGrants(grants, constraint):
  """Finds grants overlapping with protection entity.

  Grants overlapping with frequency range of the protection entity are
  considered as overlapping grants.

  Args:
    grants: An iterable of CBSD grants of type |data.CbsdGrantInfo|.
    constraint: A protection constraint of type |data.ProtectionConstraint|.
  Returns:
    grants_overlap: a list of |data.CbsdGrantInfo| grants of all CBSDs inside
      the neighborhood of the protection constraint.
  """
  grants_overlap = [grant for grant in grants
                    if grantFrequencyOverlapCheck(
                        grant, constraint.low_frequency ,
                        constraint.high_frequency, constraint.entity_type)]
  return grants_overlap

def computeInterferencePpaGwpzPoint(cbsd_grant, constraint, h_inc_ant,
                                    max_eirp, region_type='SUBURBAN'):
  """Computes interference that a grant causes to GWPZ or PPA protection area.

  Routine to compute interference neighborhood grant causes to protection
  point within GWPZ or PPA protection area.

  Args:
    cbsd_grant: A CBSD grant of type |data.CbsdGrantInfo|.
    constraint: The protection constraint of type |data.ProtectionConstraint|.
    h_inc_ant: The reference incumbent antenna height (in meters).
    max_eirp: The maximum EIRP allocated to the grant during IAP procedure
    region_type: Region type of the GWPZ or PPA area:
                    'URBAN', 'SUBURBAN' or 'RURAL'.
  Returns:
    The interference contribution (dBm).
  """
  # Get the propagation loss and incident angles for area entity
  db_loss, incidence_angles, _ = wf_hybrid.CalcHybridPropagationLoss(
                                     cbsd_grant.latitude, cbsd_grant.longitude,
                                     cbsd_grant.height_agl, constraint.latitude,
                                     constraint.longitude, h_inc_ant,
                                     cbsd_grant.indoor_deployment,
                                     reliability=-1,
                                     freq_mhz=FREQ_PROP_MODEL_MHZ,
                                     region=region_type)

  # Compute CBSD antenna gain in the direction of protection point
  ant_gain = antenna.GetStandardAntennaGains(incidence_angles.hor_cbsd,
               cbsd_grant.antenna_azimuth, cbsd_grant.antenna_beamwidth,
               cbsd_grant.antenna_gain)

  # Get the exact overlap of the grant over the GWPZ area channels
  if constraint.entity_type == data.ProtectedEntityType.GWPZ_AREA:
    grant_overlap_bandwidth = min(cbsd_grant.high_frequency, constraint.high_frequency) \
        - max(cbsd_grant.low_frequency, constraint.low_frequency)
  else:
    grant_overlap_bandwidth = RBW_HZ

  # Get the interference value for area entity
  eirp = getEffectiveSystemEirp(max_eirp, cbsd_grant.antenna_gain,
                   ant_gain, grant_overlap_bandwidth)

  interference = eirp - db_loss
  return interference

def getEscMaskLoss(constraint):
  """Returns the ESC mask loss (in dB).

  Note: This routine assumes that the constraint bandwidth is 5MHz-aligned and
  fully contained in the CBSD grant so there is no partial overlap across the
  3650MHz ESC pass-band edge. This condition is realized in current framework.

  Args:
    constraint: The protection constraint of type |data.ProtectionConstraint|.
  """
  if constraint.high_frequency <= 3650*MHZ:
    return IN_BAND_INSERTION_LOSS
  if constraint.low_frequency < 3650*MHZ:
    raise ValueError('ESC mask loss: inconsistent protection channel %r' % constraint)
  freqs = np.arange(constraint.low_frequency + 0.5*MHZ, constraint.high_frequency, MHZ)
  attens = freqs/MHZ - 3650. + IN_BAND_INSERTION_LOSS
  atten = -linearToDb(np.mean(dbToLinear(-attens)))
  return atten

def computeInterferenceEsc(cbsd_grant, constraint, esc_antenna_info, max_eirp):
  """Computes interference that a grant causes to a ESC protection point.

  Routine to compute interference neighborhood grant causes to ESC protection
  point.

  Args:
    cbsd_grant: A CBSD grant of type |data.CbsdGrantInfo|.
    constraint: The protection constraint of type |data.ProtectionConstraint|.
    esc_antenna_info: ESC antenna information of type |data.EscInformation|.
    max_eirp: The maximum EIRP allocated to the grant during IAP procedure
  Returns:
    The interference contribution(dBm).
  """
  # Get the propagation loss and incident angles for ESC entity
  db_loss, incidence_angles, _ = wf_itm.CalcItmPropagationLoss(
      cbsd_grant.latitude, cbsd_grant.longitude, cbsd_grant.height_agl,
      constraint.latitude, constraint.longitude, esc_antenna_info.antenna_height,
      cbsd_grant.indoor_deployment, reliability=-1,
      freq_mhz=FREQ_PROP_MODEL_MHZ)

  # Compute CBSD antenna gain in the direction of protection point
  ant_gain = antenna.GetStandardAntennaGains(
      incidence_angles.hor_cbsd, cbsd_grant.antenna_azimuth,
      cbsd_grant.antenna_beamwidth, cbsd_grant.antenna_gain)

  # Compute ESC antenna gain in the direction of CBSD
  esc_ant_gain = antenna.GetAntennaPatternGains(
      incidence_angles.hor_rx,
      esc_antenna_info.antenna_azimuth,
      esc_antenna_info.antenna_gain_pattern)

  # Get the total antenna gain by summing the antenna gains from CBSD to ESC
  # and ESC to CBSD
  effective_ant_gain = ant_gain + esc_ant_gain

  # Compute the interference value for ESC entity
  eirp = getEffectiveSystemEirp(max_eirp, cbsd_grant.antenna_gain,
                                effective_ant_gain)
  interference = eirp - db_loss - getEscMaskLoss(constraint)
  return interference

def computeInterferenceFssCochannel(cbsd_grant, constraint, fss_info, max_eirp):
  """Computes interference that a grant causes to a FSS protection point.

  Routine to compute interference neighborhood grant causes to FSS protection
  point for co-channel passband.

  Args:
    cbsd_grant: A CBSD grant of type |data.CbsdGrantInfo|.
    constraint: The protection constraint of type |data.ProtectionConstraint|.
    fss_info: The FSS information of type |data.FssInformation|.
    max_eirp: The maximum EIRP allocated to the grant during IAP procedure.
  Returns:
    The interference contribution(dBm).
  """
  # Get the propagation loss and incident angles for FSS entity_type
  db_loss, incidence_angles, _ = wf_itm.CalcItmPropagationLoss(cbsd_grant.latitude,
                                   cbsd_grant.longitude, cbsd_grant.height_agl,
                                   constraint.latitude, constraint.longitude,
                                   fss_info.height_agl, cbsd_grant.indoor_deployment,
                                   reliability=-1, freq_mhz=FREQ_PROP_MODEL_MHZ)

  # Compute CBSD antenna gain in the direction of protection point
  ant_gain = antenna.GetStandardAntennaGains(incidence_angles.hor_cbsd,
               cbsd_grant.antenna_azimuth, cbsd_grant.antenna_beamwidth,
               cbsd_grant.antenna_gain)

  # Compute FSS antenna gain in the direction of CBSD
  fss_ant_gain = antenna.GetFssAntennaGains(incidence_angles.hor_rx,
                   incidence_angles.ver_rx, fss_info.pointing_azimuth,
                   fss_info.pointing_elevation, fss_info.max_gain_dbi)

  # Get the total antenna gain by summing the antenna gains from CBSD to FSS
  # and FSS to CBSD
  effective_ant_gain = ant_gain + fss_ant_gain

  # Compute the interference value for Fss co-channel entity
  eff_bandwidth = (min(cbsd_grant.high_frequency, constraint.high_frequency)
                   - max(cbsd_grant.low_frequency, constraint.low_frequency))
  if eff_bandwidth <= 0:
    raise ValueError('Computing FSS co-channel on grant fully outside FSS passband')

  eirp = getEffectiveSystemEirp(max_eirp, cbsd_grant.antenna_gain,
                                effective_ant_gain, eff_bandwidth)
  interference = eirp - db_loss - IN_BAND_INSERTION_LOSS
  return interference

def getFssMaskLoss(cbsd_grant, constraint):
  """Gets the FSS mask loss for a FSS blocking protection constraint.

  Args:
    cbsd_grant: A CBSD grant of type |data.CbsdGrantInfo|.
    constraint: The protection constraint of type |data.ProtectionConstraint|.
      The constraint defines the FSS out of band left part, so typically
      [3550, min_passband_freq] or [3550, 3700].
  """
  # Sanity checks
  if (cbsd_grant.low_frequency >= cbsd_grant.high_frequency or
      cbsd_grant.low_frequency >= constraint.high_frequency):
    raise ValueError('CBSD grant frequencies incorrect')

  # Find the 50MHz edge and its rounded version
  edge_freq = constraint.high_frequency - 50*MHZ
  edge_freq_round = int(edge_freq/MHZ +0.5)*MHZ
  # Part of grant in closest mask segment
  seg1 = (max(cbsd_grant.low_frequency, edge_freq_round),
          min(cbsd_grant.high_frequency, constraint.high_frequency))
  # Part of grant in farther mask segment
  seg2 = (cbsd_grant.low_frequency,
          min(cbsd_grant.high_frequency, edge_freq_round))

  # Now compute the attenuation
  freqs1 = np.arange(seg1[0] + 0.5*MHZ, seg1[1], MHZ)
  freqs2 = np.arange(seg2[0] + 0.5*MHZ, seg2[1], MHZ)
  attens1 = (constraint.high_frequency - freqs1)/MHZ * 0.6 + 0.5
  attens2 = (edge_freq - freqs2)/MHZ * 0.25 + 30.5
  attens = np.concatenate((attens1, attens2))
  fss_mask_attenuation = -linearToDb(np.mean(dbToLinear(-attens)))
  return fss_mask_attenuation

def computeInterferenceFssBlocking(cbsd_grant, constraint, fss_info, max_eirp):
  """Computes interference that a grant causes to a FSS protection point.

  Routine to compute interference neighborhood grant causes to FSS protection
  point for blocking passband

  Args:
    cbsd_grant: A CBSD grant of type |data.CbsdGrantInfo|.
    constraint: The protection constraint of type |data.ProtectionConstraint|.
    fss_info: The FSS information of type |data.FssInformation|.
    max_eirp: The maximum EIRP allocated to the grant during IAP procedure.

  Returns:
    The interference contribution(dBm).
  """
  # Get the propagation loss and incident angles for FSS entity
  # blocking channels
  db_loss, incidence_angles, _ = wf_itm.CalcItmPropagationLoss(
                                   cbsd_grant.latitude, cbsd_grant.longitude,
                                   cbsd_grant.height_agl, constraint.latitude,
                                   constraint.longitude, fss_info.height_agl,
                                   cbsd_grant.indoor_deployment, reliability=-1,
                                   freq_mhz=FREQ_PROP_MODEL_MHZ)

  # Compute CBSD antenna gain in the direction of protection point
  ant_gain = antenna.GetStandardAntennaGains(incidence_angles.hor_cbsd,
               cbsd_grant.antenna_azimuth, cbsd_grant.antenna_beamwidth,
               cbsd_grant.antenna_gain)

  # Compute FSS antenna gain in the direction of CBSD
  fss_ant_gain = antenna.GetFssAntennaGains(incidence_angles.hor_rx,
                   incidence_angles.ver_rx, fss_info.pointing_azimuth,
                   fss_info.pointing_elevation, fss_info.max_gain_dbi)

  # Get the total antenna gain by summing the antenna gains from CBSD to FSS
  # and FSS to CBSD
  effective_ant_gain = ant_gain + fss_ant_gain

  # Compute EIRP of CBSD grant inside the frequency range of
  # protection constraint
  eff_bandwidth = (min(cbsd_grant.high_frequency, constraint.high_frequency)
                   - cbsd_grant.low_frequency)
  if eff_bandwidth <= 0:
    raise ValueError('Computing FSS blocking on grant fully inside FSS passband')
  eirp = getEffectiveSystemEirp(max_eirp, cbsd_grant.antenna_gain,
                                effective_ant_gain, eff_bandwidth)
  # Calculate the interference contribution
  interference = eirp - getFssMaskLoss(cbsd_grant, constraint) - db_loss

  return interference

def getEffectiveSystemEirp(max_eirp, cbsd_max_ant_gain, effective_ant_gain,
                           reference_bandwidth=RBW_HZ):
  """Calculates effective EIRP caused by a grant.

  Utility API to get effective EIRP caused by a grant in the
  neighborhood of the protected entity FSS/ESC/PPA/GWPZ.

  Args:
    max_eirp: The maximum EIRP allocated to the grant during IAP procedure.
    cbsd_max_ant_gain: The nominal antenna gain of the CBSD.
    effective_ant_gain: The actual total antenna gains at the CBSD and protected
      entity. This takes into account the actual antenna patterns.
    reference_bandwidth: Reference bandwidth over which effective EIRP is calculated.
  Returns:
    The effective EIRP of the CBSD(dBm)
  """

  eirp_cbsd = ((max_eirp - cbsd_max_ant_gain) + effective_ant_gain +
               linearToDb(reference_bandwidth / MHZ))

  return eirp_cbsd

def computeInterference(grant, eirp, constraint,
                        fss_info=None, esc_antenna_info=None, region_type=None):
  """Calculates interference caused by a grant.

  Utility API to get interference caused by a grant in the
  neighborhood of the protected entity FSS/ESC/PPA/GWPZ.

  Args:
    grant: A CBSD grant of type |data.CbsdGrantInfo|.
    eirp: The EIRP of the grant.
    constraint: A protection constraint of type |data.ProtectionConstraint|.
    fss_info: The FSS information of type |data.FssInformation|.
    esc_antenna_info: ESC antenna information of type |data.EscInformation|.
    region_type: Region type of the GWPZ or PPA area:
                   'URBAN', 'SUBURBAN' or 'RURAL'.
  Returns:
    interference: Interference caused by a grant(dBm)
  """
  # Compute interference to FSS Co-channel protection constraint
  if constraint.entity_type is data.ProtectedEntityType.FSS_CO_CHANNEL:
    interference = computeInterferenceFssCochannel(
                     grant, constraint, fss_info, eirp)

  # Compute interference to FSS Blocking protection constraint
  elif constraint.entity_type is data.ProtectedEntityType.FSS_BLOCKING:
    interference = computeInterferenceFssBlocking(
                     grant, constraint, fss_info, eirp)

  # Compute interference to ESC protection constraint
  elif constraint.entity_type is data.ProtectedEntityType.ESC:
    interference = computeInterferenceEsc(
                     grant, constraint, esc_antenna_info, eirp)

  # Compute interference to GWPZ or PPA protection constraint
  else:
    interference = computeInterferencePpaGwpzPoint(
                     grant, constraint, GWPZ_PPA_HEIGHT,
                     eirp, region_type)

  return interference


"""
Todo's
"""
import matplotlib.pyplot as plt
import matplotlib.cm as cm
import copy
import numpy as np
import uncertainties as uct
from uncertainties import unumpy as unp
from ..etc import isuncertainties

c_map = cm.jet  # mpl.cm.inferno #mpl.cm.CMRmap_r # mpl.cm.brg #mpl.cm.gray_r #

def thermal_data(data, figsize=(12, 4), ms_data=50,
                 v_label='Unit-cell volume $(\mathrm{\AA}^3)$',
                 pdf_filen=None, title='P-V-T data'):
    """
    plot P-V-T data before fitting

    :param data: {'p': unumpy array, 'v': unumpy array, 'temp': unumpy array}
    :param eoscurves: {'v': unumpy array, '300': unumpy array
            at the temperature ....}
    :param v_label: label for volume axis
    :param figsize: figure size
    :param ms_data: marker size for data points
    :param pdf_filen: name of pdf output file
    :param title: title of the figure
    :return: None
    """
    # basic figure setup
    f, ax = plt.subplots(1, 2, figsize=figsize, sharex=True)

    # read data to plot
    if isuncertainties([data['p'], data['v'], data['temp']]):
        p = unp.nominal_values(data['p'])
        v = unp.nominal_values(data['v'])
        temp = unp.nominal_values(data['temp'])
        sp = unp.std_devs(data['p'])
        sv = unp.std_devs(data['v'])
        stemp = unp.std_devs(data['temp'])
        ax[0].errorbar(p, v, xerr=sp, yerr=sv, marker=' ',
                       c='k', ms=0, mew=0, linestyle='None',
                       capsize=0, lw=0.5, zorder=1)
        ax[1].errorbar(p, temp, xerr=sp, yerr=stemp, marker=' ',
                       c='k', ms=0, mew=0, linestyle='None',
                       capsize=0, lw=0.5, zorder=1)
    else:
        p = data['p']
        v = data['v']
        temp = data['temp']
    points = ax[0].scatter(p, v, marker='o', s=ms_data, c=temp,
                           cmap=c_map, vmin=300., vmax=temp.max(), zorder=2)
    points = ax[1].scatter(p, temp, marker='o', s=ms_data, c=temp,
                           cmap=c_map, vmin=300., vmax=temp.max(), zorder=2)

    ax[0].set_xlabel('Pressure (GPa)')
    ax[1].set_xlabel('Pressure (GPa)')
    ax[0].set_ylabel(v_label)
    ax[1].set_ylabel('Temperature (K)')
    f.suptitle(title)
    # the parameters are the specified position you set
    position = f.add_axes([0.92, 0.11, .01, 0.75])
    f.colorbar(points, orientation="vertical", cax=position)
    # position.text(150., 0.5, 'Temperature (K)', fontsize=10,
    # rotation=270, va='center')
    if pdf_filen is not None:
        f.savefig(pdf_filen)

c_map = cm.jet

def thermal_fit_result(fit_result, v_residual=None,
                       v_label='Unit-cell volume $(\mathrm{\AA}^3)$',
                       temp_fitline=np.asarray(
                           [300., 1000., 1500., 2000., 2500., 3000.]),
                       figsize=(5, 5), height_ratios=(3, 1), ms_data=50,
                       p_err=None, v_err=None, cbar_loc=(0.99, 0.1, .01, 0.82),
                       pdf_filen=None, title='Fit result'):
    """
    plot P-V-T EOS curve fitting result

    :param fit_result: lmfit result object, see example jnb file for detail
    :param v_label: label for volume axis
    :param temp_fitline: temperatures to calculate isothermal compression
        curves, default = [300., 1000., 1500., 2000., 2500., 3000.]
    :param figsize: figure size, default = (7,7)
    :param height_ratios: height ratio between the main and residue plots,
        default = (3,1)
    :param ms_data: marker size for data points
    :param p_err: pressure error bar
    :param v_err: volume error bar
    :param cbar_loc: location of color bar
    :param pdf_filen: name of pdf output file
    :param title: title of the figure
    :return: None
    """
    # basic figure setup
    f, ax = plt.subplots(2, 1, sharex=True, figsize=figsize,
                         gridspec_kw={'height_ratios': height_ratios})
    for ax_i in ax:
        ax_i.tick_params(direction='in')
    # read data to plot
    v_data = fit_result.userkws['v']
    temp_data = fit_result.userkws['temp']
    p_data = fit_result.data
    p_datafit = fit_result.best_fit
    v0 = uct.ufloat(fit_result.params['st_v0'].value,
                    fit_result.params['st_v0'].stderr)
    sm = plt.cm.ScalarMappable(cmap=c_map,
                               norm=plt.Normalize(
                                   vmin=300., vmax=temp_data.max()))
    a = sm.to_rgba(temp_fitline)
    v_fitline = np.linspace(v0.n, min(v_data), 1000)
    fitmodel_copy = copy.deepcopy(fit_result)
    for a_i, temp_i in zip(a, temp_fitline):
        p_fitline = fitmodel_copy.eval(v=v_fitline,
                                       temp=np.ones_like(v_fitline) * temp_i)
        ax[0].plot(p_fitline, v_fitline, c=a_i)
    # error range here does not make a lot sense, so not supported
    # if (p_err is not None) and (v_err is not None):
    ax[0].errorbar(p_data, v_data, xerr=p_err, yerr=v_err, fmt=' ', c='k',
                   capsize=0, elinewidth=0.5, label='Data', zorder=0)
    points = ax[0].scatter(p_data, v_data, marker='o', s=ms_data, c=temp_data,
                           cmap=c_map, vmin=300., vmax=temp_data.max(),
                           zorder=1)
    if v_residual is None:
        ax[1].scatter(p_data, p_data - p_datafit, marker='o', s=ms_data,
                      c=temp_data, cmap=c_map, vmin=300.,
                      vmax=temp_data.max(), zorder=1)
        ax[1].errorbar(p_data, p_data - p_datafit, yerr=p_err, fmt=' ', c='k',
                       capsize=0, elinewidth=0.5, label='Data', zorder=0)
        ax[1].set_ylabel('$P_{obs} - P_{fit}$')
    else:
        ax[1].scatter(p_data, v_residual, marker='o', s=ms_data, c=temp_data,
                      cmap=c_map, vmin=300., vmax=temp_data.max(), zorder=1)
        ax[1].errorbar(p_data, v_residual, yerr=p_err, fmt=' ', c='k',
                       capsize=0, elinewidth=0.5, label='Data', zorder=0)
        ax[1].set_ylabel('$V_{obs} - V_{fit}$')
    # ax[0].legend()
    position = f.add_axes(cbar_loc)
    f.colorbar(points, orientation="vertical", cax=position,
               ticks=temp_fitline)
    ax[1].axhline(0, c='k', ls='--')
    ax[1].set_xlabel('Pressure (GPa)')
    ax[0].set_ylabel(v_label)
    ax[0].set_title(title)
    plt.tight_layout()
    if pdf_filen is not None:
        f.savefig(pdf_filen)


"""
This class implements simple policy gradient algorithm for
biasing the generation of molecules towards desired values of
properties aka Reinforcement Learninf for Structural Evolution (ReLeaSE)
as described in
Popova, M., Isayev, O., & Tropsha, A. (2018).
Deep reinforcement learning for de novo drug design.
Science advances, 4(7), eaap7885.
"""

import torch
import torch.nn.functional as F
import numpy as np
from rdkit import Chem

class Reinforcement(object):
    #def __init__(self, generator, predictor, get_reward):
    def __init__(self, generator, get_reward):
        """
        Constructor for the Reinforcement object.

        Parameters
        ----------
        generator: object of type StackAugmentedRNN
            generative model that produces string of characters (trajectories)

        predictor: object of any predictive model type
            predictor accepts a trajectory and returns a numerical
            prediction of desired property for the given trajectory

        get_reward: function
            custom reward function that accepts a trajectory, predictor and
            any number of positional arguments and returns a single value of
            the reward for the given trajectory
            Example:
            reward = get_reward(trajectory=my_traj, predictor=my_predictor,
                                custom_parameter=0.97)

        Returns
        -------
        object of type Reinforcement used for biasing the properties estimated
        by the predictor of trajectories produced by the generator to maximize
        the custom reward function get_reward.
        """

        super(Reinforcement, self).__init__()
        self.generator = generator
        #self.predictor = predictor
        self.get_reward = get_reward

    def policy_gradient(self, data, n_batch=10, gamma=0.97,
                        std_smiles=False, grad_clipping=None, **kwargs):
        """
        Implementation of the policy gradient algorithm.

        Parameters:
        -----------

        data: object of type GeneratorData
            stores information about the generator data format such alphabet, etc

        n_batch: int (default 10)
            number of trajectories to sample per batch. When training on GPU
            setting this parameter to to some relatively big numbers can result
            in out of memory error. If you encountered such an error, reduce
            n_batch.

        gamma: float (default 0.97)
            factor by which rewards will be discounted within one trajectory.
            Usually this number will be somewhat close to 1.0.

        std_smiles: bool (default False)
            boolean parameter defining whether the generated trajectories will
            be converted to standardized SMILES before running policy gradient.
            Leave this parameter to the default value if your trajectories are
            not SMILES.

        grad_clipping: float (default None)
            value of the maximum norm of the gradients. If not specified,
            the gradients will not be clipped.

        kwargs: any number of other positional arguments required by the
            get_reward function.

        Returns
        -------
        total_reward: float
            value of the reward averaged through n_batch sampled trajectories

        rl_loss: float
            value for the policy_gradient loss averaged through n_batch sampled
            trajectories

        """
        rl_loss = 0
        self.generator.optimizer.zero_grad()
        total_reward = 0

        for _ in range(n_batch):

            # Sampling new trajectory
            reward = 0
            trajectory = '<>'
            while reward == 0:
                trajectory = self.generator.evaluate(data)
                if std_smiles:
                    try:
                        mol = Chem.MolFromSmiles(trajectory[1:-1])
                        trajectory = '<' + Chem.MolToSmiles(mol) + '>'
                        #reward = self.get_reward(trajectory[1:-1],
                        #                         self.predictor,
                        #                         **kwargs)
                        reward = self.get_reward(trajectory[1:-1])
                    except:
                        reward = 0
                else:
                    #reward = self.get_reward(trajectory[1:-1],
                    #                         self.predictor,
                    #                         **kwargs)
                    reward = self.get_reward(trajectory[1:-1])

            # Converting string of characters into tensor
            trajectory_input = data.char_tensor(trajectory)
            discounted_reward = reward
            total_reward += reward

            # Initializing the generator's hidden state
            hidden = self.generator.init_hidden()
            if self.generator.has_cell:
                cell = self.generator.init_cell()
                hidden = (hidden, cell)
            if self.generator.has_stack:
                stack = self.generator.init_stack()
            else:
                stack = None

            # "Following" the trajectory and accumulating the loss
            for p in range(len(trajectory)-1):
                output, hidden, stack = self.generator(trajectory_input[p],
                                                       hidden,
                                                       stack)
                log_probs = F.log_softmax(output, dim=1)
                top_i = trajectory_input[p+1]
                rl_loss -= (log_probs[0, top_i]*discounted_reward)
                discounted_reward = discounted_reward * gamma

        # Doing backward pass and parameters update
        rl_loss = rl_loss / n_batch
        total_reward = total_reward / n_batch
        rl_loss.backward()
        if grad_clipping is not None:
            torch.nn.utils.clip_grad_norm_(self.generator.parameters(),
                                           grad_clipping)

        self.generator.optimizer.step()

        return total_reward, rl_loss.item()


"""Implements a factored random geometric graph

A factored random geometric graph is special representation of a random
geometric graph, for the case when the configuration space is the
Cartesian product of a numbor of component manifolds. The particular
case of interest is the manifold of configurations of K objects, which
is the Cartesian product of K copies SE(2) (in two dimensions) or SE(3)
in three dimensions.

The vertices can be represented by length-K tuples, where each element
is an index into a sequence of configurations in one of the component
manifolds. The precise form of edges required to ensure probabilistic
completeness and asymptotic optimality is an object of ongoing research.

The classes in this module take objects parameterizing dynamics and
geometry, and do the bookkeeping necessary to maintain a factored random
geometric graph.
"""
import math
import itertools

import numpy
from scipy.spatial import cKDTree

import metis
from metis.abstract_graphs import UndirectedGraph

class NoObjectContactBlacklist(object):
    def __init__(self, robot='robot'):
        super(NoObjectContactBlacklist, self).__init__()
        self.robot = robot

    def __contains__(self, pair):
        return pair[0] != self.robot and pair[1] is not None

def apply_transform(tfA, tfB):
    """Return the result of applying tfA to tfB

    Args:
        tfA (tuple): a transform, represented as an (x, y, theta) tuple
        tfB (tuple): a transform, represented as an (x, y, theta) tuple

    Returns:
        tuple: (x, y, theta) representing a transform with the same
            effect as first applying 'tfB' and then applying 'tfA'. If
            the transforms were represented as matrices, this would be
            tfA * tfB.
    """
    return (tfA[0] + numpy.cos(tfA[2]) * tfB[0] - numpy.sin(tfA[2]) * tfB[1],
            tfA[1] + numpy.sin(tfA[2]) * tfB[0] + numpy.cos(tfA[2]) * tfB[1],
            tfA[2] + tfB[2])

def apply_inverse_transform(tfA, tfB):
    """Return the result of applying tfA to the inverse of tfB

    Args:
        tfA (tuple): a transform, represented as an (x, y, theta) tuple
        tfB (tuple): a transform, represented as an (x, y, theta) tuple

    Returns:
        tuple: (x, y, theta) representing a transform with the same
            effect as first applying the inverse of 'tfB' and then
            applying 'tfA'. If the transforms were represented as
            matrices, this would be tfA / tfB.
    """
    dtf = tfA - tfB
    return (numpy.cos(tfB[2]) * dtf[0] + numpy.sin(tfB[2]) * dtf[1],
            -numpy.sin(tfB[2]) * dtf[0] + numpy.cos(tfB[2]) * dtf[1],
            dtf[2])

class Manifold(object):
    def __init__(self, samples):
        self.samples = numpy.array(samples)
        assert self.samples.ndim == 2, "Data must be a sequence of samples"

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

    def __iter__(self):
        return iter(xrange(len(self.samples)))

    def __contains__(self, index):
        return 0 <= index < len(self.samples)

    def __getitem__(self, index):
        return self.samples[index]

    def nearest_samples(self, point, k):
        """Return then indices of the samples near the query pose"""
        pass

    def nearby_samples(self, point, r):
        """Return then indices of the samples near the query pose"""
        pass

    def neighbors(self, index):
        """Return the neighbors of the specified sample index"""
        pass

class SE2Manifold(Manifold):
    """Random disc graph on a manifold

    A random disc graph has vertices sampled from an n dimensional
    manifold, and an edge between all vertices with minimal geodesic
    distance less than a fixed constant. We take that constant to be the
    critical distance derived in Karaman and Frazzoli (2011), which is
    minimal distance required to ensure the resulting graph includes an
    optimal path.

    Vertex labels are just indices of an array of points, so any integer
    between 0 and `count-1`.

    Args:
        samples (array-like): if supplied, include these
            configurations as vertices in the graph, in addition to
            `count` random configurations.
        epsilon (float): positive scalar multiplier by which to
            approximate nearest neighbor search (larger is faster but
            less accurate)
        eta (float): positive scalar multiplier by which to inflate the
            distance threshhold defining the graph (larger is slower but
            more accurate; graph is asymptotically optimal iff eta>=1)

    Attributes:
        geometry (metis.geometry.Geometry): object describing problem geometry
        configurations (numpy.array): array of samples corresponding to
            vertices. Each row is a sample from the configuration space
            defined by geometry.
        epsilon (float): positive scalar multiplier by which to
            approximate nearest neighbor search (larger is faster but
            less accurate)
        eta (float): positive scalar multiplier by which to inflate the
            distance threshhold defining the graph (larger is slower but
            more accurate; graph is asymptotically optimal iff eta>=1)
        mu_free (float); positive scalar approximating the Lebesgue
            measure of free space in the configuration
        search_radius (float): positive scalar defining maximum distance
            between vertices for which an edge may exist
        nearby_configurations (list(list)): list of lists of neighboring
            vertices for each vertex in the graph. Precomputed for speed
    """
    def __init__(self, samples, scale=1., epsilon=0., search_radius=None, eta=1.):
        super(SE2Manifold, self).__init__(samples)
        assert self.samples.shape[1] == 3, \
            "Data must be 'x, y, theta' sequences"

        self.scale = scale
        self.epsilon = epsilon
        self.eta = eta

        if search_radius is None:
            count, dim = self.samples.shape
            d_inv = 1. / float(dim)
            zeta_d = math.pi**(dim/2) / math.gamma(dim/2 + 1)

            mu_free = (numpy.prod(numpy.ptp(self.samples[:,:-1], axis=0))
                       * 2 * numpy.pi * scale)
            self.search_radius = 2 * self.eta * (
                d_inv * (mu_free / zeta_d) * (math.log(count) / float(count))
                )**d_inv
        else:
            self.search_radius = search_radius

        self._search = cKDTree(numpy.concatenate(
            (self.samples[:, :-1],
             self.scale * numpy.cos(self.samples[:, -1:]),
             self.scale * numpy.sin(self.samples[:, -1:])), axis=1))
        self._neighbors = self._search.query_ball_tree(
            self._search, (1 + epsilon) * self.search_radius, eps=epsilon)

    def nearest_samples(self, point, k):
        """Return the indices of the k samples nearest the query point

        Args:
            point: query location, represented as an (x, y, theta) tuple

        Returns:
            list: list of tuples (d, i), where d is the distance between
                sample i and the query point
        """
        query = (point[0], point[1],
                 self.scale * numpy.cos(point[2]),
                 self.scale * numpy.sin(point[2]))
        return zip(*self._search.query(query, k))

    def nearby_samples(self, point, r):
        """Return the indices of the samples near the query point

        Args:
            point: query location, represented as an (x, y, theta) tuple
            r: search radius

        Returns:
            list: list of indices of samples within distance r of the
                query point
        """
        query = (point[0], point[1],
                 self.scale * numpy.cos(point[2]),
                 self.scale * numpy.sin(point[2]))
        return self._search.query_ball_point(
            query, (1+self.epsilon)*r, eps=self.epsilon)

    def neighbors(self, vertex):
        """Generate all vertices less than threshhold from vertex

        This is computationally inexpensive due to precomputation.
        Note:
            Does not perform collision detection; some edges returned
            here may not be collision free. Collision detection is
            performed when the cost of an edge is evaluated.

        Args:
            vertex (int): index of a vertex in the graph

        Returns:
            generator(int): generates all neighbors of vertex
        """
        return (n for n in self._neighbors[vertex] if n != vertex)

    def distance(self, pose1, pose2):
        point1 = numpy.array((pose1[0], pose1[1],
                              self.scale * numpy.cos(pose1[2]),
                              self.scale * numpy.sin(pose1[2])))
        point2 = numpy.array((pose2[0], pose2[1],
                              self.scale * numpy.cos(pose2[2]),
                              self.scale * numpy.sin(pose2[2])))
        return numpy.sqrt(numpy.sum((point1-point2)**2))

class FactoredRandomGeometricGraph(UndirectedGraph):
    """Factored random geometric graph for manipulation

    The vertices of this graph are hashable dict-like structures mapping
    object names to tuples of (parent name, transform index), where
    parent_name is the name of the object relative to which the
    transform of the object name is computed. A vertex can be unpacked into
    configurations using the notation graph[vertex]. The graph is
    immutable after construction; vertex data cannot be changed or
    deleted.

    Args:
        geometry (ManyShapeGeometry): object describing problem geometry
        dynamics: used to determine how each object can move
        robot (str): the identifier of the object to consider the
            robot for planning purposes
        geometry (metis.geometry.ManyShapeGeometry): object describing
            problem geometry
        counts (dict): maps manifolds, specified as (object, object)
            2-tuples, to the number of random samples to draw from that
            manifold
        default_count (int): number of poses of each object to sample.
            This value is used for any manifold not specified in counts,
            and is overridden by any manifold specified in counts.
        configurations (dict): if supplied, include these
            configurations in the component set for the supplied
            objects, in addition to `count` random configurations.
        epsilon (float): positive scalar multiplier by which to
            approximate nearest neighbor search (larger is faster but
            less accurate)
        eta (float): positive scalar multiplier by which to inflate the
            distance threshhold defining the graph (larger is slower but
            more accurate; graph is asymptotically optimal iff eta>=1)

    Attributes:
        geometry (ManyShapeGeometry): object describing problem geometry
        dynamics: used to determine how each object can move
        robot (str): the identifier of the object to consider the
            robot for planning purposes
        manifolds (dict): maps manifold names to structures defining the
            manifolds themselves
        epsilon (float): positive scalar multiplier by which to
            approximate nearest neighbor search (larger is faster but
            less accurate)
        eta (float): positive scalar multiplier by which to inflate the
            distance threshhold defining the graph (larger is slower but
            more accurate; graph is asymptotically optimal iff eta>=1)
    """
    def __init__(self, geometry, dynamics, counts=None, default_count=0,
                 configurations=None, blacklist=(), robot='robot',
                 epsilon=0., eta=1.):
        # pylint: disable=too-many-arguments
        assert robot in geometry.bodies

        # Rebind default mutable arguments
        if counts is None:
            counts = {}
        if configurations is None:
            configurations = {}

        # Declare members
        self.geometry = geometry
        self.dynamics = dynamics
        self.names = [name for name in geometry.bodies]
        self.samples = {}
        self.robot = robot
        self.epsilon = epsilon
        self.eta = eta

        for name in self.names:
            transforms = configurations.get((name, None), [])
            count = counts.get((name, None), default_count)
            transforms += [geometry.sample_object_configuration(name)
                           for _ in xrange(count)]
            self.samples[(name, None)] = SE2Manifold(transforms, epsilon=epsilon, eta=eta)

            # Sample poses relative to other objects
            for other in self.names:
                if name != other and (name, other) not in blacklist:
                    count = counts.get((name, other), default_count)
                    poses = list(dynamics.sample_from_manifold((name, other), count))
                    self.samples[(name, other)] = SE2Manifold(
                        [pose for pose in poses])
                else:
                    self.samples[(name, other)] = Manifold([[]])

    @staticmethod
    def is_acyclic(vertex):
        """Determine if a directed chain graph is acyclic

        An acyclic chain is one name for a directed graph where each
        vertex is the source of at most one outward edge, and there are
        no cycles of directed edges. So for vertices (A, B, C), A=>B=>C
        would be an acyclic chain, while A=>B=>C=>A would not.

        Args:
            vertex (dict): maps vertex labels to a tuple (parent, None),
                describing the unique edge beginning at the vertex. A
                missing edge is denoted as an edge back to the vertex,
                so `{name: (name, None)}` describes a vertex with label
                `name` and no outgoing edge. The choice to represent an
                edge as a tuple with second value None was made for
                uniformity with the vertex representation, which uses
                the second argument to specify an index.

        Yields:
            bool: true if the graph has no cycles
        """
        todo = set(vertex.iterkeys())
        while todo:
            name = todo.pop()
            chain = set()
            while True:
                parent, _ = vertex[name]
                chain.add(name)
                if parent is None:
                    # We've found a root node.
                    break
                elif parent in chain:
                    return False
                elif parent in todo:
                    # If we haven't already confirmed this node is not
                    # part of a cycle, continue searching this chain
                    name = parent
                else:
                    # If we haven't already confirmed this node is not
                    # part of a cycle, continue searching this chain
                    break
            todo -= chain
        return True

    @staticmethod
    def acyclic_chains(names):
        """Generate the acyclic chains with labels drawn from names

        An acyclic chain is one name for a directed graph where each
        vertex is the source of at most one outward edge, and there are
        no cycles of directed edges. So for vertices (A, B, C), A=>B=>C
        would be an acyclic chain, while A=>B=>C=>A would not.

        Note there are very many acyclic chains: if there are $n$ names,
        there are at least $n!$ acyclic chains, and at most $n^n$. This
        function computes the acyclic chains by sequentially generating
        all directed graphs with maximal outdegree 1, then rejecting
        those graphs with cycles. The running time of this function is
        thus `O(len(names)^len(names))`, also known as really really
        slow.

        Args:
            names (list): list of vertex labels

        Yields:
            dict: maps each name to a tuple (parent, None), where parent
                is the sink of its unique edge. A missing edge is
                denoted as an edge back to the vertex, so `{name: (name,
                None)}` describes a vertex with label `name` and no
                outgoing edge. The choice to represent an edge as a
                tuple with second value None was made for uniformity
                with the vertex representation, which uses the second
                argument to specify an index..
        """
        for combination in  itertools.product(*(
                [(a, b if b != a else None) for b in names]
                for a in names)):
            vertex = {name: (parent, None) for name, parent in combination}
            if FactoredRandomGeometricGraph.is_acyclic(vertex):
                yield vertex

    def __contains__(self, vertex):
        """Check if the graph contains a vertex

        The graph contains a vertex if the keys are all the object
        names, the values are tuples (parent, index), the index is a key
        for the (object, parent) sample set, and the directed graph
        formed with edges between each object and its parent is acyclic.
        """
        for name in self.names:
            if name not in vertex:
                return False
            (parent, index) = vertex[name]
            if (name, parent) not in self.samples:
                return False
            if index not in self.samples[(name, parent)]:
                return False
        if not self.is_acyclic(vertex):
            return False
        return True

    def __len__(self):
        """Check if the graph contains a vertex

        The graph contains a vertex if the keys are all the object
        names, the values are tuples (parent, index), the index is a key
        for the (object, parent) sample set, and the directed graph
        formed with edges between each object and its parent is acyclic.
        """
        total = 0
        for vertex in self.acyclic_chains(self.names):
            prod = 1
            for name, (parent, _) in vertex.iteritems():
                prod *= len(self.samples[(name, parent)])
            total += prod
        return total

    def __iter__(self):
        """Iterate over all vertices in the graph

        Vertices are generated by sequentially generating acyclic
        chains, then generating all combinations of vertex indices for
        those acyclic chains. Note that the size of the graph is
        superexponential in the number of objects, so it can take a very
        long time to iterate over all vertices, even for small,
        low-resolution graphs.
        """
        for vertex in self.acyclic_chains(self.names):
            for indices in itertools.product(*(
                    iter(self.samples[(name, parent)])
                    for name, (parent, _) in vertex.iteritems())):
                yield {name: (parent, index)
                       for (name, (parent, _)), index
                       in zip(vertex.iteritems(), indices)}

    def get_pose_of(self, vertex, name):
        """Recursively compute the pose of an object at a vertex"""
        parent, index = vertex[name]
        transform = self.samples[(name, parent)][index]
        if parent is None:
            return transform
        else:
            return apply_transform(self.get_pose_of(vertex, parent), transform)

    def __getitem__(self, vertex):
        """Compute the configuration associated with a vertex

        Because vertices represent the pose of each object relative to
        the pose of another object, computing the absolute pose of each
        object requires first computing a topological ordering on the
        chain, then propagating an absolute pose down the ordering. This
        could be done by maintaining the vertex structure in a way that
        makes it easy to identify root poses; however, that would
        increase the size and complexity of the vertex representation.
        Instead, we compute the topological ordering implicity by
        iterating over the objects in an arbitrary order and keeping a
        'to-do list' of objects discovered whose parents had not yet
        been grounded.

        Args:
            vertex (dict)

        Returns:
            dict: maps object names to (x, y, theta) tuples
        """
        todo = {}
        result = {}
        for name in vertex:
            parent, index = vertex[name]
            if parent is None:
                result[name] = self.samples[(name, parent)][index]
            elif parent in result:
                transform = self.samples[(name, parent)][index]
                result[name] = apply_transform(result[parent], transform)
            else:
                todo.setdefault(parent, set()).add(name)

            if name in todo:
                for child in todo[name]:
                    expected_name, index = vertex[child]
                    assert expected_name == name
                    transform = self.samples[(child, name)][index]
                    result[child] = apply_transform(result[name], transform)
        return result

    def __getstate__(self):
        state = self.__dict__.copy()
        del state['geometry']
        del state['dynamics']
        return state

    def __setstate__(self, state):
        self.__dict__.update(state)
        self.geometry = None
        self.dynamics = None

    def nearest(self, configuration):
        """Return the closest collision-free vertex to configuration

        Distance is the sum of the Euclidean distances between the pose
        of each object in the configuration and in the vertex. This is
        computed in a reasonably efficient way using dynamic programming.

        Args:
            configuration (dict): maps object names to poses

        Returns:
            dict: a dictionary representing the closest collision-free
                vertex to configuration
        """
        def expand_queue(name, k):
            """Return a sorted list of (distance, index) tuples"""
            return sorted(self.samples[(name, None)].nearest_samples(
                configuration[name], k))

        initial_k = 8
        sample_queues = {}
        for name in configuration:
            sample_queues[name] = expand_queue(name, initial_k)

        distance = lambda v: sum(sample_queues[name][v[name]][0]
                                 for name in configuration)
        start = metis.hashdict.hashdict({name: 0 for name in configuration})
        queue = metis.queue.PriorityQueue([(start, distance(start))])

        searched = 0
        while searched < len(self):
            current = queue.pop()
            current_vertex = metis.hashdict.hashdict({
                name: (None, sample_queues[name][index][1])
                for name, index in current.iteritems()})
            searched += 1
            if self.geometry.configuration_is_free(
                    self[current_vertex], skip_static=True):
                return current_vertex
            else:
                for name in current:
                    child = current + {name: current[name] + 1}
                    k = len(sample_queues[name])
                    if child[name] == k:
                        k = min(len(self.samples[(name, None)]), k * 2)
                        sample_queues[name] = expand_queue(name, k)
                    if child not in queue:
                        queue[child] = distance(child)

    def neighbors(self, vertex):
        """Generate neighbors of vertex

        We can construct the adjacent vertices by considering all
        reachable acyclic chains, then identifying within each chain all
        configurations within a distance of search_radius from the
        current vertex.

        Args:
            vertex (dict): maps object names to tuples (parent, i), with
                i an index in the sample set (name, parent)

        Yields:
            hashdict: neighboring vertices on the factored graph
        """
        # Check if the configuration encoded in the vertex dict is free
        if not self.geometry.configuration_is_free(
                self[vertex], skip_static=True):
            raise StopIteration

        result = metis.hashdict.hashdict(vertex)

        # TODO: Many things in this function should be in the dynamics.
        # Instead they are hard-coded for convenience

        # Identify the parent and the root of the current vertex (the
        # root is the ancestor of 'robot' whose parent is None)
        parent, _ = vertex[self.robot]
        root_parent = root = self.robot
        while root_parent is not None:
            root = root_parent
            root_parent, root_index = vertex[root]

        # Generate neighbors on the current manifold [D]
        for sample in self.samples[(root, None)].neighbors(root_index):
            neighbor = result + {root: (None, sample)}
            if self.geometry.configuration_is_free(
                    self[neighbor], skip_static=True):
                yield neighbor

        # Generate neighbors on adjacent manifolds.  The adjacent
        # acyclic chains are hard-coded as the chains for which only the
        # robot changes its parent relative to the chain described in
        # vertex.

        # We always move through free space to change parents, so use
        # the search radius for free space
        radius = self.samples[(self.robot, None)].search_radius
        if self.robot != root:
            # The robot is not a root pose: it is currently in contact
            # with something. Look for poses not in contact with
            # anything
            pose = self.get_pose_of(vertex, self.robot)
            manifold = self.samples[(self.robot, None)]
            for sample in manifold.nearby_samples(pose, radius):
                neighbor = result + {self.robot: (None, sample)}
                if self.geometry.configuration_is_free(
                        self[neighbor], skip_static=True):
                    yield neighbor

        for new_parent in self.names:
            # Consider changing parent to non-root poses
            if new_parent == self.robot or new_parent == parent:
                continue
            else:
                # Now we find nearby poses in this cycle. Translate the
                # current pose to a pose relative to the new parent
                pose = self.get_pose_of(vertex, self.robot)
                parent_pose = self.get_pose_of(vertex, new_parent)
                relative_pose = apply_inverse_transform(pose, parent_pose)

                # Search for neighbors near the absolute pose
                manifold = self.samples[(self.robot, new_parent)]
                for sample in manifold.nearby_samples(relative_pose, radius):
                    neighbor = result + {self.robot: (new_parent, sample)}
                    if self.geometry.configuration_is_free(
                            self[neighbor], skip_static=True):
                        yield neighbor

    def cost(self, parent, child):
        parent_configuration = self[parent]
        child_configuration = {self.robot: self.get_pose_of(child, self.robot)}
        holding = child[self.robot][0]
        if holding is not None:
            child_configuration[holding] = self.get_pose_of(child, holding)
        was_holding = parent[self.robot][0]
        if was_holding is not None:
            child_configuration[was_holding] = self.get_pose_of(child, was_holding)
        # child_configuration = {
        #     name: self.get_pose_of(child, name)
        #     for name in child if parent[name] != child[name]}

        # TODO: this does not check for collision between objects. I
        # *think* what we want is for any object that doesn't move to be
        # treated as part of the background, and any object that does
        # move to be ignored for collision detection. Not sure how to
        # accomplish that.
        if self.geometry.path_is_free(parent_configuration,
                                      child_configuration,
                                      skip_configuration=True):
            return self.dynamics.cost(parent_configuration, child_configuration)
        else:
            return float('inf')


"""
Simplex optimization.
"""
from __future__ import absolute_import
from __future__ import division

import copy
import collections
import logging
import logging.config
import numpy as np
import re
import sqlite3
import textwrap

import calculate
import compare
import constants as co
import datatypes
import opt
import parameters

logger = logging.getLogger(__name__)

class Simplex(opt.Optimizer):
    """
    Optimizes force field parameters using an in-house version of the simplex
    method. See `Optimizer` for repeated documentation.

    Attributes
    ----------
    _max_cycles_wo_change : int
                            End the simplex optimization early if there have
                            been this many consecutive simplex steps without
                            improvement in the objective function.
    do_massive_contraction : bool
                             If True, allows massive contractions to be
                             performed, contracting all parameters at once.
    do_weighted_reflection : bool
                             If True, weights parameter sets based on their
                             objective function score when determining the
                             reflection point.
    max_cycles : int
                 Maximum number of simplex cycles.

    max_params : int
                 Maximum number of parameters used in a single simplex cycle.
    """
    def __init__(self,
                 direc=None,
                 ff=None,
                 ff_lines=None,
                 args_ff=None,
                 args_ref=None):
        super(Simplex, self).__init__(
            direc, ff, ff_lines, args_ff, args_ref)
        self._max_cycles_wo_change = None
        self.do_massive_contraction = True
        self.do_weighted_reflection = True
        self.max_cycles = 100
        self.max_params = 3
    @property
    def best_ff(self):
        # Typically, self.new_ffs would include the original FF, self.ff,
        # but this can be changed by massive contractions.
        if self.new_ffs:
            self.new_ffs = sorted(self.new_ffs, key=lambda x: x.score)
            # I think this is necessary after massive contraction.
            # Massive contraction can potentially make eveything worse.
            # No, it can't!!! The best FF is always retained! /Per-Ola
            # Yep, he's right. /Eric
            if self.new_ffs[0].score < self.ff.score:
                best_ff = self.new_ffs[0]
                best_ff = restore_simp_ff(best_ff, self.ff)
                return best_ff
            else:
                return self.ff
        else:
            return self.ff
    @opt.catch_run_errors
    def run(self, r_data=None):
        """
        Once all attributes are setup as you so desire, run this method to
        optimize the parameters.

        Returns
        -------
        `datatypes.FF` (or subclass)
            Contains the best parameters.
        """
        if r_data is None:
            r_data = opt.return_ref_data(self.args_ref)

        if self.ff.score is None:
            logger.log(20, '~~ CALCULATING INITIAL FF SCORE ~~'.rjust(79, '~'))
            self.ff.export_ff()
            # Could store data on self.ff.data if we wanted. Not necessary for
            # simplex. If simplex yielded no improvements, it would return this
            # FF, and then we might want the data such taht we don't have to
            # recalculate it in gradient. Let's hope simplex generally yields
            # improvements.
            data = calculate.main(self.args_ff)
            #deprecated
            #self.ff.score = compare.compare_data(r_data, data)
            r_dict = compare.data_by_type(r_data)
            c_dict = compare.data_by_type(data)
            r_dict, c_dict = compare.trim_data(r_dict,c_dict)
            self.ff.score = compare.compare_data(r_dict, c_dict)
        else:
            logger.log(20, '  -- Reused existing score and data for initial FF.')

        logger.log(20, '~~ SIMPLEX OPTIMIZATION ~~'.rjust(79, '~'))
        logger.log(20, 'INIT FF SCORE: {}'.format(self.ff.score))
        opt.pretty_ff_results(self.ff, level=20)

        # Here's what we do if there are too many parameters.
        if self.max_params and len(self.ff.params) > self.max_params:
            logger.log(20, '  -- More parameters than the maximum allowed.')
            logger.log(5, 'CURRENT PARAMS: {}'.format(len(self.ff.params)))
            logger.log(5, 'MAX PARAMS: {}'.format(self.max_params))
            # Here we select the parameters that have the lowest 2nd
            # derivatives.

            # Could fail when simplex finds improvements but restores other
            # parameters.
            # if self.ff.params[0].d1:

            if None in [x.d1 for x in self.ff.params]:
                logger.log(15, '  -- Calculating new parameter derivatives.')
                # Do central differentiation so we can calculate derivatives.
                # Another option would be to write code to determine
                # derivatives only from forward differentiation.
                ffs = opt.differentiate_ff(self.ff, central=True)
                # We have to score to get the derivatives.
                for ff in ffs:
                    ff.export_ff(path=self.ff.path, lines=self.ff_lines)
                    logger.log(20, '  -- Calculating {}.'.format(ff))
                    data = calculate.main(self.args_ff)
                    #deprecated
                    #ff.score = compare.compare_data(r_data, data)
                    r_dict = compare.data_by_type(r_data)
                    c_dict = compare.data_by_type(data)
                    r_dict, c_dict = compare.trim_data(r_dict,c_dict)
                    ff.score = compare.compare_data(r_dict, c_dict)
                    opt.pretty_ff_results(ff)
                # Add the derivatives to your original FF.
                opt.param_derivs(self.ff, ffs)
                # Only keep the forward differentiated FFs.
                ffs = opt.extract_forward(ffs)
                logger.log(5, '  -- Keeping {} forward differentiated '
                           'FFs.'.format(len(ffs)))
            else:
                logger.log(15, '  -- Reusing existing parameter derivatives.')
                # Differentiate all parameters forward. Yes, I know this is
                # counter-intuitive because we are going to only use subset of
                # the forward differentiated FFs. However, this is very
                # computationally inexpensive because we're not scoring them
                # now. We will remove the forward differentiated FFs we don't
                # want before scoring.
                ffs = opt.differentiate_ff(self.ff, central=False)

            # This sorts the parameters based upon their 2nd derivative.
            # It keeps the ones with lowest 2nd derivatives.

            # SCHEDULED FOR CHANGES. NOT A GOOD SORTING CRITERION.
            params = select_simp_params_on_derivs(
                self.ff.params, max_params=self.max_params)
            # From the entire list of forward differentiated FFs, pick
            # out the ones that have the lowest 2nd derivatives.
            self.new_ffs = opt.extract_ff_by_params(ffs, params)
            logger.log(1, '>>> len(self.new_ffs): {}'.format(len(self.new_ffs)))

            # Reduce number of parameters.
            # Will need an option that's not MM3* specific in the future.
            ff_rows = [x.mm3_row for x in params]
            ff_cols = [x.mm3_col for x in params]
            for ff in self.new_ffs:
                new_params = []
                for param in ff.params:
                    if param.mm3_row in ff_rows and param.mm3_col in ff_cols:
                        new_params.append(param)
                ff.params = new_params
            # Make a copy of your original FF that has less parameters.
            ff_copy = copy.deepcopy(self.ff)
            new_params = []
            for param in ff.params:
                if param.mm3_row in ff_rows and param.mm3_col in ff_cols:
                    new_params.append(param)
            ff_copy.params = new_params
        else:
            # In this case it's simple. Just forward differentiate each
            # parameter.
            self.new_ffs = opt.differentiate_ff(self.ff, central=False)
            logger.log(1, '>>> len(self.new_ffs): {}'.format(len(self.new_ffs)))
            # Still make that FF copy.
            ff_copy = copy.deepcopy(self.ff)
        # Double check and make sure they're all scored.
        for ff in self.new_ffs:
            if ff.score is None:
                ff.export_ff(path=self.ff.path, lines=self.ff_lines)
                logger.log(20, '  -- Calculating {}.'.format(ff))
                data = calculate.main(self.args_ff)
                #deprecated
                #ff.score = compare.compare_data(r_data, data)
                r_dict = compare.data_by_type(r_data)
                c_dict = compare.data_by_type(data)
                r_dict, c_dict = compare.trim_data(r_dict,c_dict)
                ff.score = compare.compare_data(r_dict, c_dict)
                opt.pretty_ff_results(ff)

        # Add your copy of the orignal to FF to the forward differentiated FFs.
        self.new_ffs = sorted(self.new_ffs + [ff_copy], key=lambda x: x.score)
        # Allow 3 cycles w/o change for each parameter present. Remember that
        # the initial FF was added here, hence the minus one.
        self._max_cycles_wo_change = 3 * (len(self.new_ffs) - 1)
        wrapper = textwrap.TextWrapper(width=79)
        # Shows all FFs parameters.
        opt.pretty_ff_params(self.new_ffs)

        # Start the simplex cycles.
        current_cycle = 0
        cycles_wo_change = 0
        while current_cycle < self.max_cycles \
                and cycles_wo_change < self._max_cycles_wo_change:
            current_cycle += 1

            # Save the last best in case some accidental sort goes on.
            # Plus it makes reading the code a litle easier.
            last_best_ff = copy.deepcopy(self.new_ffs[0])
            logger.log(20, '~~ START SIMPLEX CYCLE {} ~~'.format(
                    current_cycle).rjust(79, '~'))
            logger.log(20, 'ORDERED FF SCORES:')
            logger.log(20, wrapper.fill('{}'.format(
                    ' '.join('{:15.4f}'.format(x.score) for x in self.new_ffs))))

            inv_ff = self.ff.__class__()
            if self.do_weighted_reflection:
                inv_ff.method = 'WEIGHTED INVERSION'
            else:
                inv_ff.method = 'INVERSION'
            inv_ff.params = copy.deepcopy(last_best_ff.params)
            ref_ff = self.ff.__class__()
            ref_ff.method = 'REFLECTION'
            ref_ff.params = copy.deepcopy(last_best_ff.params)
            # Need score difference sum for weighted inversion.
            # Calculate this value before going into loop.
            if self.do_weighted_reflection:
                # If zero, should break.
                score_diff_sum = sum([x.score - self.new_ffs[-1].score
                                      for x in self.new_ffs[:-1]])
                if score_diff_sum == 0.:
                    logger.warning(
                        'No difference between force field scores. '
                        'Exiting simplex.')
                    # We want to raise opt.OptError such that
                    # opt.catch_run_errors will write the best FF obtained thus
                    # far.
                    raise opt.OptError(
                        'No difference between force field scores. '
                        'Exiting simplex.')
            for i in range(0, len(last_best_ff.params)):
                if self.do_weighted_reflection:
                    inv_val = (
                        sum([x.params[i].value *
                             (x.score - self.new_ffs[-1].score)
                             for x in self.new_ffs[:-1]])
                        / score_diff_sum)
                else:
                    inv_val = (
                        sum([x.params[i].value for x in self.new_ffs[:-1]])
                        /
                        len(self.new_ffs[:-1]))
                inv_ff.params[i].value = inv_val
                ref_ff.params[i].value = (
                    2 * inv_val - self.new_ffs[-1].params[i].value)
            # The inversion point does not need to be scored.
            # Calculate score for reflected parameters.
            ref_ff.export_ff(path=self.ff.path, lines=self.ff.lines)
            data = calculate.main(self.args_ff)
            #deprecated
            #ref_ff.score = compare.compare_data(r_data, data)
            r_dict = compare.data_by_type(r_data)
            c_dict = compare.data_by_type(data)
            r_dict, c_dict = compare.trim_data(r_dict,c_dict)
            ref_ff.score = compare.compare_data(r_dict, c_dict)
            opt.pretty_ff_results(ref_ff)
            if ref_ff.score < last_best_ff.score:
                logger.log(20, '~~ ATTEMPTING EXPANSION ~~'.rjust(79, '~'))
                exp_ff = self.ff.__class__()
                exp_ff.method = 'EXPANSION'
                exp_ff.params = copy.deepcopy(last_best_ff.params)
                for i in range(0, len(last_best_ff.params)):
                    exp_ff.params[i].value = (
                        3 * inv_ff.params[i].value -
                        2 * self.new_ffs[-1].params[i].value)
                exp_ff.export_ff(path=self.ff.path, lines=self.ff.lines)
                data = calculate.main(self.args_ff)
                #deprecated
                #exp_ff.score = compare.compare_data(r_data, data)
                r_dict = compare.data_by_type(r_data)
                c_dict = compare.data_by_type(data)
                r_dict, c_dict = compare.trim_data(r_dict,c_dict)
                exp_ff.score = compare.compare_data(r_dict, c_dict)
                opt.pretty_ff_results(exp_ff)
                if exp_ff.score < ref_ff.score:
                    self.new_ffs[-1] = exp_ff
                    logger.log(
                        20, '  -- Expansion succeeded. Keeping expanded '
                        'parameters.')
                else:
                    self.new_ffs[-1] = ref_ff
                    logger.log(
                        20, '  -- Expansion failed. Keeping reflected parameters.')
            elif ref_ff.score < self.new_ffs[-2].score:
                logger.log(20, '  -- Keeping reflected parameters.')
                self.new_ffs[-1] = ref_ff
            else:
                logger.log(20, '~~ ATTEMPTING CONTRACTION ~~'.rjust(79, '~'))
                con_ff = self.ff.__class__()
                con_ff.method = 'CONTRACTION'
                con_ff.params = copy.deepcopy(last_best_ff.params)
                for i in range(0, len(last_best_ff.params)):
                    if ref_ff.score > self.new_ffs[-1].score:
                        con_val = (
                            (inv_ff.params[i].value +
                             self.new_ffs[-1].params[i].value) / 2)
                    else:
                        con_val = (
                            (3 * inv_ff.params[i].value -
                             self.new_ffs[-1].params[i].value) / 2)
                    con_ff.params[i].value = con_val
                self.ff.export_ff(params=con_ff.params)
                data = calculate.main(self.args_ff)
                #deprecated
                #con_ff.score = compare.compare_data(r_data, data)
                r_dict = compare.data_by_type(r_data)
                c_dict = compare.data_by_type(data)
                r_dict, c_dict = compare.trim_data(r_dict,c_dict)
                con_ff.score = compare.compare_data(r_dict, c_dict)
                opt.pretty_ff_results(con_ff)
                # This change was made to reflect the 1998 Q2MM publication.
                # if con_ff.score < self.new_ffs[-1].score:
                if con_ff.score < self.new_ffs[-2].score:
                    logger.log(20, '  -- Contraction succeeded.')
                    self.new_ffs[-1] = con_ff
                elif self.do_massive_contraction:
                    logger.log(
                        20, '~~ DOING MASSIVE CONTRACTION ~~'.rjust(79, '~'))
                    for ff_num, ff in enumerate(self.new_ffs[1:]):
                        for i in range(0, len(last_best_ff.params)):
                            ff.params[i].value = (
                                (ff.params[i].value +
                                 self.new_ffs[0].params[i].value) / 2)
                        self.ff.export_ff(params=ff.params)
                        data = calculate.main(self.args_ff)
                        #deprecated
                        #ff.score = compare.compare_data(r_data, data)
                        r_dict = compare.data_by_type(r_data)
                        c_dict = compare.data_by_type(data)
                        r_dict, c_dict = compare.trim_data(r_dict,c_dict)
                        ff.score = compare.compare_data(r_dict, c_dict)
                        ff.method += ' MC'
                        opt.pretty_ff_results(ff)
                else:
                    logger.log(
                        20, '  -- Contraction failed. Keeping parmaeters '
                        'anyway.')
                    self.new_ffs[-1] = con_ff
            self.new_ffs = sorted(self.new_ffs, key=lambda x: x.score)
            # Keep track of the number of cycles without change. If there's
            # improvement, reset the counter.
            if self.new_ffs[0].score < last_best_ff.score:
                cycles_wo_change = 0
            else:
                cycles_wo_change += 1
                logger.log(20, '  -- {} cycles without improvement out of {} '
                           'allowed.'.format(
                        cycles_wo_change, self._max_cycles_wo_change))
            logger.log(20, 'BEST:')
            opt.pretty_ff_results(self.new_ffs[0], level=20)
            logger.log(20, '~~ END SIMPLEX CYCLE {} ~~'.format(
                    current_cycle).rjust(79, '~'))

        # This sort is likely unnecessary because it should be done at the end
        # of the last loop cycle, but I put it here just in case.
        self.new_ffs = sorted(self.new_ffs, key=lambda x: x.score)
        best_ff = self.new_ffs[0]
        if best_ff.score < self.ff.score:
            logger.log(20, '~~ SIMPLEX FINISHED WITH IMPROVEMENTS ~~'.rjust(
                    79, '~'))
            best_ff = restore_simp_ff(best_ff, self.ff)
        else:
            logger.log(20, '~~ SIMPLEX FINISHED WITHOUT IMPROVEMENTS ~~'.rjust(
                    79, '~'))
            # This restores the inital parameters, so no need to use
            # restore_simp_ff here.
            best_ff = self.ff
        opt.pretty_ff_results(self.ff, level=20)
        opt.pretty_ff_results(best_ff, level=20)
        logger.log(20, '  -- Writing best force field from simplex.')
        best_ff.export_ff(best_ff.path)
        return best_ff

def calc_simp_var(params):
    """
    Simplex variable is calculated: (2nd der.) / (1st der.)**2
    """
    logger.log(1, '>>> params: {}'.format(params))
    logger.log(1, '>>> 1st ders.: {}'.format([x.d1 for x in params]))
    logger.log(1, '>>> 2nd ders.: {}'.format([x.d2 for x in params]))
    for param in params:
        param.simp_var = param.d2 / param.d1**2.

# Sorting based upon the 2nd derivative isn't such a good criterion. This should
# be updated soon.
def select_simp_params_on_derivs(params, max_params=10):
    """
    Sorts parameter sets from lowest to highest second
    derivatives of their score in the objective function.

    Parameters
    ----------
    params : list of subclasses of `datatypes.Param`
    """
    calc_simp_var(params)
    keep = sorted(params, key=lambda x: x.simp_var)
    logger.log(1, '>>> x.simp_var: {}'.format([x.simp_var for x in keep]))

    # Eliminate all where simp_var is greater than 1. This means that the
    # correct value is bracketed by the differentiation, so gradient
    # optimization should work.
    # keep = [x for x in keep if x.simp_var < 1.]

    # Old sorting method.
    # keep = sorted(params, key=lambda x: x.d2)

    keep = keep[:max_params]
    logger.log(1, '>>> x.simp_var: {}'.format([x.simp_var for x in keep]))
    logger.log(20, 'KEEPING PARAMS FOR SIMPLEX:\n{}'.format(
            ' '.join([str(x) for x in keep])))
    return keep

def restore_simp_ff(new_ff, old_ff):
    """
    The old FF has properties that we need to copy to the new FF. We also need
    to grab all the extra parameters included in old FF and add them to the new
    FF.
    """
    old_ff.copy_attributes(new_ff)
    if len(old_ff.params) > len(new_ff.params):
        logger.log(15, '  -- Restoring {} parameters to new FF.'.format(
                len(old_ff.params) - len(new_ff.params)))

        logger.log(1, '>>> old_ff.params:')
        logger.log(1, old_ff.params)
        logger.log(1, [x.d1 for x in old_ff.params])
        logger.log(1, [x.d2 for x in old_ff.params])
        opt.pretty_derivs(old_ff.params, level=1)
        logger.log(1, '>>> new_ff.params:')
        logger.log(1, new_ff.params)
        logger.log(1, [x.d1 for x in new_ff.params])
        logger.log(1, [x.d2 for x in new_ff.params])
        opt.pretty_derivs(new_ff.params, level=1)

        # Backup new parameters.
        new_params = copy.deepcopy(new_ff.params)
        # Copy over all old parameters.
        new_ff.params = copy.deepcopy(old_ff.params)
        # Replace the old with the new.
        for i, param_o in enumerate(old_ff.params):
            for param_n in new_params:
                # Should replace this with a general index scheme.
                if param_o.mm3_row == param_n.mm3_row and \
                        param_o.mm3_col == param_n.mm3_col:
                    new_ff.params[i] = copy.deepcopy(param_n)
    return new_ff