#pragma once

#include <openvr_driver.h>
#include <random>
#include <chrono>
#include <cmath>

namespace sauna {

/**
 * @brief Utility class to generate synthetic IMU data for testing
 */
class IMUDataGenerator {
private:
    std::mt19937 rng;
    std::normal_distribution<float> accelNoise;
    std::normal_distribution<float> gyroNoise;
    
    float accelNoiseLevel;
    float gyroNoiseLevel;
    
    // Base values for acceleration (gravity)
    float baseAccelX;
    float baseAccelY;
    float baseAccelZ;
    
    // Base values for angular velocity
    float baseGyroX;
    float baseGyroY;
    float baseGyroZ;
    
    // Timestamp of the last sample
    uint64_t lastTimestamp;
    
    // Motion pattern parameters
    bool useMotionPattern;
    float motionAmplitude;
    float motionFrequency;
    float motionPhase;
    
public:
    /**
     * @brief Construct a new IMUDataGenerator
     * 
     * @param accelNoise Noise level for accelerometer data (standard deviation)
     * @param gyroNoise Noise level for gyroscope data (standard deviation)
     */
    IMUDataGenerator(float accelNoise = 0.01f, float gyroNoise = 0.005f)
        : rng(std::random_device{}())
        , accelNoise(0.0f, accelNoise)
        , gyroNoise(0.0f, gyroNoise)
        , accelNoiseLevel(accelNoise)
        , gyroNoiseLevel(gyroNoise)
        , baseAccelX(0.0f)
        , baseAccelY(0.0f)
        , baseAccelZ(-9.81f) // Default gravity along negative Z axis
        , baseGyroX(0.0f)
        , baseGyroY(0.0f)
        , baseGyroZ(0.0f)
        , lastTimestamp(0)
        , useMotionPattern(false)
        , motionAmplitude(0.0f)
        , motionFrequency(0.0f)
        , motionPhase(0.0f)
    {
    }
    
    /**
     * @brief Set the base acceleration values
     * 
     * @param x X component of acceleration (m/s^2)
     * @param y Y component of acceleration (m/s^2)
     * @param z Z component of acceleration (m/s^2)
     */
    void SetBaseAcceleration(float x, float y, float z) {
        baseAccelX = x;
        baseAccelY = y;
        baseAccelZ = z;
    }
    
    /**
     * @brief Set the base angular velocity values
     * 
     * @param x X component of angular velocity (rad/s)
     * @param y Y component of angular velocity (rad/s)
     * @param z Z component of angular velocity (rad/s)
     */
    void SetBaseAngularVelocity(float x, float y, float z) {
        baseGyroX = x;
        baseGyroY = y;
        baseGyroZ = z;
    }
    
    /**
     * @brief Set the noise levels for the IMU data
     * 
     * @param accelNoise Noise level for accelerometer data (standard deviation)
     * @param gyroNoise Noise level for gyroscope data (standard deviation)
     */
    void SetNoiseLevel(float accelNoise, float gyroNoise) {
        accelNoiseLevel = accelNoise;
        gyroNoiseLevel = gyroNoise;
        this->accelNoise = std::normal_distribution<float>(0.0f, accelNoise);
        this->gyroNoise = std::normal_distribution<float>(0.0f, gyroNoise);
    }
    
    /**
     * @brief Enable a sinusoidal motion pattern for testing
     * 
     * @param amplitude Amplitude of the motion
     * @param frequency Frequency of the motion (Hz)
     * @param phase Initial phase of the motion (radians)
     */
    void EnableMotionPattern(float amplitude, float frequency, float phase = 0.0f) {
        useMotionPattern = true;
        motionAmplitude = amplitude;
        motionFrequency = frequency;
        motionPhase = phase;
    }
    
    /**
     * @brief Disable the motion pattern
     */
    void DisableMotionPattern() {
        useMotionPattern = false;
    }
    
    /**
     * @brief Generate a new IMU sample
     * 
     * @return vr::ImuSample_t The generated IMU sample
     */
    vr::ImuSample_t GenerateSample() {
        vr::ImuSample_t sample;
        
        // Generate timestamp
        uint64_t now = std::chrono::duration_cast<std::chrono::microseconds>(
            std::chrono::high_resolution_clock::now().time_since_epoch()).count();
        
        if (lastTimestamp == 0) {
            lastTimestamp = now;
        }
        
        sample.nTimeInMicroSec = now;
        
        // Calculate motion pattern if enabled
        float motionOffset = 0.0f;
        if (useMotionPattern) {
            float timeSeconds = static_cast<float>(now - lastTimestamp) / 1000000.0f;
            motionOffset = motionAmplitude * std::sin(2.0f * M_PI * motionFrequency * timeSeconds + motionPhase);
        }
        
        // Generate accelerometer data
        sample.vAccel.v[0] = baseAccelX + motionOffset + accelNoise(rng);
        sample.vAccel.v[1] = baseAccelY + motionOffset + accelNoise(rng);
        sample.vAccel.v[2] = baseAccelZ + accelNoise(rng);
        
        // Generate gyroscope data
        sample.vGyro.v[0] = baseGyroX + gyroNoise(rng);
        sample.vGyro.v[1] = baseGyroY + gyroNoise(rng);
        sample.vGyro.v[2] = baseGyroZ + gyroNoise(rng);
        
        lastTimestamp = now;
        
        return sample;
    }
    
    /**
     * @brief Generate multiple IMU samples
     * 
     * @param count Number of samples to generate
     * @param sampleRate Sample rate in Hz
     * @return std::vector<vr::ImuSample_t> Vector of generated IMU samples
     */
    std::vector<vr::ImuSample_t> GenerateSamples(int count, float sampleRate) {
        std::vector<vr::ImuSample_t> samples;
        samples.reserve(count);
        
        uint64_t sampleInterval = static_cast<uint64_t>(1000000.0f / sampleRate);
        uint64_t now = std::chrono::duration_cast<std::chrono::microseconds>(
            std::chrono::high_resolution_clock::now().time_since_epoch()).count();
        
        if (lastTimestamp == 0) {
            lastTimestamp = now;
        }
        
        for (int i = 0; i < count; i++) {
            vr::ImuSample_t sample;
            
            // Generate timestamp
            uint64_t timestamp = lastTimestamp + i * sampleInterval;
            sample.nTimeInMicroSec = timestamp;
            
            // Calculate motion pattern if enabled
            float motionOffset = 0.0f;
            if (useMotionPattern) {
                float timeSeconds = static_cast<float>(i) / sampleRate;
                motionOffset = motionAmplitude * std::sin(2.0f * M_PI * motionFrequency * timeSeconds + motionPhase);
            }
            
            // Generate accelerometer data
            sample.vAccel.v[0] = baseAccelX + motionOffset + accelNoise(rng);
            sample.vAccel.v[1] = baseAccelY + motionOffset + accelNoise(rng);
            sample.vAccel.v[2] = baseAccelZ + accelNoise(rng);
            
            // Generate gyroscope data
            sample.vGyro.v[0] = baseGyroX + gyroNoise(rng);
            sample.vGyro.v[1] = baseGyroY + gyroNoise(rng);
            sample.vGyro.v[2] = baseGyroZ + gyroNoise(rng);
            
            samples.push_back(sample);
        }
        
        lastTimestamp = now + (count - 1) * sampleInterval;
        
        return samples;
    }
};

} // namespace sauna