from __future__ import annotations

from dials.algorithms.refinement.parameterisation.prediction_parameters import (
    PredictionParameterisation,
    SparseGradientVectorMixin,
)
from dials.array_family import flex
from dials_refinement_helpers_ext import dRq_de


class StillsPredictionParameterisation(PredictionParameterisation):
    """Concrete class that inherits functionality of the
    PredictionParameterisation parent class and provides a detector space
    implementation of the get_gradients function for still images. Gradients of
    the minimum rotation to the Ewald sphere, DeltaPsi, are calculated for use
    as a restraint in refinement"""

    _grad_names = ("dX_dp", "dY_dp", "dDeltaPsi_dp")

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        # check that a goniometer parameterisation is not passed in
        assert not self._goniometer_parameterisations
        return

    def _local_setup(self, reflections):
        """Setup additional attributes used in gradients calculation. These are
        specific to the stills-type prediction parameterisations that rotates the
        relp q onto the Ewald sphere by an angle DeltaPsi before prediction"""

        self._DeltaPsi = reflections["delpsical.rad"]

        # q is the reciprocal lattice vector, in the lab frame
        self._q = self._UB * self._h
        self._q_scalar = self._q.norms()
        self._qq = self._q_scalar * self._q_scalar

        # r is the reciprocal lattice vector rotated to the Ewald sphere
        self._r = self._s1 - self._s0

        # we also need the unit directions q0 and s0u
        self._q0 = self._q.each_normalize()
        self._s0u = self._s0.each_normalize()

        # e1 is the unit vector about which DeltaPsi rotation is defined
        self._e1 = self._q0.cross(self._s0u).each_normalize()

        # q1 completes an orthonormal set with q0 and e1
        self._q1 = self._q0.cross(self._e1).each_normalize()

        # c0 completes an orthonormal set with s0u and e1
        self._c0 = self._s0u.cross(self._e1).each_normalize()

        # we want the wavelength
        self._wavelength = 1.0 / self._s0.norms()

        return

    def _beam_derivatives(self, isel, parameterisation=None, reflections=None):
        """helper function to extend the derivatives lists by derivatives of the
        beam parameterisations"""

        # Get required data
        s0 = self._s0.select(isel)
        s0u = self._s0u.select(isel)
        wl = self._wavelength.select(isel)
        r = self._r.select(isel)
        e1 = self._e1.select(isel)
        q = self._q.select(isel)
        c0 = self._c0.select(isel)
        DeltaPsi = self._DeltaPsi.select(isel)
        D = self._D.select(isel)

        # get the derivatives of the beam vector wrt the parameters
        ds0_dbeam_p = parameterisation.get_ds_dp(use_none_as_null=True)

        dDeltaPsi_dp = []
        dpv_dp = []

        # loop through the parameters
        for der in ds0_dbeam_p:

            if der is None:
                dpv_dp.append(None)
                dDeltaPsi_dp.append(None)
                continue

            # repeat the derivative in an array
            ds0 = flex.vec3_double(len(s0u), der.elems)

            # we need the derivative of the unit beam direction too. This requires
            # scaling by the wavelength and projection onto the Ewald sphere
            scaled = ds0 * wl
            ds0u = scaled.dot(c0) * c0 + scaled.dot(e1) * e1

            # calculate the derivative of DeltaPsi for this parameter
            dDeltaPsi = -1.0 * (r.dot(ds0)) / (e1.cross(r).dot(s0))
            dDeltaPsi_dp.append(dDeltaPsi)

            # calculate (d[r]/d[e1])(d[e1]/dp)
            de1_dp = c0.cross(ds0u)
            dr_de1 = dRq_de(DeltaPsi, e1, q)
            drde_dedp = dr_de1 * de1_dp

            # dp = 1.e-8 # finite step size for the parameter
            # del_e1 = de1_dp * dp
            # e1f = e1 + del_e1 * 0.5
            # rfwd = q.rotate_around_origin(e1f, DeltaPsi)
            # e1r = e1 - del_e1 * 0.5
            # rrev = q.rotate_around_origin(e1r, DeltaPsi)
            # drde_dedp = (rfwd - rrev) * (1 / dp)

            # calculate the derivative of pv for this parameter
            dpv_dp.append(D * (ds0 + e1.cross(r) * dDeltaPsi + drde_dedp))

        return dpv_dp, dDeltaPsi_dp

    def _xl_orientation_derivatives(
        self, isel, parameterisation=None, reflections=None
    ):
        """helper function to extend the derivatives lists by derivatives of the
        crystal orientation parameterisations"""

        # Get required data
        B = self._B.select(isel)
        h = self._h.select(isel)
        e1 = self._e1.select(isel)
        DeltaPsi = self._DeltaPsi.select(isel)
        s1 = self._s1.select(isel)
        q = self._q.select(isel)
        q_scalar = self._q_scalar.select(isel)
        qq = self._qq.select(isel)
        q0 = self._q0.select(isel)
        r = self._r.select(isel)
        s0 = self._s0.select(isel)
        s0u = self._s0u.select(isel)
        D = self._D.select(isel)

        # get derivatives of the U matrix wrt the parameters
        dU_dxlo_p = parameterisation.get_ds_dp(use_none_as_null=True)

        dDeltaPsi_dp = []
        dpv_dp = []

        # loop through the parameters
        for der in dU_dxlo_p:

            if der is None:
                dpv_dp.append(None)
                dDeltaPsi_dp.append(None)
                continue

            der_mat = flex.mat3_double(len(B), der.elems)

            # calculate the derivative of q for this parameter
            dq = der_mat * B * h

            # calculate the derivative of r for this parameter
            dr = dq.rotate_around_origin(e1, DeltaPsi)

            # calculate the derivative of DeltaPsi for this parameter
            dDeltaPsi = -1.0 * (dr.dot(s1)) / (e1.cross(r).dot(s0))
            dDeltaPsi_dp.append(dDeltaPsi)

            # derivative of the axis e1
            q_dot_dq = q.dot(dq)
            dq0 = (q_scalar * dq - (q_dot_dq * q0)) / qq
            de1_dp = dq0.cross(s0u)

            # calculate (d[r]/d[e1])(d[e1]/dp)
            # dr_de1 = dRq_de(DeltaPsi, e1, q)
            # drde_dedp = dr_de1 * de1_dp

            # The above calculation is not correct. Instead, calculate the partial
            # derivative of r wrt change in rotation axis e1 by finite differences
            dp = 1.0e-8
            del_e1 = de1_dp * dp
            e1f = e1 + del_e1 * 0.5
            rfwd = q.rotate_around_origin(e1f, DeltaPsi)
            e1r = e1 - del_e1 * 0.5
            rrev = q.rotate_around_origin(e1r, DeltaPsi)
            drde_dedp = (rfwd - rrev) * (1 / dp)

            # calculate the derivative of pv for this parameter
            dpv_dp.append(D * (dr + e1.cross(r) * dDeltaPsi + drde_dedp))

        return dpv_dp, dDeltaPsi_dp

    def _xl_unit_cell_derivatives(self, isel, parameterisation=None, reflections=None):
        """helper function to extend the derivatives lists by derivatives of the
        crystal unit cell parameterisations"""

        # Get required data
        U = self._U.select(isel)
        h = self._h.select(isel)
        e1 = self._e1.select(isel)
        DeltaPsi = self._DeltaPsi.select(isel)
        s1 = self._s1.select(isel)
        q = self._q.select(isel)
        q_scalar = self._q_scalar.select(isel)
        qq = self._qq.select(isel)
        q0 = self._q0.select(isel)
        r = self._r.select(isel)
        s0 = self._s0.select(isel)
        s0u = self._s0u.select(isel)
        D = self._D.select(isel)

        # get derivatives of the B matrix wrt the parameters
        dB_dxluc_p = parameterisation.get_ds_dp(use_none_as_null=True)

        dDeltaPsi_dp = []
        dpv_dp = []

        # loop through the parameters
        for der in dB_dxluc_p:

            if der is None:
                dpv_dp.append(None)
                dDeltaPsi_dp.append(None)
                continue

            der_mat = flex.mat3_double(len(U), der.elems)

            # calculate the derivative of q for this parameter
            dq = U * der_mat * h

            # calculate the derivative of r for this parameter
            dr = dq.rotate_around_origin(e1, DeltaPsi)

            # calculate the derivative of DeltaPsi for this parameter
            dDeltaPsi = -1.0 * (dr.dot(s1)) / (e1.cross(r).dot(s0))
            dDeltaPsi_dp.append(dDeltaPsi)

            # derivative of the axis e1
            q_dot_dq = q.dot(dq)
            dq0 = (q_scalar * dq - (q_dot_dq * q0)) / qq
            de1_dp = dq0.cross(s0u)

            # calculate (d[r]/d[e1])(d[e1]/dp)
            # dr_de1 = dRq_de(DeltaPsi, e1, q)
            # drde_dedp = dr_de1 * de1_dp

            # The above calculation is not correct. Instead, calculate the partial
            # derivative of r wrt change in rotation axis e1 by finite differences
            dp = 1.0e-8
            del_e1 = de1_dp * dp
            e1f = e1 + del_e1 * 0.5
            rfwd = q.rotate_around_origin(e1f, DeltaPsi)
            e1r = e1 - del_e1 * 0.5
            rrev = q.rotate_around_origin(e1r, DeltaPsi)
            drde_dedp = (rfwd - rrev) * (1 / dp)

            # calculate the derivative of pv for this parameter
            dpv = D * (dr + e1.cross(r) * dDeltaPsi + drde_dedp)
            dpv_dp.append(dpv)

        return dpv_dp, dDeltaPsi_dp

    @staticmethod
    def _calc_dX_dp_and_dY_dp_from_dpv_dp(w_inv, u_w_inv, v_w_inv, dpv_dp):
        """helper function to calculate positional derivatives from
        dpv_dp using the quotient rule"""

        dX_dp = []
        dY_dp = []

        for der in dpv_dp:
            if der is None:
                dX_dp.append(None)
                dY_dp.append(None)
            else:
                du_dp, dv_dp, dw_dp = der.parts()

                dX_dp.append(w_inv * (du_dp - dw_dp * u_w_inv))
                dY_dp.append(w_inv * (dv_dp - dw_dp * v_w_inv))

        return dX_dp, dY_dp


class StillsPredictionParameterisationSparse(
    SparseGradientVectorMixin, StillsPredictionParameterisation
):
    """A version of StillsPredictionParameterisation that uses a sparse matrix
    data structure for memory efficiency when there are a large number of
    Experiments"""

    pass


class SphericalRelpStillsPredictionParameterisation(StillsPredictionParameterisation):
    """Modified StillsPredictionParameterisation for the model that assumes
    relps are spherical and prediction requires that this sphere intersects the
    Ewald sphere, not that the relp centre is rotated onto the Ewald sphere.
    Gradients of DeltaPsi are still calculated for use as a restraint."""

    def _local_setup(self, reflections):
        """Setup additional attributes used in gradients calculation. These are
        specific to the stills-type prediction parameterisations that leaves the
        relp q generally off the Ewald sphere but assumes it has spherical extent
        which intersects the Ewald sphere. Gradients of the minimum rotation
        angle DeltaPsi are still calculated for use as a restraint"""

        self._DeltaPsi = reflections["delpsical.rad"]

        # q is the reciprocal lattice vector, in the lab frame
        self._q = self._UB * self._h

        # quantities involving the direct beam vector
        self._q_s0 = self._q + self._s0
        s = self._q_s0.norms()
        ss = self._q_s0.dot(self._q_s0)
        sss = s * ss
        self._inv_s = 1.0 / s
        self._inv_sss = 1.0 / sss

        # r is the reciprocal lattice vector rotated to the Ewald sphere, required
        # for derivatives of DeltaPsi
        self._r = self._s1 - self._s0

        # we also need the unit directions q0 and s0u
        self._q0 = self._q.each_normalize()
        self._s0u = self._s0.each_normalize()

        # e1 is the unit vector about which DeltaPsi rotation is defined
        self._e1 = self._q0.cross(self._s0u).each_normalize()

        # q1 completes an orthonormal set with q0 and e1
        self._q1 = self._q0.cross(self._e1).each_normalize()

        # c0 completes an orthonormal set with s0u and e1
        self._c0 = self._s0u.cross(self._e1).each_normalize()

        # we want the wavenumber and wavelength
        self._nu = self._s0.norms()
        self._wavelength = 1.0 / self._nu

        return

    def _beam_derivatives(self, isel, parameterisation=None, reflections=None):
        """helper function to extend the derivatives lists by derivatives of the
        beam parameterisations"""

        # Get required data
        s0 = self._s0.select(isel)
        s0u = self._s0u.select(isel)
        nu = self._nu.select(isel)
        r = self._r.select(isel)
        e1 = self._e1.select(isel)
        q_s0 = self._q_s0.select(isel)
        inv_s = self._inv_s.select(isel)
        inv_sss = self._inv_sss.select(isel)
        D = self._D.select(isel)

        # get the derivatives of the beam vector wrt the parameters
        ds0_dbeam_p = parameterisation.get_ds_dp(use_none_as_null=True)

        dDeltaPsi_dp = []
        dpv_dp = []

        # loop through the parameters
        for der in ds0_dbeam_p:

            if der is None:
                dpv_dp.append(None)
                dDeltaPsi_dp.append(None)
                continue

            # repeat the derivative in an array
            ds0 = flex.vec3_double(len(s0u), der.elems)

            # we need the derivative of the unit beam direction too. This requires
            # scaling by the wavelength and projection onto the Ewald sphere
            # scaled = ds0 * wl
            # ds0u = scaled.dot(c0) * c0 + scaled.dot(e1) * e1

            # calculate the derivative of DeltaPsi for this parameter
            dDeltaPsi = -1.0 * (r.dot(ds0)) / (e1.cross(r).dot(s0))
            dDeltaPsi_dp.append(dDeltaPsi)

            # term 1
            term1 = (s0u.dot(ds0) * q_s0 + nu * ds0) * inv_s

            # term 2
            term2 = (nu * q_s0 * q_s0.dot(ds0)) * inv_sss

            # calculate the derivative of pv for this parameter
            dpv_dp.append(D * (term1 - term2))

        return dpv_dp, dDeltaPsi_dp

    def _xl_orientation_derivatives(
        self, isel, parameterisation=None, reflections=None
    ):
        """helper function to extend the derivatives lists by
        derivatives of the crystal orientation parameterisations"""

        # Get required data
        B = self._B.select(isel)
        h = self._h.select(isel)
        e1 = self._e1.select(isel)
        DeltaPsi = self._DeltaPsi.select(isel)
        s1 = self._s1.select(isel)
        nu = self._nu.select(isel)
        q_s0 = self._q_s0.select(isel)
        inv_s = self._inv_s.select(isel)
        inv_sss = self._inv_sss.select(isel)
        r = self._r.select(isel)
        s0 = self._s0.select(isel)
        D = self._D.select(isel)

        # get derivatives of the U matrix wrt the parameters
        dU_dxlo_p = parameterisation.get_ds_dp(use_none_as_null=True)

        dDeltaPsi_dp = []
        dpv_dp = []

        # loop through the parameters
        for der in dU_dxlo_p:

            if der is None:
                dpv_dp.append(None)
                dDeltaPsi_dp.append(None)
                continue

            der_mat = flex.mat3_double(len(B), der.elems)

            # calculate the derivative of q for this parameter
            dq = der_mat * B * h

            # calculate the derivative of r for this parameter
            dr = dq.rotate_around_origin(e1, DeltaPsi)

            # calculate the derivative of DeltaPsi for this parameter
            dDeltaPsi = -1.0 * (dr.dot(s1)) / (e1.cross(r).dot(s0))
            dDeltaPsi_dp.append(dDeltaPsi)

            # term 1
            term1 = (nu * dq) * inv_s

            # term 2
            term2 = (nu * q_s0 * q_s0.dot(dq)) * inv_sss

            # calculate the derivative of pv for this parameter
            dpv_dp.append(D * (term1 - term2))

        return dpv_dp, dDeltaPsi_dp

    def _xl_unit_cell_derivatives(self, isel, parameterisation=None, reflections=None):
        """helper function to extend the derivatives lists by
        derivatives of the crystal unit cell parameterisations"""

        # Get required data
        U = self._U.select(isel)
        h = self._h.select(isel)
        e1 = self._e1.select(isel)
        DeltaPsi = self._DeltaPsi.select(isel)
        s1 = self._s1.select(isel)
        nu = self._nu.select(isel)
        q_s0 = self._q_s0.select(isel)
        inv_s = self._inv_s.select(isel)
        inv_sss = self._inv_sss.select(isel)
        r = self._r.select(isel)
        s0 = self._s0.select(isel)
        D = self._D.select(isel)

        # get derivatives of the B matrix wrt the parameters
        dB_dxluc_p = parameterisation.get_ds_dp(use_none_as_null=True)

        dDeltaPsi_dp = []
        dpv_dp = []

        # loop through the parameters
        for der in dB_dxluc_p:

            if der is None:
                dpv_dp.append(None)
                dDeltaPsi_dp.append(None)
                continue

            der_mat = flex.mat3_double(len(U), der.elems)

            # calculate the derivative of q for this parameter
            dq = U * der_mat * h

            # calculate the derivative of r for this parameter
            dr = dq.rotate_around_origin(e1, DeltaPsi)

            # calculate the derivative of DeltaPsi for this parameter
            dDeltaPsi = -1.0 * (dr.dot(s1)) / (e1.cross(r).dot(s0))
            dDeltaPsi_dp.append(dDeltaPsi)

            # term 1
            term1 = (nu * dq) * inv_s

            # term 2
            term2 = (nu * q_s0 * q_s0.dot(dq)) * inv_sss

            # calculate the derivative of pv for this parameter
            dpv = D * (term1 - term2)
            dpv_dp.append(dpv)

        return dpv_dp, dDeltaPsi_dp


class SphericalRelpStillsPredictionParameterisationSparse(
    SparseGradientVectorMixin, SphericalRelpStillsPredictionParameterisation
):
    """A version of SphericalRelpStillsPredictionParameterisation that uses a
    sparse matrix data structure for memory efficiency when there are a large
    number of Experiments"""

    pass
