import torch
import torchquad

# Initialize Gaussian Legendre quadrature integrator
gl = torchquad.GaussLegendre()
def float_gcd(a, b, rtol = 1e-05, atol = 1e-08):
    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)
def torch_vectorize(f, inplace=False):
    def wrapper(tensor, *args):
        out = tensor if inplace else tensor.clone()
        view = out.flatten()
        for i, x in enumerate(view):
            view[i] = f(x, *args)
        return out
    return wrapper
# 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)

# Define arc length using torchquad and autodiff for derivatives
def arc(t, d, r, R):
    def integrand(t1):
        t1 = t1.requires_grad_(True)  # Ensure t1 allows gradients
        x_val = xt(t1, d, r, R)
        y_val = yt(t1, d, r, R)
        
        # Compute dx/dt and dy/dt using autograd
        dxdt_val = torch.autograd.grad(x_val, t1, grad_outputs=torch.ones_like(x_val), create_graph=True)[0]
        dydt_val = torch.autograd.grad(y_val, t1, grad_outputs=torch.ones_like(y_val), create_graph=True)[0]
        
        # Compute the integrand for the arc length formula
        integrand_val = dxdt_val ** 2 + dydt_val ** 2
        return integrand_val
    
    domain = torch.tensor([[0, t]]).requires_grad_()
    result = gl.integrate(integrand, dim=1, N=101, integration_domain=domain)
    return result
arc = torch_vectorize(arc)
# Function to estimate the derivative using Richardson extrapolation
def arc_grad(t, d, r, R):
    # Step sizes to use
    h_values = [0.01, 0.005, 0.0025]
    derivatives = []
    
    for h in h_values:
        # Central difference approximation
        arc_plus = arc(t + h, d, r, R)
        arc_minus = arc(t - h, d, r, R)
        derivative = (arc_plus - arc_minus) / (2 * h)
        derivatives.append(derivative)
    
    # Now apply Richardson extrapolation
    # Assuming the error term is proportional to h^2
    # derivative_exact = derivative_h + C * h^2
    
    # First extrapolation
    h1, h2 = h_values[0], h_values[1]
    d1, d2 = derivatives[0], derivatives[1]
    extrapolated1 = d2 + (d2 - d1) / ((h1 / h2) ** 2 - 1)
    
    # Second extrapolation
    h2, h3 = h_values[1], h_values[2]
    d2, d3 = derivatives[1], derivatives[2]
    extrapolated2 = d3 + (d3 - d2) / ((h2 / h3) ** 2 - 1)
    
    # Combine extrapolations assuming error is proportional to h^4
    final_extrapolated = extrapolated2 + (extrapolated2 - extrapolated1) / ((h1 / h3) ** 4 - 1)

    return final_extrapolated
def st(s,d,r,R):
    x = torch.Tensor([0])
    for i in range(100):
        x = x+(s-arc(x,d,r,R))/arc_grad(x,d,r,R)
        if s-arc(x,d,r,R) < 1e-9:
            print(i)
            break
    return x
RuntimeError                              Traceback (most recent call last)
Cell In[293], line 1
----> 1 st(torch.Tensor([0.6,0.6]),1,5,4)

Cell In[291], line 17, in torch_vectorize.<locals>.wrapper(tensor, *args)
     15 out = tensor if inplace else tensor.clone()
     16 for i, x in enumerate(tensor):
---> 17     tensor[i] = f(x, *args)
     18 return out

Cell In[291], line 80, in st(s, d, r, R)
     78 x = torch.Tensor([0])
     79 for i in range(100):
---> 80     x = x+(s-arc(x,d,r,R))/arc_grad(x,d,r,R)
     81     if s-arc(x,d,r,R) < 1e-9:
     82         print(i)

RuntimeError: A view was created in no_grad mode and its base or another view of its base has been modified inplace with grad mode enabled. 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.