// Copyright (c) 2019, The Emergent Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

/*
Package chans provides standard neural conductance channels for computing
a point-neuron approximation based on the standard equivalent RC circuit
model of a neuron (i.e., basic Ohms law equations).
Includes excitatory, leak, inhibition, and dynamic potassium channels.
*/
package chans

//go:generate core generate -add-types

// Chans are ion channels used in computing point-neuron activation function
type Chans struct {

        // excitatory sodium (Na) AMPA channels activated by synaptic glutamate
        E float32

        // constant leak (potassium, K+) channels -- determines resting potential (typically higher than resting potential of K)
        L float32

        // inhibitory chloride (Cl-) channels activated by synaptic GABA
        I float32

        // gated / active potassium channels -- typically hyperpolarizing relative to leak / rest
        K float32
}

// SetAll sets all the values
func (ch *Chans) SetAll(e, l, i, k float32) {
        ch.E, ch.L, ch.I, ch.K = e, l, i, k
}

// SetFromOtherMinus sets all the values from other Chans minus given value
func (ch *Chans) SetFromOtherMinus(oth Chans, minus float32) {
        ch.E, ch.L, ch.I, ch.K = oth.E-minus, oth.L-minus, oth.I-minus, oth.K-minus
}

// SetFromMinusOther sets all the values from given value minus other Chans
func (ch *Chans) SetFromMinusOther(minus float32, oth Chans) {
        ch.E, ch.L, ch.I, ch.K = minus-oth.E, minus-oth.L, minus-oth.I, minus-oth.K
}

// Copyright (c) 2019, The Emergent Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// bench runs a benchmark model with 5 layers (3 hidden, Input, Output) all of the same
// size, for benchmarking different size networks.  These are not particularly realistic
// models for actual applications (e.g., large models tend to have much more topographic
// patterns of connectivity and larger layers with fewer connections), but they are
// easy to run..
package main

import (
        "flag"
        "fmt"
        "math"
        "math/rand"
        "os"
        "time"

        "cogentcore.org/core/base/timer"
        "cogentcore.org/lab/base/randx"
        "github.com/emer/emergent/v2/params"
        "github.com/emer/emergent/v2/patgen"
        "github.com/emer/emergent/v2/paths"
        "github.com/emer/etensor/tensor/table"
        "github.com/emer/leabra/v2/leabra"
)

var Net *leabra.Network
var Pats *table.Table
var EpcLog *table.Table
var Silent = false // non-verbose mode -- just reports result

var ParamSets = params.Sets{
        "Base": {
                {Sel: "Path", Desc: "norm and momentum on works better, but wt bal is not better for smaller nets",
                        Params: params.Params{
                                "Path.Learn.Norm.On":     "true",
                                "Path.Learn.Momentum.On": "true",
                                "Path.Learn.WtBal.On":    "false",
                        }},
                {Sel: "Layer", Desc: "using default 1.8 inhib for all of network -- can explore",
                        Params: params.Params{
                                "Layer.Inhib.Layer.Gi": "1.8",
                                "Layer.Act.Gbar.L":     "0.2", // original value -- makes HUGE diff on perf!
                        }},
                {Sel: "#Output", Desc: "output definitely needs lower inhib -- true for smaller layers in general",
                        Params: params.Params{
                                "Layer.Inhib.Layer.Gi": "1.4",
                        }},
                {Sel: ".Back", Desc: "top-down back-pathways MUST have lower relative weight scale, otherwise network hallucinates",
                        Params: params.Params{
                                "Path.WtScale.Rel": "0.2",
                        }},
        },
}

func ConfigNet(net *leabra.Network, units int) {
        squn := int(math.Sqrt(float64(units)))
        shp := []int{squn, squn}

        inLay := net.AddLayer("Input", shp, leabra.InputLayer)
        hid1Lay := net.AddLayer("Hidden1", shp, leabra.SuperLayer)
        hid2Lay := net.AddLayer("Hidden2", shp, leabra.SuperLayer)
        hid3Lay := net.AddLayer("Hidden3", shp, leabra.SuperLayer)
        outLay := net.AddLayer("Output", shp, leabra.TargetLayer)

        net.ConnectLayers(inLay, hid1Lay, paths.NewFull(), leabra.ForwardPath)
        net.ConnectLayers(hid1Lay, hid2Lay, paths.NewFull(), leabra.ForwardPath)
        net.ConnectLayers(hid2Lay, hid3Lay, paths.NewFull(), leabra.ForwardPath)
        net.ConnectLayers(hid3Lay, outLay, paths.NewFull(), leabra.ForwardPath)

        net.ConnectLayers(outLay, hid3Lay, paths.NewFull(), leabra.BackPath)
        net.ConnectLayers(hid3Lay, hid2Lay, paths.NewFull(), leabra.BackPath)
        net.ConnectLayers(hid2Lay, hid1Lay, paths.NewFull(), leabra.BackPath)

        net.Defaults()
        net.ApplyParams(ParamSets["Base"], false) // no msg
        net.Build()
        net.InitWeights()
}

func ConfigPats(dt *table.Table, pats, units int) {
        squn := int(math.Sqrt(float64(units)))
        shp := []int{squn, squn}
        // fmt.Printf("shape: %v\n", shp)

        dt.AddStringColumn("Name")
        dt.AddFloat32TensorColumn("Input", shp)
        dt.AddFloat32TensorColumn("Output", shp)
        dt.SetNumRows(pats)

        // note: actually can learn if activity is .15 instead of .25
        // but C++ benchmark is for .25..
        nOn := units / 6

        patgen.PermutedBinaryRows(dt.Columns[1], nOn, 1, 0)
        patgen.PermutedBinaryRows(dt.Columns[2], nOn, 1, 0)
}

func ConfigEpcLog(dt *table.Table) {
        dt.AddIntColumn("Epoch")
        dt.AddFloat32Column("CosDiff")
        dt.AddFloat32Column("AvgCosDiff")
        dt.AddFloat32Column("SSE")
        dt.AddFloat32Column("Avg SSE")
        dt.AddFloat32Column("Count Err")
        dt.AddFloat32Column("Pct Err")
        dt.AddFloat32Column("Pct Cor")
        dt.AddFloat32Column("Hid1 ActAvg")
        dt.AddFloat32Column("Hid2 ActAvg")
        dt.AddFloat32Column("Out ActAvg")
}

func TrainNet(net *leabra.Network, pats, epcLog *table.Table, epcs int) {
        ctx := leabra.NewContext()
        net.InitWeights()
        np := pats.NumRows()
        porder := rand.Perm(np) // randomly permuted order of ints

        epcLog.SetNumRows(epcs)

        inLay := net.LayerByName("Input")
        hid1Lay := net.LayerByName("Hidden1")
        hid2Lay := net.LayerByName("Hidden2")
        outLay := net.LayerByName("Output")

        _ = hid1Lay
        _ = hid2Lay

        inPats, _ := pats.ColumnByName("Input")
        outPats, _ := pats.ColumnByName("Output")

        tmr := timer.Time{}
        tmr.Start()
        for epc := 0; epc < epcs; epc++ {
                randx.PermuteInts(porder)
                outCosDiff := float32(0)
                cntErr := 0
                sse := 0.0
                avgSSE := 0.0
                for pi := 0; pi < np; pi++ {
                        ppi := porder[pi]
                        inp := inPats.SubSpace([]int{ppi})
                        outp := outPats.SubSpace([]int{ppi})

                        inLay.ApplyExt(inp)
                        outLay.ApplyExt(outp)

                        net.AlphaCycInit(true)
                        ctx.AlphaCycStart()
                        for qtr := 0; qtr < 4; qtr++ {
                                for cyc := 0; cyc < ctx.CycPerQtr; cyc++ {
                                        net.Cycle(ctx)
                                        ctx.CycleInc()
                                }
                                net.QuarterFinal(ctx)
                                ctx.QuarterInc()
                        }
                        net.DWt()
                        net.WtFromDWt()
                        outCosDiff += outLay.CosDiff.Cos
                        pSSE, pAvgSSE := outLay.MSE(0.5)
                        sse += pSSE
                        avgSSE += pAvgSSE
                        if pSSE != 0 {
                                cntErr++
                        }
                }
                outCosDiff /= float32(np)
                sse /= float64(np)
                avgSSE /= float64(np)
                pctErr := float64(cntErr) / float64(np)
                pctCor := 1 - pctErr
                // fmt.Printf("epc: %v  \tCosDiff: %v \tAvgCosDif: %v\n", epc, outCosDiff, outLay.CosDiff.Avg)
                epcLog.SetFloat("Epoch", epc, float64(epc))
                epcLog.SetFloat("CosDiff", epc, float64(outCosDiff))
                epcLog.SetFloat("AvgCosDiff", epc, float64(outLay.CosDiff.Avg))
                epcLog.SetFloat("SSE", epc, sse)
                epcLog.SetFloat("Avg SSE", epc, avgSSE)
                epcLog.SetFloat("Count Err", epc, float64(cntErr))
                epcLog.SetFloat("Pct Err", epc, pctErr)
                epcLog.SetFloat("Pct Cor", epc, pctCor)
                epcLog.SetFloat("Hid1 ActAvg", epc, float64(hid1Lay.Pools[0].ActAvg.ActPAvgEff))
                epcLog.SetFloat("Hid2 ActAvg", epc, float64(hid2Lay.Pools[0].ActAvg.ActPAvgEff))
                epcLog.SetFloat("Out ActAvg", epc, float64(outLay.Pools[0].ActAvg.ActPAvgEff))
        }
        tmr.Stop()
        if Silent {
                fmt.Printf("%v\n", tmr.Total)
        } else {
                fmt.Printf("Took %v for %v epochs, avg per epc: m%6.4g\n", tmr.Total, epcs, float64(tmr.Total)/float64(int(time.Second)*epcs))
        }
}

func main() {
        var epochs int
        var pats int
        var units int

        flag.Usage = func() {
                fmt.Fprintf(flag.CommandLine.Output(), "Usage of %s:\n", os.Args[0])
                flag.PrintDefaults()
        }

        // process command args
        flag.IntVar(&epochs, "epochs", 2, "number of epochs to run")
        flag.IntVar(&pats, "pats", 10, "number of patterns per epoch")
        flag.IntVar(&units, "units", 100, "number of units per layer -- uses NxN where N = sqrt(units)")
        flag.BoolVar(&Silent, "silent", false, "only report the final time")
        flag.Parse()

        if !Silent {
                fmt.Printf("Running bench with: %v epochs, %v pats, %v units\n", epochs, pats, units)
        }

        Net = leabra.NewNetwork("Bench")
        ConfigNet(Net, units)

        Pats = &table.Table{}
        ConfigPats(Pats, pats, units)

        EpcLog = &table.Table{}
        ConfigEpcLog(EpcLog)

        TrainNet(Net, Pats, EpcLog, epochs)

        EpcLog.SaveCSV("bench_epc.dat", ',', table.Headers)
}

// Copyright (c) 2019, The Emergent Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// deep_fsa runs a DeepLeabra network on the classic Reber grammar
// finite state automaton problem.
package main

//go:generate core generate -add-types

import (
        "log"
        "os"

        "cogentcore.org/core/core"
        "cogentcore.org/core/enums"
        "cogentcore.org/core/icons"
        "cogentcore.org/core/math32/vecint"
        "cogentcore.org/core/tree"
        "cogentcore.org/lab/base/mpi"
        "cogentcore.org/lab/base/randx"
        "github.com/emer/emergent/v2/econfig"
        "github.com/emer/emergent/v2/egui"
        "github.com/emer/emergent/v2/elog"
        "github.com/emer/emergent/v2/emer"
        "github.com/emer/emergent/v2/env"
        "github.com/emer/emergent/v2/estats"
        "github.com/emer/emergent/v2/etime"
        "github.com/emer/emergent/v2/looper"
        "github.com/emer/emergent/v2/netview"
        "github.com/emer/emergent/v2/params"
        "github.com/emer/emergent/v2/paths"
        "github.com/emer/etensor/tensor/table"
        "github.com/emer/leabra/v2/leabra"
)

func main() {
        sim := &Sim{}
        sim.New()
        sim.ConfigAll()
        if sim.Config.GUI {
                sim.RunGUI()
        } else {
                sim.RunNoGUI()
        }
}

// ParamSets is the default set of parameters.
// Base is always applied, and others can be optionally
// selected to apply on top of that.
var ParamSets = params.Sets{
        "Base": {
                {Sel: "Path", Desc: "norm and momentum on is critical, wt bal not as much but fine",
                        Params: params.Params{
                                "Path.Learn.Norm.On":     "true",
                                "Path.Learn.Momentum.On": "true",
                                "Path.Learn.WtBal.On":    "true",
                        }},
                {Sel: "Layer", Desc: "using default 1.8 inhib for hidden layers",
                        Params: params.Params{
                                "Layer.Inhib.Layer.Gi":     "1.8",
                                "Layer.Learn.AvgL.Gain":    "1.5",  // key to lower relative to 2.5
                                "Layer.Act.Gbar.L":         "0.1",  // lower leak = better
                                "Layer.Inhib.ActAvg.Fixed": "true", // simpler to have everything fixed, for replicability
                                "Layer.Act.Init.Decay":     "0",    // essential to have all layers no decay
                        }},
                {Sel: ".SuperLayer", Desc: "fix avg act",
                        Params: params.Params{
                                "Layer.Inhib.ActAvg.Fixed": "true",
                        }},
                {Sel: ".BackPath", Desc: "top-down back-pathways MUST have lower relative weight scale, otherwise network hallucinates",
                        Params: params.Params{
                                "Path.WtScale.Rel": "0.2",
                        }},
                {Sel: ".PulvinarLayer", Desc: "standard weight is .3 here for larger distributed reps. no learn",
                        Params: params.Params{
                                "Layer.Pulvinar.DriveScale": "0.8", // using .8 for localist layer
                        }},
                {Sel: ".CTCtxtPath", Desc: "no weight balance on CT context paths -- makes a diff!",
                        Params: params.Params{
                                "Path.Learn.WtBal.On": "false", // this should be true for larger DeepLeabra models -- e.g., sg..
                        }},
                {Sel: ".CTFromSuper", Desc: "initial weight = 0.5 much better than 0.8",
                        Params: params.Params{
                                "Path.WtInit.Mean": "0.5",
                        }},
                {Sel: ".Input", Desc: "input layers need more inhibition",
                        Params: params.Params{
                                "Layer.Inhib.Layer.Gi":    "2.0",
                                "Layer.Inhib.ActAvg.Init": "0.15",
                        }},
                {Sel: "#HiddenPToHiddenCT", Desc: "critical to make this small so deep context dominates",
                        Params: params.Params{
                                "Path.WtScale.Rel": "0.05",
                        }},
                {Sel: "#HiddenCTToHiddenCT", Desc: "testing",
                        Params: params.Params{
                                "Path.Learn.WtBal.On": "false",
                        }},
        },
}

// ParamConfig has config parameters related to sim params
type ParamConfig struct {

        // network parameters
        Network map[string]any

        // size of hidden layer -- can use emer.LaySize for 4D layers
        Hidden1Size vecint.Vector2i `default:"{'X':7,'Y':7}" nest:"+"`

        // size of hidden layer -- can use emer.LaySize for 4D layers
        Hidden2Size vecint.Vector2i `default:"{'X':7,'Y':7}" nest:"+"`

        // Extra Param Sheet name(s) to use (space separated if multiple).
        // must be valid name as listed in compiled-in params or loaded params
        Sheet string

        // extra tag to add to file names and logs saved from this run
        Tag string

        // user note -- describe the run params etc -- like a git commit message for the run
        Note string

        // Name of the JSON file to input saved parameters from.
        File string `nest:"+"`

        // Save a snapshot of all current param and config settings
        // in a directory named params_<datestamp> (or _good if Good is true), then quit.
        // Useful for comparing to later changes and seeing multiple views of current params.
        SaveAll bool `nest:"+"`

        // For SaveAll, save to params_good for a known good params state.
        // This can be done prior to making a new release after all tests are passing.
        // add results to git to provide a full diff record of all params over time.
        Good bool `nest:"+"`
}

// RunConfig has config parameters related to running the sim
type RunConfig struct {
        // starting run number, which determines the random seed.
        // runs counts from there, can do all runs in parallel by launching
        // separate jobs with each run, runs = 1.
        Run int `default:"0"`

        // total number of runs to do when running Train
        NRuns int `default:"5" min:"1"`

        // total number of epochs per run
        NEpochs int `default:"100"`

        // stop run after this number of perfect, zero-error epochs.
        NZero int `default:"2"`

        // total number of trials per epoch.  Should be an even multiple of NData.
        NTrials int `default:"100"`

        // how often to run through all the test patterns, in terms of training epochs.
        // can use 0 or -1 for no testing.
        TestInterval int `default:"5"`

        // how frequently (in epochs) to compute PCA on hidden representations
        // to measure variance?
        PCAInterval int `default:"5"`

        // if non-empty, is the name of weights file to load at start
        // of first run, for testing.
        StartWts string
}

// LogConfig has config parameters related to logging data
type LogConfig struct {

        // if true, save final weights after each run
        SaveWeights bool

        // if true, save train epoch log to file, as .epc.tsv typically
        Epoch bool `default:"true" nest:"+"`

        // if true, save run log to file, as .run.tsv typically
        Run bool `default:"true" nest:"+"`

        // if true, save train trial log to file, as .trl.tsv typically. May be large.
        Trial bool `default:"false" nest:"+"`

        // if true, save testing epoch log to file, as .tst_epc.tsv typically.  In general it is better to copy testing items over to the training epoch log and record there.
        TestEpoch bool `default:"false" nest:"+"`

        // if true, save testing trial log to file, as .tst_trl.tsv typically. May be large.
        TestTrial bool `default:"false" nest:"+"`

        // if true, save network activation etc data from testing trials,
        // for later viewing in netview.
        NetData bool
}

// Config is a standard Sim config -- use as a starting point.
type Config struct {

        // specify include files here, and after configuration,
        // it contains list of include files added.
        Includes []string

        // open the GUI -- does not automatically run -- if false,
        // then runs automatically and quits.
        GUI bool `default:"true"`

        // log debugging information
        Debug bool

        // InputNames are names of input letters
        InputNames []string

        // InputNameMap has indexes of InputNames
        InputNameMap map[string]int

        // parameter related configuration options
        Params ParamConfig `display:"add-fields"`

        // sim running related configuration options
        Run RunConfig `display:"add-fields"`

        // data logging related configuration options
        Log LogConfig `display:"add-fields"`
}

func (cfg *Config) IncludesPtr() *[]string { return &cfg.Includes }

// Sim encapsulates the entire simulation model, and we define all the
// functionality as methods on this struct.  This structure keeps all relevant
// state information organized and available without having to pass everything around
// as arguments to methods, and provides the core GUI interface (note the view tags
// for the fields which provide hints to how things should be displayed).
type Sim struct {

        // simulation configuration parameters -- set by .toml config file and / or args
        Config Config `new-window:"+"`

        // the network -- click to view / edit parameters for layers, paths, etc
        Net *leabra.Network `new-window:"+" display:"no-inline"`

        // network parameter management
        Params emer.NetParams `display:"add-fields"`

        // contains looper control loops for running sim
        Loops *looper.Stacks `new-window:"+" display:"no-inline"`

        // contains computed statistic values
        Stats estats.Stats `new-window:"+"`

        // Contains all the logs and information about the logs.'
        Logs elog.Logs `new-window:"+"`

        // the training patterns to use
        Patterns *table.Table `new-window:"+" display:"no-inline"`

        // Environments
        Envs env.Envs `new-window:"+" display:"no-inline"`

        // leabra timing parameters and state
        Context leabra.Context `new-window:"+"`

        // netview update parameters
        ViewUpdate netview.ViewUpdate `display:"add-fields"`

        // manages all the gui elements
        GUI egui.GUI `display:"-"`

        // a list of random seeds to use for each run
        RandSeeds randx.Seeds `display:"-"`
}

// New creates new blank elements and initializes defaults
func (ss *Sim) New() {
        econfig.Config(&ss.Config, "config.toml")
        ss.Config.InputNames = []string{"B", "T", "S", "X", "V", "P", "E"}
        ss.Net = leabra.NewNetwork("RA25")
        ss.Params.Config(ParamSets, ss.Config.Params.Sheet, ss.Config.Params.Tag, ss.Net)
        ss.Stats.Init()
        ss.Patterns = &table.Table{}
        ss.RandSeeds.Init(100) // max 100 runs
        ss.InitRandSeed(0)
        ss.Context.Defaults()
}

//////////////////////////////////////////////////////////////////////////////
//                 Configs

// ConfigAll configures all the elements using the standard functions
func (ss *Sim) ConfigAll() {
        ss.ConfigEnv()
        ss.ConfigNet(ss.Net)
        ss.ConfigLogs()
        ss.ConfigLoops()
        if ss.Config.Params.SaveAll {
                ss.Config.Params.SaveAll = false
                ss.Net.SaveParamsSnapshot(&ss.Params.Params, &ss.Config, ss.Config.Params.Good)
                os.Exit(0)
        }
}

func (ss *Sim) ConfigEnv() {
        // Can be called multiple times -- don't re-create
        var trn, tst *FSAEnv
        if len(ss.Envs) == 0 {
                trn = &FSAEnv{}
                tst = &FSAEnv{}
        } else {
                trn = ss.Envs.ByMode(etime.Train).(*FSAEnv)
                tst = ss.Envs.ByMode(etime.Test).(*FSAEnv)
        }

        if ss.Config.InputNameMap == nil {
                ss.Config.InputNameMap = make(map[string]int, len(ss.Config.InputNames))
                for i, nm := range ss.Config.InputNames {
                        ss.Config.InputNameMap[nm] = i
                }
        }

        // note: names must be standard here!
        trn.Name = etime.Train.String()
        trn.Seq.Max = 25 // 25 sequences per epoch training
        trn.TMatReber()

        tst.Name = etime.Test.String()
        tst.Seq.Max = 10
        tst.TMatReber() // todo: random

        trn.Init(0)
        tst.Init(0)

        // note: names must be in place when adding
        ss.Envs.Add(trn, tst)
}

func (ss *Sim) ConfigNet(net *leabra.Network) {
        net.SetRandSeed(ss.RandSeeds[0]) // init new separate random seed, using run = 0

        in := net.AddLayer2D("Input", 1, 7, leabra.InputLayer)
        hid, hidct, hidp := net.AddDeep2D("Hidden", 8, 8)

        hidp.Shape.CopyShape(&in.Shape)
        hidp.Drivers.Add("Input")

        trg := net.AddLayer2D("Targets", 1, 7, leabra.InputLayer) // just for visualization

        in.AddClass("Input")
        hidp.AddClass("Input")
        trg.AddClass("Input")

        hidct.PlaceRightOf(hid, 2)
        hidp.PlaceRightOf(in, 2)
        trg.PlaceBehind(hidp, 2)

        full := paths.NewFull()
        full.SelfCon = true // unclear if this makes a diff for self cons at all

        net.ConnectLayers(in, hid, full, leabra.ForwardPath)

        // for this small localist model with longer-term dependencies,
        // these additional context pathways turn out to be essential!
        // larger models in general do not require them, though it might be
        // good to check
        net.ConnectCtxtToCT(hidct, hidct, full)
        // net.LateralConnectLayer(hidct, full) // note: this does not work AT ALL -- essential to learn from t-1
        net.ConnectCtxtToCT(in, hidct, full)

        net.Build()
        net.Defaults()
        ss.ApplyParams()
        net.InitWeights()
}

func (ss *Sim) ApplyParams() {
        ss.Params.SetAll()
        if ss.Config.Params.Network != nil {
                ss.Params.SetNetworkMap(ss.Net, ss.Config.Params.Network)
        }
}

////////////////////////////////////////////////////////////////////////////////
//             Init, utils

// Init restarts the run, and initializes everything, including network weights
// and resets the epoch log table
func (ss *Sim) Init() {
        if ss.Config.GUI {
                ss.Stats.SetString("RunName", ss.Params.RunName(0)) // in case user interactively changes tag
        }
        ss.Loops.ResetCounters()
        ss.InitRandSeed(0)
        // ss.ConfigEnv() // re-config env just in case a different set of patterns was
        // selected or patterns have been modified etc
        ss.GUI.StopNow = false
        ss.ApplyParams()
        ss.NewRun()
        ss.ViewUpdate.RecordSyns()
        ss.ViewUpdate.Update()
}

// InitRandSeed initializes the random seed based on current training run number
func (ss *Sim) InitRandSeed(run int) {
        ss.RandSeeds.Set(run)
        ss.RandSeeds.Set(run, &ss.Net.Rand)
}

// ConfigLoops configures the control loops: Training, Testing
func (ss *Sim) ConfigLoops() {
        ls := looper.NewStacks()

        trls := ss.Config.Run.NTrials

        ls.AddStack(etime.Train).
                AddTime(etime.Run, ss.Config.Run.NRuns).
                AddTime(etime.Epoch, ss.Config.Run.NEpochs).
                AddTime(etime.Trial, trls).
                AddTime(etime.Cycle, 100)

        ls.AddStack(etime.Test).
                AddTime(etime.Epoch, 1).
                AddTime(etime.Trial, trls).
                AddTime(etime.Cycle, 100)

        leabra.LooperStdPhases(ls, &ss.Context, ss.Net, 75, 99)                // plus phase timing
        leabra.LooperSimCycleAndLearn(ls, ss.Net, &ss.Context, &ss.ViewUpdate) // std algo code

        ls.Stacks[etime.Train].OnInit.Add("Init", func() { ss.Init() })

        for m, _ := range ls.Stacks {
                stack := ls.Stacks[m]
                stack.Loops[etime.Trial].OnStart.Add("ApplyInputs", func() {
                        ss.ApplyInputs()
                })
        }

        ls.Loop(etime.Train, etime.Run).OnStart.Add("NewRun", ss.NewRun)

        // Train stop early condition
        ls.Loop(etime.Train, etime.Epoch).IsDone.AddBool("NZeroStop", func() bool {
                // This is calculated in TrialStats
                stopNz := ss.Config.Run.NZero
                if stopNz <= 0 {
                        stopNz = 2
                }
                curNZero := ss.Stats.Int("NZero")
                stop := curNZero >= stopNz
                return stop
        })

        // Add Testing
        trainEpoch := ls.Loop(etime.Train, etime.Epoch)
        trainEpoch.OnStart.Add("TestAtInterval", func() {
                if (ss.Config.Run.TestInterval > 0) && ((trainEpoch.Counter.Cur+1)%ss.Config.Run.TestInterval == 0) {
                        // Note the +1 so that it doesn't occur at the 0th timestep.
                        ss.TestAll()
                }
        })

        /////////////////////////////////////////////
        // Logging

        ls.Loop(etime.Test, etime.Epoch).OnEnd.Add("LogTestErrors", func() {
                leabra.LogTestErrors(&ss.Logs)
        })
        ls.Loop(etime.Train, etime.Epoch).OnEnd.Add("PCAStats", func() {
                trnEpc := ls.Stacks[etime.Train].Loops[etime.Epoch].Counter.Cur
                if ss.Config.Run.PCAInterval > 0 && trnEpc%ss.Config.Run.PCAInterval == 0 {
                        leabra.PCAStats(ss.Net, &ss.Logs, &ss.Stats)
                        ss.Logs.ResetLog(etime.Analyze, etime.Trial)
                }
        })

        ls.AddOnEndToAll("Log", func(mode, time enums.Enum) {
                ss.Log(mode.(etime.Modes), time.(etime.Times))
        })
        leabra.LooperResetLogBelow(ls, &ss.Logs)

        ls.Loop(etime.Train, etime.Trial).OnEnd.Add("LogAnalyze", func() {
                trnEpc := ls.Stacks[etime.Train].Loops[etime.Epoch].Counter.Cur
                if (ss.Config.Run.PCAInterval > 0) && (trnEpc%ss.Config.Run.PCAInterval == 0) {
                        ss.Log(etime.Analyze, etime.Trial)
                }
        })

        ls.Loop(etime.Train, etime.Run).OnEnd.Add("RunStats", func() {
                ss.Logs.RunStats("PctCor", "FirstZero", "LastZero")
        })

        // Save weights to file, to look at later
        ls.Loop(etime.Train, etime.Run).OnEnd.Add("SaveWeights", func() {
                ctrString := ss.Stats.PrintValues([]string{"Run", "Epoch"}, []string{"%03d", "%05d"}, "_")
                leabra.SaveWeightsIfConfigSet(ss.Net, ss.Config.Log.SaveWeights, ctrString, ss.Stats.String("RunName"))
        })

        ////////////////////////////////////////////
        // GUI

        if !ss.Config.GUI {
                if ss.Config.Log.NetData {
                        ls.Loop(etime.Test, etime.Trial).OnEnd.Add("NetDataRecord", func() {
                                ss.GUI.NetDataRecord(ss.ViewUpdate.Text)
                        })
                }
        } else {
                leabra.LooperUpdateNetView(ls, &ss.ViewUpdate, ss.Net, ss.NetViewCounters)
                leabra.LooperUpdatePlots(ls, &ss.GUI)
                ls.Stacks[etime.Train].OnInit.Add("GUI-Init", func() { ss.GUI.UpdateWindow() })
                ls.Stacks[etime.Test].OnInit.Add("GUI-Init", func() { ss.GUI.UpdateWindow() })
        }

        if ss.Config.Debug {
                mpi.Println(ls.DocString())
        }
        ss.Loops = ls
}

// ApplyInputs applies input patterns from given environment.
// It is good practice to have this be a separate method with appropriate
// args so that it can be used for various different contexts
// (training, testing, etc).
func (ss *Sim) ApplyInputs() {
        ctx := &ss.Context
        net := ss.Net
        net.InitExt()

        ev := ss.Envs.ByMode(ctx.Mode).(*FSAEnv)
        ev.Step()
        ss.Stats.SetString("TrialName", ev.String())

        in := ss.Net.LayerByName("Input")
        trg := ss.Net.LayerByName("Targets")
        clrmsk, setmsk, _ := in.ApplyExtFlags()
        ns := ev.NNext.Values[0]
        for i := 0; i < ns; i++ {
                lbl := ev.NextLabels.Values[i]
                li, ok := ss.Config.InputNameMap[lbl]
                if !ok {
                        log.Printf("Input label: %v not found in InputNames list of labels\n", lbl)
                        continue
                }
                if i == 0 {
                        in.ApplyExtValue(li, 1, clrmsk, setmsk, false)
                }
                trg.ApplyExtValue(li, 1, clrmsk, setmsk, false)
        }
}

// NewRun intializes a new run of the model, using the TrainEnv.Run counter
// for the new run value
func (ss *Sim) NewRun() {
        ctx := &ss.Context
        ss.InitRandSeed(ss.Loops.Loop(etime.Train, etime.Run).Counter.Cur)
        ss.Envs.ByMode(etime.Train).Init(0)
        ss.Envs.ByMode(etime.Test).Init(0)
        ctx.Reset()
        ctx.Mode = etime.Train
        ss.Net.InitWeights()
        ss.InitStats()
        ss.StatCounters()
        ss.Logs.ResetLog(etime.Train, etime.Epoch)
        ss.Logs.ResetLog(etime.Test, etime.Epoch)
}

// TestAll runs through the full set of testing items
func (ss *Sim) TestAll() {
        ss.Envs.ByMode(etime.Test).Init(0)
        ss.Loops.ResetAndRun(etime.Test)
        ss.Loops.Mode = etime.Train // Important to reset Mode back to Train because this is called from within the Train Run.
}

////////////////////////////////////////////////////////////////////////////////////////////
//                 Stats

// InitStats initializes all the statistics.
// called at start of new run
func (ss *Sim) InitStats() {
        ss.Stats.SetFloat("UnitErr", 0.0)
        ss.Stats.SetFloat("CorSim", 0.0)
        ss.Stats.SetString("TrialName", "")
        ss.Logs.InitErrStats() // inits TrlErr, FirstZero, LastZero, NZero
}

// StatCounters saves current counters to Stats, so they are available for logging etc
// Also saves a string rep of them for ViewUpdate.Text
func (ss *Sim) StatCounters() {
        ctx := &ss.Context
        mode := ctx.Mode
        ss.Loops.Stacks[mode].CountersToStats(&ss.Stats)
        // always use training epoch..
        trnEpc := ss.Loops.Stacks[etime.Train].Loops[etime.Epoch].Counter.Cur
        ss.Stats.SetInt("Epoch", trnEpc)
        trl := ss.Stats.Int("Trial")
        ss.Stats.SetInt("Trial", trl)
        ss.Stats.SetInt("Cycle", int(ctx.Cycle))
}

func (ss *Sim) NetViewCounters(tm etime.Times) {
        if ss.ViewUpdate.View == nil {
                return
        }
        if tm == etime.Trial {
                ss.TrialStats() // get trial stats for current di
        }
        ss.StatCounters()
        ss.ViewUpdate.Text = ss.Stats.Print([]string{"Run", "Epoch", "Trial", "TrialName", "Cycle", "UnitErr", "TrlErr", "CorSim"})
}

// TrialStats computes the trial-level statistics.
// Aggregation is done directly from log data.
func (ss *Sim) TrialStats() {
        inp := ss.Net.LayerByName("HiddenP")
        trg := ss.Net.LayerByName("Targets")
        ss.Stats.SetFloat("CorSim", float64(inp.CosDiff.Cos))
        sse := 0.0
        gotOne := false
        for ni := range inp.Neurons {
                inn := &inp.Neurons[ni]
                tgn := &trg.Neurons[ni]
                if tgn.Act > 0.5 {
                        if inn.ActM > 0.4 {
                                gotOne = true
                        }
                } else {
                        if inn.ActM > 0.5 {
                                sse += float64(inn.ActM)
                        }
                }
        }
        if !gotOne {
                sse += 1
        }
        ss.Stats.SetFloat("SSE", sse)
        ss.Stats.SetFloat("AvgSSE", sse)
        if sse > 0 {
                ss.Stats.SetFloat("TrlErr", 1)
        } else {
                ss.Stats.SetFloat("TrlErr", 0)
        }
}

//////////////////////////////////////////////////////////////////////////////
//                 Logging

func (ss *Sim) ConfigLogs() {
        ss.Stats.SetString("RunName", ss.Params.RunName(0)) // used for naming logs, stats, etc

        ss.Logs.AddCounterItems(etime.Run, etime.Epoch, etime.Trial, etime.Cycle)
        ss.Logs.AddStatStringItem(etime.AllModes, etime.AllTimes, "RunName")
        ss.Logs.AddStatStringItem(etime.AllModes, etime.Trial, "TrialName")

        ss.Logs.AddStatAggItem("CorSim", etime.Run, etime.Epoch, etime.Trial)
        ss.Logs.AddStatAggItem("UnitErr", etime.Run, etime.Epoch, etime.Trial)
        ss.Logs.AddErrStatAggItems("TrlErr", etime.Run, etime.Epoch, etime.Trial)

        ss.Logs.AddCopyFromFloatItems(etime.Train, []etime.Times{etime.Epoch, etime.Run}, etime.Test, etime.Epoch, "Tst", "CorSim", "UnitErr", "PctCor", "PctErr")

        ss.Logs.AddPerTrlMSec("PerTrlMSec", etime.Run, etime.Epoch, etime.Trial)

        layers := ss.Net.LayersByType(leabra.SuperLayer, leabra.CTLayer, leabra.TargetLayer)
        leabra.LogAddDiagnosticItems(&ss.Logs, layers, etime.Train, etime.Epoch, etime.Trial)
        leabra.LogInputLayer(&ss.Logs, ss.Net, etime.Train)

        leabra.LogAddPCAItems(&ss.Logs, ss.Net, etime.Train, etime.Run, etime.Epoch, etime.Trial)

        ss.Logs.AddLayerTensorItems(ss.Net, "Act", etime.Test, etime.Trial, "InputLayer", "TargetLayer")

        ss.Logs.PlotItems("CorSim", "PctCor", "FirstZero", "LastZero")

        ss.Logs.CreateTables()
        ss.Logs.SetContext(&ss.Stats, ss.Net)
        // don't plot certain combinations we don't use
        ss.Logs.NoPlot(etime.Train, etime.Cycle)
        ss.Logs.NoPlot(etime.Test, etime.Run)
        // note: Analyze not plotted by default
        ss.Logs.SetMeta(etime.Train, etime.Run, "LegendCol", "RunName")
}

// Log is the main logging function, handles special things for different scopes
func (ss *Sim) Log(mode etime.Modes, time etime.Times) {
        ctx := &ss.Context
        if mode != etime.Analyze {
                ctx.Mode = mode // Also set specifically in a Loop callback.
        }
        dt := ss.Logs.Table(mode, time)
        if dt == nil {
                return
        }
        row := dt.Rows

        switch {
        case time == etime.Cycle:
                return
        case time == etime.Trial:
                ss.TrialStats()
                ss.StatCounters()
        }

        ss.Logs.LogRow(mode, time, row) // also logs to file, etc
}

//////// GUI

// ConfigGUI configures the Cogent Core GUI interface for this simulation.
func (ss *Sim) ConfigGUI() {
        title := "Leabra Random Associator"
        ss.GUI.MakeBody(ss, "ra25", title, `This demonstrates a basic Leabra model. See <a href="https://github.com/emer/emergent">emergent on GitHub</a>.</p>`)
        ss.GUI.CycleUpdateInterval = 10

        nv := ss.GUI.AddNetView("Network")
        nv.Options.MaxRecs = 300
        nv.SetNet(ss.Net)
        ss.ViewUpdate.Config(nv, etime.AlphaCycle, etime.AlphaCycle)
        ss.GUI.ViewUpdate = &ss.ViewUpdate

        // nv.SceneXYZ().Camera.Pose.Pos.Set(0, 1, 2.75) // more "head on" than default which is more "top down"
        // nv.SceneXYZ().Camera.LookAt(math32.Vec3(0, 0, 0), math32.Vec3(0, 1, 0))

        ss.GUI.AddPlots(title, &ss.Logs)

        ss.GUI.FinalizeGUI(false)
}

func (ss *Sim) MakeToolbar(p *tree.Plan) {
        ss.GUI.AddLooperCtrl(p, ss.Loops)

        tree.Add(p, func(w *core.Separator) {})
        ss.GUI.AddToolbarItem(p, egui.ToolbarItem{Label: "Reset RunLog",
                Icon:    icons.Reset,
                Tooltip: "Reset the accumulated log of all Runs, which are tagged with the ParamSet used",
                Active:  egui.ActiveAlways,
                Func: func() {
                        ss.Logs.ResetLog(etime.Train, etime.Run)
                        ss.GUI.UpdatePlot(etime.Train, etime.Run)
                },
        })
        ////////////////////////////////////////////////
        tree.Add(p, func(w *core.Separator) {})
        ss.GUI.AddToolbarItem(p, egui.ToolbarItem{Label: "New Seed",
                Icon:    icons.Add,
                Tooltip: "Generate a new initial random seed to get different results.  By default, Init re-establishes the same initial seed every time.",
                Active:  egui.ActiveAlways,
                Func: func() {
                        ss.RandSeeds.NewSeeds()
                },
        })
        ss.GUI.AddToolbarItem(p, egui.ToolbarItem{Label: "README",
                Icon:    icons.FileMarkdown,
                Tooltip: "Opens your browser on the README file that contains instructions for how to run this model.",
                Active:  egui.ActiveAlways,
                Func: func() {
                        core.TheApp.OpenURL("https://github.com/emer/leabra/blob/main/examples/ra25/README.md")
                },
        })
}

func (ss *Sim) RunGUI() {
        ss.Init()
        ss.ConfigGUI()
        ss.GUI.Body.RunMainWindow()
}

func (ss *Sim) RunNoGUI() {
        if ss.Config.Params.Note != "" {
                mpi.Printf("Note: %s\n", ss.Config.Params.Note)
        }
        if ss.Config.Log.SaveWeights {
                mpi.Printf("Saving final weights per run\n")
        }
        runName := ss.Params.RunName(ss.Config.Run.Run)
        ss.Stats.SetString("RunName", runName) // used for naming logs, stats, etc
        netName := ss.Net.Name

        elog.SetLogFile(&ss.Logs, ss.Config.Log.Trial, etime.Train, etime.Trial, "trl", netName, runName)
        elog.SetLogFile(&ss.Logs, ss.Config.Log.Epoch, etime.Train, etime.Epoch, "epc", netName, runName)
        elog.SetLogFile(&ss.Logs, ss.Config.Log.Run, etime.Train, etime.Run, "run", netName, runName)
        elog.SetLogFile(&ss.Logs, ss.Config.Log.TestEpoch, etime.Test, etime.Epoch, "tst_epc", netName, runName)
        elog.SetLogFile(&ss.Logs, ss.Config.Log.TestTrial, etime.Test, etime.Trial, "tst_trl", netName, runName)

        netdata := ss.Config.Log.NetData
        if netdata {
                mpi.Printf("Saving NetView data from testing\n")
                ss.GUI.InitNetData(ss.Net, 200)
        }

        ss.Init()

        mpi.Printf("Running %d Runs starting at %d\n", ss.Config.Run.NRuns, ss.Config.Run.Run)
        ss.Loops.Loop(etime.Train, etime.Run).Counter.SetCurMaxPlusN(ss.Config.Run.Run, ss.Config.Run.NRuns)

        if ss.Config.Run.StartWts != "" { // this is just for testing -- not usually needed
                ss.Loops.Step(etime.Train, 1, etime.Trial) // get past NewRun
                ss.Net.OpenWeightsJSON(core.Filename(ss.Config.Run.StartWts))
                mpi.Printf("Starting with initial weights from: %s\n", ss.Config.Run.StartWts)
        }

        mpi.Printf("Set NThreads to: %d\n", ss.Net.NThreads)

        ss.Loops.Run(etime.Train)

        ss.Logs.CloseLogFiles()

        if netdata {
                ss.GUI.SaveNetData(ss.Stats.String("RunName"))
        }
}

// Copyright (c) 2019, The Emergent Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
        "fmt"

        "cogentcore.org/lab/base/randx"
        "github.com/emer/emergent/v2/env"
        "github.com/emer/emergent/v2/etime"
        "github.com/emer/etensor/tensor"
)

// FSAEnv generates states in a finite state automaton (FSA) which is a
// simple form of grammar for creating non-deterministic but still
// overall structured sequences.
type FSAEnv struct {

        // name of this environment
        Name string

        // transition matrix, which is a square NxN tensor with outer dim being current state and inner dim having probability of transitioning to that state
        TMat tensor.Float64 `display:"no-inline"`

        // transition labels, one for each transition cell in TMat matrix
        Labels tensor.String

        // automaton state within FSA that we're in
        AState env.CurPrvInt

        // number of next states in current state output (scalar)
        NNext tensor.Int

        // next states that have non-zero probability, with actual randomly chosen next state at start
        NextStates tensor.Int

        // transition labels for next states that have non-zero probability, with actual randomly chosen one for next state at start
        NextLabels tensor.String

        // sequence counter within epoch
        Seq env.Counter `display:"inline"`

        // tick counter within sequence
        Tick env.Counter `display:"inline"`

        // trial is the step counter within sequence - how many steps taken within current sequence -- it resets to 0 at start of each sequence
        Trial env.Counter `display:"inline"`
}

func (ev *FSAEnv) Label() string { return ev.Name }

// InitTMat initializes matrix and labels to given size
func (ev *FSAEnv) InitTMat(nst int) {
        ev.TMat.SetShape([]int{nst, nst})
        ev.Labels.SetShape([]int{nst, nst})
        ev.TMat.SetZeros()
        ev.Labels.SetZeros()
        ev.NNext.SetShape([]int{1})
        ev.NextStates.SetShape([]int{nst})
        ev.NextLabels.SetShape([]int{nst})
}

// SetTMat sets given transition matrix probability and label
func (ev *FSAEnv) SetTMat(fm, to int, p float64, lbl string) {
        ev.TMat.Set([]int{fm, to}, p)
        ev.Labels.Set([]int{fm, to}, lbl)
}

// TMatReber sets the transition matrix to the standard Reber grammar FSA
func (ev *FSAEnv) TMatReber() {
        ev.InitTMat(8)
        ev.SetTMat(0, 1, 1, "B")   // 0 = start
        ev.SetTMat(1, 2, 0.5, "T") // 1 = state 0 in usu diagram (+1 for all states)
        ev.SetTMat(1, 3, 0.5, "P")
        ev.SetTMat(2, 2, 0.5, "S")
        ev.SetTMat(2, 4, 0.5, "X")
        ev.SetTMat(3, 3, 0.5, "T")
        ev.SetTMat(3, 5, 0.5, "V")
        ev.SetTMat(4, 6, 0.5, "S")
        ev.SetTMat(4, 3, 0.5, "X")
        ev.SetTMat(5, 6, 0.5, "V")
        ev.SetTMat(5, 4, 0.5, "P")
        ev.SetTMat(6, 7, 1, "E") // 7 = end
        ev.Init(0)
}

func (ev *FSAEnv) Validate() error {
        if ev.TMat.Len() == 0 {
                return fmt.Errorf("FSAEnv: %v has no transition matrix TMat set", ev.Name)
        }
        return nil
}

func (ev *FSAEnv) State(element string) tensor.Tensor {
        switch element {
        case "NNext":
                return &ev.NNext
        case "NextStates":
                return &ev.NextStates
        case "NextLabels":
                return &ev.NextLabels
        }
        return nil
}

// String returns the current state as a string
func (ev *FSAEnv) String() string {
        nn := ev.NNext.Values[0]
        lbls := ev.NextLabels.Values[0:nn]
        return fmt.Sprintf("S_%d_%v", ev.AState.Cur, lbls)
}

func (ev *FSAEnv) Init(run int) {
        ev.Tick.Scale = etime.Tick
        ev.Trial.Scale = etime.Trial
        ev.Seq.Init()
        ev.Tick.Init()
        ev.Trial.Init()
        ev.Trial.Cur = -1 // init state -- key so that first Step() = 0
        ev.AState.Cur = 0
        ev.AState.Prv = -1
}

// NextState sets NextStates including randomly chosen one at start
func (ev *FSAEnv) NextState() {
        nst := ev.TMat.DimSize(0)
        if ev.AState.Cur < 0 || ev.AState.Cur >= nst-1 {
                ev.AState.Cur = 0
        }
        ri := ev.AState.Cur * nst
        ps := ev.TMat.Values[ri : ri+nst]
        ls := ev.Labels.Values[ri : ri+nst]
        nxt := randx.PChoose64(ps) // next state chosen at random
        ev.NextStates.Set1D(0, nxt)
        ev.NextLabels.Set1D(0, ls[nxt])
        idx := 1
        for i, p := range ps {
                if i != nxt && p > 0 {
                        ev.NextStates.Set1D(idx, i)
                        ev.NextLabels.Set1D(idx, ls[i])
                        idx++
                }
        }
        ev.NNext.Set1D(0, idx)
        ev.AState.Set(nxt)
}

func (ev *FSAEnv) Step() bool {
        ev.NextState()
        ev.Trial.Incr()
        ev.Tick.Incr()
        if ev.AState.Prv == 0 {
                ev.Tick.Init()
                ev.Seq.Incr()
        }
        return true
}

func (ev *FSAEnv) Action(element string, input tensor.Tensor) {
        // nop
}

// Compile-time check that implements Env interface
var _ env.Env = (*FSAEnv)(nil)

// Copyright (c) 2024, The Emergent Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// hip runs a hippocampus model on the AB-AC paired associate learning task.
package main

//go:generate core generate -add-types

import (
        "embed"
        "fmt"
        "math"
        "math/rand"
        "reflect"
        "strings"

        "cogentcore.org/core/base/errors"
        "cogentcore.org/core/core"
        "cogentcore.org/core/enums"
        "cogentcore.org/core/icons"
        "cogentcore.org/core/tree"
        "cogentcore.org/lab/base/randx"
        "github.com/emer/emergent/v2/econfig"
        "github.com/emer/emergent/v2/egui"
        "github.com/emer/emergent/v2/elog"
        "github.com/emer/emergent/v2/emer"
        "github.com/emer/emergent/v2/env"
        "github.com/emer/emergent/v2/estats"
        "github.com/emer/emergent/v2/etime"
        "github.com/emer/emergent/v2/looper"
        "github.com/emer/emergent/v2/netview"
        "github.com/emer/emergent/v2/params"
        "github.com/emer/emergent/v2/patgen"
        "github.com/emer/emergent/v2/paths"
        "github.com/emer/etensor/plot/plotcore"
        "github.com/emer/etensor/tensor/stats/split"
        "github.com/emer/etensor/tensor/table"
        "github.com/emer/leabra/v2/leabra"
)

//go:embed train_ab.tsv train_ac.tsv test_ab.tsv test_ac.tsv test_lure.tsv
var content embed.FS

func main() {
        sim := &Sim{}
        sim.New()
        sim.ConfigAll()
        sim.RunGUI()
}

// ParamSets is the default set of parameters -- Base is always applied, and others can be optionally
// selected to apply on top of that
var ParamSets = params.Sets{
        "Base": {
                {Sel: "Path", Desc: "keeping default params for generic prjns",
                        Params: params.Params{
                                "Path.Learn.Momentum.On": "true",
                                "Path.Learn.Norm.On":     "true",
                                "Path.Learn.WtBal.On":    "false",
                        }},
                {Sel: ".EcCa1Path", Desc: "encoder projections -- no norm, moment",
                        Params: params.Params{
                                "Path.Learn.Lrate":        "0.04",
                                "Path.Learn.Momentum.On":  "false",
                                "Path.Learn.Norm.On":      "false",
                                "Path.Learn.WtBal.On":     "true",
                                "Path.Learn.XCal.SetLLrn": "false", // using bcm now, better
                        }},
                {Sel: ".HippoCHL", Desc: "hippo CHL projections -- no norm, moment, but YES wtbal = sig better",
                        Params: params.Params{
                                "Path.CHL.Hebb":          "0.05",
                                "Path.Learn.Lrate":       "0.2",
                                "Path.Learn.Momentum.On": "false",
                                "Path.Learn.Norm.On":     "false",
                                "Path.Learn.WtBal.On":    "true",
                        }},
                {Sel: ".PPath", Desc: "perforant path, new Dg error-driven EcCa1Path prjns",
                        Params: params.Params{
                                "Path.Learn.Momentum.On": "false",
                                "Path.Learn.Norm.On":     "false",
                                "Path.Learn.WtBal.On":    "true",
                                "Path.Learn.Lrate":       "0.15", // err driven: .15 > .2 > .25 > .1
                                // moss=4, delta=4, lr=0.2, test = 3 are best
                        }},
                {Sel: "#CA1ToECout", Desc: "extra strong from CA1 to ECout",
                        Params: params.Params{
                                "Path.WtScale.Abs": "4.0",
                        }},
                {Sel: "#InputToECin", Desc: "one-to-one input to EC",
                        Params: params.Params{
                                "Path.Learn.Learn": "false",
                                "Path.WtInit.Mean": "0.8",
                                "Path.WtInit.Var":  "0.0",
                        }},
                {Sel: "#ECoutToECin", Desc: "one-to-one out to in",
                        Params: params.Params{
                                "Path.Learn.Learn": "false",
                                "Path.WtInit.Mean": "0.9",
                                "Path.WtInit.Var":  "0.01",
                                "Path.WtScale.Rel": "0.5",
                        }},
                {Sel: "#DGToCA3", Desc: "Mossy fibers: strong, non-learning",
                        Params: params.Params{
                                "Path.Learn.Learn": "false",
                                "Path.WtInit.Mean": "0.9",
                                "Path.WtInit.Var":  "0.01",
                                "Path.WtScale.Rel": "4",
                        }},
                {Sel: "#CA3ToCA3", Desc: "CA3 recurrent cons",
                        Params: params.Params{
                                "Path.WtScale.Rel": "0.1",
                                "Path.Learn.Lrate": "0.1",
                        }},
                {Sel: "#ECinToDG", Desc: "DG learning is surprisingly critical: maxed out fast, hebbian works best",
                        Params: params.Params{
                                "Path.Learn.Learn":       "true", // absolutely essential to have on!
                                "Path.CHL.Hebb":          ".5",   // .5 > 1 overall
                                "Path.CHL.SAvgCor":       "0.1",  // .1 > .2 > .3 > .4 ?
                                "Path.CHL.MinusQ1":       "true", // dg self err?
                                "Path.Learn.Lrate":       "0.4",  // .4 > .3 > .2
                                "Path.Learn.Momentum.On": "false",
                                "Path.Learn.Norm.On":     "false",
                                "Path.Learn.WtBal.On":    "true",
                        }},
                {Sel: "#CA3ToCA1", Desc: "Schaffer collaterals -- slower, less hebb",
                        Params: params.Params{
                                "Path.CHL.Hebb":          "0.01",
                                "Path.CHL.SAvgCor":       "0.4",
                                "Path.Learn.Lrate":       "0.1",
                                "Path.Learn.Momentum.On": "false",
                                "Path.Learn.Norm.On":     "false",
                                "Path.Learn.WtBal.On":    "true",
                        }},
                {Sel: ".EC", Desc: "all EC layers: only pools, no layer-level",
                        Params: params.Params{
                                "Layer.Act.Gbar.L":        ".1",
                                "Layer.Inhib.ActAvg.Init": "0.2",
                                "Layer.Inhib.Layer.On":    "false",
                                "Layer.Inhib.Pool.Gi":     "2.0",
                                "Layer.Inhib.Pool.On":     "true",
                        }},
                {Sel: "#DG", Desc: "very sparse = high inibhition",
                        Params: params.Params{
                                "Layer.Inhib.ActAvg.Init": "0.01",
                                "Layer.Inhib.Layer.Gi":    "3.8",
                        }},
                {Sel: "#CA3", Desc: "sparse = high inibhition",
                        Params: params.Params{
                                "Layer.Inhib.ActAvg.Init": "0.02",
                                "Layer.Inhib.Layer.Gi":    "2.8",
                        }},
                {Sel: "#CA1", Desc: "CA1 only Pools",
                        Params: params.Params{
                                "Layer.Inhib.ActAvg.Init": "0.1",
                                "Layer.Inhib.Layer.On":    "false",
                                "Layer.Inhib.Pool.Gi":     "2.4",
                                "Layer.Inhib.Pool.On":     "true",
                        }},
        },
}

// Config has config parameters related to running the sim
type Config struct {
        // total number of runs to do when running Train
        NRuns int `default:"10" min:"1"`

        // total number of epochs per run
        NEpochs int `default:"20"`

        // stop run after this number of perfect, zero-error epochs.
        NZero int `default:"1"`

        // how often to run through all the test patterns, in terms of training epochs.
        // can use 0 or -1 for no testing.
        TestInterval int `default:"1"`

        // StopMem is the threshold for stopping learning.
        StopMem float32 `default:"1"`
}

// Sim encapsulates the entire simulation model, and we define all the
// functionality as methods on this struct.  This structure keeps all relevant
// state information organized and available without having to pass everything around
// as arguments to methods, and provides the core GUI interface (note the view tags
// for the fields which provide hints to how things should be displayed).
type Sim struct {

        // simulation configuration parameters -- set by .toml config file and / or args
        Config Config `new-window:"+"`

        // the network -- click to view / edit parameters for layers, paths, etc
        Net *leabra.Network `new-window:"+" display:"no-inline"`

        // all parameter management
        Params emer.NetParams `display:"add-fields"`

        // contains looper control loops for running sim
        Loops *looper.Stacks `new-window:"+" display:"no-inline"`

        // contains computed statistic values
        Stats estats.Stats `new-window:"+"`

        // Contains all the logs and information about the logs.'
        Logs elog.Logs `new-window:"+"`

        // if true, run in pretrain mode
        PretrainMode bool `display:"-"`

        // pool patterns vocabulary
        PoolVocab patgen.Vocab `display:"-"`

        // AB training patterns to use
        TrainAB *table.Table `new-window:"+" display:"no-inline"`

        // AC training patterns to use
        TrainAC *table.Table `new-window:"+" display:"no-inline"`

        // AB testing patterns to use
        TestAB *table.Table `new-window:"+" display:"no-inline"`

        // AC testing patterns to use
        TestAC *table.Table `new-window:"+" display:"no-inline"`

        // Lure testing patterns to use
        TestLure *table.Table `new-window:"+" display:"no-inline"`

        // TestAll has all the test items
        TestAll *table.Table `new-window:"+" display:"no-inline"`

        // Lure pretrain patterns to use
        PreTrainLure *table.Table `new-window:"+" display:"-"`

        // all training patterns -- for pretrain
        TrainAll *table.Table `new-window:"+" display:"-"`

        // Environments
        Envs env.Envs `new-window:"+" display:"no-inline"`

        // leabra timing parameters and state
        Context leabra.Context `new-window:"+"`

        // netview update parameters
        ViewUpdate netview.ViewUpdate `display:"add-fields"`

        // manages all the gui elements
        GUI egui.GUI `display:"-"`

        // a list of random seeds to use for each run
        RandSeeds randx.Seeds `display:"-"`
}

// New creates new blank elements and initializes defaults
func (ss *Sim) New() {
        // ss.Config.Defaults()
        econfig.Config(&ss.Config, "config.toml")
        // ss.Config.Hip.EC5Clamp = true      // must be true in hip.go to have a target layer
        // ss.Config.Hip.EC5ClampTest = false // key to be off for cmp stats on completion region

        ss.Net = leabra.NewNetwork("Hip")
        ss.Params.Config(ParamSets, "", "", ss.Net)
        ss.Stats.Init()
        ss.Stats.SetInt("Expt", 0)

        ss.PoolVocab = patgen.Vocab{}
        ss.TrainAB = &table.Table{}
        ss.TrainAC = &table.Table{}
        ss.TestAB = &table.Table{}
        ss.TestAC = &table.Table{}
        ss.PreTrainLure = &table.Table{}
        ss.TestLure = &table.Table{}
        ss.TrainAll = &table.Table{}
        ss.TestAll = &table.Table{}
        ss.PretrainMode = false

        ss.RandSeeds.Init(100) // max 100 runs
        ss.InitRandSeed(0)
        ss.Context.Defaults()
}

////////////////////////////////////////////////////////////////////////////////////////////
//                 Configs

// Config configures all the elements using the standard functions
func (ss *Sim) ConfigAll() {
        ss.OpenPatterns()
        // ss.ConfigPatterns()
        ss.ConfigEnv()
        ss.ConfigNet(ss.Net)
        ss.ConfigLogs()
        ss.ConfigLoops()
}

func (ss *Sim) ConfigEnv() {
        // Can be called multiple times -- don't re-create
        var trn, tst *env.FixedTable
        if len(ss.Envs) == 0 {
                trn = &env.FixedTable{}
                tst = &env.FixedTable{}
        } else {
                trn = ss.Envs.ByMode(etime.Train).(*env.FixedTable)
                tst = ss.Envs.ByMode(etime.Test).(*env.FixedTable)
        }

        // note: names must be standard here!
        trn.Name = etime.Train.String()
        trn.Config(table.NewIndexView(ss.TrainAB))
        trn.Validate()

        tst.Name = etime.Test.String()
        tst.Config(table.NewIndexView(ss.TestAll))
        tst.Sequential = true
        tst.Validate()

        trn.Init(0)
        tst.Init(0)

        // note: names must be in place when adding
        ss.Envs.Add(trn, tst)
}

func (ss *Sim) ConfigNet(net *leabra.Network) {
        net.SetRandSeed(ss.RandSeeds[0]) // init new separate random seed, using run = 0

        in := net.AddLayer4D("Input", 6, 2, 3, 4, leabra.InputLayer)
        ecin := net.AddLayer4D("ECin", 6, 2, 3, 4, leabra.SuperLayer)
        ecout := net.AddLayer4D("ECout", 6, 2, 3, 4, leabra.TargetLayer) // clamped in plus phase
        ca1 := net.AddLayer4D("CA1", 6, 2, 4, 10, leabra.SuperLayer)
        dg := net.AddLayer2D("DG", 25, 25, leabra.SuperLayer)
        ca3 := net.AddLayer2D("CA3", 30, 10, leabra.SuperLayer)

        ecin.AddClass("EC")
        ecout.AddClass("EC")

        onetoone := paths.NewOneToOne()
        pool1to1 := paths.NewPoolOneToOne()
        full := paths.NewFull()

        net.ConnectLayers(in, ecin, onetoone, leabra.ForwardPath)
        net.ConnectLayers(ecout, ecin, onetoone, leabra.BackPath)

        // EC <-> CA1 encoder pathways
        net.ConnectLayers(ecin, ca1, pool1to1, leabra.EcCa1Path)
        net.ConnectLayers(ca1, ecout, pool1to1, leabra.EcCa1Path)
        net.ConnectLayers(ecout, ca1, pool1to1, leabra.EcCa1Path)

        // Perforant pathway
        ppath := paths.NewUniformRand()
        ppath.PCon = 0.25

        net.ConnectLayers(ecin, dg, ppath, leabra.CHLPath).AddClass("HippoCHL")

        net.ConnectLayers(ecin, ca3, ppath, leabra.EcCa1Path).AddClass("PPath")
        net.ConnectLayers(ca3, ca3, full, leabra.EcCa1Path).AddClass("PPath")

        // Mossy fibers
        mossy := paths.NewUniformRand()
        mossy.PCon = 0.02
        net.ConnectLayers(dg, ca3, mossy, leabra.CHLPath).AddClass("HippoCHL")

        // Schafer collaterals
        net.ConnectLayers(ca3, ca1, full, leabra.CHLPath).AddClass("HippoCHL")

        ecin.PlaceRightOf(in, 2)
        ecout.PlaceRightOf(ecin, 2)
        dg.PlaceAbove(in)
        ca3.PlaceAbove(dg)
        ca1.PlaceRightOf(ca3, 2)

        in.Doc = "Input represents cortical processing areas for different sensory modalities, semantic categories, etc, organized into pools. It is pre-compressed in this model, to simplify and allow one-to-one projections into the EC."

        ecin.Doc = "Entorhinal Cortex (EC) input layer is the superficial layer 2 that receives from the cortex and projects into the hippocampus. It has compressed representations of cortical inputs."

        ecout.Doc = "Entorhinal Cortex (EC) output layer is the deep layers that are bidirectionally connected to the CA1, and communicate hippocampal recall back out to the cortex, while also training the CA1 to accurately represent the EC inputs"

        ca1.Doc = "CA (Cornu Ammonis = Ammon's horn) area 1, receives from CA3 and drives recalled memory output to ECout"

        ca3.Doc = "CA (Cornu Ammonis = Ammon's horn) area 3, receives inputs from ECin and DG, and is the primary site of memory encoding. Recurrent self-connections drive pattern completion of full memory representations from partial cues, along with connections to CA1 that drive memory output."

        dg.Doc = "Dentate Gyruns, which receives broad inputs from ECin and has highly sparse, pattern separated representations, which drive more separated representations in CA3"

        net.Build()
        net.Defaults()
        ss.ApplyParams()
        net.InitWeights()
        net.InitTopoScales()
}

func (ss *Sim) ApplyParams() {
        ss.Params.Network = ss.Net
        ss.Params.SetAll()
}

////////////////////////////////////////////////////////////////////////////////
//             Init, utils

// Init restarts the run, and initializes everything, including network weights
// and resets the epoch log table
func (ss *Sim) Init() {
        ss.Stats.SetString("RunName", ss.Params.RunName(0)) // in case user interactively changes tag
        ss.Loops.ResetCounters()

        ss.GUI.StopNow = false
        ss.ApplyParams()
        ss.NewRun()
        ss.ViewUpdate.RecordSyns()
        ss.ViewUpdate.Update()
}

func (ss *Sim) TestInit() {
        ss.Loops.InitMode(etime.Test)
        tst := ss.Envs.ByMode(etime.Test).(*env.FixedTable)
        tst.Init(0)
}

// InitRandSeed initializes the random seed based on current training run number
func (ss *Sim) InitRandSeed(run int) {
        rand.Seed(ss.RandSeeds[run])
        ss.RandSeeds.Set(run)
        ss.RandSeeds.Set(run, &ss.Net.Rand)
        patgen.NewRand(ss.RandSeeds[run])
}

// ConfigLoops configures the control loops: Training, Testing
func (ss *Sim) ConfigLoops() {
        ls := looper.NewStacks()

        trls := ss.TrainAB.Rows
        ttrls := ss.TestAll.Rows

        ls.AddStack(etime.Train).AddTime(etime.Run, ss.Config.NRuns).AddTime(etime.Epoch, ss.Config.NEpochs).AddTime(etime.Trial, trls).AddTime(etime.Cycle, 100)

        ls.AddStack(etime.Test).AddTime(etime.Epoch, 1).AddTime(etime.Trial, ttrls).AddTime(etime.Cycle, 100)

        leabra.LooperStdPhases(ls, &ss.Context, ss.Net, 75, 99)                // plus phase timing
        leabra.LooperSimCycleAndLearn(ls, ss.Net, &ss.Context, &ss.ViewUpdate) // std algo code
        ss.Net.ConfigLoopsHip(&ss.Context, ls)

        ls.Stacks[etime.Train].OnInit.Add("Init", func() { ss.Init() })
        ls.Stacks[etime.Test].OnInit.Add("Init", func() { ss.TestInit() })

        for _, st := range ls.Stacks {
                st.Loops[etime.Trial].OnStart.Add("ApplyInputs", func() {
                        ss.ApplyInputs()
                })
        }

        ls.Loop(etime.Train, etime.Run).OnStart.Add("NewRun", ss.NewRun)

        ls.Loop(etime.Train, etime.Run).OnEnd.Add("RunDone", func() {
                if ss.Stats.Int("Run") >= ss.Config.NRuns-1 {
                        ss.RunStats()
                        expt := ss.Stats.Int("Expt")
                        ss.Stats.SetInt("Expt", expt+1)
                }
        })

        // Add Testing
        trainEpoch := ls.Loop(etime.Train, etime.Epoch)
        trainEpoch.OnEnd.Add("TestAtInterval", func() {
                if (ss.Config.TestInterval > 0) && ((trainEpoch.Counter.Cur+1)%ss.Config.TestInterval == 0) {
                        // Note the +1 so that it doesn't occur at the 0th timestep.
                        ss.RunTestAll()

                        // switch to AC
                        trn := ss.Envs.ByMode(etime.Train).(*env.FixedTable)
                        tstEpcLog := ss.Logs.Tables[etime.Scope(etime.Test, etime.Epoch)]
                        epc := ss.Stats.Int("Epoch")
                        abMem := float32(tstEpcLog.Table.Float("ABMem", epc))
                        if (trn.Table.Table.MetaData["name"] == "TrainAB") && (abMem >= ss.Config.StopMem || epc >= ss.Config.NEpochs/2) {
                                ss.Stats.SetInt("FirstPerfect", epc)
                                trn.Config(table.NewIndexView(ss.TrainAC))
                                trn.Validate()
                        }
                }
        })

        // early stop
        ls.Loop(etime.Train, etime.Epoch).IsDone.AddBool("ACMemStop", func() bool {
                // This is calculated in TrialStats
                tstEpcLog := ss.Logs.Tables[etime.Scope(etime.Test, etime.Epoch)]
                acMem := float32(tstEpcLog.Table.Float("ACMem", ss.Stats.Int("Epoch")))
                stop := acMem >= ss.Config.StopMem
                return stop
        })

        /////////////////////////////////////////////
        // Logging

        ls.Loop(etime.Test, etime.Epoch).OnEnd.Add("LogTestErrors", func() {
                leabra.LogTestErrors(&ss.Logs)
        })

        ls.AddOnEndToAll("Log", func(mode, time enums.Enum) {
                ss.Log(mode.(etime.Modes), time.(etime.Times))
        })
        leabra.LooperResetLogBelow(ls, &ss.Logs)

        leabra.LooperUpdateNetView(ls, &ss.ViewUpdate, ss.Net, ss.NetViewCounters)
        leabra.LooperUpdatePlots(ls, &ss.GUI)

        ls.Stacks[etime.Train].OnInit.Add("GUI-Init", func() { ss.GUI.UpdateWindow() })
        ls.Stacks[etime.Test].OnInit.Add("GUI-Init", func() { ss.GUI.UpdateWindow() })

        ss.Loops = ls
        fmt.Println(ls.DocString())
}

// ApplyInputs applies input patterns from given environment.
// It is good practice to have this be a separate method with appropriate
// args so that it can be used for various different contexts
// (training, testing, etc).
func (ss *Sim) ApplyInputs() {
        ctx := &ss.Context
        net := ss.Net
        ev := ss.Envs.ByMode(ctx.Mode).(*env.FixedTable)
        ecout := net.LayerByName("ECout")
        if ctx.Mode == etime.Train {
                ecout.Type = leabra.TargetLayer // clamp a plus phase during testing
        } else {
                ecout.Type = leabra.CompareLayer // don't clamp
        }
        ecout.UpdateExtFlags() // call this after updating type
        net.InitExt()
        lays := net.LayersByType(leabra.InputLayer, leabra.TargetLayer)
        ev.Step()
        // note: must save env state for logging / stats due to data parallel re-use of same env
        ss.Stats.SetString("TrialName", ev.TrialName.Cur)
        for _, lnm := range lays {
                ly := ss.Net.LayerByName(lnm)
                pats := ev.State(ly.Name)
                if pats != nil {
                        ly.ApplyExt(pats)
                }
        }
}

// NewRun intializes a new run of the model, using the TrainEnv.Run counter
// for the new run value
func (ss *Sim) NewRun() {
        ctx := &ss.Context
        ss.InitRandSeed(ss.Loops.Loop(etime.Train, etime.Run).Counter.Cur)
        // ss.ConfigPats()
        ss.ConfigEnv()
        ctx.Reset()
        ctx.Mode = etime.Train
        ss.Net.InitWeights()
        ss.InitStats()
        ss.StatCounters()
        ss.Logs.ResetLog(etime.Train, etime.Epoch)
        ss.Logs.ResetLog(etime.Test, etime.Epoch)
}

// TestAll runs through the full set of testing items
func (ss *Sim) RunTestAll() {
        ss.Envs.ByMode(etime.Test).Init(0)
        ss.Loops.ResetAndRun(etime.Test)
        ss.Loops.Mode = etime.Train // Important to reset Mode back to Train because this is called from within the Train Run.
}

/////////////////////////////////////////////////////////////////////////
//   Pats

// OpenPatAsset opens pattern file from embedded assets
func (ss *Sim) OpenPatAsset(dt *table.Table, fnm, name, desc string) error {
        dt.SetMetaData("name", name)
        dt.SetMetaData("desc", desc)
        err := dt.OpenFS(content, fnm, table.Tab)
        if errors.Log(err) == nil {
                for i := 1; i < dt.NumColumns(); i++ {
                        dt.Columns[i].SetMetaData("grid-fill", "0.9")
                }
        }
        return err
}

func (ss *Sim) OpenPatterns() {
        ss.OpenPatAsset(ss.TrainAB, "train_ab.tsv", "TrainAB", "AB Training Patterns")
        ss.OpenPatAsset(ss.TrainAC, "train_ac.tsv", "TrainAC", "AC Training Patterns")
        ss.OpenPatAsset(ss.TestAB, "test_ab.tsv", "TestAB", "AB Testing Patterns")
        ss.OpenPatAsset(ss.TestAC, "test_ac.tsv", "TestAC", "AC Testing Patterns")
        ss.OpenPatAsset(ss.TestLure, "test_lure.tsv", "TestLure", "Lure Testing Patterns")

        ss.TestAll = ss.TestAB.Clone()
        ss.TestAll.SetMetaData("name", "TestAll")
        ss.TestAll.AppendRows(ss.TestAC)
        ss.TestAll.AppendRows(ss.TestLure)
}

func (ss *Sim) ConfigPats() {
        // hp := &ss.Config.Hip
        ecY := 3               // hp.EC3NPool.Y
        ecX := 4               // hp.EC3NPool.X
        plY := 6               // hp.EC3NNrn.Y // good idea to get shorter vars when used frequently
        plX := 2               // hp.EC3NNrn.X // makes much more readable
        npats := 10            // ss.Config.NTrials
        pctAct := float32(.15) // ss.Config.Mod.ECPctAct
        minDiff := float32(.5) // ss.Config.Pat.MinDiffPct
        nOn := patgen.NFromPct(pctAct, plY*plX)
        ctxtFlipPct := float32(0.2)
        ctxtflip := patgen.NFromPct(ctxtFlipPct, nOn)
        patgen.AddVocabEmpty(ss.PoolVocab, "empty", npats, plY, plX)
        patgen.AddVocabPermutedBinary(ss.PoolVocab, "A", npats, plY, plX, pctAct, minDiff)
        patgen.AddVocabPermutedBinary(ss.PoolVocab, "B", npats, plY, plX, pctAct, minDiff)
        patgen.AddVocabPermutedBinary(ss.PoolVocab, "C", npats, plY, plX, pctAct, minDiff)
        patgen.AddVocabPermutedBinary(ss.PoolVocab, "lA", npats, plY, plX, pctAct, minDiff)
        patgen.AddVocabPermutedBinary(ss.PoolVocab, "lB", npats, plY, plX, pctAct, minDiff)
        patgen.AddVocabPermutedBinary(ss.PoolVocab, "ctxt", 3, plY, plX, pctAct, minDiff) // totally diff

        for i := 0; i < (ecY-1)*ecX*3; i++ { // 12 contexts! 1: 1 row of stimuli pats; 3: 3 diff ctxt bases
                list := i / ((ecY - 1) * ecX)
                ctxtNm := fmt.Sprintf("ctxt%d", i+1)
                tsr, _ := patgen.AddVocabRepeat(ss.PoolVocab, ctxtNm, npats, "ctxt", list)
                patgen.FlipBitsRows(tsr, ctxtflip, ctxtflip, 1, 0)
                //todo: also support drifting
                //solution 2: drift based on last trial (will require sequential learning)
                //patgen.VocabDrift(ss.PoolVocab, ss.NFlipBits, "ctxt"+strconv.Itoa(i+1))
        }

        patgen.InitPats(ss.TrainAB, "TrainAB", "TrainAB Pats", "Input", "ECout", npats, ecY, ecX, plY, plX)
        patgen.MixPats(ss.TrainAB, ss.PoolVocab, "Input", []string{"A", "B", "ctxt1", "ctxt2", "ctxt3", "ctxt4"})
        patgen.MixPats(ss.TrainAB, ss.PoolVocab, "ECout", []string{"A", "B", "ctxt1", "ctxt2", "ctxt3", "ctxt4"})

        patgen.InitPats(ss.TestAB, "TestAB", "TestAB Pats", "Input", "ECout", npats, ecY, ecX, plY, plX)
        patgen.MixPats(ss.TestAB, ss.PoolVocab, "Input", []string{"A", "empty", "ctxt1", "ctxt2", "ctxt3", "ctxt4"})
        patgen.MixPats(ss.TestAB, ss.PoolVocab, "ECout", []string{"A", "B", "ctxt1", "ctxt2", "ctxt3", "ctxt4"})

        patgen.InitPats(ss.TrainAC, "TrainAC", "TrainAC Pats", "Input", "ECout", npats, ecY, ecX, plY, plX)
        patgen.MixPats(ss.TrainAC, ss.PoolVocab, "Input", []string{"A", "C", "ctxt5", "ctxt6", "ctxt7", "ctxt8"})
        patgen.MixPats(ss.TrainAC, ss.PoolVocab, "ECout", []string{"A", "C", "ctxt5", "ctxt6", "ctxt7", "ctxt8"})

        patgen.InitPats(ss.TestAC, "TestAC", "TestAC Pats", "Input", "ECout", npats, ecY, ecX, plY, plX)
        patgen.MixPats(ss.TestAC, ss.PoolVocab, "Input", []string{"A", "empty", "ctxt5", "ctxt6", "ctxt7", "ctxt8"})
        patgen.MixPats(ss.TestAC, ss.PoolVocab, "ECout", []string{"A", "C", "ctxt5", "ctxt6", "ctxt7", "ctxt8"})

        patgen.InitPats(ss.PreTrainLure, "PreTrainLure", "PreTrainLure Pats", "Input", "ECout", npats, ecY, ecX, plY, plX)
        patgen.MixPats(ss.PreTrainLure, ss.PoolVocab, "Input", []string{"lA", "lB", "ctxt9", "ctxt10", "ctxt11", "ctxt12"}) // arbitrary ctxt here
        patgen.MixPats(ss.PreTrainLure, ss.PoolVocab, "ECout", []string{"lA", "lB", "ctxt9", "ctxt10", "ctxt11", "ctxt12"}) // arbitrary ctxt here

        patgen.InitPats(ss.TestLure, "TestLure", "TestLure Pats", "Input", "ECout", npats, ecY, ecX, plY, plX)
        patgen.MixPats(ss.TestLure, ss.PoolVocab, "Input", []string{"lA", "empty", "ctxt9", "ctxt10", "ctxt11", "ctxt12"}) // arbitrary ctxt here
        patgen.MixPats(ss.TestLure, ss.PoolVocab, "ECout", []string{"lA", "lB", "ctxt9", "ctxt10", "ctxt11", "ctxt12"})    // arbitrary ctxt here

        ss.TrainAll = ss.TrainAB.Clone()
        ss.TrainAll.AppendRows(ss.TrainAC)
        ss.TrainAll.AppendRows(ss.PreTrainLure)
        ss.TrainAll.MetaData["name"] = "TrainAll"
        ss.TrainAll.MetaData["desc"] = "All Training Patterns"

        ss.TestAll = ss.TestAB.Clone()
        ss.TestAll.AppendRows(ss.TestAC)
        ss.TestAll.MetaData["name"] = "TestAll"
        ss.TestAll.MetaData["desc"] = "All Testing Patterns"
}

////////////////////////////////////////////////////////////////////////////////////////////
//                 Stats

// InitStats initializes all the statistics.
// called at start of new run
func (ss *Sim) InitStats() {
        ss.Stats.SetString("TrialName", "")
        ss.Stats.SetFloat("TrgOnWasOffAll", 0.0)
        ss.Stats.SetFloat("TrgOnWasOffCmp", 0.0)
        ss.Stats.SetFloat("TrgOffWasOn", 0.0)
        ss.Stats.SetFloat("ABMem", 0.0)
        ss.Stats.SetFloat("ACMem", 0.0)
        ss.Stats.SetFloat("LureMem", 0.0)
        ss.Stats.SetFloat("Mem", 0.0)
        ss.Stats.SetInt("FirstPerfect", -1) // first epoch at when AB Mem is perfect

        ss.Logs.InitErrStats() // inits TrlErr, FirstZero, LastZero, NZero
}

// StatCounters saves current counters to Stats, so they are available for logging etc
// Also saves a string rep of them for ViewUpdate.Text
func (ss *Sim) StatCounters() {
        ctx := &ss.Context
        mode := ctx.Mode
        ss.Loops.Stacks[mode].CountersToStats(&ss.Stats)
        // always use training epoch..
        trnEpc := ss.Loops.Stacks[etime.Train].Loops[etime.Epoch].Counter.Cur
        ss.Stats.SetInt("Epoch", trnEpc)
        trl := ss.Stats.Int("Trial")
        ss.Stats.SetInt("Trial", trl)
        ss.Stats.SetInt("Cycle", int(ctx.Cycle))
        ss.Stats.SetString("TrialName", ss.Stats.String("TrialName"))
}

func (ss *Sim) NetViewCounters(tm etime.Times) {
        if ss.ViewUpdate.View == nil {
                return
        }
        if tm == etime.Trial {
                ss.TrialStats() // get trial stats for current di
        }
        ss.StatCounters()
        ss.ViewUpdate.Text = ss.Stats.Print([]string{"Run", "Epoch", "Trial", "TrialName", "Cycle"})
}

// TrialStats computes the trial-level statistics.
// Aggregation is done directly from log data.
func (ss *Sim) TrialStats() {
        ss.MemStats(ss.Loops.Mode.(etime.Modes))
}

// MemStats computes ActM vs. Target on ECout with binary counts
// must be called at end of 3rd quarter so that Target values are
// for the entire full pattern as opposed to the plus-phase target
// values clamped from ECin activations
func (ss *Sim) MemStats(mode etime.Modes) {
        memthr := 0.34 // ss.Config.Mod.MemThr
        ecout := ss.Net.LayerByName("ECout")
        inp := ss.Net.LayerByName("Input") // note: must be input b/c ECin can be active
        _ = inp
        nn := ecout.Shape.Len()
        actThr := float32(0.5)
        trgOnWasOffAll := 0.0 // all units
        trgOnWasOffCmp := 0.0 // only those that required completion, missing in ECin
        trgOffWasOn := 0.0    // should have been off
        cmpN := 0.0           // completion target
        trgOnN := 0.0
        trgOffN := 0.0
        actMi, _ := ecout.UnitVarIndex("ActM")
        targi, _ := ecout.UnitVarIndex("Targ")

        ss.Stats.SetFloat("ABMem", math.NaN())
        ss.Stats.SetFloat("ACMem", math.NaN())
        ss.Stats.SetFloat("LureMem", math.NaN())

        trialnm := ss.Stats.String("TrialName")
        isAB := strings.Contains(trialnm, "ab")
        isAC := strings.Contains(trialnm, "ac")

        for ni := 0; ni < nn; ni++ {
                actm := ecout.UnitValue1D(actMi, ni, 0)
                trg := ecout.UnitValue1D(targi, ni, 0) // full pattern target
                inact := inp.UnitValue1D(actMi, ni, 0)
                if trg < actThr { // trgOff
                        trgOffN += 1
                        if actm > actThr {
                                trgOffWasOn += 1
                        }
                } else { // trgOn
                        trgOnN += 1
                        if inact < actThr { // missing in ECin -- completion target
                                cmpN += 1
                                if actm < actThr {
                                        trgOnWasOffAll += 1
                                        trgOnWasOffCmp += 1
                                }
                        } else {
                                if actm < actThr {
                                        trgOnWasOffAll += 1
                                }
                        }
                }
        }
        trgOnWasOffAll /= trgOnN
        trgOffWasOn /= trgOffN
        if mode == etime.Train { // no compare
                if trgOnWasOffAll < memthr && trgOffWasOn < memthr {
                        ss.Stats.SetFloat("Mem", 1)
                } else {
                        ss.Stats.SetFloat("Mem", 0)
                }
        } else { // test
                if cmpN > 0 { // should be
                        trgOnWasOffCmp /= cmpN
                }
                mem := 0.0
                if trgOnWasOffCmp < memthr && trgOffWasOn < memthr {
                        mem = 1.0
                }
                ss.Stats.SetFloat("Mem", mem)
                switch {
                case isAB:
                        ss.Stats.SetFloat("ABMem", mem)
                case isAC:
                        ss.Stats.SetFloat("ACMem", mem)
                default:
                        ss.Stats.SetFloat("LureMem", mem)
                }

        }
        ss.Stats.SetFloat("TrgOnWasOffAll", trgOnWasOffAll)
        ss.Stats.SetFloat("TrgOnWasOffCmp", trgOnWasOffCmp)
        ss.Stats.SetFloat("TrgOffWasOn", trgOffWasOn)

}

func (ss *Sim) RunStats() {
        dt := ss.Logs.Table(etime.Train, etime.Run)
        runix := table.NewIndexView(dt)
        spl := split.GroupBy(runix, "Expt")
        split.DescColumn(spl, "TstABMem")
        st := spl.AggsToTableCopy(table.AddAggName)
        ss.Logs.MiscTables["RunStats"] = st
        plt := ss.GUI.Plots[etime.ScopeKey("RunStats")]

        st.SetMetaData("XAxis", "RunName")

        st.SetMetaData("Points", "true")

        st.SetMetaData("TstABMem:Mean:On", "+")
        st.SetMetaData("TstABMem:Mean:FixMin", "true")
        st.SetMetaData("TstABMem:Mean:FixMax", "true")
        st.SetMetaData("TstABMem:Mean:Min", "0")
        st.SetMetaData("TstABMem:Mean:Max", "1")
        st.SetMetaData("TstABMem:Min:On", "+")
        st.SetMetaData("TstABMem:Count:On", "-")

        plt.SetTable(st)
        plt.GoUpdatePlot()
}

//////////////////////////////////////////////////////////////////////////////
//                 Logging

func (ss *Sim) AddLogItems() {
        itemNames := []string{"TrgOnWasOffAll", "TrgOnWasOffCmp", "TrgOffWasOn", "Mem", "ABMem", "ACMem", "LureMem"}
        for _, st := range itemNames {
                stnm := st
                tonm := "Tst" + st
                ss.Logs.AddItem(&elog.Item{
                        Name: tonm,
                        Type: reflect.Float64,
                        Write: elog.WriteMap{
                                etime.Scope(etime.Train, etime.Epoch): func(ctx *elog.Context) {
                                        ctx.SetFloat64(ctx.ItemFloat(etime.Test, etime.Epoch, stnm))
                                },
                                etime.Scope(etime.Train, etime.Run): func(ctx *elog.Context) {
                                        ctx.SetFloat64(ctx.ItemFloat(etime.Test, etime.Epoch, stnm)) // take the last epoch
                                        // ctx.SetAgg(ctx.Mode, etime.Epoch, stats.Max) // stats.Max for max over epochs
                                }}})
        }
}

func (ss *Sim) ConfigLogs() {
        ss.Stats.SetString("RunName", ss.Params.RunName(0)) // used for naming logs, stats, etc

        ss.Logs.AddCounterItems(etime.Run, etime.Epoch, etime.Trial, etime.Cycle)
        ss.Logs.AddStatIntNoAggItem(etime.AllModes, etime.AllTimes, "Expt")
        ss.Logs.AddStatStringItem(etime.AllModes, etime.AllTimes, "RunName")
        ss.Logs.AddStatStringItem(etime.AllModes, etime.Trial, "TrialName")

        ss.Logs.AddStatAggItem("TrgOnWasOffAll", etime.Run, etime.Epoch, etime.Trial)
        ss.Logs.AddStatAggItem("TrgOnWasOffCmp", etime.Run, etime.Epoch, etime.Trial)
        ss.Logs.AddStatAggItem("TrgOffWasOn", etime.Run, etime.Epoch, etime.Trial)
        ss.Logs.AddStatAggItem("ABMem", etime.Run, etime.Epoch, etime.Trial)
        ss.Logs.AddStatAggItem("ACMem", etime.Run, etime.Epoch, etime.Trial)
        ss.Logs.AddStatAggItem("LureMem", etime.Run, etime.Epoch, etime.Trial)
        ss.Logs.AddStatAggItem("Mem", etime.Run, etime.Epoch, etime.Trial)
        ss.Logs.AddStatIntNoAggItem(etime.Train, etime.Run, "FirstPerfect")

        // ss.Logs.AddCopyFromFloatItems(etime.Train, etime.Epoch, etime.Test, etime.Epoch, "Tst", "PhaseDiff", "UnitErr", "PctCor", "PctErr", "TrgOnWasOffAll", "TrgOnWasOffCmp", "TrgOffWasOn", "Mem")
        ss.AddLogItems()

        ss.Logs.AddPerTrlMSec("PerTrlMSec", etime.Run, etime.Epoch, etime.Trial)

        layers := ss.Net.LayersByType(leabra.SuperLayer, leabra.CTLayer, leabra.TargetLayer)
        leabra.LogAddDiagnosticItems(&ss.Logs, layers, etime.Train, etime.Epoch, etime.Trial)
        leabra.LogInputLayer(&ss.Logs, ss.Net, etime.Train)

        // leabra.LogAddPCAItems(&ss.Logs, ss.Net, etime.Train, etime.Run, etime.Epoch, etime.Trial)

        ss.Logs.AddLayerTensorItems(ss.Net, "ActM", etime.Test, etime.Trial, "TargetLayer")
        ss.Logs.AddLayerTensorItems(ss.Net, "Act", etime.Test, etime.Trial, "TargetLayer")

        ss.Logs.PlotItems("ABMem", "ACMem", "LureMem")

        // ss.Logs.PlotItems("TrgOnWasOffAll", "TrgOnWasOffCmp", "ABMem", "ACMem", "TstTrgOnWasOffAll", "TstTrgOnWasOffCmp", "TstMem", "TstABMem", "TstACMem")

        ss.Logs.CreateTables()
        ss.Logs.SetMeta(etime.Train, etime.Run, "TrgOnWasOffAll:On", "-")
        ss.Logs.SetMeta(etime.Train, etime.Run, "TrgOnWasOffCmp:On", "-")
        ss.Logs.SetMeta(etime.Train, etime.Run, "ABMem:On", "-")
        ss.Logs.SetMeta(etime.Train, etime.Run, "ACMem:On", "-")
        ss.Logs.SetMeta(etime.Train, etime.Run, "LureMem:On", "-")
        ss.Logs.SetMeta(etime.Train, etime.Run, "TstTrgOnWasOffAll:On", "-")
        ss.Logs.SetMeta(etime.Train, etime.Run, "TstTrgOnWasOffCmp:On", "-")
        ss.Logs.SetMeta(etime.Train, etime.Run, "TstABMem:On", "+")
        ss.Logs.SetMeta(etime.Train, etime.Run, "TstACMem:On", "+")
        ss.Logs.SetMeta(etime.Train, etime.Run, "TstLureMem:On", "+")
        ss.Logs.SetMeta(etime.Train, etime.Run, "Type", "Bar")
        ss.Logs.SetMeta(etime.Train, etime.Epoch, "ABMem:On", "-")
        ss.Logs.SetMeta(etime.Train, etime.Epoch, "ACMem:On", "-")
        ss.Logs.SetMeta(etime.Train, etime.Epoch, "LureMem:On", "-")
        ss.Logs.SetMeta(etime.Train, etime.Epoch, "Mem:On", "+")
        ss.Logs.SetMeta(etime.Train, etime.Epoch, "TrgOnWasOffAll:On", "+")
        ss.Logs.SetMeta(etime.Train, etime.Epoch, "TrgOffWasOn:On", "+")
        ss.Logs.SetContext(&ss.Stats, ss.Net)
        // don't plot certain combinations we don't use
        ss.Logs.NoPlot(etime.Train, etime.Cycle)
        ss.Logs.NoPlot(etime.Test, etime.Cycle)
        ss.Logs.NoPlot(etime.Test, etime.Run)
        // note: Analyze not plotted by default
        ss.Logs.SetMeta(etime.Train, etime.Run, "LegendCol", "RunName")
}

// Log is the main logging function, handles special things for different scopes
func (ss *Sim) Log(mode etime.Modes, time etime.Times) {
        ctx := &ss.Context
        if mode != etime.Analyze {
                ctx.Mode = mode // Also set specifically in a Loop callback.
        }
        dt := ss.Logs.Table(mode, time)
        if dt == nil {
                return
        }
        row := dt.Rows

        switch {
        case time == etime.Cycle:
                return
        case time == etime.Trial:
                ss.TrialStats()
                ss.StatCounters()
                ss.Logs.LogRow(mode, time, row)
                return // don't do reg below
        }

        ss.Logs.LogRow(mode, time, row) // also logs to file, etc
}

////////////////////////////////////////////////////////////////////////////////////////////
//                 Gui

// ConfigGUI configures the Cogent Core GUI interface for this simulation.
func (ss *Sim) ConfigGUI() {
        title := "Hippocampus"
        ss.GUI.MakeBody(ss, "hip", title, `runs a hippocampus model on the AB-AC paired associate learning task. See <a href="https://github.com/CompCogNeuro/sims/blob/master/ch7/hip/README.md">README.md on GitHub</a>.</p>`)
        ss.GUI.CycleUpdateInterval = 10

        nv := ss.GUI.AddNetView("Network")
        nv.Options.Raster.Max = 100
        nv.Options.MaxRecs = 300
        nv.SetNet(ss.Net)
        ss.ViewUpdate.Config(nv, etime.Phase, etime.Phase)
        ss.GUI.ViewUpdate = &ss.ViewUpdate

        // nv.SceneXYZ().Camera.Pose.Pos.Set(0, 1, 2.75)
        // nv.SceneXYZ().Camera.LookAt(math32.Vec3(0, 0, 0), math32.Vec3(0, 1, 0))

        ss.GUI.AddPlots(title, &ss.Logs)

        stnm := "RunStats"
        dt := ss.Logs.MiscTable(stnm)
        bcp, _ := ss.GUI.Tabs.NewTab(stnm + " Plot")
        plt := plotcore.NewSubPlot(bcp)
        ss.GUI.Plots[etime.ScopeKey(stnm)] = plt
        plt.Options.Title = "Run Stats"
        plt.Options.XAxis = "RunName"
        plt.SetTable(dt)

        ss.GUI.FinalizeGUI(false)
}

func (ss *Sim) MakeToolbar(p *tree.Plan) {
        ss.GUI.AddLooperCtrl(p, ss.Loops)

        tree.Add(p, func(w *core.Separator) {})
        ss.GUI.AddToolbarItem(p, egui.ToolbarItem{Label: "Reset RunLog",
                Icon:    icons.Reset,
                Tooltip: "Reset the accumulated log of all Runs, which are tagged with the ParamSet used",
                Active:  egui.ActiveAlways,
                Func: func() {
                        ss.Logs.ResetLog(etime.Train, etime.Run)
                        ss.GUI.UpdatePlot(etime.Train, etime.Run)
                },
        })
        ////////////////////////////////////////////////
        tree.Add(p, func(w *core.Separator) {})
        ss.GUI.AddToolbarItem(p, egui.ToolbarItem{Label: "New Seed",
                Icon:    icons.Add,
                Tooltip: "Generate a new initial random seed to get different results.  By default, Init re-establishes the same initial seed every time.",
                Active:  egui.ActiveAlways,
                Func: func() {
                        ss.RandSeeds.NewSeeds()
                },
        })
        ss.GUI.AddToolbarItem(p, egui.ToolbarItem{Label: "README",
                Icon:    icons.FileMarkdown,
                Tooltip: "Opens your browser on the README file that contains instructions for how to run this model.",
                Active:  egui.ActiveAlways,
                Func: func() {
                        core.TheApp.OpenURL("https://github.com/CompCogNeuro/sims/blob/master/ch7/hip/README.md")
                },
        })
}

func (ss *Sim) RunGUI() {
        ss.Init()
        ss.ConfigGUI()
        ss.GUI.Body.RunMainWindow()
}

// Copyright (c) 2019, The Emergent Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// ra25 runs a simple random-associator four-layer leabra network
// that uses the standard supervised learning paradigm to learn
// mappings between 25 random input / output patterns
// defined over 5x5 input / output layers (i.e., 25 units)
package main

//go:generate core generate -add-types

import (
        "embed"
        "log"
        "os"

        "cogentcore.org/core/core"
        "cogentcore.org/core/enums"
        "cogentcore.org/core/icons"
        "cogentcore.org/core/math32"
        "cogentcore.org/core/math32/vecint"
        "cogentcore.org/core/tree"
        "cogentcore.org/lab/base/mpi"
        "cogentcore.org/lab/base/randx"
        "github.com/emer/emergent/v2/econfig"
        "github.com/emer/emergent/v2/egui"
        "github.com/emer/emergent/v2/elog"
        "github.com/emer/emergent/v2/emer"
        "github.com/emer/emergent/v2/env"
        "github.com/emer/emergent/v2/estats"
        "github.com/emer/emergent/v2/etime"
        "github.com/emer/emergent/v2/looper"
        "github.com/emer/emergent/v2/netview"
        "github.com/emer/emergent/v2/params"
        "github.com/emer/emergent/v2/patgen"
        "github.com/emer/emergent/v2/paths"
        "github.com/emer/etensor/tensor"
        "github.com/emer/etensor/tensor/table"
        "github.com/emer/leabra/v2/leabra"
)

//go:embed *.tsv
var patsfs embed.FS

func main() {
        sim := &Sim{}
        sim.New()
        sim.ConfigAll()
        if sim.Config.GUI {
                sim.RunGUI()
        } else {
                sim.RunNoGUI()
        }
}

// ParamSets is the default set of parameters.
// Base is always applied, and others can be optionally
// selected to apply on top of that.
var ParamSets = params.Sets{
        "Base": {
                {Sel: "Path", Desc: "norm and momentum on works better, but wt bal is not better for smaller nets",
                        Params: params.Params{
                                "Path.Learn.Norm.On":     "true",
                                "Path.Learn.Momentum.On": "true",
                                "Path.Learn.WtBal.On":    "true", // no diff really
                                // "Path.Learn.WtBal.Targs": "true", // no diff here
                        }},
                {Sel: "Layer", Desc: "using default 1.8 inhib for all of network -- can explore",
                        Params: params.Params{
                                "Layer.Inhib.Layer.Gi": "1.8",
                                "Layer.Act.Init.Decay": "0.0",
                                "Layer.Act.Gbar.L":     "0.1", // set explictly, new default, a bit better vs 0.2
                        }},
                {Sel: ".BackPath", Desc: "top-down back-pathways MUST have lower relative weight scale, otherwise network hallucinates",
                        Params: params.Params{
                                "Path.WtScale.Rel": "0.2",
                        }},
                {Sel: "#Output", Desc: "output definitely needs lower inhib -- true for smaller layers in general",
                        Params: params.Params{
                                "Layer.Inhib.Layer.Gi": "1.4",
                        }},
        },
        "DefaultInhib": {
                {Sel: "#Output", Desc: "go back to default",
                        Params: params.Params{
                                "Layer.Inhib.Layer.Gi": "1.8",
                        }},
        },
        "NoMomentum": {
                {Sel: "Path", Desc: "no norm or momentum",
                        Params: params.Params{
                                "Path.Learn.Norm.On":     "false",
                                "Path.Learn.Momentum.On": "false",
                        }},
        },
        "WtBalOn": {
                {Sel: "Path", Desc: "weight bal on",
                        Params: params.Params{
                                "Path.Learn.WtBal.On": "true",
                        }},
        },
}

// ParamConfig has config parameters related to sim params
type ParamConfig struct {

        // network parameters
        Network map[string]any

        // size of hidden layer -- can use emer.LaySize for 4D layers
        Hidden1Size vecint.Vector2i `default:"{'X':7,'Y':7}" nest:"+"`

        // size of hidden layer -- can use emer.LaySize for 4D layers
        Hidden2Size vecint.Vector2i `default:"{'X':7,'Y':7}" nest:"+"`

        // Extra Param Sheet name(s) to use (space separated if multiple).
        // must be valid name as listed in compiled-in params or loaded params
        Sheet string

        // extra tag to add to file names and logs saved from this run
        Tag string

        // user note -- describe the run params etc -- like a git commit message for the run
        Note string

        // Name of the JSON file to input saved parameters from.
        File string `nest:"+"`

        // Save a snapshot of all current param and config settings
        // in a directory named params_<datestamp> (or _good if Good is true), then quit.
        // Useful for comparing to later changes and seeing multiple views of current params.
        SaveAll bool `nest:"+"`

        // For SaveAll, save to params_good for a known good params state.
        // This can be done prior to making a new release after all tests are passing.
        // add results to git to provide a full diff record of all params over time.
        Good bool `nest:"+"`
}

// RunConfig has config parameters related to running the sim
type RunConfig struct {
        // starting run number, which determines the random seed.
        // runs counts from there, can do all runs in parallel by launching
        // separate jobs with each run, runs = 1.
        Run int `default:"0"`

        // total number of runs to do when running Train
        NRuns int `default:"5" min:"1"`

        // total number of epochs per run
        NEpochs int `default:"100"`

        // stop run after this number of perfect, zero-error epochs.
        NZero int `default:"2"`

        // total number of trials per epoch.  Should be an even multiple of NData.
        NTrials int `default:"32"`

        // how often to run through all the test patterns, in terms of training epochs.
        // can use 0 or -1 for no testing.
        TestInterval int `default:"5"`

        // how frequently (in epochs) to compute PCA on hidden representations
        // to measure variance?
        PCAInterval int `default:"5"`

        // if non-empty, is the name of weights file to load at start
        // of first run, for testing.
        StartWts string
}

// LogConfig has config parameters related to logging data
type LogConfig struct {

        // if true, save final weights after each run
        SaveWeights bool

        // if true, save train epoch log to file, as .epc.tsv typically
        Epoch bool `default:"true" nest:"+"`

        // if true, save run log to file, as .run.tsv typically
        Run bool `default:"true" nest:"+"`

        // if true, save train trial log to file, as .trl.tsv typically. May be large.
        Trial bool `default:"false" nest:"+"`

        // if true, save testing epoch log to file, as .tst_epc.tsv typically.  In general it is better to copy testing items over to the training epoch log and record there.
        TestEpoch bool `default:"false" nest:"+"`

        // if true, save testing trial log to file, as .tst_trl.tsv typically. May be large.
        TestTrial bool `default:"false" nest:"+"`

        // if true, save network activation etc data from testing trials,
        // for later viewing in netview.
        NetData bool
}

// Config is a standard Sim config -- use as a starting point.
type Config struct {

        // specify include files here, and after configuration,
        // it contains list of include files added.
        Includes []string

        // open the GUI -- does not automatically run -- if false,
        // then runs automatically and quits.
        GUI bool `default:"true"`

        // log debugging information
        Debug bool

        // parameter related configuration options
        Params ParamConfig `display:"add-fields"`

        // sim running related configuration options
        Run RunConfig `display:"add-fields"`

        // data logging related configuration options
        Log LogConfig `display:"add-fields"`
}

func (cfg *Config) IncludesPtr() *[]string { return &cfg.Includes }

// Sim encapsulates the entire simulation model, and we define all the
// functionality as methods on this struct.  This structure keeps all relevant
// state information organized and available without having to pass everything around
// as arguments to methods, and provides the core GUI interface (note the view tags
// for the fields which provide hints to how things should be displayed).
type Sim struct {

        // simulation configuration parameters -- set by .toml config file and / or args
        Config Config `new-window:"+"`

        // the network -- click to view / edit parameters for layers, paths, etc
        Net *leabra.Network `new-window:"+" display:"no-inline"`

        // network parameter management
        Params emer.NetParams `display:"add-fields"`

        // contains looper control loops for running sim
        Loops *looper.Stacks `new-window:"+" display:"no-inline"`

        // contains computed statistic values
        Stats estats.Stats `new-window:"+"`

        // Contains all the logs and information about the logs.'
        Logs elog.Logs `new-window:"+"`

        // the training patterns to use
        Patterns *table.Table `new-window:"+" display:"no-inline"`

        // Environments
        Envs env.Envs `new-window:"+" display:"no-inline"`

        // leabra timing parameters and state
        Context leabra.Context `new-window:"+"`

        // netview update parameters
        ViewUpdate netview.ViewUpdate `display:"add-fields"`

        // manages all the gui elements
        GUI egui.GUI `display:"-"`

        // a list of random seeds to use for each run
        RandSeeds randx.Seeds `display:"-"`
}

// New creates new blank elements and initializes defaults
func (ss *Sim) New() {
        econfig.Config(&ss.Config, "config.toml")
        ss.Net = leabra.NewNetwork("RA25")
        ss.Params.Config(ParamSets, ss.Config.Params.Sheet, ss.Config.Params.Tag, ss.Net)
        ss.Stats.Init()
        ss.Patterns = &table.Table{}
        ss.RandSeeds.Init(100) // max 100 runs
        ss.InitRandSeed(0)
        ss.Context.Defaults()
}

//////////////////////////////////////////////////////////////////////////////
//                 Configs

// ConfigAll configures all the elements using the standard functions
func (ss *Sim) ConfigAll() {
        // ss.ConfigPatterns()
        ss.OpenPatterns()
        ss.ConfigEnv()
        ss.ConfigNet(ss.Net)
        ss.ConfigLogs()
        ss.ConfigLoops()
        if ss.Config.Params.SaveAll {
                ss.Config.Params.SaveAll = false
                ss.Net.SaveParamsSnapshot(&ss.Params.Params, &ss.Config, ss.Config.Params.Good)
                os.Exit(0)
        }
}

func (ss *Sim) ConfigEnv() {
        // Can be called multiple times -- don't re-create
        var trn, tst *env.FixedTable
        if len(ss.Envs) == 0 {
                trn = &env.FixedTable{}
                tst = &env.FixedTable{}
        } else {
                trn = ss.Envs.ByMode(etime.Train).(*env.FixedTable)
                tst = ss.Envs.ByMode(etime.Test).(*env.FixedTable)
        }

        // note: names must be standard here!
        trn.Name = etime.Train.String()
        trn.Config(table.NewIndexView(ss.Patterns))
        trn.Validate()

        tst.Name = etime.Test.String()
        tst.Config(table.NewIndexView(ss.Patterns))
        tst.Sequential = true
        tst.Validate()

        // note: to create a train / test split of pats, do this:
        // all := table.NewIndexView(ss.Patterns)
        // splits, _ := split.Permuted(all, []float64{.8, .2}, []string{"Train", "Test"})
        // trn.Table = splits.Splits[0]
        // tst.Table = splits.Splits[1]

        trn.Init(0)
        tst.Init(0)

        // note: names must be in place when adding
        ss.Envs.Add(trn, tst)
}

func (ss *Sim) ConfigNet(net *leabra.Network) {
        net.SetRandSeed(ss.RandSeeds[0]) // init new separate random seed, using run = 0

        inp := net.AddLayer2D("Input", 5, 5, leabra.InputLayer)
        inp.Doc = "Input represents sensory input, coming into the cortex via tha thalamus"
        hid1 := net.AddLayer2D("Hidden1", ss.Config.Params.Hidden1Size.Y, ss.Config.Params.Hidden1Size.X, leabra.SuperLayer)
        hid1.Doc = "First hidden layer performs initial internal processing of sensory inputs, transforming in preparation for producing appropriate responses"
        hid2 := net.AddLayer2D("Hidden2", ss.Config.Params.Hidden2Size.Y, ss.Config.Params.Hidden2Size.X, leabra.SuperLayer)
        hid2.Doc = "Another 'deep' layer of internal processing to prepare directly for Output response"
        out := net.AddLayer2D("Output", 5, 5, leabra.TargetLayer)
        out.Doc = "Output represents motor output response, via deep layer 5 neurons projecting supcortically, in motor cortex"

        // use this to position layers relative to each other
        // hid2.PlaceRightOf(hid1, 2)

        // note: see emergent/path module for all the options on how to connect
        // NewFull returns a new paths.Full connectivity pattern
        full := paths.NewFull()

        net.ConnectLayers(inp, hid1, full, leabra.ForwardPath)
        net.BidirConnectLayers(hid1, hid2, full)
        net.BidirConnectLayers(hid2, out, full)

        // net.LateralConnectLayerPath(hid1, full, &leabra.HebbPath{}).SetType(InhibPath)

        // note: if you wanted to change a layer type from e.g., Target to Compare, do this:
        // out.SetType(emer.Compare)
        // that would mean that the output layer doesn't reflect target values in plus phase
        // and thus removes error-driven learning -- but stats are still computed.

        net.Build()
        net.Defaults()
        ss.ApplyParams()
        net.InitWeights()
}

func (ss *Sim) ApplyParams() {
        ss.Params.SetAll()
        if ss.Config.Params.Network != nil {
                ss.Params.SetNetworkMap(ss.Net, ss.Config.Params.Network)
        }
}

////////////////////////////////////////////////////////////////////////////////
//             Init, utils

// Init restarts the run, and initializes everything, including network weights
// and resets the epoch log table
func (ss *Sim) Init() {
        if ss.Config.GUI {
                ss.Stats.SetString("RunName", ss.Params.RunName(0)) // in case user interactively changes tag
        }
        ss.Loops.ResetCounters()
        ss.InitRandSeed(0)
        // ss.ConfigEnv() // re-config env just in case a different set of patterns was
        // selected or patterns have been modified etc
        ss.GUI.StopNow = false
        ss.ApplyParams()
        ss.NewRun()
        ss.ViewUpdate.RecordSyns()
        ss.ViewUpdate.Update()
}

// InitRandSeed initializes the random seed based on current training run number
func (ss *Sim) InitRandSeed(run int) {
        ss.RandSeeds.Set(run)
        ss.RandSeeds.Set(run, &ss.Net.Rand)
}

// ConfigLoops configures the control loops: Training, Testing
func (ss *Sim) ConfigLoops() {
        ls := looper.NewStacks()

        trls := ss.Config.Run.NTrials

        ls.AddStack(etime.Train).
                AddTime(etime.Run, ss.Config.Run.NRuns).
                AddTime(etime.Epoch, ss.Config.Run.NEpochs).
                AddTime(etime.Trial, trls).
                AddTime(etime.Cycle, 100)

        ls.AddStack(etime.Test).
                AddTime(etime.Epoch, 1).
                AddTime(etime.Trial, trls).
                AddTime(etime.Cycle, 100)

        leabra.LooperStdPhases(ls, &ss.Context, ss.Net, 75, 99)                // plus phase timing
        leabra.LooperSimCycleAndLearn(ls, ss.Net, &ss.Context, &ss.ViewUpdate) // std algo code

        ls.Stacks[etime.Train].OnInit.Add("Init", func() { ss.Init() })

        for m, _ := range ls.Stacks {
                st := ls.Stacks[m]
                st.Loops[etime.Trial].OnStart.Add("ApplyInputs", func() {
                        ss.ApplyInputs()
                })
        }

        ls.Loop(etime.Train, etime.Run).OnStart.Add("NewRun", ss.NewRun)

        // Train stop early condition
        ls.Loop(etime.Train, etime.Epoch).IsDone.AddBool("NZeroStop", func() bool {
                // This is calculated in TrialStats
                stopNz := ss.Config.Run.NZero
                if stopNz <= 0 {
                        stopNz = 2
                }
                curNZero := ss.Stats.Int("NZero")
                stop := curNZero >= stopNz
                return stop
        })

        // Add Testing
        trainEpoch := ls.Loop(etime.Train, etime.Epoch)
        trainEpoch.OnStart.Add("TestAtInterval", func() {
                if (ss.Config.Run.TestInterval > 0) && ((trainEpoch.Counter.Cur+1)%ss.Config.Run.TestInterval == 0) {
                        // Note the +1 so that it doesn't occur at the 0th timestep.
                        ss.TestAll()
                }
        })

        /////////////////////////////////////////////
        // Logging

        ls.Loop(etime.Test, etime.Epoch).OnEnd.Add("LogTestErrors", func() {
                leabra.LogTestErrors(&ss.Logs)
        })
        ls.Loop(etime.Train, etime.Epoch).OnEnd.Add("PCAStats", func() {
                trnEpc := ls.Stacks[etime.Train].Loops[etime.Epoch].Counter.Cur
                if ss.Config.Run.PCAInterval > 0 && trnEpc%ss.Config.Run.PCAInterval == 0 {
                        leabra.PCAStats(ss.Net, &ss.Logs, &ss.Stats)
                        ss.Logs.ResetLog(etime.Analyze, etime.Trial)
                }
        })

        ls.AddOnEndToAll("Log", func(mode, time enums.Enum) {
                ss.Log(mode.(etime.Modes), time.(etime.Times))
        })
        leabra.LooperResetLogBelow(ls, &ss.Logs)

        ls.Loop(etime.Train, etime.Trial).OnEnd.Add("LogAnalyze", func() {
                trnEpc := ls.Stacks[etime.Train].Loops[etime.Epoch].Counter.Cur
                if (ss.Config.Run.PCAInterval > 0) && (trnEpc%ss.Config.Run.PCAInterval == 0) {
                        ss.Log(etime.Analyze, etime.Trial)
                }
        })

        ls.Loop(etime.Train, etime.Run).OnEnd.Add("RunStats", func() {
                ss.Logs.RunStats("PctCor", "FirstZero", "LastZero")
        })

        // Save weights to file, to look at later
        ls.Loop(etime.Train, etime.Run).OnEnd.Add("SaveWeights", func() {
                ctrString := ss.Stats.PrintValues([]string{"Run", "Epoch"}, []string{"%03d", "%05d"}, "_")
                leabra.SaveWeightsIfConfigSet(ss.Net, ss.Config.Log.SaveWeights, ctrString, ss.Stats.String("RunName"))
        })

        ////////////////////////////////////////////
        // GUI

        if !ss.Config.GUI {
                if ss.Config.Log.NetData {
                        ls.Loop(etime.Test, etime.Trial).OnEnd.Add("NetDataRecord", func() {
                                ss.GUI.NetDataRecord(ss.ViewUpdate.Text)
                        })
                }
        } else {
                leabra.LooperUpdateNetView(ls, &ss.ViewUpdate, ss.Net, ss.NetViewCounters)
                leabra.LooperUpdatePlots(ls, &ss.GUI)
                ls.Stacks[etime.Train].OnInit.Add("GUI-Init", func() { ss.GUI.UpdateWindow() })
                ls.Stacks[etime.Test].OnInit.Add("GUI-Init", func() { ss.GUI.UpdateWindow() })
        }

        if ss.Config.Debug {
                mpi.Println(ls.DocString())
        }
        ss.Loops = ls
}

// ApplyInputs applies input patterns from given environment.
// It is good practice to have this be a separate method with appropriate
// args so that it can be used for various different contexts
// (training, testing, etc).
func (ss *Sim) ApplyInputs() {
        ctx := &ss.Context
        net := ss.Net
        ev := ss.Envs.ByMode(ctx.Mode).(*env.FixedTable)
        ev.Step()
        lays := net.LayersByType(leabra.InputLayer, leabra.TargetLayer)
        net.InitExt()
        ss.Stats.SetString("TrialName", ev.TrialName.Cur)
        for _, lnm := range lays {
                ly := ss.Net.LayerByName(lnm)
                pats := ev.State(ly.Name)
                if pats != nil {
                        ly.ApplyExt(pats)
                }
        }
}

// NewRun intializes a new run of the model, using the TrainEnv.Run counter
// for the new run value
func (ss *Sim) NewRun() {
        ctx := &ss.Context
        ss.InitRandSeed(ss.Loops.Loop(etime.Train, etime.Run).Counter.Cur)
        ss.Envs.ByMode(etime.Train).Init(0)
        ss.Envs.ByMode(etime.Test).Init(0)
        ctx.Reset()
        ctx.Mode = etime.Train
        ss.Net.InitWeights()
        ss.InitStats()
        ss.StatCounters()
        ss.Logs.ResetLog(etime.Train, etime.Epoch)
        ss.Logs.ResetLog(etime.Test, etime.Epoch)
}

// TestAll runs through the full set of testing items
func (ss *Sim) TestAll() {
        ss.Envs.ByMode(etime.Test).Init(0)
        ss.Loops.ResetAndRun(etime.Test)
        ss.Loops.Mode = etime.Train // Important to reset Mode back to Train because this is called from within the Train Run.
}

/////////////////////////////////////////////////////////////////////////
//   Patterns

func (ss *Sim) ConfigPatterns() {
        dt := ss.Patterns
        dt.SetMetaData("name", "TrainPatterns")
        dt.SetMetaData("desc", "Training patterns")
        dt.AddStringColumn("Name")
        dt.AddFloat32TensorColumn("Input", []int{5, 5}, "Y", "X")
        dt.AddFloat32TensorColumn("Output", []int{5, 5}, "Y", "X")
        dt.SetNumRows(25)

        patgen.PermutedBinaryMinDiff(dt.Columns[1].(*tensor.Float32), 6, 1, 0, 3)
        patgen.PermutedBinaryMinDiff(dt.Columns[2].(*tensor.Float32), 6, 1, 0, 3)
        dt.SaveCSV("random_5x5_25_gen.tsv", table.Tab, table.Headers)
}

func (ss *Sim) OpenPatterns() {
        dt := ss.Patterns
        dt.SetMetaData("name", "TrainPatterns")
        dt.SetMetaData("desc", "Training patterns")
        err := dt.OpenFS(patsfs, "random_5x5_25.tsv", table.Tab)
        if err != nil {
                log.Println(err)
        }
}

////////////////////////////////////////////////////////////////////////////////////////////
//                 Stats

// InitStats initializes all the statistics.
// called at start of new run
func (ss *Sim) InitStats() {
        ss.Stats.SetFloat("UnitErr", 0.0)
        ss.Stats.SetFloat("CorSim", 0.0)
        ss.Stats.SetString("TrialName", "")
        ss.Logs.InitErrStats() // inits TrlErr, FirstZero, LastZero, NZero
}

// StatCounters saves current counters to Stats, so they are available for logging etc
// Also saves a string rep of them for ViewUpdate.Text
func (ss *Sim) StatCounters() {
        ctx := &ss.Context
        mode := ctx.Mode
        ss.Loops.Stacks[mode].CountersToStats(&ss.Stats)
        // always use training epoch..
        trnEpc := ss.Loops.Stacks[etime.Train].Loops[etime.Epoch].Counter.Cur
        ss.Stats.SetInt("Epoch", trnEpc)
        trl := ss.Stats.Int("Trial")
        ss.Stats.SetInt("Trial", trl)
        ss.Stats.SetInt("Cycle", int(ctx.Cycle))
}

func (ss *Sim) NetViewCounters(tm etime.Times) {
        if ss.ViewUpdate.View == nil {
                return
        }
        if tm == etime.Trial {
                ss.TrialStats() // get trial stats for current di
        }
        ss.StatCounters()
        ss.ViewUpdate.Text = ss.Stats.Print([]string{"Run", "Epoch", "Trial", "TrialName", "Cycle", "UnitErr", "TrlErr", "CorSim"})
}

// TrialStats computes the trial-level statistics.
// Aggregation is done directly from log data.
func (ss *Sim) TrialStats() {
        out := ss.Net.LayerByName("Output")

        ss.Stats.SetFloat("CorSim", float64(out.CosDiff.Cos))

        sse, avgsse := out.MSE(0.5) // 0.5 = per-unit tolerance -- right side of .5
        ss.Stats.SetFloat("SSE", sse)
        ss.Stats.SetFloat("AvgSSE", avgsse)
        if sse > 0 {
                ss.Stats.SetFloat("TrlErr", 1)
        } else {
                ss.Stats.SetFloat("TrlErr", 0)
        }
}

//////////////////////////////////////////////////////////////////////////////
//                 Logging

func (ss *Sim) ConfigLogs() {
        ss.Stats.SetString("RunName", ss.Params.RunName(0)) // used for naming logs, stats, etc

        ss.Logs.AddCounterItems(etime.Run, etime.Epoch, etime.Trial, etime.Cycle)
        ss.Logs.AddStatStringItem(etime.AllModes, etime.AllTimes, "RunName")
        ss.Logs.AddStatStringItem(etime.AllModes, etime.Trial, "TrialName")

        ss.Logs.AddStatAggItem("CorSim", etime.Run, etime.Epoch, etime.Trial)
        ss.Logs.AddStatAggItem("UnitErr", etime.Run, etime.Epoch, etime.Trial)
        ss.Logs.AddErrStatAggItems("TrlErr", etime.Run, etime.Epoch, etime.Trial)

        ss.Logs.AddCopyFromFloatItems(etime.Train, []etime.Times{etime.Epoch, etime.Run}, etime.Test, etime.Epoch, "Tst", "CorSim", "UnitErr", "PctCor", "PctErr")

        ss.Logs.AddPerTrlMSec("PerTrlMSec", etime.Run, etime.Epoch, etime.Trial)

        layers := ss.Net.LayersByType(leabra.SuperLayer, leabra.CTLayer, leabra.TargetLayer)
        leabra.LogAddDiagnosticItems(&ss.Logs, layers, etime.Train, etime.Epoch, etime.Trial)
        leabra.LogInputLayer(&ss.Logs, ss.Net, etime.Train)

        leabra.LogAddPCAItems(&ss.Logs, ss.Net, etime.Train, etime.Run, etime.Epoch, etime.Trial)

        ss.Logs.AddLayerTensorItems(ss.Net, "Act", etime.Test, etime.Trial, "InputLayer", "TargetLayer")

        ss.Logs.PlotItems("CorSim", "PctCor", "FirstZero", "LastZero")

        ss.Logs.CreateTables()
        ss.Logs.SetContext(&ss.Stats, ss.Net)
        // don't plot certain combinations we don't use
        ss.Logs.NoPlot(etime.Train, etime.Cycle)
        ss.Logs.NoPlot(etime.Test, etime.Run)
        // note: Analyze not plotted by default
        ss.Logs.SetMeta(etime.Train, etime.Run, "LegendCol", "RunName")
}

// Log is the main logging function, handles special things for different scopes
func (ss *Sim) Log(mode etime.Modes, time etime.Times) {
        ctx := &ss.Context
        if mode != etime.Analyze {
                ctx.Mode = mode // Also set specifically in a Loop callback.
        }
        dt := ss.Logs.Table(mode, time)
        if dt == nil {
                return
        }
        row := dt.Rows

        switch {
        case time == etime.Cycle:
                return
        case time == etime.Trial:
                ss.TrialStats()
                ss.StatCounters()
        }

        ss.Logs.LogRow(mode, time, row) // also logs to file, etc
}

////////////////////////////////////////////////////////////////////////////////////////////
//                 Gui

// ConfigGUI configures the Cogent Core GUI interface for this simulation.
func (ss *Sim) ConfigGUI() {
        title := "Leabra Random Associator"
        ss.GUI.MakeBody(ss, "ra25", title, `This demonstrates a basic Leabra model. See <a href="https://github.com/emer/emergent">emergent on GitHub</a>.</p>`)
        ss.GUI.CycleUpdateInterval = 10

        nv := ss.GUI.AddNetView("Network")
        nv.Options.MaxRecs = 300
        nv.SetNet(ss.Net)
        ss.ViewUpdate.Config(nv, etime.AlphaCycle, etime.AlphaCycle)
        ss.GUI.ViewUpdate = &ss.ViewUpdate

        nv.SceneXYZ().Camera.Pose.Pos.Set(0, 1, 2.75) // more "head on" than default which is more "top down"
        nv.SceneXYZ().Camera.LookAt(math32.Vec3(0, 0, 0), math32.Vec3(0, 1, 0))

        ss.GUI.AddPlots(title, &ss.Logs)

        ss.GUI.FinalizeGUI(false)
}

func (ss *Sim) MakeToolbar(p *tree.Plan) {
        ss.GUI.AddLooperCtrl(p, ss.Loops)

        ////////////////////////////////////////////////
        tree.Add(p, func(w *core.Separator) {})
        ss.GUI.AddToolbarItem(p, egui.ToolbarItem{Label: "Reset RunLog",
                Icon:    icons.Reset,
                Tooltip: "Reset the accumulated log of all Runs, which are tagged with the ParamSet used",
                Active:  egui.ActiveAlways,
                Func: func() {
                        ss.Logs.ResetLog(etime.Train, etime.Run)
                        ss.GUI.UpdatePlot(etime.Train, etime.Run)
                },
        })
        ////////////////////////////////////////////////
        tree.Add(p, func(w *core.Separator) {})
        ss.GUI.AddToolbarItem(p, egui.ToolbarItem{Label: "New Seed",
                Icon:    icons.Add,
                Tooltip: "Generate a new initial random seed to get different results.  By default, Init re-establishes the same initial seed every time.",
                Active:  egui.ActiveAlways,
                Func: func() {
                        ss.RandSeeds.NewSeeds()
                },
        })
        ss.GUI.AddToolbarItem(p, egui.ToolbarItem{Label: "README",
                Icon:    icons.FileMarkdown,
                Tooltip: "Opens your browser on the README file that contains instructions for how to run this model.",
                Active:  egui.ActiveAlways,
                Func: func() {
                        core.TheApp.OpenURL("https://github.com/emer/leabra/blob/main/examples/ra25/README.md")
                },
        })
}

func (ss *Sim) RunGUI() {
        ss.Init()
        ss.ConfigGUI()
        ss.GUI.Body.RunMainWindow()
}

func (ss *Sim) RunNoGUI() {
        if ss.Config.Params.Note != "" {
                mpi.Printf("Note: %s\n", ss.Config.Params.Note)
        }
        if ss.Config.Log.SaveWeights {
                mpi.Printf("Saving final weights per run\n")
        }
        runName := ss.Params.RunName(ss.Config.Run.Run)
        ss.Stats.SetString("RunName", runName) // used for naming logs, stats, etc
        netName := ss.Net.Name

        elog.SetLogFile(&ss.Logs, ss.Config.Log.Trial, etime.Train, etime.Trial, "trl", netName, runName)
        elog.SetLogFile(&ss.Logs, ss.Config.Log.Epoch, etime.Train, etime.Epoch, "epc", netName, runName)
        elog.SetLogFile(&ss.Logs, ss.Config.Log.Run, etime.Train, etime.Run, "run", netName, runName)
        elog.SetLogFile(&ss.Logs, ss.Config.Log.TestEpoch, etime.Test, etime.Epoch, "tst_epc", netName, runName)
        elog.SetLogFile(&ss.Logs, ss.Config.Log.TestTrial, etime.Test, etime.Trial, "tst_trl", netName, runName)

        netdata := ss.Config.Log.NetData
        if netdata {
                mpi.Printf("Saving NetView data from testing\n")
                ss.GUI.InitNetData(ss.Net, 200)
        }

        ss.Init()

        mpi.Printf("Running %d Runs starting at %d\n", ss.Config.Run.NRuns, ss.Config.Run.Run)
        ss.Loops.Loop(etime.Train, etime.Run).Counter.SetCurMaxPlusN(ss.Config.Run.Run, ss.Config.Run.NRuns)

        if ss.Config.Run.StartWts != "" { // this is just for testing -- not usually needed
                ss.Loops.Step(etime.Train, 1, etime.Trial) // get past NewRun
                ss.Net.OpenWeightsJSON(core.Filename(ss.Config.Run.StartWts))
                mpi.Printf("Starting with initial weights from: %s\n", ss.Config.Run.StartWts)
        }

        mpi.Printf("Set NThreads to: %d\n", ss.Net.NThreads)

        ss.Loops.Run(etime.Train)

        ss.Logs.CloseLogFiles()

        if netdata {
                ss.GUI.SaveNetData(ss.Stats.String("RunName"))
        }
}

// Code generated by "core generate -add-types"; DO NOT EDIT.

package main

import (
        "cogentcore.org/core/enums"
)

var _ActionsValues = []Actions{0, 1, 2, 3, 4}

// ActionsN is the highest valid value for type Actions, plus one.
const ActionsN Actions = 5

var _ActionsValueMap = map[string]Actions{`Store1`: 0, `Store2`: 1, `Ignore`: 2, `Recall1`: 3, `Recall2`: 4}

var _ActionsDescMap = map[Actions]string{0: ``, 1: ``, 2: ``, 3: ``, 4: ``}

var _ActionsMap = map[Actions]string{0: `Store1`, 1: `Store2`, 2: `Ignore`, 3: `Recall1`, 4: `Recall2`}

// String returns the string representation of this Actions value.
func (i Actions) String() string { return enums.String(i, _ActionsMap) }

// SetString sets the Actions value from its string representation,
// and returns an error if the string is invalid.
func (i *Actions) SetString(s string) error {
        return enums.SetString(i, s, _ActionsValueMap, "Actions")
}

// Int64 returns the Actions value as an int64.
func (i Actions) Int64() int64 { return int64(i) }

// SetInt64 sets the Actions value from an int64.
func (i *Actions) SetInt64(in int64) { *i = Actions(in) }

// Desc returns the description of the Actions value.
func (i Actions) Desc() string { return enums.Desc(i, _ActionsDescMap) }

// ActionsValues returns all possible values for the type Actions.
func ActionsValues() []Actions { return _ActionsValues }

// Values returns all possible values for the type Actions.
func (i Actions) Values() []enums.Enum { return enums.Values(_ActionsValues) }

// MarshalText implements the [encoding.TextMarshaler] interface.
func (i Actions) MarshalText() ([]byte, error) { return []byte(i.String()), nil }

// UnmarshalText implements the [encoding.TextUnmarshaler] interface.
func (i *Actions) UnmarshalText(text []byte) error { return enums.UnmarshalText(i, text, "Actions") }

// Copyright (c) 2024, The Emergent Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// sir illustrates the dynamic gating of information into PFC active
// maintenance, by the basal ganglia (BG). It uses a simple Store-Ignore-Recall
// (SIR) task, where the BG system learns via phasic dopamine signals
// and trial-and-error exploration, discovering what needs to be stored,
// ignored, and recalled as a function of reinforcement of correct behavior,
// and learned reinforcement of useful working memory representations.
package main

//go:generate core generate -add-types

import (
        "fmt"

        "cogentcore.org/core/core"
        "cogentcore.org/core/enums"
        "cogentcore.org/core/icons"
        "cogentcore.org/core/math32"
        "cogentcore.org/core/tree"
        "cogentcore.org/lab/base/randx"
        "github.com/emer/emergent/v2/econfig"
        "github.com/emer/emergent/v2/egui"
        "github.com/emer/emergent/v2/elog"
        "github.com/emer/emergent/v2/emer"
        "github.com/emer/emergent/v2/env"
        "github.com/emer/emergent/v2/estats"
        "github.com/emer/emergent/v2/etime"
        "github.com/emer/emergent/v2/looper"
        "github.com/emer/emergent/v2/netview"
        "github.com/emer/emergent/v2/params"
        "github.com/emer/emergent/v2/paths"
        "github.com/emer/leabra/v2/leabra"
)

func main() {
        sim := &Sim{}
        sim.New()
        sim.ConfigAll()
        sim.RunGUI()
}

// ParamSets is the default set of parameters.
// Base is always applied, and others can be optionally
// selected to apply on top of that.
var ParamSets = params.Sets{
        "Base": {
                {Sel: "Path", Desc: "no extra learning factors",
                        Params: params.Params{
                                "Path.Learn.Lrate":       "0.01", // slower overall is key
                                "Path.Learn.Norm.On":     "false",
                                "Path.Learn.Momentum.On": "false",
                                "Path.Learn.WtBal.On":    "false",
                        }},
                {Sel: "Layer", Desc: "no decay",
                        Params: params.Params{
                                "Layer.Act.Init.Decay": "0", // key for all layers not otherwise done automatically
                        }},
                {Sel: ".BackPath", Desc: "top-down back-projections MUST have lower relative weight scale, otherwise network hallucinates",
                        Params: params.Params{
                                "Path.WtScale.Rel": "0.2",
                        }},
                {Sel: ".BgFixed", Desc: "BG Matrix -> GP wiring",
                        Params: params.Params{
                                "Path.Learn.Learn": "false",
                                "Path.WtInit.Mean": "0.8",
                                "Path.WtInit.Var":  "0",
                                "Path.WtInit.Sym":  "false",
                        }},
                {Sel: ".RWPath", Desc: "Reward prediction -- into PVi",
                        Params: params.Params{
                                "Path.Learn.Lrate": "0.02",
                                "Path.WtInit.Mean": "0",
                                "Path.WtInit.Var":  "0",
                                "Path.WtInit.Sym":  "false",
                        }},
                {Sel: "#Rew", Desc: "Reward layer -- no clamp limits",
                        Params: params.Params{
                                "Layer.Act.Clamp.Range.Min": "-1",
                                "Layer.Act.Clamp.Range.Max": "1",
                        }},
                {Sel: ".PFCMntDToOut", Desc: "PFC MntD -> PFC Out fixed",
                        Params: params.Params{
                                "Path.Learn.Learn": "false",
                                "Path.WtInit.Mean": "0.8",
                                "Path.WtInit.Var":  "0",
                                "Path.WtInit.Sym":  "false",
                        }},
                {Sel: ".FmPFCOutD", Desc: "PFC OutD needs to be strong b/c avg act says weak",
                        Params: params.Params{
                                "Path.WtScale.Abs": "4",
                        }},
                {Sel: ".PFCFixed", Desc: "Input -> PFC",
                        Params: params.Params{
                                "Path.Learn.Learn": "false",
                                "Path.WtInit.Mean": "0.8",
                                "Path.WtInit.Var":  "0",
                                "Path.WtInit.Sym":  "false",
                        }},
                {Sel: ".MatrixPath", Desc: "Matrix learning",
                        Params: params.Params{
                                "Path.Learn.Lrate":         "0.04", // .04 > .1 > .02
                                "Path.WtInit.Var":          "0.1",
                                "Path.Trace.GateNoGoPosLR": "1",    // 0.1 default
                                "Path.Trace.NotGatedLR":    "0.7",  // 0.7 default
                                "Path.Trace.Decay":         "1.0",  // 1.0 default
                                "Path.Trace.AChDecay":      "0.0",  // not useful even at .1, surprising..
                                "Path.Trace.Deriv":         "true", // true default, better than false
                        }},
                {Sel: ".MatrixLayer", Desc: "exploring these options",
                        Params: params.Params{
                                "Layer.Act.XX1.Gain":       "100",
                                "Layer.Inhib.Layer.Gi":     "2.2", // 2.2 > 1.8 > 2.4
                                "Layer.Inhib.Layer.FB":     "1",   // 1 > .5
                                "Layer.Inhib.Pool.On":      "true",
                                "Layer.Inhib.Pool.Gi":      "2.1", // def 1.9
                                "Layer.Inhib.Pool.FB":      "0",
                                "Layer.Inhib.Self.On":      "true",
                                "Layer.Inhib.Self.Gi":      "0.4", // def 0.3
                                "Layer.Inhib.ActAvg.Init":  "0.05",
                                "Layer.Inhib.ActAvg.Fixed": "true",
                        }},
                {Sel: "#GPiThal", Desc: "defaults also set automatically by layer but included here just to be sure",
                        Params: params.Params{
                                "Layer.Inhib.Layer.Gi":     "1.8", // 1.8 > 2.0
                                "Layer.Inhib.Layer.FB":     "1",   // 1.0 > 0.5
                                "Layer.Inhib.Pool.On":      "false",
                                "Layer.Inhib.ActAvg.Init":  ".2",
                                "Layer.Inhib.ActAvg.Fixed": "true",
                                "Layer.Act.Dt.GTau":        "3",
                                "Layer.GPiGate.GeGain":     "3",
                                "Layer.GPiGate.NoGo":       "1.25", // was 1 default
                                "Layer.GPiGate.Thr":        "0.25", // .2 default
                        }},
                {Sel: "#GPeNoGo", Desc: "GPe is a regular layer -- needs special params",
                        Params: params.Params{
                                "Layer.Inhib.Layer.Gi":     "2.4", // 2.4 > 2.2 > 1.8 > 2.6
                                "Layer.Inhib.Layer.FB":     "0.5",
                                "Layer.Inhib.Layer.FBTau":  "3", // otherwise a bit jumpy
                                "Layer.Inhib.Pool.On":      "false",
                                "Layer.Inhib.ActAvg.Init":  ".2",
                                "Layer.Inhib.ActAvg.Fixed": "true",
                        }},
                {Sel: ".PFC", Desc: "pfc defaults",
                        Params: params.Params{
                                "Layer.Inhib.Layer.On":     "false",
                                "Layer.Inhib.Pool.On":      "true",
                                "Layer.Inhib.Pool.Gi":      "1.8",
                                "Layer.Inhib.Pool.FB":      "1",
                                "Layer.Inhib.ActAvg.Init":  "0.2",
                                "Layer.Inhib.ActAvg.Fixed": "true",
                        }},
                {Sel: "#Input", Desc: "Basic params",
                        Params: params.Params{
                                "Layer.Inhib.ActAvg.Init":  "0.25",
                                "Layer.Inhib.ActAvg.Fixed": "true",
                        }},
                {Sel: "#Output", Desc: "Basic params",
                        Params: params.Params{
                                "Layer.Inhib.Layer.Gi":     "2",
                                "Layer.Inhib.Layer.FB":     "0.5",
                                "Layer.Inhib.ActAvg.Init":  "0.25",
                                "Layer.Inhib.ActAvg.Fixed": "true",
                        }},
                {Sel: "#InputToOutput", Desc: "weaker",
                        Params: params.Params{
                                "Path.WtScale.Rel": "0.5",
                        }},
                {Sel: "#Hidden", Desc: "Basic params",
                        Params: params.Params{
                                "Layer.Inhib.Layer.Gi": "2",
                                "Layer.Inhib.Layer.FB": "0.5",
                        }},
                {Sel: "#SNc", Desc: "allow negative",
                        Params: params.Params{
                                "Layer.Act.Clamp.Range.Min": "-1",
                                "Layer.Act.Clamp.Range.Max": "1",
                        }},
                {Sel: "#RWPred", Desc: "keep it guessing",
                        Params: params.Params{
                                "Layer.RW.PredRange.Min": "0.02", // single most important param!  was .01 -- need penalty..
                                "Layer.RW.PredRange.Max": "0.95",
                        }},
        },
}

// Config has config parameters related to running the sim
type Config struct {
        // total number of runs to do when running Train
        NRuns int `default:"10" min:"1"`

        // total number of epochs per run
        NEpochs int `default:"200"`

        // total number of trials per epochs per run
        NTrials int `default:"100"`

        // stop run after this number of perfect, zero-error epochs.
        NZero int `default:"5"`

        // how often to run through all the test patterns, in terms of training epochs.
        // can use 0 or -1 for no testing.
        TestInterval int `default:"-1"`
}

// Sim encapsulates the entire simulation model, and we define all the
// functionality as methods on this struct.  This structure keeps all relevant
// state information organized and available without having to pass everything around
// as arguments to methods, and provides the core GUI interface (note the view tags
// for the fields which provide hints to how things should be displayed).
type Sim struct {

        // BurstDaGain is the strength of dopamine bursts: 1 default -- reduce for PD OFF, increase for PD ON
        BurstDaGain float32

        // DipDaGain is the strength of dopamine dips: 1 default -- reduce to siulate D2 agonists
        DipDaGain float32

        // Config contains misc configuration parameters for running the sim
        Config Config `new-window:"+" display:"no-inline"`

        // the network -- click to view / edit parameters for layers, paths, etc
        Net *leabra.Network `new-window:"+" display:"no-inline"`

        // network parameter management
        Params emer.NetParams `display:"add-fields"`

        // contains looper control loops for running sim
        Loops *looper.Stacks `new-window:"+" display:"no-inline"`

        // contains computed statistic values
        Stats estats.Stats `new-window:"+"`

        // Contains all the logs and information about the logs.'
        Logs elog.Logs `new-window:"+"`

        // Environments
        Envs env.Envs `new-window:"+" display:"no-inline"`

        // leabra timing parameters and state
        Context leabra.Context `new-window:"+"`

        // netview update parameters
        ViewUpdate netview.ViewUpdate `display:"add-fields"`

        // manages all the gui elements
        GUI egui.GUI `display:"-"`

        // a list of random seeds to use for each run
        RandSeeds randx.Seeds `display:"-"`
}

// New creates new blank elements and initializes defaults
func (ss *Sim) New() {
        ss.Defaults()
        econfig.Config(&ss.Config, "config.toml")
        ss.Net = leabra.NewNetwork("SIR")
        ss.Params.Config(ParamSets, "", "", ss.Net)
        ss.Stats.Init()
        ss.Stats.SetInt("Expt", 0)
        ss.RandSeeds.Init(100) // max 100 runs
        ss.InitRandSeed(0)
        ss.Context.Defaults()
}

func (ss *Sim) Defaults() {
        ss.BurstDaGain = 1
        ss.DipDaGain = 1
}

//////////////////////////////////////////////////////////////////////////////
//                 Configs

// ConfigAll configures all the elements using the standard functions
func (ss *Sim) ConfigAll() {
        ss.ConfigEnv()
        ss.ConfigNet(ss.Net)
        ss.ConfigLogs()
        ss.ConfigLoops()
}

func (ss *Sim) ConfigEnv() {
        // Can be called multiple times -- don't re-create
        var trn, tst *SIREnv
        if len(ss.Envs) == 0 {
                trn = &SIREnv{}
                tst = &SIREnv{}
        } else {
                trn = ss.Envs.ByMode(etime.Train).(*SIREnv)
                tst = ss.Envs.ByMode(etime.Test).(*SIREnv)
        }

        // note: names must be standard here!
        trn.Name = etime.Train.String()
        trn.SetNStim(4)
        trn.RewVal = 1
        trn.NoRewVal = 0
        trn.Trial.Max = ss.Config.NTrials

        tst.Name = etime.Test.String()
        tst.SetNStim(4)
        tst.RewVal = 1
        tst.NoRewVal = 0
        tst.Trial.Max = ss.Config.NTrials

        trn.Init(0)
        tst.Init(0)

        // note: names must be in place when adding
        ss.Envs.Add(trn, tst)
}

func (ss *Sim) ConfigNet(net *leabra.Network) {
        net.SetRandSeed(ss.RandSeeds[0]) // init new separate random seed, using run = 0

        rew, rp, da := net.AddRWLayers("", 2)
        da.Name = "SNc"

        inp := net.AddLayer2D("Input", 1, 4, leabra.InputLayer)
        ctrl := net.AddLayer2D("CtrlInput", 1, 5, leabra.InputLayer)
        out := net.AddLayer2D("Output", 1, 4, leabra.TargetLayer)
        hid := net.AddLayer2D("Hidden", 7, 7, leabra.SuperLayer)

        // args: nY, nMaint, nOut, nNeurBgY, nNeurBgX, nNeurPfcY, nNeurPfcX
        mtxGo, mtxNoGo, gpe, gpi, cin, pfcMnt, pfcMntD, pfcOut, pfcOutD := net.AddPBWM("", 4, 2, 2, 1, 5, 1, 4)
        _ = gpe
        _ = gpi
        _ = pfcMnt
        _ = pfcMntD
        _ = pfcOut
        _ = cin

        cin.CIN.RewLays.Add(rew.Name, rp.Name)

        full := paths.NewFull()
        fmin := paths.NewRect()
        fmin.Size.Set(1, 1)
        fmin.Scale.Set(1, 1)
        fmin.Wrap = true

        net.ConnectLayers(ctrl, rp, full, leabra.RWPath)
        net.ConnectLayers(pfcMntD, rp, full, leabra.RWPath)
        net.ConnectLayers(pfcOutD, rp, full, leabra.RWPath)

        net.ConnectLayers(ctrl, mtxGo, fmin, leabra.MatrixPath)
        net.ConnectLayers(ctrl, mtxNoGo, fmin, leabra.MatrixPath)
        pt := net.ConnectLayers(inp, pfcMnt, fmin, leabra.ForwardPath)
        pt.AddClass("PFCFixed")

        net.ConnectLayers(inp, hid, full, leabra.ForwardPath)
        net.ConnectLayers(ctrl, hid, full, leabra.ForwardPath)
        net.BidirConnectLayers(hid, out, full)
        pt = net.ConnectLayers(pfcOutD, hid, full, leabra.ForwardPath)
        pt.AddClass("FmPFCOutD")
        pt = net.ConnectLayers(pfcOutD, out, full, leabra.ForwardPath)
        pt.AddClass("FmPFCOutD")
        net.ConnectLayers(inp, out, full, leabra.ForwardPath)

        inp.PlaceAbove(rew)
        out.PlaceRightOf(inp, 2)
        ctrl.PlaceBehind(inp, 2)
        hid.PlaceBehind(ctrl, 2)
        mtxGo.PlaceRightOf(rew, 2)
        pfcMnt.PlaceRightOf(out, 2)

        net.Build()
        net.Defaults()

        da.AddAllSendToBut() // send dopamine to all layers..
        gpi.SendPBWMParams()

        ss.ApplyParams()
        net.InitWeights()
}

func (ss *Sim) ApplyParams() {
        if ss.Loops != nil {
                trn := ss.Loops.Stacks[etime.Train]
                trn.Loops[etime.Run].Counter.Max = ss.Config.NRuns
                trn.Loops[etime.Epoch].Counter.Max = ss.Config.NEpochs
        }
        ss.Params.SetAll()

        matg := ss.Net.LayerByName("MatrixGo")
        matn := ss.Net.LayerByName("MatrixNoGo")

        matg.Matrix.BurstGain = ss.BurstDaGain
        matg.Matrix.DipGain = ss.DipDaGain
        matn.Matrix.BurstGain = ss.BurstDaGain
        matn.Matrix.DipGain = ss.DipDaGain
}

////////////////////////////////////////////////////////////////////////////////
//             Init, utils

// Init restarts the run, and initializes everything, including network weights
// and resets the epoch log table
func (ss *Sim) Init() {
        ss.Stats.SetString("RunName", ss.Params.RunName(0)) // in case user interactively changes tag
        ss.Loops.ResetCounters()
        ss.InitRandSeed(0)
        ss.ConfigEnv() // re-config env just in case a different set of patterns was
        ss.GUI.StopNow = false
        ss.ApplyParams()
        ss.NewRun()
        ss.ViewUpdate.RecordSyns()
        ss.ViewUpdate.Update()
}

// InitRandSeed initializes the random seed based on current training run number
func (ss *Sim) InitRandSeed(run int) {
        ss.RandSeeds.Set(run)
        ss.RandSeeds.Set(run, &ss.Net.Rand)
}

// ConfigLoops configures the control loops: Training, Testing
func (ss *Sim) ConfigLoops() {
        ls := looper.NewStacks()

        trls := ss.Config.NTrials

        ls.AddStack(etime.Train).
                AddTime(etime.Run, ss.Config.NRuns).
                AddTime(etime.Epoch, ss.Config.NEpochs).
                AddTime(etime.Trial, trls).
                AddTime(etime.Cycle, 100)

        ls.AddStack(etime.Test).
                AddTime(etime.Epoch, 1).
                AddTime(etime.Trial, trls).
                AddTime(etime.Cycle, 100)

        leabra.LooperStdPhases(ls, &ss.Context, ss.Net, 75, 99)                // plus phase timing
        leabra.LooperSimCycleAndLearn(ls, ss.Net, &ss.Context, &ss.ViewUpdate) // std algo code

        ls.Stacks[etime.Train].OnInit.Add("Init", func() { ss.Init() })

        for m, _ := range ls.Stacks {
                stack := ls.Stacks[m]
                stack.Loops[etime.Trial].OnStart.Add("ApplyInputs", func() {
                        ss.ApplyInputs()
                })
        }

        ls.Loop(etime.Train, etime.Run).OnStart.Add("NewRun", ss.NewRun)

        ls.Loop(etime.Train, etime.Run).OnEnd.Add("RunDone", func() {
                if ss.Stats.Int("Run") >= ss.Config.NRuns-1 {
                        expt := ss.Stats.Int("Expt")
                        ss.Stats.SetInt("Expt", expt+1)
                }
        })

        stack := ls.Stacks[etime.Train]
        cyc, _ := stack.Loops[etime.Cycle]
        plus := cyc.EventByName("MinusPhase:End")
        plus.OnEvent.InsertBefore("MinusPhase:End", "ApplyReward", func() bool {
                ss.ApplyReward(true)
                return true
        })

        // Train stop early condition
        ls.Loop(etime.Train, etime.Epoch).IsDone.AddBool("NZeroStop", func() bool {
                // This is calculated in TrialStats
                stopNz := ss.Config.NZero
                if stopNz <= 0 {
                        stopNz = 2
                }
                curNZero := ss.Stats.Int("NZero")
                stop := curNZero >= stopNz
                return stop
        })

        // Add Testing
        trainEpoch := ls.Loop(etime.Train, etime.Epoch)
        trainEpoch.OnStart.Add("TestAtInterval", func() {
                if (ss.Config.TestInterval > 0) && ((trainEpoch.Counter.Cur+1)%ss.Config.TestInterval == 0) {
                        // Note the +1 so that it doesn't occur at the 0th timestep.
                        ss.TestAll()
                }
        })

        /////////////////////////////////////////////
        // Logging

        ls.Loop(etime.Test, etime.Epoch).OnEnd.Add("LogTestErrors", func() {
                leabra.LogTestErrors(&ss.Logs)
        })
        ls.AddOnEndToAll("Log", func(mode, time enums.Enum) {
                ss.Log(mode.(etime.Modes), time.(etime.Times))
        })
        leabra.LooperResetLogBelow(ls, &ss.Logs)
        ls.Loop(etime.Train, etime.Run).OnEnd.Add("RunStats", func() {
                ss.Logs.RunStats("PctCor", "FirstZero", "LastZero")
        })

        ////////////////////////////////////////////
        // GUI

        leabra.LooperUpdateNetView(ls, &ss.ViewUpdate, ss.Net, ss.NetViewCounters)
        leabra.LooperUpdatePlots(ls, &ss.GUI)
        ls.Stacks[etime.Train].OnInit.Add("GUI-Init", func() { ss.GUI.UpdateWindow() })
        ls.Stacks[etime.Test].OnInit.Add("GUI-Init", func() { ss.GUI.UpdateWindow() })

        ss.Loops = ls
}

// ApplyInputs applies input patterns from given environment.
// It is good practice to have this be a separate method with appropriate
// args so that it can be used for various different contexts
// (training, testing, etc).
func (ss *Sim) ApplyInputs() {
        ctx := &ss.Context
        net := ss.Net
        ev := ss.Envs.ByMode(ctx.Mode).(*SIREnv)
        ev.Step()

        lays := net.LayersByType(leabra.InputLayer, leabra.TargetLayer)
        net.InitExt()
        ss.Stats.SetString("TrialName", ev.String())
        for _, lnm := range lays {
                if lnm == "Rew" {
                        continue
                }
                ly := ss.Net.LayerByName(lnm)
                pats := ev.State(ly.Name)
                if pats != nil {
                        ly.ApplyExt(pats)
                }
        }
}

// ApplyReward computes reward based on network output and applies it.
// Call at start of 3rd quarter (plus phase).
func (ss *Sim) ApplyReward(train bool) {
        var en *SIREnv
        if train {
                en = ss.Envs.ByMode(etime.Train).(*SIREnv)
        } else {
                en = ss.Envs.ByMode(etime.Test).(*SIREnv)
        }
        if en.Act != Recall1 && en.Act != Recall2 { // only reward on recall trials!
                return
        }
        out := ss.Net.LayerByName("Output")
        mxi := out.Pools[0].Inhib.Act.MaxIndex
        en.SetReward(int(mxi))
        pats := en.State("Rew")
        ly := ss.Net.LayerByName("Rew")
        ly.ApplyExt1DTsr(pats)
}

// NewRun intializes a new run of the model, using the TrainEnv.Run counter
// for the new run value
func (ss *Sim) NewRun() {
        ctx := &ss.Context
        ss.InitRandSeed(ss.Loops.Loop(etime.Train, etime.Run).Counter.Cur)
        ss.Envs.ByMode(etime.Train).Init(0)
        ss.Envs.ByMode(etime.Test).Init(0)
        ctx.Reset()
        ctx.Mode = etime.Train
        ss.Net.InitWeights()
        ss.InitStats()
        ss.StatCounters()
        ss.Logs.ResetLog(etime.Train, etime.Epoch)
        ss.Logs.ResetLog(etime.Test, etime.Epoch)
}

// TestAll runs through the full set of testing items
func (ss *Sim) TestAll() {
        ss.Envs.ByMode(etime.Test).Init(0)
        ss.Loops.ResetAndRun(etime.Test)
        ss.Loops.Mode = etime.Train // Important to reset Mode back to Train because this is called from within the Train Run.
}

////////////////////////////////////////////////////////////////////////
//                 Stats

// InitStats initializes all the statistics.
// called at start of new run
func (ss *Sim) InitStats() {
        ss.Stats.SetFloat("SSE", 0.0)
        ss.Stats.SetFloat("DA", 0.0)
        ss.Stats.SetFloat("AbsDA", 0.0)
        ss.Stats.SetFloat("RewPred", 0.0)
        ss.Stats.SetString("TrialName", "")
        ss.Logs.InitErrStats() // inits TrlErr, FirstZero, LastZero, NZero
}

// StatCounters saves current counters to Stats, so they are available for logging etc
// Also saves a string rep of them for ViewUpdate.Text
func (ss *Sim) StatCounters() {
        ctx := &ss.Context
        mode := ctx.Mode
        ss.Loops.Stacks[mode].CountersToStats(&ss.Stats)
        // always use training epoch..
        trnEpc := ss.Loops.Stacks[etime.Train].Loops[etime.Epoch].Counter.Cur
        ss.Stats.SetInt("Epoch", trnEpc)
        trl := ss.Stats.Int("Trial")
        ss.Stats.SetInt("Trial", trl)
        ss.Stats.SetInt("Cycle", int(ctx.Cycle))
}

func (ss *Sim) NetViewCounters(tm etime.Times) {
        if ss.ViewUpdate.View == nil {
                return
        }
        if tm == etime.Trial {
                ss.TrialStats() // get trial stats for current di
        }
        ss.StatCounters()
        ss.ViewUpdate.Text = ss.Stats.Print([]string{"Run", "Epoch", "Trial", "TrialName", "Cycle", "SSE", "TrlErr"})
}

// TrialStats computes the trial-level statistics.
// Aggregation is done directly from log data.
func (ss *Sim) TrialStats() {
        params := fmt.Sprintf("burst: %g, dip: %g", ss.BurstDaGain, ss.DipDaGain)
        ss.Stats.SetString("RunName", params)

        out := ss.Net.LayerByName("Output")

        sse, avgsse := out.MSE(0.5) // 0.5 = per-unit tolerance -- right side of .5
        ss.Stats.SetFloat("SSE", sse)
        ss.Stats.SetFloat("AvgSSE", avgsse)
        if sse > 0 {
                ss.Stats.SetFloat("TrlErr", 1)
        } else {
                ss.Stats.SetFloat("TrlErr", 0)
        }

        snc := ss.Net.LayerByName("SNc")
        ss.Stats.SetFloat32("DA", snc.Neurons[0].Act)
        ss.Stats.SetFloat32("AbsDA", math32.Abs(snc.Neurons[0].Act))
        rp := ss.Net.LayerByName("RWPred")
        ss.Stats.SetFloat32("RewPred", rp.Neurons[0].Act)
}

//////////////////////////////////////////////////////////////////////
//                 Logging

func (ss *Sim) ConfigLogs() {
        ss.Stats.SetString("RunName", ss.Params.RunName(0)) // used for naming logs, stats, etc

        ss.Logs.AddCounterItems(etime.Run, etime.Epoch, etime.Trial, etime.Cycle)
        ss.Logs.AddStatIntNoAggItem(etime.AllModes, etime.AllTimes, "Expt")
        ss.Logs.AddStatStringItem(etime.AllModes, etime.AllTimes, "RunName")
        ss.Logs.AddStatStringItem(etime.AllModes, etime.Trial, "TrialName")

        ss.Logs.AddPerTrlMSec("PerTrlMSec", etime.Run, etime.Epoch, etime.Trial)

        ss.Logs.AddStatAggItem("SSE", etime.Run, etime.Epoch, etime.Trial)
        ss.Logs.AddStatAggItem("AvgSSE", etime.Run, etime.Epoch, etime.Trial)
        ss.Logs.AddErrStatAggItems("TrlErr", etime.Run, etime.Epoch, etime.Trial)

        ss.Logs.AddStatAggItem("DA", etime.Run, etime.Epoch, etime.Trial)
        ss.Logs.AddStatAggItem("AbsDA", etime.Run, etime.Epoch, etime.Trial)
        ss.Logs.AddStatAggItem("RewPred", etime.Run, etime.Epoch, etime.Trial)

        ss.Logs.PlotItems("PctErr", "AbsDA", "RewPred")

        ss.Logs.CreateTables()
        ss.Logs.SetContext(&ss.Stats, ss.Net)
        // don't plot certain combinations we don't use
        ss.Logs.NoPlot(etime.Train, etime.Cycle)
        ss.Logs.NoPlot(etime.Test, etime.Cycle)
        ss.Logs.NoPlot(etime.Test, etime.Trial)
        ss.Logs.NoPlot(etime.Test, etime.Run)
        ss.Logs.SetMeta(etime.Train, etime.Run, "LegendCol", "RunName")
}

// Log is the main logging function, handles special things for different scopes
func (ss *Sim) Log(mode etime.Modes, time etime.Times) {
        ctx := &ss.Context
        if mode != etime.Analyze {
                ctx.Mode = mode // Also set specifically in a Loop callback.
        }
        dt := ss.Logs.Table(mode, time)
        if dt == nil {
                return
        }
        row := dt.Rows

        switch {
        case time == etime.Cycle:
                return
        case time == etime.Trial:
                ss.TrialStats()
                ss.StatCounters()
        }

        ss.Logs.LogRow(mode, time, row) // also logs to file, etc

        if mode == etime.Test {
                ss.GUI.UpdateTableView(etime.Test, etime.Trial)
        }
}

//////////////////////////////////////////////////////////////////////
//                 GUI

// ConfigGUI configures the Cogent Core GUI interface for this simulation.
func (ss *Sim) ConfigGUI() {
        title := "SIR"
        ss.GUI.MakeBody(ss, "sir", title, `sir illustrates the dynamic gating of information into PFC active maintenance, by the basal ganglia (BG). It uses a simple Store-Ignore-Recall (SIR) task, where the BG system learns via phasic dopamine signals and trial-and-error exploration, discovering what needs to be stored, ignored, and recalled as a function of reinforcement of correct behavior, and learned reinforcement of useful working memory representations. See <a href="https://github.com/CompCogNeuro/sims/blob/master/ch9/sir/README.md">README.md on GitHub</a>.</p>`)
        ss.GUI.CycleUpdateInterval = 10

        nv := ss.GUI.AddNetView("Network")
        nv.Options.MaxRecs = 300
        nv.Options.Raster.Max = 100
        nv.SetNet(ss.Net)
        nv.Options.PathWidth = 0.003
        ss.ViewUpdate.Config(nv, etime.GammaCycle, etime.GammaCycle)
        ss.GUI.ViewUpdate = &ss.ViewUpdate
        nv.Current()

        // nv.SceneXYZ().Camera.Pose.Pos.Set(0, 1.15, 2.25)
        // nv.SceneXYZ().Camera.LookAt(math32.Vector3{0, -0.15, 0}, math32.Vector3{0, 1, 0})

        ss.GUI.AddPlots(title, &ss.Logs)

        ss.GUI.AddTableView(&ss.Logs, etime.Test, etime.Trial)

        ss.GUI.FinalizeGUI(false)
}

func (ss *Sim) MakeToolbar(p *tree.Plan) {
        ss.GUI.AddLooperCtrl(p, ss.Loops)

        tree.Add(p, func(w *core.Separator) {})
        ss.GUI.AddToolbarItem(p, egui.ToolbarItem{Label: "Reset RunLog",
                Icon:    icons.Reset,
                Tooltip: "Reset the accumulated log of all Runs, which are tagged with the ParamSet used",
                Active:  egui.ActiveAlways,
                Func: func() {
                        ss.Logs.ResetLog(etime.Train, etime.Run)
                        ss.GUI.UpdatePlot(etime.Train, etime.Run)
                },
        })
        ////////////////////////////////////////////////
        tree.Add(p, func(w *core.Separator) {})
        ss.GUI.AddToolbarItem(p, egui.ToolbarItem{Label: "New Seed",
                Icon:    icons.Add,
                Tooltip: "Generate a new initial random seed to get different results.  By default, Init re-establishes the same initial seed every time.",
                Active:  egui.ActiveAlways,
                Func: func() {
                        ss.RandSeeds.NewSeeds()
                },
        })
        ss.GUI.AddToolbarItem(p, egui.ToolbarItem{Label: "README",
                Icon:    icons.FileMarkdown,
                Tooltip: "Opens your browser on the README file that contains instructions for how to run this model.",
                Active:  egui.ActiveAlways,
                Func: func() {
                        core.TheApp.OpenURL("https://github.com/CompCogNeuro/sims/blob/main/ch9/sir/README.md")
                },
        })
}

func (ss *Sim) RunGUI() {
        ss.Init()
        ss.ConfigGUI()
        ss.GUI.Body.RunMainWindow()
}

// Copyright (c) 2019, The Emergent Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
        "fmt"
        "math/rand"

        "github.com/emer/emergent/v2/env"
        "github.com/emer/emergent/v2/etime"
        "github.com/emer/etensor/tensor"
)

// Actions are SIR actions
type Actions int32 //enums:enum

const (
        Store1 Actions = iota
        Store2
        Ignore
        Recall1
        Recall2
)

// SIREnv implements the store-ignore-recall task
type SIREnv struct {
        // name of this environment
        Name string

        // number of different stimuli that can be maintained
        NStim int

        // value for reward, based on whether model output = target
        RewVal float32

        // value for non-reward
        NoRewVal float32

        // current action
        Act Actions

        // current stimulus
        Stim int

        // current stimulus being maintained
        Maint1 int

        // current stimulus being maintained
        Maint2 int

        // stimulus input pattern
        Input tensor.Float64

        // input pattern with action
        CtrlInput tensor.Float64

        // output pattern of what to respond
        Output tensor.Float64

        // reward value
        Reward tensor.Float64

        // trial is the step counter within epoch
        Trial env.Counter `display:"inline"`
}

func (ev *SIREnv) Label() string { return ev.Name }

// SetNStim initializes env for given number of stimuli, init states
func (ev *SIREnv) SetNStim(n int) {
        ev.NStim = n
        ev.Input.SetShape([]int{n})
        ev.CtrlInput.SetShape([]int{int(ActionsN)})
        ev.Output.SetShape([]int{n})
        ev.Reward.SetShape([]int{1})
        if ev.RewVal == 0 {
                ev.RewVal = 1
        }
}

func (ev *SIREnv) State(element string) tensor.Tensor {
        switch element {
        case "Input":
                return &ev.Input
        case "CtrlInput":
                return &ev.CtrlInput
        case "Output":
                return &ev.Output
        case "Rew":
                return &ev.Reward
        }
        return nil
}

func (ev *SIREnv) Actions() env.Elements {
        return nil
}

// StimStr returns a letter string rep of stim (A, B...)
func (ev *SIREnv) StimStr(stim int) string {
        return string([]byte{byte('A' + stim)})
}

// String returns the current state as a string
func (ev *SIREnv) String() string {
        return fmt.Sprintf("%s_%s_mnt1_%s_mnt2_%s_rew_%g", ev.Act, ev.StimStr(ev.Stim), ev.StimStr(ev.Maint1), ev.StimStr(ev.Maint2), ev.Reward.Values[0])
}

func (ev *SIREnv) Init(run int) {
        ev.Trial.Scale = etime.Trial
        ev.Trial.Init()
        ev.Trial.Cur = -1 // init state -- key so that first Step() = 0
        ev.Maint1 = -1
        ev.Maint2 = -1
}

// SetState sets the input, output states
func (ev *SIREnv) SetState() {
        ev.CtrlInput.SetZeros()
        ev.CtrlInput.Values[ev.Act] = 1
        ev.Input.SetZeros()
        if ev.Act != Recall1 && ev.Act != Recall2 {
                ev.Input.Values[ev.Stim] = 1
        }
        ev.Output.SetZeros()
        ev.Output.Values[ev.Stim] = 1
}

// SetReward sets reward based on network's output
func (ev *SIREnv) SetReward(netout int) bool {
        cor := ev.Stim // already correct
        rw := netout == cor
        if rw {
                ev.Reward.Values[0] = float64(ev.RewVal)
        } else {
                ev.Reward.Values[0] = float64(ev.NoRewVal)
        }
        return rw
}

// Step the SIR task
func (ev *SIREnv) StepSIR() {
        for {
                ev.Act = Actions(rand.Intn(int(ActionsN)))
                if ev.Act == Store1 && ev.Maint1 >= 0 { // already full
                        continue
                }
                if ev.Act == Recall1 && ev.Maint1 < 0 { // nothing
                        continue
                }
                if ev.Act == Store2 && ev.Maint2 >= 0 { // already full
                        continue
                }
                if ev.Act == Recall2 && ev.Maint2 < 0 { // nothing
                        continue
                }
                break
        }
        ev.Stim = rand.Intn(ev.NStim)
        switch ev.Act {
        case Store1:
                ev.Maint1 = ev.Stim
        case Store2:
                ev.Maint2 = ev.Stim
        case Ignore:
        case Recall1:
                ev.Stim = ev.Maint1
                ev.Maint1 = -1
        case Recall2:
                ev.Stim = ev.Maint2
                ev.Maint2 = -1
        }
        ev.SetState()
}

func (ev *SIREnv) Step() bool {
        ev.StepSIR()
        ev.Trial.Incr()
        return true
}

func (ev *SIREnv) Action(element string, input tensor.Tensor) {
        // nop
}

// Compile-time check that implements Env interface
var _ env.Env = (*SIREnv)(nil)

// Copyright (c) 2019, The Emergent Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

/*
Package fffb provides feedforward (FF) and feedback (FB) inhibition (FFFB)
based on average (or maximum) excitatory netinput (FF) and activation (FB).

This produces a robust, graded k-Winners-Take-All dynamic of sparse
distributed representations having approximately k out of N neurons
active at any time, where k is typically 10-20 percent of N.
*/
package fffb

//go:generate core generate -add-types

// Params parameterizes feedforward (FF) and feedback (FB) inhibition (FFFB)
// based on average (or maximum) netinput (FF) and activation (FB)
type Params struct {

        // enable this level of inhibition
        On bool

        // overall inhibition gain -- this is main parameter to adjust to change overall activation levels -- it scales both the the ff and fb factors uniformly
        Gi float32 `min:"0" default:"1.8"`

        // overall inhibitory contribution from feedforward inhibition -- multiplies average netinput (i.e., synaptic drive into layer) -- this anticipates upcoming changes in excitation, but if set too high, it can make activity slow to emerge -- see also ff0 for a zero-point for this value
        FF float32 `min:"0" default:"1"`

        // overall inhibitory contribution from feedback inhibition -- multiplies average activation -- this reacts to layer activation levels and works more like a thermostat (turning up when the 'heat' in the layer is too high)
        FB float32 `min:"0" default:"1"`

        // time constant in cycles, which should be milliseconds typically (roughly, how long it takes for value to change significantly -- 1.4x the half-life) for integrating feedback inhibitory values -- prevents oscillations that otherwise occur -- the fast default of 1.4 should be used for most cases but sometimes a slower value (3 or higher) can be more robust, especially when inhibition is strong or inputs are more rapidly changing
        FBTau float32 `min:"0" default:"1.4,3,5"`

        // what proportion of the maximum vs. average netinput to use in the feedforward inhibition computation -- 0 = all average, 1 = all max, and values in between = proportional mix between average and max (ff_netin = avg + ff_max_vs_avg * (max - avg)) -- including more max can be beneficial especially in situations where the average can vary significantly but the activity should not -- max is more robust in many situations but less flexible and sensitive to the overall distribution -- max is better for cases more closely approximating single or strictly fixed winner-take-all behavior -- 0.5 is a good compromise in many cases and generally requires a reduction of .1 or slightly more (up to .3-.5) from the gi value for 0
        MaxVsAvg float32 `default:"0,0.5,1"`

        // feedforward zero point for average netinput -- below this level, no FF inhibition is computed based on avg netinput, and this value is subtraced from the ff inhib contribution above this value -- the 0.1 default should be good for most cases (and helps FF_FB produce k-winner-take-all dynamics), but if average netinputs are lower than typical, you may need to lower it
        FF0 float32 `default:"0.1"`

        // rate = 1 / tau
        FBDt float32 `edit:"-" display:"-" json:"-" xml:"-"`
}

func (fb *Params) Update() {
        fb.FBDt = 1 / fb.FBTau
}

func (fb *Params) Defaults() {
        fb.Gi = 1.8
        fb.FF = 1
        fb.FB = 1
        fb.FBTau = 1.4
        fb.MaxVsAvg = 0
        fb.FF0 = 0.1
        fb.Update()
}

func (fb *Params) ShouldDisplay(field string) bool {
        switch field {
        case "Gi", "FF", "FB", "FBTau", "MaxVsAvg", "FF0":
                return fb.On
        default:
                return true
        }
}

// FFInhib returns the feedforward inhibition value based on average and max excitatory conductance within
// relevant scope
func (fb *Params) FFInhib(avgGe, maxGe float32) float32 {
        ffNetin := avgGe + fb.MaxVsAvg*(maxGe-avgGe)
        var ffi float32
        if ffNetin > fb.FF0 {
                ffi = fb.FF * (ffNetin - fb.FF0)
        }
        return ffi
}

// FBInhib computes feedback inhibition value as function of average activation
func (fb *Params) FBInhib(avgAct float32) float32 {
        fbi := fb.FB * avgAct
        return fbi
}

// FBUpdate updates feedback inhibition using time-integration rate constant
func (fb *Params) FBUpdate(fbi *float32, newFbi float32) {
        *fbi += fb.FBDt * (newFbi - *fbi)
}

// Inhib is full inhibition computation for given inhib state, which must have
// the Ge and Act values updated to reflect the current Avg and Max of those
// values in relevant inhibitory pool.
func (fb *Params) Inhib(inh *Inhib) {
        if !fb.On {
                inh.Zero()
                return
        }

        ffi := fb.FFInhib(inh.Ge.Avg, inh.Ge.Max)
        fbi := fb.FBInhib(inh.Act.Avg)

        inh.FFi = ffi
        fb.FBUpdate(&inh.FBi, fbi)

        inh.Gi = fb.Gi * (ffi + inh.FBi)
        inh.GiOrig = inh.Gi
}

// Copyright (c) 2019, The Emergent Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package fffb

import "cogentcore.org/core/math32/minmax"

// Inhib contains state values for computed FFFB inhibition
type Inhib struct {

        // computed feedforward inhibition
        FFi float32

        // computed feedback inhibition (total)
        FBi float32

        // overall value of the inhibition -- this is what is added into the unit Gi inhibition level (along with any synaptic unit-driven inhibition)
        Gi float32

        // original value of the inhibition (before pool or other effects)
        GiOrig float32

        // for pools, this is the layer-level inhibition that is MAX'd with the pool-level inhibition to produce the net inhibition
        LayGi float32

        // average and max Ge excitatory conductance values, which drive FF inhibition
        Ge minmax.AvgMax32

        // average and max Act activation values, which drive FB inhibition
        Act minmax.AvgMax32
}

func (fi *Inhib) Init() {
        fi.Zero()
        fi.Ge.Init()
        fi.Act.Init()
}

// Zero clears inhibition but does not affect Ge, Act averages
func (fi *Inhib) Zero() {
        fi.FFi = 0
        fi.FBi = 0
        fi.Gi = 0
        fi.GiOrig = 0
        fi.LayGi = 0
}

// Decay reduces inhibition values by given decay proportion
func (fi *Inhib) Decay(decay float32) {
        fi.Ge.Max -= decay * fi.Ge.Max
        fi.Ge.Avg -= decay * fi.Ge.Avg
        fi.Act.Max -= decay * fi.Act.Max
        fi.Act.Avg -= decay * fi.Act.Avg
        fi.FFi -= decay * fi.FFi
        fi.FBi -= decay * fi.FBi
        fi.Gi -= decay * fi.Gi
}

// Inhibs is a slice of Inhib records
type Inhibs []Inhib

// Copyright (c) 2019, The Emergent Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

/*
Package knadapt provides code for sodium (Na) gated potassium (K) currents that drive
adaptation (accommodation) in neural firing.  As neurons spike, driving an influx of Na,
this activates the K channels, which, like leak channels, pull the membrane potential
back down toward rest (or even below).  Multiple different time constants have been
identified and this implementation supports 3:
M-type (fast), Slick (medium), and Slack (slow)

Here's a good reference:

Kaczmarek, L. K. (2013). Slack, Slick, and Sodium-Activated Potassium Channels.
ISRN Neuroscience, 2013. https://doi.org/10.1155/2013/354262

This package supports both spiking and rate-coded activations.
*/
package knadapt

//go:generate core generate -add-types

// Chan describes one channel type of sodium-gated adaptation, with a specific
// set of rate constants.
type Chan struct {

        // if On, use this component of K-Na adaptation
        On bool

        // Rise rate of fast time-scale adaptation as function of Na concentration -- directly multiplies -- 1/rise = tau for rise rate
        Rise float32

        // Maximum potential conductance of fast K channels -- divide nA biological value by 10 for the normalized units here
        Max float32

        // time constant in cycles for decay of adaptation, which should be milliseconds typically (roughly, how long it takes for value to change significantly -- 1.4x the half-life)
        Tau float32

        // 1/Tau rate constant
        Dt float32 `display:"-"`
}

func (ka *Chan) Defaults() {
        ka.On = true
        ka.Rise = 0.01
        ka.Max = 0.1
        ka.Tau = 100
        ka.Update()
}

func (ka *Chan) Update() {
        ka.Dt = 1 / ka.Tau
}

func (ka *Chan) ShouldDisplay(field string) bool {
        switch field {
        case "Rise", "Max", "Tau":
                return ka.On
        default:
                return true
        }
}

// GcFromSpike updates the KNa conductance based on spike or not
func (ka *Chan) GcFromSpike(gKNa *float32, spike bool) {
        if ka.On {
                if spike {
                        *gKNa += ka.Rise * (ka.Max - *gKNa)
                } else {
                        *gKNa -= ka.Dt * *gKNa
                }
        } else {
                *gKNa = 0
        }
}

// GcFromRate updates the KNa conductance based on rate-coded activation.
// act should already have the compensatory rate multiplier prior to calling.
func (ka *Chan) GcFromRate(gKNa *float32, act float32) {
        if ka.On {
                *gKNa += act*ka.Rise*(ka.Max-*gKNa) - (ka.Dt * *gKNa)
        } else {
                *gKNa = 0
        }
}

// Params describes sodium-gated potassium channel adaptation mechanism.
// Evidence supports at least 3 different time constants:
// M-type (fast), Slick (medium), and Slack (slow)
type Params struct {

        // if On, apply K-Na adaptation
        On bool

        // extra multiplier for rate-coded activations on rise factors -- adjust to match discrete spiking
        Rate float32 `default:"0.8"`

        // fast time-scale adaptation
        Fast Chan `display:"inline"`

        // medium time-scale adaptation
        Med Chan `display:"inline"`

        // slow time-scale adaptation
        Slow Chan `display:"inline"`
}

func (ka *Params) Defaults() {
        ka.Rate = 0.8
        ka.Fast.Defaults()
        ka.Med.Defaults()
        ka.Slow.Defaults()
        ka.Fast.Tau = 50
        ka.Fast.Rise = 0.05
        ka.Fast.Max = 0.1
        ka.Med.Tau = 200
        ka.Med.Rise = 0.02
        ka.Med.Max = 0.1
        ka.Slow.Tau = 1000
        ka.Slow.Rise = 0.001
        ka.Slow.Max = 1
        ka.Update()
}

func (ka *Params) Update() {
        ka.Fast.Update()
        ka.Med.Update()
        ka.Slow.Update()
}

func (ka *Params) ShouldDisplay(field string) bool {
        switch field {
        case "Rate":
                return ka.On
        default:
                return true
        }
}

// GcFromSpike updates all time scales of KNa adaptation from spiking
func (ka *Params) GcFromSpike(gKNaF, gKNaM, gKNaS *float32, spike bool) {
        ka.Fast.GcFromSpike(gKNaF, spike)
        ka.Med.GcFromSpike(gKNaM, spike)
        ka.Slow.GcFromSpike(gKNaS, spike)
}

// GcFromRate updates all time scales of KNa adaptation from rate code activation
func (ka *Params) GcFromRate(gKNaF, gKNaM, gKNaS *float32, act float32) {
        act *= ka.Rate
        ka.Fast.GcFromRate(gKNaF, act)
        ka.Med.GcFromRate(gKNaM, act)
        ka.Slow.GcFromRate(gKNaS, act)
}

/*
class STATE_CLASS(KNaAdaptMiscSpec) : public STATE_CLASS(SpecMemberBase) {
  // ##INLINE ##NO_TOKENS ##CAT_Leabra extra params associated with sodium-gated potassium channel adaptation mechanism
INHERITED(SpecMemberBase)
public:
  bool          clamp;          // #DEF_true apply adaptation even to clamped layers -- only happens if kna_adapt.on is true
  bool          invert_nd;      // #DEF_true invert the adaptation effect for the act_nd (non-depressed) value that is typically used for learning-drivng averages (avg_ss, _s, _m) -- only happens if kna_adapt.on is true
  float         max_gc;         // #CONDSHOW_ON_clamp||invert_nd #DEF_0.2 for clamp or invert_nd, maximum k_na conductance that we expect to get (prior to multiplying by g_bar.k) -- apply a proportional reduction in clamped activation and/or enhancement of act_nd based on current k_na conductance -- default is appropriate for default kna_adapt params
  float         max_adapt;      // #CONDSHOW_ON_clamp||invert_nd has opposite effects for clamp and invert_nd (and only operative when kna_adapt.on in addition): for clamp on clamped layers, this is the maximum amount of adaptation to apply to clamped activations when conductance is at max_gc -- biologically, values around .5 correspond generally to strong adaptation in primary visual cortex (V1) -- for invert_nd, this is the maximum amount of adaptation to invert, which is key for allowing learning to operate successfully despite the depression of activations due to adaptation -- values around .2 to .4 are good for g_bar.k = .2, depending on how strongly inputs are depressed -- need to experiment to find the best value for a given config
  bool          no_targ;        // #DEF_true automatically exclude units in TARGET layers and also TRC (Pulvinar) thalamic neurons from adaptation effects -- typically such layers should not be subject to these effects, so this makes it easier to not have to manually set those override params

  INLINE float Compute_Clamped(float clamp_act, float gc_kna_f, float gc_kna_m, float gc_kna_s) {
    float gc_kna = gc_kna_f + gc_kna_m + gc_kna_s;
    float pct_gc = fminf(gc_kna / max_gc, 1.0f);
    return clamp_act * (1.0f - pct_gc * max_adapt);
  }
  // apply adaptation directly to a clamped activation value, reducing in proportion to amount of k_na current

  INLINE float Compute_ActNd(float act, float gc_kna_f, float gc_kna_m, float gc_kna_s) {
    float gc_kna = gc_kna_f + gc_kna_m + gc_kna_s;
    float pct_gc = fminf(gc_kna / max_gc, 1.0f);
    return act * (1.0f + pct_gc * max_adapt);
  }
  // apply inverse of adaptation to activation value, increasing in proportion to amount of k_na current

  STATE_DECO_KEY("UnitSpec");
  STATE_TA_STD_CODE_SPEC(KNaAdaptMiscSpec);

  // STATE_UAE( UpdateDts(); );

private:
  void        Initialize()      { Defaults_init(); }
  void        Defaults_init() {
    clamp = true;  invert_nd = true;  max_gc = .2f;  max_adapt = 0.3f;  no_targ = true;
  }
};

*/

// Copyright (c) 2019, The Emergent Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package leabra

import (
        "cogentcore.org/core/math32"
        "cogentcore.org/core/math32/minmax"
        "cogentcore.org/lab/base/randx"
        "github.com/emer/leabra/v2/chans"
        "github.com/emer/leabra/v2/knadapt"
        "github.com/emer/leabra/v2/nxx1"
)

///////////////////////////////////////////////////////////////////////
//  act.go contains the activation params and functions for leabra

// leabra.ActParams contains all the activation computation params and functions
// for basic Leabra, at the neuron level .
// This is included in leabra.Layer to drive the computation.
type ActParams struct {

        // Noisy X/X+1 rate code activation function parameters
        XX1 nxx1.Params `display:"inline"`

        // optimization thresholds for faster processing
        OptThresh OptThreshParams `display:"inline"`

        // initial values for key network state variables -- initialized at start of trial with InitActs or DecayActs
        Init ActInitParams `display:"inline"`

        // time and rate constants for temporal derivatives / updating of activation state
        Dt DtParams `display:"inline"`

        // maximal conductances levels for channels
        Gbar chans.Chans `display:"inline"`

        // reversal potentials for each channel
        Erev chans.Chans `display:"inline"`

        // how external inputs drive neural activations
        Clamp ClampParams `display:"inline"`

        // how, where, when, and how much noise to add to activations
        Noise ActNoiseParams `display:"inline"`

        // range for Vm membrane potential -- by default
        VmRange minmax.F32 `display:"inline"`

        // sodium-gated potassium channel adaptation parameters -- activates an inhibitory leak-like current as a function of neural activity (firing = Na influx) at three different time-scales (M-type = fast, Slick = medium, Slack = slow)
        KNa knadapt.Params `display:"no-inline"`

        // Erev - Act.Thr for each channel -- used in computing GeThrFromG among others
        ErevSubThr chans.Chans `edit:"-" display:"-" json:"-" xml:"-"`

        // Act.Thr - Erev for each channel -- used in computing GeThrFromG among others
        ThrSubErev chans.Chans `edit:"-" display:"-" json:"-" xml:"-"`
}

func (ac *ActParams) Defaults() {
        ac.XX1.Defaults()
        ac.OptThresh.Defaults()
        ac.Init.Defaults()
        ac.Dt.Defaults()
        ac.Gbar.SetAll(1.0, 0.1, 1.0, 1.0)
        ac.Erev.SetAll(1.0, 0.3, 0.25, 0.25)
        ac.Clamp.Defaults()
        ac.VmRange.Max = 2.0
        ac.KNa.Defaults()
        ac.KNa.On = false
        ac.Noise.Defaults()
        ac.Update()
}

// Update must be called after any changes to parameters
func (ac *ActParams) Update() {
        ac.ErevSubThr.SetFromOtherMinus(ac.Erev, ac.XX1.Thr)
        ac.ThrSubErev.SetFromMinusOther(ac.XX1.Thr, ac.Erev)

        ac.XX1.Update()
        ac.OptThresh.Update()
        ac.Init.Update()
        ac.Dt.Update()
        ac.Clamp.Update()
        ac.Noise.Update()
        ac.KNa.Update()
}

///////////////////////////////////////////////////////////////////////
//  Init

// InitGinc initializes the Ge excitatory and Gi inhibitory conductance accumulation states
// including ActSent and G*Raw values.
// called at start of trial always, and can be called optionally
// when delta-based Ge computation needs to be updated (e.g., weights
// might have changed strength)
func (ac *ActParams) InitGInc(nrn *Neuron) {
        nrn.ActSent = 0
        nrn.GeRaw = 0
        nrn.GiRaw = 0
}

// DecayState decays the activation state toward initial values in proportion to given decay parameter
// Called with ac.Init.Decay by Layer during AlphaCycInit
func (ac *ActParams) DecayState(nrn *Neuron, decay float32) {
        if decay > 0 { // no-op for most, but not all..
                nrn.Act -= decay * (nrn.Act - ac.Init.Act)
                nrn.Ge -= decay * (nrn.Ge - ac.Init.Ge)
                nrn.Gi -= decay * nrn.Gi
                nrn.GiSelf -= decay * nrn.GiSelf
                nrn.Gk -= decay * nrn.Gk
                nrn.Vm -= decay * (nrn.Vm - ac.Init.Vm)
                nrn.GiSyn -= decay * nrn.GiSyn
                nrn.Burst -= decay * (nrn.Burst - ac.Init.Act)
        }
        nrn.ActDel = 0
        nrn.Inet = 0
}

// InitActs initializes activation state in neuron -- called during InitWeights but otherwise not
// automatically called (DecayState is used instead)
func (ac *ActParams) InitActs(nrn *Neuron) {
        nrn.Act = ac.Init.Act
        nrn.ActLrn = ac.Init.Act
        nrn.Ge = ac.Init.Ge
        nrn.Gi = 0
        nrn.Gk = 0
        nrn.GknaFast = 0
        nrn.GknaMed = 0
        nrn.GknaSlow = 0
        nrn.GiSelf = 0
        nrn.GiSyn = 0
        nrn.Inet = 0
        nrn.Vm = ac.Init.Vm
        nrn.Targ = 0
        nrn.Ext = 0
        nrn.ActDel = 0
        nrn.Spike = 0
        nrn.ISI = -1
        nrn.ISIAvg = -1
        nrn.CtxtGe = 0
        nrn.ActG = 0
        nrn.DALrn = 0
        nrn.Shunt = 0
        nrn.Maint = 0
        nrn.MaintGe = 0

        ac.InitActQs(nrn)
        ac.InitGInc(nrn)
}

// InitActQs initializes quarter-based activation states in neuron (ActQ0-2, ActM, ActP, ActDif)
// Called from InitActs, which is called from InitWeights, but otherwise not automatically called
// (DecayState is used instead)
func (ac *ActParams) InitActQs(nrn *Neuron) {
        nrn.ActQ0 = 0
        nrn.ActQ1 = 0
        nrn.ActQ2 = 0
        nrn.ActM = 0
        nrn.ActP = 0
        nrn.ActDif = 0
        nrn.Burst = 0
        nrn.BurstPrv = 0
}

///////////////////////////////////////////////////////////////////////
//  Cycle

// GeFromRaw integrates Ge excitatory conductance from GeRaw value
// (can add other terms to geRaw prior to calling this)
func (ac *ActParams) GeFromRaw(nrn *Neuron, geRaw float32) {
        if !ac.Clamp.Hard && nrn.HasFlag(NeurHasExt) {
                if ac.Clamp.Avg {
                        geRaw = ac.Clamp.AvgGe(nrn.Ext, geRaw)
                } else {
                        geRaw += nrn.Ext * ac.Clamp.Gain
                }
        }

        ac.Dt.GFromRaw(geRaw, &nrn.Ge)
        // first place noise is required -- generate here!
        if ac.Noise.Type != NoNoise && !ac.Noise.Fixed && ac.Noise.Dist != randx.Mean {
                nrn.Noise = float32(ac.Noise.Gen())
        }
        if ac.Noise.Type == GeNoise {
                nrn.Ge += nrn.Noise
        }
}

// GiFromRaw integrates GiSyn inhibitory synaptic conductance from GiRaw value
// (can add other terms to geRaw prior to calling this)
func (ac *ActParams) GiFromRaw(nrn *Neuron, giRaw float32) {
        ac.Dt.GFromRaw(giRaw, &nrn.GiSyn)
        nrn.GiSyn = math32.Max(nrn.GiSyn, 0) // negative inhib G doesn't make any sense
}

// InetFromG computes net current from conductances and Vm
func (ac *ActParams) InetFromG(vm, ge, gi, gk float32) float32 {
        return ge*(ac.Erev.E-vm) + ac.Gbar.L*(ac.Erev.L-vm) + gi*(ac.Erev.I-vm) + gk*(ac.Erev.K-vm)
}

// VmFromG computes membrane potential Vm from conductances Ge, Gi, and Gk.
// The Vm value is only used in pure rate-code computation within the sub-threshold regime
// because firing rate is a direct function of excitatory conductance Ge.
func (ac *ActParams) VmFromG(nrn *Neuron) {
        ge := nrn.Ge * ac.Gbar.E
        gi := nrn.Gi * ac.Gbar.I
        gk := nrn.Gk * ac.Gbar.K
        nrn.Inet = ac.InetFromG(nrn.Vm, ge, gi, gk)
        nwVm := nrn.Vm + ac.Dt.VmDt*nrn.Inet

        if ac.Noise.Type == VmNoise {
                nwVm += nrn.Noise
        }
        nrn.Vm = ac.VmRange.ClampValue(nwVm)
}

// GeThrFromG computes the threshold for Ge based on all other conductances,
// including Gk.  This is used for computing the adapted Act value.
func (ac *ActParams) GeThrFromG(nrn *Neuron) float32 {
        return ((ac.Gbar.I*nrn.Gi*ac.ErevSubThr.I + ac.Gbar.L*ac.ErevSubThr.L + ac.Gbar.K*nrn.Gk*ac.ErevSubThr.K) / ac.ThrSubErev.E)
}

// GeThrFromGnoK computes the threshold for Ge based on other conductances,
// excluding Gk.  This is used for computing the non-adapted ActLrn value.
func (ac *ActParams) GeThrFromGnoK(nrn *Neuron) float32 {
        return ((ac.Gbar.I*nrn.Gi*ac.ErevSubThr.I + ac.Gbar.L*ac.ErevSubThr.L) / ac.ThrSubErev.E)
}

// ActFromG computes rate-coded activation Act from conductances Ge, Gi, Gk
func (ac *ActParams) ActFromG(nrn *Neuron) {
        if ac.HasHardClamp(nrn) {
                ac.HardClamp(nrn)
                return
        }
        var nwAct, nwActLrn float32
        if nrn.Act < ac.XX1.VmActThr && nrn.Vm <= ac.XX1.Thr {
                // note: this is quite important -- if you directly use the gelin
                // the whole time, then units are active right away -- need Vm dynamics to
                // drive subthreshold activation behavior
                nwAct = ac.XX1.NoisyXX1(nrn.Vm - ac.XX1.Thr)
                nwActLrn = nwAct
        } else {
                ge := nrn.Ge * ac.Gbar.E
                geThr := ac.GeThrFromG(nrn)
                nwAct = ac.XX1.NoisyXX1(ge - geThr)
                geThr = ac.GeThrFromGnoK(nrn)          // excludes K adaptation effect
                nwActLrn = ac.XX1.NoisyXX1(ge - geThr) // learning is non-adapted
        }
        curAct := nrn.Act
        nwAct = curAct + ac.Dt.VmDt*(nwAct-curAct)
        nrn.ActDel = nwAct - curAct

        if ac.Noise.Type == ActNoise {
                nwAct += nrn.Noise
        }
        nrn.Act = nwAct

        nwActLrn = nrn.ActLrn + ac.Dt.VmDt*(nwActLrn-nrn.ActLrn)
        nrn.ActLrn = nwActLrn

        if ac.KNa.On {
                ac.KNa.GcFromRate(&nrn.GknaFast, &nrn.GknaMed, &nrn.GknaSlow, nrn.Act)
                nrn.Gk = nrn.GknaFast + nrn.GknaMed + nrn.GknaSlow
        }
}

// HasHardClamp returns true if this neuron has external input that should be hard clamped
func (ac *ActParams) HasHardClamp(nrn *Neuron) bool {
        return ac.Clamp.Hard && nrn.HasFlag(NeurHasExt)
}

// HardClamp clamps activation from external input -- just does it -- use HasHardClamp to check
// if it should do it.  Also adds any Noise *if* noise is set to ActNoise.
func (ac *ActParams) HardClamp(nrn *Neuron) {
        ext := nrn.Ext
        if ac.Noise.Type == ActNoise {
                ext += nrn.Noise
        }
        clmp := ac.Clamp.Range.ClampValue(ext)
        nrn.Act = clmp + nrn.Noise
        nrn.ActLrn = clmp
        nrn.Vm = ac.XX1.Thr + nrn.Act/ac.XX1.Gain
        nrn.ActDel = 0
        nrn.Inet = 0
}

//////////////////////////////////////////////////////////////////////////////////////
//  OptThreshParams

// OptThreshParams provides optimization thresholds for faster processing
type OptThreshParams struct {

        // don't send activation when act <= send -- greatly speeds processing
        Send float32 `default:"0.1"`

        // don't send activation changes until they exceed this threshold: only for when LeabraNetwork::send_delta is on!
        Delta float32 `default:"0.005"`
}

func (ot *OptThreshParams) Update() {
}

func (ot *OptThreshParams) Defaults() {
        ot.Send = .1
        ot.Delta = 0.005
}

//////////////////////////////////////////////////////////////////////////////////////
//  ActInitParams

// ActInitParams are initial values for key network state variables.
// Initialized at start of trial with Init_Acts or DecayState.
type ActInitParams struct {

        // proportion to decay activation state toward initial values at start of every trial
        Decay float32 `default:"0,1" max:"1" min:"0"`

        // initial membrane potential -- see e_rev.l for the resting potential (typically .3) -- often works better to have a somewhat elevated initial membrane potential relative to that
        Vm float32 `default:"0.4"`

        // initial activation value -- typically 0
        Act float32 `default:"0"`

        // baseline level of excitatory conductance (net input) -- Ge is initialized to this value, and it is added in as a constant background level of excitatory input -- captures all the other inputs not represented in the model, and intrinsic excitability, etc
        Ge float32 `default:"0"`
}

func (ai *ActInitParams) Update() {
}

func (ai *ActInitParams) Defaults() {
        ai.Decay = 1
        ai.Vm = 0.4
        ai.Act = 0
        ai.Ge = 0
}

//////////////////////////////////////////////////////////////////////////////////////
//  DtParams

// DtParams are time and rate constants for temporal derivatives in Leabra (Vm, net input)
type DtParams struct {

        // overall rate constant for numerical integration, for all equations at the unit level -- all time constants are specified in millisecond units, with one cycle = 1 msec -- if you instead want to make one cycle = 2 msec, you can do this globally by setting this integ value to 2 (etc).  However, stability issues will likely arise if you go too high.  For improved numerical stability, you may even need to reduce this value to 0.5 or possibly even lower (typically however this is not necessary).  MUST also coordinate this with network.time_inc variable to ensure that global network.time reflects simulated time accurately
        Integ float32 `default:"1,0.5" min:"0"`

        // membrane potential and rate-code activation time constant in cycles, which should be milliseconds typically (roughly, how long it takes for value to change significantly -- 1.4x the half-life) -- reflects the capacitance of the neuron in principle -- biological default for AdEx spiking model C = 281 pF = 2.81 normalized -- for rate-code activation, this also determines how fast to integrate computed activation values over time
        VmTau float32 `default:"3.3" min:"1"`

        // time constant for integrating synaptic conductances, in cycles, which should be milliseconds typically (roughly, how long it takes for value to change significantly -- 1.4x the half-life) -- this is important for damping oscillations -- generally reflects time constants associated with synaptic channels which are not modeled in the most abstract rate code models (set to 1 for detailed spiking models with more realistic synaptic currents) -- larger values (e.g., 3) can be important for models with higher conductances that otherwise might be more prone to oscillation.
        GTau float32 `default:"1.4,3,5" min:"1"`

        // for integrating activation average (ActAvg), time constant in trials (roughly, how long it takes for value to change significantly) -- used mostly for visualization and tracking *hog* units
        AvgTau float32 `default:"200"`

        // nominal rate = Integ / tau
        VmDt float32 `display:"-" json:"-" xml:"-"`

        // rate = Integ / tau
        GDt float32 `display:"-" json:"-" xml:"-"`

        // rate = 1 / tau
        AvgDt float32 `display:"-" json:"-" xml:"-"`
}

func (dp *DtParams) Update() {
        dp.VmDt = dp.Integ / dp.VmTau
        dp.GDt = dp.Integ / dp.GTau
        dp.AvgDt = 1 / dp.AvgTau
}

func (dp *DtParams) Defaults() {
        dp.Integ = 1
        dp.VmTau = 3.3
        dp.GTau = 1.4
        dp.AvgTau = 200
        dp.Update()
}

func (dp *DtParams) GFromRaw(geRaw float32, ge *float32) {
        *ge += dp.GDt * (geRaw - *ge)
}

//////////////////////////////////////////////////////////////////////////////////////
//  Noise

// ActNoiseType are different types / locations of random noise for activations
type ActNoiseType int //enums:enum

// The activation noise types
const (
        // NoNoise means no noise added
        NoNoise ActNoiseType = iota

        // VmNoise means noise is added to the membrane potential.
        // IMPORTANT: this should NOT be used for rate-code (NXX1) activations,
        // because they do not depend directly on the vm -- this then has no effect
        VmNoise

        // GeNoise means noise is added to the excitatory conductance (Ge).
        // This should be used for rate coded activations (NXX1)
        GeNoise

        // ActNoise means noise is added to the final rate code activation
        ActNoise

        // GeMultNoise means that noise is multiplicative on the Ge excitatory conductance values
        GeMultNoise
)

// ActNoiseParams contains parameters for activation-level noise
type ActNoiseParams struct {
        randx.RandParams

        // where and how to add processing noise
        Type ActNoiseType

        // keep the same noise value over the entire alpha cycle -- prevents noise from being washed out and produces a stable effect that can be better used for learning -- this is strongly recommended for most learning situations
        Fixed bool
}

func (an *ActNoiseParams) Update() {
}

func (an *ActNoiseParams) Defaults() {
        an.Fixed = true
}

//////////////////////////////////////////////////////////////////////////////////////
//  ClampParams

// ClampParams are for specifying how external inputs are clamped onto network activation values
type ClampParams struct {

        // whether to hard clamp inputs where activation is directly set to external input value (Act = Ext) or do soft clamping where Ext is added into Ge excitatory current (Ge += Gain * Ext)
        Hard bool `default:"true"`

        // range of external input activation values allowed -- Max is .95 by default due to saturating nature of rate code activation function
        Range minmax.F32

        // soft clamp gain factor (Ge += Gain * Ext)
        Gain float32 `default:"0.02:0.5"`

        // compute soft clamp as the average of current and target netins, not the sum -- prevents some of the main effect problems associated with adding external inputs
        Avg bool

        // gain factor for averaging the Ge -- clamp value Ext contributes with AvgGain and current Ge as (1-AvgGain)
        AvgGain float32 `default:"0.2"`
}

func (cp *ClampParams) Update() {
}

func (cp *ClampParams) Defaults() {
        cp.Hard = true
        cp.Range.Max = 0.95
        cp.Gain = 0.2
        cp.Avg = false
        cp.AvgGain = 0.2
}

func (cp *ClampParams) ShouldDisplay(field string) bool {
        switch field {
        case "Range":
                return cp.Hard
        case "Gain", "Avg":
                return !cp.Hard
        case "AvgGain":
                return !cp.Hard && cp.Avg
        default:
                return true
        }
}

// AvgGe computes Avg-based Ge clamping value if using that option.
func (cp *ClampParams) AvgGe(ext, ge float32) float32 {
        return cp.AvgGain*cp.Gain*ext + (1-cp.AvgGain)*ge
}

//////////////////////////////////////////////////////////////////////////////////////
//  WtInitParams

// WtInitParams are weight initialization parameters -- basically the
// random distribution parameters but also Symmetry flag
type WtInitParams struct {
        randx.RandParams

        // symmetrize the weight values with those in reciprocal pathway -- typically true for bidirectional excitatory connections
        Sym bool
}

func (wp *WtInitParams) Defaults() {
        wp.Mean = 0.5
        wp.Var = 0.25
        wp.Dist = randx.Uniform
        wp.Sym = true
}

//////////////////////////////////////////////////////////////////////////////////////
//  WtScaleParams

// / WtScaleParams are weight scaling parameters: modulates overall strength of pathway,
// using both absolute and relative factors
type WtScaleParams struct {

        // absolute scaling, which is not subject to normalization: directly multiplies weight values
        Abs float32 `default:"1" min:"0"`

        // relative scaling that shifts balance between different pathways -- this is subject to normalization across all other pathways into unit
        Rel float32 `min:"0"`
}

func (ws *WtScaleParams) Defaults() {
        ws.Abs = 1
        ws.Rel = 1
}

func (ws *WtScaleParams) Update() {
}

// SLayActScale computes scaling factor based on sending layer activity level (savg), number of units
// in sending layer (snu), and number of recv connections (ncon).
// Uses a fixed sem_extra standard-error-of-the-mean (SEM) extra value of 2
// to add to the average expected number of active connections to receive,
// for purposes of computing scaling factors with partial connectivity
// For 25% layer activity, binomial SEM = sqrt(p(1-p)) = .43, so 3x = 1.3 so 2 is a reasonable default.
func (ws *WtScaleParams) SLayActScale(savg, snu, ncon float32) float32 {
        ncon = math32.Max(ncon, 1) // path Avg can be < 1 in some cases
        semExtra := 2
        slayActN := int(math32.Round(savg * snu)) // sending layer actual # active
        slayActN = max(slayActN, 1)
        var sc float32
        if ncon == snu {
                sc = 1 / float32(slayActN)
        } else {
                maxActN := int(math32.Min(ncon, float32(slayActN))) // max number we could get
                avgActN := int(math32.Round(savg * ncon))           // recv average actual # active if uniform
                avgActN = max(avgActN, 1)
                expActN := avgActN + semExtra // expected
                expActN = min(expActN, maxActN)
                sc = 1 / float32(expActN)
        }
        return sc
}

// FullScale returns full scaling factor, which is product of Abs * Rel * SLayActScale
func (ws *WtScaleParams) FullScale(savg, snu, ncon float32) float32 {
        return ws.Abs * ws.Rel * ws.SLayActScale(savg, snu, ncon)
}

// Copyright (c) 2019, The Emergent Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package leabra

import "github.com/emer/emergent/v2/etime"

// leabra.Context contains all the timing state and parameter information for running a model
type Context struct {

        // accumulated amount of time the network has been running,
        // in simulation-time (not real world time), in seconds.
        Time float32

        // cycle counter: number of iterations of activation updating
        // (settling) on the current alpha-cycle (100 msec / 10 Hz) trial.
        // This counts time sequentially through the entire trial,
        // typically from 0 to 99 cycles.
        Cycle int

        // total cycle count. this increments continuously from whenever
        // it was last reset, typically this is number of milliseconds
        // in simulation time.
        CycleTot int

        // current gamma-frequency (25 msec / 40 Hz) quarter of alpha-cycle
        // (100 msec / 10 Hz) trial being processed.
        // Due to 0-based indexing, the first quarter is 0, second is 1, etc.
        // The plus phase final quarter is 3.
        Quarter Quarters

        // true if this is the plus phase (final quarter = 3), else minus phase.
        PlusPhase bool

        // amount of time to increment per cycle.
        TimePerCyc float32 `default:"0.001"`

        // number of cycles per quarter to run: 25 = standard 100 msec alpha-cycle.
        CycPerQtr int `default:"25"`

        // current evaluation mode, e.g., Train, Test, etc
        Mode etime.Modes
}

// NewContext returns a new Context struct with default parameters
func NewContext() *Context {
        tm := &Context{}
        tm.Defaults()
        return tm
}

// Defaults sets default values
func (tm *Context) Defaults() {
        tm.TimePerCyc = 0.001
        tm.CycPerQtr = 25
}

// Reset resets the counters all back to zero
func (tm *Context) Reset() {
        tm.Time = 0
        tm.Cycle = 0
        tm.CycleTot = 0
        tm.Quarter = 0
        tm.PlusPhase = false
        if tm.CycPerQtr == 0 {
                tm.Defaults()
        }
}

// AlphaCycStart starts a new alpha-cycle (set of 4 quarters)
func (tm *Context) AlphaCycStart() {
        tm.Cycle = 0
        tm.Quarter = 0
        tm.PlusPhase = false
}

// CycleInc increments at the cycle level
func (tm *Context) CycleInc() {
        tm.Cycle++
        tm.CycleTot++
        tm.Time += tm.TimePerCyc
}

// QuarterInc increments at the quarter level, updating Quarter and PlusPhase
func (tm *Context) QuarterInc() {
        tm.Quarter++
}

// QuarterCycle returns the number of cycles into current quarter
func (tm *Context) QuarterCycle() int {
        qmin := int(tm.Quarter) * tm.CycPerQtr
        return tm.Cycle - qmin
}

//////////////////////////////////////////////////////////////////////////////////////
//  Quarters

// Quarters are the different alpha trial quarters, as a bitflag,
// for use in relevant timing parameters where quarters need to be specified.
// The Q1..4 defined values are integer *bit positions* -- use Set, Has etc methods
// to set bits from these bit positions.
type Quarters int64 //enums:bitflag

// The quarters
const (
        // Q1 is the first quarter, which, due to 0-based indexing, shows up as Quarter = 0 in timer
        Q1 Quarters = iota
        Q2
        Q3
        Q4
)

// HasNext returns true if the quarter after given quarter is set.
// This wraps around from Q4 to Q1.  (qtr = 0..3 = same as Quarters)
func (qt Quarters) HasNext(qtr Quarters) bool {
        nqt := (qtr + 1) % 4
        return qt.HasFlag(nqt)
}

// HasPrev returns true if the quarter before given quarter is set.
// This wraps around from Q1 to Q4.  (qtr = 0..3 = same as Quarters)
func (qt Quarters) HasPrev(qtr Quarters) bool {
        pqt := (qtr - 1)
        if pqt < 0 {
                pqt += 4
        }
        return qt.HasFlag(pqt)
}

// Copyright (c) 2024, The Emergent Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package leabra

import (
        "fmt"
        "log"
        "math"

        "cogentcore.org/core/math32"
)

// BurstParams determine how the 5IB Burst activation is computed from
// standard Act activation values in SuperLayer. It is thresholded.
type BurstParams struct {

        // Quarter(s) when bursting occurs -- typically Q4 but can also be Q2 and Q4 for beta-frequency updating.  Note: this is a bitflag and must be accessed using its Set / Has etc routines, 32 bit versions.
        BurstQtr Quarters

        // Relative component of threshold on superficial activation value, below which it does not drive Burst (and above which, Burst = Act).  This is the distance between the average and maximum activation values within layer (e.g., 0 = average, 1 = max).  Overall effective threshold is MAX of relative and absolute thresholds.
        ThrRel float32 `max:"1" default:"0.1,0.2,0.5"`

        // Absolute component of threshold on superficial activation value, below which it does not drive Burst (and above which, Burst = Act).  Overall effective threshold is MAX of relative and absolute thresholds.
        ThrAbs float32 `min:"0" max:"1" default:"0.1,0.2,0.5"`
}

func (db *BurstParams) Defaults() {
        db.BurstQtr.SetFlag(true, Q4)
        db.ThrRel = 0.1
        db.ThrAbs = 0.1
}

func (db *BurstParams) Update() {
}

////////  Burst -- computed in CyclePost

// BurstPrv records Burst activity just prior to burst
func (ly *Layer) BurstPrv(ctx *Context) {
        if !ly.Burst.BurstQtr.HasNext(ctx.Quarter) {
                return
        }
        // if will be updating next quarter, save just prior
        // this logic works for all cases, but e.g., BurstPrv doesn't update
        // until end of minus phase for Q4 BurstQtr
        for ni := range ly.Neurons {
                nrn := &ly.Neurons[ni]
                nrn.BurstPrv = nrn.Burst
        }
}

// BurstFromAct updates Burst layer 5IB bursting value from current Act
// (superficial activation), subject to thresholding.
func (ly *Layer) BurstFromAct(ctx *Context) {
        if !ly.Burst.BurstQtr.HasFlag(ctx.Quarter) {
                return
        }
        lpl := &ly.Pools[0]
        actMax := lpl.Inhib.Act.Max
        actAvg := lpl.Inhib.Act.Avg
        thr := actAvg + ly.Burst.ThrRel*(actMax-actAvg)
        thr = math32.Max(thr, ly.Burst.ThrAbs)
        for ni := range ly.Neurons {
                nrn := &ly.Neurons[ni]
                if nrn.IsOff() {
                        continue
                }
                burst := float32(0)
                if nrn.Act > thr {
                        burst = nrn.Act
                }
                nrn.Burst = burst
        }
}

// BurstAsAct records Burst activity as the activation directly.
func (ly *Layer) BurstAsAct(ctx *Context) {
        for ni := range ly.Neurons {
                nrn := &ly.Neurons[ni]
                nrn.Burst = nrn.Act
        }
}

//////// DeepCtxt -- once after Burst quarter

// SendCtxtGe sends Burst activation over CTCtxtPath pathways to integrate
// CtxtGe excitatory conductance on CT layers.
// This must be called at the end of the Burst quarter for this layer.
// Satisfies the CtxtSender interface.
func (ly *Layer) SendCtxtGe(ctx *Context) {
        if !ly.Burst.BurstQtr.HasFlag(ctx.Quarter) {
                return
        }
        for ni := range ly.Neurons {
                nrn := &ly.Neurons[ni]
                if nrn.IsOff() {
                        continue
                }
                if nrn.Burst > ly.Act.OptThresh.Send {
                        for _, sp := range ly.SendPaths {
                                if sp.Off {
                                        continue
                                }
                                if sp.Type != CTCtxtPath {
                                        continue
                                }
                                sp.SendCtxtGe(ni, nrn.Burst)
                        }
                }
        }
}

// CTGFromInc integrates new synaptic conductances from increments
// sent during last SendGDelta.
func (ly *Layer) CTGFromInc(ctx *Context) {
        for ni := range ly.Neurons {
                nrn := &ly.Neurons[ni]
                if nrn.IsOff() {
                        continue
                }
                geRaw := nrn.GeRaw + ly.Neurons[ni].CtxtGe
                ly.Act.GeFromRaw(nrn, geRaw)
                ly.Act.GiFromRaw(nrn, nrn.GiRaw)
        }
}

// CtxtFromGe integrates new CtxtGe excitatory conductance from pathways,
// and computes overall Ctxt value, only on Deep layers.
// This must be called at the end of the DeepBurst quarter for this layer,
// after SendCtxtGe.
func (ly *Layer) CtxtFromGe(ctx *Context) {
        if ly.Type != CTLayer {
                return
        }
        if !ly.Burst.BurstQtr.HasFlag(ctx.Quarter) {
                return
        }
        for ni := range ly.Neurons {
                ly.Neurons[ni].CtxtGe = 0
        }
        for _, pt := range ly.RecvPaths {
                if pt.Off {
                        continue
                }
                if pt.Type != CTCtxtPath {
                        continue
                }
                pt.RecvCtxtGeInc()
        }
}

//////// Pulvinar

// Driver describes the source of driver inputs from cortex into Pulvinar.
type Driver struct {

        // driver layer
        Driver string

        // offset into Pulvinar pool
        Off int `inactive:"-"`
}

// Drivers are a list of drivers
type Drivers []*Driver

// Add adds new driver(s)
func (dr *Drivers) Add(laynms ...string) {
        for _, laynm := range laynms {
                d := &Driver{}
                d.Driver = laynm
                *dr = append(*dr, d)
        }
}

// PulvinarParams provides parameters for how the plus-phase (outcome) state
// of thalamic relay cell (e.g., Pulvinar) neurons is computed from the
// corresponding driver neuron Burst activation.
type PulvinarParams struct {

        // Turn off the driver inputs, in which case this layer behaves like a standard layer
        DriversOff bool `default:"false"`

        // Quarter(s) when bursting occurs -- typically Q4 but can also be Q2 and Q4 for beta-frequency updating.  Note: this is a bitflag and must be accessed using its Set / Has etc routines
        BurstQtr Quarters

        // multiplier on driver input strength, multiplies activation of driver layer
        DriveScale float32 `default:"0.3" min:"0.0"`

        // Level of Max driver layer activation at which the predictive non-burst inputs are fully inhibited.  Computationally, it is essential that driver inputs inhibit effect of predictive non-driver (CTLayer) inputs, so that the plus phase is not always just the minus phase plus something extra (the error will never go to zero then).  When max driver act input exceeds this value, predictive non-driver inputs are fully suppressed.  If there is only weak burst input however, then the predictive inputs remain and this critically prevents the network from learning to turn activation off, which is difficult and severely degrades learning.
        MaxInhib float32 `default:"0.6" min:"0.01"`

        // Do not treat the pools in this layer as topographically organized relative to driver inputs -- all drivers compress down to give same input to all pools
        NoTopo bool

        // proportion of average across driver pools that is combined with Max to provide some graded tie-breaker signal -- especially important for large pool downsampling, e.g., when doing NoTopo
        AvgMix float32 `min:"0" max:"1"`

        // Apply threshold to driver burst input for computing plus-phase activations -- above BinThr, then Act = BinOn, below = BinOff.  This is beneficial for layers with weaker graded activations, such as V1 or other perceptual inputs.
        Binarize bool

        // Threshold for binarizing in terms of sending Burst activation
        BinThr float32

        // Resulting driver Ge value for units above threshold -- lower value around 0.3 or so seems best (DriveScale is NOT applied -- generally same range as that).
        BinOn float32 `default:"0.3"`

        // Resulting driver Ge value for units below threshold -- typically 0.
        BinOff float32 `default:"0"`
}

func (tp *PulvinarParams) Update() {
}

func (tp *PulvinarParams) Defaults() {
        tp.BurstQtr.SetFlag(true, Q4)
        tp.DriveScale = 0.3
        tp.MaxInhib = 0.6
        tp.Binarize = false
        tp.BinThr = 0.4
        tp.BinOn = 0.3
        tp.BinOff = 0
}

func (tp *PulvinarParams) ShouldDisplay(field string) bool {
        switch field {
        case "BinThr", "BinOn", "BinOff":
                return tp.Binarize
        }
        return true
}

// DriveGe returns effective excitatory conductance to use for given driver
// input Burst activation
func (tp *PulvinarParams) DriveGe(act float32) float32 {
        if tp.Binarize {
                if act >= tp.BinThr {
                        return tp.BinOn
                } else {
                        return tp.BinOff
                }
        } else {
                return tp.DriveScale * act
        }
}

// GeFromMaxAvg returns the drive Ge value as function of max and average
func (tp *PulvinarParams) GeFromMaxAvg(max, avg float32) float32 {
        deff := (1-tp.AvgMix)*max + tp.AvgMix*avg
        return tp.DriveGe(deff)
}

// UnitsSize returns the dimension of the units,
// either within a pool for 4D, or layer for 2D..
func UnitsSize(ly *Layer) (x, y int) {
        if ly.Is4D() {
                y = ly.Shape.DimSize(2)
                x = ly.Shape.DimSize(3)
        } else {
                y = ly.Shape.DimSize(0)
                x = ly.Shape.DimSize(1)
        }
        return
}

// DriverLayer returns the driver layer for given Driver
func (ly *Layer) DriverLayer(drv *Driver) (*Layer, error) {
        tly := ly.Network.LayerByName(drv.Driver)
        if tly == nil {
                err := fmt.Errorf("PulvinarLayer %s: Driver Layer not found", ly.Name)
                log.Println(err)
                return nil, err
        }
        return tly, nil
}

// SetDriverOffs sets the driver offsets.
func (ly *Layer) SetDriverOffs() error {
        if ly.Type != PulvinarLayer {
                return nil
        }
        mx, my := UnitsSize(ly)
        mn := my * mx
        off := 0
        var err error
        for _, drv := range ly.Drivers {
                dl, err := ly.DriverLayer(drv)
                if err != nil {
                        continue
                }
                drv.Off = off
                x, y := UnitsSize(dl)
                off += y * x
        }
        if off > mn {
                err = fmt.Errorf("PulvinarLayer %s: size of drivers: %d is greater than units: %d", ly.Name, off, mn)
                log.Println(err)
        }
        return err
}

func DriveAct(dni int, dly *Layer, issuper bool) float32 {
        act := float32(0)
        if issuper {
                act = dly.Neurons[dni].Burst
        } else {
                act = dly.Neurons[dni].Act
        }
        lmax := dly.Pools[0].Inhib.Act.Max // normalize by drive layer max act
        if lmax > 0.1 {                    // this puts all layers on equal footing for driving..
                return act / lmax
        }
        return act
}

// SetDriverNeuron sets the driver activation for given Neuron,
// based on given Ge driving value (use DriveFromMaxAvg) from driver layer (Burst or Act)
func (ly *Layer) SetDriverNeuron(tni int, drvGe, drvInhib float32) {
        if tni >= len(ly.Neurons) {
                return
        }
        nrn := &ly.Neurons[tni]
        if nrn.IsOff() {
                return
        }
        geRaw := (1-drvInhib)*nrn.GeRaw + drvGe
        ly.Act.GeFromRaw(nrn, geRaw)
        ly.Act.GiFromRaw(nrn, nrn.GiRaw)
}

// SetDriverActs sets the driver activations, integrating across all the driver layers
func (ly *Layer) SetDriverActs() {
        nux, nuy := UnitsSize(ly)
        nun := nux * nuy
        pyn := ly.Shape.DimSize(0)
        pxn := ly.Shape.DimSize(1)
        for _, drv := range ly.Drivers {
                dly, err := ly.DriverLayer(drv)
                if err != nil {
                        continue
                }
                issuper := dly.Type == SuperLayer
                drvMax := dly.Pools[0].Inhib.Act.Max
                drvInhib := math32.Min(1, drvMax/ly.Pulvinar.MaxInhib)

                if dly.Is2D() {
                        if ly.Is2D() {
                                for dni := range dly.Neurons {
                                        tni := drv.Off + dni
                                        drvAct := DriveAct(dni, dly, issuper)
                                        ly.SetDriverNeuron(tni, ly.Pulvinar.GeFromMaxAvg(drvAct, drvAct), drvInhib)
                                }
                        } else { // copy flat to all pools -- not typical
                                for dni := range dly.Neurons {
                                        drvAct := DriveAct(dni, dly, issuper)
                                        tni := drv.Off + dni
                                        for py := 0; py < pyn; py++ {
                                                for px := 0; px < pxn; px++ {
                                                        pni := (py*pxn+px)*nun + tni
                                                        ly.SetDriverNeuron(pni, ly.Pulvinar.GeFromMaxAvg(drvAct, drvAct), drvInhib)
                                                }
                                        }
                                }
                        }
                } else { // dly is 4D
                        dpyn := dly.Shape.DimSize(0)
                        dpxn := dly.Shape.DimSize(1)
                        duxn, duyn := UnitsSize(dly)
                        dnun := duxn * duyn
                        if ly.Is2D() {
                                for dni := 0; dni < dnun; dni++ {
                                        max := float32(0)
                                        avg := float32(0)
                                        avgn := 0
                                        for py := 0; py < dpyn; py++ {
                                                for px := 0; px < dpxn; px++ {
                                                        pi := (py*dpxn + px)
                                                        pni := pi*dnun + dni
                                                        act := DriveAct(pni, dly, issuper)
                                                        max = math32.Max(max, act)
                                                        pmax := dly.Pools[1+pi].Inhib.Act.Max
                                                        if pmax > 0.5 {
                                                                avg += act
                                                                avgn++
                                                        }
                                                }
                                        }
                                        if avgn > 0 {
                                                avg /= float32(avgn)
                                        }
                                        tni := drv.Off + dni
                                        ly.SetDriverNeuron(tni, ly.Pulvinar.GeFromMaxAvg(max, avg), drvInhib)
                                }
                        } else if ly.Pulvinar.NoTopo { // ly is 4D
                                for dni := 0; dni < dnun; dni++ {
                                        max := float32(0)
                                        avg := float32(0)
                                        avgn := 0
                                        for py := 0; py < dpyn; py++ {
                                                for px := 0; px < dpxn; px++ {
                                                        pi := (py*dpxn + px)
                                                        pni := pi*dnun + dni
                                                        act := DriveAct(pni, dly, issuper)
                                                        max = math32.Max(max, act)
                                                        pmax := dly.Pools[1+pi].Inhib.Act.Max
                                                        if pmax > 0.5 {
                                                                avg += act
                                                                avgn++
                                                        }
                                                }
                                        }
                                        if avgn > 0 {
                                                avg /= float32(avgn)
                                        }
                                        drvGe := ly.Pulvinar.GeFromMaxAvg(max, avg)
                                        tni := drv.Off + dni
                                        for py := 0; py < pyn; py++ {
                                                for px := 0; px < pxn; px++ {
                                                        pni := (py*pxn+px)*nun + tni
                                                        ly.SetDriverNeuron(pni, drvGe, drvInhib)
                                                }
                                        }
                                }
                        } else { // ly is 4D
                                pyr := float64(dpyn) / float64(pyn)
                                pxr := float64(dpxn) / float64(pxn)
                                for py := 0; py < pyn; py++ {
                                        sdpy := int(math.Round(float64(py) * pyr))
                                        edpy := int(math.Round(float64(py+1) * pyr))
                                        for px := 0; px < pxn; px++ {
                                                sdpx := int(math.Round(float64(px) * pxr))
                                                edpx := int(math.Round(float64(px+1) * pxr))
                                                pni := (py*pxn + px) * nun
                                                for dni := 0; dni < dnun; dni++ {
                                                        max := float32(0)
                                                        avg := float32(0)
                                                        avgn := 0
                                                        for dpy := sdpy; dpy < edpy; dpy++ {
                                                                for dpx := sdpx; dpx < edpx; dpx++ {
                                                                        pi := (dpy*dpxn + dpx)
                                                                        dpni := pi*dnun + dni
                                                                        act := DriveAct(dpni, dly, issuper)
                                                                        max = math32.Max(max, act)
                                                                        pmax := dly.Pools[1+pi].Inhib.Act.Max
                                                                        if pmax > 0.5 {
                                                                                avg += act
                                                                                avgn++
                                                                        }
                                                                }
                                                        }
                                                        if avgn > 0 {
                                                                avg /= float32(avgn)
                                                        }
                                                        tni := pni + drv.Off + dni
                                                        ly.SetDriverNeuron(tni, ly.Pulvinar.GeFromMaxAvg(max, avg), drvInhib)
                                                }
                                        }
                                }
                        }
                }
        }
}

// Copyright (c) 2019, The Emergent Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package leabra

import (
        "github.com/emer/emergent/v2/paths"
)

// AddSuperLayer2D adds a SuperLayer of given size, with given name.
func (nt *Network) AddSuperLayer2D(name string, nNeurY, nNeurX int) *Layer {
        return nt.AddLayer2D(name, nNeurY, nNeurX, SuperLayer)
}

// AddSuperLayer4D adds a SuperLayer of given size, with given name.
func (nt *Network) AddSuperLayer4D(name string, nPoolsY, nPoolsX, nNeurY, nNeurX int) *Layer {
        return nt.AddLayer4D(name, nPoolsY, nPoolsX, nNeurY, nNeurX, SuperLayer)
}

// AddCTLayer2D adds a CTLayer of given size, with given name.
func (nt *Network) AddCTLayer2D(name string, nNeurY, nNeurX int) *Layer {
        return nt.AddLayer2D(name, nNeurY, nNeurX, CTLayer)
}

// AddCTLayer4D adds a CTLayer of given size, with given name.
func (nt *Network) AddCTLayer4D(name string, nPoolsY, nPoolsX, nNeurY, nNeurX int) *Layer {
        return nt.AddLayer4D(name, nPoolsY, nPoolsX, nNeurY, nNeurX, CTLayer)
}

// AddPulvinarLayer2D adds a PulvinarLayer of given size, with given name.
func (nt *Network) AddPulvinarLayer2D(name string, nNeurY, nNeurX int) *Layer {
        return nt.AddLayer2D(name, nNeurY, nNeurX, PulvinarLayer)
}

// AddPulvinarLayer4D adds a PulvinarLayer of given size, with given name.
func (nt *Network) AddPulvinarLayer4D(name string, nPoolsY, nPoolsX, nNeurY, nNeurX int) *Layer {
        return nt.AddLayer4D(name, nPoolsY, nPoolsX, nNeurY, nNeurX, PulvinarLayer)
}

// ConnectSuperToCT adds a CTCtxtPath from given sending Super layer to a CT layer
// This automatically sets the FromSuper flag to engage proper defaults,
// uses a OneToOne path pattern, and sets the class to CTFromSuper
func (nt *Network) ConnectSuperToCT(send, recv *Layer) *Path {
        pt := nt.ConnectLayers(send, recv, paths.NewOneToOne(), CTCtxtPath)
        pt.AddClass("CTFromSuper")
        pt.FromSuper = true
        return pt
}

// ConnectCtxtToCT adds a CTCtxtPath from given sending layer to a CT layer
// Use ConnectSuperToCT for main pathway from corresponding superficial layer.
func (nt *Network) ConnectCtxtToCT(send, recv *Layer, pat paths.Pattern) *Path {
        return nt.ConnectLayers(send, recv, pat, CTCtxtPath)
}

// ConnectSuperToCTFake adds a FAKE CTCtxtPath from given sending
// Super layer to a CT layer uses a OneToOne path pattern,
// and sets the class to CTFromSuper.
// This does NOT make a CTCtxtPath -- instead makes a regular Path -- for testing!
func (nt *Network) ConnectSuperToCTFake(send, recv *Layer) *Path {
        pt := nt.ConnectLayers(send, recv, paths.NewOneToOne(), CTCtxtPath)
        pt.AddClass("CTFromSuper")
        return pt
}

// ConnectCtxtToCTFake adds a FAKE CTCtxtPath from given sending layer to a CT layer
// This does NOT make a CTCtxtPath -- instead makes a regular Path -- for testing!
func (nt *Network) ConnectCtxtToCTFake(send, recv *Layer, pat paths.Pattern) *Path {
        return nt.ConnectLayers(send, recv, pat, CTCtxtPath)
}

//////////////////////////////////////////////////////////////////////////////////////
//  Network versions of Add Layer methods

// AddDeep2D adds a superficial (SuperLayer) and corresponding CT (CT suffix) layer
// with CTCtxtPath OneToOne pathway from Super to CT, and Pulvinar Pulvinar for Super (P suffix).
// Pulvinar projects back to Super and CT layers, type = Back, class = FromPulv
// CT is placed Behind Super, and Pulvinar behind CT.
// Drivers must be added to the Pulvinar layer, and it must be sized appropriately for those drivers.
func (nt *Network) AddDeep2D(name string, shapeY, shapeX int) (super, ct, pulv *Layer) {
        super = nt.AddSuperLayer2D(name, shapeY, shapeX)
        ct = nt.AddCTLayer2D(name+"CT", shapeY, shapeX)
        ct.PlaceBehind(super, 2)
        full := paths.NewFull()
        nt.ConnectSuperToCT(super, ct)
        pulv = nt.AddPulvinarLayer2D(name+"P", shapeY, shapeX)
        pulv.PlaceBehind(ct, 2)
        nt.ConnectLayers(ct, pulv, full, ForwardPath)
        nt.ConnectLayers(pulv, super, full, BackPath).AddClass("FromPulv")
        nt.ConnectLayers(pulv, ct, full, BackPath).AddClass("FromPulv")
        return
}

// AddDeep4D adds a superficial (SuperLayer) and corresponding CT (CT suffix) layer
// with CTCtxtPath OneToOne pathway from Super to CT, and Pulvinar Pulvinar for Super (P suffix).
// Pulvinar projects back to Super and CT layers, also PoolOneToOne, class = FromPulv
// CT is placed Behind Super, and Pulvinar behind CT.
// Drivers must be added to the Pulvinar layer, and it must be sized appropriately for those drivers.
func (nt *Network) AddDeep4D(name string, nPoolsY, nPoolsX, nNeurY, nNeurX int) (super, ct, pulv *Layer) {
        super = nt.AddSuperLayer4D(name, nPoolsY, nPoolsX, nNeurY, nNeurX)
        ct = nt.AddCTLayer4D(name+"CT", nPoolsY, nPoolsX, nNeurY, nNeurX)
        ct.PlaceBehind(super, 2)
        pone2one := paths.NewPoolOneToOne()
        nt.ConnectSuperToCT(super, ct)
        pulv = nt.AddPulvinarLayer4D(name+"P", nPoolsY, nPoolsX, nNeurY, nNeurX)
        pulv.PlaceBehind(ct, 2)
        nt.ConnectLayers(ct, pulv, pone2one, ForwardPath)
        nt.ConnectLayers(pulv, super, pone2one, BackPath).AddClass("FromPulv")
        nt.ConnectLayers(pulv, ct, pone2one, BackPath).AddClass("FromPulv")
        return
}

// AddDeep2DFakeCT adds a superficial (SuperLayer) and corresponding CT (CT suffix) layer
// with FAKE CTCtxtPath OneToOne pathway from Super to CT, and Pulvinar Pulvinar for Super (P suffix).
// Pulvinar projects back to Super and CT layers, type = Back, class = FromPulv
// CT is placed Behind Super, and Pulvinar behind CT.
// Drivers must be added to the Pulvinar layer, and it must be sized appropriately for those drivers.
// This does NOT make a CTCtxtPath -- instead makes a regular Path -- for testing!
func (nt *Network) AddDeep2DFakeCT(name string, shapeY, shapeX int) (super, ct, pulv *Layer) {
        super = nt.AddSuperLayer2D(name, shapeY, shapeX)
        ct = nt.AddCTLayer2D(name+"CT", shapeY, shapeX)
        ct.PlaceBehind(super, 2)
        full := paths.NewFull()
        nt.ConnectSuperToCTFake(super, ct)
        pulv = nt.AddPulvinarLayer2D(name+"P", shapeY, shapeX)
        pulv.PlaceBehind(ct, 2)
        nt.ConnectLayers(ct, pulv, full, ForwardPath)
        nt.ConnectLayers(pulv, super, full, BackPath).AddClass("FromPulv")
        nt.ConnectLayers(pulv, ct, full, BackPath).AddClass("FromPulv")
        return
}

// AddDeep4DFakeCT adds a superficial (SuperLayer) and corresponding CT (CT suffix) layer
// with FAKE CTCtxtPath OneToOne pathway from Super to CT, and Pulvinar Pulvinar for Super (P suffix).
// Pulvinar projects back to Super and CT layers, also PoolOneToOne, class = FromPulv
// CT is placed Behind Super, and Pulvinar behind CT.
// Drivers must be added to the Pulvinar layer, and it must be sized appropriately for those drivers.
// This does NOT make a CTCtxtPath -- instead makes a regular Path -- for testing!
func (nt *Network) AddDeep4DFakeCT(name string, nPoolsY, nPoolsX, nNeurY, nNeurX int) (super, ct, pulv *Layer) {
        super = nt.AddSuperLayer4D(name, nPoolsY, nPoolsX, nNeurY, nNeurX)
        ct = nt.AddCTLayer4D(name+"CT", nPoolsY, nPoolsX, nNeurY, nNeurX)
        ct.PlaceBehind(super, 2)
        pone2one := paths.NewPoolOneToOne()
        nt.ConnectSuperToCTFake(super, ct)
        pulv = nt.AddPulvinarLayer4D(name+"P", nPoolsY, nPoolsX, nNeurY, nNeurX)
        pulv.PlaceBehind(ct, 2)
        nt.ConnectLayers(ct, pulv, pone2one, ForwardPath)
        nt.ConnectLayers(pulv, super, pone2one, BackPath).AddClass("FromPulv")
        nt.ConnectLayers(pulv, ct, pone2one, BackPath).AddClass("FromPulv")
        return
}

// AddDeepNoPulvinar2D adds a superficial (SuperLayer) and corresponding CT (CT suffix) layer
// with CTCtxtPath OneToOne pathway from Super to CT, and NO Pulvinar Pulvinar.
// CT is placed Behind Super.
func (nt *Network) AddDeepNoPulvinar2D(name string, shapeY, shapeX int) (super, ct *Layer) {
        super = nt.AddSuperLayer2D(name, shapeY, shapeX)
        ct = nt.AddCTLayer2D(name+"CT", shapeY, shapeX)
        ct.PlaceBehind(super, 2)
        nt.ConnectSuperToCT(super, ct)
        return
}

// AddDeepNoPulvinar4D adds a superficial (SuperLayer) and corresponding CT (CT suffix) layer
// with CTCtxtPath PoolOneToOne pathway from Super to CT, and NO Pulvinar Pulvinar.
// CT is placed Behind Super.
func (nt *Network) AddDeepNoPulvinar4D(name string, nPoolsY, nPoolsX, nNeurY, nNeurX int) (super, ct *Layer) {
        super = nt.AddSuperLayer4D(name, nPoolsY, nPoolsX, nNeurY, nNeurX)
        ct = nt.AddCTLayer4D(name+"CT", nPoolsY, nPoolsX, nNeurY, nNeurX)
        ct.PlaceBehind(super, 2)
        nt.ConnectSuperToCT(super, ct)
        return
}

// Copyright (c) 2024, The Emergent Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package leabra

import (
        "cogentcore.org/core/math32"
)

func (pt *Path) CTCtxtDefaults() {
        if pt.FromSuper {
                pt.Learn.Learn = false
                pt.WtInit.Mean = 0.5 // .5 better than .8 in several cases..
                pt.WtInit.Var = 0
        }
}

// SendCtxtGe sends the full Burst activation from sending neuron index si,
// to integrate CtxtGe excitatory conductance on receivers
func (pt *Path) SendCtxtGe(si int, dburst float32) {
        scdb := dburst * pt.GScale
        nc := pt.SConN[si]
        st := pt.SConIndexSt[si]
        syns := pt.Syns[st : st+nc]
        scons := pt.SConIndex[st : st+nc]
        for ci := range syns {
                ri := scons[ci]
                pt.CtxtGeInc[ri] += scdb * syns[ci].Wt
        }
}

// RecvCtxtGeInc increments the receiver's CtxtGe from that of all the pathways
func (pt *Path) RecvCtxtGeInc() {
        rlay := pt.Recv
        for ri := range rlay.Neurons {
                rlay.Neurons[ri].CtxtGe += pt.CtxtGeInc[ri]
                pt.CtxtGeInc[ri] = 0
        }
}

// DWt computes the weight change (learning) for CTCtxt pathways.
func (pt *Path) DWtCTCtxt() {
        slay := pt.Send
        issuper := pt.Send.Type == SuperLayer
        rlay := pt.Recv
        for si := range slay.Neurons {
                sact := float32(0)
                if issuper {
                        sact = slay.Neurons[si].BurstPrv
                } else {
                        sact = slay.Neurons[si].ActQ0
                }
                nc := int(pt.SConN[si])
                st := int(pt.SConIndexSt[si])
                syns := pt.Syns[st : st+nc]
                scons := pt.SConIndex[st : st+nc]
                for ci := range syns {
                        sy := &syns[ci]
                        ri := scons[ci]
                        rn := &rlay.Neurons[ri]
                        // following line should be ONLY diff: sact for *both* short and medium *sender*
                        // activations, which are first two args:
                        err, bcm := pt.Learn.CHLdWt(sact, sact, rn.AvgSLrn, rn.AvgM, rn.AvgL)

                        bcm *= pt.Learn.XCal.LongLrate(rn.AvgLLrn)
                        err *= pt.Learn.XCal.MLrn
                        dwt := bcm + err
                        norm := float32(1)
                        if pt.Learn.Norm.On {
                                norm = pt.Learn.Norm.NormFromAbsDWt(&sy.Norm, math32.Abs(dwt))
                        }
                        if pt.Learn.Momentum.On {
                                dwt = norm * pt.Learn.Momentum.MomentFromDWt(&sy.Moment, dwt)
                        } else {
                                dwt *= norm
                        }
                        sy.DWt += pt.Learn.Lrate * dwt
                }
                // aggregate max DWtNorm over sending synapses
                if pt.Learn.Norm.On {
                        maxNorm := float32(0)
                        for ci := range syns {
                                sy := &syns[ci]
                                if sy.Norm > maxNorm {
                                        maxNorm = sy.Norm
                                }
                        }
                        for ci := range syns {
                                sy := &syns[ci]
                                sy.Norm = maxNorm
                        }
                }
        }
}

// Code generated by "core generate -add-types"; DO NOT EDIT.

package leabra

import (
        "cogentcore.org/core/enums"
)

var _ActNoiseTypeValues = []ActNoiseType{0, 1, 2, 3, 4}

// ActNoiseTypeN is the highest valid value for type ActNoiseType, plus one.
const ActNoiseTypeN ActNoiseType = 5

var _ActNoiseTypeValueMap = map[string]ActNoiseType{`NoNoise`: 0, `VmNoise`: 1, `GeNoise`: 2, `ActNoise`: 3, `GeMultNoise`: 4}

var _ActNoiseTypeDescMap = map[ActNoiseType]string{0: `NoNoise means no noise added`, 1: `VmNoise means noise is added to the membrane potential. IMPORTANT: this should NOT be used for rate-code (NXX1) activations, because they do not depend directly on the vm -- this then has no effect`, 2: `GeNoise means noise is added to the excitatory conductance (Ge). This should be used for rate coded activations (NXX1)`, 3: `ActNoise means noise is added to the final rate code activation`, 4: `GeMultNoise means that noise is multiplicative on the Ge excitatory conductance values`}

var _ActNoiseTypeMap = map[ActNoiseType]string{0: `NoNoise`, 1: `VmNoise`, 2: `GeNoise`, 3: `ActNoise`, 4: `GeMultNoise`}

// String returns the string representation of this ActNoiseType value.
func (i ActNoiseType) String() string { return enums.String(i, _ActNoiseTypeMap) }

// SetString sets the ActNoiseType value from its string representation,
// and returns an error if the string is invalid.
func (i *ActNoiseType) SetString(s string) error {
        return enums.SetString(i, s, _ActNoiseTypeValueMap, "ActNoiseType")
}

// Int64 returns the ActNoiseType value as an int64.
func (i ActNoiseType) Int64() int64 { return int64(i) }

// SetInt64 sets the ActNoiseType value from an int64.
func (i *ActNoiseType) SetInt64(in int64) { *i = ActNoiseType(in) }

// Desc returns the description of the ActNoiseType value.
func (i ActNoiseType) Desc() string { return enums.Desc(i, _ActNoiseTypeDescMap) }

// ActNoiseTypeValues returns all possible values for the type ActNoiseType.
func ActNoiseTypeValues() []ActNoiseType { return _ActNoiseTypeValues }

// Values returns all possible values for the type ActNoiseType.
func (i ActNoiseType) Values() []enums.Enum { return enums.Values(_ActNoiseTypeValues) }

// MarshalText implements the [encoding.TextMarshaler] interface.
func (i ActNoiseType) MarshalText() ([]byte, error) { return []byte(i.String()), nil }

// UnmarshalText implements the [encoding.TextUnmarshaler] interface.
func (i *ActNoiseType) UnmarshalText(text []byte) error {
        return enums.UnmarshalText(i, text, "ActNoiseType")
}

var _QuartersValues = []Quarters{0, 1, 2, 3}

// QuartersN is the highest valid value for type Quarters, plus one.
const QuartersN Quarters = 4

var _QuartersValueMap = map[string]Quarters{`Q1`: 0, `Q2`: 1, `Q3`: 2, `Q4`: 3}

var _QuartersDescMap = map[Quarters]string{0: `Q1 is the first quarter, which, due to 0-based indexing, shows up as Quarter = 0 in timer`, 1: ``, 2: ``, 3: ``}

var _QuartersMap = map[Quarters]string{0: `Q1`, 1: `Q2`, 2: `Q3`, 3: `Q4`}

// String returns the string representation of this Quarters value.
func (i Quarters) String() string { return enums.BitFlagString(i, _QuartersValues) }

// BitIndexString returns the string representation of this Quarters value
// if it is a bit index value (typically an enum constant), and
// not an actual bit flag value.
func (i Quarters) BitIndexString() string { return enums.String(i, _QuartersMap) }

// SetString sets the Quarters value from its string representation,
// and returns an error if the string is invalid.
func (i *Quarters) SetString(s string) error { *i = 0; return i.SetStringOr(s) }

// SetStringOr sets the Quarters value from its string representation
// while preserving any bit flags already set, and returns an
// error if the string is invalid.
func (i *Quarters) SetStringOr(s string) error {
        return enums.SetStringOr(i, s, _QuartersValueMap, "Quarters")
}

// Int64 returns the Quarters value as an int64.
func (i Quarters) Int64() int64 { return int64(i) }

// SetInt64 sets the Quarters value from an int64.
func (i *Quarters) SetInt64(in int64) { *i = Quarters(in) }

// Desc returns the description of the Quarters value.
func (i Quarters) Desc() string { return enums.Desc(i, _QuartersDescMap) }

// QuartersValues returns all possible values for the type Quarters.
func QuartersValues() []Quarters { return _QuartersValues }

// Values returns all possible values for the type Quarters.
func (i Quarters) Values() []enums.Enum { return enums.Values(_QuartersValues) }

// HasFlag returns whether these bit flags have the given bit flag set.
func (i *Quarters) HasFlag(f enums.BitFlag) bool { return enums.HasFlag((*int64)(i), f) }

// SetFlag sets the value of the given flags in these flags to the given value.
func (i *Quarters) SetFlag(on bool, f ...enums.BitFlag) { enums.SetFlag((*int64)(i), on, f...) }

// MarshalText implements the [encoding.TextMarshaler] interface.
func (i Quarters) MarshalText() ([]byte, error) { return []byte(i.String()), nil }

// UnmarshalText implements the [encoding.TextUnmarshaler] interface.
func (i *Quarters) UnmarshalText(text []byte) error { return enums.UnmarshalText(i, text, "Quarters") }

var _LayerTypesValues = []LayerTypes{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18}

// LayerTypesN is the highest valid value for type LayerTypes, plus one.
const LayerTypesN LayerTypes = 19

var _LayerTypesValueMap = map[string]LayerTypes{`SuperLayer`: 0, `InputLayer`: 1, `TargetLayer`: 2, `CompareLayer`: 3, `CTLayer`: 4, `PulvinarLayer`: 5, `TRNLayer`: 6, `ClampDaLayer`: 7, `RWPredLayer`: 8, `RWDaLayer`: 9, `TDPredLayer`: 10, `TDIntegLayer`: 11, `TDDaLayer`: 12, `MatrixLayer`: 13, `GPeLayer`: 14, `GPiThalLayer`: 15, `CINLayer`: 16, `PFCLayer`: 17, `PFCDeepLayer`: 18}

var _LayerTypesDescMap = map[LayerTypes]string{0: `Super is a superficial cortical layer (lamina 2-3-4) which does not receive direct input or targets. In more generic models, it should be used as a Hidden layer, and maps onto the Hidden type in LayerTypes.`, 1: `Input is a layer that receives direct external input in its Ext inputs. Biologically, it can be a primary sensory layer, or a thalamic layer.`, 2: `Target is a layer that receives direct external target inputs used for driving plus-phase learning. Simple target layers are generally not used in more biological models, which instead use predictive learning via Pulvinar or related mechanisms.`, 3: `Compare is a layer that receives external comparison inputs, which drive statistics but do NOT drive activation or learning directly. It is rarely used in axon.`, 4: `CT are layer 6 corticothalamic projecting neurons, which drive "top down" predictions in Pulvinar layers. They maintain information over time via stronger NMDA channels and use maintained prior state information to generate predictions about current states forming on Super layers that then drive PT (5IB) bursting activity, which are the plus-phase drivers of Pulvinar activity.`, 5: `Pulvinar are thalamic relay cell neurons in the higher-order Pulvinar nucleus of the thalamus, and functionally isomorphic neurons in the MD thalamus, and potentially other areas. These cells alternately reflect predictions driven by CT pathways, and actual outcomes driven by 5IB Burst activity from corresponding PT or Super layer neurons that provide strong driving inputs.`, 6: `TRNLayer is thalamic reticular nucleus layer for inhibitory competition within the thalamus.`, 7: `ClampDaLayer is an Input layer that just sends its activity as the dopamine signal.`, 8: `RWPredLayer computes reward prediction for a simple Rescorla-Wagner learning dynamic (i.e., PV learning in the PVLV framework). Activity is computed as linear function of excitatory conductance (which can be negative -- there are no constraints). Use with [RWPath] which does simple delta-rule learning on minus-plus.`, 9: `RWDaLayer computes a dopamine (DA) signal based on a simple Rescorla-Wagner learning dynamic (i.e., PV learning in the PVLV framework). It computes difference between r(t) and [RWPredLayer] values. r(t) is accessed directly from a Rew layer -- if no external input then no DA is computed -- critical for effective use of RW only for PV cases. RWPred prediction is also accessed directly from Rew layer to avoid any issues.`, 10: `TDPredLayer is the temporal differences reward prediction layer. It represents estimated value V(t) in the minus phase, and computes estimated V(t+1) based on its learned weights in plus phase. Use [TDPredPath] for DA modulated learning.`, 11: `TDIntegLayer is the temporal differences reward integration layer. It represents estimated value V(t) in the minus phase, and estimated V(t+1) + r(t) in the plus phase. It computes r(t) from (typically fixed) weights from a reward layer, and directly accesses values from [TDPredLayer].`, 12: `TDDaLayer computes a dopamine (DA) signal as the temporal difference (TD) between the [TDIntegLayer[] activations in the minus and plus phase.`, 13: `MatrixLayer represents the dorsal matrisome MSN's that are the main Go / NoGo gating units in BG driving updating of PFC WM in PBWM. D1R = Go, D2R = NoGo, and outer 4D Pool X dimension determines GateTypes per MaintN (Maint on the left up to MaintN, Out on the right after)`, 14: `GPeLayer is a Globus pallidus external layer, a key region of the basal ganglia. It does not require any additional mechanisms beyond the SuperLayer.`, 15: `GPiThalLayer represents the combined Winner-Take-All dynamic of GPi (SNr) and Thalamus. It is the final arbiter of gating in the BG, weighing Go (direct) and NoGo (indirect) inputs from MatrixLayers (indirectly via GPe layer in case of NoGo). Use 4D structure for this so it matches 4D structure in Matrix layers`, 16: `CINLayer (cholinergic interneuron) reads reward signals from named source layer(s) and sends the Max absolute value of that activity as the positively rectified non-prediction-discounted reward signal computed by CINs, and sent as an acetylcholine (ACh) signal. To handle positive-only reward signals, need to include both a reward prediction and reward outcome layer.`, 17: `PFCLayer is a prefrontal cortex layer, either superficial or output. See [PFCDeepLayer] for the deep maintenance layer.`, 18: `PFCDeepLayer is a prefrontal cortex deep maintenance layer.`}

var _LayerTypesMap = map[LayerTypes]string{0: `SuperLayer`, 1: `InputLayer`, 2: `TargetLayer`, 3: `CompareLayer`, 4: `CTLayer`, 5: `PulvinarLayer`, 6: `TRNLayer`, 7: `ClampDaLayer`, 8: `RWPredLayer`, 9: `RWDaLayer`, 10: `TDPredLayer`, 11: `TDIntegLayer`, 12: `TDDaLayer`, 13: `MatrixLayer`, 14: `GPeLayer`, 15: `GPiThalLayer`, 16: `CINLayer`, 17: `PFCLayer`, 18: `PFCDeepLayer`}

// String returns the string representation of this LayerTypes value.
func (i LayerTypes) String() string { return enums.String(i, _LayerTypesMap) }

// SetString sets the LayerTypes value from its string representation,
// and returns an error if the string is invalid.
func (i *LayerTypes) SetString(s string) error {
        return enums.SetString(i, s, _LayerTypesValueMap, "LayerTypes")
}

// Int64 returns the LayerTypes value as an int64.
func (i LayerTypes) Int64() int64 { return int64(i) }

// SetInt64 sets the LayerTypes value from an int64.
func (i *LayerTypes) SetInt64(in int64) { *i = LayerTypes(in) }

// Desc returns the description of the LayerTypes value.
func (i LayerTypes) Desc() string { return enums.Desc(i, _LayerTypesDescMap) }

// LayerTypesValues returns all possible values for the type LayerTypes.
func LayerTypesValues() []LayerTypes { return _LayerTypesValues }

// Values returns all possible values for the type LayerTypes.
func (i LayerTypes) Values() []enums.Enum { return enums.Values(_LayerTypesValues) }

// MarshalText implements the [encoding.TextMarshaler] interface.
func (i LayerTypes) MarshalText() ([]byte, error) { return []byte(i.String()), nil }

// UnmarshalText implements the [encoding.TextUnmarshaler] interface.
func (i *LayerTypes) UnmarshalText(text []byte) error {
        return enums.UnmarshalText(i, text, "LayerTypes")
}

var _DaReceptorsValues = []DaReceptors{0, 1}

// DaReceptorsN is the highest valid value for type DaReceptors, plus one.
const DaReceptorsN DaReceptors = 2

var _DaReceptorsValueMap = map[string]DaReceptors{`D1R`: 0, `D2R`: 1}

var _DaReceptorsDescMap = map[DaReceptors]string{0: `D1R primarily expresses Dopamine D1 Receptors -- dopamine is excitatory and bursts of dopamine lead to increases in synaptic weight, while dips lead to decreases -- direct pathway in dorsal striatum`, 1: `D2R primarily expresses Dopamine D2 Receptors -- dopamine is inhibitory and bursts of dopamine lead to decreases in synaptic weight, while dips lead to increases -- indirect pathway in dorsal striatum`}

var _DaReceptorsMap = map[DaReceptors]string{0: `D1R`, 1: `D2R`}

// String returns the string representation of this DaReceptors value.
func (i DaReceptors) String() string { return enums.String(i, _DaReceptorsMap) }

// SetString sets the DaReceptors value from its string representation,
// and returns an error if the string is invalid.
func (i *DaReceptors) SetString(s string) error {
        return enums.SetString(i, s, _DaReceptorsValueMap, "DaReceptors")
}

// Int64 returns the DaReceptors value as an int64.
func (i DaReceptors) Int64() int64 { return int64(i) }

// SetInt64 sets the DaReceptors value from an int64.
func (i *DaReceptors) SetInt64(in int64) { *i = DaReceptors(in) }

// Desc returns the description of the DaReceptors value.
func (i DaReceptors) Desc() string { return enums.Desc(i, _DaReceptorsDescMap) }

// DaReceptorsValues returns all possible values for the type DaReceptors.
func DaReceptorsValues() []DaReceptors { return _DaReceptorsValues }

// Values returns all possible values for the type DaReceptors.
func (i DaReceptors) Values() []enums.Enum { return enums.Values(_DaReceptorsValues) }

// MarshalText implements the [encoding.TextMarshaler] interface.
func (i DaReceptors) MarshalText() ([]byte, error) { return []byte(i.String()), nil }

// UnmarshalText implements the [encoding.TextUnmarshaler] interface.
func (i *DaReceptors) UnmarshalText(text []byte) error {
        return enums.UnmarshalText(i, text, "DaReceptors")
}

var _ValencesValues = []Valences{0, 1}

// ValencesN is the highest valid value for type Valences, plus one.
const ValencesN Valences = 2

var _ValencesValueMap = map[string]Valences{`Appetitive`: 0, `Aversive`: 1}

var _ValencesDescMap = map[Valences]string{0: `Appetititve is a positive valence US (food, water, etc)`, 1: `Aversive is a negative valence US (shock, threat etc)`}

var _ValencesMap = map[Valences]string{0: `Appetitive`, 1: `Aversive`}

// String returns the string representation of this Valences value.
func (i Valences) String() string { return enums.String(i, _ValencesMap) }

// SetString sets the Valences value from its string representation,
// and returns an error if the string is invalid.
func (i *Valences) SetString(s string) error {
        return enums.SetString(i, s, _ValencesValueMap, "Valences")
}

// Int64 returns the Valences value as an int64.
func (i Valences) Int64() int64 { return int64(i) }

// SetInt64 sets the Valences value from an int64.
func (i *Valences) SetInt64(in int64) { *i = Valences(in) }

// Desc returns the description of the Valences value.
func (i Valences) Desc() string { return enums.Desc(i, _ValencesDescMap) }

// ValencesValues returns all possible values for the type Valences.
func ValencesValues() []Valences { return _ValencesValues }

// Values returns all possible values for the type Valences.
func (i Valences) Values() []enums.Enum { return enums.Values(_ValencesValues) }

// MarshalText implements the [encoding.TextMarshaler] interface.
func (i Valences) MarshalText() ([]byte, error) { return []byte(i.String()), nil }

// UnmarshalText implements the [encoding.TextUnmarshaler] interface.
func (i *Valences) UnmarshalText(text []byte) error { return enums.UnmarshalText(i, text, "Valences") }

var _NeurFlagsValues = []NeurFlags{0, 1, 2, 3}

// NeurFlagsN is the highest valid value for type NeurFlags, plus one.
const NeurFlagsN NeurFlags = 4

var _NeurFlagsValueMap = map[string]NeurFlags{`NeurOff`: 0, `NeurHasExt`: 1, `NeurHasTarg`: 2, `NeurHasCmpr`: 3}

var _NeurFlagsDescMap = map[NeurFlags]string{0: `NeurOff flag indicates that this neuron has been turned off (i.e., lesioned)`, 1: `NeurHasExt means the neuron has external input in its Ext field`, 2: `NeurHasTarg means the neuron has external target input in its Targ field`, 3: `NeurHasCmpr means the neuron has external comparison input in its Targ field -- used for computing comparison statistics but does not drive neural activity ever`}

var _NeurFlagsMap = map[NeurFlags]string{0: `NeurOff`, 1: `NeurHasExt`, 2: `NeurHasTarg`, 3: `NeurHasCmpr`}

// String returns the string representation of this NeurFlags value.
func (i NeurFlags) String() string { return enums.BitFlagString(i, _NeurFlagsValues) }

// BitIndexString returns the string representation of this NeurFlags value
// if it is a bit index value (typically an enum constant), and
// not an actual bit flag value.
func (i NeurFlags) BitIndexString() string { return enums.String(i, _NeurFlagsMap) }

// SetString sets the NeurFlags value from its string representation,
// and returns an error if the string is invalid.
func (i *NeurFlags) SetString(s string) error { *i = 0; return i.SetStringOr(s) }

// SetStringOr sets the NeurFlags value from its string representation
// while preserving any bit flags already set, and returns an
// error if the string is invalid.
func (i *NeurFlags) SetStringOr(s string) error {
        return enums.SetStringOr(i, s, _NeurFlagsValueMap, "NeurFlags")
}

// Int64 returns the NeurFlags value as an int64.
func (i NeurFlags) Int64() int64 { return int64(i) }

// SetInt64 sets the NeurFlags value from an int64.
func (i *NeurFlags) SetInt64(in int64) { *i = NeurFlags(in) }

// Desc returns the description of the NeurFlags value.
func (i NeurFlags) Desc() string { return enums.Desc(i, _NeurFlagsDescMap) }

// NeurFlagsValues returns all possible values for the type NeurFlags.
func NeurFlagsValues() []NeurFlags { return _NeurFlagsValues }

// Values returns all possible values for the type NeurFlags.
func (i NeurFlags) Values() []enums.Enum { return enums.Values(_NeurFlagsValues) }

// HasFlag returns whether these bit flags have the given bit flag set.
func (i *NeurFlags) HasFlag(f enums.BitFlag) bool { return enums.HasFlag((*int64)(i), f) }

// SetFlag sets the value of the given flags in these flags to the given value.
func (i *NeurFlags) SetFlag(on bool, f ...enums.BitFlag) { enums.SetFlag((*int64)(i), on, f...) }

// MarshalText implements the [encoding.TextMarshaler] interface.
func (i NeurFlags) MarshalText() ([]byte, error) { return []byte(i.String()), nil }

// UnmarshalText implements the [encoding.TextUnmarshaler] interface.
func (i *NeurFlags) UnmarshalText(text []byte) error {
        return enums.UnmarshalText(i, text, "NeurFlags")
}

var _PathTypesValues = []PathTypes{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11}

// PathTypesN is the highest valid value for type PathTypes, plus one.
const PathTypesN PathTypes = 12

var _PathTypesValueMap = map[string]PathTypes{`ForwardPath`: 0, `BackPath`: 1, `LateralPath`: 2, `InhibPath`: 3, `CTCtxtPath`: 4, `CHLPath`: 5, `EcCa1Path`: 6, `RWPath`: 7, `TDPredPath`: 8, `MatrixPath`: 9, `GPiThalPath`: 10, `DaHebbPath`: 11}

var _PathTypesDescMap = map[PathTypes]string{0: `Forward is a feedforward, bottom-up pathway from sensory inputs to higher layers`, 1: `Back is a feedback, top-down pathway from higher layers back to lower layers`, 2: `Lateral is a lateral pathway within the same layer / area`, 3: `Inhib is an inhibitory pathway that drives inhibitory synaptic conductances instead of the default excitatory ones.`, 4: `CTCtxt are pathways from Superficial layers to CT layers that send Burst activations drive updating of CtxtGe excitatory conductance, at end of plus (51B Bursting) phase. Biologically, this pathway comes from the PT layer 5IB neurons, but it is simpler to use the Super neurons directly, and PT are optional for most network types. These pathways also use a special learning rule that takes into account the temporal delays in the activation states. Can also add self context from CT for deeper temporal context.`, 5: `CHLPath implements Contrastive Hebbian Learning.`, 6: `EcCa1Path implements special learning for EC <-> CA1 pathways in the hippocampus to perform error-driven learning of this encoder pathway according to the ThetaPhase algorithm. uses Contrastive Hebbian Learning (CHL) on ActP - ActQ1 Q1: ECin -> CA1 -> ECout : ActQ1 = minus phase for auto-encoder Q2, 3: CA3 -> CA1 -> ECout : ActM = minus phase for recall Q4: ECin -> CA1, ECin -> ECout : ActP = plus phase for everything`, 7: `RWPath does dopamine-modulated learning for reward prediction: Da * Send.Act Use in RWPredLayer typically to generate reward predictions. Has no weight bounds or limits on sign etc.`, 8: `TDPredPath does dopamine-modulated learning for reward prediction: DWt = Da * Send.ActQ0 (activity on *previous* timestep) Use in TDPredLayer typically to generate reward predictions. Has no weight bounds or limits on sign etc.`, 9: `MatrixPath does dopamine-modulated, gated trace learning, for Matrix learning in PBWM context.`, 10: `GPiThalPath accumulates per-path raw conductance that is needed for separately weighting NoGo vs. Go inputs.`, 11: `DaHebbPath does dopamine-modulated Hebbian learning -- i.e., the 3-factor learning rule: Da * Recv.Act * Send.Act`}

var _PathTypesMap = map[PathTypes]string{0: `ForwardPath`, 1: `BackPath`, 2: `LateralPath`, 3: `InhibPath`, 4: `CTCtxtPath`, 5: `CHLPath`, 6: `EcCa1Path`, 7: `RWPath`, 8: `TDPredPath`, 9: `MatrixPath`, 10: `GPiThalPath`, 11: `DaHebbPath`}

// String returns the string representation of this PathTypes value.
func (i PathTypes) String() string { return enums.String(i, _PathTypesMap) }

// SetString sets the PathTypes value from its string representation,
// and returns an error if the string is invalid.
func (i *PathTypes) SetString(s string) error {
        return enums.SetString(i, s, _PathTypesValueMap, "PathTypes")
}

// Int64 returns the PathTypes value as an int64.
func (i PathTypes) Int64() int64 { return int64(i) }

// SetInt64 sets the PathTypes value from an int64.
func (i *PathTypes) SetInt64(in int64) { *i = PathTypes(in) }

// Desc returns the description of the PathTypes value.
func (i PathTypes) Desc() string { return enums.Desc(i, _PathTypesDescMap) }

// PathTypesValues returns all possible values for the type PathTypes.
func PathTypesValues() []PathTypes { return _PathTypesValues }

// Values returns all possible values for the type PathTypes.
func (i PathTypes) Values() []enums.Enum { return enums.Values(_PathTypesValues) }

// MarshalText implements the [encoding.TextMarshaler] interface.
func (i PathTypes) MarshalText() ([]byte, error) { return []byte(i.String()), nil }

// UnmarshalText implements the [encoding.TextUnmarshaler] interface.
func (i *PathTypes) UnmarshalText(text []byte) error {
        return enums.UnmarshalText(i, text, "PathTypes")
}

var _GateTypesValues = []GateTypes{0, 1, 2}

// GateTypesN is the highest valid value for type GateTypes, plus one.
const GateTypesN GateTypes = 3

var _GateTypesValueMap = map[string]GateTypes{`Maint`: 0, `Out`: 1, `MaintOut`: 2}

var _GateTypesDescMap = map[GateTypes]string{0: `Maint is maintenance gating -- toggles active maintenance in PFC.`, 1: `Out is output gating -- drives deep layer activation.`, 2: `MaintOut for maint and output gating.`}

var _GateTypesMap = map[GateTypes]string{0: `Maint`, 1: `Out`, 2: `MaintOut`}

// String returns the string representation of this GateTypes value.
func (i GateTypes) String() string { return enums.String(i, _GateTypesMap) }

// SetString sets the GateTypes value from its string representation,
// and returns an error if the string is invalid.
func (i *GateTypes) SetString(s string) error {
        return enums.SetString(i, s, _GateTypesValueMap, "GateTypes")
}

// Int64 returns the GateTypes value as an int64.
func (i GateTypes) Int64() int64 { return int64(i) }

// SetInt64 sets the GateTypes value from an int64.
func (i *GateTypes) SetInt64(in int64) { *i = GateTypes(in) }

// Desc returns the description of the GateTypes value.
func (i GateTypes) Desc() string { return enums.Desc(i, _GateTypesDescMap) }

// GateTypesValues returns all possible values for the type GateTypes.
func GateTypesValues() []GateTypes { return _GateTypesValues }

// Values returns all possible values for the type GateTypes.
func (i GateTypes) Values() []enums.Enum { return enums.Values(_GateTypesValues) }

// MarshalText implements the [encoding.TextMarshaler] interface.
func (i GateTypes) MarshalText() ([]byte, error) { return []byte(i.String()), nil }

// UnmarshalText implements the [encoding.TextUnmarshaler] interface.
func (i *GateTypes) UnmarshalText(text []byte) error {
        return enums.UnmarshalText(i, text, "GateTypes")
}

// Copyright (c) 2022, The Emergent Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package leabra

import (
        "fmt"

        "cogentcore.org/core/core"
        "cogentcore.org/lab/base/mpi"
        "github.com/emer/emergent/v2/ecmd"
)

////////////////////////////////////////////////////
// Misc

// ToggleLayersOff can be used to disable layers in a Network, for example if you are doing an ablation study.
func ToggleLayersOff(net *Network, layerNames []string, off bool) {
        for _, lnm := range layerNames {
                lyi := net.LayerByName(lnm)
                if lyi == nil {
                        fmt.Printf("layer not found: %s\n", lnm)
                        continue
                }
                lyi.Off = off
        }
}

/////////////////////////////////////////////
// Weights files

// WeightsFilename returns default current weights file name,
// using train run and epoch counters from looper
// and the RunName string identifying tag, parameters and starting run,
func WeightsFilename(net *Network, ctrString, runName string) string {
        return net.Name + "_" + runName + "_" + ctrString + ".wts.gz"
}

// SaveWeights saves network weights to filename with WeightsFilename information
// to identify the weights.
// only for 0 rank MPI if running mpi
// Returns the name of the file saved to, or empty if not saved.
func SaveWeights(net *Network, ctrString, runName string) string {
        if mpi.WorldRank() > 0 {
                return ""
        }
        fnm := WeightsFilename(net, ctrString, runName)
        fmt.Printf("Saving Weights to: %s\n", fnm)
        net.SaveWeightsJSON(core.Filename(fnm))
        return fnm
}

// SaveWeightsIfArgSet saves network weights if the "wts" arg has been set to true.
// uses WeightsFilename information to identify the weights.
// only for 0 rank MPI if running mpi
// Returns the name of the file saved to, or empty if not saved.
func SaveWeightsIfArgSet(net *Network, args *ecmd.Args, ctrString, runName string) string {
        if args.Bool("wts") {
                return SaveWeights(net, ctrString, runName)
        }
        return ""
}

// SaveWeightsIfConfigSet saves network weights if the given config
// bool value has been set to true.
// uses WeightsFilename information to identify the weights.
// only for 0 rank MPI if running mpi
// Returns the name of the file saved to, or empty if not saved.
func SaveWeightsIfConfigSet(net *Network, cfgWts bool, ctrString, runName string) string {
        if cfgWts {
                return SaveWeights(net, ctrString, runName)
        }
        return ""
}

// Copyright (c) 2024, The Emergent Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package leabra

import (
        "cogentcore.org/core/base/errors"
        "cogentcore.org/core/math32"
        "github.com/emer/emergent/v2/etime"
        "github.com/emer/emergent/v2/looper"
)

// Contrastive Hebbian Learning (CHL) parameters
type CHLParams struct {

        // if true, use CHL learning instead of standard XCAL learning -- allows easy exploration of CHL vs. XCAL
        On bool

        // amount of hebbian learning (should be relatively small, can be effective at .0001)
        Hebb float32 `default:"0.001" min:"0" max:"1"`

        // amount of error driven learning, automatically computed to be 1-Hebb
        Err float32 `default:"0.999" min:"0" max:"1" edit:"-"`

        // if true, use ActQ1 as the minus phase -- otherwise ActM
        MinusQ1 bool

        // proportion of correction to apply to sending average activation for hebbian learning component (0=none, 1=all, .5=half, etc)
        SAvgCor float32 `default:"0.4:0.8" min:"0" max:"1"`

        // threshold of sending average activation below which learning does not occur (prevents learning when there is no input)
        SAvgThr float32 `default:"0.001" min:"0"`
}

func (ch *CHLParams) Defaults() {
        ch.On = true
        ch.Hebb = 0.001
        ch.SAvgCor = 0.4
        ch.SAvgThr = 0.001
        ch.Update()
}

func (ch *CHLParams) Update() {
        ch.Err = 1 - ch.Hebb
}

// MinusAct returns the minus-phase activation to use based on settings (ActM vs. ActQ1)
func (ch *CHLParams) MinusAct(actM, actQ1 float32) float32 {
        if ch.MinusQ1 {
                return actQ1
        }
        return actM
}

// HebbDWt computes the hebbian DWt value from sending, recv acts, savgCor, and linear Wt
func (ch *CHLParams) HebbDWt(sact, ract, savgCor, linWt float32) float32 {
        return ract * (sact*(savgCor-linWt) - (1-sact)*linWt)
}

// ErrDWt computes the error-driven DWt value from sending,
// recv acts in both phases, and linear Wt, which is used
// for soft weight bounding (always applied here, separate from hebbian
// which has its own soft weight bounding dynamic).
func (ch *CHLParams) ErrDWt(sactP, sactM, ractP, ractM, linWt float32) float32 {
        err := (ractP * sactP) - (ractM * sactM)
        if err > 0 {
                err *= (1 - linWt)
        } else {
                err *= linWt
        }
        return err
}

// DWt computes the overall dwt from hebbian and error terms
func (ch *CHLParams) DWt(hebb, err float32) float32 {
        return ch.Hebb*hebb + ch.Err*err
}

func (pt *Path) CHLDefaults() {
        pt.Learn.Norm.On = false     // off by default
        pt.Learn.Momentum.On = false // off by default
        pt.Learn.WtBal.On = false    // todo: experiment
}

// SAvgCor computes the sending average activation, corrected according to the SAvgCor
// correction factor (typically makes layer appear more sparse than it is)
func (pt *Path) SAvgCor(slay *Layer) float32 {
        savg := .5 + pt.CHL.SAvgCor*(slay.Pools[0].ActAvg.ActPAvgEff-0.5)
        savg = math32.Max(pt.CHL.SAvgThr, savg) // keep this computed value within bounds
        return 0.5 / savg
}

// DWtCHL computes the weight change (learning) for CHL
func (pt *Path) DWtCHL() {
        slay := pt.Send
        rlay := pt.Recv
        if slay.Pools[0].ActP.Avg < pt.CHL.SAvgThr { // inactive, no learn
                return
        }
        for si := range slay.Neurons {
                sn := &slay.Neurons[si]
                nc := int(pt.SConN[si])
                st := int(pt.SConIndexSt[si])
                syns := pt.Syns[st : st+nc]
                scons := pt.SConIndex[st : st+nc]
                snActM := pt.CHL.MinusAct(sn.ActM, sn.ActQ1)

                savgCor := pt.SAvgCor(slay)

                for ci := range syns {
                        sy := &syns[ci]
                        ri := scons[ci]
                        rn := &rlay.Neurons[ri]
                        rnActM := pt.CHL.MinusAct(rn.ActM, rn.ActQ1)

                        hebb := pt.CHL.HebbDWt(sn.ActP, rn.ActP, savgCor, sy.LWt)
                        err := pt.CHL.ErrDWt(sn.ActP, snActM, rn.ActP, rnActM, sy.LWt)

                        dwt := pt.CHL.DWt(hebb, err)
                        norm := float32(1)
                        if pt.Learn.Norm.On {
                                norm = pt.Learn.Norm.NormFromAbsDWt(&sy.Norm, math32.Abs(dwt))
                        }
                        if pt.Learn.Momentum.On {
                                dwt = norm * pt.Learn.Momentum.MomentFromDWt(&sy.Moment, dwt)
                        } else {
                                dwt *= norm
                        }
                        sy.DWt += pt.Learn.Lrate * dwt
                }
                // aggregate max DWtNorm over sending synapses
                if pt.Learn.Norm.On {
                        maxNorm := float32(0)
                        for ci := range syns {
                                sy := &syns[ci]
                                if sy.Norm > maxNorm {
                                        maxNorm = sy.Norm
                                }
                        }
                        for ci := range syns {
                                sy := &syns[ci]
                                sy.Norm = maxNorm
                        }
                }
        }
}

func (pt *Path) EcCa1Defaults() {
        pt.Learn.Norm.On = false     // off by default
        pt.Learn.Momentum.On = false // off by default
        pt.Learn.WtBal.On = false    // todo: experiment
}

// DWt computes the weight change (learning) -- on sending pathways
// Delta version
func (pt *Path) DWtEcCa1() {
        slay := pt.Send
        rlay := pt.Recv
        for si := range slay.Neurons {
                sn := &slay.Neurons[si]
                nc := int(pt.SConN[si])
                st := int(pt.SConIndexSt[si])
                syns := pt.Syns[st : st+nc]
                scons := pt.SConIndex[st : st+nc]

                for ci := range syns {
                        sy := &syns[ci]
                        ri := scons[ci]
                        rn := &rlay.Neurons[ri]

                        err := (sn.ActP * rn.ActP) - (sn.ActQ1 * rn.ActQ1)
                        bcm := pt.Learn.BCMdWt(sn.AvgSLrn, rn.AvgSLrn, rn.AvgL)
                        bcm *= pt.Learn.XCal.LongLrate(rn.AvgLLrn)
                        err *= pt.Learn.XCal.MLrn
                        dwt := bcm + err

                        norm := float32(1)
                        if pt.Learn.Norm.On {
                                norm = pt.Learn.Norm.NormFromAbsDWt(&sy.Norm, math32.Abs(dwt))
                        }
                        if pt.Learn.Momentum.On {
                                dwt = norm * pt.Learn.Momentum.MomentFromDWt(&sy.Moment, dwt)
                        } else {
                                dwt *= norm
                        }
                        sy.DWt += pt.Learn.Lrate * dwt
                }
                // aggregate max DWtNorm over sending synapses
                if pt.Learn.Norm.On {
                        maxNorm := float32(0)
                        for ci := range syns {
                                sy := &syns[ci]
                                if sy.Norm > maxNorm {
                                        maxNorm = sy.Norm
                                }
                        }
                        for ci := range syns {
                                sy := &syns[ci]
                                sy.Norm = maxNorm
                        }
                }
        }
}

// ConfigLoopsHip configures the hippocampal looper and should be included in ConfigLoops
// in model to make sure hip loops is configured correctly.
// see hip.go for an instance of implementation of this function.
func (net *Network) ConfigLoopsHip(ctx *Context, ls *looper.Stacks) {
        var tmpValues []float32
        ecout := net.LayerByName("ECout")
        ecin := net.LayerByName("ECin")
        ca1 := net.LayerByName("CA1")
        ca3 := net.LayerByName("CA3")
        ca1FromECin := errors.Log1(ca1.RecvPathBySendName("ECin")).(*Path)
        ca1FromCa3 := errors.Log1(ca1.RecvPathBySendName("CA3")).(*Path)
        ca3FromDg := errors.Log1(ca3.RecvPathBySendName("DG")).(*Path)

        dgPjScale := ca3FromDg.WtScale.Rel

        ls.AddEventAllModes(etime.Cycle, "HipMinusPhase:Start", 0, func() {
                ca1FromECin.WtScale.Abs = 1
                ca1FromCa3.WtScale.Abs = 0
                ca3FromDg.WtScale.Rel = 0
                net.GScaleFromAvgAct()
                net.InitGInc()
        })
        ls.AddEventAllModes(etime.Cycle, "Hip:Quarter1", 25, func() {
                ca1FromECin.WtScale.Abs = 0
                ca1FromCa3.WtScale.Abs = 1
                if ctx.Mode == etime.Test {
                        ca3FromDg.WtScale.Rel = 1 // weaker
                } else {
                        ca3FromDg.WtScale.Rel = dgPjScale
                }
                net.GScaleFromAvgAct()
                net.InitGInc()
        })
        for _, st := range ls.Stacks {
                ev := st.Loops[etime.Cycle].EventByCounter(75)
                ev.OnEvent.Prepend("HipPlusPhase:Start", func() bool {
                        ca1FromECin.WtScale.Abs = 1
                        ca1FromCa3.WtScale.Abs = 0
                        if ctx.Mode == etime.Train {
                                ecin.UnitValues(&tmpValues, "Act", 0)
                                ecout.ApplyExt1D32(tmpValues)
                        }
                        net.GScaleFromAvgAct()
                        net.InitGInc()
                        return true
                })
        }
}

// Copyright (c) 2019, The Emergent Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package leabra

import "github.com/emer/leabra/v2/fffb"

// leabra.InhibParams contains all the inhibition computation params and functions for basic Leabra
// This is included in leabra.Layer to support computation.
// This also includes other misc layer-level params such as running-average activation in the layer
// which is used for netinput rescaling and potentially for adapting inhibition over time
type InhibParams struct {

        // inhibition across the entire layer
        Layer fffb.Params `display:"inline"`

        // inhibition across sub-pools of units, for layers with 4D shape
        Pool fffb.Params `display:"inline"`

        // neuron self-inhibition parameters -- can be beneficial for producing more graded, linear response -- not typically used in cortical networks
        Self SelfInhibParams `display:"inline"`

        // running-average activation computation values -- for overall estimates of layer activation levels, used in netinput scaling
        ActAvg ActAvgParams `display:"inline"`
}

func (ip *InhibParams) Update() {
        ip.Layer.Update()
        ip.Pool.Update()
        ip.Self.Update()
        ip.ActAvg.Update()
}

func (ip *InhibParams) Defaults() {
        ip.Layer.Defaults()
        ip.Pool.Defaults()
        ip.Self.Defaults()
        ip.ActAvg.Defaults()
}

///////////////////////////////////////////////////////////////////////
//  SelfInhibParams

// SelfInhibParams defines parameters for Neuron self-inhibition -- activation of the neuron directly feeds back
// to produce a proportional additional contribution to Gi
type SelfInhibParams struct {

        // enable neuron self-inhibition
        On bool

        // strength of individual neuron self feedback inhibition -- can produce proportional activation behavior in individual units for specialized cases (e.g., scalar val or BG units), but not so good for typical hidden layers
        Gi float32 `default:"0.4"`

        // time constant in cycles, which should be milliseconds typically (roughly, how long it takes for value to change significantly -- 1.4x the half-life) for integrating unit self feedback inhibitory values -- prevents oscillations that otherwise occur -- relatively rapid 1.4 typically works, but may need to go longer if oscillations are a problem
        Tau float32 `default:"1.4"`

        // rate = 1 / tau
        Dt float32 `edit:"-" display:"-" json:"-" xml:"-"`
}

func (si *SelfInhibParams) Update() {
        si.Dt = 1 / si.Tau
}

func (si *SelfInhibParams) Defaults() {
        si.On = false
        si.Gi = 0.4
        si.Tau = 1.4
        si.Update()
}

func (si *SelfInhibParams) ShouldDisplay(field string) bool {
        switch field {
        case "Gi", "Tau":
                return si.On
        default:
                return true
        }
}

// Inhib updates the self inhibition value based on current unit activation
func (si *SelfInhibParams) Inhib(self *float32, act float32) {
        if si.On {
                *self += si.Dt * (si.Gi*act - *self)
        } else {
                *self = 0
        }
}

///////////////////////////////////////////////////////////////////////
//  ActAvgParams

// ActAvgParams represents expected average activity levels in the layer.
// Used for computing running-average computation that is then used for netinput scaling.
// Also specifies time constant for updating average
// and for the target value for adapting inhibition in inhib_adapt.
type ActAvgParams struct {

        // initial estimated average activity level in the layer (see also UseFirst option -- if that is off then it is used as a starting point for running average actual activity level, ActMAvg and ActPAvg) -- ActPAvg is used primarily for automatic netinput scaling, to balance out layers that have different activity levels -- thus it is important that init be relatively accurate -- good idea to update from recorded ActPAvg levels
        Init float32 `min:"0"`

        // if true, then the Init value is used as a constant for ActPAvgEff (the effective value used for netinput rescaling), instead of using the actual running average activation
        Fixed bool `default:"false"`

        // if true, then use the activation level computed from the external inputs to this layer (avg of targ or ext unit vars) -- this will only be applied to layers with Input or Target / Compare layer types, and falls back on the targ_init value if external inputs are not available or have a zero average -- implies fixed behavior
        UseExtAct bool `default:"false"`

        // use the first actual average value to override targ_init value -- actual value is likely to be a better estimate than our guess
        UseFirst bool `default:"true"`

        // time constant in trials for integrating time-average values at the layer level -- used for computing Pool.ActAvg.ActsMAvg, ActsPAvg
        Tau float32 `default:"100" min:"1"`

        // adjustment multiplier on the computed ActPAvg value that is used to compute ActPAvgEff, which is actually used for netinput rescaling -- if based on connectivity patterns or other factors the actual running-average value is resulting in netinputs that are too high or low, then this can be used to adjust the effective average activity value -- reducing the average activity with a factor < 1 will increase netinput scaling (stronger net inputs from layers that receive from this layer), and vice-versa for increasing (decreases net inputs)
        Adjust float32 `default:"1"`

        // rate = 1 / tau
        Dt float32 `edit:"-" display:"-" json:"-" xml:"-"`
}

func (aa *ActAvgParams) Update() {
        aa.Dt = 1 / aa.Tau
}

func (aa *ActAvgParams) Defaults() {
        aa.Init = 0.15
        aa.Fixed = false
        aa.UseExtAct = false
        aa.UseFirst = true
        aa.Tau = 100
        aa.Adjust = 1
        aa.Update()
}

func (aa *ActAvgParams) ShouldDisplay(field string) bool {
        switch field {
        case "UseFirst", "Tau", "Adjust":
                return !aa.Fixed
        default:
                return true
        }
}

// EffInit returns the initial value applied during InitWeights for the AvgPAvgEff effective layer activity
func (aa *ActAvgParams) EffInit() float32 {
        if aa.Fixed {
                return aa.Init
        }
        return aa.Adjust * aa.Init
}

// AvgFromAct updates the running-average activation given average activity level in layer
func (aa *ActAvgParams) AvgFromAct(avg *float32, act float32) {
        if act < 0.0001 {
                return
        }
        if aa.UseFirst && *avg == aa.Init {
                *avg += 0.5 * (act - *avg)
        } else {
                *avg += aa.Dt * (act - *avg)
        }
}

// EffFromAvg updates the effective value from the running-average value
func (aa *ActAvgParams) EffFromAvg(eff *float32, avg float32) {
        if aa.Fixed {
                *eff = aa.Init
        } else {
                *eff = aa.Adjust * avg
        }
}

// Copyright (c) 2019, The Emergent Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package leabra

import (
        "log"
        "math/rand"

        "cogentcore.org/core/enums"
        "cogentcore.org/core/math32"
        "cogentcore.org/lab/base/randx"
        "github.com/emer/etensor/tensor"
)

//////////////////////////////////////////////////////////////////////////////////////
//  Init methods

// InitWeights initializes the weight values in the network,
// i.e., resetting learning Also calls InitActs.
func (ly *Layer) InitWeights() {
        ly.UpdateParams()
        for _, pt := range ly.SendPaths {
                if pt.Off {
                        continue
                }
                pt.InitWeights()
        }
        for pi := range ly.Pools {
                pl := &ly.Pools[pi]
                pl.ActAvg.ActMAvg = ly.Inhib.ActAvg.Init
                pl.ActAvg.ActPAvg = ly.Inhib.ActAvg.Init
                pl.ActAvg.ActPAvgEff = ly.Inhib.ActAvg.EffInit()
        }
        ly.InitActAvg()
        ly.InitActs()
        ly.CosDiff.Init()
        ly.SetDriverOffs()
}

// InitActAvg initializes the running-average activation
// values that drive learning.
func (ly *Layer) InitActAvg() {
        for ni := range ly.Neurons {
                nrn := &ly.Neurons[ni]
                ly.Learn.InitActAvg(nrn)
        }
}

// InitActs fully initializes activation state.
// only called automatically during InitWeights.
func (ly *Layer) InitActs() {
        for ni := range ly.Neurons {
                nrn := &ly.Neurons[ni]
                ly.Act.InitActs(nrn)
        }
        for pi := range ly.Pools {
                pl := &ly.Pools[pi]
                pl.Init()
                pl.ActM.Init()
                pl.ActP.Init()
        }
        ly.NeuroMod.Init()
}

// UpdateActAvgEff updates the effective ActAvg.ActPAvgEff value used in netinput
// scaling, from the current ActAvg.ActPAvg and fixed Init values.
func (ly *Layer) UpdateActAvgEff() {
        for pi := range ly.Pools {
                pl := &ly.Pools[pi]
                ly.Inhib.ActAvg.EffFromAvg(&pl.ActAvg.ActPAvgEff, pl.ActAvg.ActPAvg)
        }
}

// InitWeightsSym initializes the weight symmetry.
// higher layers copy weights from lower layers.
func (ly *Layer) InitWtSym() {
        for _, pt := range ly.SendPaths {
                if pt.Off {
                        continue
                }
                if !(pt.WtInit.Sym) {
                        continue
                }
                // key ordering constraint on which way weights are copied
                if pt.Recv.Index < pt.Send.Index {
                        continue
                }
                rpt, has := ly.RecipToSendPath(pt)
                if !has {
                        continue
                }
                if !(rpt.WtInit.Sym) {
                        continue
                }
                pt.InitWtSym(rpt)
        }
}

// InitExt initializes external input state -- called prior to apply ext
func (ly *Layer) InitExt() {
        for ni := range ly.Neurons {
                nrn := &ly.Neurons[ni]
                nrn.Ext = 0
                nrn.Targ = 0
                nrn.SetFlag(false, NeurHasExt, NeurHasTarg, NeurHasCmpr)
        }
}

// ApplyExtFlags gets the flags that should cleared and set for updating neuron flags
// based on layer type, and whether input should be applied to Targ (else Ext)
func (ly *Layer) ApplyExtFlags() (clear, set []enums.BitFlag, toTarg bool) {
        clear = []enums.BitFlag{NeurHasExt, NeurHasTarg, NeurHasCmpr}
        toTarg = false
        if ly.Type == TargetLayer {
                set = []enums.BitFlag{NeurHasTarg}
                toTarg = true
        } else if ly.Type == CompareLayer {
                set = []enums.BitFlag{NeurHasCmpr}
                toTarg = true
        } else {
                set = []enums.BitFlag{NeurHasExt}
        }
        return
}

// ApplyExt applies external input in the form of an tensor.Float32.  If
// dimensionality of tensor matches that of layer, and is 2D or 4D, then each dimension
// is iterated separately, so any mismatch preserves dimensional structure.
// Otherwise, the flat 1D view of the tensor is used.
// If the layer is a Target or Compare layer type, then it goes in Targ
// otherwise it goes in Ext
func (ly *Layer) ApplyExt(ext tensor.Tensor) {
        switch {
        case ext.NumDims() == 2 && ly.Shape.NumDims() == 4: // special case
                ly.ApplyExt2Dto4D(ext)
        case ext.NumDims() != ly.Shape.NumDims() || !(ext.NumDims() == 2 || ext.NumDims() == 4):
                ly.ApplyExt1DTsr(ext)
        case ext.NumDims() == 2:
                ly.ApplyExt2D(ext)
        case ext.NumDims() == 4:
                ly.ApplyExt4D(ext)
        }
}

// ApplyExtVal applies given external value to given neuron
// using clearMask, setMask, and toTarg from ApplyExtFlags.
// Also saves Val in Exts for potential use by GPU.
func (ly *Layer) ApplyExtValue(lni int, val float32, clear, set []enums.BitFlag, toTarg bool) {
        nrn := &ly.Neurons[lni]
        if nrn.IsOff() {
                return
        }
        if toTarg {
                nrn.Targ = val
        } else {
                nrn.Ext = val
        }
        nrn.SetFlag(false, clear...)
        nrn.SetFlag(true, set...)
}

// ApplyExt2D applies 2D tensor external input
func (ly *Layer) ApplyExt2D(ext tensor.Tensor) {
        clear, set, toTarg := ly.ApplyExtFlags()
        ymx := min(ext.DimSize(0), ly.Shape.DimSize(0))
        xmx := min(ext.DimSize(1), ly.Shape.DimSize(1))
        for y := 0; y < ymx; y++ {
                for x := 0; x < xmx; x++ {
                        idx := []int{y, x}
                        vl := float32(ext.Float(idx))
                        i := ly.Shape.Offset(idx)
                        ly.ApplyExtValue(i, vl, clear, set, toTarg)
                }
        }
}

// ApplyExt2Dto4D applies 2D tensor external input to a 4D layer
func (ly *Layer) ApplyExt2Dto4D(ext tensor.Tensor) {
        clear, set, toTarg := ly.ApplyExtFlags()
        lNy, lNx, _, _ := tensor.Projection2DShape(&ly.Shape, false)

        ymx := min(ext.DimSize(0), lNy)
        xmx := min(ext.DimSize(1), lNx)
        for y := 0; y < ymx; y++ {
                for x := 0; x < xmx; x++ {
                        idx := []int{y, x}
                        vl := float32(ext.Float(idx))
                        ui := tensor.Projection2DIndex(&ly.Shape, false, y, x)
                        ly.ApplyExtValue(ui, vl, clear, set, toTarg)
                }
        }
}

// ApplyExt4D applies 4D tensor external input
func (ly *Layer) ApplyExt4D(ext tensor.Tensor) {
        clear, set, toTarg := ly.ApplyExtFlags()
        ypmx := min(ext.DimSize(0), ly.Shape.DimSize(0))
        xpmx := min(ext.DimSize(1), ly.Shape.DimSize(1))
        ynmx := min(ext.DimSize(2), ly.Shape.DimSize(2))
        xnmx := min(ext.DimSize(3), ly.Shape.DimSize(3))
        for yp := 0; yp < ypmx; yp++ {
                for xp := 0; xp < xpmx; xp++ {
                        for yn := 0; yn < ynmx; yn++ {
                                for xn := 0; xn < xnmx; xn++ {
                                        idx := []int{yp, xp, yn, xn}
                                        vl := float32(ext.Float(idx))
                                        i := ly.Shape.Offset(idx)
                                        ly.ApplyExtValue(i, vl, clear, set, toTarg)
                                }
                        }
                }
        }
}

// ApplyExt1DTsr applies external input using 1D flat interface into tensor.
// If the layer is a Target or Compare layer type, then it goes in Targ
// otherwise it goes in Ext
func (ly *Layer) ApplyExt1DTsr(ext tensor.Tensor) {
        clear, set, toTarg := ly.ApplyExtFlags()
        mx := min(ext.Len(), len(ly.Neurons))
        for i := 0; i < mx; i++ {
                vl := float32(ext.Float1D(i))
                ly.ApplyExtValue(i, vl, clear, set, toTarg)
        }
}

// ApplyExt1D applies external input in the form of a flat 1-dimensional slice of floats
// If the layer is a Target or Compare layer type, then it goes in Targ
// otherwise it goes in Ext
func (ly *Layer) ApplyExt1D(ext []float64) {
        clear, set, toTarg := ly.ApplyExtFlags()
        mx := min(len(ext), len(ly.Neurons))
        for i := 0; i < mx; i++ {
                vl := float32(ext[i])
                ly.ApplyExtValue(i, vl, clear, set, toTarg)
        }
}

// ApplyExt1D32 applies external input in the form of
//
//        a flat 1-dimensional slice of float32s.
//
// If the layer is a Target or Compare layer type, then it goes in Targ
// otherwise it goes in Ext
func (ly *Layer) ApplyExt1D32(ext []float32) {
        clear, set, toTarg := ly.ApplyExtFlags()
        mx := min(len(ext), len(ly.Neurons))
        for i := 0; i < mx; i++ {
                vl := ext[i]
                ly.ApplyExtValue(i, vl, clear, set, toTarg)
        }
}

// UpdateExtFlags updates the neuron flags for external input based on current
// layer Type field -- call this if the Type has changed since the last
// ApplyExt* method call.
func (ly *Layer) UpdateExtFlags() {
        clear, set, _ := ly.ApplyExtFlags()
        for i := range ly.Neurons {
                nrn := &ly.Neurons[i]
                if nrn.IsOff() {
                        continue
                }
                nrn.SetFlag(false, clear...)
                nrn.SetFlag(true, set...)
        }
}

// ActAvgFromAct updates the running average ActMAvg, ActPAvg, and ActPAvgEff
// values from the current pool-level averages.
// The ActPAvgEff value is used for updating the conductance scaling parameters,
// if these are not set to Fixed, so calling this will change the scaling of
// pathways in the network!
func (ly *Layer) ActAvgFromAct() {
        for pi := range ly.Pools {
                pl := &ly.Pools[pi]
                ly.Inhib.ActAvg.AvgFromAct(&pl.ActAvg.ActMAvg, pl.ActM.Avg)
                ly.Inhib.ActAvg.AvgFromAct(&pl.ActAvg.ActPAvg, pl.ActP.Avg)
                ly.Inhib.ActAvg.EffFromAvg(&pl.ActAvg.ActPAvgEff, pl.ActAvg.ActPAvg)
        }
}

// ActQ0FromActP updates the neuron ActQ0 value from prior ActP value
func (ly *Layer) ActQ0FromActP() {
        for ni := range ly.Neurons {
                nrn := &ly.Neurons[ni]
                if nrn.IsOff() {
                        continue
                }
                nrn.ActQ0 = nrn.ActP
        }
}

// AlphaCycInit handles all initialization at start of new input pattern.
// Should already have presented the external input to the network at this point.
// If updtActAvg is true, this includes updating the running-average
// activations for each layer / pool, and the AvgL running average used
// in BCM Hebbian learning.
// The input scaling is updated  based on the layer-level running average acts,
// and this can then change the behavior of the network,
// so if you want 100% repeatable testing results, set this to false to
// keep the existing scaling factors (e.g., can pass a train bool to
// only update during training).  This flag also affects the AvgL learning
// threshold
func (ly *Layer) AlphaCycInit(updtActAvg bool) {
        ly.ActQ0FromActP()
        if updtActAvg {
                ly.AvgLFromAvgM()
                ly.ActAvgFromAct()
        }
        ly.GScaleFromAvgAct() // need to do this always, in case hasn't been done at all yet
        if ly.Act.Noise.Type != NoNoise && ly.Act.Noise.Fixed && ly.Act.Noise.Dist != randx.Mean {
                ly.GenNoise()
        }
        ly.DecayState(ly.Act.Init.Decay)
        ly.InitGInc()
        if ly.Act.Clamp.Hard && ly.Type == InputLayer {
                ly.HardClamp()
        }
}

// AvgLFromAvgM updates AvgL long-term running average activation that drives BCM Hebbian learning
func (ly *Layer) AvgLFromAvgM() {
        for ni := range ly.Neurons {
                nrn := &ly.Neurons[ni]
                if nrn.IsOff() {
                        continue
                }
                ly.Learn.AvgLFromAvgM(nrn)
                if ly.Learn.AvgL.ErrMod {
                        nrn.AvgLLrn *= ly.CosDiff.ModAvgLLrn
                }
        }
}

// GScaleFromAvgAct computes the scaling factor for synaptic input conductances G,
// based on sending layer average activation.
// This attempts to automatically adjust for overall differences in raw activity
// coming into the units to achieve a general target of around .5 to 1
// for the integrated Ge value.
func (ly *Layer) GScaleFromAvgAct() {
        totGeRel := float32(0)
        totGiRel := float32(0)
        for _, pt := range ly.RecvPaths {
                if pt.Off {
                        continue
                }
                slay := pt.Send
                slpl := &slay.Pools[0]
                savg := slpl.ActAvg.ActPAvgEff
                snu := len(slay.Neurons)
                ncon := pt.RConNAvgMax.Avg
                pt.GScale = pt.WtScale.FullScale(savg, float32(snu), ncon)
                // reverting this change: if you want to eliminate a path, set the Off flag
                // if you want to negate it but keep the relative factor in the denominator
                // then set the scale to 0.
                // if pj.GScale == 0 {
                //         continue
                // }
                if pt.Type == InhibPath {
                        totGiRel += pt.WtScale.Rel
                } else {
                        totGeRel += pt.WtScale.Rel
                }
        }

        for _, pt := range ly.RecvPaths {
                if pt.Off {
                        continue
                }
                if pt.Type == InhibPath {
                        if totGiRel > 0 {
                                pt.GScale /= totGiRel
                        }
                } else {
                        if totGeRel > 0 {
                                pt.GScale /= totGeRel
                        }
                }
        }
}

// GenNoise generates random noise for all neurons
func (ly *Layer) GenNoise() {
        for ni := range ly.Neurons {
                nrn := &ly.Neurons[ni]
                nrn.Noise = float32(ly.Act.Noise.Gen())
        }
}

// DecayState decays activation state by given proportion (default is on ly.Act.Init.Decay).
// This does *not* call InitGInc -- must call that separately at start of AlphaCyc
func (ly *Layer) DecayState(decay float32) {
        for ni := range ly.Neurons {
                nrn := &ly.Neurons[ni]
                if nrn.IsOff() {
                        continue
                }
                ly.Act.DecayState(nrn, decay)
        }
        for pi := range ly.Pools { // decaying average act is essential for inhib
                pl := &ly.Pools[pi]
                pl.Inhib.Decay(decay)
        }
}

// DecayStatePool decays activation state by given proportion
// in given pool index (sub pools start at 1).
func (ly *Layer) DecayStatePool(pool int, decay float32) {
        pl := &ly.Pools[pool]
        for ni := pl.StIndex; ni < pl.EdIndex; ni++ {
                nrn := &ly.Neurons[ni]
                if nrn.IsOff() {
                        continue
                }
                ly.Act.DecayState(nrn, decay)
        }
        pl.Inhib.Decay(decay)
}

// HardClamp hard-clamps the activations in the layer.
// called during AlphaCycInit for hard-clamped Input layers.
func (ly *Layer) HardClamp() {
        for ni := range ly.Neurons {
                nrn := &ly.Neurons[ni]
                if nrn.IsOff() {
                        continue
                }
                ly.Act.HardClamp(nrn)
        }
}

//////////////////////////////////////////////////////////////////////////////////////
//  Cycle

// InitGinc initializes the Ge excitatory and Gi inhibitory conductance accumulation states
// including ActSent and G*Raw values.
// called at start of trial always, and can be called optionally
// when delta-based Ge computation needs to be updated (e.g., weights
// might have changed strength)
func (ly *Layer) InitGInc() {
        for ni := range ly.Neurons {
                nrn := &ly.Neurons[ni]
                if nrn.IsOff() {
                        continue
                }
                ly.Act.InitGInc(nrn)
        }
        for _, pt := range ly.RecvPaths {
                if pt.Off {
                        continue
                }
                pt.InitGInc()
        }
}

// SendGDelta sends change in activation since last sent, to increment recv
// synaptic conductances G, if above thresholds
func (ly *Layer) SendGDelta(ctx *Context) {
        for ni := range ly.Neurons {
                nrn := &ly.Neurons[ni]
                if nrn.IsOff() {
                        continue
                }
                if nrn.Act > ly.Act.OptThresh.Send {
                        delta := nrn.Act - nrn.ActSent
                        if math32.Abs(delta) > ly.Act.OptThresh.Delta {
                                for _, sp := range ly.SendPaths {
                                        if sp.Off {
                                                continue
                                        }
                                        sp.SendGDelta(ni, delta)
                                }
                                nrn.ActSent = nrn.Act
                        }
                } else if nrn.ActSent > ly.Act.OptThresh.Send {
                        delta := -nrn.ActSent // un-send the last above-threshold activation to get back to 0
                        for _, sp := range ly.SendPaths {
                                if sp.Off {
                                        continue
                                }
                                sp.SendGDelta(ni, delta)
                        }
                        nrn.ActSent = 0
                }
        }
}

// GFromInc integrates new synaptic conductances from increments sent during last SendGDelta.
func (ly *Layer) GFromInc(ctx *Context) {
        ly.RecvGInc(ctx)
        switch ly.Type {
        case CTLayer:
                ly.CTGFromInc(ctx)
        case PulvinarLayer:
                if ly.Pulvinar.DriversOff || !ly.Pulvinar.BurstQtr.HasFlag(ctx.Quarter) {
                        ly.GFromIncNeur(ctx)
                } else {
                        ly.SetDriverActs()
                }
        case GPiThalLayer:
                ly.GPiGFromInc(ctx)
        case PFCDeepLayer:
                ly.MaintGInc(ctx)
        default:
                ly.GFromIncNeur(ctx)
        }
}

// RecvGInc calls RecvGInc on receiving pathways to collect Neuron-level G*Inc values.
// This is called by GFromInc overall method, but separated out for cases that need to
// do something different.
func (ly *Layer) RecvGInc(ctx *Context) {
        for _, pt := range ly.RecvPaths {
                if pt.Off {
                        continue
                }
                pt.RecvGInc()
        }
}

// GFromIncNeur is the neuron-level code for GFromInc that integrates overall Ge, Gi values
// from their G*Raw accumulators.
func (ly *Layer) GFromIncNeur(ctx *Context) {
        for ni := range ly.Neurons {
                nrn := &ly.Neurons[ni]
                if nrn.IsOff() {
                        continue
                }
                // note: each step broken out here so other variants can add extra terms to Raw
                ly.Act.GeFromRaw(nrn, nrn.GeRaw)
                ly.Act.GiFromRaw(nrn, nrn.GiRaw)
        }
}

// AvgMaxGe computes the average and max Ge stats, used in inhibition
func (ly *Layer) AvgMaxGe(ctx *Context) {
        for pi := range ly.Pools {
                pl := &ly.Pools[pi]
                pl.Inhib.Ge.Init()
                for ni := pl.StIndex; ni < pl.EdIndex; ni++ {
                        nrn := &ly.Neurons[ni]
                        if nrn.IsOff() {
                                continue
                        }
                        pl.Inhib.Ge.UpdateValue(nrn.Ge, int32(ni))
                }
                pl.Inhib.Ge.CalcAvg()
        }
}

// InhibFromGeAct computes inhibition Gi from Ge and Act averages within relevant Pools
func (ly *Layer) InhibFromGeAct(ctx *Context) {
        lpl := &ly.Pools[0]
        ly.Inhib.Layer.Inhib(&lpl.Inhib)
        ly.PoolInhibFromGeAct(ctx)
        ly.InhibFromPool(ctx)
        if ly.Type == MatrixLayer {
                ly.MatrixOutAChInhib(ctx)
        }
}

// PoolInhibFromGeAct computes inhibition Gi from Ge and Act averages within relevant Pools
func (ly *Layer) PoolInhibFromGeAct(ctx *Context) {
        np := len(ly.Pools)
        if np == 1 {
                return
        }
        lpl := &ly.Pools[0]
        lyInhib := ly.Inhib.Layer.On
        for pi := 1; pi < np; pi++ {
                pl := &ly.Pools[pi]
                ly.Inhib.Pool.Inhib(&pl.Inhib)
                if lyInhib {
                        pl.Inhib.LayGi = lpl.Inhib.Gi
                        pl.Inhib.Gi = math32.Max(pl.Inhib.Gi, lpl.Inhib.Gi) // pool is max of layer
                } else {
                        lpl.Inhib.Gi = math32.Max(pl.Inhib.Gi, lpl.Inhib.Gi) // update layer from pool
                }
        }
        if !lyInhib {
                lpl.Inhib.GiOrig = lpl.Inhib.Gi // effective GiOrig
        }
}

// InhibFromPool computes inhibition Gi from Pool-level aggregated inhibition, including self and syn
func (ly *Layer) InhibFromPool(ctx *Context) {
        for ni := range ly.Neurons {
                nrn := &ly.Neurons[ni]
                if nrn.IsOff() {
                        continue
                }
                pl := &ly.Pools[nrn.SubPool]
                ly.Inhib.Self.Inhib(&nrn.GiSelf, nrn.Act)
                nrn.Gi = pl.Inhib.Gi + nrn.GiSelf + nrn.GiSyn
        }
}

// ActFromG computes rate-code activation from Ge, Gi, Gl conductances
// and updates learning running-average activations from that Act
func (ly *Layer) ActFromG(ctx *Context) {
        switch ly.Type {
        case RWDaLayer:
                ly.ActFromGRWDa(ctx)
                return
        case RWPredLayer:
                ly.ActFromGRWPred(ctx)
                return
        case TDPredLayer:
                ly.ActFromGTDPred(ctx)
                return
        case TDIntegLayer:
                ly.ActFromGTDInteg(ctx)
                return
        case TDDaLayer:
                ly.ActFromGTDDa(ctx)
                return
        case CINLayer:
                ly.ActFromGCIN(ctx)
                return
        }
        for ni := range ly.Neurons {
                nrn := &ly.Neurons[ni]
                if nrn.IsOff() {
                        continue
                }
                ly.Act.VmFromG(nrn)
                ly.Act.ActFromG(nrn)
                ly.Learn.AvgsFromAct(nrn)
        }
        switch ly.Type {
        case MatrixLayer:
                ly.DaAChFromLay(ctx)
        case PFCDeepLayer:
                ly.PFCDeepGating(ctx)
        }
}

// AvgMaxAct computes the average and max Act stats, used in inhibition
func (ly *Layer) AvgMaxAct(ctx *Context) {
        for pi := range ly.Pools {
                pl := &ly.Pools[pi]
                pl.Inhib.Act.Init()
                for ni := pl.StIndex; ni < pl.EdIndex; ni++ {
                        nrn := &ly.Neurons[ni]
                        if nrn.IsOff() {
                                continue
                        }
                        pl.Inhib.Act.UpdateValue(nrn.Act, int32(ni))
                }
                pl.Inhib.Act.CalcAvg()
        }
}

// CyclePost is called at end of Cycle, for misc updates after new Act
// value has been computed.
// SuperLayer computes Burst activity.
// GateLayer (GPiThal) computes gating, sends to other layers.
// DA, ACh neuromodulation is sent.
func (ly *Layer) CyclePost(ctx *Context) {
        switch ly.Type {
        case SuperLayer:
                ly.BurstFromAct(ctx)
        case CTLayer:
                ly.BurstAsAct(ctx)
        case GPiThalLayer:
                ly.GPiGateSend(ctx)
        case ClampDaLayer, RWDaLayer, TDDaLayer:
                ly.SendDaFromAct(ctx)
        case CINLayer:
                ly.SendAChFromAct(ctx)
        }
}

//////////////////////////////////////////////////////////////////////////////////////
//  Quarter

// QuarterFinal does updating after end of quarter.
// Calls MinusPhase and PlusPhase for quarter = 2, 3.
func (ly *Layer) QuarterFinal(ctx *Context) {
        switch ctx.Quarter {
        case 2:
                ly.MinusPhase(ctx)
        case 3:
                ly.PlusPhase(ctx)
        default:
                ly.SaveQuarterState(ctx)
        }
        switch ly.Type {
        case SuperLayer:
                ly.BurstPrv(ctx)
                ly.SendCtxtGe(ctx)
        case CTLayer:
                ly.SendCtxtGe(ctx)
        case PFCDeepLayer:
                ly.UpdateGateCnt(ctx)
                ly.DeepMaint(ctx)
        }
        if ctx.Quarter == 1 {
                ly.Quarter2DWt()
        }
}

// SaveQuarterState saves Q1, Q2 quarter states.
func (ly *Layer) SaveQuarterState(ctx *Context) {
        for ni := range ly.Neurons {
                nrn := &ly.Neurons[ni]
                if nrn.IsOff() {
                        continue
                }
                switch ctx.Quarter {
                case 0:
                        nrn.ActQ1 = nrn.Act
                case 1:
                        nrn.ActQ2 = nrn.Act
                }
        }
}

// MinusPhase is called at the end of the minus phase (quarter 3), to record state.
func (ly *Layer) MinusPhase(ctx *Context) {
        for pi := range ly.Pools {
                pl := &ly.Pools[pi]
                pl.ActM = pl.Inhib.Act
        }
        for ni := range ly.Neurons {
                nrn := &ly.Neurons[ni]
                if nrn.IsOff() {
                        continue
                }
                nrn.ActM = nrn.Act
                if nrn.HasFlag(NeurHasTarg) { // will be clamped in plus phase
                        nrn.Ext = nrn.Targ
                        nrn.SetFlag(true, NeurHasExt)
                }
        }
}

// PlusPhase is called at the end of the plus phase (quarter 4), to record state.
func (ly *Layer) PlusPhase(ctx *Context) {
        for pi := range ly.Pools {
                pl := &ly.Pools[pi]
                pl.ActP = pl.Inhib.Act
        }
        for ni := range ly.Neurons {
                nrn := &ly.Neurons[ni]
                if nrn.IsOff() {
                        continue
                }
                nrn.ActP = nrn.Act
                nrn.ActDif = nrn.ActP - nrn.ActM
                nrn.ActAvg += ly.Act.Dt.AvgDt * (nrn.Act - nrn.ActAvg)
        }
        ly.CosDiffFromActs()
}

// CosDiffFromActs computes the cosine difference in activation state between minus and plus phases.
// this is also used for modulating the amount of BCM hebbian learning
func (ly *Layer) CosDiffFromActs() {
        lpl := &ly.Pools[0]
        avgM := lpl.ActM.Avg
        avgP := lpl.ActP.Avg
        cosv := float32(0)
        ssm := float32(0)
        ssp := float32(0)
        for ni := range ly.Neurons {
                nrn := &ly.Neurons[ni]
                if nrn.IsOff() {
                        continue
                }
                ap := nrn.ActP - avgP // zero mean
                am := nrn.ActM - avgM
                cosv += ap * am
                ssm += am * am
                ssp += ap * ap
        }

        dist := math32.Sqrt(ssm * ssp)
        if dist != 0 {
                cosv /= dist
        }
        ly.CosDiff.Cos = cosv

        ly.Learn.CosDiff.AvgVarFromCos(&ly.CosDiff.Avg, &ly.CosDiff.Var, ly.CosDiff.Cos)

        if ly.IsTarget() {
                ly.CosDiff.AvgLrn = 0 // no BCM for non-hidden layers
                ly.CosDiff.ModAvgLLrn = 0
        } else {
                ly.CosDiff.AvgLrn = 1 - ly.CosDiff.Avg
                ly.CosDiff.ModAvgLLrn = ly.Learn.AvgL.ErrModFromLayErr(ly.CosDiff.AvgLrn)
        }
}

// IsTarget returns true if this layer is a Target layer.
// By default, returns true for layers of Type == TargetLayer
// Other Target layers include the PulvinarLayer in deep predictive learning.
// This is used for turning off BCM hebbian learning,
// in CosDiffFromActs to set the CosDiff.ModAvgLLrn value
// for error-modulated level of hebbian learning.
// It is also used in WtBal to not apply it to target layers.
// In both cases, Target layers are purely error-driven.
func (ly *Layer) IsTarget() bool {
        return ly.Type == TargetLayer || ly.Type == PulvinarLayer
}

//////////////////////////////////////////////////////////////////////////////////////
//  Learning

// DWt computes the weight change (learning) -- calls DWt method on sending pathways
func (ly *Layer) DWt() {
        for _, pt := range ly.SendPaths {
                if pt.Off {
                        continue
                }
                pt.DWt()
        }
}

// Quarter2DWt computes the weight change (learning), for layers that learn in Quarter 2.
func (ly *Layer) Quarter2DWt() {
        for _, pt := range ly.SendPaths {
                if pt.Off {
                        continue
                }
                rlay := pt.Recv
                if rlay.DoQuarter2DWt() {
                        pt.DWt()
                }
        }
}

func (ly *Layer) DoQuarter2DWt() bool {
        switch ly.Type {
        case MatrixLayer:
                return ly.Matrix.LearnQtr.HasFlag(Q2)
        case PFCDeepLayer:
                return ly.PFCGate.GateQtr.HasFlag(Q2)
        }
        return false
}

// WtFromDWt updates the weights from delta-weight changes -- on the sending pathways
func (ly *Layer) WtFromDWt() {
        for _, pt := range ly.SendPaths {
                if pt.Off {
                        continue
                }
                pt.WtFromDWt()
        }
}

// WtBalFromWt computes the Weight Balance factors based on average recv weights
func (ly *Layer) WtBalFromWt() {
        for _, pt := range ly.RecvPaths {
                if pt.Off {
                        continue
                }
                pt.WtBalFromWt()
        }
}

// LrateMult sets the new Lrate parameter for Paths to LrateInit * mult.
// Useful for implementing learning rate schedules.
func (ly *Layer) LrateMult(mult float32) {
        for _, pt := range ly.RecvPaths {
                // if p.Off { // keep all sync'd
                //         continue
                // }
                pt.LrateMult(mult)
        }
}

//////////////////////////////////////////////////////////////////////////////////////
//  Threading / Reports

// CostEst returns the estimated computational cost associated with this layer,
// separated by neuron-level and synapse-level, in arbitrary units where
// cost per synapse is 1.  Neuron-level computation is more expensive but
// there are typically many fewer neurons, so in larger networks, synaptic
// costs tend to dominate.  Neuron cost is estimated from TimerReport output
// for large networks.
func (ly *Layer) CostEst() (neur, syn, tot int) {
        perNeur := 300 // cost per neuron, relative to synapse which is 1
        neur = len(ly.Neurons) * perNeur
        syn = 0
        for _, pt := range ly.SendPaths {
                ns := len(pt.Syns)
                syn += ns
        }
        tot = neur + syn
        return
}

//////////////////////////////////////////////////////////////////////////////////////
//  Stats

// note: use float64 for stats as that is best for logging

// MSE returns the sum-squared-error and mean-squared-error
// over the layer, in terms of ActP - ActM (valid even on non-target layers FWIW).
// Uses the given tolerance per-unit to count an error at all
// (e.g., .5 = activity just has to be on the right side of .5).
func (ly *Layer) MSE(tol float32) (sse, mse float64) {
        nn := len(ly.Neurons)
        if nn == 0 {
                return 0, 0
        }
        sse = 0.0
        for ni := range ly.Neurons {
                nrn := &ly.Neurons[ni]
                if nrn.IsOff() {
                        continue
                }
                var d float32
                if ly.Type == CompareLayer {
                        d = nrn.Targ - nrn.ActM
                } else {
                        d = nrn.ActP - nrn.ActM
                }
                if math32.Abs(d) < tol {
                        continue
                }
                sse += float64(d * d)
        }
        return sse, sse / float64(nn)
}

// SSE returns the sum-squared-error over the layer, in terms of ActP - ActM
// (valid even on non-target layers FWIW).
// Uses the given tolerance per-unit to count an error at all
// (e.g., .5 = activity just has to be on the right side of .5).
// Use this in Python which only allows single return values.
func (ly *Layer) SSE(tol float32) float64 {
        sse, _ := ly.MSE(tol)
        return sse
}

//////////////////////////////////////////////////////////////////////////////////////
//  Lesion

// UnLesionNeurons unlesions (clears the Off flag) for all neurons in the layer
func (ly *Layer) UnLesionNeurons() {
        for ni := range ly.Neurons {
                nrn := &ly.Neurons[ni]
                nrn.SetFlag(false, NeurOff)
        }
}

// LesionNeurons lesions (sets the Off flag) for given proportion (0-1) of neurons in layer
// returns number of neurons lesioned.  Emits error if prop > 1 as indication that percent
// might have been passed
func (ly *Layer) LesionNeurons(prop float32) int {
        ly.UnLesionNeurons()
        if prop > 1 {
                log.Printf("LesionNeurons got a proportion > 1 -- must be 0-1 as *proportion* (not percent) of neurons to lesion: %v\n", prop)
                return 0
        }
        nn := len(ly.Neurons)
        if nn == 0 {
                return 0
        }
        p := rand.Perm(nn)
        nl := int(prop * float32(nn))
        for i := 0; i < nl; i++ {
                nrn := &ly.Neurons[p[i]]
                nrn.SetFlag(true, NeurOff)
        }
        return nl
}

//////////////////////////////////////////////////////////////////////////////////////
//  Layer props for gui

// var LayerProps = tree.Props{
// "ToolBar": tree.PropSlice{
//         {"Defaults", tree.Props{
//                 "icon": "reset",
//                 "desc": "return all parameters to their intial default values",
//         }},
//         {"InitWeights", tree.Props{
//                 "icon": "update",
//                 "desc": "initialize the layer's weight values according to path parameters, for all *sending* pathways out of this layer",
//         }},
//         {"InitActs", tree.Props{
//                 "icon": "update",
//                 "desc": "initialize the layer's activation values",
//         }},
//         {"sep-act", tree.BlankProp{}},
//         {"LesionNeurons", tree.Props{
//                 "icon": "close",
//                 "desc": "Lesion (set the Off flag) for given proportion of neurons in the layer (number must be 0 -- 1, NOT percent!)",
//                 "Args": tree.PropSlice{
//                         {"Proportion", tree.Props{
//                                 "desc": "proportion (0 -- 1) of neurons to lesion",
//                         }},
//                 },
//         }},
//         {"UnLesionNeurons", tree.Props{
//                 "icon": "reset",
//                 "desc": "Un-Lesion (reset the Off flag) for all neurons in the layer",
//         }},
// },
// }

// Copyright (c) 2019, The Emergent Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package leabra

import (
        "encoding/json"
        "fmt"
        "io"
        "log"
        "math"
        "strconv"
        "strings"

        "cogentcore.org/core/base/errors"
        "cogentcore.org/core/base/num"
        "cogentcore.org/core/math32"
        "github.com/emer/emergent/v2/emer"
        "github.com/emer/emergent/v2/weights"
        "github.com/emer/etensor/tensor"
)

// Layer implements the Leabra algorithm at the layer level,
// managing neurons and pathways.
type Layer struct {
        emer.LayerBase

        // our parent network, in case we need to use it to
        // find other layers etc; set when added by network.
        Network *Network `copier:"-" json:"-" xml:"-" display:"-"`

        // type of layer.
        Type LayerTypes

        // list of receiving pathways into this layer from other layers.
        RecvPaths []*Path

        // list of sending pathways from this layer to other layers.
        SendPaths []*Path

        // Activation parameters and methods for computing activations.
        Act ActParams `display:"add-fields"`

        // Inhibition parameters and methods for computing layer-level inhibition.
        Inhib InhibParams `display:"add-fields"`

        // Learning parameters and methods that operate at the neuron level.
        Learn LearnNeurParams `display:"add-fields"`

        // Burst has parameters for computing Burst from act, in Superficial layers
        // (but also needed in Deep layers for deep self connections).
        Burst BurstParams `display:"inline"`

        // Pulvinar has parameters for computing Pulvinar plus-phase (outcome)
        // activations based on Burst activation from corresponding driver neuron.
        Pulvinar PulvinarParams `display:"inline"`

        // Drivers are names of SuperLayer(s) that sends 5IB Burst driver
        // inputs to this layer.
        Drivers Drivers

        // RW are Rescorla-Wagner RL learning parameters.
        RW RWParams `display:"inline"`

        // TD are Temporal Differences RL learning parameters.
        TD TDParams `display:"inline"`

        // Matrix BG gating parameters
        Matrix MatrixParams `display:"inline"`

        // PBWM has general PBWM parameters, including the shape
        // of overall Maint + Out gating system that this layer is part of.
        PBWM PBWMParams `display:"inline"`

        // GPiGate are gating parameters determining threshold for gating etc.
        GPiGate GPiGateParams `display:"inline"`

        // CIN cholinergic interneuron parameters.
        CIN CINParams `display:"inline"`

        // PFC Gating parameters
        PFCGate PFCGateParams `display:"inline"`

        // PFC Maintenance parameters
        PFCMaint PFCMaintParams `display:"inline"`

        // PFCDyns dynamic behavior parameters -- provides deterministic control over PFC maintenance dynamics -- the rows of PFC units (along Y axis) behave according to corresponding index of Dyns (inner loop is Super Y axis, outer is Dyn types) -- ensure Y dim has even multiple of len(Dyns)
        PFCDyns PFCDyns

        // slice of neurons for this layer, as a flat list of len = Shape.Len().
        // Must iterate over index and use pointer to modify values.
        Neurons []Neuron

        // inhibition and other pooled, aggregate state variables.
        // flat list has at least of 1 for layer, and one for each sub-pool
        // if shape supports that (4D).
        // Must iterate over index and use pointer to modify values.
        Pools []Pool

        // cosine difference between ActM, ActP stats.
        CosDiff CosDiffStats

        // NeuroMod is the neuromodulatory neurotransmitter state for this layer.
        NeuroMod NeuroMod `read-only:"+" display:"inline"`

        // SendTo is a list of layers that this layer sends special signals to,
        // which could be dopamine, gating signals, depending on the layer type.
        SendTo LayerNames
}

// emer.Layer interface methods

func (ly *Layer) StyleObject() any           { return ly }
func (ly *Layer) TypeName() string           { return ly.Type.String() }
func (ly *Layer) TypeNumber() int            { return int(ly.Type) }
func (ly *Layer) NumRecvPaths() int          { return len(ly.RecvPaths) }
func (ly *Layer) RecvPath(idx int) emer.Path { return ly.RecvPaths[idx] }
func (ly *Layer) NumSendPaths() int          { return len(ly.SendPaths) }
func (ly *Layer) SendPath(idx int) emer.Path { return ly.SendPaths[idx] }

func (ly *Layer) Defaults() {
        ly.Act.Defaults()
        ly.Inhib.Defaults()
        ly.Learn.Defaults()
        ly.Burst.Defaults()
        ly.Pulvinar.Defaults()
        ly.RW.Defaults()
        ly.TD.Defaults()
        ly.Matrix.Defaults()
        ly.PBWM.Defaults()
        ly.GPiGate.Defaults()
        ly.CIN.Defaults()
        ly.PFCGate.Defaults()
        ly.PFCMaint.Defaults()
        ly.Inhib.Layer.On = true
        for _, pt := range ly.RecvPaths {
                pt.Defaults()
        }
        ly.DefaultsForType()
}

// DefaultsForType sets the default parameter values for a given layer type.
func (ly *Layer) DefaultsForType() {
        switch ly.Type {
        case ClampDaLayer:
                ly.ClampDaDefaults()
        case MatrixLayer:
                ly.MatrixDefaults()
        case GPiThalLayer:
                ly.GPiThalDefaults()
        case CINLayer:
        case PFCLayer:
        case PFCDeepLayer:
                ly.PFCDeepDefaults()
        }
}

// UpdateParams updates all params given any changes that might have been made to individual values
// including those in the receiving pathways of this layer
func (ly *Layer) UpdateParams() {
        ly.Act.Update()
        ly.Inhib.Update()
        ly.Learn.Update()
        ly.Burst.Update()
        ly.Pulvinar.Update()
        ly.RW.Update()
        ly.TD.Update()
        ly.Matrix.Update()
        ly.PBWM.Update()
        ly.GPiGate.Update()
        ly.CIN.Update()
        ly.PFCGate.Update()
        ly.PFCMaint.Update()
        for _, pt := range ly.RecvPaths {
                pt.UpdateParams()
        }
}

func (ly *Layer) ShouldDisplay(field string) bool {
        isPBWM := ly.Type == MatrixLayer || ly.Type == GPiThalLayer || ly.Type == CINLayer || ly.Type == PFCLayer || ly.Type == PFCDeepLayer
        switch field {
        case "Burst":
                return ly.Type == SuperLayer || ly.Type == CTLayer
        case "Pulvinar", "Drivers":
                return ly.Type == PulvinarLayer
        case "RW":
                return ly.Type == RWPredLayer || ly.Type == RWDaLayer
        case "TD":
                return ly.Type == TDPredLayer || ly.Type == TDIntegLayer || ly.Type == TDDaLayer
        case "PBWM":
                return isPBWM
        case "SendTo":
                return ly.Type == GPiThalLayer || ly.Type == ClampDaLayer || ly.Type == RWDaLayer || ly.Type == TDDaLayer || ly.Type == CINLayer
        case "Matrix":
                return ly.Type == MatrixLayer
        case "GPiGate":
                return ly.Type == GPiThalLayer
        case "CIN":
                return ly.Type == CINLayer
        case "PFCGate", "PFCMaint":
                return ly.Type == PFCLayer || ly.Type == PFCDeepLayer
        case "PFCDyns":
                return ly.Type == PFCDeepLayer
        default:
                return true
        }
        return true
}

// JsonToParams reformates json output to suitable params display output
func JsonToParams(b []byte) string {
        br := strings.Replace(string(b), `"`, ``, -1)
        br = strings.Replace(br, ",\n", "", -1)
        br = strings.Replace(br, "{\n", "{", -1)
        br = strings.Replace(br, "} ", "}\n  ", -1)
        br = strings.Replace(br, "\n }", " }", -1)
        br = strings.Replace(br, "\n  }\n", " }", -1)
        return br[1:] + "\n"
}

// AllParams returns a listing of all parameters in the Layer
func (ly *Layer) AllParams() string {
        str := "/////////////////////////////////////////////////\nLayer: " + ly.Name + "\n"
        b, _ := json.MarshalIndent(&ly.Act, "", " ")
        str += "Act: {\n " + JsonToParams(b)
        b, _ = json.MarshalIndent(&ly.Inhib, "", " ")
        str += "Inhib: {\n " + JsonToParams(b)
        b, _ = json.MarshalIndent(&ly.Learn, "", " ")
        str += "Learn: {\n " + JsonToParams(b)
        for _, pt := range ly.RecvPaths {
                pstr := pt.AllParams()
                str += pstr
        }
        return str
}

// RecipToSendPath finds the reciprocal pathway relative to the given sending pathway
// found within the SendPaths of this layer.  This is then a recv path within this layer:
//
//        S=A -> R=B recip: R=A <- S=B -- ly = A -- we are the sender of srj and recv of rpj.
//
// returns false if not found.
func (ly *Layer) RecipToSendPath(spj *Path) (*Path, bool) {
        for _, rpj := range ly.RecvPaths {
                if rpj.Send == spj.Recv {
                        return rpj, true
                }
        }
        return nil, false
}

// UnitVarNames returns a list of variable names available on the units in this layer
func (ly *Layer) UnitVarNames() []string {
        return NeuronVars
}

// UnitVarProps returns properties for variables
func (ly *Layer) UnitVarProps() map[string]string {
        return NeuronVarProps
}

// UnitVarIndex returns the index of given variable within the Neuron,
// according to *this layer's* UnitVarNames() list (using a map to lookup index),
// or -1 and error message if not found.
func (ly *Layer) UnitVarIndex(varNm string) (int, error) {
        return NeuronVarIndexByName(varNm)
}

// UnitVarNum returns the number of Neuron-level variables
// for this layer.  This is needed for extending indexes in derived types.
func (ly *Layer) UnitVarNum() int {
        return len(NeuronVars)
}

// UnitValue1D returns value of given variable index on given unit,
// using 1-dimensional index. returns NaN on invalid index.
// This is the core unit var access method used by other methods,
// so it is the only one that needs to be updated for derived layer types.
func (ly *Layer) UnitValue1D(varIndex int, idx int, di int) float32 {
        if idx < 0 || idx >= len(ly.Neurons) {
                return math32.NaN()
        }
        if varIndex < 0 || varIndex >= ly.UnitVarNum() {
                return math32.NaN()
        }
        nrn := &ly.Neurons[idx]
        da := NeuronVarsMap["DA"]
        if varIndex >= da {
                switch varIndex - da {
                case 0:
                        return ly.NeuroMod.DA
                case 1:
                        return ly.NeuroMod.ACh
                case 2:
                        return ly.NeuroMod.SE
                case 3:
                        return ly.Pools[nrn.SubPool].Gate.Act
                case 4:
                        return num.FromBool[float32](ly.Pools[nrn.SubPool].Gate.Now)
                case 5:
                        return float32(ly.Pools[nrn.SubPool].Gate.Cnt)
                }
        }
        return nrn.VarByIndex(varIndex)
}

// UnitValues fills in values of given variable name on unit,
// for each unit in the layer, into given float32 slice (only resized if not big enough).
// Returns error on invalid var name.
func (ly *Layer) UnitValues(vals *[]float32, varNm string, di int) error {
        nn := len(ly.Neurons)
        if *vals == nil || cap(*vals) < nn {
                *vals = make([]float32, nn)
        } else if len(*vals) < nn {
                *vals = (*vals)[0:nn]
        }
        vidx, err := ly.UnitVarIndex(varNm)
        if err != nil {
                nan := math32.NaN()
                for i := range ly.Neurons {
                        (*vals)[i] = nan
                }
                return err
        }
        for i := range ly.Neurons {
                (*vals)[i] = ly.UnitValue1D(vidx, i, di)
        }
        return nil
}

// UnitValuesTensor returns values of given variable name on unit
// for each unit in the layer, as a float32 tensor in same shape as layer units.
func (ly *Layer) UnitValuesTensor(tsr tensor.Tensor, varNm string, di int) error {
        if tsr == nil {
                err := fmt.Errorf("leabra.UnitValuesTensor: Tensor is nil")
                log.Println(err)
                return err
        }
        tsr.SetShape(ly.Shape.Sizes, ly.Shape.Names...)
        vidx, err := ly.UnitVarIndex(varNm)
        if err != nil {
                nan := math.NaN()
                for i := range ly.Neurons {
                        tsr.SetFloat1D(i, nan)
                }
                return err
        }
        for i := range ly.Neurons {
                v := ly.UnitValue1D(vidx, i, di)
                if math32.IsNaN(v) {
                        tsr.SetFloat1D(i, math.NaN())
                } else {
                        tsr.SetFloat1D(i, float64(v))
                }
        }
        return nil
}

// UnitValuesSampleTensor fills in values of given variable name on unit
// for a smaller subset of sample units in the layer, into given tensor.
// This is used for computationally intensive stats or displays that work
// much better with a smaller number of units.
// The set of sample units are defined by SampleIndexes -- all units
// are used if no such subset has been defined.
// If tensor is not already big enough to hold the values, it is
// set to a 1D shape to hold all the values if subset is defined,
// otherwise it calls UnitValuesTensor and is identical to that.
// Returns error on invalid var name.
func (ly *Layer) UnitValuesSampleTensor(tsr tensor.Tensor, varNm string, di int) error {
        nu := len(ly.SampleIndexes)
        if nu == 0 {
                return ly.UnitValuesTensor(tsr, varNm, di)
        }
        if tsr == nil {
                err := fmt.Errorf("axon.UnitValuesSampleTensor: Tensor is nil")
                log.Println(err)
                return err
        }
        if tsr.Len() != nu {
                tsr.SetShape([]int{nu}, "Units")
        }
        vidx, err := ly.UnitVarIndex(varNm)
        if err != nil {
                nan := math.NaN()
                for i, _ := range ly.SampleIndexes {
                        tsr.SetFloat1D(i, nan)
                }
                return err
        }
        for i, ui := range ly.SampleIndexes {
                v := ly.UnitValue1D(vidx, ui, di)
                if math32.IsNaN(v) {
                        tsr.SetFloat1D(i, math.NaN())
                } else {
                        tsr.SetFloat1D(i, float64(v))
                }
        }
        return nil
}

// UnitVal returns value of given variable name on given unit,
// using shape-based dimensional index
func (ly *Layer) UnitValue(varNm string, idx []int, di int) float32 {
        vidx, err := ly.UnitVarIndex(varNm)
        if err != nil {
                return math32.NaN()
        }
        fidx := ly.Shape.Offset(idx)
        return ly.UnitValue1D(vidx, fidx, di)
}

// RecvPathValues fills in values of given synapse variable name,
// for pathway into given sending layer and neuron 1D index,
// for all receiving neurons in this layer,
// into given float32 slice (only resized if not big enough).
// pathType is the string representation of the path type -- used if non-empty,
// useful when there are multiple pathways between two layers.
// Returns error on invalid var name.
// If the receiving neuron is not connected to the given sending layer or neuron
// then the value is set to math32.NaN().
// Returns error on invalid var name or lack of recv path
// (vals always set to nan on path err).
func (ly *Layer) RecvPathValues(vals *[]float32, varNm string, sendLay emer.Layer, sendIndex1D int, pathType string) error {
        var err error
        nn := len(ly.Neurons)
        if *vals == nil || cap(*vals) < nn {
                *vals = make([]float32, nn)
        } else if len(*vals) < nn {
                *vals = (*vals)[0:nn]
        }
        nan := math32.NaN()
        for i := 0; i < nn; i++ {
                (*vals)[i] = nan
        }
        if sendLay == nil {
                return fmt.Errorf("sending layer is nil")
        }
        slay := sendLay.AsEmer()
        var pt emer.Path
        if pathType != "" {
                pt, err = slay.SendPathByRecvNameType(ly.Name, pathType)
                if pt == nil {
                        pt, err = slay.SendPathByRecvName(ly.Name)
                }
        } else {
                pt, err = slay.SendPathByRecvName(ly.Name)
        }
        if pt == nil {
                return err
        }
        if pt.AsEmer().Off {
                return fmt.Errorf("pathway is off")
        }
        for ri := 0; ri < nn; ri++ {
                (*vals)[ri] = pt.AsEmer().SynValue(varNm, sendIndex1D, ri) // this will work with any variable -- slower, but necessary
        }
        return nil
}

// SendPathValues fills in values of given synapse variable name,
// for pathway into given receiving layer and neuron 1D index,
// for all sending neurons in this layer,
// into given float32 slice (only resized if not big enough).
// pathType is the string representation of the path type -- used if non-empty,
// useful when there are multiple pathways between two layers.
// Returns error on invalid var name.
// If the sending neuron is not connected to the given receiving layer or neuron
// then the value is set to math32.NaN().
// Returns error on invalid var name or lack of recv path
// (vals always set to nan on path err).
func (ly *Layer) SendPathValues(vals *[]float32, varNm string, recvLay emer.Layer, recvIndex1D int, pathType string) error {
        var err error
        nn := len(ly.Neurons)
        if *vals == nil || cap(*vals) < nn {
                *vals = make([]float32, nn)
        } else if len(*vals) < nn {
                *vals = (*vals)[0:nn]
        }
        nan := math32.NaN()
        for i := 0; i < nn; i++ {
                (*vals)[i] = nan
        }
        if recvLay == nil {
                return fmt.Errorf("receiving layer is nil")
        }
        rlay := recvLay.AsEmer()
        var pt emer.Path
        if pathType != "" {
                pt, err = rlay.RecvPathBySendNameType(ly.Name, pathType)
                if pt == nil {
                        pt, err = rlay.RecvPathBySendName(ly.Name)
                }
        } else {
                pt, err = rlay.RecvPathBySendName(ly.Name)
        }
        if pt == nil {
                return err
        }
        if pt.AsEmer().Off {
                return fmt.Errorf("pathway is off")
        }
        for si := 0; si < nn; si++ {
                (*vals)[si] = pt.AsEmer().SynValue(varNm, si, recvIndex1D)
        }
        return nil
}

// Pool returns pool at given index
func (ly *Layer) Pool(idx int) *Pool {
        return &(ly.Pools[idx])
}

// AddSendTo adds given layer name(s) to list of those to send to.
func (ly *Layer) AddSendTo(laynm ...string) {
        ly.SendTo.Add(laynm...)
}

// AddAllSendToBut adds all layers in network except those in exclude list.
func (ly *Layer) AddAllSendToBut(excl ...string) {
        ly.SendTo.AddAllBut(ly.Network, excl...)
}

// ValidateSendTo ensures that SendTo layer names are valid.
func (ly *Layer) ValidateSendTo() error {
        return ly.SendTo.Validate(ly.Network)
}

//////////////////////////////////////////////////////////////////////////////////////
//  Build

// BuildSubPools initializes neuron start / end indexes for sub-pools
func (ly *Layer) BuildSubPools() {
        if !ly.Is4D() {
                return
        }
        sh := ly.Shape.Sizes
        spy := sh[0]
        spx := sh[1]
        pi := 1
        for py := 0; py < spy; py++ {
                for px := 0; px < spx; px++ {
                        soff := ly.Shape.Offset([]int{py, px, 0, 0})
                        eoff := ly.Shape.Offset([]int{py, px, sh[2] - 1, sh[3] - 1}) + 1
                        pl := &ly.Pools[pi]
                        pl.StIndex = soff
                        pl.EdIndex = eoff
                        for ni := pl.StIndex; ni < pl.EdIndex; ni++ {
                                nrn := &ly.Neurons[ni]
                                nrn.SubPool = int32(pi)
                        }
                        pi++
                }
        }
}

// BuildPools builds the inhibitory pools structures -- nu = number of units in layer
func (ly *Layer) BuildPools(nu int) error {
        np := 1 + ly.NumPools()
        ly.Pools = make([]Pool, np)
        lpl := &ly.Pools[0]
        lpl.StIndex = 0
        lpl.EdIndex = nu
        if np > 1 {
                ly.BuildSubPools()
        }
        return nil
}

// BuildPaths builds the pathways, recv-side
func (ly *Layer) BuildPaths() error {
        emsg := ""
        for _, pt := range ly.RecvPaths {
                if pt.Off {
                        continue
                }
                err := pt.Build()
                if err != nil {
                        emsg += err.Error() + "\n"
                }
        }
        if emsg != "" {
                return errors.New(emsg)
        }
        return nil
}

// Build constructs the layer state, including calling Build on the pathways
func (ly *Layer) Build() error {
        nu := ly.Shape.Len()
        if nu == 0 {
                return fmt.Errorf("Build Layer %v: no units specified in Shape", ly.Name)
        }
        ly.Neurons = make([]Neuron, nu)
        err := ly.BuildPools(nu)
        if err != nil {
                return errors.Log(err)
        }
        err = ly.BuildPaths()
        if err != nil {
                return errors.Log(err)
        }
        err = ly.ValidateSendTo()
        if err != nil {
                return errors.Log(err)
        }
        err = ly.CIN.RewLays.Validate(ly.Network)
        if err != nil {
                return errors.Log(err)
        }
        return nil
}

// WriteWeightsJSON writes the weights from this layer from the receiver-side perspective
// in a JSON text format.  We build in the indentation logic to make it much faster and
// more efficient.
func (ly *Layer) WriteWeightsJSON(w io.Writer, depth int) {
        ly.MetaData = make(map[string]string)
        ly.MetaData["ActMAvg"] = fmt.Sprintf("%g", ly.Pools[0].ActAvg.ActMAvg)
        ly.MetaData["ActPAvg"] = fmt.Sprintf("%g", ly.Pools[0].ActAvg.ActPAvg)
        ly.LayerBase.WriteWeightsJSONBase(w, depth)
}

// SetWeights sets the weights for this layer from weights.Layer decoded values
func (ly *Layer) SetWeights(lw *weights.Layer) error {
        if ly.Off {
                return nil
        }
        if lw.MetaData != nil {
                if am, ok := lw.MetaData["ActMAvg"]; ok {
                        pv, _ := strconv.ParseFloat(am, 32)
                        ly.Pools[0].ActAvg.ActMAvg = float32(pv)
                }
                if ap, ok := lw.MetaData["ActPAvg"]; ok {
                        pv, _ := strconv.ParseFloat(ap, 32)
                        pl := &ly.Pools[0]
                        pl.ActAvg.ActPAvg = float32(pv)
                        ly.Inhib.ActAvg.EffFromAvg(&pl.ActAvg.ActPAvgEff, pl.ActAvg.ActPAvg)
                }
        }
        var err error
        rpts := ly.RecvPaths
        if len(lw.Paths) == len(rpts) { // this is essential if multiple paths from same layer
                for pi := range lw.Paths {
                        pw := &lw.Paths[pi]
                        pt := (rpts)[pi]
                        er := pt.SetWeights(pw)
                        if er != nil {
                                err = er
                        }
                }
        } else {
                for pi := range lw.Paths {
                        pw := &lw.Paths[pi]
                        pt, err := ly.RecvPathBySendName(pw.From)
                        if err == nil {
                                er := pt.SetWeights(pw)
                                if er != nil {
                                        err = er
                                }
                        }
                }
        }
        return err
}

// VarRange returns the min / max values for given variable
// todo: support r. s. pathway values
func (ly *Layer) VarRange(varNm string) (min, max float32, err error) {
        sz := len(ly.Neurons)
        if sz == 0 {
                return
        }
        vidx := 0
        vidx, err = NeuronVarIndexByName(varNm)
        if err != nil {
                return
        }

        v0 := ly.Neurons[0].VarByIndex(vidx)
        min = v0
        max = v0
        for i := 1; i < sz; i++ {
                vl := ly.Neurons[i].VarByIndex(vidx)
                if vl < min {
                        min = vl
                }
                if vl > max {
                        max = vl
                }
        }
        return
}

// Copyright (c) 2019, The Emergent Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package leabra

import (
        "cogentcore.org/core/math32"
)

///////////////////////////////////////////////////////////////////////
//  learn.go contains the learning params and functions for leabra

// leabra.LearnNeurParams manages learning-related parameters at the neuron-level.
// This is mainly the running average activations that drive learning.
type LearnNeurParams struct {

        // parameters for computing running average activations that drive learning
        ActAvg LrnActAvgParams `display:"inline"`

        // parameters for computing AvgL long-term running average
        AvgL AvgLParams `display:"inline"`

        // parameters for computing cosine diff between minus and plus phase
        CosDiff CosDiffParams `display:"inline"`
}

func (ln *LearnNeurParams) Update() {
        ln.ActAvg.Update()
        ln.AvgL.Update()
        ln.CosDiff.Update()
}

func (ln *LearnNeurParams) Defaults() {
        ln.ActAvg.Defaults()
        ln.AvgL.Defaults()
        ln.CosDiff.Defaults()
}

// InitActAvg initializes the running-average activation values that drive learning.
// Called by InitWeights (at start of learning).
func (ln *LearnNeurParams) InitActAvg(nrn *Neuron) {
        nrn.AvgSS = ln.ActAvg.Init
        nrn.AvgS = ln.ActAvg.Init
        nrn.AvgM = ln.ActAvg.Init
        nrn.AvgL = ln.AvgL.Init
        nrn.AvgSLrn = 0
        nrn.ActAvg = ln.ActAvg.Init
}

// AvgsFromAct updates the running averages based on current learning activation.
// Computed after new activation for current cycle is updated.
func (ln *LearnNeurParams) AvgsFromAct(nrn *Neuron) {
        ln.ActAvg.AvgsFromAct(nrn.ActLrn, &nrn.AvgSS, &nrn.AvgS, &nrn.AvgM, &nrn.AvgSLrn)
}

// AvgLFromAct computes long-term average activation value, and learning factor, from current AvgM.
// Called at start of new alpha-cycle.
func (ln *LearnNeurParams) AvgLFromAvgM(nrn *Neuron) {
        ln.AvgL.AvgLFromAvgM(nrn.AvgM, &nrn.AvgL, &nrn.AvgLLrn)
}

///////////////////////////////////////////////////////////////////////
//  LearnSynParams

// leabra.LearnSynParams manages learning-related parameters at the synapse-level.
type LearnSynParams struct {

        // enable learning for this pathway
        Learn bool

        // current effective learning rate (multiplies DWt values, determining rate of change of weights)
        Lrate float32

        // initial learning rate -- this is set from Lrate in UpdateParams, which is called when Params are updated, and used in LrateMult to compute a new learning rate for learning rate schedules.
        LrateInit float32

        // parameters for the XCal learning rule
        XCal XCalParams `display:"inline"`

        // parameters for the sigmoidal contrast weight enhancement
        WtSig WtSigParams `display:"inline"`

        // parameters for normalizing weight changes by abs max dwt
        Norm DWtNormParams `display:"inline"`

        // parameters for momentum across weight changes
        Momentum MomentumParams `display:"inline"`

        // parameters for balancing strength of weight increases vs. decreases
        WtBal WtBalParams `display:"inline"`
}

func (ls *LearnSynParams) Update() {
        ls.XCal.Update()
        ls.WtSig.Update()
        ls.Norm.Update()
        ls.Momentum.Update()
        ls.WtBal.Update()
}

func (ls *LearnSynParams) Defaults() {
        ls.Learn = true
        ls.Lrate = 0.04
        ls.LrateInit = ls.Lrate
        ls.XCal.Defaults()
        ls.WtSig.Defaults()
        ls.Norm.Defaults()
        ls.Momentum.Defaults()
        ls.WtBal.Defaults()
}

func (ls *LearnSynParams) ShouldDisplay(field string) bool {
        switch field {
        case "Lrate", "LrateInit", "XCal", "WtSig", "Norm", "Momentum", "WtBal":
                return ls.Learn
        default:
                return true
        }
}

// LWtFromWt updates the linear weight value based on the current effective Wt value.
// effective weight is sigmoidally contrast-enhanced relative to the linear weight.
func (ls *LearnSynParams) LWtFromWt(syn *Synapse) {
        syn.LWt = ls.WtSig.LinFromSigWt(syn.Wt / syn.Scale) // must factor out scale too!
}

// WtFromLWt updates the effective weight value based on the current linear Wt value.
// effective weight is sigmoidally contrast-enhanced relative to the linear weight.
func (ls *LearnSynParams) WtFromLWt(syn *Synapse) {
        syn.Wt = ls.WtSig.SigFromLinWt(syn.LWt)
        syn.Wt *= syn.Scale
}

// CHLdWt returns the error-driven and BCM Hebbian weight change components for the
// temporally eXtended Contrastive Attractor Learning (XCAL), CHL version
func (ls *LearnSynParams) CHLdWt(suAvgSLrn, suAvgM, ruAvgSLrn, ruAvgM, ruAvgL float32) (err, bcm float32) {
        srs := suAvgSLrn * ruAvgSLrn
        srm := suAvgM * ruAvgM
        bcm = ls.XCal.DWt(srs, ruAvgL)
        err = ls.XCal.DWt(srs, srm)
        return
}

// BCMdWt returns the BCM Hebbian weight change for AvgSLrn vs. AvgL
// long-term average floating activation on the receiver.
func (ls *LearnSynParams) BCMdWt(suAvgSLrn, ruAvgSLrn, ruAvgL float32) float32 {
        srs := suAvgSLrn * ruAvgSLrn
        return ls.XCal.DWt(srs, ruAvgL)
}

// WtFromDWt updates the synaptic weights from accumulated weight changes
// wbInc and wbDec are the weight balance factors, wt is the sigmoidal contrast-enhanced
// weight and lwt is the linear weight value
func (ls *LearnSynParams) WtFromDWt(wbInc, wbDec float32, dwt, wt, lwt *float32, scale float32) {
        if *dwt == 0 {
                return
        }
        if ls.WtSig.SoftBound {
                if *dwt > 0 {
                        *dwt *= wbInc * (1 - *lwt)
                } else {
                        *dwt *= wbDec * *lwt
                }
        } else {
                if *dwt > 0 {
                        *dwt *= wbInc
                } else {
                        *dwt *= wbDec
                }
        }
        *lwt += *dwt
        if *lwt < 0 {
                *lwt = 0
        } else if *lwt > 1 {
                *lwt = 1
        }
        *wt = scale * ls.WtSig.SigFromLinWt(*lwt)
        *dwt = 0
}

// LrnActAvgParams has rate constants for averaging over activations at different time scales,
// to produce the running average activation values that then drive learning in the XCAL learning rules
type LrnActAvgParams struct {

        // time constant in cycles, which should be milliseconds typically (roughly, how long it takes for value to change significantly -- 1.4x the half-life), for continuously updating the super-short time-scale avg_ss value -- this is provides a pre-integration step before integrating into the avg_s short time scale -- it is particularly important for spiking -- in general 4 is the largest value without starting to impair learning, but a value of 7 can be combined with m_in_s = 0 with somewhat worse results
        SSTau float32 `default:"2,4,7"  min:"1"`

        // time constant in cycles, which should be milliseconds typically (roughly, how long it takes for value to change significantly -- 1.4x the half-life), for continuously updating the short time-scale avg_s value from the super-short avg_ss value (cascade mode) -- avg_s represents the plus phase learning signal that reflects the most recent past information
        STau float32 `default:"2" min:"1"`

        // time constant in cycles, which should be milliseconds typically (roughly, how long it takes for value to change significantly -- 1.4x the half-life), for continuously updating the medium time-scale avg_m value from the short avg_s value (cascade mode) -- avg_m represents the minus phase learning signal that reflects the expectation representation prior to experiencing the outcome (in addition to the outcome) -- the default value of 10 generally cannot be exceeded without impairing learning
        MTau float32 `default:"10" min:"1"`

        // how much of the medium term average activation to mix in with the short (plus phase) to compute the Neuron AvgSLrn variable that is used for the unit's short-term average in learning. This is important to ensure that when unit turns off in plus phase (short time scale), enough medium-phase trace remains so that learning signal doesn't just go all the way to 0, at which point no learning would take place -- typically need faster time constant for updating S such that this trace of the M signal is lost -- can set SSTau=7 and set this to 0 but learning is generally somewhat worse
        LrnM float32 `default:"0.1,0" min:"0" max:"1"`

        // initial value for average
        Init float32 `default:"0.15" min:"0" max:"1"`

        // rate = 1 / tau
        SSDt float32 `display:"-" json:"-" xml:"-" edit:"-"`

        // rate = 1 / tau
        SDt float32 `display:"-" json:"-" xml:"-" edit:"-"`

        // rate = 1 / tau
        MDt float32 `display:"-" json:"-" xml:"-" edit:"-"`

        // 1-LrnM
        LrnS float32 `display:"-" json:"-" xml:"-" edit:"-"`
}

// AvgsFromAct computes averages based on current act
func (aa *LrnActAvgParams) AvgsFromAct(ruAct float32, avgSS, avgS, avgM, avgSLrn *float32) {
        *avgSS += aa.SSDt * (ruAct - *avgSS)
        *avgS += aa.SDt * (*avgSS - *avgS)
        *avgM += aa.MDt * (*avgS - *avgM)

        *avgSLrn = aa.LrnS**avgS + aa.LrnM**avgM
}

func (aa *LrnActAvgParams) Update() {
        aa.SSDt = 1 / aa.SSTau
        aa.SDt = 1 / aa.STau
        aa.MDt = 1 / aa.MTau
        aa.LrnS = 1 - aa.LrnM
}

func (aa *LrnActAvgParams) Defaults() {
        aa.SSTau = 2.0
        aa.STau = 2.0
        aa.MTau = 10.0
        aa.LrnM = 0.1
        aa.Init = 0.15
        aa.Update()

}

// AvgLParams are parameters for computing the long-term floating average value, AvgL
// which is used for driving BCM-style hebbian learning in XCAL -- this form of learning
// increases contrast of weights and generally decreases overall activity of neuron,
// to prevent "hog" units -- it is computed as a running average of the (gain multiplied)
// medium-time-scale average activation at the end of the alpha-cycle.
// Also computes an adaptive amount of BCM learning, AvgLLrn, based on AvgL.
type AvgLParams struct {

        // initial AvgL value at start of training
        Init float32 `default:"0.4" min:"0" max:"1"`

        // gain multiplier on activation used in computing the running average AvgL value that is the key floating threshold in the BCM Hebbian learning rule -- when using the DELTA_FF_FB learning rule, it should generally be 2x what it was before with the old XCAL_CHL rule, i.e., default of 5 instead of 2.5 -- it is a good idea to experiment with this parameter a bit -- the default is on the high-side, so typically reducing a bit from initial default is a good direction
        Gain float32 `default:"1.5,2,2.5,3,4,5" min:"0"`

        // miniumum AvgL value -- running average cannot go lower than this value even when it otherwise would due to inactivity -- default value is generally good and typically does not need to be changed
        Min float32 `default:"0.2" min:"0"`

        // time constant for updating the running average AvgL -- AvgL moves toward gain*act with this time constant on every alpha-cycle - longer time constants can also work fine, but the default of 10 allows for quicker reaction to beneficial weight changes
        Tau float32 `default:"10" min:"1"`

        // maximum AvgLLrn value, which is amount of learning driven by AvgL factor -- when AvgL is at its maximum value (i.e., gain, as act does not exceed 1), then AvgLLrn will be at this maximum value -- by default, strong amounts of this homeostatic Hebbian form of learning can be used when the receiving unit is highly active -- this will then tend to bring down the average activity of units -- the default of 0.5, in combination with the err_mod flag, works well for most models -- use around 0.0004 for a single fixed value (with err_mod flag off)
        LrnMax float32 `default:"0.5" min:"0"`

        // miniumum AvgLLrn value (amount of learning driven by AvgL factor) -- if AvgL is at its minimum value, then AvgLLrn will be at this minimum value -- neurons that are not overly active may not need to increase the contrast of their weights as much -- use around 0.0004 for a single fixed value (with err_mod flag off)
        LrnMin float32 `default:"0.0001,0.0004" min:"0"`

        // modulate amount learning by normalized level of error within layer
        ErrMod bool `default:"true"`

        // minimum modulation value for ErrMod-- ensures a minimum amount of self-organizing learning even for network / layers that have a very small level of error signal
        ModMin float32 `default:"0.01"`

        // rate = 1 / tau
        Dt float32 `display:"-" json:"-" xml:"-" edit:"-"`

        // (LrnMax - LrnMin) / (Gain - Min)
        LrnFact float32 `display:"-" json:"-" xml:"-" edit:"-"`
}

func (al *AvgLParams) ShouldDisplay(field string) bool {
        switch field {
        case "ModMin":
                return al.ErrMod
        default:
                return true
        }
}

// AvgLFromAvgM computes long-term average activation value, and learning factor, from given
// medium-scale running average activation avgM
func (al *AvgLParams) AvgLFromAvgM(avgM float32, avgL, lrn *float32) {
        *avgL += al.Dt * (al.Gain*avgM - *avgL)
        if *avgL < al.Min {
                *avgL = al.Min
        }
        *lrn = al.LrnFact * (*avgL - al.Min)
}

// ErrModFromLayErr computes AvgLLrn multiplier from layer cosine diff avg statistic
func (al *AvgLParams) ErrModFromLayErr(layCosDiffAvg float32) float32 {
        lmod := float32(1)
        if !al.ErrMod {
                return lmod
        }
        lmod *= math32.Max(layCosDiffAvg, al.ModMin)
        return lmod
}

func (al *AvgLParams) Update() {
        al.Dt = 1 / al.Tau
        al.LrnFact = (al.LrnMax - al.LrnMin) / (al.Gain - al.Min)
}

func (al *AvgLParams) Defaults() {
        al.Init = 0.4
        al.Gain = 2.5
        al.Min = 0.2
        al.Tau = 10
        al.LrnMax = 0.5
        al.LrnMin = 0.0001
        al.ErrMod = true
        al.ModMin = 0.01
        al.Update()
}

//////////////////////////////////////////////////////////////////////////////////////
//  CosDiffParams

// CosDiffParams specify how to integrate cosine of difference between plus and minus phase activations
// Used to modulate amount of hebbian learning, and overall learning rate.
type CosDiffParams struct {

        // time constant in alpha-cycles (roughly how long significant change takes, 1.4 x half-life) for computing running average CosDiff value for the layer, CosDiffAvg = cosine difference between ActM and ActP -- this is an important statistic for how much phase-based difference there is between phases in this layer -- it is used in standard X_COS_DIFF modulation of l_mix in LeabraConSpec, and for modulating learning rate as a function of predictability in the DeepLeabra predictive auto-encoder learning -- running average variance also computed with this: cos_diff_var
        Tau float32 `default:"100" min:"1"`

        // rate constant = 1 / Tau
        Dt float32 `edit:"-" display:"-" json:"-" xml:"-"`

        // complement of rate constant = 1 - Dt
        DtC float32 `edit:"-" display:"-" json:"-" xml:"-"`
}

func (cd *CosDiffParams) Update() {
        cd.Dt = 1 / cd.Tau
        cd.DtC = 1 - cd.Dt
}

func (cd *CosDiffParams) Defaults() {
        cd.Tau = 100
        cd.Update()
}

// AvgVarFromCos updates the average and variance from current cosine diff value
func (cd *CosDiffParams) AvgVarFromCos(avg, vr *float32, cos float32) {
        if *avg == 0 { // first time -- set
                *avg = cos
                *vr = 0
        } else {
                del := cos - *avg
                incr := cd.Dt * del
                *avg += incr
                // following is magic exponentially weighted incremental variance formula
                // derived by Finch, 2009: Incremental calculation of weighted mean and variance
                if *vr == 0 {
                        *vr = 2 * cd.DtC * del * incr
                } else {
                        *vr = cd.DtC * (*vr + del*incr)
                }
        }
}

// LrateMod computes learning rate modulation based on cos diff vals
// func (cd *CosDiffParams) LrateMod(cos, avg, vr float32) float32 {
//         if vr <= 0 {
//                 return 1
//         }
//         zval := (cos - avg) / math32.Sqrt(vr) // stdev = sqrt of var
//         // z-normal value is starting point for learning rate factor
//         //    if zval < lrmod_z_thr {
//         //         return 0
//         // }
//         return 1
// }

//////////////////////////////////////////////////////////////////////////////////////
//  CosDiffStats

// CosDiffStats holds cosine-difference statistics at the layer level
type CosDiffStats struct {

        // cosine (normalized dot product) activation difference between ActP and ActM on this alpha-cycle for this layer -- computed by CosDiffFromActs at end of QuarterFinal for quarter = 3
        Cos float32

        // running average of cosine (normalized dot product) difference between ActP and ActM -- computed with CosDiff.Tau time constant in QuarterFinal, and used for modulating BCM Hebbian learning (see AvgLrn) and overall learning rate
        Avg float32

        // running variance of cosine (normalized dot product) difference between ActP and ActM -- computed with CosDiff.Tau time constant in QuarterFinal, used for modulating overall learning rate
        Var float32

        // 1 - Avg and 0 for non-Hidden layers
        AvgLrn float32

        // 1 - AvgLrn and 0 for non-Hidden layers -- this is the value of Avg used for AvgLParams ErrMod modulation of the AvgLLrn factor if enabled
        ModAvgLLrn float32
}

func (cd *CosDiffStats) Init() {
        cd.Cos = 0
        cd.Avg = 0
        cd.Var = 0
        cd.AvgLrn = 0
        cd.ModAvgLLrn = 0
}

//////////////////////////////////////////////////////////////////////////////////////
//  XCalParams

// XCalParams are parameters for temporally eXtended Contrastive Attractor Learning function (XCAL)
// which is the standard learning equation for leabra .
type XCalParams struct {

        // multiplier on learning based on the medium-term floating average threshold which produces error-driven learning -- this is typically 1 when error-driven learning is being used, and 0 when pure Hebbian learning is used. The long-term floating average threshold is provided by the receiving unit
        MLrn float32 `default:"1" min:"0"`

        // if true, set a fixed AvgLLrn weighting factor that determines how much of the long-term floating average threshold (i.e., BCM, Hebbian) component of learning is used -- this is useful for setting a fully Hebbian learning connection, e.g., by setting MLrn = 0 and LLrn = 1. If false, then the receiving unit's AvgLLrn factor is used, which dynamically modulates the amount of the long-term component as a function of how active overall it is
        SetLLrn bool `default:"false"`

        // fixed l_lrn weighting factor that determines how much of the long-term floating average threshold (i.e., BCM, Hebbian) component of learning is used -- this is useful for setting a fully Hebbian learning connection, e.g., by setting MLrn = 0 and LLrn = 1.
        LLrn float32

        // proportional point within LTD range where magnitude reverses to go back down to zero at zero -- err-driven svm component does better with smaller values, and BCM-like mvl component does better with larger values -- 0.1 is a compromise
        DRev float32 `default:"0.1" min:"0" max:"0.99"`

        // minimum LTD threshold value below which no weight change occurs -- this is now *relative* to the threshold
        DThr float32 `default:"0.0001,0.01" min:"0"`

        // xcal learning threshold -- don't learn when sending unit activation is below this value in both phases -- due to the nature of the learning function being 0 when the sr coproduct is 0, it should not affect learning in any substantial way -- nonstandard learning algorithms that have different properties should ignore it
        LrnThr float32 `default:"0.01"`

        // -(1-DRev)/DRev -- multiplication factor in learning rule -- builds in the minus sign!
        DRevRatio float32 `edit:"-" display:"-" json:"-" xml:"-"`
}

func (xc *XCalParams) Update() {
        if xc.DRev > 0 {
                xc.DRevRatio = -(1 - xc.DRev) / xc.DRev
        } else {
                xc.DRevRatio = -1
        }
}

func (xc *XCalParams) Defaults() {
        xc.MLrn = 1
        xc.SetLLrn = false
        xc.LLrn = 1
        xc.DRev = 0.1
        xc.DThr = 0.0001
        xc.LrnThr = 0.01
        xc.Update()
}

func (xc *XCalParams) ShouldDisplay(field string) bool {
        switch field {
        case "LLrn":
                return xc.SetLLrn
        default:
                return true
        }
}

// DWt is the XCAL function for weight change -- the "check mark" function -- no DGain, no ThrPMin
func (xc *XCalParams) DWt(srval, thrP float32) float32 {
        var dwt float32
        if srval < xc.DThr {
                dwt = 0
        } else if srval > thrP*xc.DRev {
                dwt = (srval - thrP)
        } else {
                dwt = srval * xc.DRevRatio
        }
        return dwt
}

// LongLrate returns the learning rate for long-term floating average component (BCM)
func (xc *XCalParams) LongLrate(avgLLrn float32) float32 {
        if xc.SetLLrn {
                return xc.LLrn
        }
        return avgLLrn
}

//////////////////////////////////////////////////////////////////////////////////////
//  WtSigParams

// WtSigParams are sigmoidal weight contrast enhancement function parameters
type WtSigParams struct {

        // gain (contrast, sharpness) of the weight contrast function (1 = linear)
        Gain float32 `default:"1,6" min:"0"`

        // offset of the function (1=centered at .5, >1=higher, <1=lower) -- 1 is standard for XCAL
        Off float32 `default:"1" min:"0"`

        // apply exponential soft bounding to the weight changes
        SoftBound bool `default:"true"`
}

func (ws *WtSigParams) Update() {
}

func (ws *WtSigParams) Defaults() {
        ws.Gain = 6
        ws.Off = 1
        ws.SoftBound = true
}

// SigFun is the sigmoid function for value w in 0-1 range, with gain and offset params
func SigFun(w, gain, off float32) float32 {
        if w <= 0 {
                return 0
        }
        if w >= 1 {
                return 1
        }
        return (1 / (1 + math32.Pow((off*(1-w))/w, gain)))
}

// SigFun61 is the sigmoid function for value w in 0-1 range, with default gain = 6, offset = 1 params
func SigFun61(w float32) float32 {
        if w <= 0 {
                return 0
        }
        if w >= 1 {
                return 1
        }
        pw := (1 - w) / w
        return (1 / (1 + pw*pw*pw*pw*pw*pw))
}

// SigInvFun is the inverse of the sigmoid function
func SigInvFun(w, gain, off float32) float32 {
        if w <= 0 {
                return 0
        }
        if w >= 1 {
                return 1
        }
        return 1.0 / (1.0 + math32.Pow((1.0-w)/w, 1/gain)/off)
}

// SigInvFun61 is the inverse of the sigmoid function, with default gain = 6, offset = 1 params
func SigInvFun61(w float32) float32 {
        if w <= 0 {
                return 0
        }
        if w >= 1 {
                return 1
        }
        rval := 1.0 / (1.0 + math32.Pow((1.0-w)/w, 1.0/6.0))
        return rval
}

// SigFromLinWt returns sigmoidal contrast-enhanced weight from linear weight
func (ws *WtSigParams) SigFromLinWt(lw float32) float32 {
        if ws.Gain == 1 && ws.Off == 1 {
                return lw
        }
        if ws.Gain == 6 && ws.Off == 1 {
                return SigFun61(lw)
        }
        return SigFun(lw, ws.Gain, ws.Off)
}

// LinFromSigWt returns linear weight from sigmoidal contrast-enhanced weight
func (ws *WtSigParams) LinFromSigWt(sw float32) float32 {
        if ws.Gain == 1 && ws.Off == 1 {
                return sw
        }
        if ws.Gain == 6 && ws.Off == 1 {
                return SigInvFun61(sw)
        }
        return SigInvFun(sw, ws.Gain, ws.Off)
}

//////////////////////////////////////////////////////////////////////////////////////
//  DWtNormParams

// DWtNormParams are weight change (dwt) normalization parameters, using MAX(ABS(dwt)) aggregated over
// Sending connections in a given pathway for a given unit.
// Slowly decays and instantly resets to any current max(abs)
// Serves as an estimate of the variance in the weight changes, assuming zero net mean overall.
type DWtNormParams struct {

        // whether to use dwt normalization, only on error-driven dwt component, based on pathway-level max_avg value -- slowly decays and instantly resets to any current max
        On bool `default:"true"`

        // time constant for decay of dwnorm factor -- generally should be long-ish, between 1000-10000 -- integration rate factor is 1/tau
        DecayTau float32 `min:"1" default:"1000,10000"`

        // minimum effective value of the normalization factor -- provides a lower bound to how much normalization can be applied
        NormMin float32 `min:"0" default:"0.001"`

        // overall learning rate multiplier to compensate for changes due to use of normalization -- allows for a common master learning rate to be used between different conditions -- 0.1 for synapse-level, maybe higher for other levels
        LrComp float32 `min:"0" default:"0.15"`

        // record the avg, max values of err, bcm hebbian, and overall dwt change per con group and per pathway
        Stats bool `default:"false"`

        // rate constant of decay = 1 / decay_tau
        DecayDt float32 `edit:"-" display:"-" json:"-" xml:"-"`

        // complement rate constant of decay = 1 - (1 / decay_tau)
        DecayDtC float32 `edit:"-" display:"-" json:"-" xml:"-"`
}

// DWtNormParams updates the dwnorm running max_abs, slowly decaying value
// jumps up to max(abs_dwt) and slowly decays
// returns the effective normalization factor, as a multiplier, including lrate comp
func (dn *DWtNormParams) NormFromAbsDWt(norm *float32, absDwt float32) float32 {
        *norm = math32.Max(dn.DecayDtC**norm, absDwt)
        if *norm == 0 {
                return 1
        }
        return dn.LrComp / math32.Max(*norm, dn.NormMin)
}

func (dn *DWtNormParams) Update() {
        dn.DecayDt = 1 / dn.DecayTau
        dn.DecayDtC = 1 - dn.DecayDt
}

func (dn *DWtNormParams) Defaults() {
        dn.On = true
        dn.DecayTau = 1000
        dn.LrComp = 0.15
        dn.NormMin = 0.001
        dn.Stats = false
        dn.Update()
}

func (dn *DWtNormParams) ShouldDisplay(field string) bool {
        switch field {
        case "DecayTau", "NormMin", "LrComp", "Stats":
                return dn.On
        default:
                return true
        }
}

//////////////////////////////////////////////////////////////////////////////////////
//  MomentumParams

// MomentumParams implements standard simple momentum -- accentuates consistent directions of weight change and
// cancels out dithering -- biologically captures slower timecourse of longer-term plasticity mechanisms.
type MomentumParams struct {

        // whether to use standard simple momentum
        On bool `default:"true"`

        // time constant factor for integration of momentum -- 1/tau is dt (e.g., .1), and 1-1/tau (e.g., .95 or .9) is traditional momentum time-integration factor
        MTau float32 `min:"1" default:"10"`

        // overall learning rate multiplier to compensate for changes due to JUST momentum without normalization -- allows for a common master learning rate to be used between different conditions -- generally should use .1 to compensate for just momentum itself
        LrComp float32 `min:"0" default:"0.1"`

        // rate constant of momentum integration = 1 / m_tau
        MDt float32 `edit:"-" display:"-" json:"-" xml:"-"`

        // complement rate constant of momentum integration = 1 - (1 / m_tau)
        MDtC float32 `edit:"-" display:"-" json:"-" xml:"-"`
}

// MomentFromDWt updates synaptic moment variable based on dwt weight change value
// and returns new momentum factor * LrComp
func (mp *MomentumParams) MomentFromDWt(moment *float32, dwt float32) float32 {
        *moment = mp.MDtC**moment + dwt
        return mp.LrComp * *moment
}

func (mp *MomentumParams) Update() {
        mp.MDt = 1 / mp.MTau
        mp.MDtC = 1 - mp.MDt
}

func (mp *MomentumParams) Defaults() {
        mp.On = true
        mp.MTau = 10
        mp.LrComp = 0.1
        mp.Update()
}

func (mp *MomentumParams) ShouldDisplay(field string) bool {
        switch field {
        case "MTau", "LrComp":
                return mp.On
        default:
                return true
        }
}

//////////////////////////////////////////////////////////////////////////////////////
//  WtBalParams

// WtBalParams are weight balance soft renormalization params:
// maintains overall weight balance by progressively penalizing weight increases as a function of
// how strong the weights are overall (subject to thresholding) and long time-averaged activation.
// Plugs into soft bounding function.
type WtBalParams struct {

        // perform weight balance soft normalization?  if so, maintains overall weight balance across units by progressively penalizing weight increases as a function of amount of averaged receiver weight above a high threshold (hi_thr) and long time-average activation above an act_thr -- this is generally very beneficial for larger models where hog units are a problem, but not as much for smaller models where the additional constraints are not beneficial -- uses a sigmoidal function: WbInc = 1 / (1 + HiGain*(WbAvg - HiThr) + ActGain * (nrn.ActAvg - ActThr)))
        On bool

        // apply soft bounding to target layers -- appears to be beneficial but still testing
        Targs bool

        // threshold on weight value for inclusion into the weight average that is then subject to the further HiThr threshold for then driving a change in weight balance -- this AvgThr allows only stronger weights to contribute so that weakening of lower weights does not dilute sensitivity to number and strength of strong weights
        AvgThr float32 `default:"0.25"`

        // high threshold on weight average (subject to AvgThr) before it drives changes in weight increase vs. decrease factors
        HiThr float32 `default:"0.4"`

        // gain multiplier applied to above-HiThr thresholded weight averages -- higher values turn weight increases down more rapidly as the weights become more imbalanced
        HiGain float32 `default:"4"`

        // low threshold on weight average (subject to AvgThr) before it drives changes in weight increase vs. decrease factors
        LoThr float32 `default:"0.4"`

        // gain multiplier applied to below-lo_thr thresholded weight averages -- higher values turn weight increases up more rapidly as the weights become more imbalanced -- generally beneficial but sometimes not -- worth experimenting with either 6 or 0
        LoGain float32 `default:"6,0"`
}

func (wb *WtBalParams) Update() {
}

func (wb *WtBalParams) Defaults() {
        wb.On = false
        wb.AvgThr = 0.25
        wb.HiThr = 0.4
        wb.HiGain = 4
        wb.LoThr = 0.4
        wb.LoGain = 6
}

func (wb *WtBalParams) ShouldDisplay(field string) bool {
        switch field {
        case "AvgThr", "HiThr", "HiGain", "LoThr", "LoGain":
                return wb.On
        default:
                return true
        }
}

// WtBal computes weight balance factors for increase and decrease based on extent
// to which weights and average act exceed thresholds
func (wb *WtBalParams) WtBal(wbAvg float32) (fact, inc, dec float32) {
        inc = 1
        dec = 1
        if wbAvg < wb.LoThr {
                if wbAvg < wb.AvgThr {
                        wbAvg = wb.AvgThr // prevent extreme low if everyone below thr
                }
                fact = wb.LoGain * (wb.LoThr - wbAvg)
                dec = 1 / (1 + fact)
                inc = 2 - dec
        } else if wbAvg > wb.HiThr {
                fact = wb.HiGain * (wbAvg - wb.HiThr)
                inc = 1 / (1 + fact) // gets sigmoidally small toward 0 as fact gets larger -- is quick acting but saturates -- apply pressure earlier..
                dec = 2 - inc        // as inc goes down, dec goes up..  sum to 2
        }
        return fact, inc, dec
}

/*
  /////////////////////////////////////
  // CtLeabraXCAL code

  INLINE void   GetLrates(LEABRA_CON_STATE* cg, LEABRA_NETWORK_STATE* net, int thr_no,
                          float& clrate, bool& deep_on, float& bg_lrate, float& fg_lrate)  {
    LEABRA_LAYER_STATE* rlay = cg->GetRecvLayer(net);
    clrate = cur_lrate * rlay->lrate_mod;
    deep_on = deep.on;
    if(deep_on) {
      if(!rlay->deep_lrate_mod)
        deep_on = false;          // only applicable to deep_norm active layers
    }
    if(deep_on) {
      bg_lrate = deep.bg_lrate;
      fg_lrate = deep.fg_lrate;
    }
  }
  // #IGNORE get the current learning rates including layer-specific and potential deep modulations

  // todo: should go back and explore this at some point:
  // if(xcal.one_thr) {
  //   float eff_thr = ru_avg_l_lrn * ru_avg_l + (1.0f - ru_avg_l_lrn) * srm;
  //   eff_thr = fminf(eff_thr, 1.0f);
  //   dwt += clrate * xcal.dWtFun(srs, eff_thr);
  // }
  // also: fminf(ru_avg_l,1.0f) for threshold as an option..


*/

// Copyright (c) 2022, The Emergent Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package leabra

import (
        "reflect"
        "strconv"

        "cogentcore.org/core/base/errors"
        "cogentcore.org/core/math32/minmax"
        "github.com/emer/emergent/v2/egui"
        "github.com/emer/emergent/v2/elog"
        "github.com/emer/emergent/v2/estats"
        "github.com/emer/emergent/v2/etime"
        "github.com/emer/etensor/plot/plotcore"
        "github.com/emer/etensor/tensor/stats/split"
        "github.com/emer/etensor/tensor/stats/stats"
        "github.com/emer/etensor/tensor/table"
)

// LogTestErrors records all errors made across TestTrials, at Test Epoch scope
func LogTestErrors(lg *elog.Logs) {
        sk := etime.Scope(etime.Test, etime.Trial)
        lt := lg.TableDetailsScope(sk)
        ix, _ := lt.NamedIndexView("TestErrors")
        ix.Filter(func(et *table.Table, row int) bool {
                return et.Float("Err", row) > 0 // include error trials
        })
        lg.MiscTables["TestErrors"] = ix.NewTable()

        allsp := split.All(ix)
        split.AggColumn(allsp, "UnitErr", stats.Sum)
        // note: can add other stats to compute
        lg.MiscTables["TestErrorStats"] = allsp.AggsToTable(table.AddAggName)
}

// PCAStats computes PCA statistics on recorded hidden activation patterns
// from Analyze, Trial log data
func PCAStats(net *Network, lg *elog.Logs, stats *estats.Stats) {
        stats.PCAStats(lg.IndexView(etime.Analyze, etime.Trial), "ActM", net.LayersByType(SuperLayer, TargetLayer, CTLayer))
}

//////////////////////////////////////////////////////////////////////////////
//  Log items

// LogAddDiagnosticItems adds standard Axon diagnostic statistics to given logs,
// across the given time levels, in higher to lower order, e.g., Epoch, Trial
// These are useful for tuning and diagnosing the behavior of the network.
func LogAddDiagnosticItems(lg *elog.Logs, layerNames []string, mode etime.Modes, times ...etime.Times) {
        ntimes := len(times)
        for _, lnm := range layerNames {
                clnm := lnm
                itm := lg.AddItem(&elog.Item{
                        Name:   clnm + "_ActMAvg",
                        Type:   reflect.Float64,
                        FixMax: false,
                        Range:  minmax.F32{Max: 1},
                        Write: elog.WriteMap{
                                etime.Scope(mode, times[ntimes-1]): func(ctx *elog.Context) {
                                        ly := ctx.Layer(clnm).(*Layer)
                                        ctx.SetFloat32(ly.Pools[0].ActAvg.ActMAvg)
                                }}})
                lg.AddStdAggs(itm, mode, times...)

                itm = lg.AddItem(&elog.Item{
                        Name:   clnm + "_ActMMax",
                        Type:   reflect.Float64,
                        FixMax: false,
                        Range:  minmax.F32{Max: 1},
                        Write: elog.WriteMap{
                                etime.Scope(mode, times[ntimes-1]): func(ctx *elog.Context) {
                                        ly := ctx.Layer(clnm).(*Layer)
                                        ctx.SetFloat32(ly.Pools[0].ActM.Max)
                                }}})
                lg.AddStdAggs(itm, mode, times...)

                itm = lg.AddItem(&elog.Item{
                        Name:  clnm + "_CosDiff",
                        Type:  reflect.Float64,
                        Range: minmax.F32{Max: 1},
                        Write: elog.WriteMap{
                                etime.Scope(etime.Train, times[ntimes-1]): func(ctx *elog.Context) {
                                        ly := ctx.Layer(clnm).(*Layer)
                                        ctx.SetFloat32(ly.CosDiff.Cos)
                                }}})
                lg.AddStdAggs(itm, mode, times...)
        }
}

func LogInputLayer(lg *elog.Logs, net *Network, mode etime.Modes) {
        // input layer average activity -- important for tuning
        layerNames := net.LayersByType(InputLayer)
        for _, lnm := range layerNames {
                clnm := lnm
                lg.AddItem(&elog.Item{
                        Name:   clnm + "_ActAvg",
                        Type:   reflect.Float64,
                        FixMax: true,
                        Range:  minmax.F32{Max: 1},
                        Write: elog.WriteMap{
                                etime.Scope(etime.Train, etime.Epoch): func(ctx *elog.Context) {
                                        ly := ctx.Layer(clnm).(*Layer)
                                        ctx.SetFloat32(ly.Pools[0].ActM.Max)
                                }}})
        }
}

// LogAddPCAItems adds PCA statistics to log for Hidden and Target layers
// across the given time levels, in higher to lower order, e.g., Run, Epoch, Trial
// These are useful for diagnosing the behavior of the network.
func LogAddPCAItems(lg *elog.Logs, net *Network, mode etime.Modes, times ...etime.Times) {
        ntimes := len(times)
        layers := net.LayersByType(SuperLayer, TargetLayer, CTLayer)
        for _, lnm := range layers {
                clnm := lnm
                cly := net.LayerByName(clnm)
                lg.AddItem(&elog.Item{
                        Name:      clnm + "_ActM",
                        Type:      reflect.Float64,
                        CellShape: cly.GetSampleShape().Sizes,
                        FixMax:    true,
                        Range:     minmax.F32{Max: 1},
                        Write: elog.WriteMap{
                                etime.Scope(etime.Analyze, times[ntimes-1]): func(ctx *elog.Context) {
                                        ctx.SetLayerSampleTensor(clnm, "ActM")
                                }, etime.Scope(etime.Test, times[ntimes-1]): func(ctx *elog.Context) {
                                        ctx.SetLayerSampleTensor(clnm, "ActM")
                                }}})
                itm := lg.AddItem(&elog.Item{
                        Name: clnm + "_PCA_NStrong",
                        Type: reflect.Float64,
                        Write: elog.WriteMap{
                                etime.Scope(etime.Train, times[ntimes-2]): func(ctx *elog.Context) {
                                        ctx.SetStatFloat(ctx.Item.Name)
                                }}})
                lg.AddStdAggs(itm, mode, times[:ntimes-1]...)

                itm = lg.AddItem(&elog.Item{
                        Name: clnm + "_PCA_Top5",
                        Type: reflect.Float64,
                        Write: elog.WriteMap{
                                etime.Scope(etime.Train, times[ntimes-2]): func(ctx *elog.Context) {
                                        ctx.SetStatFloat(ctx.Item.Name)
                                }}})
                lg.AddStdAggs(itm, mode, times[:ntimes-1]...)

                itm = lg.AddItem(&elog.Item{
                        Name: clnm + "_PCA_Next5",
                        Type: reflect.Float64,
                        Write: elog.WriteMap{
                                etime.Scope(etime.Train, times[ntimes-2]): func(ctx *elog.Context) {
                                        ctx.SetStatFloat(ctx.Item.Name)
                                }}})
                lg.AddStdAggs(itm, mode, times[:ntimes-1]...)

                itm = lg.AddItem(&elog.Item{
                        Name: clnm + "_PCA_Rest",
                        Type: reflect.Float64,
                        Write: elog.WriteMap{
                                etime.Scope(etime.Train, times[ntimes-2]): func(ctx *elog.Context) {
                                        ctx.SetStatFloat(ctx.Item.Name)
                                }}})
                lg.AddStdAggs(itm, mode, times[:ntimes-1]...)
        }
}

// LayerActsLogConfigMetaData configures meta data for LayerActs table
func LayerActsLogConfigMetaData(dt *table.Table) {
        dt.SetMetaData("read-only", "true")
        dt.SetMetaData("precision", strconv.Itoa(elog.LogPrec))
        dt.SetMetaData("Type", "Bar")
        dt.SetMetaData("XAxis", "Layer")
        dt.SetMetaData("XAxisRot", "45")
        dt.SetMetaData("Nominal:On", "+")
        dt.SetMetaData("Nominal:FixMin", "+")
        dt.SetMetaData("ActM:On", "+")
        dt.SetMetaData("ActM:FixMin", "+")
        dt.SetMetaData("ActM:Max", "1")
        dt.SetMetaData("ActP:FixMin", "+")
        dt.SetMetaData("ActP:Max", "1")
        dt.SetMetaData("MaxGeM:FixMin", "+")
        dt.SetMetaData("MaxGeM:FixMax", "+")
        dt.SetMetaData("MaxGeM:Max", "3")
        dt.SetMetaData("MaxGeP:FixMin", "+")
        dt.SetMetaData("MaxGeP:FixMax", "+")
        dt.SetMetaData("MaxGeP:Max", "3")
}

// LayerActsLogConfig configures Tables to record
// layer activity for tuning the network inhibition, nominal activity,
// relative scaling, etc. in elog.MiscTables:
// LayerActs is current, LayerActsRec is record over trials,
// LayerActsAvg is average of recorded trials.
func LayerActsLogConfig(net *Network, lg *elog.Logs) {
        dt := lg.MiscTable("LayerActs")
        dt.SetMetaData("name", "LayerActs")
        dt.SetMetaData("desc", "Layer Activations")
        LayerActsLogConfigMetaData(dt)
        dtRec := lg.MiscTable("LayerActsRec")
        dtRec.SetMetaData("name", "LayerActsRec")
        dtRec.SetMetaData("desc", "Layer Activations Recorded")
        LayerActsLogConfigMetaData(dtRec)
        dtAvg := lg.MiscTable("LayerActsAvg")
        dtAvg.SetMetaData("name", "LayerActsAvg")
        dtAvg.SetMetaData("desc", "Layer Activations Averaged")
        LayerActsLogConfigMetaData(dtAvg)
        dts := []*table.Table{dt, dtRec, dtAvg}
        for _, t := range dts {
                t.AddStringColumn("Layer")
                t.AddFloat64Column("Nominal")
                t.AddFloat64Column("ActM")
                t.AddFloat64Column("ActP")
        }
        nlay := len(net.Layers)
        dt.SetNumRows(nlay)
        dtRec.SetNumRows(0)
        dtAvg.SetNumRows(nlay)
        for li, ly := range net.Layers {
                dt.SetString("Layer", li, ly.Name)
                dt.SetFloat("Nominal", li, float64(ly.Inhib.ActAvg.Init))
                dtAvg.SetString("Layer", li, ly.Name)
        }
}

// LayerActsLog records layer activity for tuning the network
// inhibition, nominal activity, relative scaling, etc.
// if gui is non-nil, plot is updated.
func LayerActsLog(net *Network, lg *elog.Logs, di int, gui *egui.GUI) {
        dt := lg.MiscTable("LayerActs")
        dtRec := lg.MiscTable("LayerActsRec")
        for li, ly := range net.Layers {
                lpl := &ly.Pools[0]
                dt.SetFloat("Nominal", li, float64(ly.Inhib.ActAvg.Init))
                dt.SetFloat("ActM", li, float64(lpl.ActAvg.ActMAvg))
                dt.SetFloat("ActP", li, float64(lpl.ActAvg.ActPAvg))
                dtRec.SetNumRows(dtRec.Rows + 1)
                dtRec.SetString("Layer", li, ly.Name)
                dtRec.SetFloat("Nominal", li, float64(ly.Inhib.ActAvg.Init))
                dtRec.SetFloat("ActM", li, float64(lpl.ActAvg.ActMAvg))
                dtRec.SetFloat("ActP", li, float64(lpl.ActAvg.ActPAvg))
        }
        if gui != nil {
                gui.UpdatePlotScope(etime.ScopeKey("LayerActs"))
        }
}

// LayerActsLogAvg computes average of LayerActsRec record
// of layer activity for tuning the network
// inhibition, nominal activity, relative scaling, etc.
// if gui is non-nil, plot is updated.
// if recReset is true, reset the recorded data after computing average.
func LayerActsLogAvg(net *Network, lg *elog.Logs, gui *egui.GUI, recReset bool) {
        dtRec := lg.MiscTable("LayerActsRec")
        dtAvg := lg.MiscTable("LayerActsAvg")
        if dtRec.Rows == 0 {
                return
        }
        ix := table.NewIndexView(dtRec)
        spl := split.GroupBy(ix, "Layer")
        split.AggAllNumericColumns(spl, stats.Mean)
        ags := spl.AggsToTable(table.ColumnNameOnly)
        cols := []string{"Nominal", "ActM", "ActP", "MaxGeM", "MaxGeP"}
        for li, ly := range net.Layers {
                rw := errors.Log1(ags.RowsByString("Layer", ly.Name, table.Equals, table.UseCase))[0]
                for _, cn := range cols {
                        dtAvg.SetFloat(cn, li, ags.Float(cn, rw))
                }
        }
        if recReset {
                dtRec.SetNumRows(0)
        }
        if gui != nil {
                gui.UpdatePlotScope(etime.ScopeKey("LayerActsAvg"))
        }
}

// LayerActsLogRecReset resets the recorded LayerActsRec data
// used for computing averages
func LayerActsLogRecReset(lg *elog.Logs) {
        dtRec := lg.MiscTable("LayerActsRec")
        dtRec.SetNumRows(0)
}

// LayerActsLogConfigGUI configures GUI for LayerActsLog Plot and LayerActs Avg Plot
func LayerActsLogConfigGUI(lg *elog.Logs, gui *egui.GUI) {
        pt, _ := gui.Tabs.NewTab("LayerActs Plot")
        plt := plotcore.NewPlotEditor(pt)
        gui.Plots["LayerActs"] = plt
        plt.SetTable(lg.MiscTables["LayerActs"])

        pt, _ = gui.Tabs.NewTab("LayerActs Avg Plot")
        plt = plotcore.NewPlotEditor(pt)
        gui.Plots["LayerActsAvg"] = plt
        plt.SetTable(lg.MiscTables["LayerActsAvg"])
}

// Copyright (c) 2022, The Emergent Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package leabra

import (
        "github.com/emer/emergent/v2/egui"
        "github.com/emer/emergent/v2/elog"
        "github.com/emer/emergent/v2/etime"
        "github.com/emer/emergent/v2/looper"
        "github.com/emer/emergent/v2/netview"
)

// LooperStdPhases adds the minus and plus phases of the alpha cycle,
// along with embedded beta phases which just record St1 and St2 activity in this case.
// plusStart is start of plus phase, typically 75,
// and plusEnd is end of plus phase, typically 99
// resets the state at start of trial.
// Can pass a trial-level time scale to use instead of the default etime.Trial
func LooperStdPhases(ls *looper.Stacks, ctx *Context, net *Network, plusStart, plusEnd int, trial ...etime.Times) {
        trl := etime.Trial
        if len(trial) > 0 {
                trl = trial[0]
        }
        ls.AddEventAllModes(etime.Cycle, "MinusPhase:Start", 0, func() {
                ctx.PlusPhase = false
        })
        ls.AddEventAllModes(etime.Cycle, "Quarter1", 25, func() {
                net.QuarterFinal(ctx)
                ctx.QuarterInc()
        })
        ls.AddEventAllModes(etime.Cycle, "Quarter2", 50, func() {
                net.QuarterFinal(ctx)
                ctx.QuarterInc()
        })
        ls.AddEventAllModes(etime.Cycle, "MinusPhase:End", plusStart, func() {
                net.QuarterFinal(ctx)
                ctx.QuarterInc()
        })
        ls.AddEventAllModes(etime.Cycle, "PlusPhase:Start", plusStart, func() {
                ctx.PlusPhase = true
        })

        for m, stack := range ls.Stacks {
                stack.Loops[trl].OnStart.Add("AlphaCycInit", func() {
                        net.AlphaCycInit(m == etime.Train)
                        ctx.AlphaCycStart()
                })
                stack.Loops[trl].OnEnd.Add("PlusPhase:End", func() {
                        net.QuarterFinal(ctx)
                })
        }
}

// LooperSimCycleAndLearn adds Cycle and DWt, WtFromDWt functions to looper
// for given network, ctx, and netview update manager
// Can pass a trial-level time scale to use instead of the default etime.Trial
func LooperSimCycleAndLearn(ls *looper.Stacks, net *Network, ctx *Context, viewupdt *netview.ViewUpdate, trial ...etime.Times) {
        trl := etime.Trial
        if len(trial) > 0 {
                trl = trial[0]
        }
        for m := range ls.Stacks {
                ls.Stacks[m].Loops[etime.Cycle].OnStart.Add("Cycle", func() {
                        net.Cycle(ctx)
                        ctx.CycleInc()
                })
        }
        ttrl := ls.Loop(etime.Train, trl)
        if ttrl != nil {
                ttrl.OnEnd.Add("UpdateWeights", func() {
                        net.DWt()
                        if viewupdt.IsViewingSynapse() {
                                viewupdt.RecordSyns() // note: critical to update weights here so DWt is visible
                        }
                        net.WtFromDWt()
                })
        }

        // Set variables on ss that are referenced elsewhere, such as ApplyInputs.
        for m, loops := range ls.Stacks {
                for _, loop := range loops.Loops {
                        loop.OnStart.Add("SetCtxMode", func() {
                                ctx.Mode = m.(etime.Modes)
                        })
                }
        }
}

// LooperResetLogBelow adds a function in OnStart to all stacks and loops
// to reset the log at the level below each loop -- this is good default behavior.
// Exceptions can be passed to exclude specific levels -- e.g., if except is Epoch
// then Epoch does not reset the log below it
func LooperResetLogBelow(ls *looper.Stacks, logs *elog.Logs, except ...etime.Times) {
        for m, stack := range ls.Stacks {
                for t, loop := range stack.Loops {
                        curTime := t
                        isExcept := false
                        for _, ex := range except {
                                if curTime == ex {
                                        isExcept = true
                                        break
                                }
                        }
                        if below := stack.TimeBelow(curTime); !isExcept && below != etime.NoTime {
                                loop.OnStart.Add("ResetLog"+below.String(), func() {
                                        logs.ResetLog(m.(etime.Modes), below.(etime.Times))
                                })
                        }
                }
        }
}

// LooperUpdateNetView adds netview update calls at each time level
func LooperUpdateNetView(ls *looper.Stacks, viewupdt *netview.ViewUpdate, net *Network, ctrUpdateFunc func(tm etime.Times)) {
        for m, stack := range ls.Stacks {
                for t, loop := range stack.Loops {
                        curTime := t.(etime.Times)
                        if curTime != etime.Cycle {
                                loop.OnEnd.Add("GUI:UpdateNetView", func() {
                                        ctrUpdateFunc(curTime)
                                        viewupdt.Testing = m == etime.Test
                                        viewupdt.UpdateTime(curTime)
                                })
                        }
                }
                cycLoop := ls.Loop(m, etime.Cycle)
                cycLoop.OnEnd.Add("GUI:UpdateNetView", func() {
                        cyc := cycLoop.Counter.Cur
                        ctrUpdateFunc(etime.Cycle)
                        viewupdt.Testing = m == etime.Test
                        viewupdt.UpdateCycle(cyc)
                })
        }
}

// LooperUpdatePlots adds plot update calls at each time level
func LooperUpdatePlots(ls *looper.Stacks, gui *egui.GUI) {
        for m, stack := range ls.Stacks {
                for t, loop := range stack.Loops {
                        curTime := t.(etime.Times)
                        curLoop := loop
                        if curTime == etime.Cycle {
                                curLoop.OnEnd.Add("GUI:UpdatePlot", func() {
                                        cyc := curLoop.Counter.Cur
                                        gui.GoUpdateCyclePlot(m.(etime.Modes), cyc)
                                })
                        } else {
                                curLoop.OnEnd.Add("GUI:UpdatePlot", func() {
                                        gui.GoUpdatePlot(m.(etime.Modes), curTime)
                                })
                        }
                }
        }
}

// Copyright (c) 2019, The Emergent Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package leabra

import (
        "fmt"
        "strings"
        "unsafe"

        "cogentcore.org/core/base/datasize"
        "github.com/emer/emergent/v2/paths"
        "github.com/emer/etensor/tensor"
)

///////////////////////////////////////////////////////////////////////////
//  Primary Algorithmic interface.
//
//  The following methods constitute the primary user-called API during
//  AlphaCyc method to compute one complete algorithmic alpha cycle update.

// AlphaCycInit handles all initialization at start of new input pattern.
// Should already have presented the external input to the network at this point.
// If updtActAvg is true, this includes updating the running-average
// activations for each layer / pool, and the AvgL running average used
// in BCM Hebbian learning.
// The input scaling is updated  based on the layer-level running average acts,
// and this can then change the behavior of the network,
// so if you want 100% repeatable testing results, set this to false to
// keep the existing scaling factors (e.g., can pass a train bool to
// only update during training).
// This flag also affects the AvgL learning threshold.
func (nt *Network) AlphaCycInit(updtActAvg bool) {
        for _, ly := range nt.Layers {
                if ly.Off {
                        continue
                }
                ly.AlphaCycInit(updtActAvg)
        }
}

// Cycle runs one cycle of activation updating:
// * Sends Ge increments from sending to receiving layers
// * Average and Max Ge stats
// * Inhibition based on Ge stats and Act Stats (computed at end of Cycle)
// * Activation from Ge, Gi, and Gl
// * Average and Max Act stats
// This basic version doesn't use the time info, but more specialized types do, and we
// want to keep a consistent API for end-user code.
func (nt *Network) Cycle(ctx *Context) {
        nt.SendGDelta(ctx) // also does integ
        nt.AvgMaxGe(ctx)
        nt.InhibFromGeAct(ctx)
        nt.ActFromG(ctx)
        nt.AvgMaxAct(ctx)
        nt.CyclePost(ctx)  // general post cycle actions.
        nt.RecGateAct(ctx) // Record activation state at time of gating (in ActG neuron var)
}

//////////////////////////////////////////////////////////////////////////////////////
//  Act methods

// SendGeDelta sends change in activation since last sent, if above thresholds
// and integrates sent deltas into GeRaw and time-integrated Ge values
func (nt *Network) SendGDelta(ctx *Context) {
        for _, ly := range nt.Layers {
                if ly.Off {
                        continue
                }
                ly.SendGDelta(ctx)
        }
        for _, ly := range nt.Layers {
                if ly.Off {
                        continue
                }
                ly.GFromInc(ctx)
        }
}

// AvgMaxGe computes the average and max Ge stats, used in inhibition
func (nt *Network) AvgMaxGe(ctx *Context) {
        for _, ly := range nt.Layers {
                if ly.Off {
                        continue
                }
                ly.AvgMaxGe(ctx)
        }
}

// InhibiFromGeAct computes inhibition Gi from Ge and Act stats within relevant Pools
func (nt *Network) InhibFromGeAct(ctx *Context) {
        for _, ly := range nt.Layers {
                if ly.Off {
                        continue
                }
                ly.InhibFromGeAct(ctx)
        }
}

// ActFromG computes rate-code activation from Ge, Gi, Gl conductances
func (nt *Network) ActFromG(ctx *Context) {
        for _, ly := range nt.Layers {
                if ly.Off {
                        continue
                }
                ly.ActFromG(ctx)
        }
}

// AvgMaxGe computes the average and max Ge stats, used in inhibition
func (nt *Network) AvgMaxAct(ctx *Context) {
        for _, ly := range nt.Layers {
                if ly.Off {
                        continue
                }
                ly.AvgMaxAct(ctx)
        }
}

// CyclePost is called at end of Cycle, for misc updates after new Act
// value has been computed.
// SuperLayer computes Burst activity.
// GateLayer (GPiThal) computes gating, sends to other layers.
func (nt *Network) CyclePost(ctx *Context) {
        for _, ly := range nt.Layers {
                if ly.Off {
                        continue
                }
                ly.CyclePost(ctx)
        }
}

// QuarterFinal does updating after end of a quarter, for first 2
func (nt *Network) QuarterFinal(ctx *Context) {
        for _, ly := range nt.Layers {
                if ly.Off {
                        continue
                }
                ly.QuarterFinal(ctx) // also does SendCtxtGe
        }
        for _, ly := range nt.Layers {
                if ly.Off {
                        continue
                }
                ly.CtxtFromGe(ctx)
        }
}

// MinusPhase is called at the end of the minus phase (quarter 3), to record state.
func (nt *Network) MinusPhase(ctx *Context) {
        for _, ly := range nt.Layers {
                if ly.Off {
                        continue
                }
                ly.MinusPhase(ctx)
        }
}

// PlusPhase is called at the end of the plus phase (quarter 4), to record state.
func (nt *Network) PlusPhase(ctx *Context) {
        for _, ly := range nt.Layers {
                if ly.Off {
                        continue
                }
                ly.PlusPhase(ctx)
        }
}

//////////////////////////////////////////////////////////////////////////////////////
//  Learn methods

// DWt computes the weight change (learning) based on current
// running-average activation values
func (nt *Network) DWt() {
        for _, ly := range nt.Layers {
                if ly.Off {
                        continue
                }
                ly.DWt()
        }
}

// WtFromDWt updates the weights from delta-weight changes.
// Also calls WtBalFromWt every WtBalInterval times
func (nt *Network) WtFromDWt() {
        for _, ly := range nt.Layers {
                if ly.Off {
                        continue
                }
                ly.WtFromDWt()
        }

        nt.WtBalCtr++
        if nt.WtBalCtr >= nt.WtBalInterval {
                nt.WtBalCtr = 0
                for _, ly := range nt.Layers {
                        if ly.Off {
                                continue
                        }

                        ly.WtBalFromWt()
                }
        }
}

// LrateMult sets the new Lrate parameter for Paths to LrateInit * mult.
// Useful for implementing learning rate schedules.
func (nt *Network) LrateMult(mult float32) {
        for _, ly := range nt.Layers {
                // if ly.Off { // keep all sync'd
                //         continue
                // }
                ly.LrateMult(mult)
        }
}

//////////////////////////////////////////////////////////////////////////////////////
//  Init methods

// InitWeights initializes synaptic weights and all other
// associated long-term state variables including running-average
// state values (e.g., layer running average activations etc).
func (nt *Network) InitWeights() {
        nt.WtBalCtr = 0
        for _, ly := range nt.Layers {
                if ly.Off {
                        continue
                }
                ly.InitWeights()
        }
        for _, ly := range nt.Layers {
                if ly.Off {
                        continue
                }
                ly.InitWtSym()
        }
}

// InitTopoScales initializes synapse-specific scale parameters from
// path types that support them, with flags set to support it,
// includes: paths.PoolTile paths.Circle.
// call before InitWeights if using Topo wts.
func (nt *Network) InitTopoScales() {
        scales := &tensor.Float32{}
        for _, ly := range nt.Layers {
                if ly.Off {
                        continue
                }
                rpjn := ly.RecvPaths
                for _, pt := range rpjn {
                        if pt.Off {
                                continue
                        }
                        pat := pt.Pattern
                        switch ptn := pat.(type) {
                        case *paths.PoolTile:
                                if !ptn.HasTopoWeights() {
                                        continue
                                }
                                slay := pt.Send
                                ptn.TopoWeights(&slay.Shape, &ly.Shape, scales)
                                pt.SetScalesRPool(scales)
                        case *paths.Circle:
                                if !ptn.TopoWeights {
                                        continue
                                }
                                pt.SetScalesFunc(ptn.GaussWts)
                        }
                }
        }
}

// DecayState decays activation state by given proportion
// e.g., 1 = decay completely, and 0 = decay not at all
// This is called automatically in AlphaCycInit, but is avail
// here for ad-hoc decay cases.
func (nt *Network) DecayState(decay float32) {
        for _, ly := range nt.Layers {
                if ly.Off {
                        continue
                }
                ly.DecayState(decay)
        }
}

// InitActs fully initializes activation state -- not automatically called
func (nt *Network) InitActs() {
        for _, ly := range nt.Layers {
                if ly.Off {
                        continue
                }
                ly.InitActs()
        }
}

// InitExt initializes external input state.
// call prior to applying external inputs to layers.
func (nt *Network) InitExt() {
        for _, ly := range nt.Layers {
                if ly.Off {
                        continue
                }
                ly.InitExt()
        }
}

// UpdateExtFlags updates the neuron flags for external input
// based on current layer Type field.
// call this if the Type has changed since the last
// ApplyExt* method call.
func (nt *Network) UpdateExtFlags() {
        for _, ly := range nt.Layers {
                if ly.Off {
                        continue
                }
                ly.UpdateExtFlags()
        }
}

// InitGinc initializes the Ge excitatory and Gi inhibitory
// conductance accumulation states including ActSent and G*Raw values.
// called at start of trial always (at layer level), and can be
// called optionally when delta-based Ge computation needs
// to be updated (e.g., weights might have changed strength).
func (nt *Network) InitGInc() {
        for _, ly := range nt.Layers {
                if ly.Off {
                        continue
                }
                ly.InitGInc()
        }
}

// GScaleFromAvgAct computes the scaling factor for synaptic input conductances G,
// based on sending layer average activation.
// This attempts to automatically adjust for overall differences in raw activity
// coming into the units to achieve a general target of around .5 to 1
// for the integrated Ge value.
// This is automatically done during AlphaCycInit, but if scaling parameters are
// changed at any point thereafter during AlphaCyc, this must be called.
func (nt *Network) GScaleFromAvgAct() {
        for _, ly := range nt.Layers {
                if ly.Off {
                        continue
                }
                ly.GScaleFromAvgAct()
        }
}

//////////////////////////////////////////////////////////////////////////////////////
//  Lesion methods

// LayersSetOff sets the Off flag for all layers to given setting
func (nt *Network) LayersSetOff(off bool) {
        for _, ly := range nt.Layers {
                ly.Off = off
        }
}

// UnLesionNeurons unlesions neurons in all layers in the network.
// Provides a clean starting point for subsequent lesion experiments.
func (nt *Network) UnLesionNeurons() {
        for _, ly := range nt.Layers {
                // if ly.Off { // keep all sync'd
                //         continue
                // }
                ly.UnLesionNeurons()
        }
}

//////////////////////////////////////////////////////////////////////////////////////
//  Methods used in MPI computation, which don't depend on MPI specifically

// CollectDWts writes all of the synaptic DWt values to given dwts slice
// which is pre-allocated to given nwts size if dwts is nil,
// in which case the method returns true so that the actual length of
// dwts can be passed next time around.
// Used for MPI sharing of weight changes across processors.
func (nt *Network) CollectDWts(dwts *[]float32, nwts int) bool {
        idx := 0
        made := false
        if *dwts == nil {
                // todo: if nil, compute right size right away
                *dwts = make([]float32, 0, nwts)
                made = true
        }
        for _, ly := range nt.Layers {
                for _, pt := range ly.SendPaths {
                        ns := len(pt.Syns)
                        nsz := idx + ns
                        if len(*dwts) < nsz {
                                *dwts = append(*dwts, make([]float32, nsz-len(*dwts))...)
                        }
                        for j := range pt.Syns {
                                sy := &(pt.Syns[j])
                                (*dwts)[idx+j] = sy.DWt
                        }
                        idx += ns
                }
        }
        return made
}

// SetDWts sets the DWt weight changes from given array of floats,
// which must be correct size.
func (nt *Network) SetDWts(dwts []float32) {
        idx := 0
        for _, ly := range nt.Layers {
                for _, pt := range ly.SendPaths {
                        ns := len(pt.Syns)
                        for j := range pt.Syns {
                                sy := &(pt.Syns[j])
                                sy.DWt = dwts[idx+j]
                        }
                        idx += ns
                }
        }
}

//////////////////////////////////////////////////////////////////////////////////////
//  Misc Reports

// SizeReport returns a string reporting the size of
// each layer and pathway in the network, and total memory footprint.
func (nt *Network) SizeReport() string {
        var b strings.Builder
        neur := 0
        neurMem := 0
        syn := 0
        synMem := 0
        for _, ly := range nt.Layers {
                nn := len(ly.Neurons)
                nmem := nn * int(unsafe.Sizeof(Neuron{}))
                neur += nn
                neurMem += nmem
                fmt.Fprintf(&b, "%14s:\t Neurons: %d\t NeurMem: %v \t Sends To:\n", ly.Name, nn, (datasize.Size)(nmem).String())
                for _, pt := range ly.SendPaths {
                        ns := len(pt.Syns)
                        syn += ns
                        pmem := ns*int(unsafe.Sizeof(Synapse{})) + len(pt.GInc)*4 + len(pt.WbRecv)*int(unsafe.Sizeof(WtBalRecvPath{}))
                        synMem += pmem
                        fmt.Fprintf(&b, "\t%14s:\t Syns: %d\t SynnMem: %v\n", pt.Recv.Name, ns, (datasize.Size)(pmem).String())
                }
        }
        fmt.Fprintf(&b, "\n\n%14s:\t Neurons: %d\t NeurMem: %v \t Syns: %d \t SynMem: %v\n", nt.Name, neur, (datasize.Size)(neurMem).String(), syn, (datasize.Size)(synMem).String())
        return b.String()
}

//////////////////////////////////////////////////////////////////////////////////////
//  Network props for gui

// TODO(v2): props

// var NetworkProps = tree.Props{
//         "ToolBar": tree.PropSlice{
//                 {"SaveWeightsJSON", tree.Props{
//                         "label": "Save Wts...",
//                         "icon":  "file-save",
//                         "desc":  "Save json-formatted weights",
//                         "Args": tree.PropSlice{
//                                 {"Weights File Name", tree.Props{
//                                         "default-field": "WtsFile",
//                                         "ext":           ".wts,.wts.gz",
//                                 }},
//                         },
//                 }},
//                 {"OpenWeightsJSON", tree.Props{
//                         "label": "Open Wts...",
//                         "icon":  "file-open",
//                         "desc":  "Open json-formatted weights",
//                         "Args": tree.PropSlice{
//                                 {"Weights File Name", tree.Props{
//                                         "default-field": "WtsFile",
//                                         "ext":           ".wts,.wts.gz",
//                                 }},
//                         },
//                 }},
//                 {"sep-file", tree.BlankProp{}},
//                 {"Build", tree.Props{
//                         "icon": "update",
//                         "desc": "build the network's neurons and synapses according to current params",
//                 }},
//                 {"InitWeights", tree.Props{
//                         "icon": "update",
//                         "desc": "initialize the network weight values according to path parameters",
//                 }},
//                 {"InitActs", tree.Props{
//                         "icon": "update",
//                         "desc": "initialize the network activation values",
//                 }},
//                 {"sep-act", tree.BlankProp{}},
//                 {"AddLayer", tree.Props{
//                         "label": "Add Layer...",
//                         "icon":  "new",
//                         "desc":  "add a new layer to network",
//                         "Args": tree.PropSlice{
//                                 {"Layer Name", tree.Props{}},
//                                 {"Layer Shape", tree.Props{
//                                         "desc": "shape of layer, typically 2D (Y, X) or 4D (Pools Y, Pools X, Units Y, Units X)",
//                                 }},
//                                 {"Layer Type", tree.Props{
//                                         "desc": "type of layer -- used for determining how inputs are applied",
//                                 }},
//                         },
//                 }},
//                 {"ConnectLayerNames", tree.Props{
//                         "label": "Connect Layers...",
//                         "icon":  "new",
//                         "desc":  "add a new connection between layers in the network",
//                         "Args": tree.PropSlice{
//                                 {"Send Layer Name", tree.Props{}},
//                                 {"Recv Layer Name", tree.Props{}},
//                                 {"Pattern", tree.Props{
//                                         "desc": "pattern to connect with",
//                                 }},
//                                 {"Path Type", tree.Props{
//                                         "desc": "type of pathway -- direction, or other more specialized factors",
//                                 }},
//                         },
//                 }},
//                 {"AllWtScales", tree.Props{
//                         "icon":        "file-sheet",
//                         "desc":        "AllWtScales returns a listing of all WtScale parameters in the Network in all Layers, Recv pathways.  These are among the most important and numerous of parameters (in larger networks) -- this helps keep track of what they all are set to.",
//                         "show-return": true,
//                 }},
//         },
// }

// Copyright (c) 2019, The Emergent Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package leabra

//go:generate core generate -add-types

import (
        "errors"
        "fmt"
        "log"
        "os"
        "path/filepath"
        "time"

        "cogentcore.org/core/core"
        "github.com/emer/emergent/v2/econfig"
        "github.com/emer/emergent/v2/emer"
        "github.com/emer/emergent/v2/params"
        "github.com/emer/emergent/v2/paths"
)

// leabra.Network implements the Leabra algorithm, managing the Layers.
type Network struct {
        emer.NetworkBase

        // list of layers
        Layers []*Layer

        // number of parallel threads (go routines) to use.
        NThreads int `edit:"-"`

        // how frequently to update the weight balance average
        // weight factor -- relatively expensive.
        WtBalInterval int `default:"10"`

        // counter for how long it has been since last WtBal.
        WtBalCtr int `edit:"-"`
}

func (nt *Network) NumLayers() int               { return len(nt.Layers) }
func (nt *Network) EmerLayer(idx int) emer.Layer { return nt.Layers[idx] }
func (nt *Network) MaxParallelData() int         { return 1 }
func (nt *Network) NParallelData() int           { return 1 }

// NewNetwork returns a new leabra Network
func NewNetwork(name string) *Network {
        net := &Network{}
        emer.InitNetwork(net, name)
        net.NThreads = 1
        return net
}

// LayerByName returns a layer by looking it up by name in the layer map
// (nil if not found).
func (nt *Network) LayerByName(name string) *Layer {
        ely, _ := nt.EmerLayerByName(name)
        return ely.(*Layer)
}

// LayersByType returns a list of layer names by given layer type(s).
func (nt *Network) LayersByType(layType ...LayerTypes) []string {
        var nms []string
        for _, tp := range layType {
                nm := tp.String()
                nms = append(nms, nm)
        }
        return nt.LayersByClass(nms...)
}

// KeyLayerParams returns a listing for all layers in the network,
// of the most important layer-level params (specific to each algorithm).
func (nt *Network) KeyLayerParams() string {
        return nt.AllLayerInhibs()
}

// KeyPathParams returns a listing for all Recv pathways in the network,
// of the most important pathway-level params (specific to each algorithm).
func (nt *Network) KeyPathParams() string {
        return nt.AllPathScales()
}

// SaveParamsSnapshot saves various views of current parameters
// to either `params_good` if good = true (for current good reference params)
// or `params_2006_01_02` (year, month, day) datestamp,
// providing a snapshot of the simulation params for easy diffs and later reference.
// Also saves current Config and Params state.
func (nt *Network) SaveParamsSnapshot(pars *params.Sets, cfg any, good bool) error {
        date := time.Now().Format("2006_01_02")
        if good {
                date = "good"
        }
        dir := "params_" + date
        err := os.Mkdir(dir, 0775)
        if err != nil {
                log.Println(err) // notify but OK if it exists
        }
        econfig.Save(cfg, filepath.Join(dir, "config.toml"))
        pars.SaveTOML(core.Filename(filepath.Join(dir, "params.toml")))
        nt.SaveAllParams(core.Filename(filepath.Join(dir, "params_all.txt")))
        nt.SaveNonDefaultParams(core.Filename(filepath.Join(dir, "params_nondef.txt")))
        nt.SaveAllLayerInhibs(core.Filename(filepath.Join(dir, "params_layers.txt")))
        nt.SaveAllPathScales(core.Filename(filepath.Join(dir, "params_paths.txt")))
        return nil
}

// SaveAllLayerInhibs saves list of all layer Inhibition parameters to given file
func (nt *Network) SaveAllLayerInhibs(filename core.Filename) error {
        str := nt.AllLayerInhibs()
        err := os.WriteFile(string(filename), []byte(str), 0666)
        if err != nil {
                log.Println(err)
        }
        return err
}

// SavePathScales saves a listing of all PathScale parameters in the Network
// in all Layers, Recv pathways.  These are among the most important
// and numerous of parameters (in larger networks) -- this helps keep
// track of what they all are set to.
func (nt *Network) SaveAllPathScales(filename core.Filename) error {
        str := nt.AllPathScales()
        err := os.WriteFile(string(filename), []byte(str), 0666)
        if err != nil {
                log.Println(err)
        }
        return err
}

// AllLayerInhibs returns a listing of all Layer Inhibition parameters in the Network
func (nt *Network) AllLayerInhibs() string {
        str := ""
        for _, ly := range nt.Layers {
                if ly.Off {
                        continue
                }
                ph := ly.ParamsHistory.ParamsHistory()
                lh := ph["Layer.Inhib.ActAvg.Init"]
                if lh != "" {
                        lh = "Params: " + lh
                }
                str += fmt.Sprintf("%15s\t\tNominal:\t%6.2f\t%s\n", ly.Name, ly.Inhib.ActAvg.Init, lh)
                if ly.Inhib.Layer.On {
                        lh := ph["Layer.Inhib.Layer.Gi"]
                        if lh != "" {
                                lh = "Params: " + lh
                        }
                        str += fmt.Sprintf("\t\t\t\t\t\tLayer.Gi:\t%6.2f\t%s\n", ly.Inhib.Layer.Gi, lh)
                }
                if ly.Inhib.Pool.On {
                        lh := ph["Layer.Inhib.Pool.Gi"]
                        if lh != "" {
                                lh = "Params: " + lh
                        }
                        str += fmt.Sprintf("\t\t\t\t\t\tPool.Gi: \t%6.2f\t%s\n", ly.Inhib.Pool.Gi, lh)
                }
                str += fmt.Sprintf("\n")
        }
        return str
}

// AllPathScales returns a listing of all WtScale parameters in the Network
// in all Layers, Recv pathways.  These are among the most important
// and numerous of parameters (in larger networks) -- this helps keep
// track of what they all are set to.
func (nt *Network) AllPathScales() string {
        str := ""
        for _, ly := range nt.Layers {
                if ly.Off {
                        continue
                }
                str += "\nLayer: " + ly.Name + "\n"
                for _, pt := range ly.RecvPaths {
                        if pt.Off {
                                continue
                        }
                        str += fmt.Sprintf("\t%23s\t\tAbs:\t%g\tRel:\t%g\n", pt.Name, pt.WtScale.Abs, pt.WtScale.Rel)
                }
        }
        return str
}

// Defaults sets all the default parameters for all layers and pathways
func (nt *Network) Defaults() {
        nt.WtBalInterval = 10
        nt.WtBalCtr = 0
        for li, ly := range nt.Layers {
                ly.Defaults()
                ly.Index = li
        }
}

// UpdateParams updates all the derived parameters if any have changed, for all layers
// and pathways
func (nt *Network) UpdateParams() {
        for _, ly := range nt.Layers {
                ly.UpdateParams()
        }
}

// UnitVarNames returns a list of variable names available on the units in this network.
// Not all layers need to support all variables, but must safely return 0's for
// unsupported ones.  The order of this list determines NetView variable display order.
// This is typically a global list so do not modify!
func (nt *Network) UnitVarNames() []string {
        return NeuronVars
}

// UnitVarProps returns properties for variables
func (nt *Network) UnitVarProps() map[string]string {
        return NeuronVarProps
}

func (nt *Network) VarCategories() []emer.VarCategory {
        return VarCategories
}

// SynVarNames returns the names of all the variables on the synapses in this network.
// Not all pathways need to support all variables, but must safely return 0's for
// unsupported ones.  The order of this list determines NetView variable display order.
// This is typically a global list so do not modify!
func (nt *Network) SynVarNames() []string {
        return SynapseVars
}

// SynVarProps returns properties for variables
func (nt *Network) SynVarProps() map[string]string {
        return SynapseVarProps
}

// AddLayerInit is implementation routine that takes a given layer and
// adds it to the network, and initializes and configures it properly.
func (nt *Network) AddLayerInit(ly *Layer, name string, shape []int, typ LayerTypes) {
        if nt.EmerNetwork == nil {
                log.Printf("Network EmerNetwork is nil: MUST call emer.InitNetwork on network, passing a pointer to the network to initialize properly!")
                return
        }
        emer.InitLayer(ly, name)
        ly.SetShape(shape)
        ly.Type = typ
        nt.Layers = append(nt.Layers, ly)
        nt.UpdateLayerMaps()
}

// AddLayer adds a new layer with given name and shape to the network.
// 2D and 4D layer shapes are generally preferred but not essential -- see
// AddLayer2D and 4D for convenience methods for those.  4D layers enable
// pool (unit-group) level inhibition in Leabra networks, for example.
// shape is in row-major format with outer-most dimensions first:
// e.g., 4D 3, 2, 4, 5 = 3 rows (Y) of 2 cols (X) of pools, with each unit
// group having 4 rows (Y) of 5 (X) units.
func (nt *Network) AddLayer(name string, shape []int, typ LayerTypes) *Layer {
        ly := &Layer{} // essential to use EmerNet interface here!
        nt.AddLayerInit(ly, name, shape, typ)
        return ly
}

// AddLayer2D adds a new layer with given name and 2D shape to the network.
// 2D and 4D layer shapes are generally preferred but not essential.
func (nt *Network) AddLayer2D(name string, shapeY, shapeX int, typ LayerTypes) *Layer {
        return nt.AddLayer(name, []int{shapeY, shapeX}, typ)
}

// AddLayer4D adds a new layer with given name and 4D shape to the network.
// 4D layers enable pool (unit-group) level inhibition in Leabra networks, for example.
// shape is in row-major format with outer-most dimensions first:
// e.g., 4D 3, 2, 4, 5 = 3 rows (Y) of 2 cols (X) of pools, with each pool
// having 4 rows (Y) of 5 (X) neurons.
func (nt *Network) AddLayer4D(name string, nPoolsY, nPoolsX, nNeurY, nNeurX int, typ LayerTypes) *Layer {
        return nt.AddLayer(name, []int{nPoolsY, nPoolsX, nNeurY, nNeurX}, typ)
}

// ConnectLayerNames establishes a pathway between two layers, referenced by name
// adding to the recv and send pathway lists on each side of the connection.
// Returns error if not successful.
// Does not yet actually connect the units within the layers -- that requires Build.
func (nt *Network) ConnectLayerNames(send, recv string, pat paths.Pattern, typ PathTypes) (rlay, slay *Layer, pt *Path, err error) {
        rlay = nt.LayerByName(recv)
        if rlay == nil {
                return
        }
        slay = nt.LayerByName(send)
        if slay == nil {
                return
        }
        pt = nt.ConnectLayers(slay, rlay, pat, typ)
        return
}

// ConnectLayers establishes a pathway between two layers,
// adding to the recv and send pathway lists on each side of the connection.
// Does not yet actually connect the units within the layers -- that
// requires Build.
func (nt *Network) ConnectLayers(send, recv *Layer, pat paths.Pattern, typ PathTypes) *Path {
        pt := &Path{}
        emer.InitPath(pt)
        pt.Connect(send, recv, pat, typ)
        recv.RecvPaths = append(recv.RecvPaths, pt)
        send.SendPaths = append(send.SendPaths, pt)
        return pt
}

// BidirConnectLayerNames establishes bidirectional pathways between two layers,
// referenced by name, with low = the lower layer that sends a Forward pathway
// to the high layer, and receives a Back pathway in the opposite direction.
// Returns error if not successful.
// Does not yet actually connect the units within the layers -- that requires Build.
func (nt *Network) BidirConnectLayerNames(low, high string, pat paths.Pattern) (lowlay, highlay *Layer, fwdpj, backpj *Path, err error) {
        lowlay = nt.LayerByName(low)
        if lowlay == nil {
                return
        }
        highlay = nt.LayerByName(high)
        if highlay == nil {
                return
        }
        fwdpj = nt.ConnectLayers(lowlay, highlay, pat, ForwardPath)
        backpj = nt.ConnectLayers(highlay, lowlay, pat, BackPath)
        return
}

// BidirConnectLayers establishes bidirectional pathways between two layers,
// with low = lower layer that sends a Forward pathway to the high layer,
// and receives a Back pathway in the opposite direction.
// Does not yet actually connect the units within the layers -- that
// requires Build.
func (nt *Network) BidirConnectLayers(low, high *Layer, pat paths.Pattern) (fwdpj, backpj *Path) {
        fwdpj = nt.ConnectLayers(low, high, pat, ForwardPath)
        backpj = nt.ConnectLayers(high, low, pat, BackPath)
        return
}

// LateralConnectLayer establishes a self-pathway within given layer.
// Does not yet actually connect the units within the layers -- that
// requires Build.
func (nt *Network) LateralConnectLayer(lay *Layer, pat paths.Pattern) *Path {
        return nt.ConnectLayers(lay, lay, pat, LateralPath)
}

// Build constructs the layer and pathway state based on the layer shapes
// and patterns of interconnectivity
func (nt *Network) Build() error {
        nt.MakeLayerMaps()
        var errs []error
        for li, ly := range nt.Layers {
                ly.Index = li
                ly.Network = nt
                if ly.Off {
                        continue
                }
                err := ly.Build()
                if err != nil {
                        errs = append(errs, err)
                }
        }
        nt.LayoutLayers()
        return errors.Join(errs...)
}

// VarRange returns the min / max values for given variable
// todo: support r. s. pathway values
func (nt *Network) VarRange(varNm string) (min, max float32, err error) {
        first := true
        for _, ly := range nt.Layers {
                lmin, lmax, lerr := ly.VarRange(varNm)
                if lerr != nil {
                        err = lerr
                        return
                }
                if first {
                        min = lmin
                        max = lmax
                        continue
                }
                if lmin < min {
                        min = lmin
                }
                if lmax > max {
                        max = lmax
                }
        }
        return
}

// Copyright (c) 2024, The Emergent Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package leabra

import (
        "errors"
        "fmt"
)

// LayerNames is a list of layer names, with methods to add and validate.
type LayerNames []string

// Add adds given layer name(s) to list.
func (ln *LayerNames) Add(laynm ...string) {
        *ln = append(*ln, laynm...)
}

// AddAllBut adds all layers in network except those in exclude list.
func (ln *LayerNames) AddAllBut(net *Network, excl ...string) {
        *ln = nil
        for _, l := range net.Layers {
                lnm := l.Name
                exl := false
                for _, ex := range excl {
                        if lnm == ex {
                                exl = true
                                break
                        }
                }
                if exl {
                        continue
                }
                ln.Add(lnm)
        }
}

// Validate ensures that layer names are valid.
func (ln *LayerNames) Validate(net *Network) error {
        var errs []error
        for _, lnm := range *ln {
                tly := net.LayerByName(lnm)
                if tly == nil {
                        err := fmt.Errorf("Validate: Layer name found %s", lnm)
                        errs = append(errs, err)
                }
        }
        return errors.Join(errs...)
}

// SendDA sends dopamine to SendTo list of layers.
func (ly *Layer) SendDA(da float32) {
        for _, lnm := range ly.SendTo {
                tly := ly.Network.LayerByName(lnm)
                if tly != nil {
                        tly.NeuroMod.DA = da
                }
        }
}

// SendACh sends ACh to SendTo list of layers.
func (ly *Layer) SendACh(ach float32) {
        for _, lnm := range ly.SendTo {
                tly := ly.Network.LayerByName(lnm)
                if tly != nil {
                        tly.NeuroMod.ACh = ach
                }
        }
}

////////  ClampDaLayer

// AddClampDaLayer adds a ClampDaLayer of given name
func (nt *Network) AddClampDaLayer(name string) *Layer {
        return nt.AddLayer2D(name, 1, 1, ClampDaLayer)
}

func (ly *Layer) ClampDaDefaults() {
        ly.Act.Clamp.Range.Set(-1, 1)
}

// SendDaFromAct is called in SendMods to send activity as DA.
func (ly *Layer) SendDaFromAct(ctx *Context) {
        act := ly.Neurons[0].Act
        ly.NeuroMod.DA = act
        ly.SendDA(act)
}

// NeuroMod are the neuromodulatory neurotransmitters, at the layer level.
type NeuroMod struct {

        // DA is dopamine, which primarily modulates learning, and also excitability,
        // and reflects the reward prediction error (RPE).
        DA float32

        // ACh is acetylcholine, which modulates excitability and also learning,
        // and reflects salience, i.e., reward (without discount by prediction) and
        // learned CS onset.
        ACh float32

        // SE is serotonin, which is a longer timescale neuromodulator with many
        // different effects. Currently not implemented, but here for future expansion.
        SE float32
}

func (nm *NeuroMod) Init() {
        nm.DA = 0
        nm.ACh = 0
        nm.SE = 0
}

//////// Enums

// DaReceptors for D1R and D2R dopamine receptors
type DaReceptors int32 //enums:enum

const (
        // D1R primarily expresses Dopamine D1 Receptors -- dopamine is excitatory and bursts of dopamine lead to increases in synaptic weight, while dips lead to decreases -- direct pathway in dorsal striatum
        D1R DaReceptors = iota

        // D2R primarily expresses Dopamine D2 Receptors -- dopamine is inhibitory and bursts of dopamine lead to decreases in synaptic weight, while dips lead to increases -- indirect pathway in dorsal striatum
        D2R
)

// Valences for Appetitive and Aversive valence coding
type Valences int32 //enums:enum

const (
        // Appetititve is a positive valence US (food, water, etc)
        Appetitive Valences = iota

        // Aversive is a negative valence US (shock, threat etc)
        Aversive
)

// Copyright (c) 2019, The Emergent Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package leabra

import (
        "fmt"
        "strings"
        "unsafe"

        "cogentcore.org/core/enums"
        "cogentcore.org/core/math32"
        "cogentcore.org/core/types"
        "github.com/emer/emergent/v2/emer"
)

// NeuronVarStart is the byte offset of fields in the Neuron structure
// where the float32 named variables start.
// Note: all non-float32 infrastructure variables must be at the start!
const NeuronVarStart = 12

// leabra.Neuron holds all of the neuron (unit) level variables -- this is the most basic version with
// rate-code only and no optional features at all.
// All variables accessible via Unit interface must be float32 and start at the top, in contiguous order
type Neuron struct {

        // bit flags for binary state variables
        Flags NeurFlags

        // index of the sub-level inhibitory pool that this neuron is in (only for 4D shapes, the pool (unit-group / hypercolumn) structure level) -- indicies start at 1 -- 0 is layer-level pool (is 0 if no sub-pools).
        SubPool int32

        ////////////////////////////
        // Act

        // rate-coded activation value reflecting final output of neuron communicated to other neurons, typically in range 0-1.  This value includes adaptation and synaptic depression / facilitation effects which produce temporal contrast (see ActLrn for version without this).  For rate-code activation, this is noisy-x-over-x-plus-one (NXX1) function; for discrete spiking it is computed from the inverse of the inter-spike interval (ISI), and Spike reflects the discrete spikes.
        Act float32

        // total excitatory synaptic conductance -- the net excitatory input to the neuron -- does *not* include Gbar.E
        Ge float32

        // total inhibitory synaptic conductance -- the net inhibitory input to the neuron -- does *not* include Gbar.I
        Gi float32

        // total potassium conductance, typically reflecting sodium-gated potassium currents involved in adaptation effects -- does *not* include Gbar.K
        Gk float32

        // net current produced by all channels -- drives update of Vm
        Inet float32

        // membrane potential -- integrates Inet current over time
        Vm float32

        // noise value added to unit (ActNoiseParams determines distribution, and when / where it is added)
        Noise float32

        // whether neuron has spiked or not (0 or 1), for discrete spiking neurons.
        Spike float32

        // target value: drives learning to produce this activation value
        Targ float32

        // external input: drives activation of unit from outside influences (e.g., sensory input)
        Ext float32

        ////////////////////////////
        // Learn

        // super-short time-scale average of ActLrn activation -- provides the lowest-level time integration -- for spiking this integrates over spikes before subsequent averaging, and it is also useful for rate-code to provide a longer time integral overall
        AvgSS float32

        // short time-scale average of ActLrn activation -- tracks the most recent activation states (integrates over AvgSS values), and represents the plus phase for learning in XCAL algorithms
        AvgS float32

        // medium time-scale average of ActLrn activation -- integrates over AvgS values, and represents the minus phase for learning in XCAL algorithms
        AvgM float32

        // long time-scale average of medium-time scale (trial level) activation, used for the BCM-style floating threshold in XCAL
        AvgL float32

        // how much to learn based on the long-term floating threshold (AvgL) for BCM-style Hebbian learning -- is modulated by level of AvgL itself (stronger Hebbian as average activation goes higher) and optionally the average amount of error experienced in the layer (to retain a common proportionality with the level of error-driven learning across layers)
        AvgLLrn float32

        // short time-scale activation average that is actually used for learning -- typically includes a small contribution from AvgM in addition to mostly AvgS, as determined by LrnActAvgParams.LrnM -- important to ensure that when unit turns off in plus phase (short time scale), enough medium-phase trace remains so that learning signal doesn't just go all the way to 0, at which point no learning would take place
        AvgSLrn float32

        // learning activation value, reflecting *dendritic* activity that is not affected by synaptic depression or adapdation channels which are located near the axon hillock.  This is the what drives the Avg* values that drive learning. Computationally, neurons strongly discount the signals sent to other neurons to provide temporal contrast, but need to learn based on a more stable reflection of their overall inputs in the dendrites.
        ActLrn float32

        ////////////////////////////
        // Phase

        // the activation state at end of third quarter, which is the traditional posterior-cortical minus phase activation
        ActM float32

        // the activation state at end of fourth quarter, which is the traditional posterior-cortical plus_phase activation
        ActP float32

        // ActP - ActM -- difference between plus and minus phase acts -- reflects the individual error gradient for this neuron in standard error-driven learning terms
        ActDif float32

        // delta activation: change in Act from one cycle to next -- can be useful to track where changes are taking place
        ActDel float32

        // the activation state at start of current alpha cycle (same as the state at end of previous cycle)
        ActQ0 float32

        // the activation state at end of first quarter of current alpha cycle
        ActQ1 float32

        // the activation state at end of second quarter of current alpha cycle
        ActQ2 float32

        // average activation (of final plus phase activation state) over long time intervals (time constant = DtPars.AvgTau -- typically 200) -- useful for finding hog units and seeing overall distribution of activation
        ActAvg float32

        // 5IB bursting activation value, computed by thresholding regular activation
        Burst float32

        // previous bursting activation -- used for context-based learning
        BurstPrv float32

        ////////////////////////////
        // Gmisc

        // aggregated synaptic inhibition (from Inhib pathways) -- time integral of GiRaw -- this is added with computed FFFB inhibition to get the full inhibition in Gi
        GiSyn float32

        // total amount of self-inhibition -- time-integrated to avoid oscillations
        GiSelf float32

        // last activation value sent (only send when diff is over threshold)
        ActSent float32

        // raw excitatory conductance (net input) received from sending units (send delta's are added to this value)
        GeRaw float32

        // raw inhibitory conductance (net input) received from sending units (send delta's are added to this value)
        GiRaw float32

        // conductance of sodium-gated potassium channel (KNa) fast dynamics (M-type) -- produces accommodation / adaptation of firing
        GknaFast float32

        // conductance of sodium-gated potassium channel (KNa) medium dynamics (Slick) -- produces accommodation / adaptation of firing
        GknaMed float32

        // conductance of sodium-gated potassium channel (KNa) slow dynamics (Slack) -- produces accommodation / adaptation of firing
        GknaSlow float32

        // current inter-spike-interval -- counts up since last spike.  Starts at -1 when initialized.
        ISI float32

        // average inter-spike-interval -- average time interval between spikes.  Starts at -1 when initialized, and goes to -2 after first spike, and is only valid after the second spike post-initialization.
        ISIAvg float32

        // CtxtGe is context (temporally delayed) excitatory conducances.
        CtxtGe float32

        ////////// Special algorithm vars: RL, PBWM

        // gating activation -- the activity value when gating occurred in this pool.
        ActG float32

        // per-neuron effective learning dopamine value -- gain modulated and sign reversed for D2R
        DALrn float32

        // shunting input received from Patch neurons (in reality flows through SNc DA pathways)
        Shunt float32

        // maintenance value for Deep layers = sending act at time of gating
        Maint float32

        // maintenance excitatory conductance value for Deep layers
        MaintGe float32
}

var NeuronVars = []string{
        "Act", "Ge", "Gi", "Gk", "Inet", "Vm", "Noise", "Spike", "Targ", "Ext",
        "AvgSS", "AvgS", "AvgM", "AvgL", "AvgLLrn", "AvgSLrn", "ActLrn",
        "ActM", "ActP", "ActDif", "ActDel", "ActQ0", "ActQ1", "ActQ2", "ActAvg", "Burst", "BurstPrv",
        "GiSyn", "GiSelf", "ActSent", "GeRaw", "GiRaw", "GknaFast", "GknaMed", "GknaSlow", "ISI", "ISIAvg", "CtxtGe",
        "ActG", "DALrn", "Shunt", "Maint", "MaintGe", "DA", "ACh", "SE", "GateAct", "GateNow", "GateCnt"}

var NeuronVarsMap map[string]int

var VarCategories = []emer.VarCategory{
        {"Act", "basic activation variables, including conductances, current, Vm, spiking"},
        {"Learn", "calcium-based learning variables"},
        {"Phase", "phase-based activation state"},
        {"Gmisc", "more detailed conductance (G) variables, for specific channels and computational values"},
        {"PBWM", "prefrontal cortex basal ganglia working memory model variables, including neuromodulation"},
        {"Wts", "weights and other synaptic-level variables"},
}

var NeuronVarProps = map[string]string{
        // Act vars
        "Act":   `cat:"Act"`,
        "Ge":    `cat:"Act"`,
        "Gi":    `cat:"Act"`,
        "Gk":    `cat:"Act"`,
        "Inet":  `cat:"Act"`,
        "Vm":    `cat:"Act" min:"0" max:"1"`,
        "Noise": `cat:"Act"`,
        "Spike": `cat:"Act"`,
        "Targ":  `cat:"Act"`,
        "Ext":   `cat:"Act"`,

        // Learn vars
        "AvgSS":   `cat:"Learn"`,
        "AvgS":    `cat:"Learn"`,
        "AvgM":    `cat:"Learn"`,
        "AvgL":    `cat:"Learn"`,
        "AvgLLrn": `cat:"Learn"`,
        "AvgSLrn": `cat:"Learn"`,
        "ActLrn":  `cat:"Learn"`,

        // Phase vars
        "ActM":     `cat:"Phase"`,
        "ActP":     `cat:"Phase"`,
        "ActDif":   `cat:"Phase" auto-scale:"+"`,
        "ActDel":   `cat:"Phase" auto-scale:"+"`,
        "ActQ0":    `cat:"Phase"`,
        "ActQ1":    `cat:"Phase"`,
        "ActQ2":    `cat:"Phase"`,
        "ActAvg":   `cat:"Phase"`,
        "Burst":    `cat:"Phase"`,
        "BurstPrv": `cat:"Phase"`,

        // Gmisc vars
        "GiSyn":    `cat:"Gmisc"`,
        "GiSelf":   `cat:"Gmisc"`,
        "ActSent":  `cat:"Gmisc"`,
        "GeRaw":    `cat:"Gmisc"`,
        "GiRaw":    `cat:"Gmisc"`,
        "GknaFast": `cat:"Gmisc"`,
        "GknaMed":  `cat:"Gmisc"`,
        "GknaSlow": `cat:"Gmisc"`,
        "ISI":      `cat:"Gmisc"`,
        "ISIAvg":   `cat:"Gmisc"`,
        "CtxtGe":   `cat:"Gmisc"`,

        "ActG":    `cat:"PBWM"`,
        "DALrn":   `cat:"PBWM"`,
        "Shunt":   `cat:"PBWM"`,
        "Maint":   `cat:"PBWM"`,
        "MaintGe": `cat:"PBWM"`,
        "DA":      `cat:"PBWM"`,
        "ACh":     `cat:"PBWM"`,
        "SE":      `cat:"PBWM"`,
        "GateAct": `cat:"PBWM"`,
        "GateNow": `cat:"PBWM"`,
        "GateCnt": `cat:"PBWM"`,
}

func init() {
        NeuronVarsMap = make(map[string]int, len(NeuronVars))
        for i, v := range NeuronVars {
                NeuronVarsMap[v] = i
        }
        ntyp := types.For[Neuron]()
        for _, fld := range ntyp.Fields {
                tag := NeuronVarProps[fld.Name]
                NeuronVarProps[fld.Name] = tag + ` doc:"` + strings.ReplaceAll(fld.Doc, "\n", " ") + `"`
        }
}

func (nrn *Neuron) VarNames() []string {
        return NeuronVars
}

// NeuronVarIndexByName returns the index of the variable in the Neuron, or error
func NeuronVarIndexByName(varNm string) (int, error) {
        i, ok := NeuronVarsMap[varNm]
        if !ok {
                return -1, fmt.Errorf("Neuron VarByName: variable name: %v not valid", varNm)
        }
        return i, nil
}

// VarByIndex returns variable using index (0 = first variable in NeuronVars list)
func (nrn *Neuron) VarByIndex(idx int) float32 {
        fv := (*float32)(unsafe.Pointer(uintptr(unsafe.Pointer(nrn)) + uintptr(NeuronVarStart+4*idx)))
        return *fv
}

// VarByName returns variable by name, or error
func (nrn *Neuron) VarByName(varNm string) (float32, error) {
        i, err := NeuronVarIndexByName(varNm)
        if err != nil {
                return math32.NaN(), err
        }
        return nrn.VarByIndex(i), nil
}

func (nrn *Neuron) HasFlag(f NeurFlags) bool {
        return nrn.Flags.HasFlag(f)
}

func (nrn *Neuron) SetFlag(on bool, f ...enums.BitFlag) {
        nrn.Flags.SetFlag(on, f...)
}

// IsOff returns true if the neuron has been turned off (lesioned)
func (nrn *Neuron) IsOff() bool {
        return nrn.HasFlag(NeurOff)
}

// NeurFlags are bit-flags encoding relevant binary state for neurons
type NeurFlags int64 //enums:bitflag

// The neuron flags
const (
        // NeurOff flag indicates that this neuron has been turned off (i.e., lesioned)
        NeurOff NeurFlags = iota

        // NeurHasExt means the neuron has external input in its Ext field
        NeurHasExt

        // NeurHasTarg means the neuron has external target input in its Targ field
        NeurHasTarg

        // NeurHasCmpr means the neuron has external comparison input in its Targ field -- used for computing
        // comparison statistics but does not drive neural activity ever
        NeurHasCmpr
)

// Copyright (c) 2019, The Emergent Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package leabra

import (
        "cogentcore.org/core/math32"
        "github.com/emer/etensor/tensor"
)

// note: path.go contains algorithm methods; pathbase.go has infrastructure.

//////////////////////////////////////////////////////////////////////////////////////
//  Init methods

// SetScalesRPool initializes synaptic Scale values using given tensor
// of values which has unique values for each recv neuron within a given pool.
func (pt *Path) SetScalesRPool(scales tensor.Tensor) {
        rNuY := scales.DimSize(0)
        rNuX := scales.DimSize(1)
        rNu := rNuY * rNuX
        rfsz := scales.Len() / rNu

        rsh := pt.Recv.Shape
        rNpY := rsh.DimSize(0)
        rNpX := rsh.DimSize(1)
        r2d := false
        if rsh.NumDims() != 4 {
                r2d = true
                rNpY = 1
                rNpX = 1
        }

        for rpy := 0; rpy < rNpY; rpy++ {
                for rpx := 0; rpx < rNpX; rpx++ {
                        for ruy := 0; ruy < rNuY; ruy++ {
                                for rux := 0; rux < rNuX; rux++ {
                                        ri := 0
                                        if r2d {
                                                ri = rsh.Offset([]int{ruy, rux})
                                        } else {
                                                ri = rsh.Offset([]int{rpy, rpx, ruy, rux})
                                        }
                                        scst := (ruy*rNuX + rux) * rfsz
                                        nc := int(pt.RConN[ri])
                                        st := int(pt.RConIndexSt[ri])
                                        for ci := 0; ci < nc; ci++ {
                                                // si := int(pj.RConIndex[st+ci]) // could verify coords etc
                                                rsi := pt.RSynIndex[st+ci]
                                                sy := &pt.Syns[rsi]
                                                sc := scales.Float1D(scst + ci)
                                                sy.Scale = float32(sc)
                                        }
                                }
                        }
                }
        }
}

// SetWtsFunc initializes synaptic Wt value using given function
// based on receiving and sending unit indexes.
func (pt *Path) SetWtsFunc(wtFun func(si, ri int, send, recv *tensor.Shape) float32) {
        rsh := &pt.Recv.Shape
        rn := rsh.Len()
        ssh := &pt.Send.Shape

        for ri := 0; ri < rn; ri++ {
                nc := int(pt.RConN[ri])
                st := int(pt.RConIndexSt[ri])
                for ci := 0; ci < nc; ci++ {
                        si := int(pt.RConIndex[st+ci])
                        wt := wtFun(si, ri, ssh, rsh)
                        rsi := pt.RSynIndex[st+ci]
                        sy := &pt.Syns[rsi]
                        sy.Wt = wt * sy.Scale
                        pt.Learn.LWtFromWt(sy)
                }
        }
}

// SetScalesFunc initializes synaptic Scale values using given function
// based on receiving and sending unit indexes.
func (pt *Path) SetScalesFunc(scaleFun func(si, ri int, send, recv *tensor.Shape) float32) {
        rsh := &pt.Recv.Shape
        rn := rsh.Len()
        ssh := &pt.Send.Shape

        for ri := 0; ri < rn; ri++ {
                nc := int(pt.RConN[ri])
                st := int(pt.RConIndexSt[ri])
                for ci := 0; ci < nc; ci++ {
                        si := int(pt.RConIndex[st+ci])
                        sc := scaleFun(si, ri, ssh, rsh)
                        rsi := pt.RSynIndex[st+ci]
                        sy := &pt.Syns[rsi]
                        sy.Scale = sc
                }
        }
}

// InitWeightsSyn initializes weight values based on WtInit randomness parameters
// for an individual synapse.
// It also updates the linear weight value based on the sigmoidal weight value.
func (pt *Path) InitWeightsSyn(syn *Synapse) {
        if syn.Scale == 0 {
                syn.Scale = 1
        }
        syn.Wt = float32(pt.WtInit.Gen())
        // enforce normalized weight range -- required for most uses and if not
        // then a new type of path should be used:
        if syn.Wt < 0 {
                syn.Wt = 0
        }
        if syn.Wt > 1 {
                syn.Wt = 1
        }
        syn.LWt = pt.Learn.WtSig.LinFromSigWt(syn.Wt)
        syn.Wt *= syn.Scale // note: scale comes after so LWt is always "pure" non-scaled value
        syn.DWt = 0
        syn.Norm = 0
        syn.Moment = 0
}

// InitWeights initializes weight values according to Learn.WtInit params
func (pt *Path) InitWeights() {
        for si := range pt.Syns {
                sy := &pt.Syns[si]
                pt.InitWeightsSyn(sy)
        }
        for wi := range pt.WbRecv {
                wb := &pt.WbRecv[wi]
                wb.Init()
        }
        pt.InitGInc()
        pt.ClearTrace()
}

// InitWtSym initializes weight symmetry -- is given the reciprocal pathway where
// the Send and Recv layers are reversed.
func (pt *Path) InitWtSym(rpt *Path) {
        slay := pt.Send
        ns := int32(len(slay.Neurons))
        for si := int32(0); si < ns; si++ {
                nc := pt.SConN[si]
                st := pt.SConIndexSt[si]
                for ci := int32(0); ci < nc; ci++ {
                        sy := &pt.Syns[st+ci]
                        ri := pt.SConIndex[st+ci]
                        // now we need to find the reciprocal synapse on rpt!
                        // look in ri for sending connections
                        rsi := ri
                        if len(rpt.SConN) == 0 {
                                continue
                        }
                        rsnc := rpt.SConN[rsi]
                        if rsnc == 0 {
                                continue
                        }
                        rsst := rpt.SConIndexSt[rsi]
                        rist := rpt.SConIndex[rsst]        // starting index in recv path
                        ried := rpt.SConIndex[rsst+rsnc-1] // ending index
                        if si < rist || si > ried {        // fast reject -- paths are always in order!
                                continue
                        }
                        // start at index proportional to si relative to rist
                        up := int32(0)
                        if ried > rist {
                                up = int32(float32(rsnc) * float32(si-rist) / float32(ried-rist))
                        }
                        dn := up - 1

                        for {
                                doing := false
                                if up < rsnc {
                                        doing = true
                                        rrii := rsst + up
                                        rri := rpt.SConIndex[rrii]
                                        if rri == si {
                                                rsy := &rpt.Syns[rrii]
                                                rsy.Wt = sy.Wt
                                                rsy.LWt = sy.LWt
                                                rsy.Scale = sy.Scale
                                                // note: if we support SymFromTop then can have option to go other way
                                                break
                                        }
                                        up++
                                }
                                if dn >= 0 {
                                        doing = true
                                        rrii := rsst + dn
                                        rri := rpt.SConIndex[rrii]
                                        if rri == si {
                                                rsy := &rpt.Syns[rrii]
                                                rsy.Wt = sy.Wt
                                                rsy.LWt = sy.LWt
                                                rsy.Scale = sy.Scale
                                                // note: if we support SymFromTop then can have option to go other way
                                                break
                                        }
                                        dn--
                                }
                                if !doing {
                                        break
                                }
                        }
                }
        }
}

// InitGInc initializes the per-pathway GInc threadsafe increment -- not
// typically needed (called during InitWeights only) but can be called when needed
func (pt *Path) InitGInc() {
        for ri := range pt.GInc {
                pt.GInc[ri] = 0
                pt.CtxtGeInc[ri] = 0
                pt.GeRaw[ri] = 0
        }
}

//////////////////////////////////////////////////////////////////////////////////////
//  Act methods

// SendGDelta sends the delta-activation from sending neuron index si,
// to integrate synaptic conductances on receivers
func (pt *Path) SendGDelta(si int, delta float32) {
        if pt.Type == CTCtxtPath {
                return
        }
        scdel := delta * pt.GScale
        nc := pt.SConN[si]
        st := pt.SConIndexSt[si]
        syns := pt.Syns[st : st+nc]
        scons := pt.SConIndex[st : st+nc]
        for ci := range syns {
                ri := scons[ci]
                pt.GInc[ri] += scdel * syns[ci].Wt
        }
}

// RecvGInc increments the receiver's GeRaw or GiRaw from that of all the pathways.
func (pt *Path) RecvGInc() {
        rlay := pt.Recv
        switch pt.Type {
        case CTCtxtPath:
                // nop
        case InhibPath:
                for ri := range rlay.Neurons {
                        rn := &rlay.Neurons[ri]
                        rn.GiRaw += pt.GInc[ri]
                        pt.GInc[ri] = 0
                }
        case GPiThalPath:
                for ri := range rlay.Neurons {
                        rn := &rlay.Neurons[ri]
                        ginc := pt.GInc[ri]
                        pt.GeRaw[ri] += ginc
                        rn.GeRaw += ginc
                        pt.GInc[ri] = 0
                }
        default:
                for ri := range rlay.Neurons {
                        rn := &rlay.Neurons[ri]
                        rn.GeRaw += pt.GInc[ri]
                        pt.GInc[ri] = 0
                }
        }
}

//////////////////////////////////////////////////////////////////////////////////////
//  Learn methods

// DWt computes the weight change (learning) -- on sending pathways
func (pt *Path) DWt() {
        if !pt.Learn.Learn {
                return
        }
        switch {
        case pt.Type == CHLPath && pt.CHL.On:
                pt.DWtCHL()
        case pt.Type == CTCtxtPath:
                pt.DWtCTCtxt()
        case pt.Type == EcCa1Path:
                pt.DWtEcCa1()
        case pt.Type == MatrixPath:
                pt.DWtMatrix()
        case pt.Type == RWPath:
                pt.DWtRW()
        case pt.Type == TDPredPath:
                pt.DWtTDPred()
        case pt.Type == DaHebbPath:
                pt.DWtDaHebb()
        default:
                pt.DWtStd()
        }
}

// DWt computes the weight change (learning) -- on sending pathways
func (pt *Path) DWtStd() {
        slay := pt.Send
        rlay := pt.Recv
        for si := range slay.Neurons {
                sn := &slay.Neurons[si]
                if sn.AvgS < pt.Learn.XCal.LrnThr && sn.AvgM < pt.Learn.XCal.LrnThr {
                        continue
                }
                nc := int(pt.SConN[si])
                st := int(pt.SConIndexSt[si])
                syns := pt.Syns[st : st+nc]
                scons := pt.SConIndex[st : st+nc]
                for ci := range syns {
                        sy := &syns[ci]
                        ri := scons[ci]
                        rn := &rlay.Neurons[ri]
                        err, bcm := pt.Learn.CHLdWt(sn.AvgSLrn, sn.AvgM, rn.AvgSLrn, rn.AvgM, rn.AvgL)

                        bcm *= pt.Learn.XCal.LongLrate(rn.AvgLLrn)
                        err *= pt.Learn.XCal.MLrn
                        dwt := bcm + err
                        norm := float32(1)
                        if pt.Learn.Norm.On {
                                norm = pt.Learn.Norm.NormFromAbsDWt(&sy.Norm, math32.Abs(dwt))
                        }
                        if pt.Learn.Momentum.On {
                                dwt = norm * pt.Learn.Momentum.MomentFromDWt(&sy.Moment, dwt)
                        } else {
                                dwt *= norm
                        }
                        sy.DWt += pt.Learn.Lrate * dwt
                }
                // aggregate max DWtNorm over sending synapses
                if pt.Learn.Norm.On {
                        maxNorm := float32(0)
                        for ci := range syns {
                                sy := &syns[ci]
                                if sy.Norm > maxNorm {
                                        maxNorm = sy.Norm
                                }
                        }
                        for ci := range syns {
                                sy := &syns[ci]
                                sy.Norm = maxNorm
                        }
                }
        }
}

// WtFromDWt updates the synaptic weight values from delta-weight changes -- on sending pathways
func (pt *Path) WtFromDWt() {
        if !pt.Learn.Learn {
                return
        }
        switch pt.Type {
        case RWPath, TDPredPath:
                pt.WtFromDWtLinear()
                return
        }
        if pt.Learn.WtBal.On {
                for si := range pt.Syns {
                        sy := &pt.Syns[si]
                        ri := pt.SConIndex[si]
                        wb := &pt.WbRecv[ri]
                        pt.Learn.WtFromDWt(wb.Inc, wb.Dec, &sy.DWt, &sy.Wt, &sy.LWt, sy.Scale)
                }
        } else {
                for si := range pt.Syns {
                        sy := &pt.Syns[si]
                        pt.Learn.WtFromDWt(1, 1, &sy.DWt, &sy.Wt, &sy.LWt, sy.Scale)
                }
        }
}

// WtFromDWtLinear updates the synaptic weight values from delta-weight
// changes, with no constraints or limits
func (pt *Path) WtFromDWtLinear() {
        for si := range pt.Syns {
                sy := &pt.Syns[si]
                if sy.DWt != 0 {
                        sy.Wt += sy.DWt // straight update, no limits or anything
                        sy.LWt = sy.Wt
                        sy.DWt = 0
                }
        }
}

// WtBalFromWt computes the Weight Balance factors based on average recv weights
func (pt *Path) WtBalFromWt() {
        if !pt.Learn.Learn || !pt.Learn.WtBal.On {
                return
        }

        rlay := pt.Recv
        if !pt.Learn.WtBal.Targs && rlay.IsTarget() {
                return
        }
        for ri := range rlay.Neurons {
                nc := int(pt.RConN[ri])
                if nc < 1 {
                        continue
                }
                wb := &pt.WbRecv[ri]
                st := int(pt.RConIndexSt[ri])
                rsidxs := pt.RSynIndex[st : st+nc]
                sumWt := float32(0)
                sumN := 0
                for ci := range rsidxs {
                        rsi := rsidxs[ci]
                        sy := &pt.Syns[rsi]
                        if sy.Wt >= pt.Learn.WtBal.AvgThr {
                                sumWt += sy.Wt
                                sumN++
                        }
                }
                if sumN > 0 {
                        sumWt /= float32(sumN)
                } else {
                        sumWt = 0
                }
                wb.Avg = sumWt
                wb.Fact, wb.Inc, wb.Dec = pt.Learn.WtBal.WtBal(sumWt)
        }
}

// LrateMult sets the new Lrate parameter for Paths to LrateInit * mult.
// Useful for implementing learning rate schedules.
func (pt *Path) LrateMult(mult float32) {
        pt.Learn.Lrate = pt.Learn.LrateInit * mult
}

///////////////////////////////////////////////////////////////////////
//  WtBalRecvPath

// WtBalRecvPath are state variables used in computing the WtBal weight balance function
// There is one of these for each Recv Neuron participating in the pathway.
type WtBalRecvPath struct {

        // average of effective weight values that exceed WtBal.AvgThr across given Recv Neuron's connections for given Path
        Avg float32

        // overall weight balance factor that drives changes in WbInc vs. WbDec via a sigmoidal function -- this is the net strength of weight balance changes
        Fact float32

        // weight balance increment factor -- extra multiplier to add to weight increases to maintain overall weight balance
        Inc float32

        // weight balance decrement factor -- extra multiplier to add to weight decreases to maintain overall weight balance
        Dec float32
}

func (wb *WtBalRecvPath) Init() {
        wb.Avg = 0
        wb.Fact = 0
        wb.Inc = 1
        wb.Dec = 1
}

// Copyright (c) 2019, The Emergent Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package leabra

import (
        "encoding/json"
        "errors"
        "fmt"
        "io"
        "log"
        "strconv"
        "strings"

        "cogentcore.org/core/base/indent"
        "cogentcore.org/core/math32"
        "cogentcore.org/core/math32/minmax"
        "github.com/emer/emergent/v2/emer"
        "github.com/emer/emergent/v2/paths"
        "github.com/emer/emergent/v2/weights"
        "github.com/emer/etensor/tensor"
)

// note: paths.go contains algorithm methods; pathbase.go has infrastructure.

// Path implements the Leabra algorithm at the synaptic level,
// in terms of a pathway connecting two layers.
type Path struct {
        emer.PathBase

        // sending layer for this pathway.
        Send *Layer

        // receiving layer for this pathway.
        Recv *Layer

        // type of pathway.
        Type PathTypes

        // initial random weight distribution
        WtInit WtInitParams `display:"inline"`

        // weight scaling parameters: modulates overall strength of pathway,
        // using both absolute and relative factors.
        WtScale WtScaleParams `display:"inline"`

        // synaptic-level learning parameters
        Learn LearnSynParams `display:"add-fields"`

        // For CTCtxtPath if true, this is the pathway from corresponding
        // Superficial layer.  Should be OneToOne path, with Learn.Learn = false,
        // WtInit.Var = 0, Mean = 0.8. These defaults are set if FromSuper = true.
        FromSuper bool

        // CHL are the parameters for CHL learning. if CHL is On then
        // WtSig.SoftBound is automatically turned off, as it is incompatible.
        CHL CHLParams `display:"inline"`

        // special parameters for matrix trace learning
        Trace TraceParams `display:"inline"`

        // synaptic state values, ordered by the sending layer
        // units which owns them -- one-to-one with SConIndex array.
        Syns []Synapse

        // scaling factor for integrating synaptic input conductances (G's).
        // computed in AlphaCycInit, incorporates running-average activity levels.
        GScale float32

        // local per-recv unit increment accumulator for synaptic
        // conductance from sending units. goes to either GeRaw or GiRaw
        // on neuron depending on pathway type.
        GInc []float32

        // CtxtGeInc is local per-recv unit accumulator for Ctxt excitatory
        // conductance from sending units, Not a delta, the full value.
        CtxtGeInc []float32

        // per-recv, per-path raw excitatory input, for GPiThalPath.
        GeRaw []float32

        // weight balance state variables for this pathway, one per recv neuron.
        WbRecv []WtBalRecvPath

        // number of recv connections for each neuron in the receiving layer,
        // as a flat list.
        RConN []int32 `display:"-"`

        // average and maximum number of recv connections in the receiving layer.
        RConNAvgMax minmax.AvgMax32 `edit:"-" display:"inline"`

        // starting index into ConIndex list for each neuron in
        // receiving layer; list incremented by ConN.
        RConIndexSt []int32 `display:"-"`

        // index of other neuron on sending side of pathway,
        // ordered by the receiving layer's order of units as the
        // outer loop (each start is in ConIndexSt),
        // and then by the sending layer's units within that.
        RConIndex []int32 `display:"-"`

        // index of synaptic state values for each recv unit x connection,
        // for the receiver pathway which does not own the synapses,
        // and instead indexes into sender-ordered list.
        RSynIndex []int32 `display:"-"`

        // number of sending connections for each neuron in the
        // sending layer, as a flat list.
        SConN []int32 `display:"-"`

        // average and maximum number of sending connections
        // in the sending layer.
        SConNAvgMax minmax.AvgMax32 `edit:"-" display:"inline"`

        // starting index into ConIndex list for each neuron in
        // sending layer; list incremented by ConN.
        SConIndexSt []int32 `display:"-"`

        // index of other neuron on receiving side of pathway,
        // ordered by the sending layer's order of units as the
        // outer loop (each start is in ConIndexSt), and then
        // by the sending layer's units within that.
        SConIndex []int32 `display:"-"`
}

// emer.Path interface

func (pt *Path) StyleObject() any      { return pt }
func (pt *Path) RecvLayer() emer.Layer { return pt.Recv }
func (pt *Path) SendLayer() emer.Layer { return pt.Send }
func (pt *Path) TypeName() string      { return pt.Type.String() }
func (pt *Path) TypeNumber() int       { return int(pt.Type) }

func (pt *Path) Defaults() {
        pt.WtInit.Defaults()
        pt.WtScale.Defaults()
        pt.Learn.Defaults()
        pt.CHL.Defaults()
        pt.Trace.Defaults()
        pt.GScale = 1
        pt.DefaultsForType()
}

func (pt *Path) DefaultsForType() {
        switch pt.Type {
        case CHLPath:
                pt.CHLDefaults()
        case EcCa1Path:
                pt.EcCa1Defaults()
        case TDPredPath:
                pt.TDPredDefaults()
        case RWPath:
                pt.RWDefaults()
        case MatrixPath:
                pt.MatrixDefaults()
        case DaHebbPath:
                pt.DaHebbDefaults()
        }
}

// UpdateParams updates all params given any changes that might have been made to individual values
func (pt *Path) UpdateParams() {
        pt.WtScale.Update()
        pt.Learn.Update()
        pt.Learn.LrateInit = pt.Learn.Lrate
        if pt.Type == CHLPath && pt.CHL.On {
                pt.Learn.WtSig.SoftBound = false
        }
        pt.CHL.Update()
        pt.Trace.Update()
}

func (pt *Path) ShouldDisplay(field string) bool {
        switch field {
        case "CHL":
                return pt.Type == CHLPath
        case "Trace":
                return pt.Type == MatrixPath
        default:
                return true
        }
        return true
}

// AllParams returns a listing of all parameters in the Layer
func (pt *Path) AllParams() string {
        str := "///////////////////////////////////////////////////\nPath: " + pt.Name + "\n"
        b, _ := json.MarshalIndent(&pt.WtInit, "", " ")
        str += "WtInit: {\n " + JsonToParams(b)
        b, _ = json.MarshalIndent(&pt.WtScale, "", " ")
        str += "WtScale: {\n " + JsonToParams(b)
        b, _ = json.MarshalIndent(&pt.Learn, "", " ")
        str += "Learn: {\n " + strings.Replace(JsonToParams(b), " XCal: {", "\n  XCal: {", -1)
        return str
}

func (pt *Path) SynVarNames() []string {
        return SynapseVars
}

// SynVarProps returns properties for variables
func (pt *Path) SynVarProps() map[string]string {
        return SynapseVarProps
}

// SynIndex returns the index of the synapse between given send, recv unit indexes
// (1D, flat indexes). Returns -1 if synapse not found between these two neurons.
// Requires searching within connections for receiving unit.
func (pt *Path) SynIndex(sidx, ridx int) int {
        nc := int(pt.SConN[sidx])
        st := int(pt.SConIndexSt[sidx])
        for ci := 0; ci < nc; ci++ {
                ri := int(pt.SConIndex[st+ci])
                if ri != ridx {
                        continue
                }
                return int(st + ci)
        }
        return -1
}

// SynVarIndex returns the index of given variable within the synapse,
// according to *this path's* SynVarNames() list (using a map to lookup index),
// or -1 and error message if not found.
func (pt *Path) SynVarIndex(varNm string) (int, error) {
        return SynapseVarByName(varNm)
}

// SynVarNum returns the number of synapse-level variables
// for this path.  This is needed for extending indexes in derived types.
func (pt *Path) SynVarNum() int {
        return len(SynapseVars)
}

// Numsyns returns the number of synapses for this path.
// This is the max idx for SynValue1D
// and the number of vals set by SynValues.
func (pt *Path) NumSyns() int {
        return len(pt.Syns)
}

// SynVal1D returns value of given variable index (from SynVarIndex)
// on given SynIndex.  Returns NaN on invalid index.
// This is the core synapse var access method used by other methods,
// so it is the only one that needs to be updated for derived layer types.
func (pt *Path) SynValue1D(varIndex int, synIndex int) float32 {
        if synIndex < 0 || synIndex >= len(pt.Syns) {
                return math32.NaN()
        }
        if varIndex < 0 || varIndex >= pt.SynVarNum() {
                return math32.NaN()
        }
        sy := &pt.Syns[synIndex]
        return sy.VarByIndex(varIndex)
}

// SynValues sets values of given variable name for each synapse,
// using the natural ordering of the synapses (sender based for Leabra),
// into given float32 slice (only resized if not big enough).
// Returns error on invalid var name.
func (pt *Path) SynValues(vals *[]float32, varNm string) error {
        vidx, err := pt.SynVarIndex(varNm)
        if err != nil {
                return err
        }
        ns := len(pt.Syns)
        if *vals == nil || cap(*vals) < ns {
                *vals = make([]float32, ns)
        } else if len(*vals) < ns {
                *vals = (*vals)[0:ns]
        }
        for i := range pt.Syns {
                (*vals)[i] = pt.SynValue1D(vidx, i)
        }
        return nil
}

// SynVal returns value of given variable name on the synapse
// between given send, recv unit indexes (1D, flat indexes).
// Returns math32.NaN() for access errors (see SynValTry for error message)
func (pt *Path) SynValue(varNm string, sidx, ridx int) float32 {
        vidx, err := pt.SynVarIndex(varNm)
        if err != nil {
                return math32.NaN()
        }
        synIndex := pt.SynIndex(sidx, ridx)
        return pt.SynValue1D(vidx, synIndex)
}

// SetSynVal sets value of given variable name on the synapse
// between given send, recv unit indexes (1D, flat indexes)
// returns error for access errors.
func (pt *Path) SetSynValue(varNm string, sidx, ridx int, val float32) error {
        vidx, err := pt.SynVarIndex(varNm)
        if err != nil {
                return err
        }
        synIndex := pt.SynIndex(sidx, ridx)
        if synIndex < 0 || synIndex >= len(pt.Syns) {
                return err
        }
        sy := &pt.Syns[synIndex]
        sy.SetVarByIndex(vidx, val)
        if varNm == "Wt" {
                pt.Learn.LWtFromWt(sy)
        }
        return nil
}

///////////////////////////////////////////////////////////////////////
//  Weights File

// WriteWeightsJSON writes the weights from this pathway from the receiver-side perspective
// in a JSON text format.  We build in the indentation logic to make it much faster and
// more efficient.
func (pt *Path) WriteWeightsJSON(w io.Writer, depth int) {
        slay := pt.Send
        rlay := pt.Recv
        nr := len(rlay.Neurons)
        w.Write(indent.TabBytes(depth))
        w.Write([]byte("{\n"))
        depth++
        w.Write(indent.TabBytes(depth))
        w.Write([]byte(fmt.Sprintf("\"From\": %q,\n", slay.Name)))
        w.Write(indent.TabBytes(depth))
        w.Write([]byte(fmt.Sprintf("\"MetaData\": {\n")))
        depth++
        w.Write(indent.TabBytes(depth))
        w.Write([]byte(fmt.Sprintf("\"GScale\": \"%g\"\n", pt.GScale)))
        depth--
        w.Write(indent.TabBytes(depth))
        w.Write([]byte("},\n"))
        w.Write(indent.TabBytes(depth))
        w.Write([]byte(fmt.Sprintf("\"Rs\": [\n")))
        depth++
        for ri := 0; ri < nr; ri++ {
                nc := int(pt.RConN[ri])
                st := int(pt.RConIndexSt[ri])
                w.Write(indent.TabBytes(depth))
                w.Write([]byte("{\n"))
                depth++
                w.Write(indent.TabBytes(depth))
                w.Write([]byte(fmt.Sprintf("\"Ri\": %v,\n", ri)))
                w.Write(indent.TabBytes(depth))
                w.Write([]byte(fmt.Sprintf("\"N\": %v,\n", nc)))
                w.Write(indent.TabBytes(depth))
                w.Write([]byte("\"Si\": [ "))
                for ci := 0; ci < nc; ci++ {
                        si := pt.RConIndex[st+ci]
                        w.Write([]byte(fmt.Sprintf("%v", si)))
                        if ci == nc-1 {
                                w.Write([]byte(" "))
                        } else {
                                w.Write([]byte(", "))
                        }
                }
                w.Write([]byte("],\n"))
                w.Write(indent.TabBytes(depth))
                w.Write([]byte("\"Wt\": [ "))
                for ci := 0; ci < nc; ci++ {
                        rsi := pt.RSynIndex[st+ci]
                        sy := &pt.Syns[rsi]
                        w.Write([]byte(strconv.FormatFloat(float64(sy.Wt), 'g', weights.Prec, 32)))
                        if ci == nc-1 {
                                w.Write([]byte(" "))
                        } else {
                                w.Write([]byte(", "))
                        }
                }
                w.Write([]byte("]\n"))
                depth--
                w.Write(indent.TabBytes(depth))
                if ri == nr-1 {
                        w.Write([]byte("}\n"))
                } else {
                        w.Write([]byte("},\n"))
                }
        }
        depth--
        w.Write(indent.TabBytes(depth))
        w.Write([]byte("]\n"))
        depth--
        w.Write(indent.TabBytes(depth))
        w.Write([]byte("}")) // note: leave unterminated as outer loop needs to add , or just \n depending
}

// SetWeights sets the weights for this pathway from weights.Path decoded values
func (pt *Path) SetWeights(pw *weights.Path) error {
        if pw.MetaData != nil {
                if gs, ok := pw.MetaData["GScale"]; ok {
                        pv, _ := strconv.ParseFloat(gs, 32)
                        pt.GScale = float32(pv)
                }
        }
        var err error
        for i := range pw.Rs {
                pr := &pw.Rs[i]
                for si := range pr.Si {
                        er := pt.SetSynValue("Wt", pr.Si[si], pr.Ri, pr.Wt[si]) // updates lin wt
                        if er != nil {
                                err = er
                        }
                }
        }
        return err
}

// Connect sets the connectivity between two layers and the pattern to use in interconnecting them
func (pt *Path) Connect(slay, rlay *Layer, pat paths.Pattern, typ PathTypes) {
        pt.Send = slay
        pt.Recv = rlay
        pt.Pattern = pat
        pt.Type = typ
        pt.Name = pt.Send.Name + "To" + pt.Recv.Name
}

// Validate tests for non-nil settings for the pathway -- returns error
// message or nil if no problems (and logs them if logmsg = true)
func (pt *Path) Validate(logmsg bool) error {
        emsg := ""
        if pt.Pattern == nil {
                emsg += "Pat is nil; "
        }
        if pt.Recv == nil {
                emsg += "Recv is nil; "
        }
        if pt.Send == nil {
                emsg += "Send is nil; "
        }
        if emsg != "" {
                err := errors.New(emsg)
                if logmsg {
                        log.Println(emsg)
                }
                return err
        }
        return nil
}

// Build constructs the full connectivity among the layers
// as specified in this pathway.
// Calls Validate and returns false if invalid.
// Pattern.Connect is called to get the pattern of the connection.
// Then the connection indexes are configured according to that pattern.
func (pt *Path) Build() error {
        if pt.Off {
                return nil
        }
        err := pt.Validate(true)
        if err != nil {
                return err
        }
        ssh := &pt.Send.Shape
        rsh := &pt.Recv.Shape
        sendn, recvn, cons := pt.Pattern.Connect(ssh, rsh, pt.Recv == pt.Send)
        slen := ssh.Len()
        rlen := rsh.Len()
        tcons := pt.SetNIndexSt(&pt.SConN, &pt.SConNAvgMax, &pt.SConIndexSt, sendn)
        tconr := pt.SetNIndexSt(&pt.RConN, &pt.RConNAvgMax, &pt.RConIndexSt, recvn)
        if tconr != tcons {
                log.Printf("%v programmer error: total recv cons %v != total send cons %v\n", pt.String(), tconr, tcons)
        }
        pt.RConIndex = make([]int32, tconr)
        pt.RSynIndex = make([]int32, tconr)
        pt.SConIndex = make([]int32, tcons)

        sconN := make([]int32, slen) // temporary mem needed to tracks cur n of sending cons

        cbits := cons.Values
        for ri := 0; ri < rlen; ri++ {
                rbi := ri * slen     // recv bit index
                rtcn := pt.RConN[ri] // number of cons
                rst := pt.RConIndexSt[ri]
                rci := int32(0)
                for si := 0; si < slen; si++ {
                        if !cbits.Index(rbi + si) { // no connection
                                continue
                        }
                        sst := pt.SConIndexSt[si]
                        if rci >= rtcn {
                                log.Printf("%v programmer error: recv target total con number: %v exceeded at recv idx: %v, send idx: %v\n", pt.String(), rtcn, ri, si)
                                break
                        }
                        pt.RConIndex[rst+rci] = int32(si)

                        sci := sconN[si]
                        stcn := pt.SConN[si]
                        if sci >= stcn {
                                log.Printf("%v programmer error: send target total con number: %v exceeded at recv idx: %v, send idx: %v\n", pt.String(), stcn, ri, si)
                                break
                        }
                        pt.SConIndex[sst+sci] = int32(ri)
                        pt.RSynIndex[rst+rci] = sst + sci
                        (sconN[si])++
                        rci++
                }
        }
        pt.Syns = make([]Synapse, len(pt.SConIndex))
        pt.GInc = make([]float32, rlen)
        pt.CtxtGeInc = make([]float32, rlen)
        pt.GeRaw = make([]float32, rlen)
        pt.WbRecv = make([]WtBalRecvPath, rlen)
        return nil
}

// SetNIndexSt sets the *ConN and *ConIndexSt values given n tensor from Pat.
// Returns total number of connections for this direction.
func (pt *Path) SetNIndexSt(n *[]int32, avgmax *minmax.AvgMax32, idxst *[]int32, tn *tensor.Int32) int32 {
        ln := tn.Len()
        tnv := tn.Values
        *n = make([]int32, ln)
        *idxst = make([]int32, ln)
        idx := int32(0)
        avgmax.Init()
        for i := 0; i < ln; i++ {
                nv := tnv[i]
                (*n)[i] = nv
                (*idxst)[i] = idx
                idx += nv
                avgmax.UpdateValue(float32(nv), int32(i))
        }
        avgmax.CalcAvg()
        return idx
}

// String satisfies fmt.Stringer for path
func (pt *Path) String() string {
        str := ""
        if pt.Recv == nil {
                str += "recv=nil; "
        } else {
                str += pt.Recv.Name + " <- "
        }
        if pt.Send == nil {
                str += "send=nil"
        } else {
                str += pt.Send.Name
        }
        if pt.Pattern == nil {
                str += " Pat=nil"
        } else {
                str += " Pat=" + pt.Pattern.Name()
        }
        return str
}

// Copyright (c) 2024, The Emergent Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package leabra

import (
        "fmt"

        "cogentcore.org/core/math32"
)

// MatrixParams has parameters for Dorsal Striatum Matrix computation.
// These are the main Go / NoGo gating units in BG driving updating of PFC WM in PBWM.
type MatrixParams struct {

        // Quarter(s) when learning takes place, typically Q2 and Q4, corresponding to the PFC GateQtr. Note: this is a bitflag and must be accessed using bitflag.Set / Has etc routines, 32 bit versions.
        LearnQtr Quarters

        // how much the patch shunt activation multiplies the dopamine values -- 0 = complete shunting, 1 = no shunting -- should be a factor < 1.0
        PatchShunt float32 `default:"0.2,0.5" min:"0" max:"1"`

        // also shunt the ACh value driven from CIN units -- this prevents clearing of MSNConSpec traces -- more plausibly the patch units directly interfere with the effects of CIN's rather than through ach, but it is easier to implement with ach shunting here.
        ShuntACh bool `default:"true"`

        // how much does the LACK of ACh from the CIN units drive extra inhibition to output-gating Matrix units -- gi += out_ach_inhib * (1-ach) -- provides a bias for output gating on reward trials -- do NOT apply to NoGo, only Go -- this is a key param -- between 0.1-0.3 usu good -- see how much output gating happening and change accordingly
        OutAChInhib float32 `default:"0,0.3"`

        // multiplicative gain factor applied to positive (burst) dopamine signals in computing DALrn effect learning dopamine value based on raw DA that we receive (D2R reversal occurs *after* applying Burst based on sign of raw DA)
        BurstGain float32 `default:"1"`

        // multiplicative gain factor applied to positive (burst) dopamine signals in computing DALrn effect learning dopamine value based on raw DA that we receive (D2R reversal occurs *after* applying Burst based on sign of raw DA)
        DipGain float32 `default:"1"`
}

func (mp *MatrixParams) Defaults() {
        mp.LearnQtr.SetFlag(true, Q2)
        mp.LearnQtr.SetFlag(true, Q4)
        mp.PatchShunt = 0.2
        mp.ShuntACh = true
        mp.OutAChInhib = 0.3
        mp.BurstGain = 1
        mp.DipGain = 1
}

func (mp *MatrixParams) Update() {
}

func (ly *Layer) MatrixDefaults() {
        // special inhib params
        ly.PBWM.Type = MaintOut
        ly.Inhib.Layer.Gi = 1.9
        ly.Inhib.Layer.FB = 0.5
        ly.Inhib.Pool.On = true
        ly.Inhib.Pool.Gi = 1.9
        ly.Inhib.Pool.FB = 0
        ly.Inhib.Self.On = true
        ly.Inhib.Self.Gi = 0.3
        ly.Inhib.ActAvg.Fixed = true
        ly.Inhib.ActAvg.Init = 0.2
}

// DALrnFromDA returns effective learning dopamine value from given raw DA value
// applying Burst and Dip Gain factors, and then reversing sign for D2R.
func (ly *Layer) DALrnFromDA(da float32) float32 {
        if da > 0 {
                da *= ly.Matrix.BurstGain
        } else {
                da *= ly.Matrix.DipGain
        }
        if ly.PBWM.DaR == D2R {
                da *= -1
        }
        return da
}

// MatrixOutAChInhib applies OutAChInhib to bias output gating on reward trials.
func (ly *Layer) MatrixOutAChInhib(ctx *Context) {
        if ly.Matrix.OutAChInhib == 0 {
                return
        }

        ypN := ly.Shape.DimSize(0)
        xpN := ly.Shape.DimSize(1)
        ynN := ly.Shape.DimSize(2)
        xnN := ly.Shape.DimSize(3)
        maintN := ly.PBWM.MaintN
        layAch := ly.NeuroMod.ACh // ACh comes from CIN neurons, represents reward time
        for yp := 0; yp < ypN; yp++ {
                for xp := maintN; xp < xpN; xp++ {
                        for yn := 0; yn < ynN; yn++ {
                                for xn := 0; xn < xnN; xn++ {
                                        ni := ly.Shape.Offset([]int{yp, xp, yn, xn})
                                        nrn := &ly.Neurons[ni]
                                        if nrn.IsOff() {
                                                continue
                                        }
                                        ach := layAch
                                        if ly.Matrix.ShuntACh && nrn.Shunt > 0 {
                                                ach *= ly.Matrix.PatchShunt
                                        }
                                        achI := ly.Matrix.OutAChInhib * (1 - ach)
                                        nrn.Gi += achI
                                }
                        }
                }
        }
}

// DaAChFromLay computes Da and ACh from layer and Shunt received from PatchLayer units
func (ly *Layer) DaAChFromLay(ctx *Context) {
        for ni := range ly.Neurons {
                nrn := &ly.Neurons[ni]
                if nrn.IsOff() {
                        continue
                }
                da := ly.NeuroMod.DA
                if nrn.Shunt > 0 { // note: treating Shunt as binary variable -- could multiply
                        da *= ly.Matrix.PatchShunt
                }
                nrn.DALrn = ly.DALrnFromDA(da)
        }
}

// RecGateAct records the gating activation from current activation, when gating occcurs
// based on GateState.Now
func (ly *Layer) RecGateAct(ctx *Context) {
        for pi := range ly.Pools {
                if pi == 0 {
                        continue
                }
                pl := &ly.Pools[pi]
                if !pl.Gate.Now { // not gating now
                        continue
                }
                for ni := pl.StIndex; ni < pl.EdIndex; ni++ {
                        nrn := &ly.Neurons[ni]
                        if nrn.IsOff() {
                                continue
                        }
                        nrn.ActG = nrn.Act
                }
        }
}

// GateTypes for region of striatum
type GateTypes int32 //enums:enum

const (
        // Maint is maintenance gating -- toggles active maintenance in PFC.
        Maint GateTypes = iota

        // Out is output gating -- drives deep layer activation.
        Out

        // MaintOut for maint and output gating.
        MaintOut
)

// SendToMatrixPFC adds standard SendTo layers for PBWM: MatrixGo, NoGo, PFCmntD, PFCoutD
// with optional prefix -- excludes mnt, out cases if corresp shape = 0
func (ly *Layer) SendToMatrixPFC(prefix string) {
        pfcprefix := "PFC"
        if prefix != "" {
                pfcprefix = prefix
        }
        std := []string{prefix + "MatrixGo", prefix + "MatrixNoGo", pfcprefix + "mntD", pfcprefix + "outD"}
        ly.SendTo = make([]string, 2)
        for i, s := range std {
                nm := s
                switch {
                case i < 2:
                        ly.SendTo[i] = nm
                case i == 2:
                        if ly.PBWM.MaintX > 0 {
                                ly.SendTo = append(ly.SendTo, nm)
                        }
                case i == 3:
                        if ly.PBWM.OutX > 0 {
                                ly.SendTo = append(ly.SendTo, nm)
                        }
                }
        }
}

// SendPBWMParams send PBWMParams info to all SendTo layers -- convenient config-time
// way to ensure all are consistent -- also checks validity of SendTo's
func (ly *Layer) SendPBWMParams() error {
        var lasterr error
        for _, lnm := range ly.SendTo {
                tly := ly.Network.LayerByName(lnm)
                tly.PBWM.CopyGeomFrom(&ly.PBWM)
        }
        return lasterr
}

// MatrixPaths returns the recv paths from Go and NoGo MatrixLayer pathways -- error if not
// found or if paths are not of the GPiThalPath type
func (ly *Layer) MatrixPaths() (goPath, nogoPath *Path, err error) {
        for _, p := range ly.RecvPaths {
                if p.Off {
                        continue
                }
                slay := p.Send
                if slay.Type == MatrixLayer {
                        if ly.PBWM.DaR == D1R {
                                goPath = p
                        } else {
                                nogoPath = p
                        }
                } else {
                        nogoPath = p
                }
        }
        if goPath == nil || nogoPath == nil {
                err = fmt.Errorf("GPiThalLayer must have RecvPath's from a MatrixLayer D1R (Go) and another NoGo layer")
        }
        return
}

// PBWMParams defines the shape of the outer pool dimensions of gating layers,
// organized into Maint and Out subsets which are arrayed along the X axis
// with Maint first (to the left) then Out.  Individual layers may only
// represent Maint or Out subsets of this overall shape, but all need
// to have this coordinated shape information to be able to share gating
// state information.  Each layer represents gate state information in
// their native geometry -- FullIndex1D provides access from a subset
// to full set.
type PBWMParams struct {
        // Type of gating layer
        Type GateTypes

        // dominant type of dopamine receptor -- D1R for Go pathway, D2R for NoGo
        DaR DaReceptors

        // overall shape dimensions for the full set of gating pools,
        // e.g., as present in the Matrix and GPiThal levels
        Y int

        // how many pools in the X dimension are Maint gating pools -- rest are Out
        MaintX int

        // how many pools in the X dimension are Out gating pools -- comes after Maint
        OutX int

        // For the Matrix layers, this is the number of Maint Pools in X outer
        // dimension of 4D shape -- Out gating after that. Note: it is unclear
        // how this relates to MaintX, but it is different in SIR model.
        MaintN int
}

func (pp *PBWMParams) Defaults() {

}

func (pp *PBWMParams) Update() {

}

// Set sets the shape parameters: number of Y dimension pools, and
// numbers of maint and out pools along X axis
func (pp *PBWMParams) Set(nY, maintX, outX int) {
        pp.Y = nY
        pp.MaintX = maintX
        pp.OutX = outX
}

// TotX returns the total number of X-axis pools (Maint + Out)
func (pp *PBWMParams) TotX() int {
        return pp.MaintX + pp.OutX
}

func (pp *PBWMParams) CopyGeomFrom(src *PBWMParams) {
        pp.Set(src.Y, src.MaintX, src.OutX)
        pp.Type = src.Type
}

// Index returns the index into GateStates for given 2D pool coords
// for given GateType.  Each type stores gate info in its "native" 2D format.
func (pp *PBWMParams) Index(pY, pX int, typ GateTypes) int {
        switch typ {
        case Maint:
                if pp.MaintX == 0 {
                        return 0
                }
                return pY*pp.MaintX + pX
        case Out:
                if pp.OutX == 0 {
                        return 0
                }
                return pY*pp.OutX + pX
        case MaintOut:
                return pY*pp.TotX() + pX
        }
        return 0
}

// FullIndex1D returns the index into full MaintOut GateStates
// for given 1D pool idx (0-based) *from given GateType*.
func (pp *PBWMParams) FullIndex1D(idx int, fmTyp GateTypes) int {
        switch fmTyp {
        case Maint:
                if pp.MaintX == 0 {
                        return 0
                }
                // convert to 2D and use that
                pY := idx / pp.MaintX
                pX := idx % pp.MaintX
                return pp.Index(pY, pX, MaintOut)
        case Out:
                if pp.OutX == 0 {
                        return 0
                }
                // convert to 2D and use that
                pY := idx / pp.OutX
                pX := idx%pp.OutX + pp.MaintX
                return pp.Index(pY, pX, MaintOut)
        case MaintOut:
                return idx
        }
        return 0
}

//////// GateState

// GateState is gating state values stored in layers that receive thalamic gating signals
// including MatrixLayer, PFCLayer, GPiThal layer, etc -- use GateLayer as base layer to include.
type GateState struct {

        // gating activation value, reflecting thalamic gating layer activation at time of gating (when Now = true) -- will be 0 if gating below threshold for this pool, and prior to first Now for AlphaCycle
        Act float32

        // gating timing signal -- true if this is the moment when gating takes place
        Now bool

        // unique to each layer -- not copied.  Generally is a counter for interval between gating signals -- starts at -1, goes to 0 at first gating, counts up from there for subsequent gating.  Can be reset back to -1 when gate is reset (e.g., output gating) and counts down from -1 while not gating.
        Cnt int `copy:"-"`
}

// Init initializes the values -- call during InitActs()
func (gs *GateState) Init() {
        gs.Act = 0
        gs.Now = false
        gs.Cnt = -1
}

// CopyFrom copies from another GateState -- only the Act and Now signals are copied
func (gs *GateState) CopyFrom(fm *GateState) {
        gs.Act = fm.Act
        gs.Now = fm.Now
}

// GateType returns type of gating for this layer
func (ly *Layer) GateType() GateTypes {
        switch ly.Type {
        case GPiThalLayer, MatrixLayer:
                return MaintOut
        case PFCDeepLayer:
                if ly.PFCGate.OutGate {
                        return Out
                }
                return Maint
        }
        return MaintOut
}

// SetGateStates sets the GateStates from given source states, of given gating type
func (ly *Layer) SetGateStates(src *Layer, typ GateTypes) {
        myt := ly.GateType()
        if myt < MaintOut && typ < MaintOut && myt != typ { // mismatch
                return
        }
        switch {
        case myt == typ:
                mx := min(len(src.Pools), len(ly.Pools))
                for i := 1; i < mx; i++ {
                        ly.Pool(i).Gate.CopyFrom(&src.Pool(i).Gate)
                }
        default: // typ == MaintOut, myt = Maint or Out
                mx := len(ly.Pools)
                for i := 1; i < mx; i++ {
                        gs := &ly.Pool(i).Gate
                        si := 1 + ly.PBWM.FullIndex1D(i-1, myt)
                        sgs := &src.Pool(si).Gate
                        gs.CopyFrom(sgs)
                }
        }
}

//////// GPiThalLayer

// GPiGateParams has gating parameters for gating in GPiThal layer, including threshold.
type GPiGateParams struct {
        // GateQtr is the Quarter(s) when gating takes place, typically Q1 and Q3,
        // which is the quarter prior to the PFC GateQtr when deep layer updating
        // takes place. Note: this is a bitflag and must be accessed using bitflag.
        // Set / Has etc routines, 32 bit versions.
        GateQtr Quarters

        // Cycle within Qtr to determine if activation over threshold for gating.
        // We send GateState updates on this cycle either way.
        Cycle int `default:"18"`

        // extra netinput gain factor to compensate for reduction in Ge from subtracting away NoGo -- this is *IN ADDITION* to adding the NoGo factor as an extra gain: Ge = (GeGain + NoGo) * (GoIn - NoGo * NoGoIn)
        GeGain float32 `default:"3"`

        // how much to weight NoGo inputs relative to Go inputs (which have an implied weight of 1 -- this also up-scales overall Ge to compensate for subtraction
        NoGo float32 `min:"0" default:"1,0.1"`

        // threshold for gating, applied to activation -- when any GPiThal unit activation gets above this threshold, it counts as having gated, driving updating of GateState which is broadcast to other layers that use the gating signal
        Thr float32 `default:"0.2"`

        // Act value of GPiThal unit reflects gating threshold: if below threshold, it is zeroed -- see ActLrn for underlying non-thresholded activation
        ThrAct bool `default:"true"`
}

func (gp *GPiGateParams) Defaults() {
        gp.GateQtr.SetFlag(true, Q1)
        gp.GateQtr.SetFlag(true, Q3)
        gp.Cycle = 18
        gp.GeGain = 3
        gp.NoGo = 1
        gp.Thr = 0.2
        gp.ThrAct = true
}

func (gp *GPiGateParams) Update() {
}

// GeRaw returns the net GeRaw from go, nogo specific values
func (gp *GPiGateParams) GeRaw(goRaw, nogoRaw float32) float32 {
        return (gp.GeGain + gp.NoGo) * (goRaw - gp.NoGo*nogoRaw)
}

func (ly *Layer) GPiThalDefaults() {
        ly.PBWM.Type = MaintOut
        ly.Inhib.Layer.Gi = 1.8
        ly.Inhib.Layer.FB = 0.2
        ly.Inhib.Pool.On = false
        ly.Inhib.ActAvg.Fixed = true
        ly.Inhib.ActAvg.Init = 1
}

// GPiGFromInc integrates new synaptic conductances from increments
// sent during last SendGDelta.
func (ly *Layer) GPiGFromInc(ctx *Context) {
        goPath, nogoPath, _ := ly.MatrixPaths()
        for ni := range ly.Neurons {
                nrn := &ly.Neurons[ni]
                if nrn.IsOff() {
                        continue
                }
                goRaw := goPath.GeRaw[ni]
                nogoRaw := nogoPath.GeRaw[ni]
                nrn.GeRaw = ly.GPiGate.GeRaw(goRaw, nogoRaw)
                ly.Act.GeFromRaw(nrn, nrn.GeRaw)
                ly.Act.GiFromRaw(nrn, nrn.GiRaw)
        }
}

// GPiGateSend updates gating state and sends it along to other layers
func (ly *Layer) GPiGateSend(ctx *Context) {
        ly.GPiGateFromAct(ctx)
        ly.GPiSendGateStates()
}

// GPiGateFromAct updates GateState from current activations, at time of gating
func (ly *Layer) GPiGateFromAct(ctx *Context) {
        gateQtr := ly.GPiGate.GateQtr.HasFlag(ctx.Quarter)
        qtrCyc := ctx.QuarterCycle()
        for ni := range ly.Neurons {
                nrn := &ly.Neurons[ni]
                if nrn.IsOff() {
                        continue
                }
                gs := &ly.Pool(int(nrn.SubPool)).Gate
                if ctx.Quarter == 0 && qtrCyc == 0 {
                        gs.Act = 0 // reset at start
                }
                if gateQtr && qtrCyc == ly.GPiGate.Cycle { // gating
                        gs.Now = true
                        if nrn.Act < ly.GPiGate.Thr { // didn't gate
                                gs.Act = 0 // not over thr
                                if ly.GPiGate.ThrAct {
                                        gs.Act = 0
                                }
                                if gs.Cnt >= 0 {
                                        gs.Cnt++
                                } else if gs.Cnt < 0 {
                                        gs.Cnt--
                                }
                        } else { // did gate
                                gs.Cnt = 0
                                gs.Act = nrn.Act
                        }
                } else {
                        gs.Now = false
                }
        }
}

// GPiSendGateStates sends GateStates to other layers
func (ly *Layer) GPiSendGateStates() {
        myt := MaintOut // always
        for _, lnm := range ly.SendTo {
                gl := ly.Network.LayerByName(lnm)
                gl.SetGateStates(ly, myt)
        }
}

//////// CINLayer

// CINParams (cholinergic interneuron) reads reward signals from named source layer(s)
// and sends the Max absolute value of that activity as the positively rectified
// non-prediction-discounted reward signal computed by CINs, and sent as
// an acetylcholine (ACh) signal.
// To handle positive-only reward signals, need to include both a reward prediction
// and reward outcome layer.
type CINParams struct {
        // RewThr is the threshold on reward values from RewLays,
        // to count as a significant reward event, which then drives maximal ACh.
        // Set to 0 to disable this nonlinear behavior.
        RewThr float32 `default:"0.1"`

        // Reward-representing layer(s) from which this computes ACh as Max absolute value
        RewLays LayerNames
}

func (ly *CINParams) Defaults() {
        ly.RewThr = 0.1
}

func (ly *CINParams) Update() {
}

// CINMaxAbsRew returns the maximum absolute value of reward layer activations.
func (ly *Layer) CINMaxAbsRew() float32 {
        mx := float32(0)
        for _, nm := range ly.CIN.RewLays {
                ly := ly.Network.LayerByName(nm)
                if ly == nil {
                        continue
                }
                act := math32.Abs(ly.Pools[0].Inhib.Act.Max)
                mx = math32.Max(mx, act)
        }
        return mx
}

func (ly *Layer) ActFromGCIN(ctx *Context) {
        ract := ly.CINMaxAbsRew()
        if ly.CIN.RewThr > 0 {
                if ract > ly.CIN.RewThr {
                        ract = 1
                }
        }
        for ni := range ly.Neurons {
                nrn := &ly.Neurons[ni]
                if nrn.IsOff() {
                        continue
                }
                nrn.Act = ract
                ly.Learn.AvgsFromAct(nrn)
        }
}

// SendAChFromAct sends ACh from neural activity in first unit.
func (ly *Layer) SendAChFromAct(ctx *Context) {
        act := ly.Neurons[0].Act
        ly.NeuroMod.ACh = act
        ly.SendACh(act)
}

//////// PFC

// PFCGateParams has parameters for PFC gating
type PFCGateParams struct {

        // Quarter(s) that the effect of gating on updating Deep from Super occurs -- this is typically 1 quarter after the GPiThal GateQtr
        GateQtr Quarters

        // if true, this PFC layer is an output gate layer, which means that it only has transient activation during gating
        OutGate bool

        // for output gating, only compute gating in first quarter -- do not compute in 3rd quarter -- this is typically true, and GateQtr is typically set to only Q1 as well -- does Burst updating immediately after first quarter gating signal -- allows gating signals time to influence performance within a single trial
        OutQ1Only bool `viewif:"OutGate" default:"true"`
}

func (gp *PFCGateParams) Defaults() {
        gp.GateQtr.SetFlag(true, Q2)
        gp.GateQtr.SetFlag(true, Q4)
        gp.OutQ1Only = true
}

func (gp *PFCGateParams) Update() {
}

// PFCMaintParams for PFC maintenance functions
type PFCMaintParams struct {

        // use fixed dynamics for updating deep_ctxt activations -- defined in dyn_table -- this also preserves the initial gating deep_ctxt value in Maint neuron val (view as Cust1) -- otherwise it is up to the recurrent loops between super and deep for maintenance
        UseDyn bool

        // multiplier on maint current
        MaintGain float32 `min:"0" default:"0.8"`

        // on output gating, clear corresponding maint pool.  theoretically this should be on, but actually it works better off in most cases..
        OutClearMaint bool `default:"false"`

        // how much to clear out (decay) super activations when the stripe itself gates and was previously maintaining something, or for maint pfc stripes, when output go fires and clears.
        Clear    float32 `min:"0" max:"1" default:"0"`
        MaxMaint int     `"min:"1" default:"1:100" maximum duration of maintenance for any stripe -- beyond this limit, the maintenance is just automatically cleared -- typically 1 for output gating and 100 for maintenance gating"`
}

func (mp *PFCMaintParams) Defaults() {
        mp.MaintGain = 0.8
        mp.OutClearMaint = false // theoretically should be true, but actually was false due to bug
        mp.Clear = 0
        mp.MaxMaint = 100
}

func (mp *PFCMaintParams) Update() {
}

// PFC dynamic behavior element -- defines the dynamic behavior of deep layer PFC units
type PFCDyn struct {

        // initial value at point when gating starts -- MUST be > 0 when used.
        Init float32

        // time constant for linear rise in maintenance activation (per quarter when deep is updated) -- use integers -- if both rise and decay then rise comes first
        RiseTau float32

        // time constant for linear decay in maintenance activation (per quarter when deep is updated) -- use integers -- if both rise and decay then rise comes first
        DecayTau float32

        // description of this factor
        Desc string
}

func (pd *PFCDyn) Defaults() {
        pd.Init = 1
}

func (pd *PFCDyn) Set(init, rise, decay float32, desc string) {
        pd.Init = init
        pd.RiseTau = rise
        pd.DecayTau = decay
        pd.Desc = desc
}

// Value returns dynamic value at given time point
func (pd *PFCDyn) Value(time float32) float32 {
        val := pd.Init
        if time <= 0 {
                return val
        }
        if pd.RiseTau > 0 && pd.DecayTau > 0 {
                if time >= pd.RiseTau {
                        val = 1 - ((time - pd.RiseTau) / pd.DecayTau)
                } else {
                        val = pd.Init + (1-pd.Init)*(time/pd.RiseTau)
                }
        } else if pd.RiseTau > 0 {
                val = pd.Init + (1-pd.Init)*(time/pd.RiseTau)
        } else if pd.DecayTau > 0 {
                val = pd.Init - pd.Init*(time/pd.DecayTau)
        }
        if val > 1 {
                val = 1
        }
        if val < 0.001 {
                val = 0.001
        }
        return val
}

//////////////////////////////////////////////////////////////////////////////
//  PFCDyns

// PFCDyns is a slice of dyns. Provides deterministic control over PFC
// maintenance dynamics -- the rows of PFC units (along Y axis) behave
// according to corresponding index of Dyns.
// ensure layer Y dim has even multiple of len(Dyns).
type PFCDyns []*PFCDyn

// SetDyn sets given dynamic maint element to given parameters (must be allocated in list first)
func (pd *PFCDyns) SetDyn(dyn int, init, rise, decay float32, desc string) *PFCDyn {
        dy := &PFCDyn{}
        dy.Set(init, rise, decay, desc)
        (*pd)[dyn] = dy
        return dy
}

// MaintOnly creates basic default maintenance dynamic configuration -- every
// unit just maintains over time.
// This should be used for Output gating layer.
func (pd *PFCDyns) MaintOnly() {
        *pd = make([]*PFCDyn, 1)
        pd.SetDyn(0, 1, 0, 0, "maintain stable act")
}

// FullDyn creates full dynamic Dyn configuration, with 5 different
// dynamic profiles: stable maint, phasic, rising maint, decaying maint,
// and up / down maint.  tau is the rise / decay base time constant.
func (pd *PFCDyns) FullDyn(tau float32) {
        ndyn := 5
        *pd = make([]*PFCDyn, ndyn)

        pd.SetDyn(0, 1, 0, 0, "maintain stable act")
        pd.SetDyn(1, 1, 0, 1, "immediate phasic response")
        pd.SetDyn(2, .1, tau, 0, "maintained, rising value over time")
        pd.SetDyn(3, 1, 0, tau, "maintained, decaying value over time")
        pd.SetDyn(4, .1, .5*tau, tau, "maintained, rising then falling over time")
}

// Value returns value for given dyn item at given time step
func (pd *PFCDyns) Value(dyn int, time float32) float32 {
        sz := len(*pd)
        if sz == 0 {
                return 1
        }
        dy := (*pd)[dyn%sz]
        return dy.Value(time)
}

func (ly *Layer) PFCDeepDefaults() {
        if ly.PFCGate.OutGate && ly.PFCGate.OutQ1Only {
                ly.PFCMaint.MaxMaint = 1
                ly.PFCGate.GateQtr = 0
                ly.PFCGate.GateQtr.SetFlag(true, Q1)
        }
        if len(ly.PFCDyns) > 0 {
                ly.PFCMaint.UseDyn = true
        } else {
                ly.PFCMaint.UseDyn = false
        }
}

// MaintPFC returns corresponding PFCDeep maintenance layer
// with same name but outD -> mntD; could be nil
func (ly *Layer) MaintPFC() *Layer {
        sz := len(ly.Name)
        mnm := ly.Name[:sz-4] + "mntD"
        li := ly.Network.LayerByName(mnm)
        return li
}

// SuperPFC returns corresponding PFC super layer with same name without D
// should not be nil.  Super can be any layer type.
func (ly *Layer) SuperPFC() *Layer {
        dnm := ly.Name[:len(ly.Name)-1]
        li := ly.Network.LayerByName(dnm)
        return li
}

// MaintGInc increments Ge from MaintGe, for PFCDeepLayer.
func (ly *Layer) MaintGInc(ctx *Context) {
        for ni := range ly.Neurons {
                nrn := &ly.Neurons[ni]
                if nrn.IsOff() {
                        continue
                }
                geRaw := nrn.GeRaw + nrn.MaintGe
                ly.Act.GeFromRaw(nrn, geRaw)
                ly.Act.GiFromRaw(nrn, nrn.GiRaw)
        }
}

// PFCDeepGating updates PFC Gating state.
func (ly *Layer) PFCDeepGating(ctx *Context) {
        if ly.PFCGate.OutGate && ly.PFCGate.OutQ1Only {
                if ctx.Quarter > 1 {
                        return
                }
        }

        for pi := range ly.Pools {
                if pi == 0 {
                        continue
                }
                gs := &ly.Pools[pi].Gate
                if !gs.Now { // not gating now
                        continue
                }
                if gs.Act > 0 { // use GPiThal threshold, so anything > 0
                        gs.Cnt = 0              // this is the "just gated" signal
                        if ly.PFCGate.OutGate { // time to clear out maint
                                if ly.PFCMaint.OutClearMaint {
                                        fmt.Println("clear maint")
                                        ly.ClearMaint(pi)
                                }
                        } else {
                                pfcs := ly.SuperPFC()
                                pfcs.DecayStatePool(pi, ly.PFCMaint.Clear)
                        }
                }
                // test for over-duration maintenance -- allow for active gating to override
                if gs.Cnt >= ly.PFCMaint.MaxMaint {
                        gs.Cnt = -1
                }
        }
}

// ClearMaint resets maintenance in corresponding pool (0 based) in maintenance layer
func (ly *Layer) ClearMaint(pool int) {
        pfcm := ly.MaintPFC()
        if pfcm == nil {
                return
        }
        gs := &pfcm.Pools[pool].Gate
        if gs.Cnt >= 1 { // important: only for established maint, not just gated..
                gs.Cnt = -1 // reset
                pfcs := pfcm.SuperPFC()
                pfcs.DecayStatePool(pool, pfcm.PFCMaint.Clear)
        }
}

// DeepMaint updates deep maintenance activations
func (ly *Layer) DeepMaint(ctx *Context) {
        if !ly.PFCGate.GateQtr.HasFlag(ctx.Quarter) {
                return
        }
        sly := ly.SuperPFC()
        if sly == nil {
                return
        }
        yN := ly.Shape.DimSize(2)
        xN := ly.Shape.DimSize(3)

        nn := yN * xN

        syN := sly.Shape.DimSize(2)
        sxN := sly.Shape.DimSize(3)
        snn := syN * sxN

        dper := yN / syN  // dyn per sender -- should be len(Dyns)
        dtyp := yN / dper // dyn type

        for ni := range ly.Neurons {
                nrn := &ly.Neurons[ni]
                if nrn.IsOff() {
                        continue
                }
                ui := ni % nn
                pi := ni / nn
                uy := ui / xN
                ux := ui % xN

                gs := &ly.Pool(int(nrn.SubPool)).Gate
                if gs.Cnt < 0 {
                        nrn.Maint = 0
                        nrn.MaintGe = 0
                } else if gs.Cnt <= 1 { // first gating, save first gating value
                        sy := uy % syN // inner loop is s
                        si := pi*snn + sy*sxN + ux
                        snr := &sly.Neurons[si]
                        nrn.Maint = ly.PFCMaint.MaintGain * snr.Act
                }
                if ly.PFCMaint.UseDyn {
                        nrn.MaintGe = nrn.Maint * ly.PFCDyns.Value(dtyp, float32(gs.Cnt-1))
                } else {
                        nrn.MaintGe = nrn.Maint
                }
        }
}

// UpdateGateCnt updates the gate counter
func (ly *Layer) UpdateGateCnt(ctx *Context) {
        if !ly.PFCGate.GateQtr.HasFlag(ctx.Quarter) {
                return
        }
        for pi := range ly.Pools {
                if pi == 0 {
                        continue
                }
                gs := &ly.Pools[pi].Gate
                if gs.Cnt < 0 {
                        // ly.ClearCtxtPool(gi)
                        gs.Cnt--
                } else {
                        gs.Cnt++
                }
        }
}

// Copyright (c) 2024, The Emergent Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package leabra

import (
        "github.com/emer/emergent/v2/paths"
)

// RecGateAct is called after GateSend, to record gating activations at time of gating
func (nt *Network) RecGateAct(ctx *Context) {
        for _, ly := range nt.Layers {
                if ly.Off {
                        continue
                }
                ly.RecGateAct(ctx)
        }
}

// AddMatrixLayer adds a MatrixLayer of given size, with given name.
// nY = number of pools in Y dimension, nMaint + nOut are pools in X dimension,
// and each pool has nNeurY, nNeurX neurons.  da gives the DaReceptor type (D1R = Go, D2R = NoGo)
func (nt *Network) AddMatrixLayer(name string, nY, nMaint, nOut, nNeurY, nNeurX int, da DaReceptors) *Layer {
        tX := nMaint + nOut
        mtx := nt.AddLayer4D(name, nY, tX, nNeurY, nNeurX, MatrixLayer)
        mtx.PBWM.DaR = da
        mtx.PBWM.Set(nY, nMaint, nOut)
        return mtx
}

// AddGPeLayer adds a pbwm.Layer to serve as a GPe layer, with given name.
// nY = number of pools in Y dimension, nMaint + nOut are pools in X dimension,
// and each pool has 1x1 neurons.
func (nt *Network) AddGPeLayer(name string, nY, nMaint, nOut int) *Layer {
        tX := nMaint + nOut
        gpe := nt.AddLayer4D(name, nY, tX, 1, 1, GPeLayer)
        return gpe
}

// AddGPiThalLayer adds a GPiThalLayer of given size, with given name.
// nY = number of pools in Y dimension, nMaint + nOut are pools in X dimension,
// and each pool has 1x1 neurons.
func (nt *Network) AddGPiThalLayer(name string, nY, nMaint, nOut int) *Layer {
        tX := nMaint + nOut
        gpi := nt.AddLayer4D(name, nY, tX, 1, 1, GPiThalLayer)
        gpi.PBWM.Set(nY, nMaint, nOut)
        return gpi
}

// AddCINLayer adds a CINLayer, with a single neuron.
func (nt *Network) AddCINLayer(name string) *Layer {
        cin := nt.AddLayer2D(name, 1, 1, CINLayer)
        return cin
}

// AddDorsalBG adds MatrixGo, NoGo, GPe, GPiThal, and CIN layers, with given optional prefix.
// nY = number of pools in Y dimension, nMaint + nOut are pools in X dimension,
// and each pool has nNeurY, nNeurX neurons.  Appropriate PoolOneToOne connections
// are made to drive GPiThal, with BgFixed class name set so
// they can be styled appropriately (no learning, WtRnd.Mean=0.8, Var=0)
func (nt *Network) AddDorsalBG(prefix string, nY, nMaint, nOut, nNeurY, nNeurX int) (mtxGo, mtxNoGo, gpe, gpi, cin *Layer) {
        mtxGo = nt.AddMatrixLayer(prefix+"MatrixGo", nY, nMaint, nOut, nNeurY, nNeurX, D1R)
        mtxNoGo = nt.AddMatrixLayer(prefix+"MatrixNoGo", nY, nMaint, nOut, nNeurY, nNeurX, D2R)
        gpe = nt.AddGPeLayer(prefix+"GPeNoGo", nY, nMaint, nOut)
        gpi = nt.AddGPiThalLayer(prefix+"GPiThal", nY, nMaint, nOut)
        cin = nt.AddCINLayer(prefix + "CIN")
        cin.AddSendTo(mtxGo.Name, mtxNoGo.Name)

        mtxNoGo.PlaceBehind(mtxGo, 2)
        gpe.PlaceRightOf(mtxNoGo, 2)
        gpi.PlaceRightOf(mtxGo, 2)
        cin.PlaceBehind(gpe, 2)

        one2one := paths.NewPoolOneToOne()
        pt := nt.ConnectLayers(mtxGo, gpi, one2one, GPiThalPath)
        pt.AddClass("BgFixed")
        pt = nt.ConnectLayers(mtxNoGo, gpe, one2one, ForwardPath)
        pt.AddClass("BgFixed")
        pt = nt.ConnectLayers(gpe, gpi, one2one, GPiThalPath)
        pt.AddClass("BgFixed")

        mtxGo.Doc = "Matrisome (Matrix) striatum medium spiny neuron (MSN), which is the input layer of the basal ganglia (BG), with more D1 than D2 dopamine receptors, that drives the direct pathway to disinhibit BG outputs, favoring a 'Go' response"
        mtxNoGo.Doc = "Matrisome (Matrix) striatum medium spiny neuron (MSN), which is the input layer of the basal ganglia (BG), with more D2 than D1 dopamine receptors, that drives the indirect pathway through the globus pallidus external segment (GPe) net inhibit BG outputs, favoring a 'NoGo' response"
        gpe.Doc = "Globus pallidus external segment (GPe) of the BG that is tonically active and inhibited by the Matrix NoGo pathway, causing disinhibition of the GPi, and net inhibition of overall BG output responding."
        gpi.Doc = "Globus pallidus internal segment (GPi) of the BG that is tonically active and inhibited by the Matrix Go pathway (and disinhibited by the GPe via NoGo), which then inhibits the thalamus (Thal), with the net effect of disinhibiting cortical areas on BG Go pathway activation. This layer summarizes both GPi and Thal in a net excitatory, activity-positive manner. It sends gating signals to PFC via 'SendTo' layer names, not using standard synaptic pathways."
        cin.Doc = "Cholinergic interneurons (CIN) that represent a positively rectified, non-prediction-discounted reward and overall sensory salience signal, that modulates overall BG activity and learning around salient events."
        return
}

// AddPFCLayer adds a PFCLayer, super and deep, of given size, with given name.
// nY, nX = number of pools in Y, X dimensions, and each pool has nNeurY, nNeurX neurons.
// out is true for output-gating layer, and dynmaint is true for maintenance-only dyn,
// else Full set of 5 dynamic maintenance types. Both have the class "PFC" set.
// deep is positioned behind super.
func (nt *Network) AddPFCLayer(name string, nY, nX, nNeurY, nNeurX int, out, dynMaint bool) (sp, dp *Layer) {
        sp = nt.AddLayer4D(name, nY, nX, nNeurY, nNeurX, SuperLayer)
        dym := 1
        if !dynMaint {
                dym = 5
        }
        dp = nt.AddLayer4D(name+"D", nY, nX, dym*nNeurY, nNeurX, PFCDeepLayer)
        sp.AddClass("PFC")
        dp.AddClass("PFC")
        dp.PFCGate.OutGate = out
        if dynMaint {
                dp.PFCDyns.MaintOnly()
        } else {
                dp.PFCDyns.FullDyn(10)
        }
        dp.PlaceBehind(sp, 2)
        return
}

// AddPFC adds paired PFCmnt, PFCout and associated Deep layers,
// with given optional prefix.
// nY = number of pools in Y dimension, nMaint, nOut are pools in X dimension,
// and each pool has nNeurY, nNeurX neurons.
// dynMaint is true for maintenance-only dyn, else full set of 5 dynamic maintenance types.
// Appropriate OneToOne connections are made between PFCmntD -> PFCout.
func (nt *Network) AddPFC(prefix string, nY, nMaint, nOut, nNeurY, nNeurX int, dynMaint bool) (pfcMnt, pfcMntD, pfcOut, pfcOutD *Layer) {
        if prefix == "" {
                prefix = "PFC"
        }
        if nMaint > 0 {
                pfcMnt, pfcMntD = nt.AddPFCLayer(prefix+"mnt", nY, nMaint, nNeurY, nNeurX, false, dynMaint)
                pfcMnt.Doc = "Prefrontal Cortex (PFC) maintenance (mnt) superficial layer, which receives inputs from other brain areas and drives BG (basal ganglia) gated input into the robust maintenance deep layers"
                pfcMntD.Doc = "Prefrontal Cortex (PFC) maintenance (mnt) deep layer, which has special intrinsic circuits and channels supporting robust active firing even in the absence of other inputs, and holds on to information relevant for behavioral responses, but does not directly drive those outputs"
        }
        if nOut > 0 {
                pfcOut, pfcOutD = nt.AddPFCLayer(prefix+"out", nY, nOut, nNeurY, nNeurX, true, dynMaint)
                pfcOut.Doc = "Prefrontal Cortex (PFC) output (out) superficial layer, which receives inputs from PFC maintenance and other brain areas and drives BG (basal ganglia) gated input into the output deep layers"
                pfcOutD.Doc = "Prefrontal Cortex (PFC) output (out) deep layer, which drives behavioral output pathways, either as direct motor outputs, or top-down modulation of pathways that then drive outputs"
        }

        // todo: need a Rect pathway from MntD -> out if !dynMaint, or something else..

        if pfcOut != nil && pfcMnt != nil {
                pfcOut.PlaceRightOf(pfcMnt, 2)
                pt := nt.ConnectLayers(pfcMntD, pfcOut, paths.NewOneToOne(), ForwardPath)
                pt.AddClass("PFCMntDToOut")
        }

        return
}

// AddPBWM adds a DorsalBG and PFC with given params
// Defaults to simple case of basic maint dynamics in Deep
func (nt *Network) AddPBWM(prefix string, nY, nMaint, nOut, nNeurBgY, nNeurBgX, nNeurPfcY, nNeurPfcX int) (mtxGo, mtxNoGo, gpe, gpi, cin, pfcMnt, pfcMntD, pfcOut, pfcOutD *Layer) {
        mtxGo, mtxNoGo, gpe, gpi, cin = nt.AddDorsalBG(prefix, nY, nMaint, nOut, nNeurBgY, nNeurBgX)
        pfcMnt, pfcMntD, pfcOut, pfcOutD = nt.AddPFC(prefix, nY, nMaint, nOut, nNeurPfcY, nNeurPfcX, true) // default dynmaint
        if pfcMnt != nil {
                pfcMnt.PlaceAbove(mtxGo)
        }
        gpi.SendToMatrixPFC(prefix) // sends gating to all these layers
        return
}

// Copyright (c) 2024, The Emergent Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package leabra

import (
        "cogentcore.org/core/math32"
)

// Params for for trace-based learning in the MatrixTracePath
type TraceParams struct {

        // learning rate for all not-gated stripes, which learn in the opposite direction to the gated stripes, and typically with a slightly lower learning rate -- although there are different learning logics associated with each of these different not-gated cases, in practice the same learning rate for all works best, and is simplest
        NotGatedLR float32 `default:"0.7" min:"0"`

        // learning rate for gated, NoGo (D2), positive dopamine (weights decrease) -- this is the single most important learning parameter here -- by making this relatively small (but non-zero), an asymmetry in the role of Go vs. NoGo is established, whereby the NoGo pathway focuses largely on punishing and preventing actions associated with negative outcomes, while those assoicated with positive outcomes only very slowly get relief from this NoGo pressure -- this is critical for causing the model to explore other possible actions even when a given action SOMETIMES produces good results -- NoGo demands a very high, consistent level of good outcomes in order to have a net decrease in these avoidance weights.  Note that the gating signal applies to both Go and NoGo MSN's for gated stripes, ensuring learning is about the action that was actually selected (see not_ cases for logic for actions that were close but not taken)
        GateNoGoPosLR float32 `default:"0.1" min:"0"`

        // decay driven by receiving unit ACh value, sent by CIN units, for reseting the trace
        AChDecay float32 `min:"0" default:"0"`

        // multiplier on trace activation for decaying prior traces -- new trace magnitude drives decay of prior trace -- if gating activation is low, then new trace can be low and decay is slow, so increasing this factor causes learning to be more targeted on recent gating changes
        Decay float32 `min:"0" default:"1"`

        // use the sigmoid derivative factor 2 * act * (1-act) in modulating learning -- otherwise just multiply by msn activation directly -- this is generally beneficial for learning to prevent weights from continuing to increase when activations are already strong (and vice-versa for decreases)
        Deriv bool `default:"true"`
}

func (tp *TraceParams) Defaults() {
        tp.NotGatedLR = 0.7
        tp.GateNoGoPosLR = 0.1
        tp.AChDecay = 0 // not useful at all, surprisingly.
        tp.Decay = 1
        tp.Deriv = true
}

func (tp *TraceParams) Update() {
}

// LrnFactor resturns multiplicative factor for level of msn activation.  If Deriv
// is 2 * act * (1-act) -- the factor of 2 compensates for otherwise reduction in
// learning from these factors.  Otherwise is just act.
func (tp *TraceParams) LrnFactor(act float32) float32 {
        if !tp.Deriv {
                return act
        }
        return 2 * act * (1 - act)
}

// LrateMod returns the learning rate modulator based on gating, d2r, and posDa factors
func (tp *TraceParams) LrateMod(gated, d2r, posDa bool) float32 {
        if !gated {
                return tp.NotGatedLR
        }
        if d2r && posDa {
                return tp.GateNoGoPosLR
        }
        return 1
}

func (pt *Path) MatrixDefaults() {
        pt.Learn.WtSig.Gain = 1
        pt.Learn.Norm.On = false
        pt.Learn.Momentum.On = false
        pt.Learn.WtBal.On = false
}

func (pt *Path) ClearTrace() {
        for si := range pt.Syns {
                sy := &pt.Syns[si]
                sy.NTr = 0
                sy.Tr = 0
        }
}

// DWtMatrix computes the weight change (learning) for MatrixPath.
func (pt *Path) DWtMatrix() {
        slay := pt.Send
        rlay := pt.Recv
        d2r := (rlay.PBWM.DaR == D2R)
        da := rlay.NeuroMod.DA
        ach := rlay.NeuroMod.ACh
        gateActIdx, _ := NeuronVarIndexByName("GateAct")
        for si := range slay.Neurons {
                sn := &slay.Neurons[si]
                nc := int(pt.SConN[si])
                st := int(pt.SConIndexSt[si])
                syns := pt.Syns[st : st+nc]
                scons := pt.SConIndex[st : st+nc]

                for ci := range syns {
                        sy := &syns[ci]
                        ri := scons[ci]
                        rn := &rlay.Neurons[ri]
                        daLrn := rn.DALrn
                        // da := rlay.UnitValueByIndex(DA, int(ri)) // note: more efficient to just assume same for all units
                        // ach := rlay.UnitValueByIndex(ACh, int(ri))
                        gateAct := rlay.UnitValue1D(gateActIdx, int(ri), 0)
                        achDk := math32.Min(1, ach*pt.Trace.AChDecay)
                        tr := sy.Tr

                        dwt := float32(0)
                        if da != 0 {
                                dwt = daLrn * tr
                                if d2r && da > 0 && tr < 0 {
                                        dwt *= pt.Trace.GateNoGoPosLR
                                }
                        }

                        tr -= achDk * tr

                        newNTr := pt.Trace.LrnFactor(rn.Act) * sn.Act
                        ntr := float32(0)
                        if gateAct > 0 { // gated
                                ntr = newNTr
                        } else { // not-gated
                                ntr = -pt.Trace.NotGatedLR * newNTr // opposite sign for non-gated
                        }

                        decay := pt.Trace.Decay * math32.Abs(ntr) // decay is function of new trace
                        if decay > 1 {
                                decay = 1
                        }
                        tr += ntr - decay*tr
                        sy.Tr = tr
                        sy.NTr = ntr

                        sy.DWt += pt.Learn.Lrate * dwt
                }
        }
}

//////// DaHebbPath

func (pt *Path) DaHebbDefaults() {
        pt.Learn.WtSig.Gain = 1
        pt.Learn.Norm.On = false
        pt.Learn.Momentum.On = false
        pt.Learn.WtBal.On = false
}

// DWtDaHebb computes the weight change (learning), for [DaHebbPath].
func (pt *Path) DWtDaHebb() {
        slay := pt.Send
        rlay := pt.Recv
        for si := range slay.Neurons {
                sn := &slay.Neurons[si]
                nc := int(pt.SConN[si])
                st := int(pt.SConIndexSt[si])
                syns := pt.Syns[st : st+nc]
                scons := pt.SConIndex[st : st+nc]

                for ci := range syns {
                        sy := &syns[ci]
                        ri := scons[ci]
                        rn := &rlay.Neurons[ri]
                        da := rn.DALrn
                        dwt := da * rn.Act * sn.Act
                        sy.DWt += pt.Learn.Lrate * dwt
                }
        }
}

// Copyright (c) 2019, The Emergent Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package leabra

import (
        "cogentcore.org/core/math32/minmax"
        "github.com/emer/leabra/v2/fffb"
)

// Pool contains computed values for FFFB inhibition, and various other state values for layers
// and pools (unit groups) that can be subject to inhibition, including:
// * average / max stats on Ge and Act that drive inhibition
// * average activity overall that is used for normalizing netin (at layer level)
type Pool struct {

        // starting and ending (exlusive) indexes for the list of neurons in this pool
        StIndex, EdIndex int

        // FFFB inhibition computed values, including Ge and Act AvgMax which drive inhibition
        Inhib fffb.Inhib

        // minus phase average and max Act activation values, for ActAvg updt
        ActM minmax.AvgMax32

        // plus phase average and max Act activation values, for ActAvg updt
        ActP minmax.AvgMax32

        // running-average activation levels used for netinput scaling and adaptive inhibition
        ActAvg ActAvg

        //        Gate is gating state for PBWM layers
        Gate GateState
}

func (pl *Pool) Init() {
        pl.Inhib.Init()
        pl.Gate.Init()
}

// ActAvg are running-average activation levels used for netinput scaling and adaptive inhibition
type ActAvg struct {

        // running-average minus-phase activity -- used for adapting inhibition -- see ActAvgParams.Tau for time constant etc
        ActMAvg float32

        // running-average plus-phase activity -- used for synaptic input scaling -- see ActAvgParams.Tau for time constant etc
        ActPAvg float32

        // ActPAvg * ActAvgParams.Adjust -- adjusted effective layer activity directly used in synaptic input scaling
        ActPAvgEff float32
}

// Copyright (c) 2024, The Emergent Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package leabra

import (
        "fmt"

        "cogentcore.org/core/base/errors"
        "cogentcore.org/core/math32/minmax"
        "github.com/emer/emergent/v2/paths"
)

////////  RW

type RWParams struct {
        // PredRange is the range of predictions that can be represented by the [RWRewPredLayer].
        // Having a truncated range preserves some sensitivity in dopamine at the extremes
        // of good or poor performance.
        PredRange minmax.F32

        // RewLay is the reward layer name, for [RWDaLayer], from which DA is obtained.
        // If nothing clamped, no dopamine computed.
        RewLay string

        // PredLay is the name of [RWPredLayer] layer, for [RWDaLayer], that is used for
        // subtracting prediction from the reward value.
        PredLay string
}

func (rp *RWParams) Defaults() {
        rp.PredRange.Set(0.01, 0.99)
        rp.RewLay = "Rew"
        rp.PredLay = "RWPred"
}

func (rp *RWParams) Update() {
}

// ActFromGRWPred computes linear activation for [RWPredLayer].
func (ly *Layer) ActFromGRWPred(ctx *Context) {
        for ni := range ly.Neurons {
                nrn := &ly.Neurons[ni]
                if nrn.IsOff() {
                        continue
                }
                nrn.Act = ly.RW.PredRange.ClampValue(nrn.Ge) // clipped linear
                ly.Learn.AvgsFromAct(nrn)
        }
}

// RWLayers returns the reward and RWPredLayer layers based on names.
func (ly *Layer) RWLayers() (*Layer, *Layer, error) {
        tly := ly.Network.LayerByName(ly.RW.RewLay)
        if tly == nil {
                err := fmt.Errorf("RWDaLayer %s, RewLay: %q not found", ly.Name, ly.RW.RewLay)
                return nil, nil, errors.Log(err)
        }
        ply := ly.Network.LayerByName(ly.RW.PredLay)
        if ply == nil {
                err := fmt.Errorf("RWDaLayer %s, RWPredLay: %q not found", ly.Name, ly.RW.PredLay)
                return nil, nil, errors.Log(err)
        }
        return tly, ply, nil
}

func (ly *Layer) ActFromGRWDa(ctx *Context) {
        rly, ply, _ := ly.RWLayers()
        if rly == nil || ply == nil {
                return
        }
        rnrn := &(rly.Neurons[0])
        hasRew := false
        if rnrn.HasFlag(NeurHasExt) {
                hasRew = true
        }
        ract := rnrn.Act
        pnrn := &(ply.Neurons[0])
        pact := pnrn.Act
        for ni := range ly.Neurons {
                nrn := &ly.Neurons[ni]
                if nrn.IsOff() {
                        continue
                }
                if hasRew {
                        nrn.Act = ract - pact
                } else {
                        nrn.Act = 0 // nothing
                }
                ly.Learn.AvgsFromAct(nrn)
        }
}

// AddRWLayers adds simple Rescorla-Wagner (PV only) dopamine system, with a primary
// Reward layer, a RWPred prediction layer, and a dopamine layer that computes diff.
// Only generates DA when Rew layer has external input -- otherwise zero.
func (nt *Network) AddRWLayers(prefix string, space float32) (rew, rp, da *Layer) {
        rew = nt.AddLayer2D(prefix+"Rew", 1, 1, InputLayer)
        rp = nt.AddLayer2D(prefix+"RWPred", 1, 1, RWPredLayer)
        da = nt.AddLayer2D(prefix+"DA", 1, 1, RWDaLayer)
        da.RW.RewLay = rew.Name
        rp.PlaceBehind(rew, space)
        da.PlaceBehind(rp, space)

        rew.Doc = "Reward input, activated by external rewards, e.g., the US = unconditioned stimulus"
        rp.Doc = "Reward Prediction according to Rescorla-Wagner (RW) model, representing learned estimate of Rew layer activity on each trial, using linear activation function"
        da.Doc = "Dopamine (DA)-like signal reflecting the difference Rew - RWPred, or reward prediction error (RPE), on each trial"

        return
}

func (pt *Path) RWDefaults() {
        pt.Learn.WtSig.Gain = 1
        pt.Learn.Norm.On = false
        pt.Learn.Momentum.On = false
        pt.Learn.WtBal.On = false
}

// DWtRW computes the weight change (learning) for [RWPath].
func (pt *Path) DWtRW() {
        slay := pt.Send
        rlay := pt.Recv
        lda := rlay.NeuroMod.DA
        for si := range slay.Neurons {
                sn := &slay.Neurons[si]
                nc := int(pt.SConN[si])
                st := int(pt.SConIndexSt[si])
                syns := pt.Syns[st : st+nc]
                scons := pt.SConIndex[st : st+nc]

                for ci := range syns {
                        sy := &syns[ci]
                        ri := scons[ci]
                        rn := &rlay.Neurons[ri]

                        da := lda
                        if rn.Ge > rn.Act && da > 0 { // clipped at top, saturate up
                                da = 0
                        }
                        if rn.Ge < rn.Act && da < 0 { // clipped at bottom, saturate down
                                da = 0
                        }

                        dwt := da * sn.Act // no recv unit activation
                        sy.DWt += pt.Learn.Lrate * dwt
                }
        }
}

////////  TD

// TDParams are params for TD temporal differences computation.
type TDParams struct {

        // discount factor -- how much to discount the future prediction from RewPred.
        Discount float32

        // name of [TDPredLayer] to get reward prediction from.
        PredLay string

        // name of [TDIntegLayer] from which this computes the temporal derivative.
        IntegLay string
}

func (tp *TDParams) Defaults() {
        tp.Discount = 0.9
        tp.PredLay = "Pred"
        tp.IntegLay = "Integ"
}

func (tp *TDParams) Update() {
}

// ActFromGTDPred computes linear activation for [TDPredLayer].
func (ly *Layer) ActFromGTDPred(ctx *Context) {
        for ni := range ly.Neurons {
                nrn := &ly.Neurons[ni]
                if nrn.IsOff() {
                        continue
                }
                if ctx.Quarter == 3 { // plus phase
                        nrn.Act = nrn.Ge // linear
                } else {
                        nrn.Act = nrn.ActP // previous actP
                }
                ly.Learn.AvgsFromAct(nrn)
        }
}

func (ly *Layer) TDPredLayer() (*Layer, error) {
        tly := ly.Network.LayerByName(ly.TD.PredLay)
        if tly == nil {
                err := fmt.Errorf("TDIntegLayer %s RewPredLayer: %q not found", ly.Name, ly.TD.PredLay)
                return nil, errors.Log(err)
        }
        return tly, nil
}

func (ly *Layer) ActFromGTDInteg(ctx *Context) {
        rply, _ := ly.TDPredLayer()
        if rply == nil {
                return
        }
        rpActP := rply.Neurons[0].ActP
        rpAct := rply.Neurons[0].Act
        for ni := range ly.Neurons {
                nrn := &ly.Neurons[ni]
                if nrn.IsOff() {
                        continue
                }
                if ctx.Quarter == 3 { // plus phase
                        nrn.Act = nrn.Ge + ly.TD.Discount*rpAct
                } else {
                        nrn.Act = rpActP // previous actP
                }
                ly.Learn.AvgsFromAct(nrn)
        }
}

func (ly *Layer) TDIntegLayer() (*Layer, error) {
        tly := ly.Network.LayerByName(ly.TD.IntegLay)
        if tly == nil {
                err := fmt.Errorf("TDIntegLayer %s RewIntegLayer: %q not found", ly.Name, ly.TD.IntegLay)
                return nil, errors.Log(err)
        }
        return tly, nil
}

func (ly *Layer) TDDaDefaults() {
        ly.Act.Clamp.Range.Set(-100, 100)
}

func (ly *Layer) ActFromGTDDa(ctx *Context) {
        rily, _ := ly.TDIntegLayer()
        if rily == nil {
                return
        }
        rpActP := rily.Neurons[0].Act
        rpActM := rily.Neurons[0].ActM
        da := rpActP - rpActM
        for ni := range ly.Neurons {
                nrn := &ly.Neurons[ni]
                if nrn.IsOff() {
                        continue
                }
                if ctx.Quarter == 3 { // plus phase
                        nrn.Act = da
                } else {
                        nrn.Act = 0
                }
        }
}

func (pt *Path) TDPredDefaults() {
        pt.Learn.WtSig.Gain = 1
        pt.Learn.Norm.On = false
        pt.Learn.Momentum.On = false
        pt.Learn.WtBal.On = false
}

// DWtTDPred computes the weight change (learning) for [TDPredPath].
func (pt *Path) DWtTDPred() {
        slay := pt.Send
        rlay := pt.Recv
        da := rlay.NeuroMod.DA
        for si := range slay.Neurons {
                sn := &slay.Neurons[si]
                nc := int(pt.SConN[si])
                st := int(pt.SConIndexSt[si])
                syns := pt.Syns[st : st+nc]
                // scons := pj.SConIndex[st : st+nc]

                for ci := range syns {
                        sy := &syns[ci]
                        // ri := scons[ci]
                        dwt := da * sn.ActQ0 // no recv unit activation, prior trial act
                        sy.DWt += pt.Learn.Lrate * dwt
                }
        }
}

// AddTDLayers adds the standard TD temporal differences layers, generating a DA signal.
// Pathway from Rew to RewInteg is given class TDToInteg -- should
// have no learning and 1 weight.
func (nt *Network) AddTDLayers(prefix string, space float32) (rew, rp, ri, td *Layer) {
        rew = nt.AddLayer2D(prefix+"Rew", 1, 1, InputLayer)
        rp = nt.AddLayer2D(prefix+"Pred", 1, 1, TDPredLayer)
        ri = nt.AddLayer2D(prefix+"Integ", 1, 1, TDIntegLayer)
        td = nt.AddLayer2D(prefix+"TD", 1, 1, TDDaLayer)
        ri.TD.PredLay = rp.Name
        td.TD.IntegLay = ri.Name

        rp.PlaceBehind(rew, space)
        ri.PlaceBehind(rp, space)
        td.PlaceBehind(ri, space)

        pt := nt.ConnectLayers(rew, ri, paths.NewFull(), ForwardPath)
        pt.AddClass("TDToInteg")
        pt.Learn.Learn = false
        pt.WtInit.Mean = 1
        pt.WtInit.Var = 0
        pt.WtInit.Sym = false

        rew.Doc = "Reward input, activated by external rewards, e.g., the US = unconditioned stimulus"
        rp.Doc = "Reward Prediction, representing estimated value V(t) in the minus phase, and in plus phase computes estimated V(t+1) based on learned weights"
        ri.Doc = "Integration of Pred + Rew, representing estimated value V(t) in the minus phase, and estimated V(t+1) + r(t) in the plus phase"
        td.Doc = "Temporal Difference (TD) computes a dopamine (DA)-like signal as difference between the Integ activations across plus - minus phases: [V(t+1) + r(t)] - V(t), where V are estimated cumulative discounted future reward values"

        return
}

// Copyright (c) 2019, The Emergent Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package leabra

import (
        "fmt"
        "strings"
        "unsafe"

        "cogentcore.org/core/types"
)

// leabra.Synapse holds state for the synaptic connection between neurons
type Synapse struct {

        // synaptic weight value, sigmoid contrast-enhanced version
        // of the linear weight LWt.
        Wt float32

        // linear (underlying) weight value, which learns according
        // to the lrate specified in the connection spec.
        // This is converted into the effective weight value, Wt,
        // via sigmoidal contrast enhancement (see WtSigParams).
        LWt float32

        // change in synaptic weight, driven by learning algorithm.
        DWt float32

        // DWt normalization factor, reset to max of abs value of DWt,
        // decays slowly down over time. Serves as an estimate of variance
        // in weight changes over time.
        Norm float32

        // momentum, as time-integrated DWt changes, to accumulate a
        // consistent direction of weight change and cancel out
        // dithering contradictory changes.
        Moment float32

        // scaling parameter for this connection: effective weight value
        // is scaled by this factor in computing G conductance.
        // This is useful for topographic connectivity patterns e.g.,
        // to enforce more distant connections to always be lower in magnitude
        // than closer connections.  Value defaults to 1 (cannot be exactly 0,
        // otherwise is automatically reset to 1; use a very small number to
        // approximate 0). Typically set by using the paths.Pattern Weights()
        // values where appropriate.
        Scale float32

        // NTr is the new trace, which drives updates to trace value.
        // su * (1-ru_msn) for gated, or su * ru_msn for not-gated (or for non-thalamic cases).
        NTr float32

        // Tr is the current ongoing trace of activations, which drive learning.
        // Adds NTr and clears after learning on current values, and includes both
        // thal gated (+ and other nongated, - inputs).
        Tr float32
}

func (sy *Synapse) VarNames() []string {
        return SynapseVars
}

var SynapseVars = []string{"Wt", "LWt", "DWt", "Norm", "Moment", "Scale", "NTr", "Tr"}

var SynapseVarProps = map[string]string{
        "Wt":     `cat:"Wts"`,
        "LWt":    `cat:"Wts"`,
        "DWt":    `cat:"Wts" auto-scale:"+"`,
        "Norm":   `cat:"Wts"`,
        "Moment": `cat:"Wts" auto-scale:"+"`,
        "Scale":  `cat:"Wts"`,
        "NTr":    `cat:"Wts"`,
        "Tr":     `cat:"Wts"`,
}

var SynapseVarsMap map[string]int

func init() {
        SynapseVarsMap = make(map[string]int, len(SynapseVars))
        for i, v := range SynapseVars {
                SynapseVarsMap[v] = i
        }
        styp := types.For[Synapse]()
        for _, fld := range styp.Fields {
                tag := SynapseVarProps[fld.Name]
                SynapseVarProps[fld.Name] = tag + ` doc:"` + strings.ReplaceAll(fld.Doc, "\n", " ") + `"`
        }
}

// SynapseVarByName returns the index of the variable in the Synapse, or error
func SynapseVarByName(varNm string) (int, error) {
        i, ok := SynapseVarsMap[varNm]
        if !ok {
                return -1, fmt.Errorf("Synapse VarByName: variable name: %v not valid", varNm)
        }
        return i, nil
}

// VarByIndex returns variable using index (0 = first variable in SynapseVars list)
func (sy *Synapse) VarByIndex(idx int) float32 {
        fv := (*float32)(unsafe.Pointer(uintptr(unsafe.Pointer(sy)) + uintptr(4*idx)))
        return *fv
}

// VarByName returns variable by name, or error
func (sy *Synapse) VarByName(varNm string) (float32, error) {
        i, err := SynapseVarByName(varNm)
        if err != nil {
                return 0, err
        }
        return sy.VarByIndex(i), nil
}

func (sy *Synapse) SetVarByIndex(idx int, val float32) {
        fv := (*float32)(unsafe.Pointer(uintptr(unsafe.Pointer(sy)) + uintptr(4*idx)))
        *fv = val
}

// SetVarByName sets synapse variable to given value
func (sy *Synapse) SetVarByName(varNm string, val float32) error {
        i, err := SynapseVarByName(varNm)
        if err != nil {
                return err
        }
        sy.SetVarByIndex(i, val)
        return nil
}

// Copyright (c) 2019, The Emergent Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

/*
Package nxx1 provides the Noisy-X-over-X-plus-1 activation function that well-characterizes
the neural response function empirically, as a saturating sigmoid-like nonlinear response
with an initial largely linear regime.

The basic x/(x+1) sigmoid function is convolved with a gaussian noise kernel to produce
a better approximation of the effects of noise on neural firing -- the main effect is
to create a continuous graded early level of firing even slightly below threshold, softening
the otherwise hard transition to firing at threshold.

A hand-optimized piece-wise function approximation is used to generate the NXX1 function
instead of requiring a lookup table of the gaussian convolution.  This is much easier
to use across a range of computational platforms including GPU's, and produces very similar
overall values.
*/
package nxx1

//go:generate core generate -add-types

import "cogentcore.org/core/math32"

// Params are the Noisy X/(X+1) rate-coded activation function parameters.
// This function well-characterizes the neural response function empirically,
// as a saturating sigmoid-like nonlinear response with an initial largely linear regime.
// The basic x/(x+1) sigmoid function is convolved with a gaussian noise kernel to produce
// a better approximation of the effects of noise on neural firing -- the main effect is
// to create a continuous graded early level of firing even slightly below threshold, softening
// the otherwise hard transition to firing at threshold.
// A hand-optimized piece-wise function approximation is used to generate the NXX1 function
// instead of requiring a lookup table of the gaussian convolution.  This is much easier
// to use across a range of computational platforms including GPU's, and produces very similar
// overall values.  abc.
type Params struct {

        // threshold value Theta (Q) for firing output activation (.5 is more accurate value based on AdEx biological parameters and normalization
        Thr float32 `default:"0.5"`

        // gain (gamma) of the rate-coded activation functions -- 100 is default, 80 works better for larger models, and 20 is closer to the actual spiking behavior of the AdEx model -- use lower values for more graded signals, generally in lower input/sensory layers of the network
        Gain float32 `default:"80,100,40,20" min:"0"`

        // variance of the Gaussian noise kernel for convolving with XX1 in NOISY_XX1 and NOISY_LINEAR -- determines the level of curvature of the activation function near the threshold -- increase for more graded responding there -- note that this is not actual stochastic noise, just constant convolved gaussian smoothness to the activation function
        NVar float32 `default:"0.005,0.01" min:"0"`

        // threshold on activation below which the direct vm - act.thr is used -- this should be low -- once it gets active should use net - g_e_thr ge-linear dynamics (gelin)
        VmActThr float32 `default:"0.01"`

        // multiplier on sigmoid used for computing values for net < thr
        SigMult float32 `default:"0.33" display:"-" json:"-" xml:"-"`

        // power for computing sig_mult_eff as function of gain * nvar
        SigMultPow float32 `default:"0.8" display:"-" json:"-" xml:"-"`

        // gain multipler on (net - thr) for sigmoid used for computing values for net < thr
        SigGain float32 `default:"3" display:"-" json:"-" xml:"-"`

        // interpolation range above zero to use interpolation
        InterpRange float32 `default:"0.01" display:"-" json:"-" xml:"-"`

        // range in units of nvar over which to apply gain correction to compensate for convolution
        GainCorRange float32 `default:"10" display:"-" json:"-" xml:"-"`

        // gain correction multiplier -- how much to correct gains
        GainCor float32 `default:"0.1" display:"-" json:"-" xml:"-"`

        // sig_gain / nvar
        SigGainNVar float32 `display:"-" json:"-" xml:"-"`

        // overall multiplier on sigmoidal component for values below threshold = sig_mult * pow(gain * nvar, sig_mult_pow)
        SigMultEff float32 `display:"-" json:"-" xml:"-"`

        // 0.5 * sig_mult_eff -- used for interpolation portion
        SigValAt0 float32 `display:"-" json:"-" xml:"-"`

        // function value at interp_range - sig_val_at_0 -- for interpolation
        InterpVal float32 `display:"-" json:"-" xml:"-"`
}

func (xp *Params) Update() {
        xp.SigGainNVar = xp.SigGain / xp.NVar
        xp.SigMultEff = xp.SigMult * math32.Pow(xp.Gain*xp.NVar, xp.SigMultPow)
        xp.SigValAt0 = 0.5 * xp.SigMultEff
        xp.InterpVal = xp.XX1GainCor(xp.InterpRange) - xp.SigValAt0
}

func (xp *Params) Defaults() {
        xp.Thr = 0.5
        xp.Gain = 100
        xp.NVar = 0.005
        xp.VmActThr = 0.01
        xp.SigMult = 0.33
        xp.SigMultPow = 0.8
        xp.SigGain = 3.0
        xp.InterpRange = 0.01
        xp.GainCorRange = 10.0
        xp.GainCor = 0.1
        xp.Update()
}

// XX1 computes the basic x/(x+1) function
func (xp *Params) XX1(x float32) float32 { return x / (x + 1) }

// XX1GainCor computes x/(x+1) with gain correction within GainCorRange
// to compensate for convolution effects
func (xp *Params) XX1GainCor(x float32) float32 {
        gainCorFact := (xp.GainCorRange - (x / xp.NVar)) / xp.GainCorRange
        if gainCorFact < 0 {
                return xp.XX1(xp.Gain * x)
        }
        newGain := xp.Gain * (1 - xp.GainCor*gainCorFact)
        return xp.XX1(newGain * x)
}

// NoisyXX1 computes the Noisy x/(x+1) function -- directly computes close approximation
// to x/(x+1) convolved with a gaussian noise function with variance nvar.
// No need for a lookup table -- very reasonable approximation for standard range of parameters
// (nvar = .01 or less -- higher values of nvar are less accurate with large gains,
// but ok for lower gains)
func (xp *Params) NoisyXX1(x float32) float32 {
        if x < 0 { // sigmoidal for < 0
                ex := -(x * xp.SigGainNVar)
                if ex > 50 {
                        return 0
                }
                return xp.SigMultEff / (1 + math32.FastExp(ex))
        } else if x < xp.InterpRange {
                interp := 1 - ((xp.InterpRange - x) / xp.InterpRange)
                return xp.SigValAt0 + interp*xp.InterpVal
        } else {
                return xp.XX1GainCor(x)
        }
}

// X11GainCorGain computes x/(x+1) with gain correction within GainCorRange
// to compensate for convolution effects -- using external gain factor
func (xp *Params) XX1GainCorGain(x, gain float32) float32 {
        gainCorFact := (xp.GainCorRange - (x / xp.NVar)) / xp.GainCorRange
        if gainCorFact < 0 {
                return xp.XX1(gain * x)
        }
        newGain := gain * (1 - xp.GainCor*gainCorFact)
        return xp.XX1(newGain * x)
}

// NoisyXX1Gain computes the noisy x/(x+1) function -- directly computes close approximation
// to x/(x+1) convolved with a gaussian noise function with variance nvar.
// No need for a lookup table -- very reasonable approximation for standard range of parameters
// (nvar = .01 or less -- higher values of nvar are less accurate with large gains,
// but ok for lower gains).  Using external gain factor.
func (xp *Params) NoisyXX1Gain(x, gain float32) float32 {
        if x < xp.InterpRange {
                sigMultEffArg := xp.SigMult * math32.Pow(gain*xp.NVar, xp.SigMultPow)
                sigValAt0Arg := 0.5 * sigMultEffArg

                if x < 0 { // sigmoidal for < 0
                        ex := -(x * xp.SigGainNVar)
                        if ex > 50 {
                                return 0
                        }
                        return sigMultEffArg / (1 + math32.FastExp(ex))
                } else { // else x < interp_range
                        interp := 1 - ((xp.InterpRange - x) / xp.InterpRange)
                        return sigValAt0Arg + interp*xp.InterpVal
                }
        } else {
                return xp.XX1GainCorGain(x, gain)
        }
}

// Copyright (c) 2019, The Emergent Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

//go:generate core generate -add-types

package spike

import (
        "cogentcore.org/core/math32"
        "github.com/emer/leabra/v2/leabra"
)

// ActParams is full set of activation params including those from base
// leabra and the additional Spiking-specific ones.
type ActParams struct {
        leabra.ActParams

        // spiking parameters
        Spike SpikeParams `display:"inline"`
}

func (sk *ActParams) Defaults() {
        sk.ActParams.Defaults()
        sk.Spike.Defaults()
}

func (sk *ActParams) Update() {
        sk.ActParams.Update()
        sk.Spike.Update()
}

// CopyFromAct copies ActParams from source (e.g., rate-code params)
func (sk *ActParams) CopyFromAct(act *leabra.ActParams) {
        sk.ActParams = *act
        sk.Update()
}

func (sk *ActParams) SpikeVmFromG(nrn *leabra.Neuron) {
        updtVm := true
        if sk.Spike.Tr > 0 && nrn.ISI >= 0 && nrn.ISI < float32(sk.Spike.Tr) {
                updtVm = false // don't update the spiking vm during refract
        }

        nwVm := nrn.Vm
        if updtVm {
                ge := nrn.Ge * sk.Gbar.E
                gi := nrn.Gi * sk.Gbar.I
                gk := sk.Gbar.K * (nrn.GknaFast + nrn.GknaMed + nrn.GknaSlow)
                nrn.Gk = gk
                vmEff := nrn.Vm
                // midpoint method: take a half-step in vmEff
                inet1 := sk.InetFromG(vmEff, ge, gi, gk)
                vmEff += .5 * sk.Dt.VmDt * inet1 // go half way
                inet2 := sk.InetFromG(vmEff, ge, gi, gk)
                // add spike current if relevant
                if sk.Spike.Exp {
                        inet2 += sk.Gbar.L * sk.Spike.ExpSlope *
                                math32.Exp((vmEff-sk.XX1.Thr)/sk.Spike.ExpSlope)
                }
                nwVm += sk.Dt.VmDt * inet2
                nrn.Inet = inet2
        }

        if sk.Noise.Type == leabra.VmNoise {
                nwVm += nrn.Noise
        }
        nrn.Vm = sk.VmRange.ClampValue(nwVm)
}

// SpikeActFromVm computes the discrete spiking activation
// from membrane potential Vm
func (sk *ActParams) SpikeActFromVm(nrn *leabra.Neuron) {
        var thr float32
        if sk.Spike.Exp {
                thr = sk.Spike.ExpThr
        } else {
                thr = sk.XX1.Thr
        }
        if nrn.Vm > thr {
                nrn.Spike = 1
                nrn.Vm = sk.Spike.VmR
                nrn.Inet = 0
                if nrn.ISIAvg == -1 {
                        nrn.ISIAvg = -2
                } else if nrn.ISI > 0 { // must have spiked to update
                        sk.Spike.AvgFromISI(&nrn.ISIAvg, nrn.ISI+1)
                }
                nrn.ISI = 0
        } else {
                nrn.Spike = 0
                if nrn.ISI >= 0 {
                        nrn.ISI += 1
                }
                if nrn.ISIAvg >= 0 && nrn.ISI > 0 && nrn.ISI > 1.2*nrn.ISIAvg {
                        sk.Spike.AvgFromISI(&nrn.ISIAvg, nrn.ISI)
                }
        }

        nwAct := sk.Spike.ActFromISI(nrn.ISIAvg, .001, 1) // todo: use real #'s
        if nwAct > 1 {
                nwAct = 1
        }
        nwAct = nrn.Act + sk.Dt.VmDt*(nwAct-nrn.Act)
        nrn.ActDel = nwAct - nrn.Act
        nrn.Act = nwAct
        if sk.KNa.On {
                sk.KNa.GcFromSpike(&nrn.GknaFast, &nrn.GknaMed, &nrn.GknaSlow, nrn.Spike > .5)
        }
}

// SpikeParams contains spiking activation function params.
// Implements the AdEx adaptive exponential function
type SpikeParams struct {

        // if true, turn on exponential excitatory current that drives Vm rapidly upward for spiking as it gets past its nominal firing threshold (Thr) -- nicely captures the Hodgkin Huxley dynamics of Na and K channels -- uses Brette & Gurstner 2005 AdEx formulation -- this mechanism has an unfortunate interaction with the continuous inhibitory currents generated by the standard FFFB inhibitory function, which cause this mechanism to desensitize and fail to spike
        Exp bool `default:"false"`

        // slope in Vm (2 mV = .02 in normalized units) for extra exponential excitatory current that drives Vm rapidly upward for spiking as it gets past its nominal firing threshold (Thr) -- nicely captures the Hodgkin Huxley dynamics of Na and K channels -- uses Brette & Gurstner 2005 AdEx formulation -- a value of 0 disables this mechanism
        ExpSlope float32 `default:"0.02"`

        // membrane potential threshold for actually triggering a spike when using the exponential mechanism
        ExpThr float32 `default:"1.2"`

        // post-spiking membrane potential to reset to, produces refractory effect if lower than VmInit -- 0.30 is appropriate biologically based value for AdEx (Brette & Gurstner, 2005) parameters
        VmR float32 `default:"0.3,0,0.15"`

        // post-spiking explicit refractory period, in cycles -- prevents Vm updating for this number of cycles post firing
        Tr int `default:"3"`

        // for translating spiking interval (rate) into rate-code activation equivalent (and vice-versa, for clamped layers), what is the maximum firing rate associated with a maximum activation value (max act is typically 1.0 -- depends on act_range)
        MaxHz float32 `default:"180" min:"1"`

        // constant for integrating the spiking interval in estimating spiking rate
        RateTau float32 `default:"5" min:"1"`

        // rate = 1 / tau
        RateDt float32 `display:"-"`
}

func (sk *SpikeParams) Defaults() {
        sk.Exp = false
        sk.ExpSlope = 0.02
        sk.ExpThr = 1.2
        sk.VmR = 0.3
        sk.Tr = 3
        sk.MaxHz = 180
        sk.RateTau = 5
        sk.Update()
}

func (sk *SpikeParams) Update() {
        sk.RateDt = 1 / sk.RateTau
}

func (sk *SpikeParams) ShouldDisplay(field string) bool {
        switch field {
        case "ExpSlope", "ExpThr":
                return sk.Exp
        default:
                return true
        }
}

// ActToISI compute spiking interval from a given rate-coded activation,
// based on time increment (.001 = 1msec default), Act.Dt.Integ
func (sk *SpikeParams) ActToISI(act, timeInc, integ float32) float32 {
        if act == 0 {
                return 0
        }
        return (1 / (timeInc * integ * act * sk.MaxHz))
}

// ActFromISI computes rate-code activation from estimated spiking interval
func (sk *SpikeParams) ActFromISI(isi, timeInc, integ float32) float32 {
        if isi <= 0 {
                return 0
        }
        maxInt := 1.0 / (timeInc * integ * sk.MaxHz) // interval at max hz..
        return maxInt / isi                          // normalized
}

// AvgFromISI updates spiking ISI from current isi interval value
func (sk *SpikeParams) AvgFromISI(avg *float32, isi float32) {
        if *avg <= 0 {
                *avg = isi
        } else if isi < 0.8**avg {
                *avg = isi // if significantly less than we take that
        } else { // integrate on slower
                *avg += sk.RateDt * (isi - *avg) // running avg updt
        }
}