# SPDX-License-Identifier: BSD-2-Clause
#
# Copyright (C) 2023, Raspberry Pi Ltd
#
# generate_filter.py - Resample filter kernel generation
#
# Refer to the PiSP specification, for the register details

import numpy as np
import argparse


# Lanczos (order, length)
# Order is typically 2 or 3
def lanczos(order, N):
    x = np.linspace(-order, +order, N + 6, dtype=np.float64)
    h = np.where((x > -order) & (x < order), np.sinc(x) * np.sinc(x / order), 0)
    return h[3:-3]


# Mitchell - Netravali filters (B, C, length)
# B,C = (1,0) is the cubic B-spline
# B,C = (1/3,1/3) is a good compromise
# B,C = (0, 0.5) for Catmull-Rom
def mitchell(B, C, N):
    x = np.linspace(-2, +2, N, dtype=np.float64)
    h = np.zeros(N)
    for i in range(N):
        ax = abs(x[i])
        if ax < 1:
            h[i] = (
                (12 - 9 * B - 6 * C) * ax**3
                + (-18 + 12 * B + 6 * C) * ax**2
                + (6 - 2 * B)
            ) / 6
        elif (ax >= 1) and (ax < 2):
            h[i] = (
                (-B - 6 * C) * ax**3
                + (6 * B + 30 * C) * ax**2
                + (-12 * B - 48 * C) * ax
                + (8 * B + 24 * C)
            ) / 6
    return h


# bicubic_spline(alpha, length);
# alpha is typically set to -0.5 (Hermite spline) or -0.75
def bicubic_spline(a, N):
    x = np.linspace(-2, +2, N, dtype=np.float64)
    h = np.zeros(N)
    for i in range(N):
        ax = abs(x[i])
        if ax <= 1:
            h[i] = (a + 2) * ax**3 - (a + 3) * ax**2 + 1
        elif (ax > 1) and (ax < 2):
            h[i] = (a) * ax**3 - (5 * a) * ax**2 + (8 * a) * ax - (4 * a)
    return h


def main():
    parser = argparse.ArgumentParser(formatter_class=argparse.RawTextHelpFormatter)
    parser.add_argument(
        "--phases", metavar="P", type=int, help="Number of phases.", default=16
    )
    parser.add_argument(
        "--taps",
        metavar="T",
        type=int,
        help="Number of filter taps per phase.",
        default=6,
    )
    parser.add_argument(
        "--precision",
        metavar="PR",
        type=int,
        help="Filter precision required.",
        default=10,
    )
    parser.add_argument(
        "--filter",
        type=str,
        metavar="F",
        help="Filter type and parameters, e.g.: \n"
        '"Mitchell, b = 0.333, c = 0.333"\n'
        '"Lanczos, order = 3"\n'
        '"bicubic_spline, a=-0.5"',
        required=True,
    )

    args = parser.parse_args()

    phases = args.phases
    taps = args.taps
    precision = args.precision

    # Parse the filter string and pick out the needed parameters.
    filt = args.filter.split(",")
    params = {"a": 0.0, "b": 0.0, "c": 0.0, "order": 0}
    for param in filt[1:]:
        p = param.replace(" ", "").split("=")
        params[p[0]] = type(params[p[0]])(p[1])

    # Generate the filter.
    if filt[0].lower() == "mitchell":
        filter = (
            f'"Michell - Netravali (B = {params["b"]:.3f}, C = {params["c"]:.3f})": [\n'
        )
        h = mitchell(params["b"], params["c"], phases * taps)
    elif filt[0].lower() == "lanczos":
        filter = f'"Lanczos order {params["order"]}": [\n'
        h = lanczos(params["order"], phases * taps)
    elif filt[0].lower() == "bicubic_spline":
        filter = f'"Bicubic-spline (a = {params["a"]:.3f})": [\n'
        h = bicubic_spline(params["a"], phases * taps)
    else:
        print(f"Invalid filter ({filt[0]}) selected!")
        exit()

    # Normalise and convert to fixed-point.
    h = h * phases / np.sum(h)
    h = np.rint(h * (1 << precision)).astype(np.int32)

    ppf = np.zeros((phases, taps), dtype=np.int32)
    for i in range(phases):
        # Pick out phases and flip array.
        ppf[i] = (h[i::phases])[::-1]
        # Make sure there is no DC change by adjusting the largest coefficients.
        max_index = np.nonzero(ppf[i] == ppf[i].max())[0]
        ppf[i, max_index] += (1 << precision) - np.int32(ppf[i].sum() / max_index.size)

    nl = "\n"
    for i in range(phases):
        phase = ", ".join([f"{c:>4}" for c in ppf[i, :]])
        filter += f"    {phase}{nl + ']' if i == phases - 1 else ',' + nl}"

    print(filter)


if __name__ == "__main__":
    main()
