// ChucK Brain
// Copyright 2008 Les Hall
// This software is protected by the GNU General Public License


// parameters
3 => int num_width;  // width of neural net
5 => int num_layers;  // number of layers (input + hidden + output)
0.99999 => float tau;  // time constant for convergence
0.001 => float T_limit;  // lower limit on T
440.0 => float freq;  // the frequency of the oscillator
( 2 * (1 / freq) / num_layers)::ms => dur nn_delay;  // delay of each neuron
//10::samp => dur nn_delay;  // delay of each neuron
0.1 => float k;  // change constant for training
100 => int num_population;  // number of members in population


// variables
1 => float T;  // temperature
((1 + Math.sqrt(5)) / 2) => float Phi;
float delta;  // used in updating weights
int t;  // time counter
0 => int population_index;
float foms[num_population];


// define the patch
// define the input
SinOsc kick => Gain sum;
SinOsc in => sum;
//in => dac;
adc => sum;
kick.freq (freq);
kick.gain (0.1);
sum.gain (1);
// define and hook up the neurons
Gain nn_weights[num_layers][num_width][num_width];  // define the neural net weights
float nn_population_weights[num_population][num_layers][num_width][num_width];  // the population weights
float nn_population_freqs[num_population][num_layers][num_width][num_width];  // the population frequencies
dur nn_population_delays[num_population][num_layers][num_width];  // the population delays
Dyno nn_outputs[num_layers][num_width];  // define the neural net outputs
DelayL nn_delays[num_layers][num_width];  // define the neural net delays
for (1 => int layer; layer < num_layers; layer++) {
    for (0 => int neuron; neuron < num_width; neuron++ ) {
        nn_delays[layer][neuron].max ((1 / freq)::second);  // set maximum delay
    }
}
for (0 => int neuron; neuron < num_width; neuron++ ) {
    sum => nn_outputs[0][neuron];  // connect inputs to audio source
}
for (1 => int layer; layer < num_layers; layer++) {
    for (0 => int neuron; neuron < num_width; neuron++ ) {
        for (0 => int weight; weight < num_width; weight++) {
            //nn_weights[layer][neuron][weight].Q (10);
            nn_outputs[layer-1][neuron] => nn_weights[layer][neuron][weight];  // connect weight input
            nn_weights[layer][neuron][weight] => nn_delays[layer][neuron];  // connect to delay
        }
        nn_delays[layer][neuron] => nn_outputs[layer][neuron];  // connect weight output
    }
}
for (0 => int layer; layer < num_layers; layer++) {
    for (0 => int neuron; neuron < num_width; neuron++ ) {
        nn_outputs[layer][neuron].limit ();  // the activation function
    }
}
// define the output node and hook it up
Dyno out;
out.limit ();  // set output as limiter
for (0 => int neuron; neuron < num_width; neuron++ ) {
    nn_outputs[num_layers-1][neuron] => out;  // connect all neural net outputs to out
}
//
// define the figure of merit
// the fom patch and its parameters
kick => Gain in_gain => Gain fom => dac;
sum => Gain out_gain => fom;
in_gain.gain (1);
out_gain.gain (-1);
fom.gain (1);
//
// try this:  feed the output back to the input
out => FullRect fullrect => LPF feedback => in;
feedback.gain (2 * freq);
feedback.freq (freq / 33);


// seed the neural net weights
for (1 => int layer; layer < num_layers; layer++) {
    for (0 => int neuron; neuron < num_width; neuron++) {
        for (0 => int weight; weight < num_width; weight++) {
            Std.randf () => nn_weights[layer][neuron][weight].gain;
        }
    }
}


// time loop
while (true) {
    // select a population member
    Std.rand2 (0, num_population - 1) => population_index;
    
    // update the weights
    for (1 => int layer; layer < num_layers; layer++) {
        for (0 => int neuron; neuron < num_width; neuron++) {
            for (0 => int weight; weight < num_width; weight++) {
                T * k * Std.randf () + nn_population_weights[population_index][layer][neuron][weight] => nn_weights[layer][neuron][weight].gain;
                //T * k * Std.rand2f (0, 2 * freq) + nn_population_freqs[population_index][layer][neuron][weight] => nn_weights[layer][neuron][weight].freq;
            }
            T * k * Std.randf () * nn_delay + nn_population_delays[population_index][layer][neuron] => dur duration;
            if (duration < 0::second) {
                0::second => nn_delays[layer][neuron].delay;
            } else {
                duration => nn_delays[layer][neuron].delay;
            }
        }
    }

    // update the temperature
    tau *=> T;
    if (T < T_limit) {
        1 => T;
    }
    if ((t % 500 == 0)) {
        <<<t, T * k, fom.last (), out.last (), feedback.last ()>>>;
    }
    t++;
    
    // time delay
    Phi :: ms => now;
    
    // save weights if fom is better
    if (Std.rand2f(0, fom.last () + T * foms[population_index]) < fom.last ()) {
        for (1 => int layer; layer < num_layers; layer++) {
            for (0 => int neuron; neuron < num_width; neuron++) {
                for (0 => int weight; weight < num_width; weight++) {
                    // save this set of weights since it's better
                    nn_weights[layer][neuron][weight].gain () => nn_population_weights[population_index][layer][neuron][weight];
                    //nn_weights[layer][neuron][weight].freq () => nn_population_freqs[population_index][layer][neuron][weight];
                }
                nn_delays[layer][neuron].delay () => nn_population_delays[population_index][layer][neuron];

            }
        }
        // save the fom
        fom.last () => foms[population_index];
    }

}