// Author: Simon-Pierre Boucher — contact@spboucher.ai
#pragma once

#include "core/autograd.h"
#include "ops/cpu/cpu_ops.h"

#include <random>
#include <string>
#include <utility>
#include <vector>

namespace forge::nn {

// Base module: owns nothing but the parameter registry. Children register
// their parameters at construction; parents absorb them with a name prefix.
// Tied parameters may appear under several names — consumers that update
// parameters (optimizer, checkpoint) must dedupe by Var::id().
class Module {
public:
    virtual ~Module() = default;

    const std::vector<std::pair<std::string, Var>>& named_parameters() const {
        return params_;
    }

    void zero_grad() const {
        for (const auto& [name, p] : params_) p.zero_grad();
    }

protected:
    Var register_param(const std::string& name, Tensor init) {
        Var p(std::move(init), /*requires_grad=*/true);
        params_.emplace_back(name, p);
        return p;
    }

    // Register an existing Var under a new name (weight tying).
    Var register_param(const std::string& name, const Var& shared) {
        params_.emplace_back(name, shared);
        return shared;
    }

    void absorb(const std::string& prefix, const Module& child) {
        for (const auto& [name, p] : child.named_parameters())
            params_.emplace_back(prefix + "." + name, p);
    }

    std::vector<std::pair<std::string, Var>> params_;
};

// init helper shared by all modules
inline Tensor normal_init(std::vector<int64_t> shape, float stddev, std::mt19937_64& rng) {
    Tensor t = Tensor::empty(std::move(shape));
    cpu::fill_normal(t, 0.0f, stddev, rng);
    return t;
}

} // namespace forge::nn