using Pkg; Pkg.activate("."); Pkg.resolve(), Pkg.instantiate()
using Revise
using Flux  #, CUDA
using BSON, JSON3
using MLDatasets: MNIST
using MLUtils, Images, ProgressMeter, Dates, DataFrames, Random, Statistics, LinearAlgebra, 
        BenchmarkTools, Serialization, OneHotArrays , GLMakie   # ClickHouse


# # if one need to reinstall all python packages
# # try Pkg.rm("PythonCall") catch end  # should be removed before using CondaPkg to install packages
# condapackage = ["numpy", "pytorch", "snntorch"]
# using CondaPkg  # in CondaPkg.toml file, channels = ["anaconda", "conda-forge", "pytorch"]
# CondaPkg.add_channel("pytorch")
# for i in condapackage
#     try CondaPkg.rm(i) catch end
# end
# for i in condapackage
#     CondaPkg.add(i)
# end

using PythonCall;
np = pyimport("numpy")
torch = pyimport("torch")
spikegen = pyimport("snntorch.spikegen")    # https://github.com/jeshraghian/snntorch

using Ironpen
using GeneralUtils

sep = Sys.iswindows() ? "\\" : "/"
rootDir = pwd()

# select compute device
# device = Flux.CUDA.functional() ? gpu : cpu     # Flux provide "cpu" and "gpu" keywork
# if device == gpu    CUDA.device!(3) end
# CUDA.allowscalar(false) # turn off scalar indexing
#------------------------------------------------------------------------------------------------100



""" 
    Todo:
        - [] 

    Change from version:
        - 
    
    All features
        - 
"""


# communication config --------------------------------------------------------------------------100

database_ip = "localhost"
# database_ip = "192.168.0.8"

#------------------------------------------------------------------------------------------------100

function generate_snn(filename::String, location::String)
    expect_compute_neuron_numbers = 1024  
    signalInput_portnumbers = 50
    noise_portnumbers = 1   #signalInput_portnumbers
    output_portnumbers = 10

    lif_neuron_number = Int(floor(expect_compute_neuron_numbers * 0.6))
    alif_neuron_number = expect_compute_neuron_numbers - lif_neuron_number   # from Allen Institute, ALIF is 20-40% of LIF
    computeNeuronNumber = lif_neuron_number + alif_neuron_number

    totalNeurons = computeNeuronNumber + noise_portnumbers + signalInput_portnumbers
    totalInputPort = noise_portnumbers + signalInput_portnumbers

    # kfn and neuron config 
    passthrough_neuron_params = Dict(
        :type => "passthroughNeuron"
    )

    lif_neuron_params = Dict{Symbol, Any}(
        :type => "lifNeuron",
        :v_t_default => 0.0,
        :v_th => 1.0,   # neuron firing threshold (this value is treated as maximum bound if I use auto generate)
        :tau_m => 20.0,     # membrane time constant in millisecond. 
        :eta => 1e-6,
        # Good starting value is 1/10th of tau_a
        # This is problem specific parameter. It controls how leaky the neuron is.
        # Too high(less leaky) makes learning algo harder to move model into direction that reduce error
        # resulting in model's error to explode exponantially likely because learning algo will try to 
        # exert more force (larger w_out_change) to move neuron into direction that reduce error
        # For example, model error from 7 to 2e6.

        :synapticConnectionPercent => 20,   # % coverage of total neurons in kfn
        :w_rec_generation_pattern => "random",
    )

    alif_neuron_params = Dict{Symbol, Any}(
        :type => "alifNeuron",
        :v_t_default => 0.0,
        :v_th => 1.0,   # neuron firing threshold (this value is treated as maximum bound if I use auto generate)
        :tau_m => 20.0,     # membrane time constant in millisecond.
        :eta => 1e-6,
        # Good starting value is 1/10th of tau_a
        # This is problem specific parameter. It controls how leaky the neuron is.
        # Too high(less leaky) makes learning algo harder to move model into direction that reduce error
        # resulting in model's error to explode exponantially likely because learning algo will try to 
        # exert more force (larger w_out_change) to move neuron into direction that reduce error
        # For example, model error from 7 to 2e6.

        :tau_a => 800.0,    # adaptation time constant in millisecond. it defines neuron memory length. 
        # This is problem specific parameter
        # Good starting value is 0.5 to 2 times of info STORE-RECALL length i.e. total time SNN takes to 
        # perform a task, for example, equals to episode length.  
        # From "Spike frequency adaptation supports network computations on temporally dispersed 
        # information"

        :synapticConnectionPercent => 20,   # % coverage of total neurons in kfn
        :w_rec_generation_pattern => "random",
    )

    linear_neuron_params = Dict{Symbol, Any}(
        :type => "linearNeuron",
        :v_th => 1.0,   # neuron firing threshold (this value is treated as maximum bound if I use auto generate)
        :tau_out => 50.0,   # output time constant in millisecond. 
        :synapticConnectionPercent => 20,  # % coverage of total neurons in kfn
        # Good starting value is 1/50th of tau_a
        # This is problem specific parameter. 
        # It controls how leaky the neuron is.
        # Too high(less leaky) makes learning algo harder to move model into direction that reduce error
        # resulting in model's error to explode exponantially. For example, model error from 7 to 2e6
        # One can image training output neuron is like Tetris Game.
    )

    integrate_neuron_params = Dict{Symbol, Any}(
        :type => "integrateNeuron",
        :synapticConnectionPercent => 100,  # % coverage of total neurons in kfn
        :eta => 1e-6,
        :tau_out => 50.0,
        # Good starting value is 1/50th of tau_a
        # This is problem specific parameter. 
        # It controls how leaky the neuron is.
        # Too high(less leaky) makes learning algo harder to move model into direction that reduce error
        # resulting in model's error to explode exponantially. For example, model error from 7 to 2e6
        # One can image training output neuron is like Tetris Game.
    )

    I_kfnparams = Dict{Symbol, Any}(
        :knowledgeFnName=> "I",
        :computeNeuronNumber=> computeNeuronNumber,
        :neuronFiringRateTarget=> 20.0,    # Hz
        :Bn=> "random",   # error projection coefficient for EACH neuron
        :totalNeurons=> totalNeurons,
        :totalInputPort=> totalInputPort,
        :totalComputeNeuron=> computeNeuronNumber,

        # group relavent info
        :inputPort=> Dict(
            :noise=> Dict(
                :numbers=> noise_portnumbers,
                :params=> passthrough_neuron_params,
            ),
            :signal=> Dict(
                :numbers=> signalInput_portnumbers,   # in case of GloVe word encoding, it is 300
                :params=> passthrough_neuron_params,
            ),
        ),
        :outputPort=> Dict(
            :numbers=> output_portnumbers,  # output neuron, this is also the output length
            :params=> integrate_neuron_params,
        ),
        :computeNeuron=> Dict(
            :1=> Dict(
                :numbers=> lif_neuron_number,
                :params=> lif_neuron_params,
            ),
            :2=> Dict(
                :numbers=> alif_neuron_number,
                :params=> alif_neuron_params,
            ),
        ),
    )

    #------------------------------------------------------------------------------------------------100

    I_kfn = Ironpen.kfn_1(I_kfnparams)

    model_params_1 = Dict(:knowledgeFn => Dict(
        :I => I_kfn),
        )

    model = Ironpen.model(model_params_1)

    serialize(location * sep * filename, model)
    println("SNN generated")
end

function data_loader()
    # test problem
    fullTrainDataset = MNIST(:train)
    prototypeDataset = fullTrainDataset[1:10] # use reshape(test_dataset[1], (:, 1)) to flaten matrix
    trainDataset = fullTrainDataset  # total 60000
    validateDataset = fullTrainDataset[1:100]
    labelDict = [0:9...]

    trainData = MLUtils.DataLoader(
        trainDataset;   # fullTrainDataset or trainDataset
        batchsize=100,
        collate=true,
        shuffle=true,
        buffer=true,
        partial=false,      # better for gpu memory if batchsize is fixed
        # parallel=true,    #BUG ?? causing dataloader into forever loop
    )

    validateData = MLUtils.DataLoader(
        validateDataset;
        batchsize=1,
        collate=true,
        shuffle=true,
        buffer=true,
        partial=false,      # better for gpu memory if batchsize is fixed
        # parallel=true,    #BUG ?? causing dataloader into forever loop
    )

    #CHANGE dummy data used to debug
    # trainData = [(rand(10, 10), [5]), (rand(10, 10), [2])]
    # trainData = [(rand(10, 10), [5]),]

    return trainData, validateData, labelDict
end

function train_snn(model_name::String, filename::String, location::String, 
                    trainData, validateData, labelDict::Vector)
    println("loading SNN model")

    model = deserialize(location * sep * filename)
    println("model loading completed")

    # random seed
    # rng = MersenneTwister(1234)

    logitLog = zeros(10, 2)
    firedNeurons_t1 = zeros(1)
    var1 = zeros(10, 2)
    var2 = zeros(10, 2)
    var3 = zeros(10, 2)
    var4 = zeros(10, 2)

    # ----------------------------------- plot ----------------------------------- #
    plot10 = Observable(firedNeurons_t1)

    plot20 = Observable(logitLog[1 , :])
    plot21 = Observable(logitLog[2 , :])
    plot22 = Observable(logitLog[3 , :])
    plot23 = Observable(logitLog[4 , :])
    plot24 = Observable(logitLog[5 , :])
    plot25 = Observable(logitLog[6 , :])
    plot26 = Observable(logitLog[7 , :])
    plot27 = Observable(logitLog[8 , :])
    plot28 = Observable(logitLog[9 , :])
    plot29 = Observable(logitLog[10, :])

    plot30 = Observable(var1[1 , :])
    plot31 = Observable(var1[2 , :])
    plot32 = Observable(var1[3 , :])
    plot33 = Observable(var1[4 , :])
    plot34 = Observable(var1[5 , :])
    plot35 = Observable(var1[6 , :])
    plot36 = Observable(var1[7 , :])
    plot37 = Observable(var1[8 , :])
    plot38 = Observable(var1[9 , :])
    plot39 = Observable(var1[10, :])
    
    plot40 = Observable(var2[1 , :])
    plot41 = Observable(var2[2 , :])
    plot42 = Observable(var2[3 , :])
    plot43 = Observable(var2[4 , :])
    plot44 = Observable(var2[5 , :])
    plot45 = Observable(var2[6 , :])
    plot46 = Observable(var2[7 , :])
    plot47 = Observable(var2[8 , :])
    plot48 = Observable(var2[9 , :])
    plot49 = Observable(var2[10, :])

    plot50 = Observable(var3[1 , :])
    plot51 = Observable(var3[2 , :])
    plot52 = Observable(var3[3 , :])
    plot53 = Observable(var3[4 , :])
    plot54 = Observable(var3[5 , :])
    plot55 = Observable(var3[6 , :])
    plot56 = Observable(var3[7 , :])
    plot57 = Observable(var3[8 , :])
    plot58 = Observable(var3[9 , :])
    plot59 = Observable(var3[10, :])

    plot60 = Observable(var4[1 , :])
    plot61 = Observable(var4[2 , :])
    plot62 = Observable(var4[3 , :])
    plot63 = Observable(var4[4 , :])
    plot64 = Observable(var4[5 , :])
    plot65 = Observable(var4[6 , :])
    plot66 = Observable(var4[7 , :])
    plot67 = Observable(var4[8 , :])
    plot68 = Observable(var4[9 , :])
    plot69 = Observable(var4[10, :])

    # main figure
    fig1 = Figure()

    subfig1 = GLMakie.Axis(fig1[1, 1],  # define position of this subfigure inside a figure
        title = "RSNN firedNeurons_t1",
        xlabel = "time",
        ylabel = "data"
    )
    lines!(subfig1, plot10, label = "firedNeurons_t1")
    axislegend(subfig1, position = :lb)
    
    subfig2 = GLMakie.Axis(fig1[2, 1],  # define position of this subfigure inside a figure
        title = "output neurons logit",
        xlabel = "time",
        ylabel = "data"
    )
    
    lines!(subfig2, plot20, label = "0", color = 1, colormap = :tab10, colorrange = (1, 10) )
    lines!(subfig2, plot21, label = "1", color = 2, colormap = :tab10, colorrange = (1, 10) )
    lines!(subfig2, plot22, label = "2", color = 3, colormap = :tab10, colorrange = (1, 10) )
    lines!(subfig2, plot23, label = "3", color = 4, colormap = :tab10, colorrange = (1, 10) )
    lines!(subfig2, plot24, label = "4", color = 5, colormap = :tab10, colorrange = (1, 10) )
    lines!(subfig2, plot25, label = "5", color = 6, colormap = :tab10, colorrange = (1, 10) )
    lines!(subfig2, plot26, label = "6", color = 7, colormap = :tab10, colorrange = (1, 10) )
    lines!(subfig2, plot27, label = "7", color = 8, colormap = :tab10, colorrange = (1, 10) )
    lines!(subfig2, plot28, label = "8", color = 9, colormap = :tab10, colorrange = (1, 10) )
    lines!(subfig2, plot29, label = "9", color = 10, colormap = :tab10, colorrange = (1, 10))
    axislegend(subfig2, position = :lb)


    subfig3 = GLMakie.Axis(fig1[3, 1],  # define position of this subfigure inside a figure
        title = "last RSNN wRec",
        xlabel = "time",
        ylabel = "data"
    )
    lines!(subfig3, plot30, label = "0", color = 1, colormap = :tab10, colorrange = (1, 10) )
    lines!(subfig3, plot31, label = "1", color = 2, colormap = :tab10, colorrange = (1, 10) )
    lines!(subfig3, plot32, label = "2", color = 3, colormap = :tab10, colorrange = (1, 10) )
    lines!(subfig3, plot33, label = "3", color = 4, colormap = :tab10, colorrange = (1, 10) )
    lines!(subfig3, plot34, label = "4", color = 5, colormap = :tab10, colorrange = (1, 10) )
    lines!(subfig3, plot35, label = "5", color = 6, colormap = :tab10, colorrange = (1, 10) )
    lines!(subfig3, plot36, label = "6", color = 7, colormap = :tab10, colorrange = (1, 10) )
    lines!(subfig3, plot37, label = "7", color = 8, colormap = :tab10, colorrange = (1, 10) )
    lines!(subfig3, plot38, label = "8", color = 9, colormap = :tab10, colorrange = (1, 10) )
    lines!(subfig3, plot39, label = "9", color = 10, colormap = :tab10, colorrange = (1, 10))
    axislegend(subfig3, position = :lb)

    subfig4 = GLMakie.Axis(fig1[4, 1],  # define position of this subfigure inside a figure
        title = "RSNN v_t1",
        xlabel = "time",
        ylabel = "data"
    )
    lines!(subfig4, plot40, label = "0", color = 1, colormap = :tab10, colorrange = (1, 10) )
    lines!(subfig4, plot41, label = "1", color = 2, colormap = :tab10, colorrange = (1, 10) )
    lines!(subfig4, plot42, label = "2", color = 3, colormap = :tab10, colorrange = (1, 10) )
    lines!(subfig4, plot43, label = "3", color = 4, colormap = :tab10, colorrange = (1, 10) )
    lines!(subfig4, plot44, label = "4", color = 5, colormap = :tab10, colorrange = (1, 10) )
    lines!(subfig4, plot45, label = "5", color = 6, colormap = :tab10, colorrange = (1, 10) )
    lines!(subfig4, plot46, label = "6", color = 7, colormap = :tab10, colorrange = (1, 10) )
    lines!(subfig4, plot47, label = "7", color = 8, colormap = :tab10, colorrange = (1, 10) )
    lines!(subfig4, plot48, label = "8", color = 9, colormap = :tab10, colorrange = (1, 10) )
    lines!(subfig4, plot49, label = "9", color = 10, colormap = :tab10, colorrange = (1, 10))
    axislegend(subfig4, position = :lb)

    subfig5 = GLMakie.Axis(fig1[5, 1],  # define position of this subfigure inside a figure
        title = "output neuron epsilonRec",
        xlabel = "time",
        ylabel = "data"
    )
    lines!(subfig5, plot50, label = "0", color = 1, colormap = :tab10, colorrange = (1, 10) )
    lines!(subfig5, plot51, label = "1", color = 2, colormap = :tab10, colorrange = (1, 10) )
    lines!(subfig5, plot52, label = "2", color = 3, colormap = :tab10, colorrange = (1, 10) )
    lines!(subfig5, plot53, label = "3", color = 4, colormap = :tab10, colorrange = (1, 10) )
    lines!(subfig5, plot54, label = "4", color = 5, colormap = :tab10, colorrange = (1, 10) )
    lines!(subfig5, plot55, label = "5", color = 6, colormap = :tab10, colorrange = (1, 10) )
    lines!(subfig5, plot56, label = "6", color = 7, colormap = :tab10, colorrange = (1, 10) )
    lines!(subfig5, plot57, label = "7", color = 8, colormap = :tab10, colorrange = (1, 10) )
    lines!(subfig5, plot58, label = "8", color = 9, colormap = :tab10, colorrange = (1, 10) )
    lines!(subfig5, plot59, label = "9", color = 10, colormap = :tab10, colorrange = (1, 10))
    axislegend(subfig5, position = :lb)

    subfig6 = GLMakie.Axis(fig1[6, 1],  # define position of this subfigure inside a figure
        title = "output neuron wRecChange",
        xlabel = "time",
        ylabel = "data"
    )
    lines!(subfig6, plot60, label = "0", color = 1, colormap = :tab10, colorrange = (1, 10) )
    lines!(subfig6, plot61, label = "1", color = 2, colormap = :tab10, colorrange = (1, 10) )
    lines!(subfig6, plot62, label = "2", color = 3, colormap = :tab10, colorrange = (1, 10) )
    lines!(subfig6, plot63, label = "3", color = 4, colormap = :tab10, colorrange = (1, 10) )
    lines!(subfig6, plot64, label = "4", color = 5, colormap = :tab10, colorrange = (1, 10) )
    lines!(subfig6, plot65, label = "5", color = 6, colormap = :tab10, colorrange = (1, 10) )
    lines!(subfig6, plot66, label = "6", color = 7, colormap = :tab10, colorrange = (1, 10) )
    lines!(subfig6, plot67, label = "7", color = 8, colormap = :tab10, colorrange = (1, 10) )
    lines!(subfig6, plot68, label = "8", color = 9, colormap = :tab10, colorrange = (1, 10) )
    lines!(subfig6, plot69, label = "9", color = 10, colormap = :tab10, colorrange = (1, 10))
    axislegend(subfig6, position = :lb)

    # wait(display(fig1))
    # display(fig1)
    # --------------------------------- end plot --------------------------------- #

    # model learning   
    maxRepeatRound = 1     # repeat each image
    thinkingPeriod = 16    # 1000-784 = 216   
    bestAccuracy = 0.0                       
    for epoch = 1:1000
        batchCounter = 0   
        for (imgBatch, labelBatch) in trainData
            batchCounter += 1
            println("epoch $epoch batch $batchCounter")
            @showprogress for i in eachindex(labelBatch)
                _img = (imgBatch[:, :, i])
                img = reshape(_img, (:, 1))
                row, col = size(img)
                label = labelBatch[i]

                img_tensor = torch.from_numpy( np.asarray(img) )

                # create more data for RSNN
                spike = spikegen.delta(img_tensor, threshold=0.1, off_spike=true)
                spike1 = isequal.(pyconvert(Array, spike.data.numpy()), 1)
                spike2 = isequal.(pyconvert(Array, spike.data.numpy()), 1)
                spike3 = isequal.(pyconvert(Array, spike.data.numpy()), 1)
                spike4 = isequal.(pyconvert(Array, spike.data.numpy()), 1)
                spike5 = isequal.(pyconvert(Array, spike.data.numpy()), 1)
                spike6 = isequal.(pyconvert(Array, spike.data.numpy()), -1)
                spike7 = isequal.(pyconvert(Array, spike.data.numpy()), -1)
                spike8 = isequal.(pyconvert(Array, spike.data.numpy()), -1)
                spike9 = isequal.(pyconvert(Array, spike.data.numpy()), -1)
                spike10 = isequal.(pyconvert(Array, spike.data.numpy()), -1)

                spike = spikegen.delta(img_tensor, threshold=0.2, off_spike=true)
                spike11 = isequal.(pyconvert(Array, spike.data.numpy()), 1)
                spike12 = isequal.(pyconvert(Array, spike.data.numpy()), 1)
                spike13 = isequal.(pyconvert(Array, spike.data.numpy()), 1)
                spike14 = isequal.(pyconvert(Array, spike.data.numpy()), 1)
                spike15 = isequal.(pyconvert(Array, spike.data.numpy()), 1)
                spike16 = isequal.(pyconvert(Array, spike.data.numpy()), -1)
                spike17 = isequal.(pyconvert(Array, spike.data.numpy()), -1)
                spike18 = isequal.(pyconvert(Array, spike.data.numpy()), -1)
                spike19 = isequal.(pyconvert(Array, spike.data.numpy()), -1)
                spike20 = isequal.(pyconvert(Array, spike.data.numpy()), -1)

                spike = spikegen.delta(img_tensor, threshold=0.3, off_spike=true)
                spike21 = isequal.(pyconvert(Array, spike.data.numpy()), 1)
                spike22 = isequal.(pyconvert(Array, spike.data.numpy()), 1)
                spike23 = isequal.(pyconvert(Array, spike.data.numpy()), 1)
                spike24 = isequal.(pyconvert(Array, spike.data.numpy()), 1)
                spike25 = isequal.(pyconvert(Array, spike.data.numpy()), 1)
                spike26 = isequal.(pyconvert(Array, spike.data.numpy()), -1)
                spike27 = isequal.(pyconvert(Array, spike.data.numpy()), -1)
                spike28 = isequal.(pyconvert(Array, spike.data.numpy()), -1)
                spike29 = isequal.(pyconvert(Array, spike.data.numpy()), -1)
                spike30 = isequal.(pyconvert(Array, spike.data.numpy()), -1)

                spike = spikegen.delta(img_tensor, threshold=0.4, off_spike=true)
                spike31 = isequal.(pyconvert(Array, spike.data.numpy()), 1)
                spike32 = isequal.(pyconvert(Array, spike.data.numpy()), 1)
                spike33 = isequal.(pyconvert(Array, spike.data.numpy()), 1)
                spike34 = isequal.(pyconvert(Array, spike.data.numpy()), 1)
                spike35 = isequal.(pyconvert(Array, spike.data.numpy()), 1)
                spike36 = isequal.(pyconvert(Array, spike.data.numpy()), -1)
                spike37 = isequal.(pyconvert(Array, spike.data.numpy()), -1)
                spike38 = isequal.(pyconvert(Array, spike.data.numpy()), -1)
                spike39 = isequal.(pyconvert(Array, spike.data.numpy()), -1)
                spike40 = isequal.(pyconvert(Array, spike.data.numpy()), -1)

                spike = spikegen.delta(img_tensor, threshold=0.5, off_spike=true)
                spike41 = isequal.(pyconvert(Array, spike.data.numpy()), 1)
                spike42 = isequal.(pyconvert(Array, spike.data.numpy()), 1)
                spike43 = isequal.(pyconvert(Array, spike.data.numpy()), 1)
                spike44 = isequal.(pyconvert(Array, spike.data.numpy()), 1)
                spike45 = isequal.(pyconvert(Array, spike.data.numpy()), 1)
                spike46 = isequal.(pyconvert(Array, spike.data.numpy()), -1)
                spike47 = isequal.(pyconvert(Array, spike.data.numpy()), -1)
                spike48 = isequal.(pyconvert(Array, spike.data.numpy()), -1)
                spike49 = isequal.(pyconvert(Array, spike.data.numpy()), -1)
                spike50 = isequal.(pyconvert(Array, spike.data.numpy()), -1)

                input = [spike1;; spike2;; spike3;; spike4;; spike5;; spike6;; spike7;; spike8;; spike9;; spike10;;
                spike11;; spike12;; spike13;; spike14;; spike15;; spike16;; spike17;; spike18;; spike19;; spike20;;
                spike21;; spike22;; spike23;; spike24;; spike25;; spike26;; spike27;; spike28;; spike29;; spike30;;
                spike31;; spike32;; spike33;; spike34;; spike35;; spike36;; spike37;; spike38;; spike39;; spike40;;
                spike41;; spike42;; spike43;; spike44;; spike45;; spike46;; spike47;; spike48;; spike49;; spike50
                ]'   # ' to flip 784x10 to 10x784

                predict = 0
                
                for k in 1:maxRepeatRound

                    # insert data into model sequencially
                    for i in 1:(row + thinkingPeriod)     # sMNIST ihas 784 timestep(pixel) + thinking period = 1000 timestep
                        tick = i
                        if i <= row
                            current_pixel = input[:, i]
                        else
                            current_pixel = zeros(size(input)[1])   # dummy input in "thinking" period
                        end

                        if tick == 1    # tell a model to start learning. 1-time only
                            model.learningStage = "start_learning"
                            
                        elseif tick == (row+thinkingPeriod)    
                            model.learningStage = "end_learning"
                        else
                        end

                        _firedNeurons_t1, logit, _var1, _var2, _var3, _var4 = model(current_pixel)
                        # log answer of all timestep
                        logitLog = [logitLog;; logit]
                        firedNeurons_t1 = push!(firedNeurons_t1, _firedNeurons_t1)
                        var1 = [var1;; _var1]
                        var2 = [var2;; _var2]
                        var3 = [var3;; _var3]
                        var4 = [var4;; _var4]

                        if tick < row   # online learning, 1-by-1 timestep
                            # no error calculation
                        elseif tick == row   # online learning, 1-by-1 timestep
                            correctAnswer = OneHotArrays.onehot(label, labelDict)
                            modelError = Flux.logitcrossentropy(logit, correctAnswer) * 1.0 
                            outputError = (logit - correctAnswer) * 1.0
                            Ironpen.compute_paramsChange!(model, modelError, outputError)
                        elseif tick > row && tick < row+thinkingPeriod
                            correctAnswer = OneHotArrays.onehot(label, labelDict)
                            modelError = Flux.logitcrossentropy(logit, correctAnswer) * 1.0 
                            outputError = (logit - correctAnswer) * 1.0
                            Ironpen.compute_paramsChange!(model, modelError, outputError)
                        elseif tick == row+thinkingPeriod
                            correctAnswer = OneHotArrays.onehot(label, labelDict)
                            modelError = Flux.logitcrossentropy(logit, correctAnswer) * 1.0
                            outputError = (logit - correctAnswer) * 1.0
                            Ironpen.compute_paramsChange!(model, modelError, outputError)
                            Ironpen.learn!(model)
                            _logit = round.(logit; digits=2)
                            predict = findall(isequal.(logit, maximum(logit)))[1] - 1
                            y = round.(modelError; digits=2)
                            println("")
                            println("label $label predict $predict logit $_logit model error $y")
                        else
                            error("undefined condition line $(@__LINE__)")
                        end

                        # update plot
                        plot10[] = firedNeurons_t1

                        plot20[] = view(logitLog, 1 , :)
                        plot21[] = view(logitLog, 2 , :)
                        plot22[] = view(logitLog, 3 , :)
                        plot23[] = view(logitLog, 4 , :)
                        plot24[] = view(logitLog, 5 , :)
                        plot25[] = view(logitLog, 6 , :)
                        plot26[] = view(logitLog, 7 , :)
                        plot27[] = view(logitLog, 8 , :)
                        plot28[] = view(logitLog, 9 , :)
                        plot29[] = view(logitLog, 10, :)

                        plot30[] = view(var1, 1 , :)
                        plot31[] = view(var1, 2 , :)
                        plot32[] = view(var1, 3 , :)
                        plot33[] = view(var1, 4 , :)
                        plot34[] = view(var1, 5 , :)
                        plot35[] = view(var1, 6 , :)
                        plot36[] = view(var1, 7 , :)
                        plot37[] = view(var1, 8 , :)
                        plot38[] = view(var1, 9 , :)
                        plot39[] = view(var1, 10, :)

                        plot40[] = view(var2, 1 , :)
                        plot41[] = view(var2, 2 , :)
                        plot42[] = view(var2, 3 , :)
                        plot43[] = view(var2, 4 , :)
                        plot44[] = view(var2, 5 , :)
                        plot45[] = view(var2, 6 , :)
                        plot46[] = view(var2, 7 , :)
                        plot47[] = view(var2, 8 , :)
                        plot48[] = view(var2, 9 , :)
                        plot49[] = view(var2, 10, :)

                        plot50[] = view(var3, 1 , :)
                        plot51[] = view(var3, 2 , :)
                        plot52[] = view(var3, 3 , :)
                        plot53[] = view(var3, 4 , :)
                        plot54[] = view(var3, 5 , :)
                        plot55[] = view(var3, 6 , :)
                        plot56[] = view(var3, 7 , :)
                        plot57[] = view(var3, 8 , :)
                        plot58[] = view(var3, 9 , :)
                        plot59[] = view(var3, 10, :)

                        plot60[] = view(var4, 1 , :)
                        plot61[] = view(var4, 2 , :)
                        plot62[] = view(var4, 3 , :)
                        plot63[] = view(var4, 4 , :)
                        plot64[] = view(var4, 5 , :)
                        plot65[] = view(var4, 6 , :)
                        plot66[] = view(var4, 7 , :)
                        plot67[] = view(var4, 8 , :)
                        plot68[] = view(var4, 9 , :)
                        plot69[] = view(var4, 10, :)
                    end
                    # end-thinkingPeriod+2; +2 because initialize logitLog = zeros(10, 2)
                    # _modelRespond = logitLog[:, end-thinkingPeriod+2:end]   # answer count during thinking period
                    # _modelRespond = [sum(i) for i in eachrow(_modelRespond)]
                    # modelRespond = isequal.(isequal.(_modelRespond, 0), 0)
                    
                    display(fig1)
                    # sleep(1)
                    if k % 3 == 0
                        firedNeurons_t1 = zeros(1)
                        logitLog = zeros(10, 2)
                        var1 = zeros(10, 2)
                        var2 = zeros(10, 2)
                        var3 = zeros(10, 2)
                        var4 = zeros(10, 2)
                    end
                    
                    # if predict == OneHotArrays.onehot(label, labelDict)
                    #     println("model train $label successfully, $k tries")
                    #     # wait(display(fig1))
                        
                    #     firedNeurons_t1 = zeros(1)
                    #     logitLog = zeros(10, 2)
                    #     var1 = zeros(10, 2)
                    #     var2 = zeros(10, 2)
                    #     var3 = zeros(10, 2)
                    #     var4 = zeros(10, 2)
                    #     break
                    # end

                    if k == maxRepeatRound
                        # println("model train $label unsuccessfully, $maxRepeatRound tries, skip training")
                        # display(fig1)
                        firedNeurons_t1 = zeros(1)
                        logitLog = zeros(10, 2)
                        var1 = zeros(10, 2)
                        var2 = zeros(10, 2)
                        var3 = zeros(10, 2)
                        var4 = zeros(10, 2)
                        break
                    end
                    
                    GC.gc()
                end
            end
            # check accuracy
            println("validating model")
            answerCorrectly = validate(model, validateData, labelDict)
            bestAccuracy = answerCorrectly > bestAccuracy ? answerCorrectly : bestAccuracy
            println("model accuracy is $answerCorrectly %, best accuracy is $bestAccuracy")
        end

        # # check mean error and accuracy
        # mean_error = round(mean(model_error_list), sigdigits = 3)
        # accuracy = round(model_accuracy / batch_size * 100, sigdigits = 3)
        # println("------------")
        # println(model_name)
        # println("mean error $mean_error accuracy $accuracy")
    end
end

function validate(model, dataset, labelDict)
    answerCorrectly = 0.0     # %
    thinkingPeriod = 16    # 1000-784 = 216    
    @showprogress for (image, label) in dataset
        img = reshape(image, (:, 1))
        row, col = size(img)
        label = label[1]

        img_tensor = torch.from_numpy( np.asarray(img) )

        # create more data for RSNN
        spike = spikegen.delta(img_tensor, threshold=0.1, off_spike=true)
        spike1 = isequal.(pyconvert(Array, spike.data.numpy()), 1)
        spike2 = isequal.(pyconvert(Array, spike.data.numpy()), 1)
        spike3 = isequal.(pyconvert(Array, spike.data.numpy()), 1)
        spike4 = isequal.(pyconvert(Array, spike.data.numpy()), 1)
        spike5 = isequal.(pyconvert(Array, spike.data.numpy()), 1)
        spike6 = isequal.(pyconvert(Array, spike.data.numpy()), -1)
        spike7 = isequal.(pyconvert(Array, spike.data.numpy()), -1)
        spike8 = isequal.(pyconvert(Array, spike.data.numpy()), -1)
        spike9 = isequal.(pyconvert(Array, spike.data.numpy()), -1)
        spike10 = isequal.(pyconvert(Array, spike.data.numpy()), -1)

        spike = spikegen.delta(img_tensor, threshold=0.2, off_spike=true)
        spike11 = isequal.(pyconvert(Array, spike.data.numpy()), 1)
        spike12 = isequal.(pyconvert(Array, spike.data.numpy()), 1)
        spike13 = isequal.(pyconvert(Array, spike.data.numpy()), 1)
        spike14 = isequal.(pyconvert(Array, spike.data.numpy()), 1)
        spike15 = isequal.(pyconvert(Array, spike.data.numpy()), 1)
        spike16 = isequal.(pyconvert(Array, spike.data.numpy()), -1)
        spike17 = isequal.(pyconvert(Array, spike.data.numpy()), -1)
        spike18 = isequal.(pyconvert(Array, spike.data.numpy()), -1)
        spike19 = isequal.(pyconvert(Array, spike.data.numpy()), -1)
        spike20 = isequal.(pyconvert(Array, spike.data.numpy()), -1)

        spike = spikegen.delta(img_tensor, threshold=0.3, off_spike=true)
        spike21 = isequal.(pyconvert(Array, spike.data.numpy()), 1)
        spike22 = isequal.(pyconvert(Array, spike.data.numpy()), 1)
        spike23 = isequal.(pyconvert(Array, spike.data.numpy()), 1)
        spike24 = isequal.(pyconvert(Array, spike.data.numpy()), 1)
        spike25 = isequal.(pyconvert(Array, spike.data.numpy()), 1)
        spike26 = isequal.(pyconvert(Array, spike.data.numpy()), -1)
        spike27 = isequal.(pyconvert(Array, spike.data.numpy()), -1)
        spike28 = isequal.(pyconvert(Array, spike.data.numpy()), -1)
        spike29 = isequal.(pyconvert(Array, spike.data.numpy()), -1)
        spike30 = isequal.(pyconvert(Array, spike.data.numpy()), -1)

        spike = spikegen.delta(img_tensor, threshold=0.4, off_spike=true)
        spike31 = isequal.(pyconvert(Array, spike.data.numpy()), 1)
        spike32 = isequal.(pyconvert(Array, spike.data.numpy()), 1)
        spike33 = isequal.(pyconvert(Array, spike.data.numpy()), 1)
        spike34 = isequal.(pyconvert(Array, spike.data.numpy()), 1)
        spike35 = isequal.(pyconvert(Array, spike.data.numpy()), 1)
        spike36 = isequal.(pyconvert(Array, spike.data.numpy()), -1)
        spike37 = isequal.(pyconvert(Array, spike.data.numpy()), -1)
        spike38 = isequal.(pyconvert(Array, spike.data.numpy()), -1)
        spike39 = isequal.(pyconvert(Array, spike.data.numpy()), -1)
        spike40 = isequal.(pyconvert(Array, spike.data.numpy()), -1)

        spike = spikegen.delta(img_tensor, threshold=0.5, off_spike=true)
        spike41 = isequal.(pyconvert(Array, spike.data.numpy()), 1)
        spike42 = isequal.(pyconvert(Array, spike.data.numpy()), 1)
        spike43 = isequal.(pyconvert(Array, spike.data.numpy()), 1)
        spike44 = isequal.(pyconvert(Array, spike.data.numpy()), 1)
        spike45 = isequal.(pyconvert(Array, spike.data.numpy()), 1)
        spike46 = isequal.(pyconvert(Array, spike.data.numpy()), -1)
        spike47 = isequal.(pyconvert(Array, spike.data.numpy()), -1)
        spike48 = isequal.(pyconvert(Array, spike.data.numpy()), -1)
        spike49 = isequal.(pyconvert(Array, spike.data.numpy()), -1)
        spike50 = isequal.(pyconvert(Array, spike.data.numpy()), -1)

        input = [spike1;; spike2;; spike3;; spike4;; spike5;; spike6;; spike7;; spike8;; spike9;; spike10;;
        spike11;; spike12;; spike13;; spike14;; spike15;; spike16;; spike17;; spike18;; spike19;; spike20;;
        spike21;; spike22;; spike23;; spike24;; spike25;; spike26;; spike27;; spike28;; spike29;; spike30;;
        spike31;; spike32;; spike33;; spike34;; spike35;; spike36;; spike37;; spike38;; spike39;; spike40;;
        spike41;; spike42;; spike43;; spike44;; spike45;; spike46;; spike47;; spike48;; spike49;; spike50
        ]'   # ' to flip 784x10 to 10x784

        # insert data into model sequencially
        logit = Float64[]
        for i in 1:(row + thinkingPeriod)     # sMNIST ihas 784 timestep(pixel) + thinking period = 1000 timestep
            if i <= row
                current_pixel = input[:, i]
            else
                current_pixel = zeros(size(input)[1])   # dummy input in "thinking" period
            end

            _firedNeurons_t1, logit, _var1, _var2, _var3, _var4 = model(current_pixel)
        end

        predict = findall(isequal.(logit, maximum(logit)))[1] - 1
        if predict == label
            answerCorrectly += 1
            # println("model answer $label correctly")
        else
            # println("img $label, model answer $predict")
        end
        GC.gc()
    end

    correctPercent = answerCorrectly * 100.0 / length(dataset)
            
    return  correctPercent::Float64
end


function main()
    training_start_time = Dates.now()
    println("program started ", training_start_time)

    filelocation = string(@__DIR__)

    # generate SNN
    for i = 1:1
        modelname = "v06_36"
        filename = "$modelname.jl163"
        generate_snn(filename, filelocation)
    end

    modelname = "v06_36"
    filename = "$modelname.jl163"
    # filename = "v06_31c.jl163"

    trainDataset, validateDataset, labelDict = data_loader()

    train_snn(modelname, filename, filelocation, trainDataset, validateDataset, labelDict)

    finish_training_time = Dates.now()
    println("training done,  $training_start_time ==> $finish_training_time ")
    println(" ///////////////////////////////////////////////////////////////////////")
end

# only runs main() if julia isn’t started interactively 
# https://discourse.julialang.org/t/scripting-like-a-julian/50707
!isinteractive() && main() 
#------------------------------------------------------------------------------------------------100