#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <stdbool.h>

// Function to check if input is a vector
bool isvector(double* f, int rows, int cols) {
    return (rows == 1 || cols == 1);
}

// Function to check if input is a matrix
bool ismatrix(double* f, int rows, int cols) {
    return (rows > 1 && cols > 1);
}

// Function to check if input is a scalar
bool isscalar(double* f, int rows, int cols) {
    return (rows == 1 && cols == 1);
}

// Main function
double* cossquare(double* f, int rows, int cols, double fc, double BW) {
    double f1 = fc - BW;
    double f4 = fc + BW;
    double* r = malloc(rows * cols * sizeof(double));

    if (isvector(f, rows, cols)) {
        int n = rows * cols;
        for (int i = 0; i < n; i++) {
            double ss = sin(M_PI * f[i] / (2 * fc));
            bool ri1 = (fabs(f[i]) > fc) && (fabs(f[i]) <= f4);
            bool ri2 = (fabs(f[i]) <= fc) && (fabs(f[i]) <= f4);
            double ff1 = ri1 ? fabs(f[i]) : 0;
            double ff2 = ri2 ? fabs(f[i]) : 0;
            double r1 = ri1 ? pow(cos(M_PI * (fc - ff1) / (f4 - f1)), 2) : 0;
            double r2 = ri2 ? ss * pow(cos(M_PI * (fc - ff2) / (f4 - f1)), 2) : 0;
            r[i] = r1 + (f[i] >= 0 ? 1 : -1) * r2;
        }
    } else if (ismatrix(f, rows, cols)) {
        for (int c = 0; c < cols; c++) {
            for (int i = 0; i < rows; i++) {
                double ss = sin(M_PI * f[i * cols + c] / (2 * fc));
                bool ri1 = (fabs(f[i * cols + c]) > fc) && (fabs(f[i * cols + c]) <= f4);
                bool ri2 = (fabs(f[i * cols + c]) <= fc) && (fabs(f[i * cols + c]) <= f4);
                double ff1 = ri1 ? fabs(f[i * cols + c]) : 0;
                double ff2 = ri2 ? fabs(f[i * cols + c]) : 0;
                double r1 = ri1 ? pow(cos(M_PI * (fc - ff1) / (f4 - f1)), 2) : 0;
                double r2 = ri2 ? ss * pow(cos(M_PI * (fc - ff2) / (f4 - f1)), 2) : 0;
                r[i * cols + c] = r1 + (f[i * cols + c] >= 0 ? 1 : -1) * r2;
            }
        }
    } else if (isscalar(f, rows, cols)) {
        if ((fabs(f[0]) >= f1) && (fabs(f[0]) <= f4)) {
            if (fabs(f[0]) > fc) {
                r[0] = pow(cos(M_PI * (fc - fabs(f[0])) / (f4 - f1)), 2);
            } else {
                r[0] = sin(M_PI * f[0] / (2 * fc)) * pow(cos(M_PI * (fc - fabs(f[0])) / (f4 - f1)), 2);
            }
        } else {
            r[0] = 0;
        }
    }

    return r;
}