/*
 * fractal
 * Written in 2020 by Lucas
 * CC0 1.0 Universal/Public domain - No rights reserved
 *
 * To the extent possible under law, the author(s) have dedicated all
 * copyright and related and neighboring rights to this software to the
 * public domain worldwide. This software is distributed without any
 * warranty. You should have received a copy of the CC0 Public Domain
 * Dedication along with this software. If not, see
 * <http://creativecommons.org/publicdomain/zero/1.0/>.
 */
/* gepopt visibility for GLIBC */
#if !defined(__OpenBSD__)
#define _DEFAULT_SOURCE
#endif

#include <complex.h>
#include <math.h>
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>

#include "util.h"

#define MIN(a, b) ((a) < (b) ? (a) : (b))
#define MAX(a, b) ((a) > (b) ? (a) : (b))

#define MAX_DEGREE 1000

struct poly {
	int degree;
	double complex *coefs;
	const char *expr;
};

static void
poly_init(struct poly *p, unsigned int degree, const char *expr)
{
	if ((p->coefs = malloc((degree + 1) * sizeof(*p->coefs))) == NULL)
		errx(1, "out of memory");

	p->degree = degree;
	p->expr = expr;
}

static void
poly_derivate(struct poly p, struct poly *dp, const char *expr)
{
	int i;

	poly_init(dp, p.degree == 0 ? 0 : p.degree - 1, expr);
	for (i = p.degree; i > 0; i--)
		dp->coefs[i - 1] = p.coefs[i] * i;
}

static void
poly_multiply(struct poly p, struct poly q, struct poly *f, const char *expr)
{
	int i, j;

	poly_init(f, p.degree + q.degree, expr);
	for (i = 0; i <= f->degree; i++)
		f->coefs[i] = 0.0;
	for (i = 0; i <= p.degree; i++)
		for (j = 0; j <= q.degree; j++)
			f->coefs[j + i] += p.coefs[i] * q.coefs[j];
}

static int
poly_cmp(const void *vp, const void *vq)
{
	const struct poly *p = vp, *q = vq;
	return p->degree - q->degree;
}

static void
poly_free(struct poly *p)
{
	free(p->coefs);
}

static void
print_fun_header(const char *fun_name)
{
	printf("static long double complex\n");
	printf("%s(long double complex z)\n", fun_name);
	printf("{\n");
}

static void
print_fun_body(size_t n, struct poly *polys)
{
	size_t i, j;
	int d;

	printf("\tlong double complex a");
	for (i = 0; i < n; i++)
		printf(", %s = 0.0L", polys[i].expr);
	printf(";\n\n");

	d = 0;
	for (i = 0; i < n; i++) {
		for (; d <= polys[i].degree; d++) {
			if (d == 0)
				printf("\ta = 1.0L;\n");
			else
				printf("\ta *= z;\n");
			for (j = n; j > i; j--)
				if (polys[j - 1].coefs[d] != 0.0L)
					printf("\t%s += (%.20f + I * %.20f) * a;\n",
					    polys[j - 1].expr,
					    creal(polys[j - 1].coefs[d]),
					    cimag(polys[j - 1].coefs[d]));
			printf("\n");
		}
	}
}

static void
print_fun_footer(void)
{
	printf("}\n");
}

static void
print_funs(struct poly p, struct poly q)
{
	struct poly polys[6];
	struct poly dp, ddp, dq, ddq;
	int i, j;

	for (i = 0; i <= MIN(p.degree, q.degree)
	    && p.coefs[i] == 0.0 && q.coefs[i] == 0.0; i++)
		;
	for (j = i; j <= p.degree; j++)
		p.coefs[j - i] = p.coefs[j];
	p.degree -= i;
	for (j = i; j <= q.degree; j++)
		q.coefs[j - i] = q.coefs[j];
	q.degree -= i;

	polys[0] = p;
	polys[1] = q;
	qsort(polys, 2, sizeof(struct poly), &poly_cmp);

	print_fun_header("f");
	print_fun_body(2, polys);
	printf("\treturn p / q;\n");
	print_fun_footer();

	poly_derivate(p, &dp, "dp");
	poly_derivate(q, &dq, "dq");
	polys[0] = p;
	polys[1] = q;
	polys[2] = dp;
	polys[3] = dq;
	qsort(polys, 4, sizeof(struct poly), &poly_cmp);

	printf("\n");
	print_fun_header("df");
	print_fun_body(4, polys);
	printf("\treturn (dp * q - p * dq) / (q * q);\n");
	print_fun_footer();

	poly_derivate(dp, &ddp, "ddp");
	poly_derivate(dq, &ddq, "ddq");
	polys[0] = p;
	polys[1] = q;
	polys[2] = dp;
	polys[3] = dq;
	polys[4] = ddp;
	polys[5] = ddq;
	qsort(polys, 6, sizeof(struct poly), &poly_cmp);

	printf("\n");
	print_fun_header("ddf");
	print_fun_body(6, polys);
	printf("\treturn ("
	    "q * q * ddp "
	    "- q * (2.0L * dp * dq + p * ddq) "
	    "+ 2.0L * p * dq * dq"
	    ") / (q * q * q);\n");
	print_fun_footer();

	poly_free(&dp);
	poly_free(&dq);
	poly_free(&ddp);
	poly_free(&ddq);
}

static void
usage(void)
{
	const char *p = xgetprogname();
	fprintf(stderr, "Usage:\n"
	    "\t%s p-degree q-degree p-coefs q-coefs\n"
	    "\t%s -h p-degree p-coefs\n"
	    "\t%s -N ra ia p-degree p-coefs\n"
	    "\t%s -n p-degree p-coefs\n", p, p, p, p);
	exit(1);
}

static void
do_poly_div(int argc, char *argv[])
{
	struct poly p, q;
	int i, j, p_degree, q_degree;

	if (argc < 2)
		usage();

	p_degree = parse_integer(argv[0], 0, MAX_DEGREE);
	q_degree = parse_integer(argv[1], 0, MAX_DEGREE);
	if (argc != p_degree + q_degree + 2 + 2)
		usage();

	poly_init(&p, p_degree, "p");
	poly_init(&q, q_degree, "q");

	for (j = 0, i = 1 + 1 + p.degree + 1 + q.degree;
	    i > 1 + 1 + p.degree;
	    i--, j++)
		q.coefs[j] = parse_double(argv[i]);
	for (j = 0;
	    i > 1;
	    i--, j++)
		p.coefs[j] = parse_double(argv[i]);

	print_funs(p, q);
	poly_free(&p);
	poly_free(&q);
}

static void
do_halley(int argc, char *argv[])
{
	struct poly f, df, ddf, p, q, s, t, u;
	int i, j, f_degree;

	if (argc < 1)
		usage();

	f_degree = parse_integer(argv[0], 0, MAX_DEGREE);
	if (argc != f_degree + 1 + 1)
		usage();

	poly_init(&f, f_degree, "f");
	for (j = 0, i = 1 + f.degree;
	    i > 0;
	    i--, j++)
		f.coefs[j] = parse_double(argv[i]);

	poly_derivate(f, &df, "df");
	poly_derivate(df, &ddf, "ddf");
	poly_multiply(f, df, &s, "f_df");
	poly_multiply(df, df, &t, "df_df");
	poly_multiply(f, ddf, &u, "f_ddf");

	poly_init(&p, MAX(MAX(t.degree, u.degree) + 1, s.degree), "p");
	poly_init(&q, MAX(t.degree, u.degree), "q");

	for (i = 0; i <= q.degree; i++)
		q.coefs[i] = 2.0 * t.coefs[i] - u.coefs[i];
	for (i = 0; i <= p.degree; i++)
		p.coefs[i] = (i == 0 ? 0.0 : q.coefs[i - 1])
		    - 2.0L * (i <= s.degree ? s.coefs[i] : 0.0);

	print_funs(p, q);
	poly_free(&f);
	poly_free(&df);
	poly_free(&ddf);
	poly_free(&s);
	poly_free(&t);
	poly_free(&u);
	poly_free(&p);
	poly_free(&q);
}

static void
do_nova(int argc, char *argv[])
{
	struct poly f, df, p;
	double complex a;
	double ra, ia;
	int i, j, f_degree;

	if (argc < 3)
		usage();

	ra = parse_double(argv[0]);
	ia = parse_double(argv[1]);
	a = ra + I * ia;
	argc -= 2; argv += 2;

	f_degree = parse_integer(argv[0], 0, MAX_DEGREE);
	if (argc != f_degree + 1 + 1)
		usage();

	poly_init(&f, f_degree, "f");
	for (j = 0, i = 1 + f.degree;
	    i > 0;
	    i--, j++)
		f.coefs[j] = parse_double(argv[i]);

	poly_derivate(f, &df, "q");
	poly_init(&p, f_degree, "p");
	for (i = 0; i <= f.degree; i++)
		p.coefs[i] = (i == 0 ? 0.0 : df.coefs[i - 1]) - a * f.coefs[i];

	print_funs(p, df);
	poly_free(&f);
	poly_free(&df);
	poly_free(&p);
}

static void
do_newton(int argc, char *argv[])
{
	struct poly f, df, p;
	int i, j, f_degree;

	if (argc < 1)
		usage();

	f_degree = parse_integer(argv[0], 0, MAX_DEGREE);
	if (argc != f_degree + 1 + 1)
		usage();

	poly_init(&f, f_degree, "f");
	for (j = 0, i = 1 + f.degree;
	    i > 0;
	    i--, j++)
		f.coefs[j] = parse_double(argv[i]);

	poly_derivate(f, &df, "q");
	poly_init(&p, f_degree, "p");
	for (i = 0; i <= f.degree; i++)
		p.coefs[i] = (i == 0 ? 0.0 : df.coefs[i - 1]) - f.coefs[i];

	print_funs(p, df);
	poly_free(&f);
	poly_free(&df);
	poly_free(&p);
}

int
main(int argc, char *argv[])
{
	int ch, hflag, Nflag, nflag;

	hflag = Nflag = nflag = 0;
	while ((ch = getopt(argc, argv, "hNn")) != -1)
		switch (ch) {
		case 'h':
			hflag = 1;
			break;
		case 'N':
			Nflag = 1;
			break;
		case 'n':
			nflag = 1;
			break;
		default:
			usage();
		}
	argc -= optind;
	argv += optind;

	if (hflag)
		do_halley(argc, argv);
	else if (Nflag)
		do_nova(argc, argv);
	else if (nflag)
		do_newton(argc, argv);
	else
		do_poly_div(argc, argv);

	return 0;
}