#pragma once
#include "../layers/layer.hpp"
#include "../layers/dense.hpp"
#include <vector>
#include <memory>

namespace newnet {

class Sequential {
public:
    std::vector<std::unique_ptr<Layer>> layers;
    
    // Add a layer to the network
    void add(Layer* layer) {
        layers.emplace_back(layer);
    }
    
    // Forward pass: input flows through each layer in order
    // Layer 0 output → Layer 1 input → Layer 1 output → Layer 2 input → ...
    Tensor forward(const Tensor& input) {
        Tensor current = input;
        for (auto& layer : layers) {
            current = layer->forward(current);
        }
        return current;
    }
    
    // Backward pass: gradient flows through each layer in REVERSE order
    // This is the chain rule: dL/dx1 = dL/dx3 * dx3/dx2 * dx2/dx1
    void backward(const Tensor& grad_output) {
        Tensor current_grad = grad_output;
        for (int i = (int)layers.size() - 1; i >= 0; i--) {
            current_grad = layers[i]->backward(current_grad);
        }
    }
    
    // Collect all learnable parameters from all layers
    std::vector<Tensor*> parameters() {
        std::vector<Tensor*> params;
        for (auto& layer : layers) {
            auto layer_params = layer->parameters();
            params.insert(params.end(), layer_params.begin(), layer_params.end());
        }
        return params;
    }
};

} // namespace newnet