#include "acb.h"
#include "acb_mat.h"
#include "acb_modular.h"
#include "acb_theta.h"
static void
acb_theta_naive_all_gen(acb_ptr th, acb_srcptr zs, slong nb, const acb_mat_t tau, slong prec)
{
slong g = acb_mat_nrows(tau);
slong n = 1 << g;
acb_ptr all_zs, ata, v;
acb_t c;
slong a, b, d, k;
all_zs = _acb_vec_init(g * n * nb);
ata = _acb_vec_init(n);
v = _acb_vec_init(g);
acb_init(c);
for (a = 0; a < n; a++)
{
acb_theta_char_get_acb(v, a, g);
acb_mat_vector_mul_col(v, tau, v, prec);
for (k = 0; k < nb; k++)
{
_acb_vec_add(all_zs + k * g * n + a * g, zs + k * g, v, g, prec);
}
acb_theta_char_dot_acb(&ata[a], a, v, g, prec);
}
acb_theta_naive_0b(th, all_zs, n * nb, tau, prec);
for (a = 0; a < n; a++)
{
for (k = 0; k < nb; k++)
{
acb_theta_char_dot_acb(c, a, zs + k * g, g, prec);
acb_mul_2exp_si(c, c, 1);
acb_add(c, c, &ata[a], prec);
acb_exp_pi_i(c, c, prec);
_acb_vec_scalar_mul(th + k * n * n + a * n,
th + k * n * n + a * n, n, c, prec);
}
for (b = 0; b < n; b++)
{
d = acb_theta_char_dot(a, b, g);
for (k = 0; k < nb; k++)
{
acb_mul_i_pow_si(&th[k * n * n + a * n + b],
&th[k * n * n + a * n + b], d);
}
}
}
_acb_vec_clear(all_zs, g * n * nb);
_acb_vec_clear(ata, n);
_acb_vec_clear(v, g);
acb_clear(c);
}
static void
acb_theta_naive_all_g1(acb_ptr th, acb_srcptr zs, slong nb, const acb_mat_t tau, slong prec)
{
acb_t q4, q, w;
acb_ptr res;
int w_is_unit;
slong k;
acb_init(q4);
acb_init(q);
acb_init(w);
res = _acb_vec_init(4);
acb_mul_2exp_si(q4, acb_mat_entry(tau, 0, 0), -2);
acb_exp_pi_i(q4, q4, prec);
acb_pow_ui(q, q4, 4, prec);
for (k = 0; k < nb; k++)
{
acb_exp_pi_i(w, &zs[k], prec);
w_is_unit = arb_is_zero(acb_imagref(&zs[k]));
acb_modular_theta_sum(&res[0], &res[1], &res[2], &res[3],
w, w_is_unit, q, 1, prec);
acb_set(&th[4 * k], &res[2]);
acb_set(&th[4 * k + 1], &res[3]);
acb_mul(&th[4 * k + 2], &res[1], q4, prec);
acb_mul(&th[4 * k + 3], &res[0], q4, prec);
acb_neg(&th[4 * k + 3], &th[4 * k + 3]);
}
acb_clear(q4);
acb_clear(q);
acb_clear(w);
_acb_vec_clear(res, 4);
}
void
acb_theta_naive_all(acb_ptr th, acb_srcptr zs, slong nb, const acb_mat_t tau, slong prec)
{
slong g = acb_mat_nrows(tau);
if (g == 1)
{
acb_theta_naive_all_g1(th, zs, nb, tau, prec);
}
else
{
acb_theta_naive_all_gen(th, zs, nb, tau, prec);
}
}