#include "core.h"
#include <cmath>
#include <cstdlib>

#define USE_PARALLEL            // Закомментируйте, чтобы отключить параллелизм
#define MAX_THREADS 0           // Укажите число ядер (0 — использовать все доступные)

#ifdef USE_PARALLEL
#include <omp.h>
#define OMP_PARALLEL _Pragma("omp parallel for")
#define OMP_SET_THREADS() { if (MAX_THREADS > 0) omp_set_num_threads(MAX_THREADS); }
#else
#define OMP_PARALLEL 
#define OMP_SET_THREADS() 
#endif


NeuralNetwork::NeuralNetwork(LayerStructure_t layers[], int count) : numLayers(count) {
    for (int i = 0; i < count; i++) sizes.push_back(layers[i].size);
    for (int i = 0; i < count - 1; i++) {
        std::vector<std::vector<double>> layerW;
        double scale = sqrt(2.0 / sizes[i]);
        for (int j = 0; j < sizes[i+1]; j++) {
            std::vector<double> nodeW;
            for (int k = 0; k < sizes[i]; k++) 
                nodeW.push_back(((double)rand()/RAND_MAX * 2 - 1) * scale);
            layerW.push_back(nodeW);
        }
        weights.push_back(layerW);
        biases.push_back(std::vector<double>(sizes[i+1], 0.0));
    }
}

std::vector<double> NeuralNetwork::feedForward(const std::vector<double>& input) {
    OMP_SET_THREADS(); // Применяем лимит ядер

    outputs.clear();
    outputs.push_back(input);

    std::vector<double> curr = input;

    for (int i = 0; i < numLayers - 1; i++) {
        // Заранее готовим вектор нужного размера для текущего слоя
        std::vector<double> next(sizes[i + 1]);

        // Распараллеливаем расчет каждого нейрона в слое
        OMP_PARALLEL
            for (int j = 0; j < sizes[i + 1]; j++) {
                double sum = biases[i][j];

                // Внутренний цикл обычно оставляем последовательным, 
                // так как создание потоков здесь даст больше тормозов, чем пользы
                for (int k = 0; k < (int)curr.size(); k++) {
                    sum += curr[k] * weights[i][j][k];
                }

                // Сигмоида. Теперь пишем по индексу j — это безопасно для потоков
                next[j] = 1.0 / (1.0 + exp(-sum));
            }

        curr = next;
        outputs.push_back(curr);
    }
    return curr;
}


double NeuralNetwork::train(const std::vector<double>& input, const std::vector<double>& target, double lr) {
    // 0. Устанавливаем количество потоков (если включено в define)
    OMP_SET_THREADS();

    // 1. Прямой проход (получаем предсказание)
    std::vector<double> pred = feedForward(input);

    std::vector<std::vector<double>> errors(numLayers);
    errors[numLayers - 1].resize(sizes[numLayers - 1]);

    double totalErr = 0;

    // 2. Вычисление ошибки на выходном слое
    for (int i = 0; i < sizes[numLayers - 1]; i++) {
        double e = target[i] - pred[i];
        // Производная функции активации (сигмоиды): pred * (1 - pred)
        errors[numLayers - 1][i] = e * pred[i] * (1.0 - pred[i]);
        totalErr += e * e;
    }

    // 3. Обратное распространение ошибки (Backpropagation) по скрытым слоям
    for (int i = numLayers - 2; i > 0; i--) {
        errors[i].resize(sizes[i]);

        // Параллелим вычисления для каждого нейрона в слое
        OMP_PARALLEL
            for (int j = 0; j < sizes[i]; j++) {
                double e = 0;
                for (int k = 0; k < sizes[i + 1]; k++) {
                    e += errors[i + 1][k] * weights[i][k][j];
                }
                // outputs[i][j] — это сохраненный результат активации этого нейрона из feedForward
                errors[i][j] = e * outputs[i][j] * (1.0 - outputs[i][j]);
            }
    }

    // 4. Обновление весов и смещений (Biases)
    for (int i = 0; i < numLayers - 1; i++) {

        // Параллелим обновление весов следующего слоя
        OMP_PARALLEL
            for (int j = 0; j < sizes[i + 1]; j++) {
                for (int k = 0; k < sizes[i]; k++) {
                    // Градиентный спуск: прибавляем (lr * ошибка * входной сигнал)
                    weights[i][j][k] += lr * errors[i + 1][j] * outputs[i][k];
                }
                biases[i][j] += lr * errors[i + 1][j];
            }
    }

    return totalErr;
}