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

#include "nn/module.h"
#include "ops/ops.h"

#include <string>

namespace forge::nn {

inline ops::QuantMode quant_mode_from(const std::string& s) {
    if (s == "int8") return ops::QuantMode::Int8;
    if (s == "ternary") return ops::QuantMode::Ternary;
    return ops::QuantMode::None;
}

// y = x ⋅ Wᵀ (+ b). Weight stored [out, in] (PyTorch convention), x is
// [N, in] 2-D — callers flatten [B,T,C] with Var::reshaped.
//
// Research seams (see CLAUDE.md):
//  - set_mask(): optional binary mask multiplied into the weight each
//    forward (lottery-ticket / pruning). Gradients of masked weights are
//    zeroed by the same multiply.
//  - forward() is virtual so quantized-weight variants (BitNet-style) can
//    override without touching call sites.
class Linear : public Module {
public:
    Linear(int64_t in_features, int64_t out_features, bool has_bias, float weight_std,
           std::mt19937_64& rng) {
        weight_ = register_param("weight",
                                 normal_init({out_features, in_features}, weight_std, rng));
        if (has_bias) bias_ = register_param("bias", Tensor::zeros({out_features}));
    }

    virtual ~Linear() = default;

    virtual Var forward(const Var& x) const;

    // mask: [out, in], 0/1 f32. Undefined tensor clears the mask.
    void set_mask(Tensor mask) {
        mask_ = mask.defined() ? Var(std::move(mask), /*requires_grad=*/false) : Var();
    }

    // QAT: fake-quantize the weight each forward (STE backward). The master
    // weight and optimizer state stay f32.
    void set_quant(ops::QuantMode mode) { quant_ = mode; }

    const Var& weight() const { return weight_; }

protected:
    Var weight_;
    Var bias_; // undefined if bias disabled
    Var mask_; // undefined if no mask
    ops::QuantMode quant_ = ops::QuantMode::None;
};

} // namespace forge::nn