import torch
import torchquad

def float_gcd(a, b, rtol=1e-5, atol=1e-8):
    t = min(abs(a), abs(b))
    while abs(b) > rtol * t + atol:
        a, b = b, a % b
    return a

def float_lcm(a, b):
    return abs(a * b) / float_gcd(a, b)

# Define derivative for a function f with respect to x
def grad(f, x):
    x = x.clone().detach().requires_grad_(True)
    y = f(x)
    y.backward()
    return x.grad

# Vectorize a function to work on tensors
def torch_vectorize(f):
    def wrapper(tensor, *args):
        out = tensor.clone()
        view = out.flatten()
        for i, x in enumerate(view):
            view[i] = f(x, *args).item()
        return out
    return wrapper

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

# Parametric equations and their derivatives
xt = lambda t, d, r, R: (R - r) * torch.cos(t) + d * torch.cos((R - r) / r * t)
dxdt = lambda t, d, r, R: grad(lambda t1: xt(t1, d, r, R), t)
dxdt = torch_vectorize(dxdt)

yt = lambda t, d, r, R: (R - r) * torch.sin(t) + d * torch.sin((R - r) / r * t)
dydt = lambda t, d, r, R: grad(lambda t1: yt(t1, d, r, R), t)
dydt = torch_vectorize(dydt)

# Define arc length using torchquad
def arc(t, d, r, R):
    integrand = lambda t1: torch.sqrt(dxdt(t1, d, r, R) ** 2 + dydt(t1, d, r, R) ** 2)
    domain = torch.Tensor([[0, t]]).requires_grad_()
    result = gl.integrate(integrand, dim=1, N=101, integration_domain=domain)
    return result

arc = torch_vectorize(arc)

# 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)
        darc_val = grad(lambda x1: arc(x1, d, r, R), x)
        x = x - (arc_val - s) / darc_val
        if torch.abs(arc_val - s) < 1e-9:
            break
    return x
#2*np.pi*float_lcm(r,R)/R
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Cell In[98], line 1
----> 1 arc(torch.Tensor([5,6]),5,5,2)

Cell In[97], line 26, in torch_vectorize.<locals>.wrapper(tensor, *args)
     24 view = out.flatten()
     25 for i, x in enumerate(view):
---> 26     view[i] = f(x, *args)
     27 return out

Cell In[97], line 49, in arc(t, d, r, R)
     47 t_var = t.clone().detach().requires_grad_(True)
     48 integrand = lambda t1: torch.sqrt(dxdt(t1, d, r, R) ** 2 + dydt(t1, d, r, R) ** 2)
---> 49 result = gl.integrate(integrand, dim=1, N=101, integration_domain=torch.tensor([[0, t_var]]).requires_grad_())
     50 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:49, in GridIntegrator.integrate(self, fn, dim, N, integration_domain, backend)
     46 grid_points, hs, n_per_dim = self.calculate_grid(N, integration_domain)
     48 logger.debug("Evaluating integrand on the grid.")
---> 49 function_values, num_points = self.evaluate_integrand(
     50     fn, grid_points, weights=self._weights(n_per_dim, dim, backend)
     51 )
     52 self._nr_of_fevals = num_points
     54 return self.calculate_result(
     55     function_values, dim, n_per_dim, hs, integration_domain
     56 )

File ~/.local/lib/python3.11/site-packages/torchquad/integration/base_integrator.py:65, in BaseIntegrator.evaluate_integrand(fn, points, weights, args)
     62 if args is None:
     63     args = ()
---> 65 result = fn(points, *args)
     67 if infer_backend(result) != infer_backend(points):
     68     warnings.warn(
     69         "The passed function's return value has a different numerical backend than the passed points. Will try to convert. Note that this may be slow as it results in memory transfers between CPU and GPU, if torchquad uses the GPU."
     70     )

Cell In[97], line 48, in arc.<locals>.<lambda>(t1)
     46 def arc(t, d, r, R):
     47     t_var = t.clone().detach().requires_grad_(True)
---> 48     integrand = lambda t1: torch.sqrt(dxdt(t1, d, r, R) ** 2 + dydt(t1, d, r, R) ** 2)
     49     result = gl.integrate(integrand, dim=1, N=101, integration_domain=torch.tensor([[0, t_var]]).requires_grad_())
     50     return result

Cell In[97], line 26, in torch_vectorize.<locals>.wrapper(tensor, *args)
     24 view = out.flatten()
     25 for i, x in enumerate(view):
---> 26     view[i] = f(x, *args)
     27 return out

Cell In[97], line 35, in <lambda>(t, d, r, R)
     33 # Parametric equations and their derivatives
     34 xt = lambda t, d, r, R: (R - r) * torch.cos(t) + d * torch.cos((R - r) / r * t)
---> 35 dxdt = lambda t, d, r, R: grad(lambda t1: xt(t1, d, r, R), t)
     36 dxdt = torch_vectorize(dxdt)
     38 yt = lambda t, d, r, R: (R - r) * torch.sin(t) + d * torch.sin((R - r) / r * t)

Cell In[97], line 15, in grad(f, x)
     14 def grad(f, x):
---> 15     x = x.clone().detach().requires_grad_(True)
     16     y = f(x)
     17     y.backward()

RuntimeError: Output 1 of UnbindBackward0 is a view and its base or another view of its base has been modified inplace. This view is the output of a function that returns multiple views. Such functions do not allow the output views to be modified inplace. You should replace the inplace operation by an out-of-place one.