#pragma once

/**
 * @file spectral.hpp
 * @brief Spectral Sensor and Controller — grokking awareness via rank tracking.
 *
 * Port of: packages/fces/core/spectral_sensor.py + spectral_controller.py
 */

#include <string>
#include <unordered_map>
#include <vector>

#include <torch/torch.h>

namespace fces {

/**
 * SpectralSensor — tracks the effective rank of weight matrices.
 *
 * Used for grokking detection: when the model transitions from
 * memorization to generalization, the effective rank drops.
 */
class SpectralSensor {
public:
    SpectralSensor() = default;
    explicit SpectralSensor(torch::nn::Module& model);

    /// Track a layer's weight tensor
    void track_layer(const std::string& name, const torch::Tensor& weight);

    /// Get the global (average) effective rank
    float get_global_rank() const;

    /// Reset all tracked layers
    void reset();

private:
    std::unordered_map<std::string, float> layer_ranks_;

    /// Compute effective rank via SVD
    static float compute_effective_rank(const torch::Tensor& weight);
};

/**
 * SpectralController — converts spectral rank into optimizer decisions.
 */
class SpectralController {
public:
    /// Compute the spectral alpha (gating factor for rank-aware updates)
    float compute_alpha(float global_rank, float grokking_coefficient) const;
};

}  // namespace fces