#include "fces/spectral.hpp"
#include <numeric>

namespace fces {

SpectralSensor::SpectralSensor(torch::nn::Module& /*model*/) {}

void SpectralSensor::track_layer(const std::string& name, const torch::Tensor& weight) {
    if (weight.dim() >= 2) {
        layer_ranks_[name] = compute_effective_rank(weight);
    }
}

float SpectralSensor::get_global_rank() const {
    if (layer_ranks_.empty()) return 0.0f;
    float sum = 0.0f;
    for (const auto& [_, rank] : layer_ranks_) {
        sum += rank;
    }
    return sum / static_cast<float>(layer_ranks_.size());
}

void SpectralSensor::reset() {
    layer_ranks_.clear();
}

float SpectralSensor::compute_effective_rank(const torch::Tensor& weight) {
    // SVD-based effective rank (Shannon entropy of normalized singular values)
    auto svd_result = torch::svd(weight.to(torch::kFloat32));
    auto svd = std::get<1>(svd_result);
    auto s = svd / svd.sum();
    auto log_s = (s + 1e-10f).log();
    float entropy = -s.mul(log_s).sum().item<float>();
    return std::exp(entropy);
}

float SpectralController::compute_alpha(float global_rank, float grokking_coefficient) const {
    return global_rank * grokking_coefficient;
}

}  // namespace fces