from __future__ import annotations

import math
from typing import Any

import numpy as np
from numpy.typing import NDArray

from LabPy import Constants, Core
from LabPy.Constants import OmegaToWl, WlToOmega

__all__ = ["Chi2Tensor", "SecondOrderNLTmm"]

X = 0
Y = 1
Z = 2

# ===============================================================================
# Chi2Tensor
# ===============================================================================


class Chi2Tensor:
    def __init__(self, dtype: type = complex, distinctFields: bool = True, **kwargs: Any) -> None:
        self._dIndices = [(0, 0), (1, 1), (2, 2), (1, 2), (0, 2), (0, 1)]
        self._chi2 = np.zeros((3, 3, 3), dtype=dtype)
        self._chi2Rotated = None
        self._chi2RotatedChanged = True

        self.distinctFields = distinctFields
        self.phiX = 0.0
        self.phiY = 0.0
        self.phiZ = 0.0

        self.Update(**kwargs)

    def Update(self, **kwargs: Any) -> None:
        self._chi2RotatedChanged = True
        self.phiX = float(kwargs.pop("phiX", self.phiX))
        self.phiY = float(kwargs.pop("phiY", self.phiY))
        self.phiZ = float(kwargs.pop("phiZ", self.phiZ))

        # d-values
        for k in list(kwargs):
            if (not k.startswith("d")) or (len(k) != 3):
                continue
            i1, (i2, i3) = int(k[1]) - 1, self._dIndices[int(k[2]) - 1]
            self._chi2[i1, i2, i3] = 2.0 * kwargs.pop(k)
            self._chi2[i1, i3, i2] = self._chi2[i1, i2, i3]

        # chi2-values
        for k in list(kwargs):
            if (not k.startswith("chi")) or (len(k) != 6):
                continue
            i1, i2, i3 = int(k[3]) - 1, int(k[4]) - 1, int(k[5]) - 1
            self._chi2[i1, i2, i3] = kwargs.pop(k)

        # Errors
        for k in list(kwargs):
            raise ValueError(f"Unknown kwarg {k}")

    def GetChi2(self) -> NDArray:
        if not self._chi2RotatedChanged:
            return self._chi2Rotated

        # Compute chi2Rotated
        self._chi2Rotated = Chi2Tensor._RotateTensorO3(self._chi2, self.phiX, self.phiY, self.phiZ)
        self._chi2RotatedChanged = False
        return self._chi2Rotated

    def GetD(self) -> NDArray:
        res = np.zeros((3, 6), dtype=self._chi2.dtype)
        chi2 = self.GetChi2()

        for i in range(3):
            for j in range(6):
                a, b = self._dIndices[j]
                if abs(chi2[i, a, b] - chi2[i, b, a]) > 1e-20:
                    raise ValueError(f"Not equal elements {(i, a, b)}, {(i, b, a)}")

                res[i, j] = chi2[i, a, b]
        return res

    def GetNonlinearPolarization(self, E1: NDArray, E2: NDArray) -> NDArray:
        chi2 = self.GetChi2()
        fieldTensor = np.outer(E1, E2)
        res = Constants.eps0 * np.tensordot(chi2, fieldTensor).flatten()

        if self.distinctFields:
            res *= 2.0
        return res

    def GetChi2Eff(self, E1: NDArray, E2: NDArray, E3: NDArray) -> complex:
        E1Norm = (E1 / np.linalg.norm(E1)).flatten()
        E2Norm = (E2 / np.linalg.norm(E2)).flatten()
        E3Norm = (E3 / np.linalg.norm(E3)).flatten()
        p = self.GetNonlinearPolarization(E1Norm, E2Norm)

        res = np.dot(E3Norm, p / Constants.eps0)
        if self.distinctFields:
            res *= 2.0
        return res

    def GetDEff(self, E1: NDArray, E2: NDArray, E3: NDArray) -> complex:
        chi2Eff = self.GetChi2Eff(E1, E2, E3)
        res = chi2Eff / 2.0
        return res

    @staticmethod
    def _RotateTensorO3(
        tensor: NDArray,
        phiX: float = 0.0,
        phiY: float = 0.0,
        phiZ: float = 0.0,
    ) -> NDArray:

        def ApplyRotationMatrix(inp: NDArray, R: NDArray) -> NDArray:
            res = np.zeros_like(inp)
            for i1 in range(3):
                for i2 in range(3):
                    for i3 in range(3):
                        for j1 in range(3):
                            for j2 in range(3):
                                for j3 in range(3):
                                    res[i1, i2, i3] += R[i1, j1] * R[i2, j2] * R[i3, j3] * inp[j1, j2, j3]

            return res

        Rx = Core.RotationMatrixX(phiX)
        Ry = Core.RotationMatrixY(phiY)
        Rz = Core.RotationMatrixZ(phiZ)

        res = ApplyRotationMatrix(tensor, Rx)
        res = ApplyRotationMatrix(res, Ry)
        res = ApplyRotationMatrix(res, Rz)

        return res


# ===============================================================================
# _HomogeneousWave
# ===============================================================================


class _HomogeneousWave:
    def __init__(self, layer: _NonlinearLayer) -> None:
        self.layer = layer

    def Solve(self, kx: complex) -> None:
        layer = self.layer

        self.kx = kx
        self.k = self.layer.k0 * self.layer.n(self.layer.wl)
        self.kzF = np.sqrt(self.k**2.0 - self.kx**2.0, dtype=complex)
        self.kzB = -self.kzF
        self.kz = np.array([[self.kzF, self.kzB]]).T

        if self.kzF.imag < 0.0:
            raise ValueError(f"kzF imaginary part negative {self.kzF}")

        # Propagation matrix

        if layer.d == math.inf:
            self.propMatrix = np.identity(2, dtype=complex)
        else:
            phase = np.exp(1.0j * self.kzF * layer.d)
            self.propMatrix = np.array([[phase, 0.0], [0.0, 1.0 / phase]], dtype=complex)

    def GetMainFields(self, zs: NDArray) -> NDArray:
        U0 = np.array([[self.layer.fieldAmps[0, 0], self.layer.fieldAmps[1, 0]]]).T
        U = U0 * np.exp(1.0j * self.kz * zs)
        return U


# ===============================================================================
# _InhomogeneosWave
# ===============================================================================


class _InhomogeneosWave:
    def __init__(self, layer: _NonlinearLayer) -> None:
        self.layer = layer

    def Solve(self, kx: complex, kpS: tuple | None) -> None:
        layer = self.layer
        kSzF, pF, pB = [0.0, [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]] if kpS is None else kpS

        self.kx = kx
        self.kSzF = kSzF
        self.kSzB = -kSzF
        self.kSz = np.array([[self.kSzF, self.kSzB]]).T
        self.kS = np.sqrt(self.kx**2.0 + self.kSzF**2.0, dtype=complex)

        self.pxF, self.pyF, self.pzF = [0.0, 0.0, 0.0] if pF is None else pF
        self.pxB, self.pyB, self.pzB = [0.0, 0.0, 0.0] if pB is None else pB
        self.px = np.array([[self.pxF, self.pxB]]).T
        self.py = np.array([[self.pyF, self.pyB]]).T
        self.pz = np.array([[self.pzF, self.pzB]]).T

        if self.kS.imag < 0.0:
            # print ValueError("kS imaginary part negative %s" % (self.kS))
            self.kS *= -1.0

        # Check if layer is nonlinear

        if layer.pol == "s":
            self.isNonlinearF = abs(self.pyF) > 1e-200
            self.isNonlinearB = abs(self.pyB) > 1e-200
        elif layer.pol == "p":
            self.isNonlinearF = abs(self.pxF) > 1e-200 or abs(self.pzF) > 1e-200
            self.isNonlinearB = abs(self.pxB) > 1e-200 or abs(self.pzB) > 1e-200
        else:
            raise NotImplementedError("Unknown polarization.")
        self.isNonlinear = self.isNonlinearF or self.isNonlinearB

        # Calc B-s

        if not self.isNonlinear:
            self.ByF = 0.0
            self.ByB = 0.0
        else:
            if layer.pol == "s":
                BFunc = lambda py: -(py / Constants.eps0) * layer.k0**2.0 / (layer.k**2.0 - self.kS**2.0)
                self.ByF = BFunc(self.pyF)
                self.ByB = BFunc(self.pyB) * np.exp(1.0j * self.kSzF * layer.d)
            elif layer.pol == "p":
                BFunc = lambda px, pz, kSz: -layer.omega / (layer.k**2.0 - self.kS**2.0) * (kSz * px - self.kx * pz)
                self.ByF = BFunc(self.pxF, self.pzF, self.kSzF)
                self.ByB = BFunc(self.pxB, self.pzB, -self.kSzF) * np.exp(1.0j * self.kSzF * layer.d)
            else:
                raise NotImplementedError("Unknown polarization.")
        self.By = np.array([[self.ByF, self.ByB]]).T

        # Propagation Matrix NL

        if layer.d == math.inf:
            self.propMatrixNL = np.array([[0.0, 0.0]], dtype=complex).T
            if self.isNonlinear:
                raise NotImplementedError("First and last medium must be linear.")
        else:
            if self.isNonlinear:
                self.propMatrixNL = self.By * (np.exp(1.0j * self.kSz * layer.d) - np.exp(1.0j * layer.hw.kz * layer.d))
            else:
                self.propMatrixNL = np.array([[0.0, 0.0]], dtype=complex).T

    def GetMainFields(self, zs: NDArray) -> NDArray:
        U = self.By * (np.exp(1.0j * self.kSz * zs) - np.exp(1.0j * self.layer.hw.kz * zs))
        return U


# ===============================================================================
# _NonlinearLayer
# ===============================================================================


class _NonlinearLayer(Core.ParamsBaseClass):
    def __init__(self, d: float, n: object, **kwargs: Any) -> None:
        self._params = ["d", "n", "kpS", "kpA"]
        self.d = d
        self.n = n
        self.kpS = None
        self.kpA = None

        self.hw = _HomogeneousWave(self)
        self.iws = _InhomogeneosWave(self)
        self.iwa = _InhomogeneosWave(self)

        super().__init__(**kwargs)

    def Solve(self, wl: float, beta: float, pol: str, lastLayer: bool = False) -> None:
        self.wl = wl
        self.beta = beta
        self.pol = pol
        self.omega = Constants.WlToOmega(wl)
        self.lastLayer = lastLayer

        # Find wavevectors

        self.eps = self.n(wl) ** 2.0
        self.k0 = 2.0 * np.pi / wl
        self.k = self.k0 * self.n(self.wl)
        kx = beta * self.k0

        # Solve waves

        self.hw.Solve(kx)
        self.iws.Solve(kx, self.kpS)
        self.iwa.Solve(kx, self.kpA)

        # Propagation matrix

        self.propMatrix = self.hw.propMatrix
        self.propMatrixNL = self.iws.propMatrixNL + self.iwa.propMatrixNL

    def GetMainFields(self, zs: NDArray) -> NDArray:
        return self.hw.GetMainFields(zs) + self.iws.GetMainFields(zs) + self.iwa.GetMainFields(zs)

    def GetFields(self, zs: NDArray, outBackAndForward: bool = False) -> tuple[NDArray, NDArray]:
        E = np.zeros((2, len(zs), 3), dtype=complex)
        H = np.zeros((2, len(zs), 3), dtype=complex)

        expKSz = np.exp(1.0j * self.iws.kSz * zs)
        expKAz = np.exp(1.0j * self.iwa.kSz * zs)
        if self.pol == "p":
            c = -1.0 / (self.omega * self.eps * Constants.eps0)
            Hy = self.GetMainFields(zs)

            H[:, :, Y] = Hy
            E[:, :, X] = c * (
                -self.hw.kz * Hy
                + self.iws.By * (self.hw.kz - self.iws.kSz) * expKSz
                + self.omega * self.iws.px * expKSz
                + self.iwa.By * (self.hw.kz - self.iwa.kSz) * expKAz
                + self.omega * self.iwa.px * expKAz
            )
            E[:, :, Z] = c * (self.hw.kx * Hy + self.omega * self.iws.pz * expKSz + self.omega * self.iwa.pz * expKAz)
        elif self.pol == "s":
            c = 1.0 / (self.omega * Constants.mu0)
            Ey = self.GetMainFields(zs)

            E[:, :, Y] = Ey
            H[:, :, X] = c * (
                -self.hw.kz * Ey
                + self.iws.By * (self.hw.kz - self.iws.kSz) * expKSz
                + self.iwa.By * (self.hw.kz - self.iwa.kSz) * expKAz
            )
            H[:, :, Z] = c * (self.hw.kx * Ey)
        else:
            raise NotImplementedError("Unknown polarization.")

        if outBackAndForward:
            return E, H
        else:
            return (E[0, :, :] + E[1, :, :]), (H[0, :, :] + H[1, :, :])

    def GetAmplitudes(self, z: float = 0.0) -> tuple[tuple[NDArray, NDArray], tuple[NDArray, NDArray]]:
        (E0FTemp, E0BTemp), (H0FTemp, H0BTemp) = self.GetFields(np.array([z]), outBackAndForward=True)

        E0F, E0B = E0FTemp[0, :], E0BTemp[0, :]
        H0F, H0B = H0FTemp[0, :], H0BTemp[0, :]
        return (E0F, E0B), (H0F, H0B)

    def GetPowerFlow(self, z: float) -> float:
        (E0F, E0B), (H0F, H0B) = self.GetAmplitudes(z=z)
        E, H = E0F + E0B, H0F + H0B
        Sz = 0.5 * (E[0] * np.conj(H[1]) - E[1] * np.conj(H[0])).real
        return Sz

    def GetAbsorbedPower(self) -> float:
        if self.eps.imag == 0.0:
            # If layer is not absorbing
            return 0.0

        (E0F, E0B), (_, __) = self.GetAmplitudes()
        kz = self.hw.kz[0, 0]
        intValue1 = -1.0j * (np.exp(1.0j * (kz - np.conj(kz)) * self.d) - 1.0) / (kz - np.conj(kz))
        intValue2 = -1.0j * (np.exp(1.0j * (kz + np.conj(kz)) * self.d) - 1.0) / (kz + np.conj(kz))
        intValue3 = 1.0j * (np.exp(-1.0j * (kz + np.conj(kz)) * self.d) - 1.0) / (kz + np.conj(kz))
        intValue4 = 1.0j * (np.exp(-1.0j * (kz - np.conj(kz)) * self.d) - 1.0) / (kz - np.conj(kz))

        intValue = (
            np.dot(E0F, np.conj(E0F)) * intValue1
            + np.dot(E0F, np.conj(E0B)) * intValue2
            + np.dot(E0B, np.conj(E0F)) * intValue3
            + np.dot(E0B, np.conj(E0B)) * intValue4
        )

        res = 0.5 * Constants.eps0 * self.eps.imag * self.omega * intValue
        return res.real

    def GetSrcPower(self) -> float:
        absorbed = self.GetAbsorbedPower()
        if self.d == math.inf:
            deltaS = 0.0
        else:
            deltaS = self.GetPowerFlow(self.d) - self.GetPowerFlow(0.0)

        res = absorbed + deltaS
        return res


# ===============================================================================
# _NonlinearTmm
# ===============================================================================


class _NonlinearTmm(Core.ParamsBaseClass):
    MODE_INCIDENT = "incident"
    MODE_NONLINEAR = "nonlinear"

    def __init__(self, **kwargs: Any) -> None:
        self._params = ["wl", "beta", "pol", "I0", "overrideE0", "mode"]
        self.I0 = 1e10
        self.overrideE0 = None  # E0 specified in vacuum, not in first layer
        self.layers: list[_NonlinearLayer] = []

        super().__init__(**kwargs)

    def AddLayer(self, *args: Any, **kwargs: Any) -> None:
        self.layers.append(_NonlinearLayer(*args, **kwargs))

    def Solve(self, **kwargs: Any) -> tuple:
        self.SetParams(**kwargs)

        # Solve all layers

        for i, layer in enumerate(self.layers):
            last = i == len(self.layers) - 1
            layer.Solve(self.wl, self.beta, self.pol, last)

        # Solve system matrix

        mSysL, mSysNL = self.__SystemMatrix(len(self.layers))

        # Calc reflection and transmission

        l0 = self.layers[0]
        lL = self.layers[-1]
        omega = Constants.WlToOmega(self.wl)
        if self.mode == self.MODE_INCIDENT:
            if self.overrideE0 is None:
                cosTh0 = l0.hw.kz[0, 0].real / l0.k.real
                if self.pol == "p":
                    inc = np.sqrt(
                        abs(2.0 * omega * Constants.eps0 / l0.hw.kzF.real * l0.eps.real * self.I0 * cosTh0),
                        dtype=complex,
                    )
                elif self.pol == "s":
                    inc = np.sqrt(abs(2.0 * omega * Constants.mu0 / l0.hw.kzF.real * self.I0 * cosTh0), dtype=complex)
                else:
                    raise NotImplementedError("Unknown polarization.")
            else:
                if self.pol == "p":
                    inc = self.overrideE0 * np.sqrt(l0.n(self.wl).real) * Constants.eps0 * Constants.c
                elif self.pol == "s":
                    inc = self.overrideE0 / np.sqrt(l0.n(self.wl).real)
                else:
                    raise NotImplementedError("Unknown polarization.")

            r = -inc * mSysL[1, 0] / mSysL[1, 1]
            t = inc * np.linalg.det(mSysL) / mSysL[1, 1]
        elif self.mode == self.MODE_NONLINEAR:
            inc = 0.0
            r = -mSysNL[1, 0] / mSysL[1, 1]
            t = mSysNL[0, 0] - mSysL[0, 1] / mSysL[1, 1] * mSysNL[1, 0]
        else:
            raise NotImplementedError("Unknown mode.")

        # Calc power flows

        if self.pol == "p":
            c = 1.0 / (2.0 * omega * Constants.eps0)
            I = c * l0.hw.kzF.real / l0.eps.real * abs(inc) ** 2.0
            R = c * l0.hw.kzF.real / l0.eps.real * abs(r) ** 2.0
            T = c * lL.hw.kzF.real / lL.eps.real * abs(t) ** 2.0
        elif self.pol == "s":
            c = 1.0 / (2.0 * omega * Constants.mu0)
            I = c * l0.hw.kzF.real * abs(inc) ** 2.0
            R = c * l0.hw.kzF.real * abs(r) ** 2.0
            T = c * lL.hw.kzF.real * abs(t) ** 2.0
        else:
            raise NotImplementedError("Unknown polarization.")

        # Calc field coefs in all layers

        self.layers[0].fieldAmps = np.array([[inc, r]], dtype=complex).T
        for i in range(1, len(self.layers)):
            mL, mNL = self.__SystemMatrix(i + 1)
            self.layers[i].fieldAmps = np.dot(mL, self.layers[0].fieldAmps) + mNL

        # Absorbed energy
        A = self.GetAbsorbedPower()

        return inc, r, t, I, R, T, A

    def __TransferMatrix(self, layerNr: int) -> tuple[NDArray, NDArray]:

        def TransferMatrixNL(w1: _InhomogeneosWave, w2: _InhomogeneosWave) -> NDArray:
            exp1F = np.exp(1.0j * w1.kSzF * l1.d) if l1.d != math.inf else 0.0
            exp1B = np.exp(-1.0j * w1.kSzF * l1.d) if l1.d != math.inf else 0.0
            # exp1B = 1.0 / exp1F if l1.d != float("inf") else 0.0

            f1 = (w1.kSzF - l1.hw.kzF) * (w1.ByF * exp1F - w1.ByB * exp1B)
            f2 = (w2.kSzF - l2.hw.kzF) * (w2.ByB - w2.ByF)

            if self.pol == "p":
                f3 = -l1.omega * l2.eps * (w1.pxB * exp1B + w1.pxF * exp1F) + l1.omega * l1.eps * (w2.pxB + w2.pxF)
                cNL = 1.0 / (2.0 * l1.eps * l2.hw.kzF) * (l2.eps * f1 + l1.eps * f2 + f3)
            elif self.pol == "s":
                cNL = 1.0 / (2.0 * l2.hw.kzF) * (f1 + f2)
            else:
                raise NotImplementedError("Unknown polarization.")

            transferMatrixNL = np.array([[cNL, -cNL]], dtype=complex).T
            return transferMatrixNL

        l1, l2 = self.layers[layerNr], self.layers[layerNr + 1]

        # Polarization specific

        if self.pol == "p":
            a = (l2.eps * l1.hw.kzF) / (l1.eps * l2.hw.kzF)
        elif self.pol == "s":
            a = l1.hw.kzF / l2.hw.kzF
        else:
            raise NotImplementedError("Unknown polarization.")

        # Transfer matrices
        transferMatrix = 0.5 * np.array([[1.0 + a, 1.0 - a], [1.0 - a, 1.0 + a]], dtype=complex)
        transferMatrixNL = TransferMatrixNL(l1.iws, l2.iws) + TransferMatrixNL(l1.iwa, l2.iwa)

        return transferMatrix, transferMatrixNL

    def __SystemMatrix(self, nLayers: int) -> tuple[NDArray, NDArray]:
        # Independent of polarization

        mSysL = np.identity(2, dtype=complex)
        mSysNL = np.array([[0.0, 0.0]], dtype=complex).T

        for i in range(nLayers - 1):
            layer = self.layers[i]
            transferL, transferNL = self.__TransferMatrix(i)

            mSysL = np.dot(layer.propMatrix, mSysL)
            mSysL = np.dot(transferL, mSysL)

            mSysNL = np.dot(layer.propMatrix, mSysNL) + layer.propMatrixNL
            mSysNL = np.dot(transferL, mSysNL) + transferNL

        return mSysL, mSysNL

    def GetFields(self, zs: NDArray, outBackAndForward: bool = False) -> tuple[NDArray, NDArray]:
        if not outBackAndForward:
            resE = np.zeros((len(zs), 3), dtype=complex)
            resH = np.zeros((len(zs), 3), dtype=complex)
        else:
            resE = np.zeros((2, len(zs), 3), dtype=complex)
            resH = np.zeros((2, len(zs), 3), dtype=complex)
        layerDsIndexes = self.__GetLayerIndexesAndDistances(zs)

        for layerId, layerDist, bIndex, eIndex in layerDsIndexes:
            if bIndex == eIndex:
                continue

            E, H = self.layers[layerId].GetFields(zs[bIndex:eIndex] - layerDist, outBackAndForward=outBackAndForward)
            if not outBackAndForward:
                resE[bIndex:eIndex, :] = E
                resH[bIndex:eIndex, :] = H
            else:
                resE[:, bIndex:eIndex, :] = E
                resH[:, bIndex:eIndex, :] = H

        return resE, resH

    def GetFields2D(self, zs: NDArray, xs: NDArray, outBackAndForward: bool = False) -> tuple[NDArray, NDArray]:
        E1D, H1D = self.GetFields(zs, outBackAndForward=outBackAndForward)
        phaseX = np.exp(1.0j * 2.0 * np.pi / self.wl * self.beta * xs)

        if not outBackAndForward:
            resE = np.zeros((len(zs), len(xs), 3), dtype=complex)
            resH = np.zeros((len(zs), len(xs), 3), dtype=complex)

            for i in range(3):
                resE[:, :, i] = np.outer(E1D[:, i], phaseX)
                resH[:, :, i] = np.outer(H1D[:, i], phaseX)

        else:
            resE = np.zeros((2, len(zs), len(xs), 3), dtype=complex)
            resH = np.zeros((2, len(zs), len(xs), 3), dtype=complex)

            for j in range(2):
                for i in range(3):
                    resE[j, :, :, i] = np.outer(E1D[j, :, i], phaseX)
                    resH[j, :, :, i] = np.outer(H1D[j, :, i], phaseX)

        return resE, resH

    def GetAbsorbedPower(self) -> float:
        return sum(layer.GetAbsorbedPower() for layer in self.layers)

    def __GetLayerIndexesAndDistances(self, zs: NDArray) -> list[tuple[int, float, int, int]]:
        layerDsIndexes = []
        curLayer, curDist, layerDist, beginIndex = 0, 0.0, 0.0, 0
        for i in range(len(zs) + 1):
            if (i == len(zs)) or (zs[i] >= curDist):
                layerDsIndexes.append((curLayer, layerDist, beginIndex, i))
                if i == len(zs):
                    break
                layerDist += self.layers[curLayer].d if curLayer > 0 else 0.0
                curLayer += 1
                curDist += self.layers[curLayer].d
                beginIndex = i
        return layerDsIndexes


# ===============================================================================
# SecondOrderNLTmm
# ===============================================================================


class SecondOrderNLTmm(Core.ParamsBaseClass):
    SFG = "sfg"
    DFG = "dfg"

    def __init__(self, **kwargs: Any) -> None:
        self._params = [
            "I0P1",
            "wlP1",
            "betaP1",
            "polP1",
            "overrideE0P1",
            "I0P2",
            "wlP2",
            "betaP2",
            "polP2",
            "overrideE0P2",
            "polGen",
            "process",
        ]
        self.layers: list[tuple] = []
        self.I0P1 = 1e10
        self.I0P2 = 1e10
        self.overrideE0P1 = None
        self.overrideE0P2 = None
        self.process = self.SFG
        super().__init__(**kwargs)

    def AddLayer(self, d: float, n: object, chi2: Chi2Tensor | None = None) -> None:
        if chi2 is None:
            chi2 = Chi2Tensor()
        self.layers.append((d, n, chi2))

    def Solve(self, **kwargs: Any) -> tuple:
        self.SetParams(**kwargs)

        self._SolveFundamentalFields()

        if self.process == self.SFG:
            self.wlGen = OmegaToWl(WlToOmega(self.wlP1) + WlToOmega(self.wlP2))
            self.betaGen = self.wlGen * (self.betaP1 / self.wlP1 + self.betaP2 / self.wlP2)
            kSzFunc = lambda kzFP1, kzFP2: kzFP1 + kzFP2
            kAzFunc = lambda kzFP1, kzFP2: kzFP1 - kzFP2
            pFunc = lambda chi2, E1, E2: chi2.GetNonlinearPolarization(E1, E2)
        elif self.process == self.DFG:
            self.wlGen = OmegaToWl(WlToOmega(self.wlP1) - WlToOmega(self.wlP2))
            self.betaGen = self.wlGen * (self.betaP1 / self.wlP1 - self.betaP2 / self.wlP2)
            kSzFunc = lambda kzFP1, kzFP2: kzFP1 - np.conj(kzFP2)
            kAzFunc = lambda kzFP1, kzFP2: kzFP1 + np.conj(kzFP2)
            pFunc = lambda chi2, E1, E2: chi2.GetNonlinearPolarization(E1, np.conj(E2))
        else:
            raise ValueError(f"Unknown process: {self.process}")

        self.tmmGen = _NonlinearTmm(wl=self.wlGen, pol=self.polGen, mode=_NonlinearTmm.MODE_NONLINEAR)

        for i, (d, n, chi2) in enumerate(self.layers):
            (E0P1F, E0P1B), _ = self.tmmP1.layers[i].GetAmplitudes()
            (E0P2F, E0P2B), _ = self.tmmP2.layers[i].GetAmplitudes()
            kzFP1, kzFP2 = self.tmmP1.layers[i].hw.kzF, self.tmmP2.layers[i].hw.kzF

            kSz, kAz = kSzFunc(kzFP1, kzFP2), kAzFunc(kzFP1, kzFP2)
            pSF, pSB = pFunc(chi2, E0P1F, E0P2F), pFunc(chi2, E0P1B, E0P2B)
            pAF, pAB = pFunc(chi2, E0P1F, E0P2B), pFunc(chi2, E0P1B, E0P2F)

            if kSz.imag < 0.0:
                # print "reversed kSz"
                kSz, pSF, pSB = -kSz, pSB, pSF

            if kAz.imag < 0.0:
                # print "reversed kAz"
                kAz, pAF, pAB = -kAz, pAB, pAF

            self.tmmGen.AddLayer(d, n, kpS=(kSz, pSF, pSB), kpA=(kAz, pAF, pAB))

        self.resGen = self.tmmGen.Solve(beta=self.betaGen)

        return self.resP1, self.resP2, self.resGen

    def Sweep(
        self,
        paramNames: list[str],
        paramValues: list[NDArray],
        enhpos: tuple[int, float] | None = None,
    ) -> dict[str, NDArray]:
        keysP1 = ["iP1", "rP1", "tP1", "IP1", "RP1", "TP1", "AP1"]
        keysP2 = ["iP2", "rP2", "tP2", "IP2", "RP2", "TP2", "AP2"]
        keysGen = ["iGen", "rGen", "tGen", "IGen", "RGen", "TGen", "AGen"]
        res = {}

        res["betasGen"] = np.zeros(len(paramValues[0]))

        for k in keysP1 + keysP2 + keysGen:
            res[k] = np.zeros((len(paramValues[0])), dtype=complex)

        if enhpos is not None:
            res["enhP1"] = np.zeros(len(paramValues[0]))
            res["enhP2"] = np.zeros(len(paramValues[0]))

        for i in range(len(paramValues[0])):
            for name, values in zip(paramNames, paramValues):
                self.SetParams(**{name: values[i]})

            rrP1, rrP2, rrGen = self.Solve()
            res["betasGen"][i] = self.tmmGen.beta.real

            for key, val in zip(keysP1, rrP1):
                res[key][i] = val
            for key, val in zip(keysP2, rrP2):
                res[key][i] = val
            for key, val in zip(keysGen, rrGen):
                res[key][i] = val

            if enhpos is not None:
                enhLayer, enhDist = enhpos
                EP1, _ = self.tmmP1.layers[enhLayer].GetFields(np.array([enhDist]))
                (E0P1, _), __ = self.tmmP1.layers[0].GetFields(np.array([0.0]), outBackAndForward=True)
                enhP1 = np.linalg.norm(EP1[0, :]) / (
                    np.linalg.norm(E0P1[0, :]) * np.sqrt(self.tmmP1.layers[0].n(self.wlP1).real)
                )

                EP2, _ = self.tmmP2.layers[enhLayer].GetFields(np.array([enhDist]))
                (E0P2, _), __ = self.tmmP2.layers[0].GetFields(np.array([0.0]), outBackAndForward=True)
                enhP2 = np.linalg.norm(EP2[0, :]) / (
                    np.linalg.norm(E0P2[0, :]) * np.sqrt(self.tmmP2.layers[0].n(self.wlP2).real)
                )

                res["enhP1"][i] = enhP1
                res["enhP2"][i] = enhP2

        return res

    def _SolveFundamentalFields(self) -> None:
        # Solve fundametal field 1
        self.tmmP1 = _NonlinearTmm(
            I0=self.I0P1, overrideE0=self.overrideE0P1, wl=self.wlP1, pol=self.polP1, mode=_NonlinearTmm.MODE_INCIDENT
        )

        for d, n, _ in self.layers:
            self.tmmP1.AddLayer(d, n)
        self.resP1 = self.tmmP1.Solve(beta=self.betaP1)

        # Solve fundametal field 2
        self.tmmP2 = _NonlinearTmm(
            I0=self.I0P2, overrideE0=self.overrideE0P2, wl=self.wlP2, pol=self.polP2, mode=_NonlinearTmm.MODE_INCIDENT
        )

        for d, n, _ in self.layers:
            self.tmmP2.AddLayer(d, n)
        self.resP2 = self.tmmP2.Solve(beta=self.betaP2)
