/********************************************************
 *  ██████╗  ██████╗████████╗██╗
 * ██╔════╝ ██╔════╝╚══██╔══╝██║
 * ██║  ███╗██║        ██║   ██║
 * ██║   ██║██║        ██║   ██║
 * ╚██████╔╝╚██████╗   ██║   ███████╗
 *  ╚═════╝  ╚═════╝   ╚═╝   ╚══════╝
 * Geophysical Computational Tools & Library (GCTL)
 *
 * Copyright (c) 2023  Yi Zhang (yizhang-geo@zju.edu.cn)
 *
 * GCTL is distributed under a dual licensing scheme. You can redistribute 
 * it and/or modify it under the terms of the GNU Lesser General Public 
 * License as published by the Free Software Foundation, either version 2 
 * of the License, or (at your option) any later version. You should have 
 * received a copy of the GNU Lesser General Public License along with this 
 * program. If not, see <http://www.gnu.org/licenses/>.
 * 
 * If the terms and conditions of the LGPL v.2. would prevent you from using 
 * the GCTL, please consider the option to obtain a commercial license for a 
 * fee. These licenses are offered by the GCTL's original author. As a rule, 
 * licenses are provided "as-is", unlimited in time for a one time fee. Please 
 * send corresponding requests to: yizhang-geo@zju.edu.cn. Please do not forget 
 * to include some description of your company and the realm of its activities. 
 * Also add information on how to contact you by electronic and paper mail.
 ******************************************************/

#include "kde.h"

gctl::kde::kde(){}

gctl::kde::~kde(){}

gctl::kde::kde(double h, const array<double> &x)
{
    init(h, x);
}

void gctl::kde::init(double h, const array<double> &x)
{
    if (h <= 0) throw std::runtime_error("[gctl::kde2d] Invalid averaging width.");
    if (x.size() < 2) throw std::runtime_error("[gctl::kde2d] Invalid sample size.");

    h_ = h;
    x_ = x;
    xs_ = x.size();
    return;
}

void gctl::kde::get_distribution(const array<double> &m, array<double> &d, 
    kde_kernel_e k_type)
{
    double out;
    int ms = m.size();
    d.resize(xs_);

    if (k_type == KDE_Gaussian)
    {
        for (size_t i = 0; i < xs_; i++)
        {
            out = 0;
            for (size_t j = 0; j < ms; j++)
            {
                out += gaussian_kernel((x_[i] - m[j])/h_);
            } 
            
            d[i] = out/(h_*ms);
        }
    }
    else if (k_type == KDE_Epanechnikov)
    {
        for (size_t i = 0; i < xs_; i++)
        {
            out = 0;
            for (size_t j = 0; j < ms; j++)
            {
                out += epanechnikov_kernel((x_[i] - m[j])/h_);
            } 
            
            d[i] = out/(h_*ms);
        }
    }
    else if (k_type == KDE_Rectangular)
    {
        for (size_t i = 0; i < xs_; i++)
        {
            out = 0;
            for (size_t j = 0; j < ms; j++)
            {
                out += rectangular_kernel((x_[i] - m[j])/h_);
            } 
            
            d[i] = out/(h_*ms);
        }
    }
    else
    {
        for (size_t i = 0; i < xs_; i++)
        {
            out = 0;
            for (size_t j = 0; j < ms; j++)
            {
                out += triangular_kernel((x_[i] - m[j])/h_);
            } 
            
            d[i] = out/(h_*ms);
        }
    }
    return;
}

void gctl::kde::get_gradient_at(size_t m_id, const array<double> &m, 
    array<double> &dm, kde_kernel_e k_type)
{
    dm.resize(xs_);
    int ms = m.size();

    if (k_type == KDE_Gaussian)
    {
        for (size_t i = 0; i < xs_; i++)
        {
            dm[i] = ((x_[i] - m[m_id])/h_)*gaussian_kernel((x_[i] - m[m_id])/h_)/(h_*h_*ms);
        }
    }
    else if (k_type == KDE_Epanechnikov)
    {
        for (size_t i = 0; i < xs_; i++)
        {
            dm[i] = -1.0*epanechnikov_kernel((x_[i] - m[m_id])/h_, true)/(h_*h_*ms);
        }
    }
    else if (k_type == KDE_Rectangular)
    {
        for (size_t i = 0; i < xs_; i++)
        {
            dm[i] = -1.0*rectangular_kernel((x_[i] - m[m_id])/h_, true)/(h_*h_*ms);
        }
    }
    else
    {
        for (size_t i = 0; i < xs_; i++)
        {
            dm[i] = -1.0*triangular_kernel((x_[i] - m[m_id])/h_, true)/(h_*h_*ms);
        }
    }
    return;
}

double gctl::kde::gaussian_kernel(double x)
{
    return exp(-0.5*x*x)/sqrt(2*M_PI);
}

double gctl::kde::epanechnikov_kernel(double x, bool gradient)
{
    if (gradient)
    {
        if (fabs(x) >= 1) return 0;
        else return -1.5*x;
    }

    if (fabs(x) >= 1) return 0;
    else return 0.75*(1 - x*x);
}

double gctl::kde::rectangular_kernel(double x, bool gradient)
{
    if (gradient) return 0;

    if (fabs(x) >= 1) return 0;
    else return 0.5;
}

double gctl::kde::triangular_kernel(double x, bool gradient)
{
    if (gradient)
    {
        if (fabs(x) >= 1) return 0;
        else if (x >= 0) return -1.0;
        else return 1.0;
    }

    if (fabs(x) >= 1) return 0;
    else return (1 - fabs(x));
}

gctl::kde2d::kde2d(){}
    
gctl::kde2d::~kde2d(){}

gctl::kde2d::kde2d(double hx, double hy, const array<double> &x, const array<double> &y)
{
    init(hx, hy, x, y);
}

void gctl::kde2d::init(double hx, double hy, const array<double> &x, const array<double> &y)
{
    if (hx <= 0 || hy <= 0) throw std::runtime_error("[gctl::kde2d] Invalid averaging width.");
    if (x.size() < 2 || y.size() < 2) throw std::runtime_error("[gctl::kde2d] Invalid sample size.");

    hx_ = hx;
    hy_ = hy;
    xs_ = x.size();
    ys_ = y.size();
    x_ = x;
    y_ = y;
    return;
}

void gctl::kde2d::get_distribution(const array<double> &mx, 
    const array<double> &my, array<double> &dxy, kde_kernel_e k_type)
{
    if (mx.size() != my.size()) throw std::runtime_error("[gctl::kde2d] Invalid evaluating size.");;
    int ms = mx.size();

    dxy.resize(xs_*ys_);

    double out;
    if (k_type == KDE_Gaussian)
    {
        for (size_t i = 0; i < ys_; i++)
        {
            for (size_t j = 0; j < xs_; j++)
            {
                out = 0.0;
                for (size_t k = 0; k < ms; k++)
                {
                    out += gaussian_kernel((x_[j] - mx[k])/hx_, (y_[i] - my[k])/hy_);
                }
                
                dxy[i*xs_ + j] = out/(hx_*hy_*ms);
            }
        }
    }
    else throw std::runtime_error("[gctl::kde2d] Invalid kernel type.");
    return;
}

void gctl::kde2d::get_gradient_x_at(size_t m_id, const array<double> &mx, 
    const array<double> &my, array<double> &dmx, kde_kernel_e k_type)
{
    if (mx.size() != my.size()) throw std::runtime_error("[gctl::kde2d] Invalid evaluating size.");;
    int ms = mx.size();

    dmx.resize(xs_*ys_);

    if (k_type == KDE_Gaussian)
    {
        for (size_t i = 0; i < ys_; i++)
        {
            for (size_t j = 0; j < xs_; j++)
            {
                dmx[i*xs_ + j] = ((x_[j] - mx[m_id])/hx_)*gaussian_kernel((x_[j] - mx[m_id])/hx_, (y_[i] - my[m_id])/hy_)/(hx_*hx_*hy_*ms);
            }
        }
    }
    return;
}

void gctl::kde2d::get_gradient_y_at(size_t m_id, const array<double> &mx, 
    const array<double> &my, array<double> &dmy, kde_kernel_e k_type)
{
    if (mx.size() != my.size()) throw std::runtime_error("[gctl::kde2d] Invalid evaluating size.");;
    int ms = mx.size();

    dmy.resize(xs_*ys_);

    if (k_type == KDE_Gaussian)
    {
        for (size_t i = 0; i < ys_; i++)
        {
            for (size_t j = 0; j < xs_; j++)
            {
                dmy[i*xs_ + j] = ((y_[i] - my[m_id])/hy_)*gaussian_kernel((x_[j] - mx[m_id])/hx_, (y_[i] - my[m_id])/hy_)/(hy_*hy_*hx_*ms);
            }
        }
    }
    return;
}

double gctl::kde2d::gaussian_kernel(double x, double y)
{
    return exp(-0.5*(x*x + y*y))/(2*M_PI);
}