package com.naaturel.ANN.domain.model.neuron;

import com.naaturel.ANN.domain.abstraction.Model;
import com.naaturel.ANN.domain.abstraction.Network;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.function.Consumer;

/**
 * Represents a fully connected neural network
 */
public class FullyConnectedNetwork implements Model {

    private final List<Layer> layers;;
    private final Map<Neuron, List<Neuron>> connectionMap;

    public FullyConnectedNetwork(List<Layer> layers) {
        this.layers = layers;
        this.connectionMap = this.createConnectionMap();
    }

    @Override
    public List<Float> predict(List<Input> inputs) {
        List<Input> previousLayerOutputs = new ArrayList<>(inputs);
        for(Layer layer : this.layers){
            List<Float> currentLayerOutputs = layer.predict(previousLayerOutputs);
            previousLayerOutputs = currentLayerOutputs.stream().map(Input::new).toList();
        }
        return previousLayerOutputs.stream().map(Input::getValue).toList();
    }

    @Override
    public int synCount() {
        int res = 0;
        for(Layer layer : this.layers){
            res += layer.synCount();
        }
        return res;
    }

    @Override
    public int neuronCount() {
        int res = 0;
        for(Layer layer : this.layers){
            res += layer.neuronCount();
        }
        return res;
    }

    @Override
    public void forEachSynapse(Consumer<Synapse> consumer) {
        this.layers.forEach(layer -> layer.forEachSynapse(consumer));
    }

    @Override
    public void forEachNeuron(Consumer<Neuron> consumer) {
        this.layers.forEach(layer -> layer.forEachNeuron(consumer));
    }

    @Override
    public void forEachOutputNeurons(Consumer<Neuron> consumer) {
        this.layers.getLast().forEachNeuron(consumer);
    }

    @Override
    public void forEachNeuronConnectedTo(Neuron n, Consumer<Neuron> consumer) {
        this.connectionMap.get(n).forEach(consumer);
    }

    private Map<Neuron, List<Neuron>> createConnectionMap() {
        Map<Neuron, List<Neuron>> res = new HashMap<>();

        for (int i = 0; i < this.layers.size() - 1; i++) {
            List<Neuron> nextLayerNeurons = new ArrayList<>();
            this.layers.get(i + 1).forEachNeuron(nextLayerNeurons::add);
            this.layers.get(i).forEachNeuron(n -> res.put(n, nextLayerNeurons));
        }

        return res;
    }
}
