#include "mocca/kernel.hpp"

#include "config_bridge.hpp"
#include "embedded_tables.hpp"
#include "line_codec.hpp"
#include "physics_engine.hpp"

#include <array>
#include <cmath>
#include <optional>
#include <utility>
#include <vector>

namespace mocca {

class ModernCascadeKernel::Impl {
public:
    explicit Impl(const SimulationConfig& config)
        : config_(config),
          tables_(detail::bundled_numeric_tables()),
          primitives_(tables_, config_.numerics.matrix_element_precision_digits) {}

    [[nodiscard]] SimulationResult run() const {
        detail::KernelState cfg = detail::build_kernel_config(tables_, config_);
        primitives_.prepare_case(cfg);
        auto [lyman_sum, states] = propagate_cascade(cfg);
        auto lines = detail::collect_lines(cfg);
        return SimulationResult{
            lyman_sum,
            static_cast<int>(lines.size()),
            std::move(lines),
            std::move(states),
            cfg.warnings,
        };
    }

private:
    [[nodiscard]] static std::array<double, 5> popj(int l1, int l2, int ll) {
        std::array<double, 5> p{};
        const int li = ll + 1;
        const int l = std::max(l1, l2);
        if (li == 1) {
            p[1] = static_cast<double>(l + 1) / static_cast<double>(2 * l + 1);
            p[2] = 1.0 - p[1];
            return p;
        }
        if (li == 2) {
            p[1] = static_cast<double>((l + 1) * (2 * l - 1)) / static_cast<double>(4 * l * l - 1);
            p[2] = 1.0 / static_cast<double>(4 * l * l - 1);
            p[3] = 1.0 - (p[1] + p[2]);
            return p;
        }
        if (li == 3) {
            if (l1 != l2) {
                p[1] = static_cast<double>(l + 1) / static_cast<double>(2 * l + 1);
                p[2] = 2.0 / static_cast<double>((2 * l - 3) * (2 * l + 1));
                p[3] = 1.0 - (p[1] + p[2]);
                return p;
            }
            const double d = static_cast<double>((2 * l + 1) * (2 * l + 1));
            p[1] = static_cast<double>((l + 2) * (2 * l - 1)) / d;
            p[2] = static_cast<double>((l - 1) * (2 * l + 3)) / d;
            p[3] = 3.0 / d;
            p[4] = p[3];
            return p;
        }
        if (std::abs(l1 - l2) != 1) {
            p[1] = static_cast<double>(l + 1) / static_cast<double>(2 * l + 1);
            p[2] = 3.0 / static_cast<double>((2 * l - 5) * (2 * l + 1));
            p[3] = 1.0 - (p[1] + p[2]);
            return p;
        }
        const double d = static_cast<double>(4 * l * l - 1);
        p[1] = static_cast<double>((2 * l - 3) * (l + 2)) / d;
        p[2] = 5.0 / d;
        p[3] = 6.0 / d;
        p[4] = 1.0 - (p[1] + p[2] + p[3]);
        return p;
    }

    [[nodiscard]] static double beta(int l1, int j1, int l2, int j2, int l) {
        const double a1 = 0.25 * static_cast<double>(j1 * (j1 + 2));
        const double a2 = 0.25 * static_cast<double>(j2 * (j2 + 2));
        return (
            (a2 - static_cast<double>(l2 * (l2 + 1)) + 0.750) /
            (a1 - static_cast<double>(l1 * (l1 + 1)) + 0.750)) *
            ((a1 + a2 - static_cast<double>(l * (l + 1))) / (2.0 * a2));
    }

    [[nodiscard]] std::pair<double, std::vector<StateSummary>> propagate_cascade(
        detail::KernelState& cfg) const {
        // The kernel keeps the original one-based workspace layout to preserve
        // the validated cascade-update order while expressing it in explicit
        // C++ containers instead of the monolithic runner entry point.
        std::vector<double> popt(4, 0.0);
        const std::vector<double> pop1 = cfg.pop;
        std::vector<double> pop2(7, 0.0);
        popt[1] = 2.0 * cfg.pop[1];
        popt[2] = 2.0 * cfg.pop[2] + 6.0 * cfg.pop[3];
        popt[3] = 2.0 * cfg.pop[4] + 6.0 * cfg.pop[5] + 10.0 * cfg.pop[6];
        pop2[1] = pop1[1];
        pop2[2] = (popt[2] <= 1.0e-20) ? 1.0 : 1.0 / popt[2];
        pop2[3] = (popt[2] <= 1.0e-20) ? 1.0 : 1.0 / popt[2];
        pop2[4] = (popt[3] <= 1.0e-20) ? 1.0 : 1.0 / popt[3];
        pop2[5] = (popt[3] <= 1.0e-20) ? 1.0 : 1.0 / popt[3];
        pop2[6] = (popt[3] <= 1.0e-20) ? 1.0 : 1.0 / popt[3];

        const int ms = 3;
        const double rk = cfg.widthk / cfg.hbar;
        const int nu = cfg.nmax * (cfg.nmax + 1) / 2;
        std::vector<double> pnl(nu + 1, 0.0);
        std::vector<double> polpos(nu + 1, 0.0);
        std::vector<double> polneg(nu + 1, 0.0);
        std::vector<double> width(nu + 1, 0.0);
        std::vector<double> convc(nu + 1, 0.0);
        std::vector<double> sporb(nu + 1, 0.0);
        std::vector<double> pc0(nu + 1, 0.0);
        std::vector<double> pc1(nu + 1, 0.0);
        std::vector<double> pc2(nu + 1, 0.0);
        std::vector<std::vector<double>> pc(ms + 1, std::vector<double>(nu + 1, 0.0));
        const int mu = cfg.nmax * (cfg.nmax - 1) / 2;

        for (int j = 1; j <= nu; ++j) {
            if (cfg.ip8 != 0) {
                pnl[j] = cfg.pln[j];
            }
        }
        for (int j = 1; j <= cfg.nmax; ++j) {
            const int jj = mu + j;
            if (cfg.ip8 == 0) {
                pnl[jj] = cfg.pl[j];
            }
            pc1[jj] = 2.0 - 2.0 * cfg.pop[1];
            pc2[jj] = 2.0 * cfg.pop[1] - 1.0;
            if (cfg.pop[1] <= 0.5) {
                pc0[jj] = 1.0 - 2.0 * cfg.pop[1];
                pc1[jj] = 2.0 * cfg.pop[1];
                pc2[jj] = 0.0;
            }
            if (pnl[jj] > 1.0e-20) {
                polpos[jj] = (1.0 + 2.0 / static_cast<double>(2 * j - 1)) / 3.0;
            }
            if (pnl[jj] > 1.0e-20 && j != 1) {
                polneg[jj] = (1.0 - 2.0 / static_cast<double>(2 * j - 1)) / 3.0;
            }
            for (int shell = 1; shell <= ms; ++shell) {
                pc[shell][jj] = popt[shell];
            }
        }

        double rlyman = 0.0;
        std::vector<double> zt(131, 0.0);
        std::vector<std::vector<double>> zt_shell(4, std::vector<double>(131, 0.0));
        std::vector<double> zr(131, 0.0);
        std::vector<double> za(131, 0.0);
        std::vector<double> za0(131, 0.0);
        std::vector<double> za1(131, 0.0);
        std::vector<double> za2(131, 0.0);
        for (int i1 = 1; i1 <= cfg.nmax; ++i1) {
            const int n1 = cfg.nmax + 1 - i1;
            for (int i2 = 1; i2 <= n1; ++i2) {
                const int l1 = i2 - 1;
                double rategt = 0.0;
                double rate0 = 2.0 * rk;
                double rate1 = rk;
                double rateauger = 0.0;
                double raterd = 0.0;
                const int k1 = detail::state_index(n1, l1);
                if (i1 > 1) {
                    if (pnl[k1] < 1.0e-20) {
                        pc[2][k1] = popt[2];
                        pc[3][k1] = popt[3];
                        pc0[k1] = 0.0;
                        pc1[k1] = 2.0 - 2.0 * pop2[1];
                        pc2[k1] = 2.0 * pop2[1] - 1.0;
                        if (pop2[1] < 0.5) {
                            pc0[k1] = 1.0 - 2.0 * pop2[1];
                            pc1[k1] = 2.0 * pop2[1];
                            pc2[k1] = 0.0;
                        }
                        pc[1][k1] = pc1[k1] + 2.0 * pc2[k1];
                    } else {
                        const double xnorm = 1.0 / pnl[k1];
                        pc0[k1] *= xnorm;
                        pc1[k1] *= xnorm;
                        pc2[k1] *= xnorm;
                        for (int shell = 2; shell <= ms; ++shell) {
                            pc[shell][k1] = std::max(pc[shell][k1] * xnorm, 0.0);
                        }
                        pc[1][k1] = pc1[k1] + 2.0 * pc2[k1];
                        pc[1][k1] = std::min(std::max(pc[1][k1], 0.0), 2.0);
                        for (int shell = 1; shell <= ms; ++shell) {
                            if (cfg.ipc[shell] != 0) {
                                pc[shell][k1] = popt[shell];
                            }
                        }
                        if (cfg.ipc[1] != 0) {
                            pc2[k1] = 2.0 * pop1[1] - 1.0;
                            pc1[k1] = 2.0 - 2.0 * pop1[1];
                            pc0[k1] = 0.0;
                            if (pop1[1] < 0.5) {
                                pc2[k1] = 0.0;
                                pc1[k1] = 2.0 * pop1[1];
                                pc0[k1] = 1.0 - 2.0 * pop1[1];
                            }
                        }
                        polpos[k1] = polpos[k1] * xnorm * static_cast<double>(2 * l1 + 1) /
                                     static_cast<double>(l1 + 1);
                        if (cfg.npol[n1] - l1 == 0 || cfg.npol[n1] - l1 == 1) {
                            polpos[k1] = (1.0 + 2.0 / static_cast<double>(2 * l1 + 1)) / 3.0;
                        }
                        if (l1 > 0) {
                            polneg[k1] = polneg[k1] * xnorm * static_cast<double>(2 * l1 + 1) /
                                         static_cast<double>(l1);
                            if (cfg.npol[n1] - l1 == 0 || cfg.npol[n1] - l1 == 1) {
                                polneg[k1] = (1.0 - 2.0 / static_cast<double>(2 * l1 + 1)) / 3.0;
                            }
                        }
                    }
                }
                if (n1 == 1) {
                    continue;
                }
                double pc012 = pc0[k1] + pc1[k1] + pc2[k1];
                if (std::abs(pc012) < 1.0e-20) {
                    pc012 = 1.0;
                }
                pc0[k1] /= pc012;
                pc1[k1] /= pc012;
                pc2[k1] /= pc012;

                cfg.pop[1] = 0.5 * pc[1][k1];
                cfg.pop[2] = pop2[2] * pc[2][k1];
                cfg.pop[3] = pop2[3] * pc[2][k1];
                cfg.pop[4] = pop2[4] * pc[3][k1];
                cfg.pop[5] = pop2[5] * pc[3][k1];
                cfg.pop[6] = pop2[6] * pc[3][k1];
                std::fill(zt.begin(), zt.end(), 0.0);
                for (auto& shell : zt_shell) {
                    std::fill(shell.begin(), shell.end(), 0.0);
                }
                std::fill(zr.begin(), zr.end(), 0.0);
                std::fill(za.begin(), za.end(), 0.0);
                std::fill(za0.begin(), za0.end(), 0.0);
                std::fill(za1.begin(), za1.end(), 0.0);
                std::fill(za2.begin(), za2.end(), 0.0);

                int k = 0;
                for (int i3 = 1; i3 <= 7; ++i3) {
                    const int l2 = l1 - 4 + i3;
                    if (l2 < 0) {
                        continue;
                    }
                    for (int i4 = 1; i4 <= n1; ++i4) {
                        const int n2 = n1 - i4 + 1;
                        if (n2 <= l2) {
                            continue;
                        }
                        if (n2 == n1 && (n2 != 2 || l1 != 0 || l2 != 1)) {
                            continue;
                        }
                        if (n1 == n2 && cfg.espm <= 1.0e-20) {
                            continue;
                        }
                        k += 1;
                        if (n1 == n2) {
                            cfg.ijk = 1;
                            cfg.energ = cfg.espm;
                        } else {
                            cfg.ijk = 0;
                        }
                        const double popq = cfg.pop[1];
                        cfg.pop[1] = 1.0;
                        zt[k] = primitives_.transition_rate(cfg, n1, l1, n2, l2) + 1.0e-10;
                        rategt += zt[k];
                        if (cfg.pop[1] < 1.0e-20) {
                            cfg.pop[1] = 1.0;
                        }
                        cfg.rsa[1] = cfg.rsa[1] / cfg.pop[1];
                        cfg.pop[1] = popq;
                        zr[k] = cfg.rad;
                        raterd += cfg.rad;
                        za[k] = cfg.rsa[1];
                        // `rad_to_auger` used to reconstruct the effective Auger
                        // width from `effective_total - radiative`, which loses
                        // several digits when the state is almost purely
                        // radiative. Accumulating the Auger contribution directly
                        // keeps the same physics while avoiding that cancellation.
                        const double transition_auger = cfg.rau + 1.0e-10;
                        rateauger +=
                            pc2[k1] * transition_auger +
                            pc1[k1] * (transition_auger - 0.5 * za[k]) +
                            pc0[k1] * (transition_auger - za[k]);
                        const double xr = zt[k] - cfg.rsa[1];
                        double t0 = xr + 0.5 * cfg.rsa[1] + rk;
                        double g0 = xr + 2.0 * rk + 1.0e-10;
                        rate1 = rate1 + t0 - rk + 1.0e-10;
                        rate0 = rate0 + xr + 1.0e-10;
                        double xm = zt[k] - za[k] * (pc0[k1] + 0.5 * pc1[k1]);
                        if (std::abs(xm) < 1.0e-10) {
                            xm = 1.0;
                        }
                        if (std::abs(t0) < 1.0e-10) {
                            t0 = 1.0;
                        }
                        if (std::abs(g0) < 1.0e-10) {
                            g0 = 1.0;
                        }
                        if (std::abs(zt[k]) < 1.0e-10) {
                            zt[k] = 1.0;
                        }
                        zt_shell[3][k] = pc[3][k1] - cfg.rsa[3] / xm;
                        if (zt_shell[3][k] < 0.0) {
                            zt_shell[3][k] = 0.0;
                        }
                        za2[k] = xr / zt[k] *
                                 (pc2[k1] + rk * pc1[k1] / t0 + 2.0 * rk * rk * pc0[k1] / t0 / g0);
                        za1[k] = za[k] * pc2[k1] / zt[k] +
                                 (xr + rk * za[k] / zt[k]) / t0 *
                                     (pc1[k1] + 2.0 * rk * pc0[k1] / g0);
                        za0[k] = pc1[k1] * za[k] / (2.0 * t0) +
                                 pc0[k1] / g0 * (xr + rk * za[k] / t0);
                        const double ym = za[k] * (
                            pc2[k1] / zt[k] +
                            pc1[k1] * (0.5 + rk / zt[k]) / t0 +
                            pc0[k1] * rk / t0 / g0 * (1.0 + 2.0 * rk / zt[k]));
                        const double xl =
                            rk * (pc1[k1] / t0 + 2.0 * pc0[k1] / t0 / g0 * (t0 + rk));
                        zt_shell[1][k] = pc[1][k1] - ym + xl;
                        zt_shell[2][k] = pc[2][k1] - cfg.rsa[2] / xm - xl;
                        if (zt_shell[1][k] < 0.0) {
                            zt_shell[1][k] = 0.0;
                        }
                        if (zt_shell[2][k] < 0.0) {
                            zt_shell[2][k] = 0.0;
                        }
                    }
                }

                width[k1] = (rategt * pc2[k1] + (rate1 - rk) * pc1[k1] +
                             (rate0 - 2.0 * rk) * pc0[k1]) *
                            cfg.hbar;
                if (width[k1] < 0.0) {
                    width[k1] = 0.0;
                }
                if (std::abs(rate0) < 1.0e-10) {
                    rate0 = 1.0;
                }
                if (std::abs(rate1) < 1.0e-10) {
                    rate1 = 1.0;
                }
                if (std::abs(rategt) < 1.0e-10) {
                    rategt = 1.0;
                }
                convc[k1] = raterd / (rateauger + 1.0e-10);
                if (convc[k1] < 0.0) {
                    convc[k1] = 9.999e99;
                }
                if (l1 != 0) {
                    sporb[k1] = 0.150 * std::pow(cfg.z, 4) /
                                static_cast<double>(n1 * n1 * n1 * l1 * (l1 + 1));
                }

                k = 0;
                for (int i3 = 1; i3 <= 7; ++i3) {
                    const int l2 = l1 - 4 + i3;
                    if (l2 < 0) {
                        continue;
                    }
                    for (int i4 = 1; i4 <= n1; ++i4) {
                        const int n2 = n1 - i4 + 1;
                        if (n2 <= l2) {
                            continue;
                        }
                        if (n2 == n1 && (n2 != 2 || l1 != 0 || l2 != 1)) {
                            continue;
                        }
                        if (n2 == n1 && cfg.espm <= 1.0e-20) {
                            continue;
                        }
                        k += 1;
                        const int k2 = detail::state_index(n2, l2);
                        const double bnorm = pnl[k1] * (
                            pc2[k1] * zt[k] / rategt +
                            pc1[k1] *
                                (zt[k] - 0.5 * za[k] + rk * zt[k] / rategt) / rate1 +
                            pc0[k1] *
                                (zt[k] - za[k] + 2.0 * rk / rate1 *
                                    (zt[k] - 0.5 * za[k] + rk * zt[k] / rategt)) /
                                rate0);
                        for (int shell = 1; shell <= ms; ++shell) {
                            pc[shell][k2] += zt_shell[shell][k] * bnorm;
                        }
                        pnl[k2] += bnorm;
                        pc2[k2] += za2[k] * bnorm;
                        pc1[k2] += za1[k] * bnorm;
                        pc0[k2] += za0[k] * bnorm;

                        const double radint =
                            pnl[k1] * zr[k] *
                            (pc2[k1] / rategt +
                             pc1[k1] * (1.0 + rk / rategt) / rate1 +
                             pc0[k1] *
                                 (1.0 + 2.0 * rk / rate1 * (1.0 + rk / rategt)) /
                                 rate0);
                        const int ll = std::abs(l1 - l2);
                        const auto p = popj(l1, l2, ll);
                        if (cfg.ipol == 0) {
                            const int li = ll + 1;
                            const int j1u = 2 * l1 + 1;
                            const int j2u = 2 * l2 + 1;
                            const int j1d = j1u - 2;
                            const int j2d = j2u - 2;
                            if (li <= 1) {
                                polpos[k2] += bnorm * p[1] * polpos[k1];
                                polneg[k2] += bnorm * p[2] * polneg[k1];
                            } else if (l2 <= l1) {
                                polpos[k2] +=
                                    bnorm *
                                    (p[1] * polpos[k1] * beta(l1, j1u, l2, j2u, ll) +
                                     p[2] * polneg[k1] * beta(l1, j1d, l2, j2u, ll));
                                if (j1d != -1 && j2d != -1) {
                                    polneg[k2] +=
                                        bnorm * p[3] * polneg[k1] *
                                        beta(l1, j1d, l2, j2d, ll);
                                }
                            } else {
                                polpos[k2] +=
                                    bnorm * p[1] * polpos[k1] *
                                    beta(l1, j1u, l2, j2u, ll);
                                polneg[k2] +=
                                    bnorm *
                                    (p[2] * polpos[k1] * beta(l1, j1u, l2, j2d, ll) +
                                     p[3] * polneg[k1] * beta(l1, j1d, l2, j2d, ll));
                            }
                        }
                        const int li = (ll != 0) ? ll : 2;
                        detail::record_lines(cfg, n1, l1, n2, l2, li, radint);
                        if (l1 == 1 && l2 == 0 && n2 == 1) {
                            rlyman += radint;
                        }
                    }
                }
            }
        }

        std::vector<StateSummary> states;
        states.reserve(static_cast<std::size_t>(cfg.nmax * (cfg.nmax + 1) / 2));
        pc[1][1] = pc1[1] + 2.0 * pc2[1];
        for (int m1 = 1; m1 <= cfg.nmax; ++m1) {
            const int n1 = cfg.nmax + 1 - m1;
            for (int ll1 = 1; ll1 <= n1; ++ll1) {
                const int l1 = ll1 - 1;
                const int k1 = detail::state_index(n1, l1);
                std::optional<double> polar_down;
                if (l1 > 0) {
                    polar_down =
                        -polneg[k1] * (static_cast<double>(l1) + 0.5) /
                        (static_cast<double>(l1) - 0.5);
                }
                std::optional<double> polar_up =
                    (cfg.ipol == 1) ? std::optional<double>{} : std::optional<double>{polpos[k1]};
                std::optional<double> width_ev =
                    (n1 == 1 || (l1 == 0 && pnl[k1] > 1.0e-20))
                        ? std::optional<double>{}
                        : std::optional<double>{width[k1]};
                std::optional<double> rad_to_auger =
                    (n1 == 1 || (l1 == 0 && pnl[k1] > 1.0e-20))
                        ? std::optional<double>{}
                        : std::optional<double>{convc[k1]};
                std::optional<double> spin_orbit_ev =
                    (cfg.ipol == 1 || n1 == 1 || (l1 == 0 && pnl[k1] > 1.0e-20))
                        ? std::optional<double>{}
                        : std::optional<double>{sporb[k1]};
                states.push_back(StateSummary{
                    n1,
                    l1,
                    pnl[k1],
                    polar_up,
                    polar_down,
                    width_ev,
                    rad_to_auger,
                    spin_orbit_ev,
                    pc[1][k1],
                    pc[2][k1],
                    pc[3][k1],
                });
            }
        }
        for (int i = 1; i <= 6; ++i) {
            cfg.pop[i] = pop1[i];
        }
        return {rlyman, std::move(states)};
    }

    SimulationConfig config_;
    detail::NumericTables tables_;
    detail::PhysicsEngine primitives_;
};

ModernCascadeKernel::ModernCascadeKernel(const SimulationConfig& config)
    : impl_(std::make_unique<Impl>(config)) {}

ModernCascadeKernel::~ModernCascadeKernel() = default;

ModernCascadeKernel::ModernCascadeKernel(ModernCascadeKernel&&) noexcept = default;

ModernCascadeKernel& ModernCascadeKernel::operator=(ModernCascadeKernel&&) noexcept = default;

SimulationResult ModernCascadeKernel::run() const {
    return impl_->run();
}

SimulationResult run_modern_kernel(const SimulationConfig& config) {
    return ModernCascadeKernel(config).run();
}

}  // namespace mocca