"""Radar module, containing the :class:`Radar` class."""
import math
import numpy as np
import numpy.ma as ma
from scipy import constants
from cloudnetpy.categorize import ProfileDataSource
from cloudnetpy import utils

class Radar(ProfileDataSource):
    """Radar class, child of ProfileDataSource.

    Args:
        radar_file (str): File name of the calibrated radar netCDF file.

    Attributes:
        radar_frequency (float): Radar frequency (GHz).
        wl_band (int): Int corresponding to frequency 0 = 35.5 GHz, 1 = 94 GHz.
        folding_velocity (float): Radar's folding velocity (m/s).
        location (str): Location of the radar, copied from the global attribute
            `location` of the *radar_file*.
        sequence_indices (list): Indices denoting the different altitude
            regimes of the radar.
        type (str): Type of the radar, e.g. 'RPG-FMCW', 'METEK MIRA-36'

    """
    def __init__(self, radar_file):
        super().__init__(radar_file, radar=True)
        self.radar_frequency = float(self.getvar('radar_frequency', 'frequency'))
        self.wl_band = utils.get_wl_band(self.radar_frequency)
        self.folding_velocity = self._get_folding_velocity()
        self.sequence_indices = self._get_sequence_indices()
        self.location = getattr(self.dataset, 'location', '')
        self.type = getattr(self.dataset, 'source', '')
        self._netcdf_to_cloudnet(('v', 'width', 'ldr'))
        self._unknown_to_cloudnet(('Zh', 'Zv', 'Ze'), 'Z', units='dBZ')
        self._init_sigma_v()

    def rebin_to_grid(self, time_new):
        """Rebins radar data in time using mean.

        Args:
            time_new (ndarray): Target time array as fraction hour. Updates
                *time* attribute.

        """
        for key in self.data:
            if key in ('Z', 'ldr'):
                self.data[key].db2lin()
                self.data[key].rebin_data(self.time, time_new)
                self.data[key].lin2db()
            elif key == 'v':
                # This has some problems with RPG data when folding is present.
                self.data[key].rebin_velocity(self.time, time_new,
                                              self.folding_velocity,
                                              self.sequence_indices)
            elif key == 'v_sigma':
                self.data[key].calc_linear_std(self.time, time_new)
            else:
                self.data[key].rebin_data(self.time, time_new)
        self.time = time_new

    def remove_incomplete_pixels(self):
        """Removes pixels where some of the (required) variables are existing."""
        good_ind = (~ma.getmaskarray(self.data['Z'][:])
                    & ~ma.getmaskarray(self.data['width'][:])
                    & ~ma.getmaskarray(self.data['v'][:]))
        for key in ('Z', 'v', 'width', 'ldr', 'v_sigma'):
            self.data[key].mask_indices(~good_ind)

    def filter_speckle_noise(self):
        """Removes speckle noise from radar data."""
        for key in ('Z', 'v', 'width', 'ldr', 'v_sigma'):
            if key in self.data.keys():
                self.data[key].filter_isolated_pixels()

    def correct_atten(self, attenuations):
        """Corrects radar echo for liquid and gas attenuation.

        Args:
            attenuations (dict): 2-D attenuations due to atmospheric gases
                and liquid: `radar_gas_atten`, `radar_liquid_atten`.

        """
        z_corrected = self.data['Z'][:] + attenuations['radar_gas_atten']
        ind = ma.where(attenuations['radar_liquid_atten'])
        z_corrected[ind] += attenuations['radar_liquid_atten'][ind]
        self.append_data(z_corrected, 'Z')

    def calc_errors(self, attenuations, classification):
        """Calculates uncertainties of radar echo.

        Calculates and adds `Z_error`, `Z_sensitivity` and `Z_bias`
        :class:`CloudnetArray` instances to `data` attribute.

        Args:
            attenuations (dict): 2-D attenuations due to atmospheric gases.
            classification (ClassificationResult): The
                :class:`ClassificationResult` instance.

        """
        def _calc_sensitivity():
            """Returns sensitivity of radar as function of altitude."""
            mean_gas_atten = ma.mean(attenuations['radar_gas_atten'], axis=0)
            z_sensitivity = z_power_min + log_range + mean_gas_atten
            zc = ma.median(ma.array(z, mask=~classification.is_clutter), axis=0)
            z_sensitivity[~zc.mask] = zc[~zc.mask]
            return z_sensitivity

        def _calc_error():
            z_precision = 4.343 * (1 / np.sqrt(_number_of_pulses())
                                   + utils.db2lin(z_power_min - z_power) / 3)
            gas_error = attenuations['radar_gas_atten'] * 0.1
            liq_error = attenuations['liquid_atten_err'].filled(0)
            z_error = utils.l2norm(gas_error, liq_error, z_precision)
            z_error[attenuations['liquid_uncorrected']] = ma.masked
            return z_error

        def _number_of_pulses():
            """Returns number of independent pulses."""
            seconds_in_hour = 3600
            dwell_time = utils.mdiff(self.time) * seconds_in_hour
            return (dwell_time * self.radar_frequency * 1e9 * 4
                    * np.sqrt(math.pi) * self.data['width'][:] / 3e8)

        def _calc_z_power_min():
            if ma.all(z_power.mask):
                return 0
            return np.percentile(z_power.compressed(), 0.1)

        z = self.data['Z'][:]
        radar_range = self.km2m(self.dataset.variables['range'])
        log_range = utils.lin2db(radar_range, scale=20)
        z_power = z - log_range
        z_power_min = _calc_z_power_min()
        self.append_data(_calc_error(), 'Z_error')
        self.append_data(_calc_sensitivity(), 'Z_sensitivity')
        self.append_data(1, 'Z_bias')

    def add_meta(self):
        """Copies misc. metadata from the input file."""
        for key in ('latitude', 'longitude', 'altitude'):
            self.append_data(self.getvar(key), key)
        for key in ('time', 'height', 'radar_frequency'):
            self.append_data(getattr(self, key), key)

    def _init_sigma_v(self):
        """Initializes std of the velocity field. The std will be calculated
        later when re-binning the data."""
        self.append_data(self.getvar('v'), 'v_sigma')

    def _get_sequence_indices(self):
        """Mira has only one sequence and one folding velocity. RPG has
        several sequences with different folding velocities."""
        all_indices = np.arange(len(self.height))
        if not utils.isscalar(self.folding_velocity):
            starting_indices = self.getvar('chirp_start_indices')
            return np.split(all_indices, starting_indices[1:])
        return [all_indices]

    def _get_folding_velocity(self):
        for key in ('nyquist_velocity', 'NyquistVelocity'):
            if key in self.dataset.variables:
                return self.getvar(key)
        if 'prf' in self.dataset.variables:
            prf = self.getvar('prf')
            return _prf_to_folding_velocity(prf, self.radar_frequency)
        raise RuntimeError('Unable to determine folding velocity')

def _prf_to_folding_velocity(prf, radar_frequency):
    ghz_to_hz = 1e9
    return float(prf * constants.c / (4*radar_frequency*ghz_to_hz))


import warnings
import numpy as np
import copy

from astropy.time import Time
from astropy.utils.iers.iers import IERSRangeError

__all__ = ["ModifiedJulianDate", "MJDWarning", "UTCtoUT1Warning"]

# Filter out ERFA's complaints that we are simulating dates which
# are in the future
warnings.filterwarnings("ignore",
                        message='.*taiutc.*dubious.year.*')

class MJDWarning(Warning):
    """
    A sub-class of Warning.  All of the warnings raised by ModifiedJulianDate
    will be of this class (or its sub-classes), so that users can filter them
    out by creating a simple filter targeted at category=MJDWarning.
    """
    pass

class UTCtoUT1Warning(MJDWarning):
    """
    A sub-class of MJDWarning meant for use when astropy.Time cannot interpolate
    UT1-UTC as a function of UTC because UTC is out of bounds of the data.
    This class exists so that users can filter these warnings out by creating
    a simple filter targeted at category=UTCtoUT1Warning.
    """
    pass

class ModifiedJulianDate(object):

    @classmethod
    def _get_ut1_from_utc(cls, UTC):
        """
        Take a numpy array of UTC values and return a numpy array of UT1 and dut1 values
        """

        time_list = Time(UTC, scale='utc', format='mjd')

        try:
            dut1_out = time_list.delta_ut1_utc
            ut1_out = time_list.ut1.mjd
        except IERSRangeError:
            ut1_out = np.copy(UTC)
            dut1_out = np.zeros(len(UTC))
            warnings.warn("ModifiedJulianData.get_list() was given date values that are outside "
                          "astropy's range of interpolation for converting from UTC to UT1. "
                          "We will treat UT1=UTC for those dates, lacking a better alternative.",
                          category=UTCtoUT1Warning)
            from astropy.utils.iers import TIME_BEFORE_IERS_RANGE, TIME_BEYOND_IERS_RANGE
            dut1_test, status = time_list.get_delta_ut1_utc(return_status=True)
            good_dexes = np.where(np.logical_and(status != TIME_BEFORE_IERS_RANGE,
                                                 status != TIME_BEYOND_IERS_RANGE))

            if len(good_dexes[0]) > 0:
                time_good = Time(UTC[good_dexes], scale='utc', format='mjd')
                dut1_good = time_good.delta_ut1_utc
                ut1_good = time_good.ut1.mjd

                ut1_out[good_dexes] = ut1_good
                dut1_out[good_dexes] = dut1_good

        return ut1_out, dut1_out

    @classmethod
    def get_list(cls, TAI=None, UTC=None):
        """
        Instantiate a list of ModifiedJulianDates from a numpy array of either TAI
        or UTC values.

        @param[in] TAI (optional) a numpy array of MJD' in TAI

        @param[in] UTC (optional) a numpy array of MJDs in UTC

        @param[out] a list of ModifiedJulianDate instantiations with all of their
        properties already set (so the code does not waste time converting from TAI
        to TT, TDB, etc. when those time scales are called for).
        """

        if TAI is None and UTC is None:
            return None

        if TAI is not None and UTC is not None:
            raise RuntimeError("You should not specify both TAI and UTC in ModifiedJulianDate.get_list()")

        if TAI is not None:
            time_list = Time(TAI, scale='tai', format='mjd')
            tai_list = TAI
            utc_list = time_list.utc.mjd
        elif UTC is not None:
            time_list = Time(UTC, scale='utc', format='mjd')
            utc_list = UTC
            tai_list = time_list.tai.mjd

        tt_list = time_list.tt.mjd
        tdb_list = time_list.tdb.mjd

        ut1_list, dut1_list = cls._get_ut1_from_utc(utc_list)

        values = np.array([tai_list, utc_list, tt_list, tdb_list,
                           ut1_list, dut1_list]).transpose()

        output = []
        for vv in values:
            mjd = ModifiedJulianDate(TAI=40000.0)
            mjd._force_values(vv)
            output.append(mjd)

        return output

    def __init__(self, TAI=None, UTC=None):
        """
        Must specify either:

        @param [in] TAI = the International Atomic Time as an MJD

        or

        @param [in] UTC = Universal Coordinate Time as an MJD
        """

        if TAI is None and UTC is None:
            raise RuntimeError("You must specify either TAI or UTC to "
                               "instantiate ModifiedJulianDate")

        if TAI is not None:
            self._time = Time(TAI, scale='tai', format='mjd')
            self._tai = TAI
            self._utc = None
            self._initialized_with = 'TAI'
        else:
            self._time = Time(UTC, scale='utc', format='mjd')
            self._utc = UTC
            self._tai = None
            self._initialized_with = 'UTC'

        self._tt = None
        self._tdb = None
        self._ut1 = None
        self._dut1 = None

    def _force_values(self, values):
        """
        Force the properties of this ModifiedJulianDate to have specific values.

        values is a list of [TAI, UTC, TT, TDB, UT1, UT1-UTC] values.

        This method exists so that, when instantiating lists of ModifiedJulianDates,
        we can use astropy.time.Time's vectorized methods to quickly perform many
        conversions at once.  Users should not try to use this method by hand.
        """
        self._tai = values[0]
        self._utc = values[1]
        self._tt = values[2]
        self._tdb = values[3]
        self._ut1 = values[4]
        self._dut1 = values[5]

    def __eq__(self, other):
        return self._time == other._time

    def __ne__(self, other):
        return not self.__eq__(other)

    def __deepcopy__(self, memo):
        if self._initialized_with == 'TAI':
            new_mjd = ModifiedJulianDate(TAI=self.TAI)
        else:
            new_mjd = ModifiedJulianDate(UTC=self.UTC)

        new_mjd._tai = copy.deepcopy(self._tai, memo)
        new_mjd._utc = copy.deepcopy(self._utc, memo)
        new_mjd._tt = copy.deepcopy(self._tt, memo)
        new_mjd._tdb = copy.deepcopy(self._tdb, memo)
        new_mjd._ut1 = copy.deepcopy(self._ut1, memo)
        new_mjd._dut1 = copy.deepcopy(self._dut1, memo)

        return new_mjd

    def _warn_utc_out_of_bounds(self, method_name):
        """
        Raise a standard warning if UTC is outside of the range that can
        be interpolated on the IERS tables.

        method_name is the name of the method that caused this warning.
        """
        warnings.warn("UTC is outside of IERS table for UT1-UTC.\n"
                      "Returning UT1 = UTC for lack of a better idea\n"
                      "This warning was caused by calling ModifiedJulianDate.%s\n" % method_name,
                      category=UTCtoUT1Warning)

    @property
    def TAI(self):
        """
        International Atomic Time as an MJD
        """
        if self._tai is None:
            self._tai = self._time.tai.mjd

        return self._tai

    @property
    def UTC(self):
        """
        Universal Coordinate Time as an MJD
        """
        if self._utc is None:
            self._utc = self._time.utc.mjd

        return self._utc

    @property
    def UT1(self):
        """
        Universal Time as an MJD
        """
        if self._ut1 is None:
            try:
                self._ut1 = self._time.ut1.mjd
            except IERSRangeError:
                self._warn_utc_out_of_bounds('UT1')
                self._ut1 = self.UTC

        return self._ut1

    @property
    def dut1(self):
        """
        UT1-UTC in seconds
        """

        if self._dut1 is None:
            try:
                self._dut1 = self._time.delta_ut1_utc
            except IERSRangeError:
                self._warn_utc_out_of_bounds('dut1')
                self._dut1 = 0.0

        return self._dut1

    @property
    def TT(self):
        """
        Terrestrial Time (aka Terrestrial Dynamical Time) as an MJD
        """
        if self._tt is None:
            self._tt = self._time.tt.mjd

        return self._tt

    @property
    def TDB(self):
        """
        Barycentric Dynamical Time as an MJD
        """
        if self._tdb is None:
            self._tdb = self._time.tdb.mjd

        return self._tdb


#!/usr/bin/env python3
# encoding: utf-8

import numpy as np

from config import THRESHOLD
from engine.datastore.models.section import IMRaDType
from engine.preprocessing.chapter_classifier.classifier_nn import ClassifierNN
from engine.preprocessing.chapter_classifier.classifier_simple import ClassifierSimple
from engine.utils.exceptions.import_exceptions import ClassificationError

class IMRaDDetection(object):
    def __init__(self):
        self.classifierNN = ClassifierNN()
        self.classifierSimple = ClassifierSimple()

    @staticmethod
    def __print_chapters_and_values(chapter_names, y):
        for i in range(len(chapter_names)):
            tmp = np.round(y[i], 3)
            print("{:30}: {}".format(chapter_names[i], tmp))

    @staticmethod
    def __set_section_in_paper(paper, positions, imrad_type):
        for pos in positions:
            paper.sections[pos].add_to_imrad(imrad_type)

    def proceed(self, paper):
        chapter_names = [name.heading_proceed for name in paper.sections]
        if not len(chapter_names):
            raise ClassificationError("Chapters can't be extracted from pdf-file.")

        imrad = self.proceed_chapters_simple(chapter_names)

        for imrad_type, chapters in imrad.items():
            self.__set_section_in_paper(paper, chapters, imrad_type)

    def proceed_chapters_simple(self, chapters):
        y_simple = self.classifierSimple.predict_chapter(chapters)
        imrad = {}

        intro = np.array([x[IMRaDType.INTRODUCTION.value] for x in y_simple])
        intro_pos = np.where(intro >= THRESHOLD)[0]
        imrad[IMRaDType.INTRODUCTION] = intro_pos

        # if no Background -> Background is in Indroduction
        background = np.array([x[IMRaDType.BACKGROUND.value] for x in y_simple])
        background_pos = np.where(background >= THRESHOLD)[0]
        if not len(background_pos):
            background_pos = intro_pos
        imrad[IMRaDType.BACKGROUND] = background_pos

        abstract = np.array([x[IMRaDType.ABSTRACT.value] for x in y_simple])
        abstract_pos = np.where(abstract >= THRESHOLD)[0]
        imrad[IMRaDType.ABSTRACT] = abstract_pos

        acknowledge = np.array([x[IMRaDType.ACKNOWLEDGE.value] for x in y_simple])
        acknowledge_pos = np.where(acknowledge >= THRESHOLD)[0]
        imrad[IMRaDType.ACKNOWLEDGE] = acknowledge_pos

        discussion = np.array([x[IMRaDType.DISCUSSION.value] for x in y_simple])
        discussion_pos = np.where(discussion >= THRESHOLD)[0]
        imrad[IMRaDType.DISCUSSION] = discussion_pos

        result = np.array([x[IMRaDType.RESULTS.value] for x in y_simple])
        result_pos = np.where(result >= THRESHOLD)[0]
        imrad[IMRaDType.RESULTS] = result_pos

        try:
            disc_res_max = np.amax(np.append(discussion_pos, result_pos))
        except ValueError as e:
            raise ClassificationError("Result and Discussion not set: {}".format(e))

        try:
            intro_min = np.amin(intro_pos)
        except ValueError as e:
            raise ClassificationError("Introduction not set: {}".format(e))

        # Unset chapters between INTRO and DISCUSSION/RESULT are methods
        all_pos = np.concatenate((intro_pos, background_pos, abstract_pos, acknowledge_pos, discussion_pos, result_pos))
        method_pos = [x for x in np.arange(intro_min + 1, disc_res_max) if x not in all_pos]
        imrad[IMRaDType.METHODS] = np.array(method_pos)

        return imrad

    def proceed_chapters(self, chapters):
        # By now uses simple classification as helper
        imrad = {}

        y = self.classifierNN.predict_chapter(chapters)
        y_simple = self.classifierSimple.predict_chapter(chapters)

        intro = np.array([x[IMRaDType.INTRODUCTION.value] for x in y])
        intro_pos = np.where(intro >= THRESHOLD)[0]

        if not len(intro_pos):
            intro = np.array([x[IMRaDType.INTRODUCTION.value] for x in y_simple])
            intro_pos = np.where(intro >= THRESHOLD)[0]

        imrad[IMRaDType.INTRODUCTION] = intro_pos

        # if no Background -> Background is in Indroduction
        background = np.array([x[IMRaDType.BACKGROUND.value] for x in y])
        background_pos = np.where(background >= THRESHOLD)[0]
        if not len(background_pos):
            background_pos = intro_pos

        imrad[IMRaDType.BACKGROUND] = background_pos

        abstract = np.array([x[IMRaDType.ABSTRACT.value] for x in y])
        abstract_pos = np.where(abstract >= THRESHOLD)[0]
        imrad[IMRaDType.ABSTRACT] = abstract_pos

        acknowledge = np.array([x[IMRaDType.ACKNOWLEDGE.value] for x in y])
        acknowledge_pos = np.where(acknowledge >= THRESHOLD)[0]
        imrad[IMRaDType.ACKNOWLEDGE] = acknowledge_pos

        discussion = np.array([x[IMRaDType.DISCUSSION.value] for x in y])
        discussion_pos = np.where(discussion >= THRESHOLD)[0]

        result = np.array([x[IMRaDType.RESULTS.value] for x in y])
        result_pos = np.where(result >= THRESHOLD)[0]

        if not len(list(discussion_pos) + list(result_pos)):
            discussion = np.array([x[IMRaDType.DISCUSSION.value] for x in y_simple])
            discussion_pos = np.where(discussion >= THRESHOLD)[0]

            result = np.array([x[IMRaDType.RESULTS.value] for x in y_simple])
            result_pos = np.where(result >= THRESHOLD)[0]

        imrad[IMRaDType.DISCUSSION] = discussion_pos
        imrad[IMRaDType.RESULTS] = result_pos

        try:
            disc_res_max = np.amax(np.append(discussion_pos, result_pos))
        except ValueError as e:
            raise ClassificationError("Result and Discussion not set: {}".format(e))

        try:
            intro_min = np.amin(intro_pos)
        except ValueError as e:
            raise ClassificationError("Introduction not set: {}".format(e))

        # Unset chapters between INTRO and DISCUSSION/RESULT are methods
        all_pos = np.concatenate((intro_pos, background_pos, abstract_pos, acknowledge_pos, discussion_pos, result_pos))
        method_pos = [x for x in np.arange(intro_min + 1, disc_res_max) if x not in all_pos]
        imrad[IMRaDType.METHODS] = np.array(method_pos)

        return imrad


# --------------------------------------------------------
# Fast R-CNN
# Copyright (c) 2015 Microsoft
# Licensed under The MIT License [see LICENSE for details]
# Written by Ross Girshick
# --------------------------------------------------------

"""Test a Fast R-CNN network on an imdb (image database)."""

from fast_rcnn.config import cfg, get_output_dir
from fast_rcnn.bbox_transform import clip_boxes, bbox_transform_inv
import argparse
from utils.timer import Timer
import numpy as np
import cv2
import caffe
from fast_rcnn.nms_wrapper import nms
import cPickle
from utils.blob import im_list_to_blob,prep_im_and_label_for_blob,label_list_to_blob
import os
import PIL
from datetime import datetime
import pdb
from score import fast_hist
#from __future__ import print_function

def _get_image_blob(im):
    """Converts an image into a network input.

    Arguments:
        im (ndarray): a color image in BGR order

    Returns:
        blob (ndarray): a data blob holding an image pyramid
        im_scale_factors (list): list of image scales (relative to im) used
            in the image pyramid
    """
    im_orig = im.astype(np.float32, copy=True)
    im_orig -= cfg.PIXEL_MEANS

    im_shape = im_orig.shape
    im_size_min = np.min(im_shape[0:2])
    im_size_max = np.max(im_shape[0:2])

    processed_ims = []
    im_scale_factors = []

    for target_size in cfg.TEST.SCALES:
        im_scale = float(target_size) / float(im_size_min)
        #import pdb;pdb.set_trace()
        # Prevent the biggest axis from being more than MAX_SIZE
        if np.round(im_scale * im_size_max) > cfg.TEST.MAX_SIZE:
            im_scale = float(cfg.TEST.MAX_SIZE) / float(im_size_max)
        im = cv2.resize(im_orig, None, None, fx=im_scale, fy=im_scale,
                        interpolation=cv2.INTER_LINEAR)
        im_scale_factors.append(im_scale)
        processed_ims.append(im)

    # Create a blob to hold the input images
    blob = im_list_to_blob(processed_ims)

    return blob, np.array(im_scale_factors)

def _get_image_label_blob(im, label):
    """Converts an image into a network input.

    Arguments:
        im (ndarray): a color image in BGR order
        label (ndarray): the pixel-wise label of im

    Returns:
        blob (ndarray): a data blob holding an image pyramid
        im_scale_factors (list): list of image scales (relative to im) used
            in the image pyramid
        label_blob (ndarray): a label blob holding resized labels
    """
    im_orig = im.astype(np.float32, copy=True)
    label_orig = label.astype(np.uint8, copy=True)

    processed_ims = []
    processed_labels = []
    im_scale_factors = []

    for target_size in cfg.TEST.SCALES:
        im, im_scale, label = prep_im_and_label_for_blob(im, label, cfg.PIXEL_MEANS, target_size, cfg.TEST.MAX_SIZE)
        im_scale_factors.append(im_scale)
        processed_ims.append(im)
        processed_labels.append(label)
        # recover from the original img and label
        im = np.copy(im_orig)
        label = np.copy(label_orig)

    # Create a blob to hold the input images
    blob = im_list_to_blob(processed_ims)
    label_blob = label_list_to_blob(processed_labels)

    return blob, np.array(im_scale_factors), label_blob

def _get_rois_blob(im_rois, im_scale_factors):
    """Converts RoIs into network inputs.

    Arguments:
        im_rois (ndarray): R x 4 matrix of RoIs in original image coordinates
        im_scale_factors (list): scale factors as returned by _get_image_blob

    Returns:
        blob (ndarray): R x 5 matrix of RoIs in the image pyramid
    """
    rois, levels = _project_im_rois(im_rois, im_scale_factors)
    rois_blob = np.hstack((levels, rois))
    return rois_blob.astype(np.float32, copy=False)

def _project_im_rois(im_rois, scales):
    """Project image RoIs into the image pyramid built by _get_image_blob.

    Arguments:
        im_rois (ndarray): R x 4 matrix of RoIs in original image coordinates
        scales (list): scale factors as returned by _get_image_blob

    Returns:
        rois (ndarray): R x 4 matrix of projected RoI coordinates
        levels (list): image pyramid levels used by each projected RoI
    """
    im_rois = im_rois.astype(np.float, copy=False)

    if len(scales) > 1:
        widths = im_rois[:, 2] - im_rois[:, 0] + 1
        heights = im_rois[:, 3] - im_rois[:, 1] + 1

        areas = widths * heights
        scaled_areas = areas[:, np.newaxis] * (scales[np.newaxis, :] ** 2)
        diff_areas = np.abs(scaled_areas - 224 * 224)
        levels = diff_areas.argmin(axis=1)[:, np.newaxis]
    else:
        levels = np.zeros((im_rois.shape[0], 1), dtype=np.int)

    rois = im_rois * scales[levels]

    return rois, levels

def _get_blobs(im, rois, label=None):
    """Convert an image and RoIs within that image into network inputs."""
    blobs = {'data' : None, 'rois' : None}
    if label is not None:
        # add img_labels if has labels
        blobs['img_labels'] = label
    blobs['data'], im_scale_factors = _get_image_blob(im)
    if not cfg.TEST.HAS_RPN:
        blobs['rois'] = _get_rois_blob(rois, im_scale_factors)
    return blobs, im_scale_factors

#def _get_blobs_label(im, label):
#    blobs = {'data' : None, 'img_labels': None}
#    blobs['data'], im_scale_factors, blobs['img_labels'] = _get_image_label_blob(im, label)
#    return blobs, im_scale_factors

def im_seg(net, im, label=None):
    """
        Arguments:
            net (caffe.Net): Fast R-CNN network to use
            im (ndarray): color image to test (in BGR order)
            label(ndarray): pixel-wise label for the input image
        Returns:
            scores (ndarray): pixel-wise prediction in object proposals
            boxes (ndarray): R x (4*K) array of predicted bounding boxes
            boxes_score (ndarray): R x 1 array of scores for the predicted bounding boxes
    """
    blobs, im_scales = _get_blobs(im, None, label)
    if cfg.TEST.HAS_RPN:
        im_blob = blobs['data']
        # blobs['im_info']: H x W x scale_factor(transform the input image to
        # canonical size)
        blobs['im_info'] = np.array(
            [[im_blob.shape[2], im_blob.shape[3], im_scales[0]]],
            dtype=np.float32)
    # reshape network inputs
    net.blobs['data'].reshape(*(blobs['data'].shape))
    if label is not None:
        net.blobs['img_labels'].reshape(*(blobs['img_labels'].shape))
    if cfg.TEST.HAS_RPN:
        net.blobs['im_info'].reshape(*(blobs['im_info'].shape))
    # do forward
    forward_kwargs = {'data': blobs['data'].astype(np.float32, copy=False)}
    if label is not None:
        forward_kwargs['img_labels'] = blobs['img_labels']
    if cfg.TEST.HAS_RPN:
        forward_kwargs['im_info'] = blobs['im_info'].astype(np.float32, copy=False)
    blobs_out = net.forward(**forward_kwargs)

    if cfg.TEST.HAS_RPN:
        assert len(im_scales) == 1, "Only single-image batch implemented"
        rois = net.blobs['rois'].data.copy()
        # DEBUG
        #if blobs['img_labels'] > 0:
        #    import pdb;pdb.set_trace()
        # unscale back to raw image space
        #boxes = rois[:, 1:5] * 16 / im_scales[0] # remove after fix
        boxes = rois[:, 1:5] / im_scales[0] # now the change of rois has been fixed
        #import pdb;pdb.set_trace()
        boxes_score = net.blobs['rois_score'].data
    scores = net.blobs['score'].data # (1, 21, 562, 1000), resized image,

    # pick the label for maximum class
    scores = np.argmax(scores, axis=1).astype(np.uint8)
    # resize this scores to map back to the original size of the image
    # only resize if img had diff size
    if im_scales[0] != 1:
        scores = cv2.resize(scores[0,:,:], None, None, fx=1./im_scales[0], fy=1./im_scales[0],
                    interpolation=cv2.INTER_NEAREST)

    return scores, boxes, boxes_score

def im_detect(net, im, boxes=None):
    """Detect object classes in an image given object proposals.

    Arguments:
        net (caffe.Net): Fast R-CNN network to use
        im (ndarray): color image to test (in BGR order)
        boxes (ndarray): R x 4 array of object proposals or None (for RPN)

    Returns:
        scores (ndarray): R x K array of object class scores (K includes
            background as object category 0)
        boxes (ndarray): R x (4*K) array of predicted bounding boxes
    """
    blobs, im_scales = _get_blobs(im, boxes)

    # When mapping from image ROIs to feature map ROIs, there's some aliasing
    # (some distinct image ROIs get mapped to the same feature ROI).
    # Here, we identify duplicate feature ROIs, so we only compute features
    # on the unique subset.
    if cfg.DEDUP_BOXES > 0 and not cfg.TEST.HAS_RPN:
        v = np.array([1, 1e3, 1e6, 1e9, 1e12])
        hashes = np.round(blobs['rois'] * cfg.DEDUP_BOXES).dot(v)
        _, index, inv_index = np.unique(hashes, return_index=True,
                                        return_inverse=True)
        blobs['rois'] = blobs['rois'][index, :]
        boxes = boxes[index, :]

    if cfg.TEST.HAS_RPN:
        im_blob = blobs['data']
        # blobs['im_info']: H x W x scale_factor(transform the input image to
        # canonical size)
        blobs['im_info'] = np.array(
            [[im_blob.shape[2], im_blob.shape[3], im_scales[0]]],
            dtype=np.float32)

    # reshape network inputs
    net.blobs['data'].reshape(*(blobs['data'].shape))
    if cfg.TEST.HAS_RPN:
        net.blobs['im_info'].reshape(*(blobs['im_info'].shape))
    else:
        net.blobs['rois'].reshape(*(blobs['rois'].shape))

    # do forward
    forward_kwargs = {'data': blobs['data'].astype(np.float32, copy=False)}
    if cfg.TEST.HAS_RPN:
        forward_kwargs['im_info'] = blobs['im_info'].astype(np.float32, copy=False)
    else:
        forward_kwargs['rois'] = blobs['rois'].astype(np.float32, copy=False)
    blobs_out = net.forward(**forward_kwargs)

    if cfg.TEST.HAS_RPN:
        assert len(im_scales) == 1, "Only single-image batch implemented"
        rois = net.blobs['rois'].data.copy()
        # unscale back to raw image space
        boxes = rois[:, 1:5] / im_scales[0]

    if cfg.TEST.SVM:
        # use the raw scores before softmax under the assumption they
        # were trained as linear SVMs
        scores = net.blobs['cls_score'].data
    else:
        # use softmax estimated probabilities
        scores = blobs_out['cls_prob']

    if cfg.TEST.BBOX_REG: # Default
        # Apply bounding-box regression deltas
        box_deltas = blobs_out['bbox_pred']
        pred_boxes = bbox_transform_inv(boxes, box_deltas)
        pred_boxes = clip_boxes(pred_boxes, im.shape)
    else:
        # Simply repeat the boxes, once for each class
        pred_boxes = np.tile(boxes, (1, scores.shape[1]))

    return scores, pred_boxes

def vis_detections(im, class_name, dets, thresh=0.3):
    """Visual debugging of detections."""
    import matplotlib.pyplot as plt
    im = im[:, :, (2, 1, 0)]
    for i in xrange(np.minimum(10, dets.shape[0])):
        bbox = dets[i, :4]
        score = dets[i, -1]
        if score > thresh:
            plt.cla()
            plt.imshow(im)
            plt.gca().add_patch(
                plt.Rectangle((bbox[0], bbox[1]),
                              bbox[2] - bbox[0],
                              bbox[3] - bbox[1], fill=False,
                              edgecolor='g', linewidth=3)
                )
            plt.title('{}  {:.3f}'.format(class_name, score))
            plt.show()

def apply_nms(all_boxes, thresh):
    """Apply non-maximum suppression to all predicted boxes output by the
    test_net method.
    """
    num_classes = len(all_boxes)
    num_images = len(all_boxes[0])
    nms_boxes = [[[] for _ in xrange(num_images)]
                 for _ in xrange(num_classes)]
    for cls_ind in xrange(num_classes):
        for im_ind in xrange(num_images):
            dets = all_boxes[cls_ind][im_ind]
            if dets == []:
                continue
            # CPU NMS is much faster than GPU NMS when the number of boxes
            # is relative small (e.g., < 10k)
            # TODO(rbg): autotune NMS dispatch
            keep = nms(dets, thresh, force_cpu=True)
            if len(keep) == 0:
                continue
            nms_boxes[cls_ind][im_ind] = dets[keep, :].copy()
    return nms_boxes

def test_net(net, imdb, max_per_image=100, thresh=0.05, vis=False):
    """Test a Fast R-CNN network on an image database."""
    num_images = len(imdb.image_index)
    # all detections are collected into:
    #    all_boxes[cls][image] = N x 5 array of detections in
    #    (x1, y1, x2, y2, score)
    all_boxes = [[[] for _ in xrange(num_images)]
                 for _ in xrange(imdb.num_classes)]

    output_dir = get_output_dir(imdb, net)

    # timers
    _t = {'im_detect' : Timer(), 'misc' : Timer()}

    if not cfg.TEST.HAS_RPN:
        roidb = imdb.roidb

    for i in xrange(num_images):
        # filter out any ground truth boxes
        if cfg.TEST.HAS_RPN:
            box_proposals = None
        else:
            # The roidb may contain ground-truth rois (for example, if the roidb
            # comes from the training or val split). We only want to evaluate
            # detection on the *non*-ground-truth rois. We select those the rois
            # that have the gt_classes field set to 0, which means there's no
            # ground truth.
            box_proposals = roidb[i]['boxes'][roidb[i]['gt_classes'] == 0]

        im = cv2.imread(imdb.image_path_at(i))
        _t['im_detect'].tic()
        scores, boxes = im_detect(net, im, box_proposals)
        _t['im_detect'].toc()

        _t['misc'].tic()
        # skip j = 0, because it's the background class
        for j in xrange(1, imdb.num_classes):
            inds = np.where(scores[:, j] > thresh)[0]
            cls_scores = scores[inds, j]
            cls_boxes = boxes[inds, j*4:(j+1)*4]
            cls_dets = np.hstack((cls_boxes, cls_scores[:, np.newaxis])) \
                .astype(np.float32, copy=False)
            keep = nms(cls_dets, cfg.TEST.NMS)
            cls_dets = cls_dets[keep, :]
            if vis:
                vis_detections(im, imdb.classes[j], cls_dets)
            all_boxes[j][i] = cls_dets

        # Limit to max_per_image detections *over all classes*
        if max_per_image > 0:
            image_scores = np.hstack([all_boxes[j][i][:, -1]
                                      for j in xrange(1, imdb.num_classes)])
            if len(image_scores) > max_per_image:
                image_thresh = np.sort(image_scores)[-max_per_image]
                for j in xrange(1, imdb.num_classes):
                    keep = np.where(all_boxes[j][i][:, -1] >= image_thresh)[0]
                    all_boxes[j][i] = all_boxes[j][i][keep, :]
        _t['misc'].toc()

        print('im_detect: {:d}/{:d} {:.3f}s {:.3f}s' \
              .format(i + 1, num_images, _t['im_detect'].average_time,
                      _t['misc'].average_time))

    det_file = os.path.join(output_dir, 'detections.pkl')
    with open(det_file, 'wb') as f:
        cPickle.dump(all_boxes, f, cPickle.HIGHEST_PROTOCOL)

    print('Evaluating detections')
    imdb.evaluate_detections(all_boxes, output_dir)

def test_net_seg(net, imdb, suffix='',bbox=None):
    """Test a RPN-FCN network on an image database."""
    print('Evaluating segmentations')
    # This is for saving the segmentation prediction to disk
    #do_seg_tests(net, 0, os.path.join(imdb.data_path, 'pred{}'), imdb, suffix)
    # This is for testing for sorted dice on individual slices
    #do_seg_tests_on_slices(net, 0, False, imdb, suffix)
    # This is for not saving the pred
    #do_seg_tests(net, 0, False, imdb, suffix)
    if bbox is not None:
        do_seg_tests(net, 0, False, imdb, suffix, save_bbox = os.path.join(imdb.data_path,
            'bbox_pred'))
    else:
        do_seg_tests(net, 0, False, imdb, suffix)

def do_seg_tests_on_slices(net, iter, save_format, imdb, suffix, layer='score', gt='label',
        save_bbox = False):
    compute_hist_imdb_for_each_slice(net, save_format, imdb, layer, gt, save_bbox)

def do_seg_tests(net, iter, save_format, imdb, suffix, layer='score', gt='label',
        save_bbox = False):
    if save_format:
        if 'socket' in imdb.name:
            save_format = save_format.format('_socket_'+suffix)
        else:
            save_format = save_format.format(str(iter)+'_'+suffix)
    hist = compute_hist_imdb(net, save_format, imdb, layer, gt, save_bbox)
    #compute_metrics(hist, suffix)
    compute_metrics_flat(hist, suffix)
    #compute_mean_metrics(hist, suffix)
    return hist

def compute_mean_metrics(hist, iter):
    # overall accuracy
    acc = np.diag(hist).sum() / hist.sum()
    print '>>>', datetime.now(), 'Iteration', iter, 'overall accuracy', acc
    # per-class accuracy
    acc = np.diag(hist) / hist.sum(1)
    print '>>>', datetime.now(), 'Iteration', iter, 'mean accuracy', np.nanmean(acc)
    # per-class IU
    iu = np.diag(hist) / (hist.sum(1) + hist.sum(0) - np.diag(hist))
    print '>>>', datetime.now(), 'Iteration', iter, 'mean IU', np.nanmean(iu)
    freq = hist.sum(1) / hist.sum()
    print '>>>', datetime.now(), 'Iteration', iter, 'fwavacc', \
            (freq[freq > 0] * iu[freq > 0]).sum()
    # per-class Dice Score / F1, dice = 2*iu/(1+iu)
    dice = 2 * np.diag(hist) / (hist.sum(1) + hist.sum(0))
    print '>>>', datetime.now(), 'Iteration', iter, 'mean Dice Score', \
            np.nanmean(dice)

def compute_metrics(hist, iter):
    # overall accuracy
    acc = np.diag(hist).sum() / hist.sum()
    print '>>>', datetime.now(), 'Iteration', iter, 'overall accuracy', acc
    # positive-class accuracy
    acc = np.diag(hist)[1] / hist.sum()
    print '>>>', datetime.now(), 'Iteration', iter, 'accuracy', acc
    # positive-class IU
    iu = np.diag(hist)[1] / (hist[:,1].sum() + hist[1,:].sum() - np.diag(hist)[1])
    print '>>>', datetime.now(), 'Iteration', iter, 'IU', iu
    #freq = hist.sum(1) / hist.sum()
    #print '>>>', datetime.now(), 'Iteration', iter, 'fwavacc', \
    #        (freq[freq > 0] * iu[freq > 0]).sum()
    # positive class Dice Score / F1, dice = 2*iu/(1+iu)
    dice = 2 * np.diag(hist)[1] / (hist[:,1].sum() + hist[1,:].sum())
    print '>>>', datetime.now(), 'Iteration', iter, 'Dice Score', \
            dice
    dice0 = 2*iu/(1+iu)
    print '>>>', datetime.now(), 'Iteration', iter, 'Dice Score from IU', \
            dice0

def compute_metrics_flat(hist, iter):
    eps = 1e-10
    # overall accuracy
    acc = np.diag(hist).sum() / hist.sum()
    # recall
    recall = np.diag(hist)[1] / (hist[1,:].sum() + eps)
    # precision
    prec = np.diag(hist)[1] / (hist[:,1].sum() + eps)
    # mean IOU
    #iu = np.diag(hist) / (hist.sum(1) + hist.sum(0) - np.diag(hist))
    iu = np.diag(hist)[1] / (hist[:,1].sum() + hist[1,:].sum() -
            np.diag(hist)[1] + eps)
    # positive class Dice Score / F1, dice = 2*iu/(1+iu)
    if hist.shape[0] == 2:
        dice = 2 * np.diag(hist)[1] / (hist[:,1].sum() + hist[1,:].sum() + eps)
        print('>>>{}, Iteration {}, accuracy: {}, precision: {}, recall: {}, IOU: {}, Dice: {}'.format(datetime.now(), iter, acc, prec, recall, iu, dice))
    else:
        dice = 2 * np.diag(hist) / (hist.sum(1) + hist.sum(0) + eps)
        print('>>>{}, Iteration {}, accuracy: {}, precision: {}, recall: {}, mean IOU: {}, mean Dice: {}'.format(datetime.now(), iter, acc, prec, recall, np.nanmean(iu), np.nanmean(dice)))

def compute_hist_imdb(net, save_dir, imdb, layer='score', gt='label',
        save_bbox_dir = False,
        loss_layer='loss_cls'):
    n_cl = net.blobs[layer].channels
    if save_dir and not os.path.exists(save_dir):
        os.mkdir(save_dir)
    if save_bbox_dir and not os.path.exists(save_bbox_dir):
        os.mkdir(save_bbox_dir)
    if save_dir:
        # show where to save the image
        print("Pred will be saved to {}".format(save_dir))
    if save_bbox_dir:
        print("Predicted bbox will be saved to {}".format(save_bbox_dir))
    hist = np.zeros((n_cl, n_cl))
    #output_dir = get_output_dir(imdb, net)
    # timers
    _t = {'im_seg' : Timer()}
    num_images = len(imdb.image_index)
    is_jpg = '.jpg' in imdb.image_path_at(0) or '.png' in imdb.image_path_at(0)
    #import pdb;pdb.set_trace()
    for i in xrange(num_images):
        if is_jpg:
            # Load the demo image
            im = cv2.imread(imdb.image_path_at(i))
            # Load the label, in jpeg format
            #label = imdb.get_label(i)
            label = np.asarray(PIL.Image.open(imdb.label_path_at(i)))
            # binarize the label
            #label = np.where(label > 127, 1, 0)
        else:
            # h5 file
            vol_name, sliceidx = imdb.image_path_at(i).rsplit('_',1)
            im = imdb.get_image(vol_name, int(sliceidx))
            im = np.dstack((im, im, im))
            # this label is already binary
            label = imdb.get_label(vol_name, int(sliceidx))

        # One forward pass
        _t['im_seg'].tic()
        im_pred, boxes, boxes_score = im_seg(net, im, label)
        _t['im_seg'].toc()

        if i % 100 == 0:
            print 'im_seg: {:d}/{:d} {:.3f}s' \
                .format(i + 1, num_images, _t['im_seg'].average_time)

        hist += fast_hist(label.flatten(), im_pred.flatten(), n_cl)

        if save_dir:
            # from {0,1} to {0,255} for image display purpose
            im = PIL.Image.fromarray(im_pred*255, mode='L')
            im.save(os.path.join(save_dir, imdb.image_index[i] + '.png'))
        if save_bbox_dir:
            vol_name, sliceidx = imdb.image_path_at(i).rsplit('_',1)
            with open(os.path.join(save_bbox_dir, vol_name+'.txt'),"a") as text_file:
                # format of the line:
                # volname_sliceidx, num_of_boxes,
                num_box = boxes.shape[0]
                #pdb.set_trace()
                text_file.write("{}_{},{}".format(vol_name,sliceidx,num_box))
                for box_id in xrange(num_box):
                    text_file.write(",{},{},{},{},{}".format(boxes_score[box_id,0],
                        boxes[box_id,0],boxes[box_id,1], boxes[box_id,2],
                        boxes[box_id,3]))
                text_file.write("\n")
        # compute the loss as well
        #loss += net.blobs[loss_layer].data.flat[0]
    if save_dir:
        hist_dir = os.path.join(save_dir, 'hist.txt')
    else:
        hist_dir = os.path.join(imdb.data_path, 'hist.txt')
    with open(hist_dir, 'w') as hist_file:
        np.savetxt(hist_file, hist, '%f')
    return hist

def compute_hist_imdb_for_each_slice(net, save_dir, imdb, layer='score', gt='label',
        save_bbox_dir = False,
        loss_layer='loss_cls'):
    n_cl = net.blobs[layer].channels
    if save_dir and not os.path.exists(save_dir):
        os.mkdir(save_dir)
    if save_dir:
        # show where to save the image
        print("Pred will be saved to {}".format(save_dir))
    # timers
    _t = {'im_seg' : Timer()}
    num_images = len(imdb.image_index)
    for i in xrange(num_images):
        # h5 file
        vol_name, sliceidx = imdb.image_path_at(i).rsplit('_',1)
        im = imdb.get_image(vol_name, int(sliceidx))
        im = np.dstack((im, im, im))
        # this label is already binary
        label = imdb.get_label(vol_name, int(sliceidx))

        # One forward pass
        _t['im_seg'].tic()
        im_pred, boxes, boxes_score = im_seg(net, im, label)
        _t['im_seg'].toc()

        if i % 100 == 0:
            print 'im_seg: {:d}/{:d} {:.3f}s' \
                .format(i + 1, num_images, _t['im_seg'].average_time)

        hist = fast_hist(label.flatten(), im_pred.flatten(), n_cl)
        compute_metrics_flat(hist.astype(np.float32),
                "RPNFCN-{}-{}-{}".format(i, vol_name, sliceidx))

        if save_dir:
            # from {0,1} to {0,255} for image display purpose
            im = PIL.Image.fromarray(im_pred*255, mode='L')
            im.save(os.path.join(save_dir, imdb.image_index[i] + '.png'))
    return hist


import base64
import cv2
import io

import matplotlib.pyplot as plt
import numpy as np

from FER.predict import Model

# ----------------------------------------------------------------------------------
# Basic Processing/Conversions of image
# ----------------------------------------------------------------------------------
def process_image(file):
    # because imghdr.what() reads file to end, must set file's position 0.
    # some other function also reads the file to end
    file.seek(0)

    # Read Image
    image = cv2.imdecode(np.fromstring(
        file.read(), np.uint8), cv2.IMREAD_UNCHANGED)

    # Resizing Factor
    f1, f2 = 1200 / image.shape[1], 600 / image.shape[0]
    f = min(f1, f2, 1)

    # Resize Image
    dim = (int(image.shape[1] * f), int(image.shape[0] * f))
    image = cv2.resize(image, dim)

    # Return Image
    return image

# ----------------------------------------------------------------------------------
# Detect faces using OpenCV
# ----------------------------------------------------------------------------------
def detect_faces(file):
    """Detect face(s) in an image"""

    faces_dict = {}

    # Read image
    image = process_image(file)

    # Convert the image into gray scale
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

    # Load OpenCV face detector (LBP is faster)
    face_cascade = cv2.CascadeClassifier(
        'FER/haarcascade_frontalface_default.xml')

    # Detect multiple faces (some images may be closer to camera than others)
    # result is a list of faces
    faces = face_cascade.detectMultiScale(
        gray, scaleFactor=1.2, minNeighbors=5)

    # If not face detected, return empty list
    if len(faces) == 0:
        return faces, gray

    coords = ['x', 'y', 'w', 'h']

    for i in range(0, len(faces)):
        faces_dict[i] = {'coordinates': dict(zip(coords, faces[i].tolist()))}

    # Return the face image area and the face rectangle
    return faces_dict, gray

# ----------------------------------------------------------------------------------
# Plot the Graph for Analysis
# ----------------------------------------------------------------------------------
def plot(probs, width=0.8):
    """ plotting a bar chart of prediction probability """
    plt.switch_backend('Agg')

    # Emotion List
    emotions = Model.EMOTIONS_LIST

    # no of Probabilities
    n = len(probs)
    # numpy array as long as no of emotions
    X = np.arange(len(emotions))
    # width of a bar
    width = width/n

    # plot bar graph for 'n' faces
    for i in range(n):
        plt.bar(X + (i*width), probs[i], width=width, label=f'Face {i}')

    # place the x ticks
    plt.xticks(X + (n-1)*width/2, emotions)

    # y axis range
    plt.ylim(0, 1)

    # naming the axis
    plt.ylabel('Probability')
    plt.xlabel('Emotions')

    # Add Legend
    plt.legend(loc='upper center', bbox_to_anchor=(0.5, 1.17),
               fancybox=True, shadow=True, ncol=5)

    # Convert Graph into Numpy Array
    buf = io.BytesIO()
    plt.savefig(buf, format="png")
    buf.seek(0)
    graph_arr = np.frombuffer(buf.getvalue(), dtype=np.uint8)
    buf.close()

    # Encode the graph into base64
    graph = base64.encodebytes(graph_arr)
    graph = 'data:image/png;base64, ' + str(graph, 'utf-8')
    return graph

# ----------------------------------------------------------------------------------
# Predict Emotion for face(s)
# ----------------------------------------------------------------------------------
def predict(faces_dict, img):
    """ Draw a rectangle(s) on the image and write their suitable emotions """
    # Load the Model
    predictor = Model("FER/model.json", "FER/model_weights.h5")
    emotions = predictor.EMOTIONS_LIST

    for f, detail in faces_dict.items():
        # decalaring face coordiantes as variables
        (x, y, w, h) = detail['coordinates'].values()

        # Select only Face
        fc = img[y:y + h, x:x + w]

        # Convert image into 48x48 and predict Emotion
        roi = cv2.resize(fc, (48, 48))
        emotion, prob = predictor.predict_emotion(
            roi[np.newaxis, :, :, np.newaxis])

        # add emotion and probability to faces_dict
        faces_dict[f]['emotion'] = emotion
        faces_dict[f]['probability'] = dict(zip(emotions, prob.tolist()))

# ----------------------------------------------------------------------------------
# Draw rectangle on image
# Write Emotion on Image
# according to given (x, y) coordinates and given width and height
# ----------------------------------------------------------------------------------
def edit_image(file, data):
    """ Draw a rectangle(s) on the image and write their suitable emotions """
    # Process Image
    image = process_image(file)

    probabilities = []

    # Font type
    font = cv2.FONT_HERSHEY_SIMPLEX

    for i, detail in data.items():
        # Variable Declaration
        (h, w, x, y) = detail['coordinates'].values()
        emotion = detail['emotion']

        # Text Modification
        font_size = h/300 if h/300 > 0.6 else 0.6

        # Draw Rectangle and Write Text
        cv2.putText(image, emotion, (x, y-4), font, font_size, (0, 255, 255), 2)
        cv2.putText(image, i, (x+2, y+h-4), font, font_size, (237, 162, 0), 2)
        cv2.rectangle(image, (x, y), (x+w, y+h), (0, 165, 255), 2)

        # append probabilities
        probabilities.append(detail['probability'].values())

    # Convert Image into base64
    image_content = cv2.imencode('.jpg', image)[1].tostring()
    encoded_image = base64.encodebytes(image_content)
    image = 'data:image/jpg;base64, ' + str(encoded_image, 'utf-8')

    # Plot Graph
    graph = plot(probabilities)

    # Return Image and Graph
    return image, graph


"""
Samplers for perses automated molecular design.

TODO
----
* Determine where `System` object should be stored: In `SamplerState` or in `Thermodynamic State`, or both, or neither?
* Can we create a generalized, extensible `SamplerState` that also stores chemical/thermodynamic state information?
* Can we create a generalized log biasing weight container class that gracefully handles new chemical states that have yet to be explored?

"""

__author__ = 'John D. Chodera'

################################################################################
# IMPORTS
################################################################################

import mdtraj as md
import numpy as np
import time
from openmmtools.states import SamplerState, ThermodynamicState
from openmmtools import cache, utils
from perses.dispersed.utils import configure_platform
cache.global_context_cache.platform = configure_platform(utils.get_fastest_platform().getName())

from perses.annihilation.ncmc_switching import NCMCEngine
from perses.dispersed import feptasks
from perses.storage import NetCDFStorageView
from perses.utils.openeye import smiles_to_oemol

################################################################################
# LOGGER
################################################################################

import logging
_logger = logging.getLogger()
_logger.setLevel(logging.INFO)
_logger = logging.getLogger("samplers")

################################################################################
# UTILITY FUNCTIONS
################################################################################

def log_sum_exp(a_n):
    """
    Compute log(sum(exp(a_n)))

    Parameters
    ----------
    a_n : dict of objects : floats

    """
    a_n = np.array(list(a_n.values()))
    return np.log( np.sum( np.exp(a_n - a_n.max() ) ) )

################################################################################
# EXPANDED ENSEMBLE SAMPLER
################################################################################

class ExpandedEnsembleSampler(object):
    """
    Method of expanded ensembles sampling engine.

    The acceptance criteria is given in the reference document. Roughly, the proposal scheme is:

    * Draw a proposed chemical state k', and calculate reverse proposal probability
    * Conditioned on k' and the current positions x, generate new positions with the GeometryEngine
    * With new positions, jump to a hybrid system at lambda=0
    * Anneal from lambda=0 to lambda=1, accumulating work
    * Jump from the hybrid system at lambda=1 to the k' system, and compute reverse GeometryEngine proposal
    * Add weight of chemical states k and k' to acceptance probabilities

    Properties
    ----------
    sampler : MCMCSampler
        The MCMC sampler used for updating positions.
    proposal_engine : ProposalEngine
        The ProposalEngine to use for proposing new sampler states and topologies.
    system_generator : SystemGenerator
        The SystemGenerator to use for creating System objects following proposals.
    state : hashable object
        The current sampler state. Can be any hashable object.
    states : set of hashable object
        All known states.
    iteration : int
        Iterations completed.
    naccepted : int
        Number of accepted thermodynamic/chemical state changes.
    nrejected : int
        Number of rejected thermodynamic/chemical state changes.
    number_of_state_visits : dict of state_key
        Cumulative counts of visited states.
    verbose : bool
        If True, verbose output is printed.

    References
    ----------
    [1] Lyubartsev AP, Martsinovski AA, Shevkunov SV, and Vorontsov-Velyaminov PN. New approach to Monte Carlo calculation of the free energy: Method of expanded ensembles. JCP 96:1776, 1992
    http://dx.doi.org/10.1063/1.462133

    Examples
    --------
    >>> # Create a test system
    >>> test = testsystems.AlanineDipeptideVacuum()
    >>> # Create a SystemGenerator and rebuild the System.
    >>> from perses.rjmc.topology_proposal import SystemGenerator
    >>> system_generator = SystemGenerator(['amber99sbildn.xml'], forcefield_kwargs={'implicitSolvent' : None, 'constraints' : None }, nonperiodic_forcefield_kwargs={'nonbondedMethod' : app.NoCutoff})
    >>> test.system = system_generator.build_system(test.topology)
    >>> # Create a sampler state.
    >>> sampler_state = SamplerState(system=test.system, positions=test.positions)
    >>> # Create a thermodynamic state.
    >>> thermodynamic_state = ThermodynamicState(system=test.system, temperature=298.0*unit.kelvin)
    >>> # Create an MCMC sampler
    >>> mcmc_sampler = MCMCSampler(thermodynamic_state, sampler_state)
    >>> # Turn off verbosity
    >>> mcmc_sampler.verbose = False
    >>> # Create an Expanded Ensemble sampler
    >>> from perses.rjmc.topology_proposal import PointMutationEngine
    >>> from perses.rjmc.geometry import FFAllAngleGeometryEngine
    >>> geometry_engine = FFAllAngleGeometryEngine(metadata={})
    >>> allowed_mutations = [[('2','ALA')],[('2','VAL'),('2','LEU')]]
    >>> proposal_engine = PointMutationEngine(test.topology, system_generator, max_point_mutants=1, chain_id='1', proposal_metadata=None, allowed_mutations=allowed_mutations)
    >>> exen_sampler = ExpandedEnsembleSampler(mcmc_sampler, test.topology, 'ACE-ALA-NME', proposal_engine, geometry_engine)
    >>> # Run the sampler
    >>> exen_sampler.run()

    """
    def __init__(self, sampler, topology, state_key, proposal_engine, geometry_engine, log_weights=None, options=None, platform=None, envname=None, storage=None, ncmc_write_interval=1):
        r"""
        Create an expanded ensemble sampler.

        p(x,k) \propto \exp[-u_k(x) + g_k]

        where g_k is the log weight.

        Parameters
        ----------
        sampler : MCMCSampler
            MCMCSampler initialized with current SamplerState
        topology : simtk.openmm.app.Topology
            Current topology
        state : hashable object
            Current chemical state
        proposal_engine : ProposalEngine
            ProposalEngine to use for proposing new chemical states
        geometry_engine : GeometryEngine
            GeometryEngine to use for dimension matching
        log_weights : dict of object : float
            Log weights to use for expanded ensemble biases.
        options : dict, optional, default=dict()
            Options for initializing switching scheme, such as 'timestep', 'nsteps', 'functions' for NCMC
        platform : simtk.openmm.Platform, optional, default=None
            Platform to use for NCMC switching.  If `None`, default (fastest) platform is used.
        storage : NetCDFStorageView, optional, default=None
            If specified, use this storage layer.
        ncmc_write_interval : int, default 1
            How frequently to write out NCMC protocol steps.
        """
        # Keep copies of initializing arguments.
        # TODO: Make deep copies?
        self.sampler = sampler
        self._pressure = sampler.thermodynamic_state.pressure
        self._temperature = sampler.thermodynamic_state.temperature
        self._omm_topology = topology
        self.topology = md.Topology.from_openmm(topology)
        self.state_key = state_key
        self.proposal_engine = proposal_engine
        self.log_weights = log_weights
        if self.log_weights is None: self.log_weights = dict()

        self.storage = None
        if storage is not None:
            self.storage = NetCDFStorageView(storage, modname=self.__class__.__name__)

        # Initialize
        self.iteration = 0
        option_names = ['timestep', 'nsteps', 'functions', 'nsteps_mcmc', 'splitting']

        if options is None:
            options = dict()
        for option_name in option_names:
            if option_name not in options:
                options[option_name] = None

        if options['splitting']:
            self._ncmc_splitting = options['splitting']
        else:
            self._ncmc_splitting = "V R O H R V"

        if options['nsteps']:
            self._switching_nsteps = options['nsteps']
            self.ncmc_engine = NCMCEngine(temperature=self.sampler.thermodynamic_state.temperature,
                                          timestep=options['timestep'], nsteps=options['nsteps'],
                                          functions=options['functions'], integrator_splitting=self._ncmc_splitting,
                                          platform=platform, storage=self.storage,
                                          write_ncmc_interval=ncmc_write_interval)
        else:
            self._switching_nsteps = 0

        if options['nsteps_mcmc']:
            self._n_iterations_per_update = options['nsteps_mcmc']
        else:
            self._n_iterations_per_update = 100

        self.geometry_engine = geometry_engine
        self.naccepted = 0
        self.nrejected = 0
        self.number_of_state_visits = dict()
        self.verbose = False
        self.pdbfile = None # if not None, write PDB file
        self.geometry_pdbfile = None # if not None, write PDB file of geometry proposals
        self.accept_everything = False # if True, will accept anything that doesn't lead to NaNs
        self.logPs = list()
        self.sampler.minimize(max_iterations=40)

    @property
    def state_keys(self):
        return self.log_weights.keys()

    def get_log_weight(self, state_key):
        """
        Get the log weight of the specified state.

        Parameters
        ----------
        state_key : hashable object
            The state key (e.g. chemical state key) to look up.

        Returns
        -------
        log_weight : float
            The log weight of the provided state key.

        Notes
        -----
        This adds the key to the self.log_weights dict.

        """
        if state_key not in self.log_weights:
            self.log_weights[state_key] = 0.0
        return self.log_weights[state_key]

    def _system_to_thermodynamic_state(self, system):
        """
        Given an OpenMM system object, create a corresponding ThermodynamicState that has the same
        temperature and pressure as the current thermodynamic state.

        Parameters
        ----------
        system : openmm.System
            The OpenMM system for which to create the thermodynamic state

        Returns
        -------
        new_thermodynamic_state : openmmtools.states.ThermodynamicState
            The thermodynamic state object representing the given system
        """
        return ThermodynamicState(system, temperature=self._temperature, pressure=self._pressure)

    def _geometry_forward(self, topology_proposal, old_sampler_state):
        """
        Run geometry engine to propose new positions and compute logP

        Parameters
        ----------
        topology_proposal : TopologyProposal
            Contains old/new Topology and System objects and atom mappings.
        old_sampler_state : openmmtools.states.SamplerState
            Configurational properties of the old system atoms.

        Returns
        -------
        new_sampler_state : openmmtools.states.SamplerState
            Configurational properties of new atoms proposed by geometry engine calculation.
        geometry_logp_propose : float
            The log probability of the forward-only proposal
        """
        if self.verbose: print("Geometry engine proposal...")
        # Generate coordinates for new atoms and compute probability ratio of old and new probabilities.
        initial_time = time.time()
        new_positions, geometry_logp_propose = self.geometry_engine.propose(topology_proposal, old_sampler_state.positions, self.sampler.thermodynamic_state.beta)
        if self.verbose: print('proposal took %.3f s' % (time.time() - initial_time))

        if self.geometry_pdbfile is not None:
            print("Writing proposed geometry...")
            from simtk.openmm.app import PDBFile
            PDBFile.writeFile(topology_proposal.new_topology, new_positions, file=self.geometry_pdbfile)
            self.geometry_pdbfile.flush()

        new_sampler_state = SamplerState(new_positions, box_vectors=old_sampler_state.box_vectors)

        return new_sampler_state, geometry_logp_propose

    def _geometry_reverse(self, topology_proposal, new_sampler_state, old_sampler_state):
        """
        Run geometry engine reverse calculation to determine logP
        of proposing the old positions based on the new positions

        Parameters
        ----------
        topology_proposal : TopologyProposal
            Contains old/new Topology and System objects and atom mappings.
        new_sampler_state : openmmtools.states.SamplerState
            Configurational properties of the new atoms.
        old_sampler_state : openmmtools.states.SamplerState
            Configurational properties of the old atoms.

        Returns
        -------
        geometry_logp_reverse : float
            The log probability of the proposal for the given transformation
        """
        if self.verbose: print("Geometry engine logP_reverse calculation...")
        initial_time = time.time()
        geometry_logp_reverse = self.geometry_engine.logp_reverse(topology_proposal, new_sampler_state.positions, old_sampler_state.positions, self.sampler.thermodynamic_state.beta)
        if self.verbose: print('calculation took %.3f s' % (time.time() - initial_time))
        return geometry_logp_reverse

    def _ncmc_hybrid(self, topology_proposal, old_sampler_state, new_sampler_state):
        """
        Run a hybrid NCMC protocol from lambda = 0 to lambda = 1

        Parameters
        ----------
        topology_proposal : TopologyProposal
            Contains old/new Topology and System objects and atom mappings.
        old_sampler_State : openmmtools.states.SamplerState
            SamplerState of old system at the beginning of NCMCSwitching
        new_sampler_state : openmmtools.states.SamplerState
            SamplerState of new system at the beginning of NCMCSwitching

        Returns
        -------
        old_final_sampler_state : openmmtools.states.SamplerState
            SamplerState of old system at the end of switching
        new_final_sampler_state : openmmtools.states.SamplerState
            SamplerState of new system at the end of switching
        logP_work : float
            The NCMC work contribution to the log acceptance probability (Eq. 44)
        logP_energy : float
            The contribution of switching to and from the hybrid system to the acceptance probability (Eq. 45)
        """
        if self.verbose: print("Performing NCMC switching")
        initial_time = time.time()
        [ncmc_old_sampler_state, ncmc_new_sampler_state, logP_work, logP_initial_hybrid, logP_final_hybrid] = self.ncmc_engine.integrate(topology_proposal, old_sampler_state, new_sampler_state, iteration=self.iteration)
        if self.verbose: print('NCMC took %.3f s' % (time.time() - initial_time))
        # Check that positions are not NaN
        if new_sampler_state.has_nan():
            raise Exception("Positions are NaN after NCMC insert with %d steps" % self._switching_nsteps)
        return ncmc_old_sampler_state, ncmc_new_sampler_state, logP_work, logP_initial_hybrid, logP_final_hybrid

    def _geometry_ncmc_geometry(self, topology_proposal, sampler_state, old_log_weight, new_log_weight):
        """
        Use a hybrid NCMC protocol to switch from the old system to new system
        Will calculate new positions for the new system first, then give both
        sets of positions to the hybrid NCMC integrator, and finally use the
        final positions of the old and new systems to calculate the reverse
        geometry probability

        Parameters
        ----------
        topology_proposal : TopologyProposal
            Contains old/new Topology and System objects and atom mappings.
        sampler_state : openmmtools.states.SamplerState
            Configurational properties of old atoms at the beginning of the NCMC switching.
        old_log_weight : float
            Chemical state weight from SAMSSampler
        new_log_weight : float
            Chemical state weight from SAMSSampler

        Returns
        -------
        logP_accept : float
            Log of acceptance probability of entire Expanded Ensemble switch (Eq. 25 or 46)
        ncmc_new_sampler_state : openmmtools.states.SamplerState
            Configurational properties of new atoms at the end of the NCMC switching.
        """
        if self.verbose: print("Updating chemical state with geometry-ncmc-geometry scheme...")

        logP_chemical_proposal = topology_proposal.logp_proposal

        old_thermodynamic_state = self.sampler.thermodynamic_state
        new_thermodynamic_state = self._system_to_thermodynamic_state(topology_proposal.new_system)

        initial_reduced_potential = feptasks.compute_reduced_potential(old_thermodynamic_state, sampler_state)
        logP_initial_nonalchemical = - initial_reduced_potential

        new_geometry_sampler_state, logP_geometry_forward = self._geometry_forward(topology_proposal, sampler_state)

        #if we aren't doing any switching, then skip running the NCMC engine at all.
        if self._switching_nsteps == 0:
            ncmc_old_sampler_state = sampler_state
            ncmc_new_sampler_state = new_geometry_sampler_state
            logP_work = 0.0
            logP_initial_hybrid = 0.0
            logP_final_hybrid = 0.0
        else:
            ncmc_old_sampler_state, ncmc_new_sampler_state, logP_work, logP_initial_hybrid, logP_final_hybrid = self._ncmc_hybrid(topology_proposal, sampler_state, new_geometry_sampler_state)

        if logP_work > -np.inf and logP_initial_hybrid > -np.inf and logP_final_hybrid > -np.inf:
            logP_geometry_reverse = self._geometry_reverse(topology_proposal, ncmc_new_sampler_state, ncmc_old_sampler_state)
            logP_to_hybrid = logP_initial_hybrid - logP_initial_nonalchemical

            final_reduced_potential = feptasks.compute_reduced_potential(new_thermodynamic_state, ncmc_new_sampler_state)
            logP_final_nonalchemical = -final_reduced_potential
            logP_from_hybrid = logP_final_nonalchemical - logP_final_hybrid
            logP_sams_weight = new_log_weight - old_log_weight

            # Compute total log acceptance probability according to Eq. 46
            logP_accept = logP_to_hybrid - logP_geometry_forward + logP_work + logP_from_hybrid + logP_geometry_reverse + logP_sams_weight
        else:
            logP_geometry_reverse = 0.0
            logP_final = 0.0
            logP_to_hybrid = 0.0
            logP_from_hybrid = 0.0
            logP_sams_weight = new_log_weight - old_log_weight
            logP_accept = logP_to_hybrid - logP_geometry_forward + logP_work + logP_from_hybrid + logP_geometry_reverse + logP_sams_weight
            #TODO: mark failed proposals as unproposable

        if self.verbose:
            print("logP_accept = %+10.4e [logP_to_hybrid = %+10.4e, logP_chemical_proposal = %10.4e, logP_reverse = %+10.4e, -logP_forward = %+10.4e, logP_work = %+10.4e, logP_from_hybrid = %+10.4e, logP_sams_weight = %+10.4e]"
                % (logP_accept, logP_to_hybrid, logP_chemical_proposal, logP_geometry_reverse, -logP_geometry_forward, logP_work, logP_from_hybrid, logP_sams_weight))
        # Write to storage.
        if self.storage:
            self.storage.write_quantity('logP_accept', logP_accept, iteration=self.iteration)
            # Write components to storage
            self.storage.write_quantity('logP_ncmc_work', logP_work, iteration=self.iteration)
            self.storage.write_quantity('logP_from_hybrid', logP_from_hybrid, iteration=self.iteration)
            self.storage.write_quantity('logP_to_hybrid', logP_to_hybrid, iteration=self.iteration)
            self.storage.write_quantity('logP_chemical_proposal', logP_chemical_proposal, iteration=self.iteration)
            self.storage.write_quantity('logP_reverse', logP_geometry_reverse, iteration=self.iteration)
            self.storage.write_quantity('logP_forward', logP_geometry_forward, iteration=self.iteration)
            self.storage.write_quantity('logP_sams_weight', logP_sams_weight, iteration=self.iteration)
            # Write some aggregate statistics to storage to make contributions to acceptance probability easier to analyze
            self.storage.write_quantity('logP_groups_chemical', logP_chemical_proposal, iteration=self.iteration)
            self.storage.write_quantity('logP_groups_geometry', logP_geometry_reverse - logP_geometry_forward, iteration=self.iteration)

        return logP_accept, ncmc_new_sampler_state

    def update_positions(self, n_iterations=1):
        """
        Sample new positions.
        """
        self.sampler.run(n_iterations=n_iterations)

    def update_state(self):
        """
        Sample the thermodynamic state.
        """

        initial_time = time.time()

        # Propose new chemical state.
        if self.verbose: print("Proposing new topology...")
        [system, positions] = [self.sampler.thermodynamic_state.get_system(remove_thermostat=True), self.sampler.sampler_state.positions]

        #omm_topology = topology.to_openmm() #convert to OpenMM topology for proposal engine
        self._omm_topology.setPeriodicBoxVectors(self.sampler.sampler_state.box_vectors) #set the box vectors because in OpenMM topology has these...
        topology_proposal = self.proposal_engine.propose(system, self._omm_topology)
        if self.verbose: print("Proposed transformation: %s => %s" % (topology_proposal.old_chemical_state_key, topology_proposal.new_chemical_state_key))

        # Determine state keys
        old_state_key = self.state_key
        new_state_key = topology_proposal.new_chemical_state_key

        # Determine log weight
        old_log_weight = self.get_log_weight(old_state_key)
        new_log_weight = self.get_log_weight(new_state_key)

        logp_accept, ncmc_new_sampler_state = self._geometry_ncmc_geometry(topology_proposal, self.sampler.sampler_state, old_log_weight, new_log_weight)

        # Accept or reject.
        if np.isnan(logp_accept):
            accept = False
            print('logp_accept = NaN')
        else:
            accept = ((logp_accept>=0.0) or (np.random.uniform() < np.exp(logp_accept)))
            if self.accept_everything:
                print('accept_everything option is turned on; accepting')
                accept = True

        if accept:
            self.sampler.thermodynamic_state.set_system(topology_proposal.new_system, fix_state=True)
            self.sampler.sampler_state.system = topology_proposal.new_system
            self.topology = md.Topology.from_openmm(topology_proposal.new_topology)
            self.sampler.sampler_state = ncmc_new_sampler_state
            self.sampler.topology = self.topology
            self.state_key = topology_proposal.new_chemical_state_key
            self.naccepted += 1
            if self.verbose: print("    accepted")
        else:
            self.nrejected += 1
            if self.verbose: print("    rejected")

        if self.storage:
            self.storage.write_configuration('positions', self.sampler.sampler_state.positions, self.topology, iteration=self.iteration)
            self.storage.write_object('state_key', self.state_key, iteration=self.iteration)
            self.storage.write_object('proposed_state_key', topology_proposal.new_chemical_state_key, iteration=self.iteration)
            self.storage.write_quantity('naccepted', self.naccepted, iteration=self.iteration)
            self.storage.write_quantity('nrejected', self.nrejected, iteration=self.iteration)
            self.storage.write_quantity('logp_accept', logp_accept, iteration=self.iteration)
            self.storage.write_quantity('logp_topology_proposal', topology_proposal.logp_proposal, iteration=self.iteration)

        # Update statistics.
        self.update_statistics()

    def update(self):
        """
        Update the sampler with one step of sampling.
        """
        if self.verbose:
            print("-" * 80)
            print("Expanded Ensemble sampler iteration %8d" % self.iteration)
        self.update_positions(n_iterations=self._n_iterations_per_update)
        self.update_state()
        self.iteration += 1
        if self.verbose:
            print("-" * 80)

        if self.pdbfile is not None:
            print("Writing frame...")
            from simtk.openmm.app import PDBFile
            PDBFile.writeModel(self.topology.to_openmm(), self.sampler.sampler_state.positions, self.pdbfile, self.iteration)
            self.pdbfile.flush()

        if self.storage:
            self.storage.sync()

    def run(self, niterations=1):
        """
        Run the sampler for the specified number of iterations

        Parameters
        ----------
        niterations : int, optional, default=1
            Number of iterations to run the sampler for.
        """
        for iteration in range(niterations):
            self.update()

    def update_statistics(self):
        """
        Update sampler statistics.
        """
        if self.state_key not in self.number_of_state_visits:
            self.number_of_state_visits[self.state_key] = 0
        self.number_of_state_visits[self.state_key] += 1

################################################################################
# SAMS SAMPLER
################################################################################

class SAMSSampler(object):
    """
    Self-adjusted mixture sampling engine.

    Properties
    ----------
    state_keys : set of objects
        The names of states sampled by the sampler.
    logZ : dict() of keys : float
        logZ[key] is the log partition function (up to an additive constant) estimate for chemical state `key`
    update_method : str
        Update method.  One of ['default']
    iteration : int
        Iterations completed.
    verbose : bool
        If True, verbose debug output is printed.

    References
    ----------
    [1] Tan, Z. (2015) Optimally adjusted mixture sampling and locally weighted histogram analysis, Journal of Computational and Graphical Statistics, to appear. (Supplement)
    http://www.stat.rutgers.edu/home/ztan/Publication/SAMS_redo4.pdf

    Examples
    --------
    >>> # Create a test system
    >>> test = testsystems.AlanineDipeptideVacuum()
    >>> # Create a SystemGenerator and rebuild the System.
    >>> from perses.rjmc.topology_proposal import SystemGenerator
    >>> system_generator = SystemGenerator(['amber99sbildn.xml'], forcefield_kwargs={'implicitSolvent' : None, 'constraints' : None }, nonperiodic_forcefield_kwargs={'nonbondedMethod' : app.NoCutoff})
    >>> test.system = system_generator.build_system(test.topology)
    >>> # Create a sampler state.
    >>> sampler_state = SamplerState(system=test.system, positions=test.positions)
    >>> # Create a thermodynamic state.
    >>> thermodynamic_state = ThermodynamicState(system=test.system, temperature=298.0*unit.kelvin)
    >>> # Create an MCMC sampler
    >>> mcmc_sampler = MCMCSampler(thermodynamic_state, sampler_state)
    >>> # Turn off verbosity
    >>> mcmc_sampler.verbose = False
    >>> from perses.rjmc.geometry import FFAllAngleGeometryEngine
    >>> geometry_engine = FFAllAngleGeometryEngine(metadata={})
    >>> # Create an Expanded Ensemble sampler
    >>> from perses.rjmc.topology_proposal import PointMutationEngine
    >>> allowed_mutations = [[('2','ALA')],[('2','VAL'),('2','LEU')]]
    >>> proposal_engine = PointMutationEngine(test.topology, system_generator, max_point_mutants=1, chain_id='1', proposal_metadata=None, allowed_mutations=allowed_mutations)
    >>> exen_sampler = ExpandedEnsembleSampler(mcmc_sampler, test.topology, 'ACE-ALA-NME', proposal_engine, geometry_engine)
    >>> # Create a SAMS sampler
    >>> sams_sampler = SAMSSampler(exen_sampler)
    >>> # Run the sampler
    >>> sams_sampler.run() # doctest: +ELLIPSIS
    ...
    """
    def __init__(self, sampler, logZ=None, log_target_probabilities=None, update_method='two-stage', storage=None, second_stage_start=1000):
        """
        Create a SAMS Sampler.

        Parameters
        ----------
        sampler : ExpandedEnsembleSampler
            The expanded ensemble sampler used to sample both configurations and discrete thermodynamic states.
        logZ : dict of key : float, optional, default=None
            If specified, the log partition functions for each state will be initialized to the specified dictionary.
        log_target_probabilities : dict of key : float, optional, default=None
            If specified, unnormalized target probabilities; default is all 0.
        update_method : str, optional, default='default'
            SAMS update algorithm
        storage : NetCDFStorageView, optional, default=None
        second_state_start : int, optional, default None
            At what iteration number to switch to the optimal gain decay

        """
        from scipy.special import logsumexp

        # Keep copies of initializing arguments.
        # TODO: Make deep copies?
        self.sampler = sampler
        self.chemical_states = None
        self._reference_state = None
        try:
            self.chemical_states = self.sampler.proposal_engine.chemical_state_list
        except NotImplementedError:
            _logger.warning("The proposal engine has not properly implemented the chemical state property; SAMS will add states on the fly.")

        if self.chemical_states:
            # Select a reference state that will always be subtracted (ensure that dict ordering does not change)
            self._reference_state = self.chemical_states[0]

            # Initialize the logZ dictionary with scores based on the number of atoms
            # This is not the negative because the weights are set to the negative of the initial weights
            self.logZ = {chemical_state: self._num_dof_compensation(chemical_state) for chemical_state in self.chemical_states}

            #Initialize log target probabilities with log(1/n_states)
            self.log_target_probabilities = {chemical_state : np.log(len(self.chemical_states)) for chemical_state in self.chemical_states}

            #If initial weights are specified, override any weight with what is provided
            #However, if the chemical state is not in the reachable chemical state list,throw an exception
            if logZ is not None:
                for (chemical_state, logZ_value) in logZ:
                    if chemical_state not in self.chemical_states:
                        raise ValueError("Provided a logZ initial value for an un-proposable chemical state")
                    self.logZ[chemical_state] = logZ_value

            if log_target_probabilities is not None:
                for (chemical_state, log_target_probability) in log_target_probabilities:
                    if chemical_state not in self.chemical_states:
                        raise ValueError("Provided a log target probability for an un-proposable chemical state.")
                    self.log_target_probabilities[chemical_state] = log_target_probability

                #normalize target probabilities
                #this is likely not necessary, but it is copying the algorithm in Ref 1
                log_sum_target_probabilities = logsumexp((list(self.log_target_probabilities.values())))
                self.log_target_probabilities = {chemical_state : log_target_probability - log_sum_target_probabilities for chemical_state, log_target_probability in self.log_target_probabilities}
        else:
            self.logZ = dict()
            self.log_target_probabilities = dict()

        self.update_method = update_method

        self.storage = None
        if storage is not None:
            self.storage = NetCDFStorageView(storage, modname=self.__class__.__name__)

        # Initialize.
        self.iteration = 0
        self.verbose = False
        self.sampler.log_weights = {state_key: - self.logZ[state_key] for state_key in self.logZ.keys()}

        self.second_stage_start = 0
        if second_stage_start is not None:
            self.second_stage_start = second_stage_start

    @property
    def state_keys(self):
        return self.logZ.keys()

    def _num_dof_compensation(self, smiles):
        """
        Compute an approximate compensating factor for a chemical state based on the number of degrees of freedom that it has.

        The formula is:
        (num_heavy*heavy_factor) + (num_hydrogen*h_factor) where
        heavy_factor = 4.5 and
        light_factor = 3.8

        Parameters
        ----------
        smiles : str
            The SMILES string of the molecule

        Returns
        -------
        correction_factor : float
        """
        mol = smiles_to_oemol(smiles)
        num_heavy = 0
        num_light = 0

        heavy_factor = 4.5
        light_factor = 3.8

        for atom in mol.GetAtoms():
            if atom.GetAtomicNum() == 1:
                num_light += 1
            else:
                num_heavy += 1

        correction_factor = num_heavy*heavy_factor + num_light*light_factor

        return correction_factor

    def update_sampler(self):
        """
        Update the underlying expanded ensembles sampler.
        """
        self.sampler.update()

    def update_logZ_estimates(self):
        """
        Update the logZ estimates according to self.update_method.
        """
        state_key = self.sampler.state_key

        # Add state key to dictionaries if we haven't visited this state before.
        if state_key not in self.logZ:
            _logger.warning("A new state key is being added to the logZ; note that this makes the resultant algorithm different from SAMS")
            self.logZ[state_key] = 0.0
        if state_key not in self.log_target_probabilities:
            _logger.warning("A new state key is being added to the target probabilities; note that this makes the resultant algorithm different from SAMS")
            self.log_target_probabilities[state_key] = 0.0

        # Update estimates of logZ.
        if self.update_method == 'one-stage':
            # Based on Eq. 9 of Ref. [1]
            gamma = 1.0 / float(self.iteration+1)
        elif self.update_method == 'two-stage':
            # Keep gamma large until second stage is activated.
            if self.iteration < self.second_stage_start:
                # First stage.
                gamma = 1.0
                # TODO: Determine when to switch to second stage
            else:
                # Second stage.
                gamma = 1.0 / float(self.iteration - self.second_stage_start + 1)
        else:
            raise Exception("SAMS update method '%s' unknown." % self.update_method)

        #get the (t-1/2) update from equation 9 in ref 1
        self.logZ[state_key] += gamma / np.exp(self.log_target_probabilities[state_key])

        if self._reference_state:
            #the second step of the (t-1/2 update), subtracting the reference state from everything else.
            #we can only do this for cases where all states have been enumerated
            self.logZ = {state_key : logZ_estimate - self.logZ[self._reference_state] for state_key, logZ_estimate in self.logZ.items()}

        # Update log weights for sampler.
        self.sampler.log_weights = { state_key : - self.logZ[state_key] for state_key in self.logZ.keys()}

        if self.storage:
            self.storage.write_object('logZ', self.logZ, iteration=self.iteration)
            self.storage.write_object('log_weights', self.sampler.log_weights, iteration=self.iteration)

    def update(self):
        """
        Update the sampler with one step of sampling.
        """
        if self.verbose:
            print("=" * 80)
            print("SAMS sampler iteration %5d" % self.iteration)
        self.update_sampler()
        self.update_logZ_estimates()
        if self.storage: self.storage.sync()
        self.iteration += 1
        if self.verbose:
            print("=" * 80)

    def run(self, niterations=1):
        """
        Run the sampler for the specified number of iterations

        Parameters
        ----------
        niterations : int, optional, default=1
            Number of iterations to run the sampler for.
        """
        for iteration in range(niterations):
            self.update()

################################################################################
# MULTITARGET OPTIMIZATION SAMPLER
################################################################################

class MultiTargetDesign(object):
    """
    Multi-objective design using self-adjusted mixture sampling with additional recursion steps
    that update target weights on the fly.

    Parameters
    ----------
    samplers : list of SAMSSampler
        The SAMS samplers whose relative partition functions go into the design objective computation.
    sampler_exponents : dict of SAMSSampler : float
        samplers.keys() are the samplers, and samplers[key]
    log_target_probabilities : dict of hashable object : float
        log_target_probabilities[key] is the computed log objective function (target probability) for chemical state `key`
    verbose : bool
        If True, verbose output is printed.

    """
    def __init__(self, target_samplers, storage=None, verbose=False):
        r"""
        Initialize a multi-objective design sampler with the specified target sampler powers.

        Parameters
        ----------
        target_samplers : dict
            target_samplers[sampler] is the exponent associated with SAMS sampler `sampler` in the multi-objective design.
        storage : NetCDFStorage, optional, default=None
            If specified, will use the storage layer to write trajectory data.
        verbose : bool, optional, default=False
            If true, will print verbose output

        The target sampler weights for N samplers with specified exponents \alpha_n are given by

        \pi_{nk} \propto \prod_{n=1}^N Z_{nk}^{alpha_n}

        where \pi_{nk} is the target weight for sampler n state k,
        and Z_{nk} is the relative partition function of sampler n among states k.

        Examples
        --------
        Set up a mutation sampler to maximize implicit solvent hydration free energy.
        >>> from perses.tests.testsystems import AlanineDipeptideTestSystem
        >>> testsystem = AlanineDipeptideTestSystem()
        >>> # Set up target samplers.
        >>> target_samplers = { testsystem.sams_samplers['implicit'] : 1.0, testsystem.sams_samplers['vacuum'] : -1.0 }
        >>> # Set up the design sampler.
        >>> designer = MultiTargetDesign(target_samplers)

        """
        # Store target samplers.
        self.sampler_exponents = target_samplers
        self.samplers = list(target_samplers.keys())

        self.storage = None
        if storage is not None:
            self.storage = NetCDFStorageView(storage, modname=self.__class__.__name__)

        # Initialize storage for target probabilities.
        self.log_target_probabilities = dict()
        self.verbose = verbose
        self.iteration = 0

    @property
    def state_keys(self):
        return self.log_target_probabilities.keys()

    def update_samplers(self):
        """
        Update all samplers.
        """
        for sampler in self.samplers:
            sampler.update()

    def update_target_probabilities(self):
        """
        Update all target probabilities.
        """
        # Gather list of all keys.
        state_keys = set()
        for sampler in self.samplers:
            for key in sampler.state_keys:
                state_keys.add(key)

        # Compute unnormalized log target probabilities.
        log_target_probabilities = { key : 0.0 for key in state_keys }
        for (sampler, log_weight) in self.sampler_exponents.items():
            for key in sampler.state_keys:
                log_target_probabilities[key] += log_weight * sampler.logZ[key]

        # Normalize
        log_sum = log_sum_exp(log_target_probabilities)
        for key in log_target_probabilities:
            log_target_probabilities[key] -= log_sum

        # Store.
        self.log_target_probabilities = log_target_probabilities

        if self.verbose:
            print("log_target_probabilities = %s" % str(self.log_target_probabilities))

        if self.storage:
            self.storage.write_object('log_target_probabilities', self.log_target_probabilities, iteration=self.iteration)

    def update(self):
        """
        Run one iteration of the sampler.
        """
        if self.verbose:
            print("*" * 80)
            print("MultiTargetDesign sampler iteration %8d" % self.iteration)
        self.update_samplers()
        self.update_target_probabilities()
        self.iteration += 1
        if self.storage: self.storage.sync()
        if self.verbose:
            print("*" * 80)

    def run(self, niterations=1):
        """
        Run the multi-target design sampler for the specified number of iterations.

        Parameters
        ----------
        niterations : int
            The number of iterations to run the sampler for.

        """
        # Update all samplers.
        for iteration in range(niterations):
            self.update()

################################################################################
# CONSTANT PH SAMPLER
################################################################################

class ProtonationStateSampler(object):
    """
    Protonation state sampler with given fixed target probabilities for ligand in solvent.

    Parameters
    ----------
    samplers : list of SAMSSampler
        The SAMS samplers whose relative partition functions go into the design objective computation.
    sampler_exponents : dict of SAMSSampler : float
        samplers.keys() are the samplers, and samplers[key]
    log_target_probabilities : dict of hashable object : float
        log_target_probabilities[key] is the computed log objective function (target probability) for chemical state `key`
    verbose : bool
        If True, verbose output is printed.

    """
    def __init__(self, complex_sampler, solvent_sampler, log_state_penalties, storage=None, verbose=False):
        """
        Initialize a protonation state sampler with fixed target probabilities for ligand in solvent.

        Parameters
        ----------
        complex_sampler : ExpandedEnsembleSampler
            Ligand in complex sampler
        solvent_sampler : SAMSSampler
            Ligand in solution sampler
        log_state_penalties : dict
            log_state_penalties[smiles] is the log state free energy (in kT) for ligand state 'smiles'
        storage : NetCDFStorage, optional, default=None
            If specified, will use the storage layer to write trajectory data.
        verbose : bool, optional, default=False
            If true, will print verbose output

        """
        # Store target samplers.
        self.log_state_penalties = log_state_penalties
        self.samplers = [complex_sampler, solvent_sampler]
        self.complex_sampler = complex_sampler
        self.solvent_sampler = solvent_sampler

        self.storage = None
        if storage is not None:
            self.storage = NetCDFStorageView(storage, modname=self.__class__.__name__)

        # Initialize storage for target probabilities.
        self.log_target_probabilities = { key : - log_state_penalties[key] for key in log_state_penalties }
        self.verbose = verbose
        self.iteration = 0

    @property
    def state_keys(self):
        return self.log_target_probabilities.keys()

    def update_samplers(self):
        """
        Update all samplers.
        """
        for sampler in self.samplers:
            sampler.update()

    def update_target_probabilities(self):
        """
        Update all target probabilities.
        """
        # Update the complex sampler log weights using the solvent sampler log weights
        for key in self.solvent_sampler.state_keys:
            self.complex_sampler.log_weights[key] = self.solvent_sampler.sampler.log_weights[key]

        if self.verbose:
            print("log_weights = %s" % str(self.solvent_sampler.sampler.log_weights))

    def update(self):
        """
        Run one iteration of the sampler.
        """
        if self.verbose:
            print("*" * 80)
            print("ProtonationStateSampler iteration %8d" % self.iteration)
        self.update_samplers()
        self.update_target_probabilities()
        if self.storage: self.storage.sync()
        self.iteration += 1
        if self.verbose:
            print("*" * 80)

    def run(self, niterations=1):
        """
        Run the protonation state sampler for the specified number of iterations.

        Parameters
        ----------
        niterations : int
            The number of iterations to run the sampler for.

        """
        # Update all samplers.
        for iteration in range(niterations):
            self.update()