"""
A mixture consisting of two substructures.

One substructure is assumed to be known and fixed, up to a transformation. The other substructure is completely unknown and needs to be estimated from data.
"""

import numpy as np
cimport numpy as np
cimport cython

from cryoem.model.isomixture cimport IsoMixture3D
from cryoem.model.isodiagonalmixture cimport IsoDiagonalMixture3D
from cryoem.model.transformation cimport Transformation3D



cdef class IsoSubstructure:
    """Mixture with fixed substructure."""

    def __cinit__(self, IsoMixture3D mixture_fixed,
                 IsoMixture3D mixture_variable,
                 double ratio,
                 Transformation3D substructure_transformation):
        """
        Args:
            mixture_fixed: Known substructure to be fixed during sampling.
            mixture_variable: Unknown submixture.
            ratio: Total weight of substructure relative to combined mixture.
            substructure_transformation: Needs to be applied to
                mixture_fixed to align it with mixture_variable.
        """

        self._mixture_fixed = mixture_fixed
        self._mixture_variable = mixture_variable
        self._ratio = ratio
        self._substructure_transformation = substructure_transformation

        # Create mixture to contain given mixtures as submixtures.
        self._num_fixed_components = self._mixture_fixed._num_components
        self._num_variable_components = self._mixture_variable._num_components
        K = self._num_fixed_components + self._num_variable_components
        self._num_components = K
        self._mixture_full = IsoDiagonalMixture3D(
            np.empty((K, 3), dtype=np.double),
            np.empty(K, dtype=np.double),
            np.empty(K, dtype=np.double))

        self._mixture_transformed = mixture_fixed.copy()

    @property
    def variance(self):
        return self._mixture_variable._variance

    @property
    def mixture_variable(self):
        return self._mixture_variable

    @property
    def mixture_transformed(self):
        return self._mixture_transformed

    @property
    def mixture_full(self):
        return self._mixture_full

    @property
    def substructure_transformation(self):
        return self._substructure_transformation

    @property
    def ratio(self):
        return self._ratio

    def set_substructure_transformation(self, t):
        self._substructure_transformation = t

    cpdef update_mixture_full(self):
        """Copy two smaller mixtures to unified mixture."""

        m0 = self._mixture_fixed
        m0T = self._mixture_transformed
        m1 = self._mixture_variable
        m = self._mixture_full

        K0 = m0._num_components
        K1 = m1._num_components

        # Update parameters from fixed mixture.
        self._substructure_transformation.transform(m0._means, m0T._means)
        for k in range(K0):
            for i in range(3):
                m._means[k, i] = m0T._means[k, i]
            m._variances[k] = m0T._variance
            m._weights[k] = self._ratio * m0T._weights[k]

        # Update parameters from variable mixture.
        for k in range(K1):
            for i in range(3):
                m._means[k + K0, i] = m1._means[k, i]
            m._variances[k + K0] = m1._variance
            m._weights[k + K0] = (1.0 - self._ratio) * m1._weights[k]

