#include "Common.h"

namespace TMM {
TMMParamType GetParamType(TMMParam param) {
    switch (param) {
    case TMMParam::PARAM_WL:
    case TMMParam::PARAM_BETA:
    case TMMParam::PARAM_POL:
    case TMMParam::PARAM_I0:
    case TMMParam::PARAM_OVERRIDE_E0:
    case TMMParam::PARAM_E0:
    case TMMParam::PARAM_MODE:
        return TMMParamType::PTYPE_NONLINEAR_TMM;
    case TMMParam::PARAM_WAVE_W0:
        return TMMParamType::PTYPE_WAVE;
    case TMMParam::PARAM_LAYER_D:
        return TMMParamType::PTYPE_NONLINEAR_LAYER;
    default:
        throw std::invalid_argument("Param has no type");
    }
}
Matrix3d RotationMatrixX(double phi) {
    Matrix3d res;
    res << 1.0, 0.0, 0.0, 0.0, std::cos(phi), -std::sin(phi), 0.0, std::sin(phi), std::cos(phi);
    return res;
}

Matrix3d RotationMatrixY(double phi) {
    Matrix3d res;
    res << std::cos(phi), 0.0, std::sin(phi), 0.0, 1.0, 0.0, -std::sin(phi), 0.0, std::cos(phi);
    return res;
}

Matrix3d RotationMatrixZ(double phi) {
    Matrix3d res;
    res << std::cos(phi), -std::sin(phi), 0.0, std::sin(phi), std::cos(phi), 0.0, 0.0, 0.0, 1.0;
    return res;
}

Tensor3d ApplyRotationMatrixToTensor(const Tensor3d& input, const Matrix3d& R) {
    Tensor3d res(3, 3, 3);
    res.setZero();

    for (int i1 = 0; i1 < 3; i1++) {
        for (int i2 = 0; i2 < 3; i2++) {
            for (int i3 = 0; i3 < 3; i3++) {
                for (int j1 = 0; j1 < 3; j1++) {
                    for (int j2 = 0; j2 < 3; j2++) {
                        for (int j3 = 0; j3 < 3; j3++) {
                            res(i1, i2, i3) += R(i1, j1) * R(i2, j2) * R(i3, j3) * input(j1, j2, j3);
                        }
                    }
                }
            }
        }
    }
    return res;
}

Tensor3d RotateTensor(Tensor3d& input, double phiX, double phiY, double phiZ) {
    Tensor3d output(3, 3, 3);
    output = input;

    Matrix3d Rx = RotationMatrixX(phiX);
    Matrix3d Ry = RotationMatrixY(phiY);
    Matrix3d Rz = RotationMatrixZ(phiZ);

    output = ApplyRotationMatrixToTensor(output, Rx);
    output = ApplyRotationMatrixToTensor(output, Ry);
    output = ApplyRotationMatrixToTensor(output, Rz);
    return output;
}

template <typename T> T Interpolate(double x, const ArrayXd& xs, const Eigen::Array<T, Eigen::Dynamic, 1>& ys) {
    // xs must be sorted

    // Check range
    if (x < xs(0) || x >= xs(xs.size() - 1)) {
        throw std::runtime_error("Interpolation out of range");
    }

    if (xs(0) >= xs(xs.size() - 1)) {
        throw std::runtime_error("Interpolation: xs must be sorted");
    }

    // Binary search (last element, that is less or equal than x)
    int b = 0, e = xs.size() - 1;
    while (b < e) {
        int a = (b + e) / 2;
        if (xs(b) >= xs(e)) {
            throw std::runtime_error("Interpolation: xs must be sorted");
        }

        if (xs(a) > x) {
            // [b..a[
            e = a - 1;

            if (xs(e) <= x) {
                b = e;
            }
        } else {
            // [a..e]
            b = a;
            if (xs(a + 1) > x) {
                e = a;
            }
        }
    }
    // Linear interpolation in range x[b]..x[b+1]
    double dx = xs(b + 1) - xs(b);
    T dy = ys(b + 1) - ys(b);
    T res = ys(b) + dy / dx * (x - xs(b));
    return res;
}

template dcomplex Interpolate<dcomplex>(double, const ArrayXd&, const ArrayXcd&);
template double Interpolate<double>(double, const ArrayXd&, const ArrayXd&);

double GetDifferential(const ArrayXd& intVar, int nr) {
    double dIntVar;
    if (intVar.size() <= 1) {
        dIntVar = 1.0;
    } else if (nr == 0) {
        dIntVar = intVar(1) - intVar(0);
    } else if (nr + 1 == intVar.size()) {
        dIntVar = intVar(intVar.size() - 1) - intVar(intVar.size() - 2);
    } else {
        dIntVar = 0.5 * (intVar(nr + 1) - intVar(nr - 1));
    }
    return dIntVar;
}

/*
inline dcomplex CalcFx(dcomplex dk, double x0, double x1) {
    dcomplex exp1 = FastExp(constI * dk * x1);
    dcomplex exp2 = FastExp(constI * dk * x0);
    dcomplex res  = -constI * (exp1 - exp2) / dk;
    return res;
}
*/

inline double CalcFxSym(double dk, double x0) {
    double res = -2.0 * std::sin(x0 * dk) / dk;
    return res;
}

/*
double const IntegrateWavePower(int layerNr, Polarization pol, double wl, dcomplex epsLayer0, const ArrayXcd & Us, const
ArrayXd & kxs, const ArrayXcd & kzs, double x0, double x1, double z, double Ly) {

    if (x0 != -x1) {
        std::cerr << "Currently only x0 = -x1 supported." << std::endl;
        throw std::invalid_argument("Currently only x0 = -x1 supported.");
    }

    int m = kxs.size();
    bool needFz = bool(z != 0.0);

    //Precalc Fx
    MatrixXd Fxprecalc(m, m);
    double *FxPrecalcPtr = Fxprecalc.data();

    for (int i = 0; i < m; i++) {
        Fxprecalc(i, i) = x1 - x0;
    }

    for (int i = 0; i < m; i++) {
        double kx = kxs(i);
        for (int j = i + 1; j < m; j++) {
            double kxP = kxs(j);
            double dk = kx - kxP;
            double Fx = CalcFxSym(dk, x0);
            FxPrecalcPtr[i * m + j] = Fx;
            FxPrecalcPtr[j * m + i] = Fx;
            //Fxprecalc(j, i) = Fx;
            //Fxprecalc(i, j) = Fx;
        }
    }

    // Precalc differentials
    ArrayXd dkxs(m);
    for (int i = 0; i < m; i++) {
        dkxs(i) = GetDifferential(kxs, i);
    }

    double integValue2dRe = 0.0;
    double integValue2dIm = 0.0;
    //#pragma omp parallel for default(shared) reduction(+:integValue2dRe,integValue2dIm)
    for (int i = 0; i < m; i++) {
        double kx = kxs(i);
        dcomplex kz = kzs(i);
        dcomplex U = Us(i);
        dcomplex integValue1d = 0.0;
        for (int j = 0; j < m; j++) {
            double kxP = kxs(j);
            dcomplex kzP = kzs(j);
            dcomplex UPc = std::conj(Us(j));
            double dkxP = dkxs(j);

            double Fx = FxPrecalcPtr[i * m + j];// Fxprecalc(j, i);
            dcomplex coef = dkxP * UPc * Fx;
            if (needFz) {
                coef *= FastExp(constI * (kz - kzP) * z);
            }

            // Integrate

            switch (pol)
            {
            case TMM::P_POL:
                integValue1d += coef * kz;
                break;
            case TMM::S_POL:
                integValue1d += coef * kzP;
                break;
            default:
                throw std::invalid_argument("Invalid polarization");
                break;
            }
        }
        double dkx = dkxs(i);
        dcomplex tmpRes = U * integValue1d * dkx;
        integValue2dRe += real(tmpRes);
        integValue2dIm += imag(tmpRes);
    }

    // Polarization specific multipliers
    dcomplex integValue2d(integValue2dRe, integValue2dIm);
    double omega = WlToOmega(wl);
    double res;

    switch (pol)
    {
    case TMM::P_POL:
        res = Ly / (2.0 * omega * constEps0) * std::real(integValue2d / epsLayer0);
        break;
    case TMM::S_POL:
        res = Ly / (2.0 * omega * constMu0) * real(integValue2d);
        break;
    default:
        throw std::invalid_argument("Invalid polarization");
        break;
    }

    return res;
}
*/

pairdd IntegrateWavePower([[maybe_unused]] int layerNr, Polarization pol, double wl, dcomplex epsLayer,
                          const Eigen::MatrixX2cd& Us, const ArrayXd& kxs, const Eigen::MatrixX2cd& kzs, double x0,
                          double x1, double z, double Ly) {
    if (x0 != -x1 || z != 0) {
        throw std::invalid_argument("Currently only x0 = -x1, z = 0 supported.");
    }

    constexpr int iF = static_cast<int>(WaveDirection::F);
    constexpr int iB = static_cast<int>(WaveDirection::B);

    int m = kxs.size();
    // Precalc differentials
    ArrayXd dkxs(m);
    for (int i = 0; i < m; i++) {
        dkxs(i) = GetDifferential(kxs, i);
    }

    // Not polarization specific
    dcomplex integValue2dF = 0.0;
    dcomplex integValue2dB = 0.0;
    for (int i = 0; i < m; i++) {
        dcomplex UF = Us(i, iF);
        dcomplex kzF = kzs(i, iF);
        dcomplex UB = Us(i, iB);
        dcomplex kzB = kzs(i, iB);
        // double kx = kxs(i);
        double dkx = dkxs(i);
        integValue2dF += dkx * dkx * UF * std::conj(UF) * kzF * (x1 - x0);
        integValue2dB += dkx * dkx * UB * std::conj(UB) * kzB * (x1 - x0);
    }

    // Polarization spetcific
    double omega = WlToOmega(wl);
    double resF, resB;
    if (pol == Polarization::P_POL) {
        for (int i = 0; i < m; i++) {
            double kx = kxs(i);
            dcomplex kzF = kzs(i, iF);
            dcomplex UF = Us(i, iF);
            dcomplex kzB = kzs(i, iB);
            dcomplex UB = Us(i, iB);
            dcomplex integValue1dF = 0.0, integValue1dB = 0.0;
            for (int j = i + 1; j < m; j++) {
                double kxP = kxs(j);
                dcomplex kzPF = kzs(j, iF);
                dcomplex UPF = Us(j, iF);
                dcomplex kzPB = kzs(j, iB);
                dcomplex UPB = Us(j, iB);
                double dkxP = dkxs(j);
                double dk = kx - kxP;

                double cc = (std::sin(x0 * dk) * dkxP / dk);
                integValue1dF += cc * (UF * std::conj(UPF) * kzF + std::conj(UF) * UPF * kzPF);
                integValue1dB += cc * (UB * std::conj(UPB) * kzB + std::conj(UB) * UPB * kzPB);
            }
            double dkx = dkxs(i);
            integValue2dF -= (2.0 * dkx) * integValue1dF;
            integValue2dB -= (2.0 * dkx) * integValue1dB;
        }
        resF = Ly / (2.0 * omega * constEps0) * std::real(integValue2dF / epsLayer);
        resB = -Ly / (2.0 * omega * constEps0) * std::real(integValue2dB / epsLayer);
    } else if (pol == Polarization::S_POL) {
        // Almost identical to p-polarization, copy for performance
        for (int i = 0; i < m; i++) {
            double kx = kxs(i);
            dcomplex kzF = kzs(i, iF);
            dcomplex UF = Us(i, iF);
            dcomplex kzB = kzs(i, iB);
            dcomplex UB = Us(i, iB);
            dcomplex integValue1dF = 0.0, integValue1dB = 0.0;
            for (int j = i + 1; j < m; j++) {
                double kxP = kxs(j);
                dcomplex kzPF = kzs(j, iF);
                dcomplex UPF = Us(j, iF);
                dcomplex kzPB = kzs(j, iB);
                dcomplex UPB = Us(j, iB);
                double dkxP = dkxs(j);
                double dk = kx - kxP;

                double cc = (std::sin(x0 * dk) * dkxP / dk);
                integValue1dF += cc * (UF * std::conj(UPF) * kzPF + std::conj(UF) * UPF * kzF);
                integValue1dB += cc * (UB * std::conj(UPB) * kzPB + std::conj(UB) * UPB * kzB);
            }
            double dkx = dkxs(i);
            integValue2dF -= (2.0 * dkx) * integValue1dF;
            integValue2dB -= (2.0 * dkx) * integValue1dB;
        }
        resF = Ly / (2.0 * omega * constMu0) * real(integValue2dF);
        resB = -Ly / (2.0 * omega * constMu0) * real(integValue2dB);
    } else {
        throw std::invalid_argument("Unknown polarization.");
    }

    return pairdd(resF, resB);
}

WaveDirection GetWaveDirection(dcomplex kzF, dcomplex eps, Polarization pol) {
    double pwrFlow;
    switch (pol) {
    case Polarization::P_POL:
        pwrFlow = real(kzF / eps);
        break;
    case Polarization::S_POL:
        pwrFlow = real(kzF);
        break;
    default:
#pragma omp critical
        throw std::invalid_argument("Unknown polarization.");
    }

    WaveDirection res;

    if (real(kzF) > 0.0) {
        res = WaveDirection::F;
    } else if (real(kzF) < 0.0) {
        res = WaveDirection::B;
    } else {
        if (imag(kzF) > 0.0) {
            res = WaveDirection::F;
        } else if (imag(kzF) < 0.0) {
            res = WaveDirection::B;
        } else {
// Only if kz = 0.0
#pragma omp critical
            throw std::runtime_error("Could not determine wave direction.");
        }
    }

    if (res == WaveDirection::F && (/*imag(kzF) < 0.0 || */ pwrFlow < 0.0)) {
#pragma omp critical
        std::cerr << "F: " << kzF << " " << eps << ": " << pwrFlow << std::endl;
        // throw std::runtime_error("Could not determine wave direction.");
    }

    if (res == WaveDirection::B && (/*imag(kzF) > 0.0 || */ pwrFlow > 0.0)) {
#pragma omp critical
        std::cerr << "B: " << kzF << " " << eps << ": " << pwrFlow << std::endl;
        // throw std::runtime_error("Could not determine wave direction.");
    }
    return res;
}

ArrayXcd FFTShift(ArrayXcd data) {
    int p2 = (data.size() + 1) / 2;

    ArrayXcd res(data.size());
    for (int i = p2; i < data.size(); i++) {
        res(i - p2) = data(i);
    }
    for (int i = 0; i < p2; i++) {
        res(data.size() - p2 + i) = data(i);
    }
    return res;
}

ArrayXd IFFTShift(ArrayXd data) {
    int p2 = data.size() - (data.size() + 1) / 2;
    ArrayXd res(data.size());

    for (int i = p2; i < data.size(); i++) {
        res(i - p2) = data(i);
    }
    for (int i = 0; i < p2; i++) {
        res(data.size() - p2 + i) = data(i);
    }
    return res;
}

ArrayXd FFTFreq(int n, double dx) {
    ArrayXd res(n);
    int startI = -(n - (n % 2)) / 2;
    double step = 1.0 / (dx * n);
    for (int i = 0; i < n; i++) {
        res(i) = (startI + i) * step;
    }
    return res;
}

}; // namespace TMM
