from cryoem.model.distribution cimport Distribution

cdef class Rotation:
    cdef double[:, :] _matrix
    cpdef rotate(self, double[:, :] x, double[:, :] y)
    cpdef compose(self, other)

cdef class Rotation3D(Rotation):
    cpdef copy_in_place(self, Rotation3D other)
    cpdef compose_in_place(self, Rotation3D other, Rotation3D result)

cdef class Rotation2D(Rotation):
    pass

cdef class ExpTrace(Distribution):
    cdef:
        int _n_iter
        double[:, :] U, V, rotation, rotation2, rotation3
        double[:] L

    cdef double clip(self, double x, double a, double b)
    cdef double det(self, double[:, :] x)
    cdef double _sample_beta(self, double kappa)
    cpdef double[:, :] sample(self, double[:, ::1] A)

cdef euler_angles_to_rotation(
    double alpha, double beta, double gamma, double[:, :] matrix)

cpdef exponential_coordinates_to_rotation(
    double[:] v, double[:, :] matrix)

cdef dot3x3(double[:, :] x, double[:, :] y, double[:, :] result)

cdef transpose3x3(double[:, :] x, double[:, :] result)

cdef class ExpTrace2D(Distribution):
    cpdef double sample(self, double[:, :] A)

cdef class UniformRotations:
    cdef:
        int _num_rotations
    cpdef double[:, ::1] generate_random(self)
    cpdef double[:, ::1] generate_variable_stepping(self)

cdef class RotationGradients:
    cdef:
        double[:, ] _vector
        double[:, :] _identity, _term1, _term2, _R, _dvx, _dvy, _dvz
    cpdef get_dvx(self)
    cpdef get_dvy(self)
    cpdef get_dvz(self)
    cpdef get_dv(self, int index)
    cpdef compute(self, double[:] v)
    cpdef scalar_derivative(self, double[:, :] dEdR, double[:] dEdv)    
    cpdef relative_scalar_derivative(
        self, double[:, :] dEdR, double[:, :] R0,
        double[:] dEdv)
