from mdt import LibraryFunctionTemplate


class RotateOrthogonalVector(LibraryFunctionTemplate):
    """Uses Rodrigues' rotation formula to rotate the given vector v by psi around k.

    This assumes that the vectors we are rotating are orthogonal to each other.

    Args:
        basis: the unit vector defining the rotation axis (k)
        to_rotate: the vector to rotate by the angle psi (v)
        psi: the rotation angle (psi)

    Returns:
        vector: the rotated vector
    """
    return_type = 'float4'
    parameters = ['float4 basis',
                  'float4 to_rotate',
                  'mot_float_type psi']
    cl_code = '''
        float cos_psi;
        float sin_psi = sincos(psi, &cos_psi);
        
        return to_rotate * cos_psi + (cross(basis, to_rotate) * sin_psi);
    '''


class RotateVector(LibraryFunctionTemplate):
    """Uses Rodrigues' rotation formula to rotate the given vector v by psi around k.

    This function makes no assumptions on the orthogonality of the vectors.

    Args:
        basis: the unit vector defining the rotation axis (k)
        to_rotate: the vector to rotate by the angle psi (v)
        psi: the rotation angle (psi)

    Returns:
        vector: the rotated vector
    """
    return_type = 'float4'
    parameters = ['float4 basis',
                  'float4 to_rotate',
                  'mot_float_type psi']
    cl_code = '''
        float cos_psi;
        float sin_psi = sincos(psi, &cos_psi);

        return to_rotate * cos_psi
                + (cross(basis, to_rotate) * sin_psi)
                + (basis * dot(basis, to_rotate) * (1 - cos_psi));
    '''


class SphericalToCartesian(LibraryFunctionTemplate):
    """Convert polar coordinates in 3d space to cartesian unit coordinates.

    .. code-block:: python

        x = sin(theta) * cos(phi)
        y = sin(theta) * sin(phi)
        z = cos(theta)

    Args:
        theta: polar angle of the first vector
        phi: azimuth angle of the first vector
    """
    return_type = 'float4'
    parameters = ['mot_float_type theta',
                  'mot_float_type phi']
    cl_code = '''
        mot_float_type cos_theta;
        mot_float_type sin_theta = sincos(theta, &cos_theta);
        mot_float_type cos_phi;
        mot_float_type sin_phi = sincos(phi, &cos_phi);
        
        return (float4)(cos_phi * sin_theta, sin_phi * sin_theta, cos_theta, 0.0);
    '''


class TensorSphericalToCartesian(LibraryFunctionTemplate):
    """Generates the D matrix for a Tensor compartment.

    The angles ``theta`` and ``phi`` are used for creating the first vector, ``vec0``.
    Next, ``vec0`` is rotated 90 degrees in the (x, z) plane to form a vector perpendicular to the
    principal direction. This vector is then rotated around ``psi`` to generate the first perpendicular
    orientation, ``vec1``. The third vector is generated by being perpendicular to the other two vectors.

    Args:
        theta: polar angle of the first vector
        phi: azimuth angle of the first vector
        psi: rotation around the first vector, used to generate the perpendicular vectors.
    """
    dependencies = ['RotateOrthogonalVector', 'SphericalToCartesian']
    parameters = ['mot_float_type theta',
                  'mot_float_type phi',
                  'mot_float_type psi',
                  'float4* vec0',
                  'float4* vec1',
                  'float4* vec2']
    cl_code = '''
        *vec0 = SphericalToCartesian(theta, phi);
        *vec1 = RotateOrthogonalVector(*vec0, SphericalToCartesian(theta + M_PI_2_F, phi), psi);
        *vec2 = cross(*vec0, *vec1);
    '''


class CartesianPolarDotProduct(LibraryFunctionTemplate):
    """Calculates the dot product between a cartesian and a polar coordinate vector.

    Prior to taking the dot product it will convert the polar coordinate vector to cartesian coordinates.

    Args:
        v0_x: the x coordinate of the first vector
        v0_y: the y coordinate of the first vector
        v0_z: the z coordinate of the first vector
        v1_theta: the polar angle of the second vector
        v1_phi: the azimuth angle of the second vector

    Returns:
        dot product between the two vectors
    """
    return_type = 'mot_float_type'
    parameters = ['mot_float_type v0_x',
                  'mot_float_type v0_y',
                  'mot_float_type v0_z',
                  'mot_float_type v1_theta',
                  'mot_float_type v1_phi']
    cl_code = '''
        mot_float_type cos_theta;
        mot_float_type sin_theta = sincos(v1_theta, &cos_theta);
        mot_float_type cos_phi;
        mot_float_type sin_phi = sincos(v1_phi, &cos_phi);
        
        return (v0_x * (cos_phi * sin_theta)) + (v0_y * (sin_phi * sin_theta)) + (v0_z * cos_theta);
    '''


class SquareMatrixMultiplication(LibraryFunctionTemplate):
    """Matrix multiplication of two square matrices.

    Having this as a special function is slightly faster than a more generic matrix multiplication algorithm.

    All matrices are expected to be in c/row-major order.

    Args:
        n: the rectangular size of the matrix (the number of rows / the number of columns).
        A: the left matrix
        B: the right matrix
        C: the output matrix
    """
    parameters = ['uint n',
                  'mot_float_type* A',
                  'mot_float_type* B',
                  'mot_float_type* C']
    cl_code = '''
        int i, j, z;
        double sum;
        for(int ind = 0; ind < n*n; ind++){
            i = ind / n;
            j = ind - n * i;
            
            sum = 0;
            for(z = 0; z < n; z++){
                sum += A[i * n + z] * B[z * n + j];
            }
            C[ind] = sum;
        }
    '''

