/* See {nmsim_neuron_net.h} */
/* Last edited on 2018-05-25 08:26:56 by jstolfi */

#define _GNU_SOURCE
#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <assert.h>
#include <math.h>

#include <affirm.h>
#include <bool.h>
#include <jsrandom.h>

#include <nmsim_neuron_parms.h>
#include <nmsim_neuron_state.h>

#include <nmsim_neuron_net.h>

nmsim_neuron_net_t *nmsim_neuron_net_new(nmsim_neuron_count_t nn)
  { 
    demand(nn <= nmsim_neuron_net_MAX_NEURONS, "too many neurons");
    
    nmsim_neuron_net_t *net = notnull(malloc(sizeof(nmsim_neuron_net_t)), "no mem");
    net->nn = nn;
    net->ns = 0;
    /* Allocate the neuron property and state tables: */
    net->state = notnull(malloc(nn*sizeof(nmsim_neuron_state_t)), "no mem");
    net->parms = notnull(malloc(nn*sizeof(nmsim_neuron_parms_t *)), "no mem");
    net->deg_in = notnull(malloc(nn*sizeof(nmsim_neuron_count_t)), "no mem");
    net->deg_ot = notnull(malloc(nn*sizeof(nmsim_neuron_count_t)), "no mem");
    net->syn = nmsim_synapse_vec_new(0); /* To be expanded. */
    /* Initialize the data of each neuron with null/undefined values: */
    for (nmsim_neuron_ix_t i = 0; i < nn; i ++)
      { net->parms[i] = NULL;
        net->deg_in[i] = 0;
        net->deg_ot[i] = 0;
        net->state[i] = (nmsim_neuron_state_t){ .V = NAN, .age = (~(uint64_t)0) };
      }
    return net;
  }

void nmsim_neuron_net_add_random_synapses
  ( nmsim_neuron_net_t *net,
    nmsim_neuron_ix_t ia0,
    nmsim_neuron_ix_t ia1,
    nmsim_neuron_ix_t ib0,
    nmsim_neuron_ix_t ib1,
    double pr,
    double swt_tot_avg,
    double swt_tot_dev
  )
  {
    if (pr == 0.0)
      { /* Nothing to do: */
        demand(swt_tot_avg == 0.0, "with no synapses the total weight avg must be zero");
        demand(swt_tot_dev >= 0.0, "with no synapses the total weight dev must be zero");
        return;
      }
    
    /* We must add at least one synapse for each neuron in {ib0..ib1}. */
    demand((pr > 0.0) && (pr <= 1.0), "invalid synapse probability");
    demand(fabs(swt_tot_avg) > 1.0e-6, "avg of total synaptic weight is too small");
    demand(swt_tot_dev >= 0.0, "dev of total synaptic weight is negative");

    /* Check pre- and post-synaptic neuron index ranges and count them: */
    demand((ia0 <= ia1) && (ia1 < nmsim_neuron_net_MAX_NEURONS), "invalid presynaptic index range");
    demand((ib0 <= ib1) && (ib1 < nmsim_neuron_net_MAX_NEURONS), "invalid postynaptic index range");

    /* Get and check the number {ns} of previously defined synapses: */
    nmsim_synapse_count_t ns = net->ns;
    nmsim_synapse_vec_t *syn = &(net->syn);
    demand(ns <= syn->ne, "invalid synapse count");

    auto void add_synapse(nmsim_neuron_ix_t ia, nmsim_neuron_ix_t ib);
      /* Appends a synapse from neuron {ia} to neuron {ib} to
        {syn[0..ns-1]}, expanding the vectors as needed and
        incrementing {ns}. The weight is set to {NAN} for the time
        being. */
    
    /* Add the new synapses: */
    for (nmsim_neuron_ix_t ib = ib0; ib <= ib1; ib++)
      { /* Create synapses from {ia0..ia1} to {ib}: */
        nmsim_synapse_ix_t Mb0 = ns; /* Index of first new synapse into {ib}. */
        /* !!! This loop is executed {nn^2} times. Must improve for low {pr}. !!! */
        for (nmsim_neuron_ix_t ia = ia0; ia <= ia1; ia++)
          { if (drandom() < pr) { add_synapse(ia, ib); } }
        if (Mb0 == ns)
          { /* No synapses were added.  Force one synapse into {ib}: */
            nmsim_neuron_ix_t ia = uint64_abrandom(ia0,ia1);
            add_synapse(ia, ib);
          }
        /* Compute the desired average and deviation of weights: */
        nmsim_neuron_count_t din = ns - Mb0; /* New synapses into to {ib} */
        assert(din >= 1);
        double swt_avg = swt_tot_avg/((double)din); /* Avg of each synapse weight. */
        double swt_dev = swt_tot_dev/sqrt((double)din); /* Dev of each synapse weight. */
        if (swt_dev == 0)
          { /* All synapse weights are equal to {swt_avg}: */
            for (nmsim_synapse_ix_t kb = Mb0; kb < ns; kb++) 
              { syn->e[kb].swt = swt_avg; }
          }
        else
          { /* Synapse weights are independent log-normal deviates. */
            /* Convert {swt_avg,swt_dev} to params of log normal distrib: */
            double r = swt_dev/swt_avg;
            double s = 1 + r*r;
            double zwt_avg = log(fabs(swt_avg)/sqrt(s));
            double zwt_dev = sqrt(log(s));
            for (nmsim_synapse_ix_t kb = Mb0; kb < ns; kb++)
              { /* Generate a random log-weight: */
                double zwt = zwt_avg + zwt_dev*dgaussrand();
                /* Convert to actual weight: */
                double swt_abs = exp(zwt);
                syn->e[kb].swt =  (swt_avg < 0 ? -swt_abs : + swt_abs);
              }
          }
      }
            
    net->ns = ns;
    return;
    
    /* Internal procedures: */
    
    void add_synapse(nmsim_neuron_ix_t ia, nmsim_neuron_ix_t ib)
      {
        demand(ns < nmsim_neuron_net_MAX_SYNAPSES, "too many synapses");
        nmsim_synapse_vec_expand(syn, (int)ns);
        syn->e[ns] = (nmsim_synapse_t){ .pre = ia, .pos = ib, .swt = NAN };
        net->deg_ot[ia]++;
        net->deg_in[ib]++;
        ns++;
      }
  }

void nmsim_neuron_net_write(FILE *wr, nmsim_neuron_net_t *net)
  {
    fprintf(wr, "neurons = %ld\n", net->nn);
    for (nmsim_neuron_ix_t i = 0; i < net->nn; i++)
      { nmsim_neuron_count_t din = net->deg_in[i];
        nmsim_neuron_count_t dot = net->deg_ot[i];
        fprintf(wr, "  %ld %ld %ld\n", i, din, dot);
        nmsim_neuron_parms_write(wr, "    parms = {", net->parms[i], "}\n");
        nmsim_neuron_state_write(wr, "    state = {", &(net->state[i]), "}\n");
      }
    fprintf(wr, "synapses = %ld\n", net->ns);
    for (nmsim_synapse_ix_t k = 0; k < net->nn; k++)
      { nmsim_synapse_t *syn = &(net->syn.e[k]);
        nmsim_neuron_ix_t ipre = syn->pre;
        nmsim_neuron_ix_t ipos = syn->pos;
        double swt = syn->swt;
        fprintf(wr, "  %ld %ld %ld %10.8f\n", k, ipre, ipos, swt);
      }
    fflush(wr);
  }

void nmsim_neuron_net_tot_inputs(nmsim_neuron_net_t *net, bool_t X[], double J[], double dV[])
  { 
    /* Initialize the potential increments {dV[0..nn-1]} with external inputs: */
    /* !!! Should attenuate by input modulator? !!! */
    for (nmsim_neuron_ix_t i = 0; i < net->nn; i++) 
      { dV[i] = (J == NULL ? 0.0 : J[i]); }
    
    /* Scan the synapses and accumulate firing contributions: */
    /* !!! Should precompute the modulators of the neurons. */
    for (nmsim_synapse_ix_t k = 0; k < net->ns; k++)
      { nmsim_synapse_t *syn = &(net->syn.e[k]);
        nmsim_neuron_ix_t ia = syn->pre;
        nmsim_neuron_ix_t ib = syn->pos;
        
        if (X[ia])
          { /* Grab properties of neurons {ia,ib}: */
            nmsim_neuron_parms_t *parms_a = net->parms[ia];
            nmsim_neuron_state_t *state_a = &(net->state[ia]);
            nmsim_neuron_parms_t *parms_b = net->parms[ib];
            nmsim_neuron_state_t *state_b = &(net->state[ib]);

            /* Output modulator of neuron {ia}: */
            double H_a = nmsim_neuron_state_output_modulator(parms_a, state_a->age);

            /* Input modulator of neuron {ib}: */
            double G_b = nmsim_neuron_state_input_modulator(parms_b, state_b->age);

            /* Accumulate the increment: */
            dV[ib] += H_a * syn->swt * G_b;
          }
      }
  }

void nmsim_neuron_net_free(nmsim_neuron_net_t *net)
  {
    if(net->state  != NULL) { free(net->state ); net->state  = NULL; }
    if(net->parms  != NULL) { free(net->parms ); net->parms  = NULL; }
    if(net->deg_in != NULL) { free(net->deg_in); net->deg_in = NULL; }
    if(net->deg_ot != NULL) { free(net->deg_ot); net->deg_ot = NULL; }

    if(net->syn.e != NULL) 
      { free(net->syn.e); net->syn.e = NULL; net->syn.ne = 0; }

    free(net);
  }