text: # Import all dependencies.
  # Note: Please import matplotlib before importing any psyneulink dependencies.
  import numpy as np
  
  from matplotlib import pyplot as plt
  # from scipy.special import erfinv # need to import this to make us of the UniformToNormalDist function.
  import psyneulink as pnl
  
  # --------------------------------- Global Variables ----------------------------------------
  # Now, we set the global variables, weights and initial values as in the paper.
  # WATCH OUT !!! In the paper the weight "Mutual inhibition among response units" is not defined, but needs to be set to
  # 0 in order to reproduce the paper.
  import psyneulink.core.components.functions.transferfunctions
  
  SD = 0.15       # noise determined by standard deviation (SD)
  a = 0.50        # Parameter describing shape of the FitzHugh–Nagumo cubic nullcline for the fast excitation variable v
  d = 0.5         # Uncorrelated Activity
  k = 1.5         # Scaling factor for transforming NE release (u ) to gain (g ) on potentiated units
  G = 0.5         # Base level of gain applied to decision and response units
  dt = 0.02       # time step size
  C = 0.90        # LC coherence (see Gilzenrat et al. (2002) on more details on LC coherence
  
  initial_hv = 0.07                     # Initial value for h(v)
  initial_w = 0.14                      # initial value u
  initial_v = (initial_hv - (1 - C) * d) / C  # get initial v from initial h(v)
  
  # Weights:
  inpwt = 1.5       # inpwt (Input to decision layer)
  crswt = 1 / 3       # crswt (Crosstalk input to decision layer)
  inhwt = 1.0       # inhwt (Mutual inhibition among decision units)
  respinhwt = 0     # respinhwt (Mutual inhibition among response units)  !!! WATCH OUT: this parameter is not mentioned
  # in the original paper, most likely since it was set 0
  decwt = 3.5       # decwt (Target decision unit to response unit)
  selfdwt = 2.5     # selfdwt (Self recurrent conn. for each decision unit)
  selfrwt = 2.0     # selfrwt (Self recurrent conn. for response unit)
  lcwt = 0.3        # lcwt (Target decision unit to LC)
  decbias = 1.75    # decbias (Bias input to decision units)
  respbias = 1.75   # respbias (Bias input to response units)
  tau_v = 0.05    # Time constant for fast LC excitation variable v
  tau_u = 5.00    # Time constant for slow LC recovery variable (‘NE release’) u
  trials = 1100   # number of trials to reproduce Figure 3 from Nieuwenhuis et al. (2005)
  
  # Setting seed (if noise is applied)
  # np.random.seed(22)    # Set noise with seed generator to 22 to use the UniformToNormalDist function to have a python
  # seed compatible with MATLAB random seed generator (rsg=22)
  # Please see https://github.com/jonasrauber/randn-matlab-python for further documentation
  
  # Create mechanisms ---------------------------------------------------------------------------------------------------
  
  # Input Layer --- [ Target 1, Target 2, Distractor ]
  
  # First, we create the 3 layers of the behavioral network, i.e. INPUT LAYER, DECISION LAYER, and RESPONSE LAYER.
  input_layer = pnl.TransferMechanism(
      size=3,                      # Number of units in input layer
      initial_value=[[0.0, 0.0, 0.0]],     # Initial input values
      name='INPUT LAYER'                  # Define the name of the layer; this is optional,
  )                 # but will help you to overview your model later on
  
  # Create Decision Layer  --- [ Target 1, Target 2, Distractor ]
  decision_layer = pnl.LCAMechanism(
      size=3,                            # Number of units in input layer
      initial_value=[[0.0, 0.0, 0.0]],    # Initial input values
      time_step_size=dt,                 # Integration step size
      leak=-1.0,                         # Sets off diagonals to negative values
      self_excitation=selfdwt,           # Set diagonals to self excitate
      competition=inhwt,                 # Set off diagonals to inhibit
      function=psyneulink.core.components.functions.transferfunctions.Logistic(x_0=decbias),   # Set the Logistic function with bias = decbias
      # noise=pnl.UniformToNormalDist(standard_deviation = SD).function, # The UniformToNormalDist function will
      integrator_mode=True,               # set the noise with a seed generator that is compatible with
      name='DECISION LAYER'               # MATLAB random seed generator 22 (rsg=22)
  )
  
  # decision_layer.set_log_conditions('RESULT')  # Log RESULT of the decision layer
  decision_layer.set_log_conditions('value')  # Log value of the decision layer
  
  for output_port in decision_layer.output_ports:
      output_port.value *= 0.0                                       # Set initial output values for decision layer to 0
  
  # Create Response Layer  --- [ Target1, Target2 ]
  response_layer = pnl.LCAMechanism(
      size=2,                                        # Number of units in input layer
      initial_value=[[0.0, 0.0]],                    # Initial input values
      time_step_size=dt,                             # Integration step size
      leak=-1.0,                                     # Sets off diagonals to negative values
      self_excitation=selfrwt,                       # Set diagonals to self excitate
      competition=respinhwt,                         # Set off diagonals to inhibit
      function=psyneulink.core.components.functions.transferfunctions.Logistic(x_0=respbias),          # Set the Logistic function with bias = decbias
      # noise=pnl.UniformToNormalDist(standard_deviation = SD).function,
      integrator_mode=True,
      name='RESPONSE LAYER'
  )
  
  response_layer.set_log_conditions('RESULT')     # Log RESULT of the response layer
  for output_port in response_layer.output_ports:
      output_port.value *= 0.0                   # Set initial output values for response layer to 0
  
  # Connect mechanisms --------------------------------------------------------------------------------------------------
  # Weight matrix from Input Layer --> Decision Layer
  input_weights = np.array([
      [inpwt, crswt, crswt],                    # Input weights are diagonals, cross weights are off diagonals
      [crswt, inpwt, crswt],
      [crswt, crswt, inpwt]
  ])
  
  # Weight matrix from Decision Layer --> Response Layer
  output_weights = np.array([
      [decwt, 0.0],  # Weight from T1 and T2 but not distractor unit (row 3 set to all zeros) to response layer
      [0.0, decwt],  # Need a 3 by 2 matrix, to project from decision layer with 3 units to response layer with 2 units
      [0.0, 0.0]
  ])
  
  decision_process = pnl.Process(
      pathway=[
          input_layer,
          input_weights,
          decision_layer,
          output_weights,
          response_layer
      ],
      name='DECISION PROCESS'
  )
  
  # Abstracted LC to modulate gain --------------------------------------------------------------------
  
  # This LCControlMechanism modulates gain.
  LC = pnl.LCControlMechanism(
      integration_method="EULER",       # We set the integration method to Euler like in the paper
      threshold_FitzHughNagumo=a,                  # Here we use the Euler method for integration and we want to set the parameters,
      uncorrelated_activity_FitzHughNagumo=d,      # for the FitzHugh–Nagumo system.
      time_step_size_FitzHughNagumo=dt,
      mode_FitzHughNagumo=C,
      time_constant_v_FitzHughNagumo=tau_v,
      time_constant_w_FitzHughNagumo=tau_u,
      a_v_FitzHughNagumo=-1.0,
      b_v_FitzHughNagumo=1.0,
      c_v_FitzHughNagumo=1.0,
      d_v_FitzHughNagumo=0.0,
      e_v_FitzHughNagumo=-1.0,
      f_v_FitzHughNagumo=1.0,
      a_w_FitzHughNagumo=1.0,
      b_w_FitzHughNagumo=-1.0,
      c_w_FitzHughNagumo=0.0,
      t_0_FitzHughNagumo=0.0,
      base_level_gain=G,                # Additionally, we set the parameters k and G to compute the gain equation
      scaling_factor_gain=k,
      initial_v_FitzHughNagumo=initial_v,          # Initialize v
      initial_w_FitzHughNagumo=initial_w,          # Initialize w
      objective_mechanism=pnl.ObjectiveMechanism(
          function=psyneulink.core.components.functions.transferfunctions.Linear,
          monitored_output_ports=[(
              decision_layer,  # Project output of T1 and T2 but not distractor from decision layer to LC
              np.array([[lcwt], [lcwt], [0.0]])
          )],
          name='Combine values'
      ),
      modulated_mechanisms=[decision_layer, response_layer],  # Modulate gain of decision & response layers
      name='LC'
  )
  
  # Log value of LC
  LC.set_log_conditions('value')
  
  # Set initial gain to G + k*initial_w, when the System runs the very first time,
  # since the decison layer executes before the LC and hence needs one initial gain value to start with.
  for output_port in LC.output_ports:
      output_port.value *= G + k * initial_w
  
  LC_process = pnl.Process(pathway=[LC])
  
  # Now, we specify the processes of the System, which in this case is just the decision_process
  task = pnl.System(processes=[decision_process, LC_process],
                    reinitialize_mechanisms_when=pnl.Never(),)
  
  # Create Stimulus -----------------------------------------------------------------------------------------------------
  
  # In the paper, each period has 100 time steps, so we will create 11 time periods.
  # As described in the paper in figure 3, during the first 3 time periods input to distractor units is fixed to 1.
  # Then T1 gets turned on during time period 4 with an input of 1.
  # T2 gets turns on with some lag from T1 onset on, in this example we turn T2 on with Lag 2 and an input of 1
  # Between T1 and T2 and after T2 the distractor unit is on.
  # We create one array with 3 numbers, one for each input unit and repeat this array 100 times for one time period
  # We do this 11 times. T1 is on for time4, T2 is on for time7 to model Lag3
  num_time_steps = 100  # Each stimulus is presented for two units of time which is equivalent to 100 time steps
  stimulus_T1 = np.repeat(np.array([[0, 0, 1]]), num_time_steps, axis=0)
  stimulus_T2 = np.repeat(np.array([[0, 0, 1]]), num_time_steps, axis=0)
  stimulus_T3 = np.repeat(np.array([[0, 0, 1]]), num_time_steps, axis=0)
  stimulus_T4 = np.repeat(np.array([[1, 0, 0]]), num_time_steps, axis=0)    # Turn T1 on
  stimulus_T5 = np.repeat(np.array([[0, 0, 1]]), num_time_steps, axis=0)
  stimulus_T6 = np.repeat(np.array([[0, 1, 0]]), num_time_steps, axis=0)    # Turn T2 on --> example for Lag 2
  stimulus_T7 = np.repeat(np.array([[0, 0, 1]]), num_time_steps, axis=0)
  stimulus_T8 = np.repeat(np.array([[0, 0, 1]]), num_time_steps, axis=0)
  stimulus_T9 = np.repeat(np.array([[0, 0, 1]]), num_time_steps, axis=0)
  stimulus_T10 = np.repeat(np.array([[0, 0, 1]]), num_time_steps, axis=0)
  stimulus_T11 = np.repeat(np.array([[0, 0, 1]]), num_time_steps, axis=0)
  
  # Concatenate the 11 arrays to one array with 1100 rows and 3 columns.
  time = np.concatenate((stimulus_T1, stimulus_T2, stimulus_T3, stimulus_T4, stimulus_T5, stimulus_T6,
                         stimulus_T7, stimulus_T8, stimulus_T9, stimulus_T10, stimulus_T11), axis=0)
  
  # assign inputs to input_layer (Origin Mechanism) for each trial
  stim_list_dict = {input_layer: time}
  
  # show the system
  task.show()
  # run the system
  task.run(stim_list_dict, num_trials=trials)
  
  
  # This displays a diagram of the System
  # task.show_graph()
  
  LC_results = LC.log.nparray()[1][1]        # get logged results
  LC_results_w = np.zeros([trials])          # get LC_results_w
  for i in range(trials):
      LC_results_w[i] = LC_results[4][i + 1][3][0][0]
  LC_results_v = np.zeros([trials])          # get LC_results_v
  for i in range(trials):
      LC_results_v[i] = LC_results[4][i + 1][2][0][0]
  
  
  def h_v(v, C, d):                   # Compute h(v)
      return C * v + (1 - C) * d
  
  
  LC_results_hv = np.zeros([trials])    # get LC_results_hv
  for i in range(trials):
      LC_results_hv[i] = h_v(LC_results_v[i], C, d)
  
  
  # Plot the Figure 3 from the paper
  t = np.linspace(0, trials, trials)            # Create array for x axis with same length then LC_results_v
  fig = plt.figure()                          # Instantiate figure
  ax = plt.gca()                              # Get current axis for plotting
  ax2 = ax.twinx()                            # Create twin axis with a different y-axis on the right side of the figure
  ax.plot(t, LC_results_hv, label="h(v)")      # Plot h(v)
  ax2.plot(t, LC_results_w, label="w", color='red')  # Plot w
  h1, l1 = ax.get_legend_handles_labels()
  h2, l2 = ax2.get_legend_handles_labels()
  ax.legend(h1 + h2, l1 + l2, loc=2)          # Create legend on one side
  ax.set_xlabel('Time (ms)')                  # Set x axis lable
  ax.set_ylabel('LC Activity')                # Set left y axis label
  ax2.set_ylabel('NE Output')                 # Set right y axis label
  plt.title('Nieuwenhuis 2005 PsyNeuLink Lag 2 without noise', fontweight='bold')  # Set title
  ax.set_ylim((-0.2, 1.0))                     # Set left y axis limits
  ax2.set_ylim((0.0, 0.4))                    # Set right y axis limits
  plt.show()
  