#include "profile_providers.h"
#include <cmath>

namespace{
    template<typename T>
    static T safe_div(const T el1, const T el2)
    {
        if(el1 == 0 && el2 == 0) return 0.0;
        return el1 / el2;
    }

    template<typename T>
    mosca::image weight_image(const mosca::image& slit_image,
            const mosca::image& slit_image_weight,
            T& total_flux, T& total_weight){

        mosca::image slit_weighted = slit_image;
        std::transform (slit_image.get_data<T>(),
                      slit_image.get_data<T>() + slit_image.size_x() * slit_image.size_y(),
                      slit_image_weight.get_data<T>(),
                      slit_weighted.get_data<T>(), std::multiplies<T>());

        T * p_ima = slit_weighted.get_data<T>();
        total_flux =
           std::accumulate(p_ima, p_ima + slit_image.size_x() * slit_image.size_y(), T(0));
        const T * p_weight = slit_image_weight.get_data<T>();
        total_weight =
           std::accumulate(p_weight, p_weight + slit_image.size_x() * slit_image.size_y(), T(0));
        return slit_weighted;
    }
}

mosca::profile_smoother::profile_smoother(const int smooth_radius) : m_smooth_radius(smooth_radius){}

bool mosca::profile_smoother::is_enabled() const {return m_smooth_radius > 0;}

template<typename T>
void mosca::profile_smoother::smooth(std::vector<T>& profile, std::vector<T>& weight_profile) const{

    if (!is_enabled()) return;

    std::vector<bool> mask;
    std::transform(weight_profile.begin(), weight_profile.end(),
                   std::back_inserter(mask), std::bind1st(std::not_equal_to<T>(), T(0)));

    const int available_pixels = (int)std::count(mask.cbegin(), mask.cend(), true);
    int smooth_radius = m_smooth_radius;
    if (available_pixels / 2 < smooth_radius)
        smooth_radius = available_pixels / 2;

    mosca::vector_smooth<T>(profile, mask, smooth_radius);
}


mosca::profile_spatial_fitter::profile_spatial_fitter(const int fit_polyorder, const double fit_threshold)
    : m_fit_polyorder(fit_polyorder),
      m_fit_threshold(fit_threshold){}

bool mosca::profile_spatial_fitter::is_enabled()const {return m_fit_polyorder > 0;}

template<typename T>
void mosca::profile_spatial_fitter::fit(std::vector<T>& profile, std::vector<T>& weight_profile) const{

    if (!is_enabled() ) return;

    std::vector<bool> mask;
    const double max_el = *std::max_element(profile.begin(), profile.end());
    const double th_this = m_fit_threshold * max_el;
    std::transform(profile.begin(), profile.end(),
                   std::back_inserter(mask), std::bind2nd(std::greater_equal<T>(), th_this));

    size_t used_spa_fit_polyorder = m_fit_polyorder;
    mosca::vector_polynomial polfit;
    polfit.fit<T>(profile, mask, used_spa_fit_polyorder);

}

mosca::profile_dispersion_fitter::profile_dispersion_fitter(const int disp_fit_nknots,
        const double fit_threshold)
        : m_disp_fit_nknots(disp_fit_nknots),
          m_fit_threshold(fit_threshold){}

bool mosca::profile_dispersion_fitter::is_enabled()  const {return m_disp_fit_nknots > 0;}

template<typename T>
void mosca::profile_dispersion_fitter::fit(std::vector<T>& profile, std::vector<T>& weight_profile) const{

    if (!is_enabled() ) return;

      std::vector<bool> mask;
      std::transform(weight_profile.begin(), weight_profile.end(),
                     std::back_inserter(mask), std::bind1st(std::not_equal_to<T>(), T(0)));
      size_t used_disp_fit_nknots = m_disp_fit_nknots;
      mosca::vector_cubicspline splfit;
      splfit.fit<T>(profile, mask, used_disp_fit_nknots);
}

template<typename T>
template<typename profile_smootherType, typename FitterType>
mosca::profile_provider_base<T>::profile_provider_base(const mosca::image& slit_image,
        const mosca::image& slit_image_weight,
        profile_smootherType profile_smoother,
        FitterType fitter,
        const mosca::axis profile_axis,
        const mosca::axis collapse_axis)
        : m_profile_axis(profile_axis),
          m_collapse_axis(collapse_axis){

    mosca::image slit_weighted = weight_image(slit_image, slit_image_weight,
            m_total_flux, m_total_weight);
    //Collapsing the data to get the profiles in each direction
    std::vector<T> slit_profile =
            slit_weighted.collapse<T>(m_collapse_axis);

    if(m_total_flux == T(0) || m_total_weight == T(0))
    {
        m_total_flux = m_total_weight = T(1);
        m_slit_norm_profile.resize(slit_profile.size());
        return;
    }

    //Collapsing the weights
    std::vector<T> weight_profile =
            slit_image_weight.collapse<T>(m_collapse_axis);

    //Getting the profiles properly weighted
    std::vector<T> slit_profile_w;
    std::transform (slit_profile.begin(), slit_profile.end(),
                    weight_profile.begin(), std::back_inserter(slit_profile_w),
                    safe_div<T>);

    //If we are doing any fitting/smoothing in that direction,
    //initialise it to the current profile, if not initialise it to a constant
    if (profile_smoother.is_enabled() || fitter.is_enabled())
        m_slit_norm_profile = slit_profile_w;
    else
        m_slit_norm_profile = std::vector<T>(slit_profile_w.size(),
                T(m_total_flux / m_total_weight));

    profile_smoother.smooth(m_slit_norm_profile, weight_profile);
    fitter.fit(m_slit_norm_profile, weight_profile);
}

template<typename T>
T mosca::profile_provider_base<T>::get(size_t x, size_t y) const noexcept{

    size_t coord = x;
    if(m_profile_axis == mosca::Y_AXIS)
        coord = y;

    const T norm_factor = get_total_flux() / get_total_weight();
    const T val = m_slit_norm_profile[coord] /  sqrt(norm_factor);
    if(std::isnan(val))
        return T(0);
    return val;
}

template<typename T>
T mosca::profile_provider_base<T>::get_total_flux() const noexcept{
    return m_total_flux;
}

template<typename T>
T mosca::profile_provider_base<T>::get_total_weight() const noexcept{
    return m_total_weight;
}

template<typename T>
const std::vector<T>& mosca::profile_provider_base<T>::get_average_linear_profile()  const noexcept{
    return m_slit_norm_profile;
}

template<typename T>
mosca::profile_provider_base<T>::~profile_provider_base() = default;



template<typename T>
mosca::dispersion_profile_provider<T>::dispersion_profile_provider(const mosca::image& slit_image,
                            const mosca::image& slit_image_weight,
                            const int disp_smooth_radius,
                            const int disp_fit_nknots,
                            const double fit_threshold) :

                        profile_provider_base<T>(slit_image, slit_image_weight,
                        profile_smoother{disp_smooth_radius}, profile_dispersion_fitter{disp_fit_nknots, fit_threshold},
                        slit_image.dispersion_axis(),
                        mosca::SPATIAL_AXIS){}



template<typename T>
mosca::spatial_profile_provider<T>::spatial_profile_provider(const mosca::image& slit_image,
                            const mosca::image& slit_image_weight,
                            const int spa_smooth_radius,
                            const int spa_fit_polyorder,
                            const double fit_threshold) :

                profile_provider_base<T>(slit_image, slit_image_weight,
                        profile_smoother{spa_smooth_radius},
                        profile_spatial_fitter{spa_fit_polyorder, fit_threshold},
                        slit_image.spatial_axis(),
                        mosca::DISPERSION_AXIS){}


static mosca::image extract_spatial_line(const mosca::image& img, const cpl_size spa_line){
    return img.trim(spa_line, 1, spa_line, img.size_spatial());
}

template<typename T>
static void sum_vecs(std::vector<T>& dest, const std::vector<T>& source){
    std::transform(dest.cbegin(), dest.cend(),
                   source.cbegin(),
                   dest.begin(),
                   std::plus<T>());
}

template<typename T>
mosca::local_spatial_profile_provider<T>::local_spatial_profile_provider(
                            const mosca::image& slit_image,
                            const mosca::image& slit_image_weight,
                            const int spa_smooth_radius,
                            const int spa_fit_nknots,
                            const double spa_fit_threshold) :
                            m_dispersion_profile_axis(slit_image.dispersion_axis()),
                            m_average_profile(slit_image.size_dispersion(), T(0)),
                            m_avg_flux_image(0){

    const cpl_size disp_size = slit_image.size_dispersion();

    for(cpl_size spa_line  = 1; spa_line <= disp_size; ++spa_line   ){
        auto cropped_img = extract_spatial_line(slit_image, spa_line);
        auto cropped_weight_img = extract_spatial_line(slit_image_weight, spa_line);
        m_providers.emplace_back(cropped_img, cropped_weight_img, spa_smooth_radius, spa_fit_nknots, spa_fit_threshold);
        sum_vecs(m_average_profile, m_providers.crbegin()->get_average_linear_profile());
    }

    if(m_providers.empty()) return;

    T num_profiles = static_cast<T>(m_providers.size());
    auto op = [num_profiles](const T v){return v / num_profiles;};
    std::transform(m_average_profile.cbegin(), m_average_profile.cend(),
                   m_average_profile.begin(), op);


    T total_flux{0};
    T total_weight{0};
    weight_image(slit_image, slit_image_weight,
            total_flux, total_weight);

    m_avg_flux_image = total_flux / total_weight;
}

template<typename T>
T mosca::local_spatial_profile_provider<T>::get(size_t x, size_t y) const noexcept{

    if(m_providers.empty()) return T(0);

    cpl_size provider_idx = x;
    if(m_dispersion_profile_axis == mosca::Y_AXIS)
        provider_idx = y;

    const T avg_flux_line = m_providers[provider_idx].get_total_flux() / m_providers[provider_idx].get_total_weight();
    const T factor = sqrt(m_avg_flux_image / avg_flux_line);
    return m_providers[provider_idx].get(x, y) * factor;
}

template<typename T>
const std::vector<T>& mosca::local_spatial_profile_provider<T>::get_average_linear_profile() const noexcept{

    return m_average_profile;
}


