import torch
import torchquad

# Initialize Gaussian Legendre quadrature integrator
gl = torchquad.GaussLegendre()

# Parametric equations
def xt(t, d, r, R):
    return (R - r) * torch.cos(t) + d * torch.cos((R - r) / r * t)

def yt(t, d, r, R):
    return (R - r) * torch.sin(t) + d * torch.sin((R - r) / r * t)

# Analytical derivatives of the parametric equations
def dxdt(t, d, r, R):
    return - (R - r) * torch.sin(t) - d * ((R - r) / r) * torch.sin((R - r) / r * t)

def dydt(t, d, r, R):
    return (R - r) * torch.cos(t) + d * ((R - r) / r) * torch.cos((R - r) / r * t)

# Define arc length using torchquad
def arc(t, d, r, R):
    def integrand(t1):
        dxdt_val = dxdt(t1, d, r, R)
        dydt_val = dydt(t1, d, r, R)
        return torch.sqrt(dxdt_val ** 2 + dydt_val ** 2)
    
    # Ensure t is a tensor with the correct shape for integration_domain
    if t.dim() == 0:
        domain = torch.tensor([[0, t]])
    else:
        domain = torch.stack([torch.zeros_like(t), t], dim=1)
    
    result = gl.integrate(integrand, dim=1, N=101, integration_domain=domain)
    return result

# Solve for t given s, d, r, R using Newton's method
def ts(s, d, r, R):
    x = torch.tensor(0.0, requires_grad=True)
    for i in range(50):
        arc_val = arc(x, d, r, R)
        arc_val.backward()
        darc_val = x.grad
        with torch.no_grad():
            x_new = x - (arc_val - s) / darc_val
        x.grad.zero_()
        x = x_new.detach().requires_grad_()
        if torch.abs(arc_val - s) < 1e-9:
            break
    return x

# Example usage
d = 5
r = 5
R = 2
t_values = torch.tensor([5.0, 6.0], requires_grad=True)
arc_lengths = arc(t_values, d, r, R)
print(arc_lengths)
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[159], line 57
     55 R = 2
     56 t_values = torch.tensor([5.0, 6.0], requires_grad=True)
---> 57 arc_lengths = arc(t_values, d, r, R)
     58 print(arc_lengths)

Cell In[159], line 34, in arc(t, d, r, R)
     31 else:
     32     domain = torch.stack([torch.zeros_like(t), t], dim=1)
---> 34 result = gl.integrate(integrand, dim=1, N=101, integration_domain=domain)
     35 return result

File ~/.local/lib/python3.11/site-packages/torchquad/integration/gaussian.py:42, in Gaussian.integrate(self, fn, dim, N, integration_domain, backend)
     29 def integrate(self, fn, dim, N=8, integration_domain=None, backend=None):
     30     """Integrates the passed function on the passed domain using a Gaussian rule (Gauss-Legendre on [-1,1] as a default).
     31 
     32     Args:
   (...)
     40         backend-specific number: Integral value
     41     """
---> 42     return super().integrate(fn, dim, N, integration_domain, backend)

File ~/.local/lib/python3.11/site-packages/torchquad/integration/grid_integrator.py:42, in GridIntegrator.integrate(self, fn, dim, N, integration_domain, backend)
     39 if N is None:
     40     N = self._get_minimal_N(dim)
---> 42 integration_domain = _setup_integration_domain(dim, integration_domain, backend)
     43 backend = infer_backend(integration_domain)
     44 self._check_inputs(dim=dim, N=N, integration_domain=integration_domain)

File ~/.local/lib/python3.11/site-packages/torchquad/integration/utils.py:151, in _setup_integration_domain(dim, integration_domain, backend)
    148         integration_domain = anp.array(integration_domain, like=backend)
    150 if integration_domain.shape != (dim, 2):
--> 151     raise ValueError(
    152         "The integration domain has an unexpected shape. "
    153         f"Expected {(dim, 2)}, got {integration_domain.shape}"
    154     )
    155 return integration_domain

ValueError: The integration domain has an unexpected shape. Expected (1, 2), got torch.Size([2, 2])