from minimal_example_interface import *


def multiplicative_term(kappa_val,r_array,cue_array):
    assert(type(kappa_val)==float or type(kappa_val)==int)
    assert(type(r_array)==np.ndarray)
    assert(type(cue_array)==np.ndarray)
    
    return({"kappa": kappa_val, "r": r_array,"cue": cue_array})

set_term_function(multiplicative_term)

def SaltSugarModelWithCue_SingleTerm(term):
    return term['kappa']*term['r']*term['cue']

set_singleterm_calculation_function(SaltSugarModelWithCue_SingleTerm)

def SaltSugarModelWithCue(term_list):
    return sum([SaltSugarModelWithCue_SingleTerm(term) for term in term_list])

set_model_calculation_function(SaltSugarModelWithCue)



"""
Explanation: Rat model minimal example
Model objectives
This minimal example sets out to solve the following model objectives:

The rat prefers the moderate solution to the strong solution under salt depletion; it may approach the strong solution sometimes, but less often than the moderate solution
The rat prefers the moderate solution to the strong solution under salt satiation (i.e., finds it less aversive)
There will be some wanting for the aversive solution under sufficiently strong salt depletion
The rat will take the strong solution under salt depletion if it has no other option
Something representing reward signal will fire for a sugar solution but not a salt solution when there is a need for salt; however, when there is a need for salt, the reward signal will fire regardless

Objectives 1 and 2: multiple stimulus features
We can specify the basic model we need to model salt like:
\begin{align}\label{eq:BasicSaltModelZBVector}
\tilde r(\mathbf{r},\mathbf{\kappa})=& \mathbf{\kappa}{\text{Na}} \mathbf{r}{\text{Na}}+\mathbf{\kappa}{h} \mathbf{r}{h}
\end{align}
Objective 3: environmental cues
To meet Objective 3 above--to show the rat will drink the strong salty solution if it is sufficiently salt-depleted, though still prefer the moderate solution--we can assume a different weighting of hydration to deprivation. For instance, we can do that by increasing the deprivation/interoceptive state value $\kappa_{\text{Na}}$ from 1 to 1.5.
Objective 4: environmental cues
To meet Objective 4 above--to show the rat will drink the strong salty solution consistently if it is the only one available, we need to introduce an extra vector describing environmental cues that elicit responses, $\mathbf{c}$. Following the S-R-O expectancy valence account given in \textcite{zhang2009neural}, the extra parameter describes environmental features eliciting a response. 
Using the new model with our vector representing environmental cues $\mathbf{c}$, we can predict behavior where there is no cue for a moderate salt solution (Note the use of elementwise multiplication between the matrices using the symbol $\circ$):
$$
\begin{align}\label{eq:WithCues}
\tilde r(\mathbf{r},\mathbf{\kappa})= & \mathbf{\kappa}{\textrm{Na}} \mathbf{r}\textrm{Na}\mathbf{c}+\mathbf{\kappa}{h} \mathbf{r}{h}\mathbf{c} % \
% = & 1.5 \times \begin{bmatrix}0.5 \ 1.0\end{bmatrix}\circ \begin{bmatrix}0.0 \ 1.0\end{bmatrix} + 1 \times \begin{bmatrix}0.0 \ -1.0\end{bmatrix}\circ \begin{bmatrix}0.0 \ 1.0\end{bmatrix} \nonumber \
% = & \begin{bmatrix}0.0 \ 0.5\end{bmatrix} \nonumber
\end{align}
$$
Here, I've modeled $\mathbf{c}$ as a single vector without a subscript to indicate that we do not expect the cue value to systematically differ for each interoceptive value. This differs from the other documents produced--there, I'd included a subscript, more because I didn't think it through than any other reason.
Objective 5: Sugar
The model in Equation~\ref{eq:BasicSaltModelZBVector} can be extended slightly to account for the evidence discussed in the sugar example in Tindell et al. (2009):
\begin{equation}\label{eq:SaltSugarModelTindell}
\tilde r(\mathbf{r},\mathbf{\kappa})=\kappa_{\text{Na}} r_{\text{Na}}+\kappa_{h} r_{h} + \kappa_{\text{Glc}} r_{\text{Glc}}
\end{equation}
Combined Model
Combining the models presented above, we obtain the model 
\begin{align}
\begin{split}\label{eq:SaltSugarModelWithCue}
\tilde r(\mathbf{r},\mathbf{\kappa})= & \kappa_{\text{Na}} r_{\text{Na}}\mathbf{c}+\kappa_{h} r_{h}\mathbf{c} + \kappa_{\text{Glc}} r_{\text{Glc}}\mathbf{c}
\end{split}
\end{align}
We can generalize this further by describing an arbitrary number of cue-stimulus-interoception multiplicative terms:
\begin{align}
\tilde r(\mathbf{r},\mathbf{\kappa},\mathbf{c})= & \kappa_{1} \mathbf{r}{1}\mathbf{c}+\dots+\kappa{\text{N}} \mathbf{r}_{\text{N}}\mathbf{c}
\end{align}
We can then write a python function to model this equation, and then demonstrate how it handles each situation described above.
The function multiplicative_term forms a python dictionary representing a term with kappa value, r array, and cue array. The function SaltSugarModelWithCue is simply a python implementation of the $r(\mathbf{r},\mathbf{\kappa},\mathbf{c})$ function above. It calls SaltSugarModelWithCue_SingleTerm, which is written separately so that it can be used directly to show the working in the equation below.
End of explanation
"""


cue=np.array([1.0,1.0,0.0])
term_Na=multiplicative_term(1.5,np.array([1.0,1.0,0.0]),cue)
term_h=multiplicative_term(1,np.array([0.0,-1.0,0.0]),cue)
term_Glc=multiplicative_term(1,np.array([0.0,0.0,1.0]),cue)


"""
Explanation: We can then create variables to represent each term and then calculate their value:
End of explanation
"""


Markdown(get_and_format_main_equation_text(term_Na,term_h,term_Glc))

myb= widgets.Button(description="the fun button")
def add_text(b):
    txtArea.value = txtArea.value + txtArea.value


myb.on_click(add_text)

display(myb)

"""
Explanation: We can write this out mathematically:
End of explanation
"""


math_widget.value=get_and_format_main_equation_text(term_Na,term_h,term_Glc)

bDeprivationPreset

box_bordered=widgets.Layout(border='solid 1pt black')
#     def widget_header(value):
#         return widgets.Label(value='### ' + value)
    
reward_box=widgets.VBox([widgets.Label(value='Expected reward ($ \\tilde r$)'),
                                widgets.HBox(ftRewardVals)],layout=box_bordered)
cue_box=widgets.VBox([widgets.Label('Cue Accessibility ($\mathbf{c}$)')]+ftCueVals, layout=box_bordered)


"""
Explanation: You can play with the values yourself here or use the presets to see how values change.
End of explanation
"""


widgets.HBox(ftKappaVals)


widgets.VBox([
    widgets.HBox([bDeprivationPreset,bSatiationPreset,tbSugarPresence,bResetAll]),
    widgets.VBox([widgets.Label('Interoception ($\kappa$)'),
                 widgets.HBox(ftKappaVals)],layout=box_bordered),
    widgets.HBox([reward_box,cue_box]),
    bRecalculate,
    math_widget])

reward_box

display(reward_box)

widgets.HBox([bDeprivationPreset,bSatiationPreset,tbSugarPresence,bResetAll])

display(widgets.HBox([bDeprivationPreset,bSatiationPreset,tbSugarPresence,bResetAll]))

"""
Explanation: (Hint: Press the 'Recalculate' button to recalculate custom values you entered)
End of explanation
"""


