import numpy as np


class ols:
    """
    Simple algorithm for dimensionality reduction.
    """
            
    def __init__(self, n_components):
        self._fitted_parameters = None
        self._n_components = n_components
    
    def fit(self, fvs, labels):
        fvs = np.matrix(fvs)
        
        num_features = fvs.shape[1]
        
        """
        if num_features > 8*self._n_components:
            
            sub = [ols(self._n_components) for i in range(0, self._divide)]
            self._sub = sub
            sub_fvs = np.split(fvs, self._divide, 1) # Split into 'divide' parts.
            
            # Train sub machines.
            sub_reduced = []
            print('%d sub machines. %d'%(len(sub), num_features))
            for s,f in zip(sub, sub_fvs):
                print('fit')
                print(f.shape)
                s.fit(f, labels)
                
                # Apply sub transforms
                print('apply tf')
                r = s.transform(f)
                sub_reduced.append(r)
            
            fvs = np.concatenate(sub_reduced, 1)
        """
        
        # Train main transform.
        self._fitted_parameters = _train_ols_transform(fvs, labels, self._n_components)
            
    
    def transform(self, fvs):
        assert self._fitted_parameters is not None, "OLS is not yet fitted."
        """
        if self._sub is not None:
            sub = self._sub
            sub_fvs = np.split(fvs, self._divide, 1) # Split into 'divide' parts.
            
            sub_reduced = [s.transform(f) for s,f in zip(sub, sub_fvs)]
            fvs = np.concatenate(sub_reduced, 1)
        """
        
        idx,tf = self._fitted_parameters
        return _apply_ols_transform(self.fitted_indices, tf, fvs)
        
    @property
    def fitted_indices(self):
        assert self._fitted_parameters is not None, "OLS is not yet fitted."
        return self._fitted_parameters[0]
        

def _orthogonalize_inplace(to, vectors):
    """
    orthogonalize to one vector
    
    to: The vector to orthogonalize the others to.
    vectors: List of vectors that get orthogonalized to `to`.
    
    Returns: Vector of np.dot(to, v[i]) = <to, v[i]>.
    """
    to = np.array(to)
    
    num_features = vectors.shape[0]
    
    tf = np.zeros(num_features)
    
    # normalize
    to_norm = np.sqrt(np.sum(to**2))
    if to_norm == 0:
        to_norm = 1
    to /= to_norm
    
    for c in range(0, num_features):
        
        v = vectors[c, :] # column

        beta = np.dot(to, v) # inner product
        v -= to*beta
        
        tf[c] = beta/to_norm
    
    return tf

def _variance_scores(fvs, labels):
    """
    fvs: row: feature, column: feature vector
    """
    assert fvs.shape[1] == len(labels)
    
    classes = np.unique(labels)
    
    var_all_classes = np.var(fvs, 1)
    
    vars_single_class = np.array([np.var(fvs[:,labels==c], 1) for c in classes])
    var_single_class = np.sum(vars_single_class, 0)
    
    scores = var_all_classes / (var_single_class + 1e-3)
    
    return scores

def _train_ols_transform(fvs, labels, num_features):
    """
    fvs: List of vectors. Each vector contains a single feature of all samples.
    labels: The class labels of the samples.
    """
    
    fvs = np.array(fvs).transpose()
    
    picked_indices = []     # Indices of chosen features.
    picked_features = []    # The chosen features.
    tfs = []                # List of transformation to reproduce orthogonalization.
    
    for i in range(0, num_features):
        scores = _variance_scores(fvs, labels)
        
        # Pick vector with maximum score.
        max_idx = np.argmax(scores)
        picked_feature = fvs[max_idx].copy()
        
        picked_indices.append(max_idx)
        picked_features.append(picked_feature)
        
        # Orthogonalize all vectors to picked vector.
        tf = _orthogonalize_inplace(picked_feature, fvs)
        
        # Set chosen feature to zero such that we will ignore it in next iteration.
        fvs[max_idx] = 0
        
        tfs.append(tf)
        print('ols %6.2f%%'%(100*(i+1)/num_features))
    
    picked_indices = np.array(picked_indices)
    tfs = np.array(tfs)
    tfs = tfs[:,picked_indices]
    
    return picked_indices, tfs
        
def _apply_ols_transform(indices, tf, vector):
    
    def _tf_single_vector(vector):
        # Only pick certain indices.
        r = vector[indices]
        assert(len(indices) == len(r))

        """
        # Replay orthogonalization.
        for i, betas in zip(range(1, len(r)), tf):
            r[i:] -= r[i-1]*betas[i:]
        """
        return r
    
    if len(vector.shape) == 1:
        return _tf_single_vector(vector)
    
    assert(len(indices) == tf.shape[0])
    assert(len(indices) == tf.shape[1])
    
    r = np.array(list(map(_tf_single_vector, vector)))
    
    return r
