import torch
from utils import linear_solve

torch.set_default_dtype(torch.float64)

def compute_adjoint(solid ,dJdx_U):
    '''
    This assumes that S is at equilibrium when called
    
    Input:
    - solid : an elastic solid at equilibrium
    - dJdx_U  : array of shape (3*#unpinned,)
    
    Output:
    - adjoint : array of shape (3*#v,)
    '''
    y = torch.zeros_like(solid.v_rest)
    dx0s = torch.zeros_like(solid.v_rest)
    def LHS(dx):
        '''
        Should implement the Hessian-Vector product (taking into account pinning constraints) as described in the handout. 
        '''
        dx0s[solid.free_idx] = dx.reshape(-1,3)
        df0s = - solid.compute_force_differentials(dx0s)
        return df0s[solid.free_idx].reshape(-1,)
    
    RHS = dJdx_U.T.reshape(-1,)

    y[solid.free_idx] = linear_solve(LHS, RHS).reshape(-1,3)
    
    return y.reshape(-1,)
