#pragma once

#include <complex>
#include <vector>

#include "SimdTypes.h"
#include "VectorOps.h"

namespace staffpad {
template<typename T = float>
class SamplesFloat
{
public:
    ~SamplesFloat()
    {
        for (int ch = 0; ch < num_channels; ch++) {
            dealloc(ch);
        }
    }

    void setSize(int32_t numChannels, int32_t samples)
    {
        for (int ch = 0; ch < num_channels; ch++) {
            dealloc(ch);
        }

        num_channels = numChannels;
        num_samples = samples;
        data.resize(num_channels);
        for (int ch = 0; ch < num_channels; ch++) {
            alloc(ch, num_samples);
        }
    }

    int32_t getNumChannels() const
    {
        return num_channels;
    }

    int32_t getNumSamples() const
    {
        return num_samples;
    }

    float** getPtrs()
    {
        return data.data();
    }

    T* getPtr(int32_t channel)
    {
        assert(channel < num_channels);
        assert(data[channel]);
        return data[channel];
    }

    const T* getPtr(int32_t channel) const
    {
        assert(channel < num_channels);
        assert(data[channel]);
        return data[channel];
    }

    void assignSamples(int32_t channel, const T* input)
    {
        assert(channel < num_channels);
        assert(data[channel]);
        vo::copy(input, data[channel], num_samples);
    }

    void assignSamples(const SamplesFloat& rhs)
    {
        assert(num_channels == rhs.num_channels);
        assert(num_samples == rhs.num_samples);
        for (int ch = 0; ch < num_channels; ch++) {
            assert(data[ch]);
            vo::copy(rhs.getPtr(ch), getPtr(ch), num_samples);
        }
    }

    void zeroOut()
    {
        for (int ch = 0; ch < num_channels; ch++) {
            vo::setToZero(data[ch], num_samples);
        }
    }

private:
    int32_t num_channels{ 0 };
    int32_t num_samples{ 0 };
    std::vector<T*> data;

    void alloc(int32_t channel, int32_t samples)
    {
        assert(channel < num_channels);
        if (data[channel]) {
            dealloc(channel);
        }
        data[channel] = (T*)audio::simd::aligned_malloc(samples * sizeof(T), 64);
    }

    void dealloc(int32_t channel)
    {
        assert(channel < num_channels);
        if (data[channel]) {
            audio::simd::aligned_free(data[channel]);
            data[channel] = nullptr;
        }
    }
};

typedef SamplesFloat<float> SamplesReal;
typedef SamplesFloat<std::complex<float> > SamplesComplex;
} // namespace staffpad