#!/usr/bin/env python3
# class Worm sets PD['GJ_scale']=.02
#	repr prints it (both still the same)
# setup_Bitsey_network copies PD['GJ_scale'] to GJ_connects['scale']
#	now correctly done differently
# worm_lab1 sets GJ_density=.1, then W.PD['GJ_scale']=GJ_density (still OK)
#	then backs off sim.GJ_connects['scale'] /= 10
#	then in each loop, sim.GJ_connects['scale'] *= 1.1
#	and prints sim.GJ_connects['scale'][0]

# Copyright 2018 Alexis Pietak and Joel Grodstein
# See "LICENSE" for further details.

#    import pdb; pdb.set_trace()
import numpy as np
import sim as sim
import eplot as eplt
import edebug as edb
import copy

SIM_TIME = 20000000
GRADIENT = .030		# How much delta-Vmem counts as growing a gradient.
GD_FVAL=1		# Generation/decay -> [M] eventually becomes 1.

# What makes worms different? It's all in Worm.PD, the parameter dictionary:
# - num_cells, GJ_scale, ...: these affect how the network is built
# - kM, N, ...: these affect how ion-channel gating arrays are populated.
# - diff_IC_K, ...: these affect global arrays (Dm_array, GJ_connects, ...)
# - GJ_len: affects the global parameter object GP.
class Worm:
  def __init__(self):
    self._fitness = -1	 # Save the fitness value to avoid recomputes
    self._n_good_children = -1 # Part of fitness; how many child pieces lived.

    # Now set the default parameters (which match the Worm regression tests).
    # Any instance typically overrules some/most of these.
    self.PD = {}	  # Dictionary of the parameters for this Worm object.
    self.PD['num_cells'] = 5	# the obvious
    self.PD['GJ_scale'] = .02	# controls how strong the GJs are
    self.PD['GJ_len'] = 15e-9	# GJ length
    self.PD['kM'] = .2
    self.PD['N'] = 5		# for the ion-channel Hill model cooperativity
    self.PD['diff_IC_K'] =1.7e-17	# ion-channel diff const for K
                                        #    (i.e., sim.Dm_array[K,:])
    self.PD['diff_GJ_Na'] =1.33e-17	# GJ diff const
    self.PD['diff_GJ_K'] = 1.96e-17	#    (i.e., sim.GJ_diffusion[])
    self.PD['diff_GJ_M'] = 1e-14	#    for Na, K and M
    self.PD['env_Na'] = 145		# Environmental [Na], [K] and [P]; then
    self.PD['env_K'] = 5		# choose [Cl] automatically for charge
    self.PD['env_P'] = 10		# neutrality.
    self.PD['int_Na'] = 12		# Ditto, but cell-internal. Note these
    self.PD['int_K'] = 139		#    are just initial values, and may
    self.PD['int_P'] = 135		#    change over time.
    self.PD['M_valence'] =-1		# valence of the mystery ion M
    self.PD['GD_tau'] = 100		# time constant for gen/decay of M
    self.PD['stronger_at_end'] =1	# ion channels weaker in the body
    self.PD['scale_pumps'] = 0		# ion pumps commensurately weaker also

  # The official Python print() calls this to print objects.
  def __repr__ (self):
      child = " ("+str(self._n_good_children)+" / 20)" \
              if (self._n_good_children>=0) else ""
      return ("[Fitness={:.2f}{}: shape={}, {}{} {} cells; GJ_scale={:.3f}, kM={:.3f}, N={}, qM={}, GD_tau={:.3f}, SAE={}({}), head/tail pres={}/{}]".\
              format (self._fitness, child, self.shape(),
                      "F" if self.PD['full_body'] else "E",
                      "G" if self.PD['do_gen_decay'] else "",
                      self.PD['num_cells'], self.PD['GJ_scale'], self.PD['kM'],
                      self.PD['N'], self.PD['M_valence'], self.PD['GD_tau'],
                      self.PD['stronger_at_end'],
                      ("pumps too" if self.PD['scale_pumps'] else "normal pumps"),
                      self.PD['VNChead_pres'], self.PD['VNCtail_pres']))

  # Use my set of parameters (i.e., self.PD[...]) and create a Bitsey
  # simulation network.
  def setup_Bitsey_network(self):
    self._fitness = -1; self._n_good_children=0
    p=sim.Params()
    p.sim_dump_interval = 20000	# Not much debug printout during the sim
    num_cells = self.PD['num_cells']
    n_GJs = num_cells-1

    sim.init_big_arrays (num_cells, n_GJs, p, ['M']) # sim.GP gets set here
    Na=sim.ion_i['Na']; K=sim.ion_i['K']; Cl=sim.ion_i['Cl']
    P=sim.ion_i['P']; M=sim.ion_i['M']

    # Alexis uses feedback on every cell; usually I only do it at the head/tail.
    # Cells without feedback will not have any ion channels at all.
    fb_cells = range(num_cells) if self.PD['full_body'] else [0,num_cells-1]

    # We let evolution (and anyone else, for that matter) change the area of
    # K ion channels, but not of Na or Cl ion channels. This doesn't necessarily
    # make sense biologically, but it was done since we've typically built
    # positive feedback with K-ion-channel gates and not Na ones.
    # Default ion-channel diffusion coef: Na=K=M = 1e-18 m2/s, and P=0.

    # Set the K ion-channel diffusion rate (really, the cross-sec area).
    # If we're not in full-body mode, then kill *all* body-interior ICs (not
    # just K).
    sim.Dm_array[K, fb_cells] = self.PD['diff_IC_K']	# Set the K channels
    if (not self.PD['full_body']):			# Kill all interior
        sim.Dm_array[:,1:-1]= 0				# ICs if not full-body.

    # The "end" cells may have stronger ion channels. If so, then define "end"
    # as the 10% of cells on either end (but at least one cell if num_cells=5!).
    end = max (round(num_cells/10), 1)
    SAE = self.PD['stronger_at_end']
    sim.Dm_array[:,end:-end] = sim.Dm_array[:,end:-end] / SAE
    if self.PD['scale_pumps']:
        sim.pump_scaling[end:-end] = 1.0/SAE

    # Ligand-gated K channels at feedback cells (head, tail and maybe body).
    # First initialize to all off, then turn on the feedback cells.
    # Big picture (remember that K has Vnernst<0): the head is positive, which
    # attracts M, which turns off the K gate, which makes the head still more
    # positive. And the tail is negative, which repels M, which turns on the K
    # gate and makes the tail more negative.
    g=sim.Hill_gate (sim.Gate.GATE_IC, out_ion=K, in_ion=M)
    for c in fb_cells:
        g.change_Hill_gate (c, inv=True, kM=self.PD['kM'], N=self.PD['N'])

    # Dummy ligand-gated Na Hill channels everywhere, that really just serve to
    # allow an offset for VNC head pressure (tail pressure is through the K
    # channels just above).
    sim.Hill_gate (sim.Gate.GATE_IC, out_ion=Na, in_ion=Na)

    # Straight-line GJ connections along the worm. All are simple & ungated.
    sim.GJ_connects['from']  = range(n_GJs)
    sim.GJ_connects['to']    = sim.GJ_connects['from'] + 1

    # Physical parameters of the 'mystery' ion.
    sim.z_array[M] = self.PD['M_valence']	# fixed ion valence

    # Initial concentrations: external {Na=145 mM, K=5mM, P=10mM}
    sim.cc_env[Na] = self.PD['env_Na']; sim.cc_env[K]=self.PD['env_K']
    sim.cc_env[P]=self.PD['env_P']
    # Adjust [Cl]ext to reach charge neutrality
    sim.cc_env[Cl] = self.PD['env_Na'] + self.PD['env_K'] - self.PD['env_P']

    # Initial concentrations: internal {Na=12, K=139, Cl=16, P=135}
    sim.cc_cells[Na]=self.PD['int_Na']; sim.cc_cells[K]=self.PD['int_K']
    sim.cc_cells[P] =self.PD['int_P']
    # Adjust [Cl]int to reach charge neutrality
    sim.cc_cells[Cl]=self.PD['int_Na'] + self.PD['int_K'] - self.PD['int_P']

    # Ligand creation/decay/gating: creation at .1uM/s, and decaying at
    # (.1/s)*(its concentration)). This yields concentration=1uM. However, we'll
    # spread it from head to tail to seed the gradient. Note that 1uM is 1
    # umole/liter, or 1 mmole/m3. We're actually doing 1000 times that big!
    if (self.PD['do_gen_decay']):
        gen_rate = GD_FVAL/self.PD['GD_tau']	# moles/m3 per s
        sim.GD_const_gate (sim.Gate.GATE_GD, M, gen_rate, 1/self.PD['GD_tau'])

    # Change the gap-junction length from 100nm to 15nm.
    p.GJ_len = self.PD['GJ_len']

    # Assign the GJ diffusion consants however this worm has been set up.
    sim.GJ_diffusion[Na,:] = self.PD['diff_GJ_Na']
    sim.GJ_diffusion[K,:] = self.PD['diff_GJ_K']
    sim.GJ_diffusion[M,:] = self.PD['diff_GJ_M']
    sim.GJ_diffusion[P,:] = 0
    sim.GJ_diffusion *= self.PD['GJ_scale']

    #Na/K-ATPase pump with a maximum rate of 1.0x10-7 mol/(m^2 s)
    #print ("Worm setup is done.")

  # Simulate and return any of three statuses:
  # True  => successful simulation and built a gradient
  # False => successful simulation but no gradient
  # None  => simulation did not complete.
  # Note that both False and None act as False in an if() statement.
  # The big side effect is that the cc_cells and Vm arrays get filled in;
  # we will check them to find our fitness and our body shape.
  def sim (self, end_time):
      sim.GP.use_implicit = True
      sim.GP.adaptive_timestep = False

      try:
          self.t_shots, self.cc_shots = sim.sim (end_time)
      except ValueError:
          print ('Simulation: valueError encountered!!!')
          return (None)

      #np.set_printoptions (formatter={'float': "{:.3f}".format})
      plots=False
      if (plots):
          eplt.plot_Vmem (self.t_shots, self.cc_shots)
          #eplt.plot_ion (self.t_shots, self.cc_shots, 'Na')
          #eplt.plot_ion (self.t_shots, self.cc_shots, 'K')
          #eplt.plot_ion (self.t_shots, self.cc_shots, 'Cl')
          eplt.plot_ion (self.t_shots, self.cc_shots, 'M')
      print ('    Vm={}mV\n    [M]={} => {}'.format (
		sim.compute_Vm(sim.cc_cells)*1000,self.cc_shots[-1][4],
                self.shape()))
      grad = self.gradient()
      print ('Simulation: gradient={:.1f}mV => {} gradient'.format (grad*1000,
                                          "good" if (grad>=GRADIENT) else "no"))
      return (grad >= GRADIENT)

  # Return the Vmem gradient -- i.e., (max Vmem) minus (min Vmem).
  # Assume somebody has already run a simulation.
  def gradient(self, only_TH=False):
      Vm = sim.compute_Vm (sim.cc_cells)
      if (only_TH):
          return (max (Vm[-1] - Vm[0], 0.0))
      return (Vm.max() - Vm.min())

  # Look at the profile of [M] and deduce where the worm's heads and tails are.
  # Return a string such as "TH" (hopefully), or "HTH" (for a two-headed worm).
  # Note that Vmem is irrelevant; we only look at [M].
  # The general idea is we walk the cells from left to right; we end a rising
  # segment when we see a new [M] that is at least DELTA smaller than the
  # largest value on the segment (which is usually but not always the [M] at
  # the segment's start). A falling segment is similar.
  def shape(self, sim_results=True):
    if (sim_results==None):
        return ("DNF")		# Sim did not finish
    if (sim_results==False):
        return ("NoGrd")	# No gradient of Vmem

    DELTA=.02	# Ignore any changes in [M] that are less then DELTA.
    shape = ''	# Start with an empty string and add to it as we go.
    mode = 'X'	# Is the current segment falling(F), rising(R) or not set(X)?
    M = sim.ion_i['M']
    mMin = sim.cc_cells[M,0]	# Running min & max [M] on the current segment.
    mMax = sim.cc_cells[M,0]
    n_cells=sim.cc_cells.shape[1]
    for cell in range (1,n_cells):	# Walk from left to right, and
        mNew = sim.cc_cells[M,cell]	# analyze the current cell's [M]
        # If we're walking a rising segment, then we've already outputted the
        # "H" corresponding to the end of the segment that we've not reached
        # yet, but that has to come at some point. (Startup is a special
        # case... see below).

        # print ('M[{}]={}...'.format(cell,mNew))
        # Rising segment ends because we see a fall in [M]. This means that the
        # segment's H (which we've already outputted) is the cell immediately
        # to our left. We output a T, for the tail at the end of the segment we
        # now start (we'll find the actual tail cell eventually).
        if (mode=='R') and (mNew < mMax-DELTA):
            mode='F'; shape=shape+'T'
            mMin=mNew; mMax=mNew
            # print ('R->F, output T')

        if (mode=='F') and (mNew > mMin+DELTA):	# Ditto, for a falling segment
            mode='R'; shape=shape+'H'
            mMin=mNew; mMax=mNew
            # print ('F->R, output H')

        # Two special cases for the start of a worm. When we see the initial
        # rise in [M], we output both the T (for cell #0), and also the H (for
        # the eventual head that we haven't seen yet).
        # Note that we don't decide whether this initial segment is rising or
        # falling until we see >DELTA of [M] change... which may take a few
        # cells. So [M]= [.999 .987 1.012  0.999 1.001] will be TH.
        if (mode=='X') and (mNew>mMin+DELTA):
            mode='R'; shape=shape+'TH'
            mMin=mNew; mMax=mNew
            # print ('Initial TH, mode=R')

        if (mode=='X') and (mNew<mMax-DELTA):	# Ditto for an initial
            mode='F'; shape=shape+'HT'		# falling segment.
            mMin=mNew; mMax=mNew
            # print ('Initial HT, mode=F')

        mMin = min (mMin, mNew)	# Update the running min & max [M]
        mMax = max (mMax, mNew)	# along the current segment.

    return (shape)

# This function produces versions of cc_init and Vm_init for spread_M().
# An initial simulation may come up with a 2H Vmem, or have a normal Vmem but
# a 2H [M] profile (see worm_M_hooks.txt for an example). Sanitize() changes
# that sim result into a nice monotonic cc_init and Vm_init.
# Note our algorithm:
# - Note the min and max values of Vmem, [M] and [Na] across the entire worm.
# - Stuff the min(max) value of Vmem at the tail(head); i.e., Vm_init[0] (-1).
#   Linearly interpolate Vmem in the other cells. Stuff/interpolate [M] and
#   [Na] similarly.
# - The only information we really provide to spread_M() is the Na and M rows;
#   spread_M() sets the Cl and K rows itself, as follows.
# - Always let [Cl] be dictated by Vmem (i.e., the Nernst voltage). Physics
#   demands this.
# - Then pick [K] in each cell to be whatever it needs to be to come up with
#   the desired Vmem. Typically, cell-internal [Na] and [M] are much smaller
#   than [K], and so any sanitizing we do to [Na] and [M], while making
#   relatively large changes in *those* concentrations, will really not affect
#   [K] very much. I.e., whatever we set [K] to is going to be pretty much what
#   it started as anyway.
# - Set [P] to whatever the sim said -- which had better also be the value in
#   self.PD['int_P']. We bother to do this because 
def sanitize (cc_cells):
    np.set_printoptions (formatter={'float': "{:.3f}".format})
    print ("Sanitizing...")
    global cc_init, Vm_init
    cc_init = cc_cells.copy()
    Vm_init = sim.compute_Vm (cc_cells).copy()

    Na=sim.ion_i['Na']; M=sim.ion_i['M']

    # If we're already nice and monotonic, then we're all done.
    if (np.all ((np.diff(Vm_init)>=0) & (np.diff(cc_cells[M])>=0) \
				      & (np.diff(cc_cells[Na])<=0))):
        return

    # Otherwise, just do a linear spread
    n_cells = Vm_init.size
    Vm_init    = np.linspace (Vm_init.min(), Vm_init.max(), n_cells)
    cc_init[M] = np.linspace (cc_cells[M].min(),  cc_cells[M].max(),  n_cells)
    cc_init[Na]= np.linspace (cc_cells[Na].max(), cc_cells[Na].min(), n_cells)
    #print ("Sanitized => ", cc_init)
    assert (np.all(np.diff(cc_init[Na])<=0) and np.all(np.diff(cc_init[M])>=0))

# These help spread() initialize cc_cells. They're just the end results from
# the test1 benchmark; once our own sim passes someone usually calls sanitize()
# to use those numbers instead.
def set_default_spread_M():
    global cc_init, Vm_init
    cc_init = np.array ([[8.1, 7.2, 5.3, 4.5, 4.2], [0,0,0,0,0], [0,0,0,0,0],
        		 [0,0,0,0,0], [.17, .25, .7, 1.4, 2.4]])
    Vm_init = np.array ([-64, 9, 28, 34, 36])/1000

# Spread [M], [Na], [K], [Cl] according to the pattern string 'shape'. Roughly:
# - "TH"  would set cc_cells[M,0:5] to cc_init[M,0:5]
# - "THT" would set cc_cells[M,0:5] to cc_init[M,[0,2,4,2,0]]
# I.e., we're assuming that the [M]-vs-cell_number profile in cc_init[] is
# correct, and using whatever piece(s) of it we need to build the desired shape.
# (Arguably, this is a bad assumption if cc_init came from a 2H worm).
# Specifically, for each worm cell w, we:
# - use 'shape' to map worm cell[w] into a corresponding cc_init index i.
# - set cc_cells[M,w]  = cc_init[M,i] (where i might be a fraction)
# - set cc_cells[Na,w] = cc_init[Na,i]
# - set cc_cells[Cl,w] from Vm_init[i] (since [Cl] is at equilibrium)
# - set cc_cells[K,w] to the correct charge so that the cell is at Vm_init[i]
# See the comment above sanitize() for our global strategy.
# The parameters piece,n_pieces are mostly used by fitness(); if, e.g., piece=1
# and n_pieces=10 then you will build an entire worm that has the very narrow
# range of [M] that would typically be found near the tail (i.e., near cell 0).
def spread_M (shape, piece=0, n_pieces=1):
    global cc_init, Vm_init

    # Note num_cellsI (for cc_init) vs. num_cellsW (for the worm). We can use
    # a 5-cell cc_init to initialize a 30-cell worm! (Only for the first such
    # sim; after which we push the 30-cell sim results into cc_init).
    num_cellsI = cc_init.shape[1]
    Na=sim.ion_i['Na']; M=sim.ion_i['M']; K=sim.ion_i['K']; Cl=sim.ion_i['Cl']

    # Usually, a "T" will mean cell #0 and "H"=cell #4. But if we're doing,
    # e.g., piece #1/5 then we'll only spread from T=1 to H=2. Worm.fitness()
    # uses this capability to create fragments.
    tail = (num_cellsI-1) * (piece/n_pieces)
    head = tail + (num_cellsI-1)/n_pieces

    # Indices[w] is the (potentially fractional) cc_init index for worm[w]
    # So worm_cell[w] will take its [M] from cc_init[M, indices[w]] (and
    # similar for Na and Vmem).
    num_cellsW = sim.cc_cells.shape[1]
    indices = np.empty(num_cellsW)

    # Fill in indices[]. For, e.g., THTH, we take each of the three segments
    # TH, HT and TH from it, and fill in the corresponding part of indices[].
    for pos in range (len(shape)-1):	# One letter at a time...
        # This single segment of the pattern will go into worm cells[w0,w1].
        w0 = round (pos    *(num_cellsW-1)/(len(shape)-1))
        w1 = round ((pos+1)*(num_cellsW-1)/(len(shape)-1))
        # Are we putting a TH pattern into cells[w0,w1), or a HT?
        seg = shape[pos:pos+2]
        assert ((seg=="TH") or (seg=="HT"))
        i0 = tail if (seg=="TH") else head	# cc_init idx for worm cell[w0]
        i1 = head if (seg=="TH") else tail	# cc_init idx for worm cell[w1]
        spread = np.linspace (i0, i1, w1-w0+1)	# Spreads this entire segment.
        indices[w0:w1+1] = spread

    #print ("For shape {} piece {}/{}, spreading to indices {}"\
    #		.format (shape,piece,n_pieces,indices))

    # Now set sim.cc_cells accordingly.
    sim.cc_cells[Na] = np.interp (indices, range(num_cellsI), cc_init[Na])
    sim.cc_cells[M]  = np.interp (indices, range(num_cellsI), cc_init[M])
    Vm_target = np.interp (indices, range(num_cellsI), Vm_init)
    #print ("Spread: desired final Vmem=", Vm_target)

    # Compute [Cl]_int. That's easy -- just make Vnernst match Vmem.
    # Vn = 26mV * ln(Cint / Cext), so Cint = Cext * exp(Vn / 26mv)
    k26mV = sim.GP.R * sim.GP.T / sim.GP.F	# a.k.a. kT/q
    sim.cc_cells[Cl] = sim.cc_env[Cl] * np.exp (Vm_target / k26mV)

    # Pick [K] so that the desired Vmem is correct. We could do the
    # simple physics. But interpolation is even easier.
    sim.cc_cells[K]=0
    K0 = sim.compute_Vm(sim.cc_cells)
    sim.cc_cells[K]=10
    K10 = sim.compute_Vm(sim.cc_cells)
    # Simple interpolation: v = k0 + (k10-k0)*k/10, or k = (v-k0)*10/(k10-k0)
    sim.cc_cells[K] = (Vm_target-K0)*10/(K10-K0)

    np.set_printoptions (formatter={'float': "{:.3f}".format})
    print ("Spread {} set Vm(V)={}mV\n\tand M = {}".format (shape,
		sim.compute_Vm(sim.cc_cells)*1000, sim.cc_cells[M]))

############################################################
# Basic code to help simulate worms.
############################################################

# This is typically the first simulation we do for any parameter set.
# We do a simple TH simulation from a default [M] spread. If the sim completes
# and builds a gradient, then great -- use the results to call sanitize() and
# tailor spread_M() for *this* parameter set. Then return True.
# But if the sim dies or does not create a gradient, then just return False.
def sim_start(W):
    print ("Starting simulation for a new worm")
    set_VNC_WT (W)				# Assume a normal WT VNC

    # See if TH is stable. If not, then return -- nothing else is worth checking
    set_default_spread_M()
    spread_M ("TH")
    sim_results = W.sim (SIM_TIME)
    if (sim_results):
        # We converged. Save our final cc_cells[] to help init future worms
        print ("Initial simulation converged!")
        sanitize (sim.cc_cells)
        return (True)

    if (sim_results==None):
        print ("Initial simulation did not complete")
    if (sim_results==False):
        print ("Initial simulation did not produce enough gradient")
    return (False)

# Implement the VNC pressure on Vmem. The VNC's "head"("tail") tries to
# de(hyper)polarize Vmem. The parameters:
#	'head' is True(False) if we're working with the VNC's head(tail).
#	'cell' is which cell to put pressure on.
# For a normal WT worm, we would be called as (False, 0) and also as
# (True, num_cells-1). However, e.g., a HTH worm would call us as (True, 0),
# (False, num_cells/2) and (True, num_cells-1).
def set_VNC_pres (W, amount, head, cell):
    # What range of cells should we touch? In principle, 10% of the worm's
    # length, centered on 'cell'.
    n_cells = max (round(W.PD['num_cells']/10), 1)	# num of cells to touch
    n_cells = max (round(W.PD['num_cells']/5), 1)	# num of cells to touch
    if (cell < 0):			# in case we're given
        cell = cell + W.PD['num_cells']	# cell=-1 to mean the head.
    c0 = max (round(cell-n_cells/2), 0, cell+1-n_cells)	# [c0,c1) is the range
    c1 = min (c0 + n_cells, W.PD['num_cells'])		# of cells to touch.

    # VNC pushes the head(tail) more positive(negative), so we implement
    # head(tail) pressure by increasing Na(K), whose Vnernst is
    # positive(negative).
    # Assume that setup_Bitsey_network() built K and Na ion-channel gates as the
    # first two ion-channel gates. We'll tweek their offset fields.
    Na=sim.ion_i['Na']; K=sim.ion_i['K']
    OFFSET=4
    if (head):
        print ("\tSet Na IC_gate[{}:{}) offset to head pres={}".format(c0,c1,amount))
        gate = sim.IC_gates[1]	# Assume the Na gate was declared second
    else:
        print ("\tSet K  IC_gate[{}:{}) offset to tail pres={}".format(c0,c1,amount))
        gate = sim.IC_gates[0]	# Assume the K gate was declared first
    gate.params[OFFSET,c0:c1] = amount

def clear_VNC_pres (W):
    OFFSET=4
    K_gate = sim.IC_gates[0]
    K_gate.params[OFFSET,:]=0		# Turn off the K offset
    Na_gate = sim.IC_gates[1]
    Na_gate.params[OFFSET,:] = 0	# And ditto for Na

def set_VNC_WT (W):
    print ("Setting WT VNC pressure")
    clear_VNC_pres (W)			# Make sure the previous shape is gone!
    set_VNC_pres (W, W.PD['VNChead_pres'], head=True,  cell=-1)
    set_VNC_pres (W, W.PD['VNCtail_pres'], head=False, cell=0)

############################################################
# EVERYTHING ABOVE THIS IS COMMON INFRASTRUCTURE.
# NOW HERE'S THE FUNCTION THAT WE USE FOR THIS SIMULATION.
############################################################

def worm_lab1 ():
    # Build a 5-cell worm.
    num_cells = 5

    # The main parameter we'll modify.
    GJ_density=.1

    # We said the mystery morphagen M has negative charge.
    M_valence = -2

    #  A bunch of parameters we've not discussed.
    kM = .5; N = 2
    GD_tau = 1000
    stronger_at_end = 1	# stronger ion chans at the ends
    scale_pumps = 0	# reduce pump strength with SAE
    VNChead_pres = 0	# Extra D_Na at VNC head
    VNCtail_pres = 0 	# Extra D_K at VNC tail
    do_gen_decay = 1; do_full_body=1
    print ("Num_cells={}, GJ_density={}, M_val={}".format (num_cells, GJ_density, M_valence))

    # Create a worm object.
    W = Worm()
    print ("Starting initial worm setup...")

    # Set its parameters from what we decided above.
    W.PD['do_gen_decay'] = do_gen_decay; W.PD['full_body'] = do_full_body
    W.PD['num_cells']=num_cells; W.PD['GJ_scale']=GJ_density; W.PD['kM']=kM
    W.PD['N']=N; W.PD['M_valence']=M_valence; W.PD['GD_tau']=GD_tau
    W.PD['stronger_at_end']=stronger_at_end; W.PD['scale_pumps']=scale_pumps
    W.PD['VNChead_pres']=VNChead_pres; W.PD['VNCtail_pres']=VNCtail_pres

    # We have a good worm object; now build a Bitsey network.
    W.setup_Bitsey_network ()

    # Initial sim: Check it holds a gradient.
    print ("Starting initial simulation...")
    if (not sim_start (W)):
        return		# Failed sim or not enough gradient

    # Back off to a really low GJ density so we can make it higher.
    GJ_density /= 10
    sim.GJ_diffusion /= 10

    # Set [M] to a HTHTH pattern, which will then drive each cell to turn ion
    # channels on/off to make roughly that Vmem pattern.
    spread_M ("HTHTH")

    print ("\nStarting actual lab simulations...\n")
    if (W.sim (SIM_TIME) == None):	# Failed simulation
        return
    print ('At GJ density={:.4f}, deltaV={:.3f}mV and shape={}\n'.format (
           GJ_density, W.gradient()*1000, W.shape()))

    # Slowly increase the GJ density; resim each time.
    # Each sim starts at the previous sim's ending concentrations.
    # Stop when there is no more gradient.
    while (W.gradient() > .005):
        GJ_density *= 1.1
        sim.GJ_diffusion *= 1.1
        if (W.sim (SIM_TIME) == None):	# Failed simulation
            return
        print ('At GJ density={:.4f}, deltaV={:.3f}mV and shape={}\n'.format (
               GJ_density, W.gradient()*1000, W.shape()))

worm_lab1()