import numpy as np
import pickle

class MFA:
    """
    Gaussian Mixture Model with optimization for High-dimensional Data
    In each component, the Covariance Matrix is approximated using:
    - A: a rectangular d x l matrix, where l <= d (typically l << d)
    - s: variance (sigma^2) of the added isotropic noise
    - Data is generated by transforming samples drawn from a standard normal random variable:
      - x = z_l @ A.T + z_d * s*np.eye() + mu
      - Two random sources are used: z_l (in the lower dimension l) and z_d in the full dimension
    - The covariance matrix Sigma = A @ A.T + s*np.eye()
    """

    def __init__(self, components=None):
        self.components = components
        self.eps = 1e-16
        self.max_l = 8

    @staticmethod
    def _draw_from_component(num_samples, c, add_noise=True):
        d, l = c['A'].shape
        z_l = np.random.normal(size=[num_samples, l])
        X = z_l @ c['A'].T + c['mu'].T

        if add_noise:
            z_d = np.random.normal(size=[num_samples, d])
            X += z_d * np.sqrt(c['D'])

        return X

    def draw_samples(self, num_samples, add_noise=True):
        pi = [c['pi'] for c in self.components.values()]
        pi /= sum(pi)

        # Choose components and then sample relevant points from each components
        samples = np.zeros((num_samples, self.components[0]['mu'].size), dtype=float)
        s_k = np.random.choice(len(pi), p=pi, size=num_samples)
        for k, c in self.components.items():
            s_k_i = (s_k == k)
            samples[s_k_i, :] = MFA._draw_from_component(sum(s_k_i), c, add_noise=add_noise)
        return samples

    def save(self, file_name):
        full_name = file_name if file_name.endswith('.pkl') else file_name + '.pkl'
        with open(full_name, 'wb') as f:
            pickle.dump(self.components, f, pickle.HIGHEST_PROTOCOL)

    def load(self, file_name):
        full_name = file_name if file_name.endswith('.pkl') else file_name + '.pkl'
        with open(full_name, 'rb') as f:
            self.components = pickle.load(f)