#include "NonlinearLayer.h"

namespace TMM {

Fields::Fields() {
    E = Vector3cd::Zero();
    H = Vector3cd::Zero();
    dir = WaveDirection::TOT;
}

void Fields::SetFieldsPpol(const Array2cd& Hy, const Array2cd& Ex, const Array2cd& Ez, WaveDirection dir_) {
    constexpr int iF = static_cast<int>(WaveDirection::F);
    constexpr int iB = static_cast<int>(WaveDirection::B);
    E = Vector3cd::Zero();
    H = Vector3cd::Zero();
    dir = dir_;

    switch (dir) {
    case WaveDirection::F:
        E(0) = Ex(iF);
        E(2) = Ez(iF);
        H(1) = Hy(iF);
        break;
    case WaveDirection::B:
        E(0) = Ex(iB);
        E(2) = Ez(iB);
        H(1) = Hy(iB);
        break;
    case WaveDirection::TOT:
        E(0) = Ex(iF) + Ex(iB);
        E(2) = Ez(iF) + Ez(iB);
        H(1) = Hy(iF) + Hy(iB);
        break;
    default:
        throw std::invalid_argument("Unknown wave direction.");
    }
}
void Fields::SetFieldsSpol(const Array2cd& Ey, const Array2cd& Hx, const Array2cd& Hz, WaveDirection dir_) {
    constexpr int iF = static_cast<int>(WaveDirection::F);
    constexpr int iB = static_cast<int>(WaveDirection::B);
    E = Vector3cd::Zero();
    H = Vector3cd::Zero();
    dir = dir_;

    switch (dir) {
    case WaveDirection::F:
        E(1) = Ey(iF);
        H(0) = Hx(iF);
        H(2) = Hz(iF);
        break;
    case WaveDirection::B:
        E(1) = Ey(iB);
        H(0) = Hx(iB);
        H(2) = Hz(iB);
        break;
    case WaveDirection::TOT:
        E(1) = Ey(iF) + Ey(iB);
        H(0) = Hx(iF) + Hx(iB);
        H(2) = Hz(iF) + Hz(iB);
        break;
    default:
        throw std::invalid_argument("Unknown wave direction.");
    }
}

InhomogeneosWaveParams::InhomogeneosWaveParams() {
    kSzF = 0.0;
    pF = Vector3cd::Zero();
    pB = Vector3cd::Zero();
}

HomogeneousWave::HomogeneousWave(NonlinearLayer* layer_) {
    layer = layer_;
    kz = Array2cd::Zero();
    phase = Array2cd::Zero();
    propMatrix = Matrix2cd::Zero();
    solved = false;
}

void HomogeneousWave::Solve(NonlinearLayer* layer_) {
    layer = layer_;
    constexpr int iF = static_cast<int>(WaveDirection::F);
    constexpr int iB = static_cast<int>(WaveDirection::B);
    kz(iF) = std::sqrt(dcomplex(layer->kNorm - sqr(layer->kx))); // Slowest part (AVX?)
    kz(iB) = -kz(iF);

    phase = Array2cd::Ones();
    if (!std::isinf(layer->d)) {
        phase = (constI * kz * layer->d).exp();
    }

    if (imag(kz(iF)) < 0.0) {
        // Should not happen, checking
        throw std::runtime_error("kzF imaginary part negative");
    }

    if (real(kz(iF)) < 0.0) {
        // Should not happen, checking
        throw std::runtime_error("kzF real part negative");
    }

    // Propagation matrix
    if (std::isinf(layer->d)) {
        propMatrix = Matrix2cd::Identity();
    } else {
        propMatrix << phase(iF), 0.0, 0.0, phase(iB);
    }
    solved = true;
}

Array2cd HomogeneousWave::GetKz() const {
    if (!solved) {
        throw std::runtime_error("Homogeneous wave must be solved first");
    }
    return kz;
}

dcomplex HomogeneousWave::GetKzF() const {
    if (!solved) {
        throw std::runtime_error("Homogeneous wave must be solved first");
    }
    return kz(static_cast<int>(WaveDirection::F));
}

double HomogeneousWave::GetKx() const {
    if (!solved) {
        throw std::runtime_error("Homogeneous wave must be solved first");
    }
    return layer->kx;
}

Array2cd HomogeneousWave::GetMainFields(double z) const {
    if (!solved) {
        throw std::runtime_error("Homogeneous wave must be solved first");
    }
    Array2cd U = layer->U0;
    if (z == 0.0) {
        // No phase correction
    } else if (z == layer->d) {
        // Phase at the end of the layer is precalculated
        U *= phase;
    } else {
        U *= (constI * kz * z).exp();
    }
    return U;
}

InhomogeneousWave::InhomogeneousWave(NonlinearLayer* layer_) {
    layer = layer_;
    kSz = Array2cd::Zero();
    kSNorm = 0.0;
    phaseS = Array2cd::Zero();
    px = Array2cd::Zero();
    py = Array2cd::Zero();
    pz = Array2cd::Zero();
    By = Array2cd::Zero();
    propMatrixNL = Array2cd::Zero();
    solved = false;
}

void InhomogeneousWave::Solve(NonlinearLayer* layer_, InhomogeneosWaveParams& params) {
    layer = layer_;
    constexpr int iF = static_cast<int>(WaveDirection::F);
    constexpr int iB = static_cast<int>(WaveDirection::B);

    // kS
    kSz(iF) = params.kSzF;
    kSz(iB) = -kSz(iF);
    kSNorm = sqr(layer->kx) + sqr(kSz(iF));

    // Phase
    phaseS = Array2cd::Ones();
    if (!std::isinf(layer->d) && layer->IsNonlinear()) {
        phaseS = (constI * kSz * layer->d).exp();
    }

    // px, py, pz
    px << params.pF(0), params.pB(0);
    py << params.pF(1), params.pB(1);
    pz << params.pF(2), params.pB(2);

    // Calc By
    if (!layer->IsNonlinear()) {
        By << 0.0, 0.0;
    } else {
        if (layer->pol == Polarization::S_POL) {
            By = -(py / constEps0) * sqr(layer->k0) / (layer->kNorm - kSNorm);
            By(iB) *= phaseS(iF);
        } else if (layer->pol == Polarization::P_POL) {
            By = -layer->omega / (layer->kNorm - kSNorm) * (kSz * px - layer->kx * pz);
            By(iB) *= phaseS(iF);
        } else {
            throw std::runtime_error("Unknown polarization");
        }
    }

    // NL Propagation Matrix
    propMatrixNL = Array2cd::Zero();
    if (layer->IsNonlinear()) {
        if (std::isinf(layer->d)) {
            throw std::runtime_error("First and last medium must be nonlinear.");
        } else {
            propMatrixNL = By * (phaseS - layer->hw.phase);
        }
    }
    solved = true;
}

Array2cd InhomogeneousWave::GetMainFields(double z) const {
    if (!solved) {
        throw std::runtime_error("InhomogeneousWave wave must be solved first");
    }

    if (!layer->IsNonlinear()) {
        return Array2cd::Zero();
    }

    Array2cd U = By;
    if (z == 0.0) {
        // Phase is one
    } else if (z == layer->d) {
        // Phase at the end of the layer is precalculated
        U *= (phaseS - layer->hw.phase);
    } else {
        U *= ((constI * kSz * z).exp() - (constI * layer->hw.kz * z).exp());
    }
    return U;
}

NonlinearLayer::NonlinearLayer(double d_, Material* material_) {
    d = d_;
    material = material_;
    wl = 0.0;
    beta = 0.0;
    pol = Polarization::NOT_DEFINED_POL;
    omega = 0.0;
    n = 0.0;
    eps = 0.0;
    k0 = 0.0;
    k = 0.0;
    kNorm = 0.0;
    kx = 0.0;
    propMatrix = Matrix2cd::Zero();
    propMatrixNL = Array2cd::Zero();
    isNonlinear = false;
    U0 = Array2cd::Zero();
    solved = false;
}

void NonlinearLayer::Solve(double wl_, double beta_, Polarization pol_) {
    wl = wl_;
    beta = beta_;
    pol = pol_;
    omega = WlToOmega(wl);

    // Find wavevectors
    n = material->GetN(wl);
    eps = sqr(n);
    k0 = 2.0 * PI / wl;
    k = k0 * n;
    kNorm = sqr(k);
    kx = beta * k0;

    // Solve waves
    hw.Solve(this);
    iws.Solve(this, kpS);
    iwa.Solve(this, kpA);

    // Propagation matrix
    propMatrix = hw.propMatrix;
    propMatrixNL = iwa.propMatrixNL + iws.propMatrixNL;
    solved = true;
}

bool NonlinearLayer::IsNonlinear() const noexcept {
    return isNonlinear;
}

void NonlinearLayer::SetNonlinearity(InhomogeneosWaveParams kpS_, InhomogeneosWaveParams kpA_) {
    isNonlinear = true;
    kpS = kpS_;
    kpA = kpA_;
}

void NonlinearLayer::ClearNonlinearity() {
    isNonlinear = false;
    kpS = InhomogeneosWaveParams();
    kpA = InhomogeneosWaveParams();
}

void NonlinearLayer::SetThickness(double d_) {
    d = d_;
}

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

double NonlinearLayer::GetThickness() const noexcept {
    return d;
}

Material* NonlinearLayer::GetMaterial() const noexcept {
    // It is the responsibility of the user to ensure, that material ptr stays valid
    return material;
}

HomogeneousWave* NonlinearLayer::GetHw() noexcept {
    return &hw;
}

double NonlinearLayer::GetKx() const {
    if (!solved) {
        throw std::runtime_error("NonlinearLayer must be solved first.");
    }
    return kx;
}

double NonlinearLayer::GetK0() const {
    if (!solved) {
        throw std::runtime_error("NonlinearLayer must be solved first.");
    }
    return k0;
}

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

Array2cd NonlinearLayer::GetMainFields(double z) const {
    if (!solved) {
        throw std::runtime_error("NonlinearLayer must be solved first.");
    }

    Array2cd U = hw.GetMainFields(z) + iws.GetMainFields(z) + iwa.GetMainFields(z);
    return U;
}

Fields NonlinearLayer::GetFields(double z, WaveDirection dir) const {
    if (!solved) {
        throw std::runtime_error("NonlinearLayer must be solved first.");
    }

    Fields res;

    Array2cd expKSz = Array2cd::Ones();
    Array2cd expKAz = Array2cd::Ones();

    if (z == 0.0) {
        // Phase is 1
    } else if (z == d) {
        // End of layer phases are precalculated
        expKSz = iws.phaseS;
        expKAz = iwa.phaseS;
    } else {
        expKSz = (constI * iws.kSz * z).exp();
        expKAz = (constI * iwa.kSz * z).exp();
    }

    if (pol == Polarization::P_POL) {
        dcomplex c = -1.0 / (omega * eps * constEps0);
        Array2cd Hy = GetMainFields(z);
        Array2cd Ex = (-hw.kz * Hy) * c;
        Array2cd Ez = (kx * Hy) * c;
        if (IsNonlinear()) {
            Ex += (iws.By * (hw.kz - iws.kSz) * expKSz + iws.px * expKSz * omega + iwa.By * (hw.kz - iwa.kSz) * expKAz +
                   iwa.px * expKAz * omega) *
                  c;
            Ez += (iws.pz * expKSz + iwa.pz * expKAz) * c * omega;
        }
        res.SetFieldsPpol(Hy, Ex, Ez, dir);
    } else if (pol == Polarization::S_POL) {
        double c = 1.0 / (omega * constMu0);
        Array2cd Ey = GetMainFields(z);
        Array2cd Hx = (-hw.kz * Ey) * c;
        Array2cd Hz = (kx * Ey) * c;
        if (IsNonlinear()) {
            Hx += (iws.By * (hw.kz - iws.kSz) * expKSz + iwa.By * (hw.kz - iwa.kSz) * expKAz) * c;
        }
        res.SetFieldsSpol(Ey, Hx, Hz, dir);
    } else {
        throw std::runtime_error("Unknown polarization.");
    }
    return res;
}

double NonlinearLayer::GetIntensity(double z) const {
    if (!solved) {
        throw std::runtime_error("NonlinearLayer must be solved first.");
    }
    Fields f = GetFields(z, WaveDirection::TOT);
    double Sz = 0.5 * real(f.E(0) * std::conj(f.H(1)) - f.E(1) * std::conj(f.H(0)));
    return Sz;
}

double NonlinearLayer::GetAbsorbedIntensity() const {
    if (!solved) {
        throw std::runtime_error("NonlinearLayer must be solved first.");
    }

    if (std::abs(imag(eps)) < 1e-14) {
        // Layer not absorbing
        return 0.0;
    }

    if (IsNonlinear()) {
        throw std::runtime_error("Absorbed power calculation is not allowed in nonlinear media.");
    }

    constexpr int iF = static_cast<int>(WaveDirection::F);
    dcomplex kzF = hw.kz(iF);
    dcomplex dk1 = kzF - std::conj(kzF);
    dcomplex dk2 = kzF + std::conj(kzF);
    dcomplex intValue1 = -constI * (std::exp(constI * dk1 * d) - 1.0) / dk1;
    dcomplex intValue2 = -constI * (std::exp(constI * dk2 * d) - 1.0) / dk2;
    dcomplex intValue3 = constI * (std::exp(-constI * dk2 * d) - 1.0) / dk2;
    dcomplex intValue4 = constI * (std::exp(-constI * dk1 * d) - 1.0) / dk1;

    Vector3cd E0F = GetFields(0.0, WaveDirection::F).E;
    Vector3cd E0B = GetFields(0.0, WaveDirection::B).E;
    Vector3cd E0FC = E0F.conjugate();
    Vector3cd E0BC = E0B.conjugate();

    // E0F.dot(E0FC) provided wrong results
    dcomplex dot1 = E0F(0) * E0FC(0) + E0F(1) * E0FC(1) + E0F(2) * E0FC(2);
    dcomplex dot2 = E0F(0) * E0BC(0) + E0F(1) * E0BC(1) + E0F(2) * E0BC(2);
    dcomplex dot3 = E0B(0) * E0FC(0) + E0B(1) * E0FC(1) + E0B(2) * E0FC(2);
    dcomplex dot4 = E0B(0) * E0BC(0) + E0B(1) * E0BC(1) + E0B(2) * E0BC(2);
    dcomplex intValue = dot1 * intValue1 + dot2 * intValue2 + dot3 * intValue3 + dot4 * intValue4;
    double res = 0.5 * constEps0 * imag(eps) * omega * real(intValue);

    return res;
}

double NonlinearLayer::GetSrcIntensity() const {
    if (!solved) {
        throw std::runtime_error("NonlinearTMM must be solved first.");
    }

    double absorbed = GetAbsorbedIntensity();
    double deltaS = 0.0;
    if (!std::isinf(d)) {
        deltaS = GetIntensity(d) - GetIntensity(0.0);
    }

    double res = absorbed + deltaS;
    return res;
}

double NonlinearLayer::GetENorm(double z, WaveDirection dir) const {
    Fields f = GetFields(z, dir);
    double res = std::sqrt(norm(f.E(0)) + norm(f.E(1)) + norm(f.E(2))); // f.E.norm();
    return res;
}

}; // namespace TMM
