#include "Waves.h"

namespace TMM {

ArrayXd TukeyFunc(const ArrayXd& xs, double w0, double a) {
    ArrayXd res(xs.size());
    for (int i = 0; i < xs.size(); i++) {
        double x = xs(i);

        if (x < -0.5 * w0 * (2.0 - a)) {
            res(i) = 0.0;
        } else if (x > 0.5 * w0 * (2.0 - a)) {
            res(i) = 0.0;
        } else if (x > 0.5 * a * w0) {
            res(i) = 0.5 * std::cos(PI / (w0 * (1.0 - a)) * (x - 0.5 * a * w0)) + 0.5;
        } else if (x < -0.5 * a * w0) {
            res(i) = 0.5 * std::cos(PI / (w0 * (1.0 - a)) * (x + 0.5 * (2.0 - a) * w0) + PI) + 0.5;
        } else {
            res(i) = 1.0;
        }
    }
    return res;
}

void Wave::SolvePlaneWave() {
    maxXThis = (0.5 * w0 / std::cos(thLayer0));

    // Plane waves
    phis = ArrayXd(1);
    kxs = ArrayXd(1);
    kzs = ArrayXd(1);
    expansionCoefsKx = ArrayXcd(1);
    fieldProfileXs = ArrayXd(0);
    fieldProfile = ArrayXd(0);
    phis(0) = 0.0;
    kxs(0) = k * std::sin(thLayer0);
    kzs(0) = k * std::cos(thLayer0);
    expansionCoefsKx(0) = E0;
}

void Wave::SolveFFTWave() {
    if (nPointsInteg < 2) {
        throw std::runtime_error("nPointsInteg too small");
    }

    double maxXByPhi = (nPointsInteg / 2) * PI / (std::sin(maxPhi) * k);
    // Calc maxX
    if (dynamicMaxX) {
        maxXThis = dynamicMaxXCoef * (0.5 * w0 / std::cos(thLayer0)) + dynamicMaxXAddition;
    } else {
        maxXThis = maxX;
    }
    maxXThis = std::max(maxXThis, maxXByPhi);

    // Init xs
    ArrayXd xs(nPointsInteg);
    xs = ArrayXd::LinSpaced(nPointsInteg, -maxXThis, maxXThis);
    double dx = xs(1) - xs(0);

    // Field profile
    fieldProfileXs = xs;
    switch (waveType) {
    case WaveType::GAUSSIANWAVE:
        fieldProfile = E0 * (-(xs * xs) / sqr(w0 / 1.2533141369277430)).exp();
        break;
    case WaveType::TUKEYWAVE:
        fieldProfile = E0 * TukeyFunc(xs, w0 / (0.25 * a + 0.75), a);
        break;
    default:
        throw std::invalid_argument("Unknown FFT wave type.");
    }

    // FFT spectrum
    ArrayXd shiftedFieldProfile = IFFTShift(fieldProfile);
    ArrayXcd fieldProfileSpectrum(fieldProfile.size());
    fft.fwd(&fieldProfileSpectrum(0), &shiftedFieldProfile(0), fieldProfile.size());
    fieldProfileSpectrum = FFTShift(fieldProfileSpectrum);
    fieldProfileSpectrum *= dx / (2.0 * PI);
    ArrayXd kxPs = 2.0 * PI * FFTFreq(fieldProfile.size(), dx);
    phis = (kxPs / k).asin();

    // Check for backward-propagating waves
    if (phis.maxCoeff() + thLayer0 >= PI / 2.0) {
        throw std::runtime_error("Phi larger than 90 deg.");
    }

    // Save result
    ArrayXd kzPs = phis.cos() * k;
    kxs = kxPs * std::cos(thLayer0) + kzPs * std::sin(thLayer0);
    kzs = -kxPs * std::sin(thLayer0) + kzPs * std::cos(thLayer0);
    expansionCoefsKx = fieldProfileSpectrum / std::cos(thLayer0);
}

void Wave::SolveSpdcWave() {
    // Checks
    if (std::isnan(deltaKxSpdc)) {
        throw std::runtime_error("SPDC wave can only be used through SecondOrderNLTMM");
    }
    maxXThis = (0.5 * w0 / std::cos(thLayer0));

    // Init arrays
    phis = ArrayXd(1);
    fieldProfileXs = ArrayXd(0);
    fieldProfile = ArrayXd(0);

    double kxMid = k * std::sin(thLayer0);
    kxs = ArrayXd::LinSpaced(nPointsInteg, kxMid - deltaKxSpdc, kxMid + deltaKxSpdc);
    kzs = (k * k - kxs.pow(2)).sqrt();
    expansionCoefsKx = ArrayXcd::Ones(nPointsInteg);
    phis(0) = 0.0; // Currently not calculated (not needed)
}

Wave::Wave() : phis(0), kxs(0), kzs(0), expansionCoefsKx(0), fieldProfileXs(0), fieldProfile(0) {
    waveType = WaveType::PLANEWAVE;
    pwr = 1.0;
    overrideE0 = false;
    E0OverrideValue = 1.0;
    w0 = 100e-6;
    materialLayer0 = nullptr;
    materialLayerThis = nullptr;
    Ly = 1e-3;
    a = 0.7;
    nPointsInteg = 100;
    maxX = 1e-3;
    dynamicMaxX = true;
    dynamicMaxXCoef = 2.0;
    dynamicMaxXAddition = 0.0;
    maxXThis = 0.0;
    solved = false;
    maxPhi = 0.17;
    deltaKxSpdc = constNAN;

    wl = 0.0;
    beta = 0.0;
    k0 = 0.0;
    E0 = 0.0;
    k = 0.0;
    nLayer0 = 0.0;
    thLayer0 = 0.0;
}

void Wave::SetWaveType(WaveType waveType_) {
    waveType = waveType_;
}

void Wave::SetPwr(double pwr_) {
    pwr = pwr_;
}

void Wave::SetOverrideE0(bool overrideE0_) {
    overrideE0 = overrideE0_;
}

void Wave::SetE0(double E0_) {
    E0OverrideValue = E0_;
}

void Wave::SetW0(double w0_) {
    w0 = w0_;
}

void Wave::SetLy(double Ly_) {
    Ly = Ly_;
}

void Wave::SetA(double a_) {
    a = a_;
}

void Wave::SetNPointsInteg(int nPointsInteg_) {
    nPointsInteg = nPointsInteg_;
}

void Wave::SetMaxX(double maxX_) {
    maxX = maxX_;
}

void Wave::EnableDynamicMaxX(bool dynamicMaxX_) {
    dynamicMaxX = dynamicMaxX_;
}

void Wave::SetDynamicMaxXCoef(double dynamicMaxXCoef_) {
    dynamicMaxXCoef = dynamicMaxXCoef_;
}

void Wave::SetDynamicMaxXAddition(double dynamicMaxXAddition_) {
    dynamicMaxXAddition = dynamicMaxXAddition_;
}

void Wave::SetMaxPhi(double maxPhi_) {
    maxPhi = maxPhi_;
}

void Wave::SetParam(TMMParam param, double value) {
    switch (param) {
    case TMMParam::PARAM_WAVE_W0:
        SetW0(value);
        break;
    default:
        throw std::invalid_argument("Param not in list.");
    }
}

void Wave::Solve(double wl_, double beta_, Material* materialLayer0_, Material* materialLayerThis_,
                 double deltaKxSpdc_) {
    wl = wl_;
    beta = beta_;
    materialLayer0 = materialLayer0_;
    materialLayerThis = materialLayerThis_;
    deltaKxSpdc = deltaKxSpdc_;
    maxXThis = maxX;

    // Calc E0
    if (!overrideE0) {
        double I0 = pwr / (Ly * w0);
        E0 = std::sqrt(2.0 * constMu0 * constC * I0);
    } else {
        E0 = E0OverrideValue;
    }

    // Precalc variables
    k0 = 2.0 * PI / wl;
    nLayer0 = real(materialLayer0->GetN(wl));
    k = k0 * nLayer0;
    thLayer0 = std::asin(beta / nLayer0);
    beamArea = w0 * Ly / std::cos(thLayer0);

    // Solve
    if (waveType == WaveType::PLANEWAVE) {
        SolvePlaneWave();
    } else if (waveType == WaveType::GAUSSIANWAVE || waveType == WaveType::TUKEYWAVE) {
        SolveFFTWave();
    } else if (waveType == WaveType::SPDCWAVE) {
        SolveSpdcWave();
    } else {
        throw std::invalid_argument("Unknown wave type.");
    }
    solved = true;
}


// Getters

WaveType Wave::GetWaveType() const noexcept {
    return waveType;
}

double Wave::GetPwr() const noexcept {
    return pwr;
}

bool Wave::GetOverrideE0() const noexcept {
    return overrideE0;
}

double Wave::GetE0() const noexcept {
    return E0OverrideValue;
}

double Wave::GetW0() const noexcept {
    return w0;
}

double Wave::GetLy() const noexcept {
    return Ly;
}

double Wave::GetA() const noexcept {
    return a;
}

int Wave::GetNPointsInteg() const noexcept {
    return nPointsInteg;
}

double Wave::GetMaxX() const noexcept {
    return maxX;
}

bool Wave::IsDynamicMaxXEnabled() const noexcept {
    return dynamicMaxX;
}

double Wave::GetDynamicMaxXCoef() const noexcept {
    return dynamicMaxXCoef;
}

double Wave::GetDynamicMaxXAddition() const noexcept {
    return dynamicMaxXAddition;
}

double Wave::GetMaxPhi() const noexcept {
    return maxPhi;
}

double Wave::GetDouble(TMMParam param) const {
    switch (param) {
    case TMMParam::PARAM_WAVE_W0:
        return GetW0();
    default:
        throw std::invalid_argument("Param not in list.");
    }
}

pairdd Wave::GetXRange() const {
    if (!solved) {
        throw std::runtime_error("Wave must be solved first.");
    }
    return pairdd(-maxXThis, maxXThis);
}

ArrayXd Wave::GetBetas() const {
    if (!solved) {
        throw std::runtime_error("Wave must be solved first.");
    }
    return kxs / k0;
}

ArrayXd Wave::GetPhis() const noexcept {
    return phis;
}

ArrayXd Wave::GetKxs() const {
    if (!solved) {
        throw std::runtime_error("Wave must be solved first.");
    }
    return kxs;
}

ArrayXd Wave::GetKzs() const {
    if (!solved) {
        throw std::runtime_error("Wave must be solved first.");
    }
    return kzs;
}

ArrayXd Wave::GetFieldProfileXs() const {
    if (!solved) {
        throw std::runtime_error("Wave must be solved first.");
    }
    return fieldProfileXs;
}

ArrayXd Wave::GetFieldProfile() const {
    if (!solved) {
        throw std::runtime_error("Wave must be solved first.");
    }
    return fieldProfile;
}

ArrayXcd Wave::GetExpansionCoefsKx() const {
    if (!solved) {
        throw std::runtime_error("Wave must be solved first.");
    }
    return expansionCoefsKx;
}

double Wave::GetBeamArea() const {
    if (!solved) {
        throw std::runtime_error("Wave must be solved first.");
    }
    return beamArea;
}

bool Wave::IsCoherent() const noexcept {
    if (waveType == WaveType::SPDCWAVE) {
        return false;
    }
    return true;
}

} // namespace TMM
