class MDSPlusDataFetcher(BaseDataFetcher):
     """Determine which access mode should be used, and fetch the MDSplus data."""

     def setup(self):
          self.mds_path_components = get_tree_path(self.mds_path)
          if hasattr(self.acq, '%s_path' %self.mds_path_components['tree']):
               self.tree = MDSplus.Tree(self.mds_path_components['tree'],
                                        self.shot)
               self.fetch_mode = 'path'
          elif self.acq.server_mode == 'mds':
               self.acq.connection.openTree(self.mds_path_components['tree'],
                                            self.shot)
               self.fetch_mode = 'thin client'
          elif self.acq.server_mode == 'http':
               self.fetch_mode = 'http'
          else:
               raise Exception('Cannot determine MDSPlus fetch mode')

     def do_fetch(self):
          # TODO support non-signal datatypes
          if self.fetch_mode == 'thin client':
               ch = Channel(self.mds_path_components['nodepath'],
                            Coords('dummy', (0,0,0)))
               data = self.acq.connection.get(self.mds_path_components['nodepath'])
               dim = self.acq.connection.get('dim_of(%s)' %self.mds_path_components['nodepath'])
               # TODO: fix this hack (same hack as when getting signal from node)
               if len(data.shape) > 1:
                    data = np.array(data)[0,]
               if len(dim.shape) > 1:
                    dim = np.array(dim)[0,]
               output_data = TimeseriesData(timebase=Timebase(dim),
                                            signal=Signal(data), channels=ch)
               output_data.meta.update({'shot':self.shot})
               return output_data

          elif self.fetch_mode == 'http':
               data_url = self.acq.server + '/'.join([self.mds_path_components['tree'],
                                                      str(self.shot),
                                                      self.mds_path_components['tagname'],
                                                      self.mds_path_components['nodepath']])
               
               data = mdsweb.data_from_url(data_url)
               ch = Channel(self.mds_path_components['nodepath'], Coords('dummy', (0,0,0)))
               t = Timebase(data.data.dim)
               s = Signal(data.data.signal)
               output_data = TimeseriesData(timebase=t, signal=s, channels=ch)
               output_data.meta.update({'shot':self.shot})
               return output_data

          else:
               #print self.mds_path
               node = self.tree.getNode(self.mds_path)
               #print node.dtype
               #if int(node.dtype) == 195:
               if (int(node.dtype) == 195) or (int(node.dtype) == 192):#Shaun hack to work with PMT data
                    return get_tsd_from_node(self, node)
               else:
                    raise Exception('Unsupported MDSplus node type')

     def pulldown(self):
          if self.fetch_mode == 'thin client':
               self.acq.connection.closeTree(self.mds_path_components['tree'], self.shot) def get_tree_path(path_string):
     """Use regex to extract mdsplus tree, tag and node from full path"""
     components = mds_path_regex.search(path_string)
     ret_dict = {'tree':components.group('tree'),
                 'tagname':components.group('tagname'),
                 'nodepath':components.group('nodepath')}
     return ret_dict class MultiChannelFetcher(BaseDataFetcher):
    """Fetch data from a diagnostic with multiple timeseries channels.

    This fetcher requres a multichannel configuration section such as::

     [Diagnostic:H1_mirnov_array_1]
     data_fetcher = pyfusion.acquisition.base.MultiChannelFetcher
     channel_1 = H1_mirnov_array_1_coil_1
     channel_2 = H1_mirnov_array_1_coil_2
     channel_3 = H1_mirnov_array_1_coil_3
     channel_4 = H1_mirnov_array_1_coil_4

    The channel  names must be  `channel\_` followed by an  integer, and
    the channel  values must correspond to  other configuration sections
    (for        example       ``[Diagnostic:H1_mirnov_array_1_coil_1]``,
    ``[Diagnostic:H1_mirnov_array_1_coil_1]``, etc)  which each return a
    single               channel               instance               of
    :py:class:`~pyfusion.data.timeseries.TimeseriesData`.
    """
    def ordered_channel_names(self):
        """Get an ordered list of the channel names in the diagnostic

        :rtype: list
        """
        channel_list = []
        for k in self.__dict__.keys():
            if k.startswith('channel_'):
                channel_list.append(
                    [int(k.split('channel_')[1]), self.__dict__[k]]
                    )
        channel_list.sort()
        return [i[1] for i in channel_list]
    
    def fetch(self, interp_if_diff = True):
        """Fetch each  channel and combine into  a multichannel instance
        of :py:class:`~pyfusion.data.timeseries.TimeseriesData`.

        :rtype: :py:class:`~pyfusion.data.timeseries.TimeseriesData`
        """
 
        ## initially, assume only single channel signals
        ordered_channel_names = self.ordered_channel_names()
        data_list = []
        channels = ChannelList()
        timebase = None
        meta_dict={}
        for chan in ordered_channel_names:
            fetcher_class = import_setting('Diagnostic', chan, 'data_fetcher')
            tmp_data = fetcher_class(self.acq, self.shot,
                                     config_name=chan).fetch()
            channels.append(tmp_data.channels)
            meta_dict.update(tmp_data.meta)
            if timebase == None:
                timebase = tmp_data.timebase
                data_list.append(tmp_data.signal)
            else:
                try:
                    assert_array_almost_equal(timebase, tmp_data.timebase)
                    data_list.append(tmp_data.signal)
                except:
                    if interp_if_diff:
                        data_list.append(np.interp(timebase, tmp_data.timebase, tmp_data.signal))
                    else:
                        raise
        signal=Signal(data_list)
        output_data = TimeseriesData(signal=signal, timebase=timebase,
                                     channels=channels)
        #output_data.meta.update({'shot':self.shot})
        output_data.meta.update(meta_dict)
        return output_data 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 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 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 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 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 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) def get_coords_for_channel(channel_name=None, **kwargs):
    config_dict = kwargs.copy()
    if channel_name:
        config_dict.update(get_config_as_dict('Diagnostic', channel_name))
    coord_name = 'dummy'
    coord_values = (0.0,0.0,0.0)
    transforms = []
    for k in config_dict.keys():
        if k.startswith('coords_'):
            coord_name = k[7:]
            coord_values = tuple(map(float,config_dict[k].split(',')))
    coords_instance = Coords(coord_name, coord_values)
    if config_dict.has_key('coord_transform'):
        transform_list = pyfusion.config.pf_options('CoordTransform', config_dict['coord_transform'])
        for transform_name in transform_list:
            transform_class_str = pyfusion.config.pf_get('CoordTransform', config_dict['coord_transform'], transform_name)
            transform_class = import_from_str(transform_class_str)
            coords_instance.load_transform(transform_class)

    return coords_instance def import_setting(component, component_name, setting):
    """Attempt to import and return a config setting."""
    value_str = pyfusion.config.pf_get(component, component_name, setting)
    return import_from_str(value_str) def kwarg_config_handler(component_type, component_name, **kwargs):
    for config_var in pyfusion.config.pf_options(component_type, component_name):
            if not config_var in kwargs.keys():
                kwargs[config_var] = pyfusion.config.pf_get(component_type,
                                                   component_name, config_var)
    return kwargs def get_config_as_dict(component_type, component_name):
    config_option_list = pyfusion.config.pf_options(component_type, component_name)
    config_map = lambda x: (x, pyfusion.config.pf_get(component_type, component_name, x))
    return dict(map(config_map, config_option_list)) def import_from_str(string_value):
    # TODO: make shortcuts for loading from within pyfusion
    split_val = string_value.split('.')
    val_module = __import__('.'.join(split_val[:-1]),
                            globals(), locals(),
                            [split_val[-1]])
    return val_module.__dict__[split_val[-1]]
import pyfusion, os
import pyfusion.acquisition.MDSPlus.h1ds as mdsweb
import traceback
import numpy as np
from pyfusion.acquisition.MDSPlus.fetch import MDSPlusDataFetcher, get_tree_path
from pyfusion.acquisition.base import MultiChannelFetcher
from pyfusion.data.timeseries import TimeseriesData, Signal, Timebase
from pyfusion.data.base import Coords, Channel, ChannelList, \
    get_coords_for_channel
from pyfusion.conf.utils import import_setting, kwarg_config_handler, \
     get_config_as_dict, import_from_str
        from scipy.io import netcdf
"""Subclass of MDSplus data fetcher to grab additional H1-specific metadata."""

class DIIIDDataFetcherPTdata(MDSPlusDataFetcher):
    """Subclass of MDSplus fetcher to get additional H1-specific metadata."""
    def setup(self):
        pass

    def do_fetch(self):
        print(self.pointname)
        print(self.shot)
        if self.NC!=None:
            print(self.NC)
            t_name = '{}_time'.format(self.pointname)
            NC_vars = self.NC.variables.keys()
            if self.pointname in NC_vars:
                print('Reading cache!!!!')
                t_axis = self.NC.variables[t_name].data[:].copy()
                data = self.NC.variables[self.pointname].data[:].copy()
        else:
            tmp = self.acq.connection.get('ptdata2("{}",{})'.format(self.pointname, self.shot))
            data = tmp.data()
            tmp = self.acq.connection.get('dim_of(ptdata2("{}",{}))'.format(self.pointname, self.shot))
            t_axis = tmp.data()
            self.write_cache = True
        print(t_axis)
        print(data)

        coords = get_coords_for_channel(**self.__dict__)