// SPDX-License-Identifier: MPL-2.0
#include <cstdint>
#include <filesystem>
#include <pybind11/pybind11.h>
#include <pybind11/stl.h>
#include <pybind11/stl_bind.h>

#include "aare/Fit.hpp"
#include "aare/FitModel.hpp"
#include "aare/Models.hpp"

namespace py = pybind11;
using namespace pybind11::literals;

template <typename Model>
py::object
fit_dispatch(const aare::FitModel<Model> &model,
             py::array_t<double, py::array::c_style | py::array::forcecast> x,
             py::array_t<double, py::array::c_style | py::array::forcecast> y,
             py::object y_err_obj, int n_threads);

template <typename Model> void bind_fit_model(py::module &m, const char *name) {
    using FM = aare::FitModel<Model>;
    py::class_<FM>(m, name)
        .def(py::init<unsigned int, unsigned int, double, bool>(),
             py::arg("strategy") = 0, py::arg("max_calls") = 100,
             py::arg("tolerance") = 0.5, py::arg("compute_errors") = false)
        .def("SetParLimits",
             py::overload_cast<unsigned int, double, double>(&FM::SetParLimits),
             py::arg("idx"), py::arg("lo"), py::arg("hi"))
        .def("SetParLimits",
             py::overload_cast<const std::string &, double, double>(
                 &FM::SetParLimits),
             py::arg("idx"), py::arg("lo"), py::arg("hi"))
        .def("FixParameter",
             py::overload_cast<unsigned int, double>(&FM::FixParameter),
             py::arg("idx"), py::arg("val"))
        .def("FixParameter",
             py::overload_cast<const std::string &, double>(&FM::FixParameter),
             py::arg("idx"), py::arg("val"))
        .def("ReleaseParameter",
             py::overload_cast<unsigned int>(&FM::ReleaseParameter),
             py::arg("idx"))
        .def("ReleaseParameter",
             py::overload_cast<const std::string &>(&FM::ReleaseParameter),
             py::arg("idx"))
        .def("SetParameter",
             py::overload_cast<unsigned int, double>(&FM::SetParameter),
             py::arg("idx"), py::arg("val"))
        .def("SetParameter",
             py::overload_cast<const std::string &, double>(&FM::SetParameter),
             py::arg("idx"), py::arg("val"))
        .def("GetParName", &FM::GetParName, py::arg("idx"))
        .def("GetParNames", &FM::GetParNames)
        .def_property_readonly("par_names", &FM::GetParNames)
        .def_property_readonly("n_par",
                               [](py::object /*cls*/) { return Model::npar; })
        .def_property("max_calls", &FM::max_calls, &FM::SetMaxCalls)
        .def_property("tolerance", &FM::tolerance, &FM::SetTolerance)
        .def_property("compute_errors", &FM::compute_errors,
                      &FM::SetComputeErrors)
        .def(
            "__call__", // conversion ok, we want to be able to call with any
                        // dtype
            [](const FM & /*self*/,
               py::array_t<double, py::array::c_style | py::array::forcecast> x,
               py::array_t<double, py::array::c_style | py::array::forcecast>
                   par) {
                auto x_view = make_view_1d(x);
                auto p_view = make_view_1d(par);

                std::vector<double> pvec(p_view.begin(), p_view.end());

                auto *result = new aare::NDArray<double, 1>({x_view.size()});
                for (ssize_t i = 0; i < x_view.size(); ++i)
                    (*result)(i) = Model::eval(x_view[i], pvec);

                return return_image_data(result);
            },
            py::arg("x"), py::arg("par"))
        .def(
            "fit", // conversion ok
            [](const FM &self,
               py::array_t<double, py::array::c_style | py::array::forcecast> x,
               py::array_t<double, py::array::c_style | py::array::forcecast> y,
               py::object y_err_obj, int n_threads) -> py::object {
                return fit_dispatch<Model>(self, x, y, y_err_obj, n_threads);
            },
            R"doc(
            Fit this model to 1D or 3D data using Minuit2.

            Parameters
            ----------
            x : array_like, shape (n_scan,)
                Scan points.
            y : array_like, shape (n_scan,) or (rows, cols, n_scan)
                Measured data.
            y_err : array_like or None
                Per-point uncertainties. None for unweighted fit.
            n_threads : int
                Number of threads for 3D parallel loop.
            )doc",
            py::arg("x"), py::arg("y"), py::arg("y_err") = py::none(),
            py::arg("n_threads") = 4);
}

template <typename Model>
py::dict pack_1d_result_dict(const aare::NDArray<double, 1> &result,
                             bool compute_errors) {
    constexpr std::size_t npar = Model::npar;

    auto res = result.view();

    auto par_out = new NDArray<double, 1>({npar}, 0.0);
    auto chi2_out = new NDArray<double, 1>({1}, 0.0);

    auto par_view = par_out->view();
    auto chi2_view = chi2_out->view();

    for (std::size_t i = 0; i < npar; ++i) {
        par_view(i) = res(i);
    }

    if (compute_errors) {
        auto err_out = new NDArray<double, 1>({npar}, 0.0);
        auto err_view = err_out->view();

        for (std::size_t i = 0; i < npar; ++i) {
            err_view(i) = res(npar + i);
        }

        chi2_view(0) = res(2 * npar);

        return py::dict("par"_a = return_image_data(par_out),
                        "par_err"_a = return_image_data(err_out),
                        "chi2"_a = return_image_data(chi2_out));
    } else {
        chi2_view(0) = res(npar);

        return py::dict("par"_a = return_image_data(par_out),
                        "chi2"_a = return_image_data(chi2_out));
    }
}

// Helper: typed dispatch for one Model, handles 1D/3D + y_err logic
template <typename Model>
py::object
fit_dispatch(const aare::FitModel<Model> &model,
             py::array_t<double, py::array::c_style | py::array::forcecast> x,
             py::array_t<double, py::array::c_style | py::array::forcecast> y,
             py::object y_err_obj, int n_threads) {
    constexpr std::size_t npar = Model::npar;

    if (y.ndim() == 3) {
        auto par_out =
            new NDArray<double, 3>({y.shape(0), y.shape(1), npar}, 0.0);
        auto chi2_out = new NDArray<double, 2>({y.shape(0), y.shape(1)}, 0.0);

        auto x_view = make_view_1d(x);
        auto y_view = make_view_3d(y);

        if (!y_err_obj.is_none()) {
            auto y_err = py::cast<
                py::array_t<double, py::array::c_style | py::array::forcecast>>(
                y_err_obj);

            auto err_out =
                new NDArray<double, 3>({y.shape(0), y.shape(1), npar}, 0.0);
            auto y_view_err = make_view_3d(y_err);

            aare::fit_3d<Model>(model, x_view, y_view, y_view_err,
                                par_out->view(), err_out->view(),
                                chi2_out->view(), n_threads);

            if (model.compute_errors()) {
                return py::dict("par"_a = return_image_data(par_out),
                                "par_err"_a = return_image_data(err_out),
                                "chi2"_a = return_image_data(chi2_out));
            } else {
                delete err_out;
                return py::dict("par"_a = return_image_data(par_out),
                                "chi2"_a = return_image_data(chi2_out));
            }
        } else {

            NDView<double, 3> dummy_err{};
            NDView<double, 3> dummy_err_out{};

            aare::fit_3d<Model>(model, x_view, y_view, dummy_err,
                                par_out->view(), dummy_err_out,
                                chi2_out->view(), n_threads);

            return py::dict("par"_a = return_image_data(par_out),
                            "chi2"_a = return_image_data(chi2_out));
        }
    } else if (y.ndim() == 1) {
        NDArray<double, 1> result{};

        auto x_view = make_view_1d(x);
        auto y_view = make_view_1d(y);

        if (!y_err_obj.is_none()) {
            auto y_err = py::cast<
                py::array_t<double, py::array::c_style | py::array::forcecast>>(
                y_err_obj);

            auto y_view_err = make_view_1d(y_err);
            result = aare::fit_pixel<Model>(model, x_view, y_view, y_view_err);
        } else {
            result = aare::fit_pixel<Model>(model, x_view, y_view);
        }

        return pack_1d_result_dict<Model>(result, model.compute_errors());

    } else {
        throw std::runtime_error("Data must be 1D or 3D.");
    }
}

void define_fit_bindings(py::module &m) {
    // ── Bind model classes ──────────────────────────────────────────
    bind_fit_model<aare::model::Gaussian>(m, "Gaussian");
    bind_fit_model<aare::model::GaussianErfcPlateau>(m, "GaussianErfcPlateau");
    bind_fit_model<aare::model::GaussianChargeSharing>(m,
                                                       "GaussianChargeSharing");
    bind_fit_model<aare::model::GaussianChargeSharingKb>(
        m, "GaussianChargeSharingKb");
    bind_fit_model<aare::model::RisingScurve>(m, "RisingScurve");
    bind_fit_model<aare::model::FallingScurve>(m, "FallingScurve");
    bind_fit_model<aare::model::Pol1>(m, "Pol1");
    bind_fit_model<aare::model::Pol2>(m, "Pol2");

    m.def(
        "fit",
        [](py::object model_obj, // conversion ok, we want to be able to call
                                 // with any dtype
           py::array_t<double, py::array::c_style | py::array::forcecast> x,
           py::array_t<double, py::array::c_style | py::array::forcecast> y,
           py::object y_err_obj, int n_threads) -> py::object {
            using namespace aare::model;

            // ── Polynomial of degree 1 ───────
            if (py::isinstance<aare::FitModel<Pol1>>(model_obj)) {
                const auto &mdl =
                    model_obj.cast<const aare::FitModel<Pol1> &>();
                return fit_dispatch<Pol1>(mdl, x, y, y_err_obj, n_threads);
            }

            // ── Polynomial of degree 2 ───────
            if (py::isinstance<aare::FitModel<Pol2>>(model_obj)) {
                const auto &mdl =
                    model_obj.cast<const aare::FitModel<Pol2> &>();
                return fit_dispatch<Pol2>(mdl, x, y, y_err_obj, n_threads);
            }

            // ── Gaussian ───────
            if (py::isinstance<aare::FitModel<Gaussian>>(model_obj)) {
                const auto &mdl =
                    model_obj.cast<const aare::FitModel<Gaussian> &>();
                return fit_dispatch<Gaussian>(mdl, x, y, y_err_obj, n_threads);
            }

            // ── GaussianErfcPlateau ───────
            if (py::isinstance<aare::FitModel<GaussianErfcPlateau>>(
                    model_obj)) {
                const auto &mdl =
                    model_obj
                        .cast<const aare::FitModel<GaussianErfcPlateau> &>();
                return fit_dispatch<GaussianErfcPlateau>(mdl, x, y, y_err_obj,
                                                         n_threads);
            }

            // ── GaussianChargeSharing ───────
            if (py::isinstance<aare::FitModel<GaussianChargeSharing>>(
                    model_obj)) {
                const auto &mdl =
                    model_obj
                        .cast<const aare::FitModel<GaussianChargeSharing> &>();
                return fit_dispatch<GaussianChargeSharing>(mdl, x, y, y_err_obj,
                                                           n_threads);
            }

            // ── GaussianChargeSharingKb ───────
            if (py::isinstance<aare::FitModel<GaussianChargeSharingKb>>(
                    model_obj)) {
                const auto &mdl = model_obj.cast<
                    const aare::FitModel<GaussianChargeSharingKb> &>();
                return fit_dispatch<GaussianChargeSharingKb>(
                    mdl, x, y, y_err_obj, n_threads);
            }

            // ── Rising Scurve ───────
            if (py::isinstance<aare::FitModel<RisingScurve>>(model_obj)) {
                const auto &mdl =
                    model_obj.cast<const aare::FitModel<RisingScurve> &>();
                return fit_dispatch<RisingScurve>(mdl, x, y, y_err_obj,
                                                  n_threads);
            }

            // ── Falling Scurve ───────
            if (py::isinstance<aare::FitModel<FallingScurve>>(model_obj)) {
                const auto &mdl =
                    model_obj.cast<const aare::FitModel<FallingScurve> &>();
                return fit_dispatch<FallingScurve>(mdl, x, y, y_err_obj,
                                                   n_threads);
            }

            throw std::runtime_error(
                "Unknown model type. Expected Pol1, Pol2, Gaussian, "
                "RisingScurve or FallingScurve.");
        },
        R"(
        Fit a model to 1D or 3D data using Minuit2.
 
        Parameters
        ----------
        model : Pol1, Pol2, Gaussian, RisingScurve, or FallingScurve
            Configured model object.  User-set limits, fixed parameters,
            and start values take precedence over automatic estimates.
        x : array_like, shape (n_scan,)
            Scan points (e.g. energy or threshold values).
        y : array_like, shape (n_scan,) or (rows, cols, n_scan)
            Measured data.  1D for a single pixel, 3D for a detector image.
        y_err : array_like or None
            Per-point uncertainties.  Same shape as y.  None → unweighted fit.
        n_threads : int
            Number of threads for the 3D parallel loop.
 
        Returns
        -------
        For 1D input:
            numpy array of shape (2*npar+1,) if compute_errors else (npar+1,).
            Layout: [params..., (errors...,) chi2].
 
        For 3D input:
            dict with keys:
              "par"     : (rows, cols, npar) fitted parameters.
              "par_err" : (rows, cols, npar) parameter errors (if compute_errors).
              "chi2"    : (rows, cols)       chi-squared per pixel.
        )",
        py::arg("model"), py::arg("x"), py::arg("y"),
        py::arg("y_err") = py::none(), py::arg("n_threads") = 4);
}