#include "waveform.h"
#include "CSVhelper.hpp"
#include <iostream>
#include <fstream>
#include <vector>
#include <string>
#include <numeric>
#include <fcntl.h>
#include <sys/mman.h>
#include <unistd.h>

ArrayWaveform::ArrayWaveform() {}

ArrayWaveform::~ArrayWaveform() { unbindBuffer(); }

void* ArrayWaveform::getDataBuffer() { return this->pDataBuffer; }

int64_t ArrayWaveform::getDataLen() { return this->dataLen; }

void ArrayWaveform::setSamplingRate(ulong sr) {
    this->wfmParam.samplingRate = sr;
}

ulong ArrayWaveform::getSamplingRate() { return this->wfmParam.samplingRate; }

void ArrayWaveform::setFreqResolution(ulong fr) {
    this->wfmParam.freqResolution = fr;
}

ulong ArrayWaveform::getFreqResolution() { return this->wfmParam.freqResolution; }

void ArrayWaveform::setFreqTone(
    int centerFreq, 
    int freqSpacing,
    int numTones
) {
    int freqStart = centerFreq - 
        freqSpacing * int(std::floor(numTones / 2));
    this->wfmParam.freqTones.setLinSpaced(
        numTones,
        freqStart,
        freqStart + freqSpacing * (numTones-1)
    );
}

void ArrayWaveform::setFreqTone(const std::vector<int>& tones) {
    this->wfmParam.freqTones.resize(tones.size());
    for (auto i = 0; i < tones.size(); i++) {
        this->wfmParam.freqTones[i] = tones[i];
    }
}

std::vector<double> ArrayWaveform::getFreqTone() {
    std::vector<double> v(
        this->wfmParam.freqTones.data(),
        this->wfmParam.freqTones.data() + this->wfmParam.freqTones.size()
    );
    return v;
}

void ArrayWaveform::setPhase(const std::vector<double>& phases) {
    this->wfmParam.phases.resize(phases.size());
    for (auto i = 0; i < phases.size(); i++) {
        this->wfmParam.phases[i] = phases[i];
    }
}

std::vector<double> ArrayWaveform::getPhase() {
    std::vector<double> v(
        this->wfmParam.phases.data(),
        this->wfmParam.phases.data() + this->wfmParam.phases.size()
    );
  return v;
}

void ArrayWaveform::setAmplitude(const std::vector<double>& amplitudes) {
    this->wfmParam.amplitudes.resize(amplitudes.size());
    for (auto i = 0; i < amplitudes.size(); i++) {
        this->wfmParam.amplitudes[i] = amplitudes[i];
    }
}

std::vector<double> ArrayWaveform::getAmplitude() {
    std::vector<double> v(
        this->wfmParam.amplitudes.data(),
        this->wfmParam.amplitudes.data() + this->wfmParam.amplitudes.size()
    );
  return v;
}

void ArrayWaveform::setDefaultParam() {
    if (this->wfmParam.freqTones.size() == 0) {return;}
    auto numTones = this->wfmParam.freqTones.size();
    this->setAmplitude(std::vector(numTones, 2000.0));
    this->setPhase(std::vector(numTones, 0.0));
    this->wfmParam.samplingRate = 614.4e6;
    this->wfmParam.freqResolution = 1e3;
}

void ArrayWaveform::saveParam(std::string fileName) {
    const static Eigen::IOFormat csvFormat(
        Eigen::FullPrecision,
        Eigen::DontAlignCols,
        ",",
        ",",
        ""
    );
    std::ofstream saveFile(fileName);
    if (saveFile.is_open()) {
        saveFile << this->wfmParam.samplingRate << "\n";
        saveFile << this->wfmParam.freqResolution << "\n";
        saveFile
            << this->wfmParam.freqTones.format(csvFormat)
            << "\n";
        saveFile
            << this->wfmParam.phases.format(csvFormat)
            << "\n";
        saveFile
            << this->wfmParam.amplitudes.format(csvFormat)
            << "\n";
        saveFile.close();
    }
}

void ArrayWaveform::loadParam(std::string fileName) {
    std::ifstream file(fileName);
    if (!file.is_open()) {
        std::cout << fileName
            << " does not exists in loadParam()"
            << std::endl;
        return;
    }
    int lineCounter = 0;
    for (auto& line : CSVRange(file)) {
        std::vector<int> lineDataI;
        std::vector<double> lineDataD;
        switch (lineCounter) {
        case 0:
            this->wfmParam.samplingRate = std::stoi(std::string(line[0]));
            break;
        case 1:
            this->wfmParam.freqResolution = std::stoi(std::string(line[0]));
            break;
        case 2:
            for (auto i = 0; i < line.size(); i++) {
                lineDataI.push_back(std::stoi(std::string(line[i])));
            }
            setFreqTone(lineDataI);
            break;
        case 3:
            for (auto i = 0; i < line.size(); i++) {
                lineDataD.push_back(std::stod(std::string(line[i])));
            }
            setPhase(lineDataD);
            break;
        case 4:
            for (auto i = 0; i < line.size(); i++) {
                lineDataD.push_back(std::stod(std::string(line[i])));
            }
            setAmplitude(lineDataD);
            break;
        }
        lineCounter++;
    }
}

void ArrayWaveform::printParam() {
    std::cout << "sampling rate: " 
        << this->wfmParam.samplingRate << "\n";
    std::cout << "frequency resolution: "
        << this->wfmParam.freqResolution << "\n";
    std::cout << "frequency tones (MHz): " 
        << this->wfmParam.freqTones.transpose() / int(1e6) 
        << "\n";
    std::cout << "phases: " 
        << this->wfmParam.phases.transpose() << "\n";
    std::cout << "amplitudes: "
        << this->wfmParam.amplitudes.transpose() << "\n";
}

void ArrayWaveform::bindBuffer(int64_t bytes) {
    this->unbindBuffer();
    int fd = open("/dev/zero", O_RDONLY);
    this->pDataBuffer = mmap(
        NULL,
        bytes,
        PROT_READ | PROT_WRITE,
        MAP_PRIVATE,
        fd,
        0
    );
    if (this->pDataBuffer != MAP_FAILED) {
        memset(this->pDataBuffer, 0, bytes);
        this->dataLen = bytes / 2;
        int16_t* pData = (int16_t*) this->pDataBuffer;
    }
    close(fd);    
}

void ArrayWaveform::unbindBuffer() {
    if (this->dataLen != 0 and this->pDataBuffer != nullptr) {
        munmap(this->pDataBuffer, this->dataLen * 2);
        this->pDataBuffer = nullptr;
        this->dataLen = 0;
    }
}

ulong ArrayWaveform::getMinSampleLen(ulong samplingRate, ulong freqResolution) {
    return 2 * samplingRate / std::gcd(samplingRate, freqResolution);
}

ulong ArrayWaveform::getSampleLen(
    double tau,
    ulong samplingRate,
    ulong freqResolution
) {
    ulong sampleLen = ulong(tau * samplingRate);
    ulong minLen = getMinSampleLen(samplingRate, freqResolution);
    return sampleLen % minLen == 0 ? sampleLen : (sampleLen / minLen + 1) * minLen;
}

double ArrayWaveform::setStaticSegment(
    Eigen::Ref<Eigen::VectorXd> timeSeries,
    double f,
    double initPhase,
    double amp
) {
  Eigen::Ref<Eigen::ArrayXd> t = timeSeries.array();
  t -= t(0);
  double nextPhase = 2 * M_PI * f * (t(t.size() - 1) + t(1)) + initPhase;
  timeSeries.array() = amp * Eigen::sin(2 * M_PI * f * t + initPhase);
  return nextPhase;
}

double ArrayWaveform::setMovingSegment(
    Eigen::Ref<Eigen::VectorXd> timeSeries,
    double fInit,
    double fFinal,
    double initPhase,
    double amp
) {
    auto df = (fFinal - fInit) * 2 * M_PI;
    auto sampleLen = timeSeries.size();
    auto dt = timeSeries(1) - timeSeries(0);
    auto tau = timeSeries(sampleLen - 1) - timeSeries(0) + dt;
    auto accel = 4 * df / (tau * tau);
    unsigned int midIdx = int(sampleLen / 2.0);
    Eigen::Ref<Eigen::ArrayXd> segl = timeSeries(Eigen::seq(0, midIdx));
    Eigen::Ref<Eigen::ArrayXd> segr = timeSeries(Eigen::seq(midIdx+1, sampleLen-1));
    segr -= segl(0) + tau / 2;  // it may seem more intuitive to use
                                // segr -= segr(0) here, but doing so
                                // will cause phase jump
    segl -= segl(0);
    segl = initPhase
        + 2 * M_PI * fInit * segl
        + accel / 6 * Eigen::pow(segl,3);
    segr = segl(segl.size()-1)
        + (
            2 * M_PI * fInit + accel / 2 * std::pow((tau / 2),2)
        ) * segr
        + accel / 2 * tau / 2 * Eigen::pow(segr, 2)
        - accel / 6 * Eigen::pow(segr, 3);
    double nextPhase = segr(segr.size() - 1) + dt * 2*M_PI * fFinal;
    timeSeries = amp * Eigen::sin(timeSeries.array());
    return nextPhase;
}

std::pair<void*, int64_t> ArrayWaveform::getStaticWaveform() {
    auto minSampleLen = 2 * this->wfmParam.samplingRate
        / std::gcd(
            this->wfmParam.samplingRate,
            this->wfmParam.freqResolution
        );
    Eigen::VectorXd t = Eigen::VectorXd::LinSpaced(
        minSampleLen,
        0,
        minSampleLen - 1
    ) / this->wfmParam.samplingRate;
    Eigen::MatrixXd wfmMatrixrix = Eigen::sin(
        (
            (this->wfmParam.freqTones * t.transpose() * M_PI * 2)
        .array().colwise()
        + this->wfmParam.phases.array())
    ).array().colwise() * this->wfmParam.amplitudes.array();
    bindBuffer(minSampleLen * 2); // bytes
    auto pData = (int16_t*) this->pDataBuffer;
    Eigen::Map<EigenVectorXi16, Eigen::Aligned16> dataMap(pData, minSampleLen);
    dataMap = wfmMatrixrix.colwise().sum().cast<int16_t>();
    return std::pair(this->pDataBuffer, int64_t(minSampleLen));
}

std::pair<void*, int64_t> ArrayWaveform::getTrickWaveform(
    std::set<int> siteIndex,
    double df,
    double tauMove,
    double tauStay
) {
    if ((tauMove == 0 and tauStay == 0)
        or siteIndex.empty()) {
        return getStaticWaveform();
    }
    if (tauMove != 0 and tauStay != 0) {
        tauStay = std::ceil((tauMove + tauStay) * df) / df - tauMove;
    } else if (tauMove != 0) {
        tauMove = std::ceil(tauMove * df) / df;
    } else {
        tauStay = std::ceil(tauStay * df) / df;
    }
    auto tauTotal = tauMove * 2 + tauStay;
    ulong sampleLen = getSampleLen(
        tauTotal,
        this->wfmParam.samplingRate,
        this->wfmParam.freqResolution
    );
    ulong moveLen = tauMove * this->wfmParam.samplingRate;
    ulong stayLen = tauStay * this->wfmParam.samplingRate;
    auto idxSeg0 = moveLen;
    auto idxSeg1 = moveLen + stayLen;
    auto idxSeg2 = moveLen + stayLen + moveLen;
    auto t = Eigen::ArrayXd::LinSpaced(sampleLen, 0, sampleLen - 1) / this->wfmParam.samplingRate;
    Eigen::MatrixXd wfmMatrix(sampleLen, this->wfmParam.freqTones.size()); // site wfm stored column by column
    for (auto i = 0; i < this->wfmParam.freqTones.size(); i++) {
        wfmMatrix.col(i) = t;
        Eigen::Ref<Eigen::ArrayXd> siteWfm = wfmMatrix.col(i);
        auto fInit = this->wfmParam.freqTones[i];
        auto fFinal = fInit + df;
        auto phi = this->wfmParam.phases[i];
        auto amp = this->wfmParam.amplitudes[i];
        if (siteIndex.contains(i)) {
            if (tauMove != 0) {
                phi = setMovingSegment(
                    siteWfm(Eigen::seq(0, idxSeg0 - 1)),
                    fInit, fFinal, phi, amp
                );
            }
            if (tauStay != 0) {
                phi = setStaticSegment(
                    siteWfm(Eigen::seq(idxSeg0, idxSeg1 - 1)),
                    fFinal, phi, amp
                );
            }
            if (tauMove != 0) {
                phi = setMovingSegment(
                    siteWfm(Eigen::seq(idxSeg1, idxSeg2 - 1)),
                    fFinal, fInit, phi, amp
                );
            }
            setStaticSegment(
                siteWfm(Eigen::seq(idxSeg2, siteWfm.size() - 1)),
                fInit, phi, amp
            );
        } else {
            setStaticSegment(siteWfm, fInit, phi, amp);
        }
    }
    bindBuffer(sampleLen * 2);
    auto pData = (int16_t*) this->pDataBuffer;
    Eigen::Map<EigenVectorXi16, Eigen::Aligned16> dataMap(pData, sampleLen);
    dataMap = wfmMatrix.rowwise().sum().cast<int16_t> ();
    return std::pair(this->pDataBuffer, int64_t(sampleLen));
}

void ArrayWaveform::saveWaveform(std::string fileName) {
    const static Eigen::IOFormat csvFormat(
        Eigen::FullPrecision,
        Eigen::DontAlignCols,
        "",
        ",",
        ""
    );
    auto pData = (int16_t*) this->pDataBuffer;
    auto wfm = Eigen::Map<EigenVectorXi16>(pData, this->dataLen);
    std::ofstream saveFile(fileName);
    if (saveFile.is_open()) {
        saveFile << wfm.format(csvFormat);
        saveFile.close();
    }
}