// SPDX-License-Identifier: BSD-2-Clause

// This code is part of the sfizz library and is licensed under a BSD 2-clause
// license. You should have receive a LICENSE.md file along with the code.
// If not, contact the sfizz maintainers at https://github.com/sfztools/sfizz

#include <hiir/PolyphaseIir2Designer.h>
#include <vector>
#include <memory>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <cmath>

using FD = hiir::PolyphaseIir2Designer;

struct Stage {
    int factor;
    double tbw;
    int nbr_coefs;
    std::unique_ptr<double[]> coefs;
};

static std::vector<Stage> calculate_stages(int oversampling, double attenuation, double transition);
static void generate_cpp_prologue(int argc, char *argv[]);
static void generate_cpp_epilogue();
static void generate_cpp_coefs(const Stage *stages, int num_stages);
static void generate_cpp_upsampler(const Stage *stages, int num_stages);
static void generate_cpp_downsampler(const Stage *stages, int num_stages);

int main(int argc, char *argv[])
{
    double attenuation = 0.0;
    double transition = 0.0;
    int oversampling = 16;
    bool have_a = false;
    bool have_t = false;

    for (int argi = 1; argi < argc; ++argi) {
        const char *arg = argv[argi];
        if (!strcmp(arg, "-a")) {
            if (++argi >= argc) {
                fprintf(stderr, "The option %s expects a value.\n", arg);
                return 1;
            }
            arg = argv[argi];
            attenuation = atof(arg);
            have_a = true;
        }
        else if (!strcmp(arg, "-t")) {
            if (++argi >= argc) {
                fprintf(stderr, "The option %s expects a value.\n", arg);
                return 1;
            }
            arg = argv[argi];
            transition = atof(arg);
            have_t = true;
        }
        else if (!strcmp(arg, "-o")) {
            if (++argi >= argc) {
                fprintf(stderr, "The option %s expects a value.\n", arg);
                return 1;
            }
            arg = argv[argi];
            oversampling = atoi(arg);
        }
        else {
            fprintf(stderr, "Unrecognized argument: %s\n", arg);
            return 1;
        }
    }

    if (!have_a) {
        fprintf(stderr, "No attenuation given (-a)\n");
        return 1;
    }
    if (!have_t) {
        fprintf(stderr, "No transition bandwidth given (-t)\n");
        return 1;
    }
    else if (attenuation < 0.0) {
        fprintf(stderr, "Invalid attenuation\n");
        return 1;
    }
    else if (transition <= 0.0 || transition >= 0.5) {
        fprintf(stderr, "Invalid transition bandwidth\n");
        return 1;
    }
    else if (oversampling < 2) {
        fprintf(stderr, "Invalid oversampling\n");
        return 1;
    }

    std::vector<Stage> stages = calculate_stages(oversampling, attenuation, transition);
    int num_stages = (int)stages.size();

    // generate the coeffs
    generate_cpp_prologue(argc, argv);
    printf("\n");
    generate_cpp_coefs(stages.data(), num_stages);
    printf("\n");
    generate_cpp_upsampler(stages.data(), num_stages);
    printf("\n");
    generate_cpp_downsampler(stages.data(), num_stages);
    printf("\n");
    generate_cpp_epilogue();

    return 0;
}

static std::vector<Stage> calculate_stages(int oversampling, double attenuation, double transition)
{
    std::vector<Stage> stages;
    stages.reserve(8);

    bool done = false;
    for (int num_stage = 0; !done; ++num_stage) {
        if (num_stage > 0)
            printf("\n");

        Stage stage;
        stage.factor = 2 << num_stage;
        stage.tbw = transition *
            std::pow(0.5, num_stage) + 0.5 * (1 - std::pow(0.5, num_stage));

        stage.nbr_coefs = FD::compute_nbr_coefs_from_proto(attenuation, stage.tbw);
        double *coefs = new double[stage.nbr_coefs]{};
        stage.coefs.reset(coefs);

        FD::compute_coefs(coefs, attenuation, stage.tbw);

        done = stage.factor >= oversampling;

        stages.push_back(std::move(stage));
    }

    return stages;
}

static void generate_cpp_prologue(int argc, char *argv[])
{
    printf("// This is generated by the Sfizz HIIR designer\n");
    printf("// Using options:");
    for (int i = 1; i < argc; ++i)
        printf(" %s", argv[i]);
    printf("\n");

    printf("\n");

    printf(
        "#pragma once\n"
        "#include \"OversamplerHelpers.h\"\n"
        "\n"
        "namespace sfz {\n"
        );
}

static void generate_cpp_epilogue()
{
    printf("} // namespace sfz\n");
}

static void generate_cpp_coefs(const Stage *stages, int num_stages)
{
    for (int i = 0; i < num_stages; ++i) {
        const Stage &stage = stages[i];
        const double *coefs = stage.coefs.get();
        printf("// %dx <-> %dx: TBW = %g\n", stage.factor, stage.factor / 2, stage.tbw);
        printf("static constexpr double OSCoeffs%dx[%d] = {\n", stage.factor, stage.nbr_coefs);
        for (int i = 0; i < stage.nbr_coefs; ++i) {
            printf("\t" "%.18f,\n", coefs[i]);
        }
        printf("};\n");
    }
}

static void generate_cpp_upsampler(const Stage *stages, int num_stages)
{
    printf("class Upsampler {\n");
    printf("public:\n");

    printf("\t" "Upsampler()\n");
    printf("\t" "{\n");
    for (int i = 0; i < num_stages; ++i) {
        const Stage &stage = stages[i];
        printf("\t\t" "up%d_.set_coefs(OSCoeffs%dx);\n", stage.factor, stage.factor);
    }
    printf("\t" "}\n");

    printf("\t" "void clear()\n");
    printf("\t" "{\n");
    for (int i = 0; i < num_stages; ++i) {
        const Stage &stage = stages[i];
        printf("\t\t" "up%d_.clear_buffers();\n", stage.factor);
    }
    printf("\t" "}\n");

    printf("\t" "static int recommendedBuffer(int factor, int spl)\n");
    printf("\t" "{\n");
    printf("\t\t" "return factor * spl;\n");
    printf("\t" "}\n");

    printf("\t" "static bool canProcess(int factor)\n");
    printf("\t" "{\n");
    printf("\t\t" "switch (factor) {\n");
    printf("\t\t" "case 1:\n");
    for (int i = 0; i < num_stages; ++i) {
        const Stage &stage = stages[i];
        printf("\t\t" "case %d:\n", stage.factor);
    }
    printf("\t\t\t" "return true;\n");
    printf("\t\t" "default:\n");
    printf("\t\t\t" "return false;\n");
    printf("\t\t" "}\n");
    printf("\t" "}\n");

    printf("\t" "void process(int factor, const float *in, float *out, int spl, float *temp, int ntemp)\n");
    printf("\t" "{\n");
    printf("\t\t" "switch (factor) {\n");
    printf("\t\t" "case 1:\n");
    printf("\t\t\t" "if (in != out) std::memcpy(out, in, spl * sizeof(float));\n");
    printf("\t\t\t" "break;\n");
    for (int i = 0; i < num_stages; ++i) {
        const Stage &stage = stages[i];
        printf("\t\t" "case %d:\n", stage.factor);
        printf("\t\t\t" "process%dx(in, out, spl, temp, ntemp);\n", stage.factor);
        printf("\t\t\t" "break;\n");
    }
    printf("\t\t" "default:\n");
    printf("\t\t\t" "ASSERTFALSE;\n");
    printf("\t\t\t" "break;\n");
    printf("\t\t" "}\n");
    printf("\t" "}\n");

    for (int n = 1; n <= num_stages; ++n) {
        printf("\t" "void process%dx(const float *in, float *out, int spl, float *temp, int ntemp)\n", stages[n - 1].factor);
        printf("\t" "{\n");
        printf("\t\t" "int maxspl = ntemp / %d;\n", stages[n - 1].factor);
        printf("\t\t" "ASSERT(maxspl >= 0);\n");
        printf("\t\t" "float *t1 = temp;\n");
        printf("\t\t" "float *t2 = temp + %d * maxspl;\n", stages[n - 1].factor / 2);
        printf("\t\t" "(void)t1;\n");
        printf("\t\t" "(void)t2;\n");
        printf("\t\t" "while (spl > 0) {\n");
        printf("\t\t\t" "int curspl = (spl < maxspl) ? spl : maxspl;\n");
        for (int i = 0; i < n; ++i) {
            const Stage &stage = stages[i];
            const char *tempnames[] = {"t1", "t2"};
            const char *outname = tempnames[i & 1];
            const char *inname = tempnames[1 - (i & 1)];
            if (i == 0)
                inname = "in";
            if (i + 1 == n)
                outname = "out";
            printf("\t\t\t" "up%d_.process_block(%s, %s, %d * curspl);\n", stage.factor, outname, inname, stage.factor / 2);
        }
        printf("\t\t\t" "in += curspl;\n");
        printf("\t\t\t" "out += curspl;\n");
        printf("\t\t\t" "spl -= curspl;\n");
        printf("\t\t" "}\n");
        printf("\t" "}\n");
    }

    printf("private:\n");
    for (int i = 0; i < num_stages; ++i) {
        const Stage &stage = stages[i];
        printf("\t" "hiir::Upsampler2x<%d> up%d_;\n", stage.nbr_coefs, stage.factor);
    }
    printf("};\n");
}

static void generate_cpp_downsampler(const Stage *stages, int num_stages)
{
    printf("class Downsampler {\n");
    printf("public:\n");

    printf("\t" "Downsampler()\n");
    printf("\t" "{\n");
    for (int i = 0; i < num_stages; ++i) {
        const Stage &stage = stages[num_stages - 1 - i];
        printf("\t\t" "down%d_.set_coefs(OSCoeffs%dx);\n", stage.factor, stage.factor);
    }
    printf("\t" "}\n");

    printf("\t" "void clear()\n");
    printf("\t" "{\n");
    for (int i = 0; i < num_stages; ++i) {
        const Stage &stage = stages[num_stages - 1 - i];
        printf("\t\t" "down%d_.clear_buffers();\n", stage.factor);
    }
    printf("\t" "}\n");

    printf("\t" "static int recommendedBuffer(int factor, int spl)\n");
    printf("\t" "{\n");
    printf("\t\t" "return factor * spl;\n");
    printf("\t" "}\n");

    printf("\t" "static bool canProcess(int factor)\n");
    printf("\t" "{\n");
    printf("\t\t" "switch (factor) {\n");
    printf("\t\t" "case 1:\n");
    for (int i = 0; i < num_stages; ++i) {
        const Stage &stage = stages[i];
        printf("\t\t" "case %d:\n", stage.factor);
    }
    printf("\t\t\t" "return true;\n");
    printf("\t\t" "default:\n");
    printf("\t\t\t" "return false;\n");
    printf("\t\t" "}\n");
    printf("\t" "}\n");

    printf("\t" "void process(int factor, const float *in, float *out, int spl, float *temp, int ntemp)\n");
    printf("\t" "{\n");
    printf("\t\t" "switch (factor) {\n");
    for (int i = 0; i < num_stages; ++i) {
        const Stage &stage = stages[num_stages - 1 - i];
        printf("\t\t" "case %d:\n", stage.factor);
        printf("\t\t\t" "process%dx(in, out, spl, temp, ntemp);\n", stage.factor);
        printf("\t\t\t" "break;\n");
    }
    printf("\t\t" "case 1:\n");
    printf("\t\t\t" "if (in != out) std::memcpy(out, in, spl * sizeof(float));\n");
    printf("\t\t\t" "break;\n");
    printf("\t\t" "default:\n");
    printf("\t\t\t" "ASSERTFALSE;\n");
    printf("\t\t\t" "break;\n");
    printf("\t\t" "}\n");
    printf("\t" "}\n");

    for (int n = 1; n <= num_stages; ++n) {
        printf("\t" "void process%dx(const float *in, float *out, int spl, float *temp, int ntemp)\n", stages[n - 1].factor);
        printf("\t" "{\n");
        printf("\t\t" "int maxspl = ntemp / %d;\n", stages[n - 1].factor);
        printf("\t\t" "ASSERT(maxspl >= 0);\n");
        printf("\t\t" "float *t1 = temp;\n");
        printf("\t\t" "float *t2 = temp + %d * maxspl;\n", stages[n - 1].factor / 2);
        printf("\t\t" "(void)t1;\n");
        printf("\t\t" "(void)t2;\n");
        printf("\t\t" "while (spl > 0) {\n");
        printf("\t\t\t" "int curspl = (spl < maxspl) ? spl : maxspl;\n");
        for (int i = 0; i < n; ++i) {
            const Stage &stage = stages[n - 1 - i];
            const char *tempnames[] = {"t1", "t2"};
            const char *outname = tempnames[i & 1];
            const char *inname = tempnames[1 - (i & 1)];
            if (i == 0)
                inname = "in";
            if (i + 1 == n)
                outname = "out";
            printf("\t\t\t" "down%d_.process_block(%s, %s, %d * curspl);\n", stage.factor, outname, inname, stage.factor / 2);
        }
        printf("\t\t\t" "in += curspl;\n");
        printf("\t\t\t" "out += curspl;\n");
        printf("\t\t\t" "spl -= curspl;\n");
        printf("\t\t" "}\n");
        printf("\t" "}\n");
    }

    printf("private:\n");
    for (int i = 0; i < num_stages; ++i) {
        const Stage &stage = stages[num_stages - 1 - i];
        printf("\t" "hiir::Downsampler2x<%d> down%d_;\n", stage.nbr_coefs, stage.factor);
    }
    printf("};\n");
}