#include "waveform.h"
#include "helper.hpp"
#include <iostream>
#include <fstream>
#include <vector>
#include <string>
#include <numeric>

ArrayWaveform::ArrayParam::ArrayParam() {
    this->samplingRate = 0;
    this->freqResolution = 0;
}

ArrayWaveform::ArrayParam::ArrayParam(const ArrayParam& other) {
    this->samplingRate = other.samplingRate;
    this->freqResolution = other.freqResolution;
    this->freqTones = other.freqTones;
    this->phases = other.phases;
    this->amplitudes = other.amplitudes;
}

void ArrayWaveform::ArrayParam::setSamplingRate(unsigned long sr) {
    this->samplingRate = sr;
}

void ArrayWaveform::ArrayParam::setFreqResolution(unsigned long fr) {
    this->freqResolution = fr;
}

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

void ArrayWaveform::ArrayParam::setFreqTone(const Eigen::VectorXd& tones) {
    this->freqTones.resize(0);
    this->freqTones.resize(tones.size());
    this->freqTones = tones;
}

Eigen::VectorXd ArrayWaveform::ArrayParam::getFreqTone() {
  return this->freqTones;
}

void ArrayWaveform::ArrayParam::setPhase(const Eigen::VectorXd& phases) {
    this->phases.resize(0);
    this->phases.resize((phases.size()));
    this->phases = phases;
}

Eigen::VectorXd ArrayWaveform::ArrayParam::getPhase() {
  return this->phases;
}

void ArrayWaveform::ArrayParam::setAmplitude(const Eigen::VectorXd& amplitudes) {
    this->amplitudes.resize(0);
    this->amplitudes.resize(amplitudes.size());
    this->amplitudes = amplitudes;
}

Eigen::VectorXd ArrayWaveform::ArrayParam::getAmplitude() {
  return this->amplitudes;
}

void ArrayWaveform::ArrayParam::setDefaultParam() {
    if (this->freqTones.size() == 0) {return;}
    auto numTones = this->freqTones.size();
    this->setAmplitude(Eigen::VectorXd::Ones(numTones) * std::pow(2, 12));
    this->setPhase(Eigen::VectorXd::Zero(numTones));
    this->samplingRate = 614.4e6;
    this->freqResolution = 1e3;
}

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

void ArrayWaveform::ArrayParam::loadParam(std::string fileName) {
    std::ifstream file(fileName);
    int lineCounter = 0;
    for (auto& line : CSVRange(file)) {
        std::vector<int> lineDataI;
        std::vector<double> lineDataD;
        switch (lineCounter) {
        case 0:
            this->samplingRate = std::stoi(std::string(line[1]));
            break;
        case 1:
            this->freqResolution = std::stoi(std::string(line[1]));
            break;
        case 2:
            for (auto i = 1; i < line.size(); i++) {
                lineDataI.push_back(std::stoi(std::string(line[i])));
            }
            this->freqTones.resize(lineDataI.size());
            this->freqTones(lineDataI.data());
            break;
        case 3:
            for (auto i = 1; i < line.size(); i++) {
                lineDataD.push_back(std::stod(std::string(line[i])));
            }
            this->phases.resize(lineDataD.size());
            this->phases(lineDataD.data());
            break;
        case 4:
            for (auto i = 1; i < line.size(); i++) {
                lineDataD.push_back(std::stod(std::string(line[i])));
            }
            this->amplitudes.resize(lineDataD.size());
            this->amplitudes(lineDataD.data());
            break;
        }
        lineCounter++;
    }
}

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

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;
}

EigenVectorXi16 ArrayWaveform::getStaticWaveform(const ArrayParam& param) {
    auto minSampleLen = 2 * param.samplingRate
        / std::gcd(param.samplingRate, param.freqResolution);
    Eigen::VectorXd t = Eigen::VectorXd::LinSpaced(
        minSampleLen,
        0,
        minSampleLen - 1
    ) / param.samplingRate;
    Eigen::MatrixXd wfmMatrix = Eigen::sin(
        ((param.freqTones * t.transpose() * M_PI * 2).array().colwise()
        + param.phases.array())
    ).array().colwise() * param.amplitudes.array();
    return wfmMatrix.colwise().sum().cast<int16_t>();
}

EigenVectorXi16 ArrayWaveform::getTrickWaveform(
    const ArrayParam& param,
    std::set<int> siteIndex,
    double df,
    double tauMove,
    double tauStay
) {
    if ((tauMove == 0 and tauStay == 0)
        or siteIndex.empty()) {
        return getStaticWaveform(param);
    }
    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, param.samplingRate, param.freqResolution);
    ulong moveLen = tauMove * param.samplingRate;
    ulong stayLen = tauStay * param.samplingRate;
    auto idxSeg0 = moveLen;
    auto idxSeg1 = moveLen + stayLen;
    auto idxSeg2 = moveLen + stayLen + moveLen;
    auto t = Eigen::ArrayXd::LinSpaced(sampleLen, 0, sampleLen - 1) / param.samplingRate;
    Eigen::MatrixXd wfmMat(sampleLen, param.freqTones.size()); // site wfm stored column by column
    for (auto i = 0; i < param.freqTones.size(); i++) {
        wfmMat.col(i) = t;
        Eigen::Ref<Eigen::ArrayXd> siteWfm = wfmMat.col(i);
        auto fInit = param.freqTones[i];
        auto fFinal = fInit + df;
        auto phi = param.phases[i];
        auto amp = param.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);
        }
    }
    return wfmMat.rowwise().sum().cast<int16_t> ();
}

void ArrayWaveform::saveCSV(std::string fileName, EigenVectorXi16 wfm) {
    const static Eigen::IOFormat csvFormat(
        Eigen::FullPrecision,
        Eigen::DontAlignCols,
        "",
        ",",
        ""
    );
    std::ofstream saveFile(fileName);
    if (saveFile.is_open()) {
        saveFile << wfm.format(csvFormat);
        saveFile.close();
    }
}