# ridge_kernel_gaussian.py
#
# Gaussian-kernel ridge regression.
#
# This script is a companion example to ridge_gaussian_basis.py.
# In this script, each input data point is used as a kernel center.
#
# Usage:
#   python ridge_kernel_gaussian.py
#   python ridge_kernel_gaussian.py input.xlsx
#   python ridge_kernel_gaussian.py input.xlsx 0.1
#   python ridge_kernel_gaussian.py input.xlsx 0.1 1.0e-6
#
# Arguments:
#   argv[1] : input Excel file
#   argv[2] : Gaussian kernel width wG
#   argv[3] : ridge regularization parameter alpha

import sys

import numpy as np
from numpy import exp
import pandas as pd
import matplotlib.pyplot as plt


# ================================================================
# Default parameters
# ================================================================
infile = 'random-poly-Gauss.xlsx'

# Gaussian kernel width
wG = 0.1

# Ridge regularization parameter
# alpha = 0.0 gives interpolation-like fitting when the matrix is solvable.
# A small positive value is usually more stable.
alpha = 1.0e-8

xcal0 = None
xcal1 = None
ncal = 101

fontsize = 16


# ================================================================
# Read command-line arguments
# ================================================================
argv = sys.argv
narg = len(argv)

if narg >= 2:
    infile = argv[1]

if narg >= 3:
    wG = float(argv[2])

if narg >= 4:
    alpha = float(argv[3])


# ================================================================
# Kernel function
# ================================================================
def kernel(xi, xj):
    """
    Gaussian kernel function.

    Parameters
    ----------
    xi, xj : float
        Input coordinates.

    Returns
    -------
    float
        Gaussian kernel value

            exp(-((xi - xj) / wG)**2)
    """
    a = (xi - xj) / wG
    return exp(-a**2)


# ================================================================
# Kernel ridge regression
# ================================================================
def KernelRidge(x, y, alpha=0.0):
    """
    Perform Gaussian-kernel ridge regression.

    This function constructs the Gaussian kernel matrix K for the
    input data points and solves the regularized linear system

        (K + alpha I) c = y

    to obtain the kernel expansion coefficients.

    Parameters
    ----------
    x : array-like
        Input coordinates.
    y : array-like
        Target values.
    alpha : float, optional
        Ridge regularization parameter.

    Returns
    -------
    list of float
        Kernel expansion coefficients.
    """
    n = len(x)

    Si = np.empty([n, 1])
    Sij = np.empty([n, n])

    # Right-hand-side vector
    for i in range(0, n):
        Si[i, 0] = y[i]

    # Kernel matrix
    for j in range(0, n):
        for l in range(j, n):
            Sij[j, l] = Sij[l, j] = kernel(x[j], x[l])

        # Ridge regularization term
        Sij[j, j] += alpha

    print("Vector and Matrix:")
    print("Si=")
    print(Si)
    print("Sij=")
    print(Sij)
    print("")

    # Solve the linear system:
    #
    #     Sij @ ci = Si
    #
    # This is numerically better than explicitly calculating inv(Sij).
    ci = np.linalg.solve(Sij, Si)

    ci = ci.transpose().tolist()
    return ci[0]


# ================================================================
# Main routine
# ================================================================
def main():
    global wG

    print("Kernel ridge regression with Gaussian kernel")
    print(f"infile={infile}")
    print(f"Gaussian kernel width={wG}")
    print(f"Ridge alpha={alpha}")

    print("")
    print(f"Read [{infile}]")

    df = pd.read_excel(infile, engine='openpyxl')
    labels = df.columns.to_list()

    # Convert pandas Series to NumPy arrays.
    # This makes x[i] and y[i] simple array indexing.
    x = df[labels[0]].to_numpy(dtype=float)
    y = df[labels[1]].to_numpy(dtype=float)

    ndata = len(x)

    xcal0 = min(x)
    xcal1 = max(x)
    xcalstep = (xcal1 - xcal0) / (ncal - 1)

    print("")
    print("Execute kernel ridge regression")
    ci = KernelRidge(x, y, alpha)

    print("Expansion coefficients:")
    for i in range(ndata):
        print(f" x0[{i}]={x[i]:6.3g}: c[{i}]={ci[i]:6.3g}")

    # Calculate fitted curve
    xcal = [xcal0 + i * xcalstep for i in range(ncal)]
    ycal = []

    for i in range(ncal):
        _x = xcal[i]
        yl = 0.0

        for k in range(ndata):
            yl += ci[k] * kernel(_x, x[k])

        ycal.append(yl)

    # ================================================================
    # Plot
    # ================================================================
    fig, axes = plt.subplots(1, 2, figsize=(8, 6))

    axes[0].plot(x, y, label='input', linestyle='', marker='o')
    axes[0].plot(xcal, ycal, label='fit', linestyle='-')
    axes[0].tick_params(labelsize=fontsize)
    axes[0].set_xlabel('$x$', fontsize=fontsize)
    axes[0].set_ylabel('$y$', fontsize=fontsize)
    axes[0].legend(fontsize=fontsize)

    axes[1].plot(range(ndata), ci, label='coeff', linestyle='-', linewidth=0.5, marker='o')
    axes[1].tick_params(labelsize=fontsize)
    axes[1].set_xlabel('$i$', fontsize=fontsize)
    axes[1].set_ylabel('$c_i$', fontsize=fontsize)

    plt.tight_layout()
    plt.pause(0.1)

    print("")
    print("Press ENTER to terminate")
    input()


# ================================================================
# Entry point
# ================================================================
if __name__ == "__main__":
    main()