text: # -*- coding: utf-8 -*-
  """
  Created on Tue Apr  3 08:57:55 2018
  
  @author: garrettsmith
  
  Beginning of a module for general-use SOSP simulations.
  
  Will import dynamics functions from a dynamics module.
  Will import plotting fns. from a plotting module.
  
  Needs attributes of curr_state, state_hist, curr_harmony, harmony_hist,
  centers, center_labels, n_dim, dim_names, max_sent_length, n_words,
  prob. more.
  """
  
  
  import numpy as np
  import matplotlib.pyplot as plt
  import pandas as pd
  from scipy.optimize import minimize
  from .dynamics import iterate, euclid_stop, vel_stop, cheb_stop, calc_harmony
  
  
  class SimpleModel(object):
      def __init__(self, centers=None, local_harmonies=None,
                   stopping_crit='euclid_stop'):
  
          if centers is not None:
              # Not sure why I wanted this to be 2D... Leaving for now.
              assert len(centers.shape) > 1, 'centers must be a 2D array.'
              self.centers = centers
              self.n_dim = self.centers.shape[1]
          if local_harmonies is not None:
              assert len(local_harmonies) == self.centers.shape[0], \
               'Number of local harmonies does not match number of centers.'
              self.local_harmonies = local_harmonies
          if stopping_crit == 'vel_stop':
              self.stopping_crit = vel_stop
          elif stopping_crit == 'cheb_stop':
              self.stopping_crit = cheb_stop
          else:
              self.stopping_crit = euclid_stop
  
          self.tau = 0.01  # Time step for discretized dynamics
          self.max_time = 1000  # Max. number of time steps
          self._zero_state_hist()
          self.noise_mag = 0.001  # default
          self.gamma = [0.25] * centers.shape[1]  # default
          self.tol = 0.1  # Stopping tolerance on each dim.
  
      def _zero_state_hist(self):
          self.state_hist = np.zeros((self.max_time, self.n_dim))
  
      def set_noise_mag(self, noise_mag):
          self.noise_mag = noise_mag
  
      def set_gamma(self, gamma):
          self.gamma = gamma
  
      def set_tol(self, tol):
          self.tol = tol
  
      def set_max_time(self, max_time):
          self.max_time = max_time
  
      def set_local_harmonies(self, local_harmonies):
          self.local_harmonies = local_harmonies
  
      def set_params(self, **kwargs):
          for param, val in kwargs.items():
              setattr(self, param, val)
  
      def single_run(self, init_cond=None):
          """Run the model once until stopping criterion is met or 
          time runs out.
  
          Parameters
          ----------
          init_cond : array
              Point in state space at which to initialize the system
          
          Returns
          -------
          Nothing. Updates the state_hist attribute of the object.
          """
          assert init_cond is not None, 'Must set initial conditions.'
          assert len(init_cond) == self.state_hist.shape[1], \
              'Shape mismatch for initial conditions.'
          self.state_hist[0,] = init_cond
          noise = (np.sqrt(2 * self.noise_mag * self.tau)
                   * np.random.normal(0, 1, self.state_hist.shape))
          t = 0
          while t < self.max_time-1:
              if self.stopping_crit(self.state_hist[t], self.attrs, self.tol):
  #            if self.stopping_crit(self.state_hist[t], self.centers, self.tol):
                  self.state_hist[t+1,] = (self.state_hist[t,]
                          + self.tau * iterate(self.state_hist[t,], self.centers,
                                               self.local_harmonies, self.gamma)
                          + noise[t,])
                  t += 1
              else:
                  break
  
      def many_runs(self, n_runs=100, init_cond=None, verbose=False):
          """Do repeated Monte Carlo runs on a single harmony landscape.
  
          Parameters
          ----------
          n_runs : scalar
              The number of runs to do per condition
          init_cond : array
              Point in state space at which to initialize the system
  
          Returns
          -------
          A Pandas data frame with the center number and settling time.
          """
          #print('Run number:')
          data_list = []
          for run in range(n_runs):
              if run % 100 == 0 and verbose:
                  print('[{}] '.format(run), end='')
              self._zero_state_hist()
              self.single_run(init_cond)
              trunc = self.state_hist[~np.all(self.state_hist == 0, axis=1)]
              # Should probably use the actual attractors instead of the centers
              for center in range(self.centers.shape[0]):
                  if np.all(np.round(trunc[-1,]) == self.centers[center,]):
                      data_list.append([center, trunc.shape[0]])
                      break
              # Is executed after the for loop, but not if it found a center
              # and used the break statement
              else:
                  data_list.append([np.NaN, trunc.shape[0]])
          return pd.concat([pd.DataFrame([i], columns=('CenterNr', 'Time',))
                            for i in data_list])
  
      def run_multiple_conditions(self, n_runs=100, conditions=None,
                                  init_cond=None, gamma=None, D=None):
          """Do many runs of multiple conditions by changing the harmony
          landscape between conditions.
  
          Parameters
          ----------
          n_runs : scalar
              The number of runs to do per condition
          conditions : array
              An array with the local harmonies of each attractor in the
              columns. Each row is a different condition
          init_cond : array
              Point in state space at which to initialize the system
          gamma : array
              Width parameter for RBFs.
          D : array
              Noise magnitude.
  
          Returns
          -------
          all_data
              A pandas dataframe with columns for which center was chosen,
              how long the settling took, the condition number, gamma, and 
              D values.
          """
          if gamma is None:
              gamma = [self.gamma] * self.centers.shape[0]
          else:
              #gamma = [gamma] * self.centers.shape[0]
              gamma = np.tile(gamma, self.centers.shape[0]).T
          if D is None:
              D = [self.noise_mag]
          state_init = np.zeros((len(conditions), self.state_hist.shape[1]))
          if init_cond is not None:
              state_init = init_cond
          all_data = []
          for curr_g in gamma:
              for curr_D in D:
                  for cond in range(conditions.shape[0]):
                      #self.set_local_harmonies(conditions[cond,])
                      self.set_params(gamma=curr_g,
                                      local_harmonies=conditions[cond,],
                                      noise_mag=curr_D)
                      self.locate_attrs()
                      #print('Condition {}'.format(cond))
                      cond_data = self.many_runs(n_runs, state_init[cond,])
  #                   cond_data['Condition'] = [cond] * n_runs
                      cond_data['Condition'] = cond  # Make all rows have same value
                      cond_data['gamma'] = curr_g
                      cond_data['NoiseMag'] = curr_D
                      all_data.append(cond_data)
          return pd.concat(all_data)
  
      def neg_harmony(self, x, centers, local_harmonies, gamma):
          """Returns the negative of the harmony function.
          """
          return -1 * calc_harmony(x, centers, local_harmonies, gamma)
  
      def jac_neg_harmony(self, x, centers, local_harmonies, gamma):
          """Returns the negative of the dynamics for use as the Jacobian
          when finding the exact locations of the attractors.
          """
          return -1 * iterate(x, centers, local_harmonies, gamma)
  
      def locate_attrs(self):
          """Finds actual locations of attractors in the full harmony landscape
          using the L-BFGS-B algorithm on the negative of the harmony fn. Finds
          the unique attractors by rounding to two decimal places.
          """
          attrs = np.zeros(self.centers.shape)
          for c in range(self.centers.shape[0]):
              extremum = minimize(self.neg_harmony, self.centers[c],
                                  args=(self.centers, self.local_harmonies,
                                        self.gamma), method='L-BFGS-B',
                                  jac=self.jac_neg_harmony)
              attrs[c] = extremum.x
          unique_attrs = np.unique(np.round(attrs, 2), axis=0)
          self.attrs = unique_attrs
          print('Found {} unique attractors from {} centers'.format(
                  self.attrs.shape[0], self.centers.shape[0]))
          return
  
      def plot_trace(self):
          """Plots the trajectory of the latest run.
          """
          trunc = self.state_hist[~np.all(self.state_hist == 0, axis=1)]
          plt.plot(trunc)
          plt.ylim(-0.01, 1.01)
          plt.xlabel('Time')
          plt.ylabel('Activation')
          plt.show()
  
  
  if __name__ == '__main__':
      centers = np.array([[0., 1.], [1., 0.]])
      harmonies = np.array([1.0, 1.0])
      xinit = np.array([0., 0.])
      test = SimpleModel(centers, harmonies, 'cheb_stop')
      # test.set_tol(0.01)
      # test.single_run(xinit)
      # test.plot_trace()
      # results = test.many_runs(50, init_cond=xinit)
      # print(results.groupby('CenterNr').agg(['count', 'mean', 'std']))
  
      test.set_noise_mag(0.005)
      conds = np.array([[1.0, 1.0], [0.5, 1.0]])
      init = np.array([[0.5, 0.5], [0.5, 0.5]])
      many_results = test.run_multiple_conditions(n_runs=10, conditions=conds,
                                                  init_cond=init,
                                                  gamma=[0.15, 0.4],
                                                  D=[0.001, 0.005])
      print(many_results.groupby(['Condition', 'CenterNr']).agg(
              ['count', 'mean', 'std']))
      print(many_results)
  