#include "Material.h"
#include <array>

namespace TMM {

Chi2Tensor::Chi2Tensor() : chi2(3, 3, 3), chi2Rotated(3, 3, 3) {
    Clear();
    SetDistinctFields(true);
    SetRotation(0.0, 0.0, 0.0);
}

void Chi2Tensor::SetDistinctFields(bool distinctFields_) {
    distinctFields = distinctFields_;
}

void Chi2Tensor::SetRotation(double phiX_, double phiY_, double phiZ_) {
    phiX = phiX_;
    phiY = phiY_;
    phiZ = phiZ_;
    needRotationRecalc = true;
}

void Chi2Tensor::SetChi2(int i1, int i2, int i3, double value) {
    // indices are from 1..3
    isNonlinear = true;

    if (std::max<int>(i1, std::max<int>(i2, i3)) > 3 || std::min<int>(i1, std::min<int>(i2, i3)) < 1) {
        throw std::runtime_error("chi2 index not in range 1..3");
    }

    needRotationRecalc = true;
    chi2(i1 - 1, i2 - 1, i3 - 1) = value;
}

void Chi2Tensor::SetD(int i1, int i2, double value) {
    // indices are from i1=1..3, i2=1..6
    isNonlinear = true;

    if (std::min<int>(i1, i2) < 1 || i1 > 3 || i2 > 6) {
        throw std::runtime_error("d index not in range 1..3, 1..6");
    }
    needRotationRecalc = true;
    constexpr std::array<int, 6> dIndex2 = {0, 1, 2, 1, 0, 0};
    constexpr std::array<int, 6> dIndex3 = {0, 1, 2, 2, 2, 1};
    chi2(i1 - 1, dIndex2[i2 - 1], dIndex3[i2 - 1]) = 2.0 * value;
    chi2(i1 - 1, dIndex3[i2 - 1], dIndex2[i2 - 1]) = 2.0 * value;
}

void Chi2Tensor::Clear() {
    chi2.setZero();
    chi2Rotated.setZero();
    isNonlinear = false;
    needRotationRecalc = false;
}

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

const Tensor3d& Chi2Tensor::GetChi2Tensor() {
#pragma omp critical
    if (needRotationRecalc) {
        chi2Rotated = RotateTensor(chi2, phiX, phiY, phiZ);
        needRotationRecalc = false;
    }
    return chi2Rotated;
}

double Chi2Tensor::GetChi2Element(int i1, int i2, int i3) {
    // Indices form 1..3
    if (std::max<int>(i1, std::max<int>(i2, i3)) > 3 || std::min<int>(i1, std::min<int>(i2, i3)) < 1) {
        throw std::invalid_argument("Indices must be in range 1..3");
    }

    double res = GetChi2Tensor()(i1 - 1, i2 - 1, i3 - 1);
    return res;
}

Vector3cd Chi2Tensor::GetNonlinearPolarization(const Vector3cd& E1, const Vector3cd& E2) {
    Vector3cd res = Vector3cd::Zero();
    Tensor3d chi2 = GetChi2Tensor();

    for (int i = 0; i < 3; i++) {
        res(i) = E1(0) * (E2(0) * chi2(i, 0, 0) + E2(1) * chi2(i, 0, 1) + E2(2) * chi2(i, 0, 2)) +
                 E1(1) * (E2(0) * chi2(i, 1, 0) + E2(1) * chi2(i, 1, 1) + E2(2) * chi2(i, 1, 2)) +
                 E1(2) * (E2(0) * chi2(i, 2, 0) + E2(1) * chi2(i, 2, 1) + E2(2) * chi2(i, 2, 2));
    }
    res *= constEps0;
    if (distinctFields) {
        res *= 2.0;
    }
    return res;
}

Material::Material(dcomplex n_) {
    isStatic = true;
    staticN = n_;
}

Material::Material(ArrayXd wlsExp_, ArrayXcd nsExp_) : wlsExp(wlsExp_), nsExp(nsExp_) {
    // Copy of wlsExp and nsExp is made intentionally
    isStatic = false;
    staticN = 1.0;

    if (wlsExp.size() != nsExp.size()) {
        throw std::invalid_argument("wls and ns must have the same length");
    }

    if (wlsExp.size() < 2) {
        throw std::invalid_argument("The length of wls and ns must be at least 2");
    }
}

dcomplex Material::GetN(double wl) const {
    if (isStatic) {
        return staticN;
    }

    // Interpolate
    dcomplex res = Interpolate(wl, wlsExp, nsExp);
    return res;
}

bool Material::IsNonlinear() const noexcept {
    return chi2.IsNonlinear();
}

} // namespace TMM
