import tectosaur.util.quadrature as quad
from tectosaur.util.paget import paget
import numpy as np

def vertex_interior_quad(n_theta, n_rho, a):
    # return quad.gauss2d_tri(10)
    theta_quad = quad.map_to(quad.gaussxw(n_theta), [0, np.pi / 2])
    if a == 0:
        # rho_quad = quad.telles_singular(n_rho, -1)
        # rho_quad = quad.gaussxw(n_rho)
        rho_quad = quad.tanh_sinh(n_rho)
    else:
        rho_quad = paget(n_rho, a)
    pts = []
    wts = []
    for t, tw in zip(*theta_quad):
        rho_max = 1.0 / (np.cos(t) + np.sin(t))
        q = quad.map_to(rho_quad, [0, rho_max])
        for r, rw in zip(*q):
            assert(r > 0)
            srcxhat = r * np.cos(t)
            srcyhat = r * np.sin(t)
            pts.append((srcxhat, srcyhat))
            # The r**2 comes from two factors:
            # 1) The first r comes from the jacobian for the polar transform
            # 2) The second r corrects for the implicit 1/r in the quad rule
            weight = tw * rw * (r ** (a + 1))
            wts.append(weight)
    return np.array(pts), np.array(wts)

def vertex_adj_quad(n_theta, n_beta, n_alpha):
    theta_quad = quad.gaussxw(n_theta)
    beta_quad = quad.gaussxw(n_beta)
    alpha_quad = quad.gaussxw(n_alpha)

    def rho_max(theta):
        return 1.0 / (np.cos(theta) + np.sin(theta))

    pts = []
    wts = []
    def alpha_integral(w, tp, tq, b, alpha_max):
        q = quad.map_to(alpha_quad, [0, alpha_max])
        for a, aw in zip(*q):