"""
Tokdata MDSPLUS  MdsPlus interface to tokamak data
"""
import MDSplus
import numpy as np
from v2d import V2d, NamedV2d

class _Connections(object):
    """Class that keeps a dictionary of all the open connections.
    It contains a class attribute "_connection" which is a dictionary with the
    presumably open connections, and a member attribute "server" with the
    name of the actual server for the current object. 
    
    During initialization the server name is stored in the object. 
    When a connection is opened that server name is stored in the dictionary
    (which is the same for alla objects) and the connection is opened.
    """
    _connection = {}
    
    def __init__(self, server):
        self.server = server
        
    def openconnection(self):
        """Open the connecion to the server. . 
        """
        #return self._openconnection(self.server)
        if self.server not in self._connection:
            self._connection[self.server] = MDSplus.Connection(self.server)

        return self._connection[self.server]

    def closeconnection(self):
        """Close the connection to the server
        """
        del self._connection[self.server]
        
class Ftudata(NamedV2d):
    """This class implement the reading of FTU data. Additional fields with
    the additional attributes present on FTU data  are also exposed
    """
    @staticmethod
    def _decodelabel2d(xl):
        """Decode the label for 2D FTU channels
        """
        xl = xl.split('(')
        rlabel = xl[0][0:10]
        tlabel = xl[0][10:]
        runit = xl[1][0:6]
        tunit = xl[1][6:-1]
        return (rlabel, runit, tlabel, tunit)
        
    def __init__(self, shot=None, channel=None, imode=0):
        super(Ftudata, self).__init__()
        self.shot = shot
        self.channel = channel
        
        self.imode = imode

        if shot is None or channel is None:
            return

#        self.connection = _Connections('ftucluster.frascati.enea.it:8100')
        self.connection = _Connections('fusc06.frascati.enea.it:8100')
        con = self.connection.openconnection()
        
        status = con.get('ftustatus($1,$2)', shot, channel)
        self.status = status
        if status == 0 or status == 2:
            sig = con.get('_sig=ftu($1,$2,0)', shot, channel).value
            x_data = con.get('dim_of(_sig)').value
            expr_xl = con.get('ftuxl($1,$2)', shot, channel)
            expr_yl = con.get('ftuyl($1,$2)', shot, channel)
            if status == 0:
                self.v = sig
                self.t = x_data
                self.r = None

                self.tax.parse1d(expr_xl)
                self.vax.parse1d(expr_yl)
            else:
                nr = x_data[-2]
                nt = x_data[-1]
                r = x_data[0:nr]
                t = x_data[nr:nr+nt]
                self.v = np.reshape(sig, [nt, nr])
                self.r = r
                self.t = t
                
                decoded_xl = self._decodelabel2d(expr_xl)
                self.rax.addlu(*decoded_xl[0:2])
                self.tax.addlu(*decoded_xl[2:])
                self.vax.parse1d(expr_yl)
        else:
            self.v = np.array([])
            self.t = np.array([])
            self.r = np.array([])
                           
    def _getacopy(self, *arg, **args):
        other = super(Ftudata, self)._getacopy(*arg, **args)
        if other is not self:
            other.shot = self.shot 
            other.channel = self.channel
            other.imode = self.imode
        return other
        
    def __repr__(self):
        basestr = super(Ftudata, self).__repr__()
        rstr = []
        rstr += [" shot:    %d" % self.shot]
        rstr += [" channel: %s" % self.channel]
        
        return "\n".join(rstr)+"\n"+basestr
            
class Jetdata(NamedV2d):
    """This class implement the reading of JET data. Additional fields with
    the additional attributes present on JET data  are also exposed (only if
    they are available ny MDSplus)
    """
        
    def __init__(self, shot=None, channel=None):
        super(Jetdata, self).__init__()
        self.shot = shot
        self.channel = channel
        
        if shot is None or channel is None:
            return
            
        self.connection = _Connections('mdsplus.jet.efda.org')
        con = self.connection.openconnection()

        _channel = 'ppf/'+channel.lower()
        
        sig = con.get('_sig=jet($1,$2)', _channel, shot).value
        
        if len(sig.shape)==2:
            r = con.get('dim_of(_sig,0)').value
            t = con.get('dim_of(_sig,1)').value
        elif len(sig.shape)==1:
            t = con.get('dim_of(_sig,0)').value
            r = None
        else:
            t = None
            r = None
                
        self.v = sig
        self.r = r
        self.t = t

    def _getacopy(self, *arg, **args):
        other = super(Jetdata, self)._getacopy(*arg, **args)
        if other is not self:
            other.shot = self.shot 
            other.channel = self.channel
        return other
        
    def __repr__(self):
        basestr = super(Jetdata, self).__repr__()
        rstr = []

        rstr += [" shot: {0}".format(self.shot)]
        rstr += [" channel: {0}".format(self.channel)]
        return "\n".join(rstr)+"\n\n"+basestr
        
class Asdexdata(NamedV2d):
    """This class implement the reading of JET data. Additional fields with
    the additional attributes present on JET data  are also exposed (only if
    they are available ny MDSplus)
    """
        
    def __init__(self, shot=None, channel=None):
        super(Asdexdata, self).__init__()
        self.shot = shot
        self.channel = channel
        
        if shot is None or channel is None:
            return
            
        self.connection = _Connections('localhost:8001')
        con = self.connection.openconnection()

        _channel = channel
        
        sig = con.get('_sig=augsignal($1,$2)', shot, _channel).value
        
        if len(sig.shape)==2:
            r = con.get('dim_of(_sig,0)').value
            t = con.get('dim_of(_sig,1)').value
        elif len(sig.shape)==1:
            t = con.get('dim_of(_sig,0)').value
            r = None
        else:
            t = None
            r = None
                
        self.v = sig
        self.r = r
        self.t = t

    def _getacopy(self, *arg, **args):
        other = super(Jetdata, self)._getacopy(*arg, **args)
        if other is not self:
            other.shot = self.shot 
            other.channel = self.channel
        return other
        
    def __repr__(self):
        basestr = super(Jetdata, self).__repr__()
        rstr = []

        rstr += [" shot: {0}".format(self.shot)]
        rstr += [" channel: {0}".format(self.channel)]
        return "\n".join(rstr)+"\n\n"+basestr
        
class Eastdata(NamedV2d):
    """This class implement the reading of JET data. Additional fields with
    the additional attributes present on JET data  are also exposed (only if
    they are available ny MDSplus)
    """
    _open_tree = []
    
    def __init__(self, shot=None, channel=None):
        super(Eastdata, self).__init__()
        self.shot = shot
        self.channel = channel
        
        if shot is None or channel is None:
            return
            
        self.connection = _Connections('202.127.204.41')
        con = self.connection.openconnection()

        if self._open_tree != [shot]:
            #print "Here A"
            if len(self._open_tree)>0:
                #print "Here B"
                con.closeTree('pcs_east', self._open_tree[0])
                self._open_tree[0] = shot
            else:
                self._open_tree.append(shot)
            con.openTree('pcs_east',shot)
            
            
        _channel = '\\'+channel
        
        sig = con.get(_channel).value
        
        if len(sig.shape)==2:
            r = con.get('dim_of('+_channel+',0)').value
            t = con.get('dim_of('+_channel+',1)').value
        elif len(sig.shape)==1:
            t = con.get('dim_of('+_channel+',0)').value
            r = None
        else:
            t = None
            r = None

        #con.closeTree('pcs_east',shot)
        
        self.v = sig
        self.r = r
        self.t = t

    def _getacopy(self, *arg, **args):
        other = super(Eastdata, self)._getacopy(*arg, **args)
        if other is not self:
            other.shot = self.shot 
            other.channel = self.channel
        return other
        
    def __repr__(self):
        basestr = super(Eastdata, self).__repr__()
        rstr = []

        rstr += [" shot: {0}".format(self.shot)]
        rstr += [" channel: {0}".format(self.channel)]
        return "\n".join(rstr)+"\n\n"+basestr