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

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

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 last = 2*M_PI * f * (t(t.size() - 1) + t(1)) + initPhase;
    timeSeries.array() = amp * Eigen::sin(2*M_PI * f * t + initPhase);
    return last;
}

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;
    segl -= segl(0);
    segl = initPhase
        + 2 * M_PI * fInit * segl
        + accel / 6 * Eigen::pow(segl,3);
    // segr.array() -= tau / 2;
    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 last = segr(segr.size() - 1) + dt * 2*M_PI * fFinal;
    timeSeries = amp * Eigen::sin(timeSeries.array());
    return last;
}

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