#include "SecondOrderNLTMM.h"

namespace TMM {

SweepResultSecondOrderNLTMM::SweepResultSecondOrderNLTMM(int n, int outmask_, int layerNr_, double layerZ_)
    : P1(n, outmask_, layerNr_, layerZ_), P2(n, outmask_, layerNr_, layerZ_),
      Gen(n, outmask_ & ~(SWEEP_ENH), layerNr_, layerZ_), wlsGen(n), betasGen(n) {
    outmask = outmask_;
}

void SweepResultSecondOrderNLTMM::SetValues(int nr, SecondOrderNLTMM& tmm) {
    wlsGen(nr) = tmm.GetGen()->GetWl();
    betasGen(nr) = tmm.GetGen()->GetBeta();

    if (outmask & SWEEP_P1) {
        P1.SetValues(nr, *tmm.GetP1());
    }

    if (outmask & SWEEP_P2) {
        P2.SetValues(nr, *tmm.GetP2());
    }

    if (outmask & SWEEP_GEN) {
        Gen.SetValues(nr, *tmm.GetGen());
    }
}

void SecondOrderNLTMM::CheckPrerequisites(TMMParam toIgnore) {
    // Sub TMMs
    tmmP1.CheckPrerequisites(toIgnore);
    tmmP2.CheckPrerequisites(toIgnore);
    tmmGen.CheckPrerequisites(toIgnore);
}

void SecondOrderNLTMM::UpdateGenParams() {
    double wlP1 = tmmP1.GetWl();
    double wlP2 = tmmP2.GetWl();
    double betaP1 = tmmP1.GetBeta();
    double betaP2 = tmmP2.GetBeta();

    switch (process) {
    case NonlinearProcess::SHG:
        throw std::runtime_error("Not implemented");
    case NonlinearProcess::SFG:
        wlGen = OmegaToWl(WlToOmega(wlP1) + WlToOmega(wlP2));
        betaGen = wlGen * (betaP1 / wlP1 + betaP2 / wlP2);
        break;
    case NonlinearProcess::DFG:
        wlGen = OmegaToWl(WlToOmega(wlP1) - WlToOmega(wlP2));
        betaGen = wlGen * (betaP1 / wlP1 - betaP2 / wlP2);
        break;
    case NonlinearProcess::SPDC:
        wlGen = OmegaToWl(WlToOmega(wlP1) - WlToOmega(wlP2));
        betaGen = wlGen * (betaP1 / wlP1 - betaP2 / wlP2);
        tmmP2.UpdateSPDCParams(deltaWlSpdc, solidAngleSpdc, deltaThetaSpdc, tmmP1.GetWl(), tmmP1.GetBeta());
        break;
    default:
        throw std::runtime_error("Unknown process.");
    }

    if (wlGen <= 0.0) {
        throw std::runtime_error("Wavelength cannot be negative.");
    }

    tmmGen.SetWl(wlGen);
    tmmGen.SetBeta(betaGen);
}

void SecondOrderNLTMM::CalcInhomogeneosWaveParams(int layerNr, Material* material, InhomogeneosWaveParams* kpS,
                                                  InhomogeneosWaveParams* kpA) {
    constexpr int iF = static_cast<int>(WaveDirection::F);
    dcomplex kzFP1 = tmmP1.GetLayer(layerNr)->GetHw()->GetKz()(iF);
    dcomplex kzFP2 = tmmP2.GetLayer(layerNr)->GetHw()->GetKz()(iF);
    Fields FP1F = tmmP1.GetLayer(layerNr)->GetFields(0.0, WaveDirection::F);
    Fields FP2F = tmmP2.GetLayer(layerNr)->GetFields(0.0, WaveDirection::F);
    Fields FP1B = tmmP1.GetLayer(layerNr)->GetFields(0.0, WaveDirection::B);
    Fields FP2B = tmmP2.GetLayer(layerNr)->GetFields(0.0, WaveDirection::B);
    Chi2Tensor& chi2 = material->chi2;

    switch (process) {
    case NonlinearProcess::SHG:
        throw std::runtime_error("Not implemented");
    case NonlinearProcess::SFG:
        kpS->kSzF = kzFP1 + kzFP2;
        kpA->kSzF = kzFP1 - kzFP2;
        kpS->pF = chi2.GetNonlinearPolarization(FP1F.E, FP2F.E);
        kpS->pB = chi2.GetNonlinearPolarization(FP1B.E, FP2B.E);
        kpA->pF = chi2.GetNonlinearPolarization(FP1F.E, FP2B.E);
        kpA->pB = chi2.GetNonlinearPolarization(FP1B.E, FP2F.E);
        break;
    case NonlinearProcess::DFG:
    case NonlinearProcess::SPDC:
        kpS->kSzF = kzFP1 - std::conj(kzFP2);
        kpA->kSzF = kzFP1 + std::conj(kzFP2);
        kpS->pF = chi2.GetNonlinearPolarization(FP1F.E, FP2F.E.conjugate());
        kpS->pB = chi2.GetNonlinearPolarization(FP1B.E, FP2B.E.conjugate());
        kpA->pF = chi2.GetNonlinearPolarization(FP1F.E, FP2B.E.conjugate());
        kpA->pB = chi2.GetNonlinearPolarization(FP1B.E, FP2F.E.conjugate());
        break;
    default:
        throw std::runtime_error("Unknown process.");
    }

    // Check inhomogenous wave directionality, correct if necessary
    /*
    dcomplex epsL = sqr(material->GetN(tmmGen.GetDouble(PARAM_WL)));
    Polarization pol = (Polarization)tmmGen.GetInt(PARAM_POL);
    WaveDirection dS = F, dA = F;
    if (kpS->kSzF != 0.0) {
        dS = GetWaveDirection(kpS->kSzF, epsL, pol);
    }
    if (kpA->kSzF != 0.0) {
        dA = GetWaveDirection(kpA->kSzF, epsL, pol);
    }


    if (dS != F) {
        kpS->kSzF *= -1.0;
        swap(kpS->pF, kpS->pB);
    }

    if (dA != F) {
        kpA->kSzF *= -1.0;
        swap(kpA->pF, kpA->pB);
    }
    */
}

void SecondOrderNLTMM::SolveFundamentalFields() {
    // Solve pump 1
    tmmP1.Solve();
    tmmP2.Solve();
}
void SecondOrderNLTMM::SolveGeneratedField() {
    // Insert nonlinearities
    for (int i = 0; i < tmmGen.LayersCount(); i++) {
        Material* material = tmmGen.GetLayer(i)->GetMaterial();
        if (material->IsNonlinear()) {
            InhomogeneosWaveParams kpS, kpA;
            CalcInhomogeneosWaveParams(i, material, &kpS, &kpA);
            tmmGen.GetLayer(i)->SetNonlinearity(kpS, kpA);
        } else {
            tmmGen.GetLayer(i)->ClearNonlinearity();
        }
    }

    // Solve generated fields
    tmmGen.Solve();
}

void SecondOrderNLTMM::SolveWaves(ArrayXd* betasP1, ArrayXcd* E0sP1, ArrayXd* betasP2, ArrayXcd* E0sP2) {
    Material* matLayer0 = tmmP1.GetLayer(0)->GetMaterial();
    Material* matLayerThis = nullptr;

    // Solve wave P1
    Wave* waveP1 = tmmP1.GetWave();
    waveP1->Solve(tmmP1.GetWl(), tmmP1.GetBeta(), matLayer0, matLayerThis);
    *betasP1 = waveP1->GetBetas();
    *E0sP1 = waveP1->GetExpansionCoefsKx();

    // Solve wave P2
    Wave* waveP2 = tmmP2.GetWave();
    if (waveP2->GetWaveType() == WaveType::SPDCWAVE && process != NonlinearProcess::SPDC) {
        throw std::invalid_argument("SecondOrderNLTMM must be in SPDC mode to use SPDC wave.");
    }
    double deltaKxSpdc = tmmP2.CalcDeltaKxSpdc();
    waveP2->Solve(tmmP2.GetWl(), tmmP2.GetBeta(), matLayer0, matLayerThis, deltaKxSpdc);
    *betasP2 = waveP2->GetBetas();
    *E0sP2 = waveP2->GetExpansionCoefsKx();
}
SecondOrderNLTMM::SecondOrderNLTMM() {
    SetProcess(NonlinearProcess::SFG);
    tmmP1.SetMode(NonlinearTmmMode::MODE_INCIDENT);
    tmmP2.SetMode(NonlinearTmmMode::MODE_INCIDENT);
    tmmGen.SetMode(NonlinearTmmMode::MODE_NONLINEAR);
    deltaWlSpdc = constNAN;
    solidAngleSpdc = constNAN;
    deltaThetaSpdc = constNAN;
}

void SecondOrderNLTMM::SetProcess(NonlinearProcess process_) {
    process = process_;
    if (process == NonlinearProcess::SPDC) {
        tmmP2.SetMode(NonlinearTmmMode::MODE_VACUUM_FLUCTUATIONS);
    } else {
        tmmP2.SetMode(NonlinearTmmMode::MODE_INCIDENT);
    }
}

void SecondOrderNLTMM::SetDeltaWlSpdc(double value) {
    deltaWlSpdc = value;
}

void SecondOrderNLTMM::SetSolidAngleSpdc(double value) {
    solidAngleSpdc = value;
}

void SecondOrderNLTMM::SetDeltaThetaSpdc(double value) {
    deltaThetaSpdc = value;
}

void SecondOrderNLTMM::AddLayer(double d_, Material* material_) {
    tmmP1.AddLayer(d_, material_);
    tmmP2.AddLayer(d_, material_);
    tmmGen.AddLayer(d_, material_);
}

void SecondOrderNLTMM::Solve() {
    UpdateGenParams();
    CheckPrerequisites();
    SolveFundamentalFields();
    SolveGeneratedField();
}

SecondOrderNLIntensities SecondOrderNLTMM::GetIntensities() const {
    SecondOrderNLIntensities res;
    res.P1 = tmmP1.GetIntensities();
    res.P2 = tmmP2.GetIntensities();
    res.Gen = tmmGen.GetIntensities();
    return res;
}

NonlinearTMM* SecondOrderNLTMM::GetP1() noexcept {
    return &tmmP1;
}

NonlinearTMM* SecondOrderNLTMM::GetP2() noexcept {
    return &tmmP2;
}

NonlinearTMM* SecondOrderNLTMM::GetGen() noexcept {
    return &tmmGen;
}

double SecondOrderNLTMM::GetDeltaWlSpdc() const noexcept {
    return deltaWlSpdc;
}

double SecondOrderNLTMM::GetSolidAngleSpdc() const noexcept {
    return solidAngleSpdc;
}

double SecondOrderNLTMM::GetDeltaThetaSpdc() const noexcept {
    return deltaThetaSpdc;
}

std::unique_ptr<SweepResultSecondOrderNLTMM> SecondOrderNLTMM::Sweep(TMMParam param,
                                                                     const Eigen::Map<ArrayXd>& valuesP1,
                                                                     const Eigen::Map<ArrayXd>& valuesP2, int outmask,
                                                                     int paramLayer, int layerNr, double layerZ) {
    if (valuesP1.size() != valuesP2.size()) {
        throw std::invalid_argument("Value arrays must have the same size.");
    }

    // Alloc memory for result
    auto res = std::make_unique<SweepResultSecondOrderNLTMM>(valuesP1.size(), outmask, layerNr, layerZ);

#pragma omp parallel
    {
        // int this_thread = omp_get_thread_num();
        // std::cout << "Thread started: " << this_thread << std::endl;

        // Make a copy
        SecondOrderNLTMM nlTMMCopy = *this;

// Sweep
#pragma omp for
        for (int i = 0; i < valuesP1.size(); i++) {
            // Set sweep param
            nlTMMCopy.GetP1()->SetParam(param, valuesP1(i), paramLayer);
            nlTMMCopy.GetP2()->SetParam(param, valuesP2(i), paramLayer);
            if (GetParamType(param) == TMMParamType::PTYPE_NONLINEAR_LAYER) {
                // NB: Gen param is set from P2 value; a more accurate approach
                // would recompute Gen from both P1 and P2 (e.g. for layer-d sweeps).
                nlTMMCopy.GetGen()->SetParam(param, valuesP2(i), paramLayer);
            }

            // Solver
            nlTMMCopy.Solve();
            res->SetValues(i, nlTMMCopy);
        }
    }
    return res;
}

std::unique_ptr<FieldsZX> SecondOrderNLTMM::WaveGetFields2D(const Eigen::Map<ArrayXd>& zs,
                                                            const Eigen::Map<ArrayXd>& xs, WaveDirection dir) {
    // Do checking
    UpdateGenParams();
    CheckPrerequisites();

    // Solve waves
    ArrayXd betasP1, betasP2;
    ArrayXcd E0sP1, E0sP2;
    SolveWaves(&betasP1, &E0sP1, &betasP2, &E0sP2);

    // kxs
    ArrayXd kxsP1 = betasP1 * 2.0 * PI / tmmP1.GetWl();
    ArrayXd kxsP2 = betasP2 * 2.0 * PI / tmmP2.GetWl();

    // Allocate space
    auto res = std::make_unique<FieldsZX>(zs.size(), xs.size(), tmmGen.GetPolarization());
    res->SetZero();

    if (tmmP1.GetWave()->IsCoherent() && tmmP2.GetWave()->IsCoherent()) {
#pragma omp parallel
        {
            SecondOrderNLTMM tmmThread = *this;
            tmmThread.tmmP1.SetOverrideE0(true);
            tmmThread.tmmP2.SetOverrideE0(true);
#pragma omp for
            for (int i = 0; i < betasP1.size(); i++) {
                tmmThread.tmmP1.SetBeta(betasP1(i));
                tmmThread.tmmP1.SetE0(E0sP1(i));
                double dkxP1 = GetDifferential(kxsP1, i);
                for (int j = 0; j < betasP2.size(); j++) {
                    tmmThread.tmmP2.SetBeta(betasP2(j));
                    tmmThread.tmmP2.SetE0(E0sP2(j));
                    double dkxP2 = GetDifferential(kxsP2, j);

                    // Solve
                    tmmThread.Solve();

                    // Integrate fields
                    double kxGen = tmmThread.tmmGen.GetLayer(0)->GetKx();
                    auto fGen = tmmThread.tmmGen.GetFields(zs, dir);
                    ArrayXcd phaseXGen = (constI * kxGen * xs).exp() * dkxP1 * dkxP2;
                    res->AddFields(*fGen, phaseXGen);
                }
            }
        }
    } else if (tmmP1.GetWave()->IsCoherent() && !tmmP2.GetWave()->IsCoherent()) {
#pragma omp parallel
        {
            SecondOrderNLTMM tmmThread = *this;
            tmmThread.tmmP1.SetOverrideE0(true);
            tmmThread.tmmP2.SetOverrideE0(true);
// Sum field incoherently
#pragma omp for
            for (int j = 0; j < betasP2.size(); j++) {
                tmmThread.tmmP2.SetBeta(betasP2(j));
                tmmThread.tmmP2.SetE0(E0sP2(j));
                double dkxP2 = GetDifferential(kxsP2, j);

                FieldsZX coherentFields(zs.size(), xs.size(), tmmGen.GetPolarization());
                coherentFields.SetZero();
                // Sum field coherently
                for (int i = 0; i < betasP1.size(); i++) {
                    tmmThread.tmmP1.SetBeta(betasP1(i));
                    tmmThread.tmmP1.SetE0(E0sP1(i));
                    double dkxP1 = GetDifferential(kxsP1, i);

                    // Solve
                    tmmThread.Solve();

                    // Integrate fields
                    double kxGen = tmmThread.tmmGen.GetLayer(0)->GetKx();
                    auto fGen = tmmThread.tmmGen.GetFields(zs, dir);
                    ArrayXcd phaseXGen = (constI * kxGen * xs).exp() * dkxP1 * dkxP2;
                    coherentFields.AddFields(*fGen, phaseXGen);
                }

                res->AddSquaredFields(&coherentFields);
            }
        }
        res->TakeSqrt();
    } else {
        throw std::invalid_argument("Only P2 can be incoherent.");
    }

    return res;
}

pairdd SecondOrderNLTMM::WaveGetPowerFlows(int layerNr, double x0, double x1, double z) {
    // Do checking
    if (layerNr < 0 || layerNr > tmmP1.LayersCount()) {
        throw std::invalid_argument("Invalid layer index.");
    }
    UpdateGenParams();
    CheckPrerequisites();

    // Solve waves
    ArrayXd betasP1, betasP2;
    ArrayXcd E0sP1, E0sP2;
    SolveWaves(&betasP1, &E0sP1, &betasP2, &E0sP2);
    Wave* waveP1 = tmmP1.GetWave();
    Wave* waveP2 = tmmP2.GetWave();
    double LyP1 = tmmP1.GetWave()->GetLy();
    double LyP2 = waveP2->GetLy();

    // Wider P2 is required so that the integration x-range derived from P1
    // is fully covered by P2.  Relaxing this would need separate x-ranges.
    if (waveP1->GetW0() > waveP2->GetW0()) {
        throw std::invalid_argument("Currently P2 must be at least as wide as P1.");
    }

    // Default x-range is derived from P1 beam extent
    auto [xMinP1, xMaxP1] = waveP1->GetXRange();
    if (std::isnan(x0)) {
        x0 = xMinP1;
    }
    if (std::isnan(x1)) {
        x1 = xMaxP1;
    }

    // Ly
    if (LyP1 != LyP2) {
        throw std::invalid_argument("Both pump waves must have the same Ly.");
    }

    // Init memory
    int nTot = betasP1.size() * betasP2.size();
    Eigen::MatrixX2cd UsUnsorted(nTot, 2);
    Eigen::MatrixX2cd kzsUnsorted(nTot, 2);
    ArrayXd kxsUnsorted(nTot);

    // Solve for every beta
    bool oldOverrideE0P1 = tmmP1.GetOverrideE0();
    bool oldOverrideE0P2 = tmmP2.GetOverrideE0();
    double oldBetaP1 = tmmP1.GetBeta();
    double oldBetaP2 = tmmP2.GetBeta();
    tmmP1.SetOverrideE0(true);
    tmmP2.SetOverrideE0(true);
    for (int i = 0; i < betasP1.size(); i++) {
        tmmP1.SetBeta(betasP1(i));
        tmmP1.SetE0(E0sP1(i));
        for (int j = 0; j < betasP2.size(); j++) {
            tmmP2.SetBeta(betasP2(j));
            tmmP2.SetE0(E0sP2(j));
            Solve();

            int index = i * betasP2.size() + j;
            UsUnsorted.row(index) = tmmGen.GetLayer(layerNr)->GetMainFields(0.0);
            kzsUnsorted.row(index) = tmmGen.GetLayer(layerNr)->GetHw()->GetKz();
            kxsUnsorted(index) = tmmGen.GetLayer(layerNr)->GetKx();
        }
    }
    tmmP1.SetOverrideE0(oldOverrideE0P1);
    tmmP2.SetOverrideE0(oldOverrideE0P2);
    tmmP1.SetBeta(oldBetaP1);
    tmmP2.SetBeta(oldBetaP2);


    double Ly = LyP1;
    double wlGen = tmmGen.GetWl();
    Polarization polGen = tmmGen.GetPolarization();
    dcomplex epsLayer = tmmGen.GetLayer(layerNr)->eps;

    pairdd res(0.0, 0.0);
    if (waveP1->IsCoherent() && waveP2->IsCoherent()) {
        // Sort values by kxs
        Eigen::ArrayXi indices(nTot);
        for (int i = 0; i < nTot; i++) {
            indices(i) = i;
        }

        std::sort(indices.data(), indices.data() + nTot,
                  [&kxsUnsorted](const int a, int b) -> bool { return kxsUnsorted(a) < kxsUnsorted(b); });

        Eigen::MatrixX2cd Us(nTot, 2);
        Eigen::MatrixX2cd kzs(nTot, 2);
        ArrayXd kxs(nTot);

        for (int i = 0; i < nTot; i++) {
            Us.row(i) = UsUnsorted.row(indices(i));
            kzs.row(i) = kzsUnsorted.row(indices(i));
            kxs(i) = kxsUnsorted(indices(i));
        }

        // Integrate powers
        res = IntegrateWavePower(layerNr, polGen, wlGen, epsLayer, Us, kxs, kzs, x0, x1, z, Ly);

    } else if (waveP1->IsCoherent() && !waveP2->IsCoherent()) {
        ArrayXd kxsP2 = betasP2 * (2.0 * PI / tmmP2.GetWl());
        for (int j = 0; j < betasP2.size(); j++) {
            // Init
            Eigen::MatrixX2cd Us(betasP1.size(), 2);
            Eigen::MatrixX2cd kzs(betasP1.size(), 2);
            ArrayXd kxs(betasP1.size());

            // Collect amplitutes
            for (int i = 0; i < betasP1.size(); i++) {
                int index = i * betasP2.size() + j;
                Us.row(i) = UsUnsorted.row(index);
                kzs.row(i) = kzsUnsorted.row(index);
                kxs.row(i) = kxsUnsorted.row(index);
            }

            // Sum powers
            double dkxP2 = GetDifferential(kxsP2, j);
            auto [rFwd, rBwd] = IntegrateWavePower(layerNr, polGen, wlGen, epsLayer, Us, kxs, kzs, x0, x1, z, Ly);
            res.first += rFwd * dkxP2;
            res.second += rBwd * dkxP2;
        }
    } else {
        throw std::invalid_argument("Only P2 is allowed to be incoherent.");
    }


    return res;
}

std::unique_ptr<WaveSweepResultSecondOrderNLTMM>
SecondOrderNLTMM::WaveSweep(TMMParam param, const Eigen::Map<ArrayXd>& valuesP1, const Eigen::Map<ArrayXd>& valuesP2,
                            int outmask, int paramLayer, int layerNr, double layerZ) {
    if (valuesP1.size() != valuesP2.size()) {
        throw std::invalid_argument("Value arrays must have the same size.");
    }

    CheckPrerequisites(param);

    // Alloc memory for result
    auto res = std::make_unique<WaveSweepResultSecondOrderNLTMM>(valuesP1.size(), outmask, layerNr, layerZ);

#pragma omp parallel
    {
        // Make a copy
        SecondOrderNLTMM nlTMMCopy = *this;

// Sweep
#pragma omp for
        for (int i = 0; i < valuesP1.size(); i++) {
            // Set sweep param
            nlTMMCopy.GetP1()->SetParam(param, valuesP1(i), paramLayer);
            nlTMMCopy.GetP2()->SetParam(param, valuesP2(i), paramLayer);
            if (GetParamType(param) == TMMParamType::PTYPE_NONLINEAR_LAYER) {
                // NB: Gen param is set from P2 value; a more accurate approach
                // would recompute Gen from both P1 and P2 (e.g. for layer-d sweeps).
                nlTMMCopy.GetGen()->SetParam(param, valuesP2(i), paramLayer);
            }
            nlTMMCopy.UpdateGenParams();
            res->SetValues(i, nlTMMCopy);
        }
    }
    return res;
}

WaveSweepResultSecondOrderNLTMM::WaveSweepResultSecondOrderNLTMM(int n, int outmask_, int layerNr_, double layerZ_)
    : P1(n, outmask_, layerNr_, layerZ_), P2(n, outmask_, layerNr_, layerZ_),
      Gen(n, outmask_ & ~(SWEEP_ENH), layerNr_, layerZ_), wlsGen(n), betasGen(n) {
    outmask = outmask_;
}

void WaveSweepResultSecondOrderNLTMM::SetValues(int nr, SecondOrderNLTMM& tmm) {
    wlsGen(nr) = tmm.GetGen()->GetWl();
    betasGen(nr) = tmm.GetGen()->GetBeta();

    if (outmask & SWEEP_P1) {
        P1.SetValues(nr, *tmm.GetP1());
    }

    if (outmask & SWEEP_P2) {
        P2.SetValues(nr, *tmm.GetP2());
    }

    if (outmask & SWEEP_GEN) {
        // First layer
        if ((outmask & SWEEP_I) || (outmask & SWEEP_R)) {
            auto [fwd0, bwd0] = tmm.WaveGetPowerFlows(0);
            Gen.Pi(nr) = fwd0;
            Gen.Pr(nr) = bwd0;
        }

        // Last layer
        if (outmask & SWEEP_T) {
            auto [fwdL, bwdL] = tmm.WaveGetPowerFlows(tmm.GetGen()->LayersCount() - 1);
            Gen.Pt(nr) = fwdL;
        }

        Gen.beamArea(nr) = tmm.GetP1()->GetWave()->GetBeamArea();
    }
}

} // namespace TMM
