#include <galprop/Particle.h>
#include "galpropy.hpp"


class PyParticle : public Particle {
public:
    using Particle::Particle;

    bool depends_on(const Particle& p, Galprop& g) const override {
        PYBIND11_OVERRIDE_PURE(bool, Particle, depends_on, p, g);
    }
    void fill_transport_arrays(Galprop& g) override {
        PYBIND11_OVERRIDE(void, Particle, fill_transport_arrays, g);
    }
    void fill_source_functions(Galprop& g) override {
        PYBIND11_OVERRIDE_PURE(void, Particle, fill_source_functions, g);
    }
    void fill_diffusion_coeffs(Galprop& g) override {
        PYBIND11_OVERRIDE(void, Particle, fill_diffusion_coeffs, g);
    }
    void fill_energy_loss_rate(Galprop& g) override {
        PYBIND11_OVERRIDE(void, Particle, fill_energy_loss_rate, g);
    }
    void fill_destruction_rate(Galprop& g) override {
        PYBIND11_OVERRIDE(void, Particle, fill_destruction_rate, g);
    }
    void propagate(Galprop& g) override {
        PYBIND11_OVERRIDE(void, Particle, propagate, g);
    }
    double abundance(void) const override {
        PYBIND11_OVERRIDE(double, Particle, abundance);
    }
};

void bind_Particle(py::module_ &m)
{
    py::bind_vector<std::vector<std::shared_ptr<Particle>>>(
        m, "VecPtrParticle");
    py::implicitly_convertible<
        py::list, std::vector<std::shared_ptr<Particle>>>();
    py::implicitly_convertible<
        py::tuple, std::vector<std::shared_ptr<Particle>>>();

    py::bind_map<std::map<std::string, std::shared_ptr<Particle>>>(
        m, "MapPtrParticle");

    py::class_<Particle, PyParticle, std::shared_ptr<Particle>>(m, "Particle")
        .def(py::init<std::shared_ptr<const ParticleInfo>,
                      std::shared_ptr<const Galaxy>,
                      const std::string &>())
        .def("info", &Particle::info)
        .def("galaxy", &Particle::galaxy)
        .def("dims", &Particle::dims)
        .def("Z", &Particle::Z)
        .def("A", &Particle::A)
        .def("A_pos", &Particle::A_pos)
        .def("is_lepton", &Particle::is_lepton)
        .def("is_isotope", &Particle::is_isotope)
        .def("depends_on", &Particle::depends_on)
        .def("fill_transport_arrays", &Particle::fill_transport_arrays)
        .def("fill_source_functions", &Particle::fill_source_functions)
        .def("fill_diffusion_coeffs", &Particle::fill_diffusion_coeffs)
        .def("fill_energy_loss_rate", &Particle::fill_energy_loss_rate)
        .def("fill_destruction_rate", &Particle::fill_destruction_rate)
        .def("propagate", &Particle::propagate)
        .def("abundance", &Particle::abundance)
        .def("local_flux", py::overload_cast<double, double>(
             &Particle::local_flux, py::const_),
             "2-D local flux", "r"_a=Rsun, "z"_a=0)
        .def("local_flux", py::overload_cast<double, double, double>(
             &Particle::local_flux, py::const_),
             "3-D local flux", "x"_a, "y"_a, "z"_a)
        .def("r", py::overload_cast<>(&Particle::r, py::const_))
        .def("r", py::overload_cast<size_t>(&Particle::r, py::const_))
        .def("x", py::overload_cast<>(&Particle::x, py::const_))
        .def("x", py::overload_cast<size_t>(&Particle::x, py::const_))
        .def("y", py::overload_cast<>(&Particle::y, py::const_))
        .def("y", py::overload_cast<size_t>(&Particle::y, py::const_))
        .def("z", py::overload_cast<>(&Particle::z, py::const_))
        .def("z", py::overload_cast<size_t>(&Particle::z, py::const_))
        .def("Ekin", py::overload_cast<>(&Particle::Ekin, py::const_))
        .def("Ekin", py::overload_cast<size_t>(&Particle::Ekin, py::const_))
        .def("Etot", py::overload_cast<>(&Particle::Etot, py::const_))
        .def("Etot", py::overload_cast<size_t>(&Particle::Etot, py::const_))
        .def("p", py::overload_cast<>(&Particle::p, py::const_))
        .def("p", py::overload_cast<size_t>(&Particle::p, py::const_))
        .def("rigidity", py::overload_cast<>(&Particle::rigidity, py::const_))
        .def("rigidity", py::overload_cast<size_t>(&Particle::rigidity, py::const_))
        .def("gamma", py::overload_cast<>(&Particle::gamma, py::const_))
        .def("gamma", py::overload_cast<size_t>(&Particle::gamma, py::const_))
        .def("beta", py::overload_cast<>(&Particle::beta, py::const_))
        .def("beta", py::overload_cast<size_t>(&Particle::beta, py::const_))
        .def("Egamma", py::overload_cast<>(&Particle::Egamma, py::const_))
        .def("Egamma", py::overload_cast<size_t>(&Particle::Egamma, py::const_))
        .def("Eradio", py::overload_cast<>(&Particle::Eradio, py::const_))
        .def("Eradio", py::overload_cast<size_t>(&Particle::Eradio, py::const_))
        .def("Ekin_GeV", &Particle::Ekin_GeV)
        .def("t_half", &Particle::t_half)
        .def_readwrite("name", &Particle::name)
        .def_readwrite("K_electron", &Particle::K_electron)
        .def_readwrite("cr_density", &Particle::cr_density)
        .def_readwrite("source_function", &Particle::source_function)
        .def_readwrite("destruct_rate", &Particle::destruct_rate)
        .def_readwrite("Dxx", &Particle::Dxx)
        .def_readwrite("Dpp", &Particle::Dpp)
        .def_readwrite("dpdt", &Particle::dpdt)
        .def_readwrite("normalization_factor", &Particle::normalization_factor)
        .def_readwrite("primary_source_function", &Particle::primary_source_function)
        .def_readwrite("secondary_source_function", &Particle::secondary_source_function);
}