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
"""Tests for the Data and related classes."""





DEFAULT_N_CHANNELS = 10
DEFAULT_T0 = 0.00
DEFAULT_T1 = 0.01
DEFAULT_SAMPLE = 1.0e-5
DEFAULT_TIMEBASE = Timebase(np.arange(DEFAULT_T0, DEFAULT_T1, DEFAULT_SAMPLE))
DEFAULT_NOISE = 0.2

mode_1 = {'amp': 0.7, 'freq': 24.0e3, 'mode_number':3, 'phase':0.2}
mode_2 = {'amp': 0.5, 'freq': 37.0e3, 'mode_number':4, 'phase':0.3}
mode_3 = {'amp': 0.5, 'freq': 27.0e3, 'mode_number':2, 'phase':0.1}
mode_4 = {'amp': 0.7, 'freq': 24.0e3, 'mode_number':7, 'phase':3.2}
mode_5 = {'amp': 1.0, 'freq': 37.0e3, 'mode_number':9, 'phase':1.3}

###############################################################################
## Convenience functions                                                     ##
###############################################################################

def get_n_channels(n_ch):
    """Return a list of n_ch channels."""
    poloidal_coords = 2*np.pi*np.arange(n_ch)/n_ch
    channel_gen = (Channel('ch_%02d' %i, Coords('cylindrical', (1.0,i,0.0)))
                   for i in poloidal_coords)
    return ChannelList(*channel_gen)


def get_multimode_test_data(channels = get_n_channels(DEFAULT_N_CHANNELS),
                            timebase = DEFAULT_TIMEBASE,
                            modes = [mode_1, mode_2], noise = DEFAULT_NOISE):
    """Generate synthetic multi-channel data for testing."""
    n_channels = len(channels)
    data_size = (n_channels, timebase.size)
    data_array = noise*2*(np.random.random(data_size)-0.5)
    timebase_matrix = np.resize(timebase, data_size)
    angle_matrix = np.resize(np.array([i.coords.cylindrical[1] for i in channels]),
                          data_size[::-1]).T
    for m in modes:
        data_array += m['amp']*np.cos(2*np.pi*m['freq']*timebase_matrix +
                                   m['mode_number']*angle_matrix + m['phase'])
    output = TimeseriesData(timebase=timebase, signal=Signal(data_array),
                            channels=channels)
    return output

###############################################################################
## End of convenience functions                                              ##
###############################################################################

###############################################################################
## Tests for pyfusion/data/utils.py                                          ##
###############################################################################

class CheckUtils(PfTestBase):
    """Test the helper functions in pyfusion.data.utils.py"""
    def test_remap_periodic(self):
        data = np.array([-3,-2,-1,0,1,2,2.5, 3])
        output = remap_periodic(data, min_val = 0, period=3)
        expected = np.array([0, 1, 2, 0, 1, 2, 2.5, 0])
        assert_array_almost_equal(output, expected)

    def test_peak_freq(self):
        timebase = Timebase(np.arange(0.0,0.01, 1.e-6))
        single_mode_signal = get_multimode_test_data(channels=get_n_channels(1),
                                                     timebase=timebase,
                                                     modes = [mode_3])
        p_f, p_f_elmt = peak_freq(single_mode_signal.signal[0],
                                  single_mode_signal.timebase)
        # Check that we get mode_3 frequency of 27.0 kHz (to 1 decimal place).
        self.assertAlmostEqual(1.e-3*p_f, 1.e-3*mode_3['freq'], 1)


###############################################################################
## End of tests for pyfusion/data/utils.py                                   ##
###############################################################################

###############################################################################
## Tests for pyfusion/data/base.py                                           ##
###############################################################################

class DummyCoordTransform(BaseCoordTransform):
    """Minimal coordinate transform class for testing.

    Transforms cylindrical coordinates (a,b,c) to dummy coordinates (2a,3b,4c)

    """
    input_coords = 'cylindrical'
    output_coords = 'dummy'

    def transform(self, coords):
        return (2*coords[0], 3*coords[1], 4*coords[2])



class CheckCoordinates(PfTestBase):