class PfTestBase(object):
    """Base class for generated sql and non-sql test cases."""
    pass class BaseData(object):
    """Base class for handling processed data.

    In general, specialised subclasses of BaseData will be used
    to handle processed data rather than BaseData itself.

    Usage: ..........
    """
    __metaclass__ = MetaMethods

    def __init__(self):
        self.meta = PfMetaData()
        self.history = "%s > New %s" %(datetime.now(), self.__class__.__name__)
        if not hasattr(self, 'channels'):
            self.channels = ChannelList()
        
    def save(self):
        if pyfusion.orm_manager.IS_ACTIVE:
            # this may be inefficient: get it working, then get it fast
            self.channels.save()
            session = pyfusion.orm_manager.Session()
            session.add(self)
            session.commit()
            session.close() class DataSet(BaseDataSet):
    pass class Coords(object):
    """Stores coordinates with an interface for coordinate transforms."""
    def __init__(self, default_coords_name, default_coords_tuple,  **kwargs):
        self.default_name = default_coords_name
        self.default_value_1 = default_coords_tuple[0]
        self.default_value_2 = default_coords_tuple[1]
        self.default_value_3 = default_coords_tuple[2]
        kwargs.update(((default_coords_name, default_coords_tuple),))
        self.__dict__.update(kwargs)

    def add_coords(self, **kwargs):
        self.__dict__.update(kwargs)
    
    def load_from_config(self, **kwargs):
        for kw in kwargs.iteritems():
            if kw[0] == 'coord_transform':
                transform_list = pyfusion.config.pf_options('CoordTransform', kw[1])
                for transform_name in transform_list:
                    transform_class_str = pyfusion.config.pf_get('CoordTransform', kw[1], transform_name)
                    transform_class = import_from_str(transform_class_str)
                    self.load_transform(transform_class)
            elif kw[0].startswith('coords_'):
                coord_values = tuple(map(float,kw[1].split(',')))
                self.add_coords(**{kw[0][7:]: coord_values})
    
    def load_transform(self, transform_class):
        def _new_transform_method(**kwargs):
            return transform_class().transform(self.__dict__.get(transform_class.input_coords),**kwargs)
        self.__dict__.update({transform_class.output_coords:_new_transform_method})

    def save(self):
        if pyfusion.orm_manager.IS_ACTIVE:
            # this may be inefficient: get it working, then get it fast
            session = pyfusion.orm_manager.Session()
            session.add(self)
            session.commit()
            session.close() class ChannelList(list):
    def __init__(self, *args):
        self.extend(args)

    def save(self):
        if pyfusion.orm_manager.IS_ACTIVE:
            self._channels.extend(self)
            session = pyfusion.orm_manager.Session()
            session.add(self)
            session.commit()
            session.close()

            
    @reconstructor
    def repopulate(self):
        for i in self._channels:
            if not i in self: self.append(i) class Channel(object):
    def __init__(self, name, coords):
        self.name = name
        self.coords = coords

    def save(self):
        if pyfusion.orm_manager.IS_ACTIVE:
            # this may be inefficient: get it working, then get it fast
            self.coords.save()
            session = pyfusion.orm_manager.Session()
            session.add(self)
            session.commit()
            session.close() class FloatDelta(BaseData):
#  Note: channel_1, 2 are channels, not channel names
    def __init__(self, channel_1, channel_2, delta, **kwargs):
        self.channel_1 = channel_1
        self.channel_2 = channel_2
        self.delta = delta
        super(FloatDelta, self).__init__(**kwargs) class BaseOrderedDataSet(object):
    __metaclass__ = MetaMethods

    def __init__(self, label=''):
        #SH added this to make it more compatible with DataSet
        self.meta = PfMetaData()
        self.created = datetime.now()
        self.label = label
        self.history = "%s > New %s" %(self.created, self.__class__.__name__)
        if label == '':
            label = unique_id()
        if not pyfusion.orm_manager.IS_ACTIVE:
            self.data_items = []
        
    def save(self):
        if pyfusion.orm_manager.IS_ACTIVE:
            session = pyfusion.orm_manager.Session()
            session.add(self)
            session.commit()
            session.close()

    def append(self, item):
        if pyfusion.orm_manager.IS_ACTIVE:
            self.data_items[len(self)] = OrderedDataSetItem(item, len(self))
        else:
            self.data_items.append(item)
    def __len__(self):
        #if pyfusion.orm_manager.IS_ACTIVE:
        #    return self.data_items.count()
        #else:
        return self.data_items.__len__()

    def __getitem__(self, key):
        if pyfusion.orm_manager.IS_ACTIVE:
            return self.data_items[key].item
        else:
            return self.data_items.__getitem__(key) class TimeseriesData(BaseData):
    def __init__(self, timebase = None, signal=None, channels=None, bypass_length_check=False, **kwargs):
        self.timebase = timebase
        self.channels = channels
        self.scales = np.array([[1]])  # retain amplitude scaling info for later
        if bypass_length_check == True:
            self.signal = signal
        else:
            if signal.n_samples() == len(timebase):
                self.signal = signal
            else:
                raise ValueError, "signal has different number of samples to timebase"
        super(TimeseriesData, self).__init__(**kwargs)


    def generate_frequency_series(self, NFFT, step, window='hamming'):
        w =scipy.hamming(NFFT)
        if len(self.signal.shape)==2:
            signal = np.array([np.fft.rfft(w*self.signal[:,i:i+NFFT]/NFFT) for i in range(0, self.signal.shape[1]-NFFT, step)])
        else:
            signal = np.array([np.fft.rfft(w*self.signal[i:i+NFFT]/NFFT) for i in range(0, self.signal.shape[1]-NFFT, step)])
        timebase = np.array([np.average(self.timebase[i:i+NFFT]) 
                             for i in range(0, self.signal.shape[1]-NFFT, step)])

        #This is borrowed from here:
        #http://docs.scipy.org/doc/numpy-dev/reference/generated/numpy.fft.rfftfreq.html
        d = (self.timebase[1] - self.timebase[0])
        val = 1.0/(NFFT*d)
        N = NFFT//2 + 1
        frequency_base = np.round((np.arange(0, N, dtype=int)) * val,4)
        return FrequencyseriesData(frequency_base = frequency_base, timebase=timebase, 
                                   signal=signal,channels=self.channels,NFFT=NFFT, window='hamming',step=step) class Timebase(np.ndarray):
    """Timebase vector with parameterised internal representation.

    see doc/subclassing.py in numpy code for details on subclassing ndarray
    """
    def __new__(cls, input_array):
        # should this follow the example in doc/subclassing.py?... (it doesn't)
        obj = np.asarray(input_array).view(cls).copy()
        obj.sample_freq = 1.0/(obj[1]-obj[0])
        obj.meta = PfMetaData()
        return obj

    def is_contiguous(self):
        return max(((self[1:]-self[:-1])-1.0/self.sample_freq)**2) < (0.1/self.sample_freq)**2

    def normalise_freq(self, input_freq):
        """Normalise input frequencies to [0,1] where 1 is pi*sample_freq"""
        try:
            return input_freq/(0.5*self.sample_freq)
        except:
            # sample_freq should maybe be self.sample_freq? i don't remember
            return [i/(0.5*self.sample_freq) for i in sample_freq]

    def __array_finalize__(self, obj):
        # ``self`` is a new object resulting from
        # ndarray.__new__(InfoArray, ...), therefore it only has
        # attributes that the ndarray.__new__ constructor gave it -
        # i.e. those of a standard ndarray.
        #
        # We could have got to the ndarray.__new__ call in 3 ways:
        # From an explicit constructor - e.g. InfoArray():
        #    obj is None
        #    (we're in the middle of the InfoArray.__new__
        #    constructor, and self.info will be set when we return to
        #    InfoArray.__new__)
        if obj is None: return
        # From view casting - e.g arr.view(InfoArray):
        #    obj is arr
        #    (type(obj) can be InfoArray)
        # From new-from-template - e.g infoarr[:3]
        #    type(obj) is InfoArray
        #
        # Note that it is here, rather than in the __new__ method,
        # that we set the default value for 'info', because this
        # method sees all creation of default objects - with the
        # InfoArray.__new__ constructor, but also with
        # arr.view(InfoArray).
        self.sample_freq = getattr(obj, 'sample_freq', None)
        self.meta = getattr(obj, 'meta', None)
        # We do not need to return anything class Signal(np.ndarray):
    """Timeseries signal class with (not-yet-implemented) configurable
    digitisation.

    see doc/subclassing.py in numpy code for details on subclassing ndarray
    """
    def __new__(cls, input_array):
        obj = np.asarray(input_array).view(cls).copy()
        obj.meta = PfMetaData()
        return obj

    def __array_finalize__(self,obj):
        self.meta = getattr(obj, 'meta', None)


    def n_channels(self):
        if self.ndim == 1:
            return 1
        else:
            return self.shape[0]

    def n_samples(self):
        if self.ndim == 1:
            return self.shape[0]
        else:
            return self.shape[1]

    def get_channel(self, channel_number):
        """allows us to use get_channel(0) no matter what ndim is"""
        if self.ndim == 1 and channel_number == 0:
            return self
        elif self.ndim > 1:
            return self[channel_number,:]
        else:
            raise ValueError

        
    #def add_channel(self, additional_array):
    ##### -- problems with resizing subclass of ndarray, need to be careful 
    #    additional_signal = Signal(additional_array)
    #    new_n_channels = self.n_channels() + additional_signal.n_channels()
    #    self.resize((new_n_channels, self.n_samples()))
    #    self[-additional_signal.n_channels():,:] = additional_signal class FlucStruc(BaseData):
    def __init__(self, svd_data, sv_list, timebase, min_dphase = -np.pi):
        # NOTE I'd prefer not to duplicate info here which is in svd_data - should be able to refer to that, once sqlalchemy is hooked in
        #self.topo_channels = svd_data.topo_channels
        self.channels = svd_data.channels
        #self.svs = sv_list
        if len(sv_list) == 1:
            self.a12 = 0
        else:
            # this ratio was inverted accidentally at first
            self.a12 = svd_data.svs[sv_list[1]]/svd_data.svs[sv_list[0]]
        self._binary_svs = list2bin(sv_list)
        # peak frequency for fluctuation structure
        self.timebase = timebase
        self.freq, self.freq_elmt = peak_freq(svd_data.chronos[sv_list[0]], self.timebase)
        self.t0 = timebase[0]
        # singular value filtered signals
        self.signal = np.dot(np.transpose(svd_data.topos[sv_list,:]),
                           np.dot(np.diag(svd_data.svs.take(sv_list)), svd_data.chronos[sv_list,:]))
        # phase differences between nearest neighbour channels
        self.dphase, self.fourier_values = self._get_dphase(min_dphase=min_dphase, get_fourier_value = 1)
        self.p = np.sum(svd_data.svs.take(sv_list)**2)/svd_data.E
        self.H = svd_data.H
        self.E = svd_data.E
        super(FlucStruc, self).__init__()

    def save(self):
        self.dphase.save()
        super(FlucStruc, self).save()

    def svs(self):
        return bin2list(self._binary_svs)

    def _get_dphase(self, min_dphase = -np.pi, get_fourier_value=0):
        """
        remap to [min_dphase, min_dphase+2pi]
        #SRH modified 24June2013 to return the complex value also
        """
        if get_fourier_value==0:
            phases = np.array([self._get_single_channel_phase(i) for i in range(self.signal.shape[0])])
        else:
            fourier_vals = []; phases = []
            for i in range(self.signal.shape[0]):
                phases_tmp, fourier_vals_tmp = self._get_single_channel_phase(i, get_fourier_value = 1)
                phases.append(phases_tmp)
                fourier_vals.append(fourier_vals_tmp)
            phases = np.array(phases)
            fourier_vals = np.array(fourier_vals)
        d_phases = remap_periodic(phases[1:]-phases[:-1], min_val = min_dphase)
        #d_phase_dataset = OrderedDataSet(ordered_by="channel_1.name")
        d_phase_dataset = BaseOrderedDataSet('d_phase_%s' %datetime.now())
        ## append then sort should be faster than ordereddataset.add() [ fewer sorts()]
        for i, d_ph in enumerate(d_phases):
            d_phase_dataset.append(FloatDelta(self.channels[i], self.channels[i+1], d_ph))

        #d_phase_dataset.sort()
        if get_fourier_value==0:
            return d_phase_dataset
        else:
            return d_phase_dataset, fourier_vals


    def _get_single_channel_phase(self, ch_id, get_fourier_value = 0):
        #SRH modified 24June2013 to return the complex value also
        data_fft = np.fft.fft(self.signal[ch_id])
        d_val = data_fft[self.freq_elmt]
        phase_val = np.arctan2(d_val.imag,d_val.real)
        if get_fourier_value == 0:
            return phase_val
        else:
            return phase_val, d_val def generate_timebase(t0=0.0, n_samples=1.e4, sample_freq=1.e6):
    sample_time = 1./sample_freq
    return Timebase(np.arange(t0, t0+n_samples*sample_time, sample_time)) class SVDData(BaseData):
    def __init__(self, chrono_labels, topo_channels, svd_input):
        """
        svd_input is a tuple as outputted by numpy.linalg.svd(data, 0)
        """
        self.channels = topo_channels
        self.chrono_labels = chrono_labels
        self.topos = np.transpose(svd_input[0])
        self.svs = svd_input[1]
        self.chronos = svd_input[2]
        self.E = sum(self.svs*self.svs)
        self.p = self.svs**2/self.E
        self.H = float((-1./np.log(len(self.svs)))*sum(self.p*np.log(self.p)))
        super(SVDData, self).__init__()

        
    def self_cps(self):
        try:
            return self._self_cps
        except AttributeError:
            self._self_cps = np.array([np.mean(cps(i,i)) for i in self.chronos])
            return self._self_cps def cps(a,b):
    return fft.fft(a)*conjugate(fft.fft(b)) def remap_periodic(input_array, min_val, period = 2*pi):
    while len(input_array[input_array<min_val]) > 0:
        input_array[input_array<min_val] += period
    while len(input_array[input_array>=min_val+period]) > 0:
        input_array[input_array>=min_val+period] -= period
    return input_array def peak_freq(signal,timebase,minfreq=0,maxfreq=1.e18):
    """
    TODO: old code: needs review
    this function only has a basic unittest to make sure it returns
    the correct freq in a simple case.
    """
    timebase = array(timebase)
    sig_fft = fft.fft(signal)
    sample_time = float(mean(timebase[1:]-timebase[:-1]))

    #SRH modification, frequencies seemed a little bit off because of the -1 in the denominator
    #Here we are trusting numpy....
    #fft_freqs = (1./sample_time)*arange(len(sig_fft)).astype(float)/(len(sig_fft)-1)
    fft_freqs = fft.fftfreq(len(sig_fft),d=sample_time)
    # only show up to nyquist freq
    new_len = len(sig_fft)/2
    sig_fft = sig_fft[:new_len]
    fft_freqs = fft_freqs[:new_len]
    [minfreq_elmt,maxfreq_elmt] = searchsorted(fft_freqs,[minfreq,maxfreq])
    sig_fft = sig_fft[minfreq_elmt:maxfreq_elmt]
    fft_freqs = fft_freqs[minfreq_elmt:maxfreq_elmt]
    
    peak_elmt = (argsort(abs(sig_fft)))[-1]
    return [fft_freqs[peak_elmt], peak_elmt] class BaseCoordTransform(object):
    """Base class does nothing useful at the moment"""
    input_coords = 'base_input'
    output_coords = 'base_output'

    def transform(self, coords):
        return coords @register("TimeseriesData", "DataSet")
def reduce_time(input_data, new_time_range):
    from pyfusion.data.base import DataSet
    if isinstance(input_data, DataSet):
        #output_dataset = input_data.copy()
        #output_dataset.clear()
        output_dataset = DataSet(input_data.label+'_reduce_time')
        for data in input_data:
            try:
                output_dataset.append(data.reduce_time(new_time_range))
            except AttributeError:
                pyfusion.logger.warning("Data filter 'reduce_time' not applied to item in dataset")
        return output_dataset

    new_time_args = searchsorted(input_data.timebase, new_time_range)
    input_data.timebase =input_data.timebase[new_time_args[0]:new_time_args[1]]
    if input_data.signal.ndim == 1:
        input_data.signal = input_data.signal[new_time_args[0]:new_time_args[1]]
    else:
        input_data.signal = input_data.signal[:,new_time_args[0]:new_time_args[1]]
    return input_data def orm_register():
    def reg_item(orm_func):
        pyfusion.orm_manager.add_reg_func(orm_func)
        return orm_func
    return reg_item class FakeDataAcquisition(BaseAcquisition):
    """Acquisition class for generating fake data for testing purposes."""
    pass class PfTestBase(object):
    """Base class for generated sql and non-sql test cases."""
    pass
from numpy.testing import assert_array_almost_equal, assert_almost_equal
from pyfusion.test.tests import PfTestBase
from pyfusion.data.base import BaseData, DataSet, Coords, ChannelList, Channel, FloatDelta, BaseOrderedDataSet
from pyfusion.data.timeseries import TimeseriesData, Timebase, Signal, FlucStruc, generate_timebase, SVDData
from pyfusion.data.utils import cps, remap_periodic, peak_freq
from pyfusion.data.base import BaseCoordTransform
from pyfusion.data.filters import reduce_time
from pyfusion.orm.utils import orm_register
from pyfusion.acquisition.FakeData.acq import FakeDataAcquisition
from pyfusion.test.tests import PfTestBase
                from sqlalchemy import Table, Column, Integer, ForeignKey
                from sqlalchemy.orm import mapper
import numpy as np
import pyfusion
        self.assertEqual(fs_data.meta, multichannel_data.meta)

        # save our dataset to the database
        fs_data.save()

        if pyfusion.orm_manager.IS_ACTIVE:
            session = pyfusion.orm_manager.Session()
            some_ds = session.query(DataSet).all().pop()
            self.assertEqual(some_ds.meta, multichannel_data.meta)

        #print some_ds.meta
        #assert False
    
class CheckStoredMetaData(PfTestBase):
    def test_stored_metadata_data(self):
        """ metadata should be stored to data instances, rather than datasets - this might be slower, but more likely to guarantee data is kept track of."""
        n_ch = 3
        n_samples = 1024
        multichannel_data = get_multimode_test_data(channels=get_n_channels(n_ch),
                                                    timebase = Timebase(np.arange(n_samples)*1.e-6),
                                                    noise = 0.01)




        # put in some fake metadata 
        multichannel_data.meta = {'hello':'world'}
        print multichannel_data.meta

        # produce a dataset of flucstrucs
        fs_data = multichannel_data.flucstruc(min_dphase = -2*np.pi)

        # check that metadata is carried to the individual flucstrucs
        for fs in fs_data:
            self.assertEqual(fs.meta, multichannel_data.meta)

        # save our dataset to the database
        fs_data.save()

        ## now test to make sure metadata is saved in database

        if pyfusion.orm_manager.IS_ACTIVE:
            session = pyfusion.orm_manager.Session()
            some_fs = session.query(FlucStruc).all().pop()
            self.assertEqual(some_fs.meta, multichannel_data.meta)





class CheckSciPyFilters(PfTestBase):
    def test_sp_filter_butterworth_bandpass(self):
        n_ch = 3
        n_samples = 1024
        sample_period = 1.e-6
        # Let's generate a test signal with strong peaks at 20kHz and
        # 60kHz and a weaker peak at 40kHz. 
        
        mode_20kHz = {'amp': 10.0, 'freq': 20.0e3, 'mode_number':3, 'phase':0.2}
        mode_40kHz = {'amp':  2.0, 'freq': 40.0e3, 'mode_number':4, 'phase':0.3}
        mode_60kHz = {'amp': 10.0, 'freq': 60.0e3, 'mode_number':5, 'phase':0.4}

        multichannel_data = get_multimode_test_data(channels=get_n_channels(n_ch),
                                                    timebase = Timebase(np.arange(n_samples)*sample_period),
                                                    modes = [mode_20kHz, mode_40kHz, mode_60kHz],
                                                    noise = 0.1)

        filtered_data = multichannel_data.sp_filter_butterworth_bandpass([35.e3, 45.e3], [25.e3, 55.e3], 1.0, 10.0)


class CheckFlucstrucPhases(PfTestBase):
    """Test code to replicate bug found by Shaun

    Pyfusion uses different syntax depending if sql is enabled.
    """

    def test_flucstruc_phases(PfTestCase):
        
        n_ch = 10
        n_samples = 5000
        timebase = Timebase(np.arange(n_samples)*1.e-6)
        channels = ChannelList(*(Channel('ch_%d' %i, Coords('cylindrical',(1.0,i,0.0))) for i in 2*np.pi*np.arange(n_ch)/n_ch))
        multichannel_data = get_multimode_test_data(channels = channels,
                                                    timebase = timebase,
                                                    noise = 0.5)

        data_reduced_time=multichannel_data.reduce_time([0,0.002]).subtract_mean().normalise(method='v',separate=True)

        fs_set=data_reduced_time.flucstruc()
        phases = []
        for fs in fs_set:
            for j in range(0,len(fs.dphase)):
                phases.append(fs.dphase[j].delta)



class CheckFilterCopy(PfTestBase):
    """Check that by default, data filters alter a copy of the input data object not the object itself."""

    def test_timeseries_filter_copy(self):
        # Use reduce_time filter for testing...
        n_ch = 10
        n_samples = 5000
        timebase = Timebase(np.arange(n_samples)*1.e-6)
        channels = ChannelList(*(Channel('ch_%d' %i, Coords('cylindrical',(1.0,i,0.0))) for i in 2*np.pi*np.arange(n_ch)/n_ch))
        multichannel_data = get_multimode_test_data(channels = channels,
                                                    timebase = timebase,
                                                    noise = 0.5)
        new_data = multichannel_data.reduce_time([0,1.e-3])
        self.assertFalse(new_data is multichannel_data)

    def test_timeseries_filter_nocopy(self):
        # Use reduce_time filter for testing...
        n_ch = 10
        n_samples = 5000
        timebase = Timebase(np.arange(n_samples)*1.e-6)
        channels = ChannelList(*(Channel('ch_%d' %i, Coords('cylindrical',(1.0,i,0.0))) for i in 2*np.pi*np.arange(n_ch)/n_ch))
        multichannel_data = get_multimode_test_data(channels = channels,
                                                    timebase = timebase,
                                                    noise = 0.5)

        new_data = multichannel_data.reduce_time([0,1.e-3], copy=False)