"""
This module provides functionality for manipulating proteomics data sets.

Functionality includes merging data sets and interfacing with attributes in a
structured format.
"""

# Built-ins
from __future__ import absolute_import, division

import logging

# Core data analysis libraries
import numpy as np

import pyproteome as pyp


LOGGER = logging.getLogger("pyproteome.constand")


def constand(ds, inplace=False, n_iters=50, tol=1e-5):
    """
    Normalize channels to given levels for intra-run comparisons.

    Divides all channel values by a given level.

    Parameters
    ----------
    lvls : dict of str, float or
    ds : :class:`.DataSet`
        Mapping of channel names to normalized levels. Alternatively,
        a data set to pass to levels.get_channel_levels() or use
        pre-calculated levels from.
    inplace : bool, optional
        Modify this data set in place.

    Returns
    -------
    ds : :class:`.DataSet`
    """
    new = ds

    if not inplace:
        new = new.copy()

    channels = list(new.channels.values())

    new_channels = pyp.utils.norm(new.channels)
    # new = new.dropna()

    k = new[channels].values
    err = np.inf
    m, n = k.shape

    for ind in range(1, n_iters * 2 + 1):
        if ind % 2 == 1:
            # In the odd step the rows are fitted to match the row marginals
            # (i.e. constraints):
            #
            # K^(2t + 1) = R^(t + 1) * K ^ (2t)
            #
            # The row multipliers R^(t + 1) are computed such that the mean
            # of the reporter ion intensities equals 1/n:
            r = np.diag([
                # Nan-hiding here!
                1 / (
                    np.count_nonzero(~np.isnan(k[i])) *
                    np.nanmean(k[i])
                )
                for i in range(m)
            ])

            # k^(2t + 1)
            k = np.matmul(r, np.nan_to_num(k))
            err = sum([
                abs(np.nanmean(k[:, i]) - 1 / n)
                for i in range(n)
            ]) / 2
        else:
            # In the even step the columns are fitted to match the column
            # marginals (i.e. constraints):
            #
            # K^(2t + 2) = K^(2t + 1) * S^(t + 1)
            #
            # The column multipliers S^(t + 1) are computed such that the
            # mean of the reporter ion intensities equals 1/n:
            s = np.diag([
                # Nan-hiding here!
                1 / (
                    np.count_nonzero(~np.isnan(k[i])) *
                    np.nanmean(k[:, i])
                )
                for i in range(n)
            ])

            # k^(2t + 2)
            k = np.matmul(k, s)
            err = sum([
                abs(np.nanmean(k[i]) - 1 / n)
                for i in range(m)
            ]) / 2

        if err < tol:
            break

    for i, (key, norm_key) in enumerate(zip(
        new.channels.values(),
        new_channels.values(),
    )):
        new.psms[norm_key] = k[:, i]
        del new.psms[key]

    new.intra_normalized = True
    new.channels = new_channels
    new.groups = new.groups.copy()

    new.update_group_changes()

    return new

@@ -1611,44 +1611,8 @@
 iag(
-%5B%0A                # Nan-hiding here!
 %0A   
@@ -1670,36 +1670,41 @@
 zero(~np.isnan(k
-%5Bi%5D)
+), axis=1
 ) *%0A            
@@ -1727,29 +1727,16 @@
 an(k
-%5Bi%5D)%0A                
+, axis=1
 )%0A  
@@ -1749,32 +1749,16 @@
         
-for i in range(m
 )%0A      
@@ -1755,33 +1755,32 @@
   )%0A            
-%5D
 )%0A%0A            #
@@ -1848,37 +1848,33 @@
           err = 
-sum(%5B
+(
 %0A               
@@ -1890,22 +1890,25 @@
 anmean(k
-%5B:, i%5D
+, axis=0)
 ) - 1 / 
@@ -1900,37 +1900,33 @@
 axis=0)) - 1 / n
-)%0A   
+%0A
              for
@@ -1925,40 +1925,14 @@
     
- for i in range(n)%0A            %5D
+).sum(
 ) / 
@@ -2308,44 +2308,8 @@
 iag(
-%5B%0A                # Nan-hiding here!
 %0A   
@@ -2375,20 +2375,25 @@
 .isnan(k
-%5Bi%5D)
+), axis=0
 ) *%0A    
@@ -2424,32 +2424,16 @@
 an(k
-%5B:, i%5D)%0A                
+, axis=0
 )%0A  
@@ -2450,24 +2450,8 @@
     
-for i in range(n
 )%0A  
@@ -2460,17 +2460,16 @@
         
-%5D
 )%0A%0A     
@@ -2542,13 +2542,9 @@
 r = 
-sum(%5B
+(
 %0A   
@@ -2576,11 +2576,17 @@
 an(k
-%5Bi%5D
+, axis=1)
 ) - 
@@ -2590,21 +2590,17 @@
  - 1 / n
-)%0A   
+%0A
         
@@ -2607,40 +2607,14 @@
     
- for i in range(m)%0A            %5D
+).sum(
 ) / 

Use numpy transformations for calculations