# Released under The MIT License (MIT)
# http://opensource.org/licenses/MIT
# Copyright (c) 2013 SCoT Development Team

"""common spatial patterns (CSP) implementation"""

import numpy as np
from scipy.linalg import eig
    

def csp(x, cl, numcomp=None):
    """ Calculate common spatial patterns (CSP)

    Parameters
    ----------
    x : array-like, shape = [n_trials, n_channels, n_samples] or [n_channels, n_samples]
        EEG data set
    cl : list of valid dict keys
        Class labels associated with each trial. Currently only two classes are supported.
    numcomp : {int}, optional
        Number of patterns to keep after applying the CSP. If `numcomp` is
        greater than n_channels or None, all n_channels patterns are returned.

    Returns
    -------
    w : array, shape = [n_channels, n_components]
        CSP weight matrix
    v : array, shape = [n_components, n_channels]
        CSP projection matrix
    """

    x = np.asarray(x)
    cl = np.asarray(cl).ravel()

    if x.ndim != 3 or x.shape[0] < 2:
        raise AttributeError('CSP needs at least two trials.')

    t, m, n = x.shape
    
    if t != cl.size:
        raise AttributeError('CSP only works with multiple classes. Number of'
                             ' elements in cl (%d) must equal 3rd dimension of X (%d)' % (cl.size, t))

    labels = np.unique(cl)
    
    if labels.size != 2:
        raise AttributeError('CSP is currently implemented for 2 classes (got %d)' % labels.size)
        
    x1 = x[cl == labels[0], :, :]
    x2 = x[cl == labels[1], :, :]
    
    sigma1 = np.zeros((m, m))
    for t in range(x1.shape[0]):
        sigma1 += np.cov(x1[t, :, :]) / x1.shape[0]
    sigma1 /= sigma1.trace()
    
    sigma2 = np.zeros((m, m))
    for t in range(x2.shape[0]):
        sigma2 += np.cov(x2[t, :, :]) / x2.shape[0]
    sigma2 /= sigma2.trace()

    e, w = eig(sigma1, sigma1 + sigma2, overwrite_a=True, overwrite_b=True, check_finite=False)

    order = np.argsort(e)[::-1]
    w = w[:, order]
    # e = e[order]
        
    v = np.linalg.inv(w)
   
    # subsequently remove unwanted components from the middle of w and v
    if numcomp is None:
        numcomp = w.shape[1]
    while w.shape[1] > numcomp:
        i = int(np.floor(w.shape[1]/2))
        w = np.delete(w, i, 1)
        v = np.delete(v, i, 0)
        
    return w, v

@@ -89,16 +89,21 @@
 (c) 2013
+-2016
  SCoT De
@@ -121,17 +121,17 @@
 eam%0A%0A%22%22%22
-c
+C
 ommon sp
@@ -165,16 +165,17 @@
 entation
+.
 %22%22%22%0A%0Aimp
@@ -262,17 +262,16 @@
 %0A    %22%22%22
- 
 Calculat
@@ -301,16 +301,17 @@
 ns (CSP)
+.
 %0A%0A    Pa
@@ -351,21 +351,16 @@
 rray
--like
 , shape 
 = %5Bn
@@ -359,21 +359,17 @@
 ape 
-= %5Bn_
+(
 trials, 
 n_ch
@@ -364,18 +364,16 @@
 trials, 
-n_
 channels
@@ -378,25 +378,21 @@
 ls, 
-n_
 samples
-%5D
+)
  or 
-%5Bn_
+(
 chan
@@ -401,18 +401,16 @@
 ls, 
-n_
 samples
-%5D
+)
 %0A   
@@ -426,16 +426,17 @@
 data set
+.
 %0A    cl 
@@ -519,16 +519,17 @@
 urrently
+,
  only tw
@@ -537,16 +537,24 @@
  classes
+%0A       
  are sup
@@ -579,13 +579,11 @@
 p : 
-%7B
 int
-%7D
 , op
@@ -639,20 +639,16 @@
 pplying 
-the 
 CSP. If 
@@ -663,24 +663,16 @@
 %60 is
-%0A       
  greater
  tha
@@ -671,16 +671,22 @@
 ater
+%0A       
  than 
-n_
 chan
@@ -707,19 +707,8 @@
 all 
-n_channels 
 patt
@@ -772,21 +772,17 @@
 , shape 
-= %5Bn_
+(
 channels
@@ -783,18 +783,16 @@
 annels, 
-n_
 componen
@@ -793,17 +793,17 @@
 mponents
-%5D
+)
 %0A       
@@ -816,24 +816,25 @@
 eight matrix
+.
 %0A    v : arr
@@ -847,13 +847,9 @@
 ape 
-= %5Bn_
+(
 comp
@@ -860,18 +860,16 @@
 ts, 
-n_
 channels
 %5D%0A  
@@ -864,17 +864,17 @@
 channels
-%5D
+)
 %0A       
@@ -895,16 +895,17 @@
 n matrix
+.
 %0A    %22%22%22
@@ -1037,12 +1037,15 @@
 CSP 
-need
+require
 s at
@@ -1191,16 +1191,17 @@
 umber of
+ 
 '%0A      
@@ -1224,17 +1224,16 @@
        '
- 
 elements
@@ -1240,18 +1240,18 @@
  in cl (
-%25d
+%7B%7D
 ) must e
@@ -1259,12 +1259,50 @@
 ual 
-3rd 
+the first '%0A                             '
 dime
@@ -1314,18 +1314,22 @@
  of 
-X (%25d)' %25 
+x (%7B%7D)'.format
 (cl.
@@ -1459,17 +1459,19 @@
 ted for 
-2
+two
  classes
@@ -1475,20 +1475,63 @@
 ses 
-(got %25d)' %25 
+'%0A                             'only (got %7B%7D).'.format(
 labe
@@ -1538,16 +1538,17 @@
 ls.size)
+)
 %0A       
@@ -1991,16 +1991,31 @@
 _b=True,
+%0A              
  check_f
@@ -2084,36 +2084,8 @@
 er%5D%0A
-    # e = e%5Border%5D%0A        %0A
     

Clean up csp.py