from xml.etree import ElementTree
from pySAXS.tools import xmltools
import pySAXS.LS.SAXSparametersXML as SAXSparameters
from pySAXS.models import model
from numpy import *
from scipy import interpolate
from pySAXS.guisaxs import pySaxsColors
import numpy

'''
Dataset class for datas
'''

class dataset:
    def __init__(self,name=None,q=None,i=None,filename=None,\
                 checked=True,model=None,rawdata_ref=None,\
                 type=None,parameters=None,error=None,comment=None,\
                 parent=None,child=None,parentformula=None,variableDict=None,color=None,\
                 image=None,background_data=None,background_value=None,rpt=None,abs=None):
        self.name=name
        if filename is None:
            self.filename=name
        else:
            self.filename=filename
        
        self.checked=checked
        if checked is not None:
            #print checked
            self.checked=(str(checked).lower()=='true')
        
        self.q=q
        self.i=i
        self.error=error
        self.model=model
        self.rawdata_ref=rawdata_ref
        self.type=type
        self.parameters=parameters
        self.comment=comment
        self.parent=parent#[data1,data2,...]
        self.child=child 
        self.parentformula=parentformula #"i1+i2+i3"
        #self.parentvariables=parentvariablesdict #['i1','i2',...]
        self.variableDict=variableDict#{'i0':'data1','i1':'data2',...}
        
        #if color is None:
        #    color=pySaxsColors.pySaxsColors().getColor() #get a new color
            
        self.color=color
        self.image=image
        self.background_data=background_data
        self.background_value=background_value
        self.rpt=rpt
        self.abs=abs

    def getQIE(self):
        return self.q, self.i, self.error

    def _check(self):
        if self.checked is not None:
            self.checked=(str(self.checked).lower()=='true')
        

    def _copy(self):
        """
        return a copy without model and parameters
        use in saveDataDictRaw
        """
        new=dataset(name=self.name,q=self.q,i=self.i,filename=self.filename,checked=self.checked,\
                    rawdata_ref=self.rawdata_ref,type=self.type,error=self.error,comment=self.comment,\
                    model=None, parameters=None)
        #new=copy(self)
        return new

    def _deepcopy(self):
        """
        return a deepcopy of the dataset object
        """
        from copy import  deepcopy
        return deepcopy(self)
    
    def _evaluateFromParent(self,data_dict,qref=None):
        '''
        evaluate i from parent
        datadict : a dictionnary of dataset 
        '''
        #from numpy import *
        if qref is None:
            qref=self.q
        #print self.name
        if self.parentformula is None:
            return
        newdict={}
        newerror=zeros(shape(qref))
        #--convert variableDict
        for var in self.variableDict:
            name=self.variableDict[var]
            #print name
            if not(name in data_dict):
                return "Data are missing for "+name+" with formula : "+self.parentformula+". Data will not be recalculated."
                #return data_dict[name].i
                
            #variableDict contain variable name and dataset name
            i=data_dict[name].i
            q=data_dict[name].q
            if str(q)!=str(qref):
                #self.printTXT("trying interpolation for ",name)
                newf=interpolate.interp1d(q,i,kind='linear',bounds_error=0)
                newi=newf(qref)
            else:
                newi=i
                #addition for errors
                error=data_dict[name].error
                if error is not None and newerror is not None:
                    newerror+=error
                else:
                    newerror=None
            newdict[var]=newi
        #--evaluate
        #self.printTXT("trying evaluation of ",formula)
        safe_list = ['math','acos', 'asin', 'atan', 'atan2', 'ceil', 'cos', 'cosh', 'degrees', \
                     'e', 'exp', 'fabs', 'floor', 'fmod', 'frexp', 'hypot', 'ldexp', 'log',\
                     'log10', 'modf', 'pi', 'pow', 'radians', 'sin', 'sinh', 'sqrt', 'tan', 'tanh'] #use the list to filter the local namespace safe_dict = dict([ (k, locals().get(k, None)) for k in safe_list ])
        for k in safe_list:
            newdict[k]=locals().get(k,None)
        iout=array(eval(self.parentformula,newdict))
        self.q=qref
        self.i=iout
        self.error=newerror
        return ""
    
    def _xml(self):
        '''
        return a xml element
        '''
        #create the root </root><root>
        root_element = ElementTree.Element("data",name=self.name)
        for attributeName in dir(self):
            #print attributeName
            if attributeName[0]!='_':
                #class attribute
                #child = ElementTree.SubElement(root_element, attributeName)
                val=getattr(self,attributeName)
                if val is not None:
                    if attributeName=="parameters":
                        child=val.xml()
                    elif attributeName=="model":
                        child=val.xml()
                    else:
                        child = ElementTree.Element(attributeName,datatype=xmltools.getDatatype(val))
                        if xmltools.getDatatype(val).find('array')>=0:#val.__class__==array(1).__class__:
                            #attribute is an array
                            child.text = xmltools.data2string(val)
                        else:
                            child.text=str(val)
                    #now append
                    root_element.append(child)
        return root_element

# ---------------------------------------------------------------------------------- 
    
def getDatasetFromXML(element):
    '''
    set the dataset with the parsed xml as Element
    
    <root>
        <name> mqlkqsmldk</name>
        <checked>true</checked>
        <q datatype='numpy.array'>0,0.1,0.2,0.3</q> if it is an array datatype must contain 'array'
        ...
    </root>
    '''
    dt=dataset()
    for subelement in element:
        #get all element in root element
        tag=subelement.tag #get tag ie <filename> or <q>
        #print tag
        if hasattr(dt,tag) and tag[0]!="_":
            if tag=="parameters":
                #print "TAG parameters"
                val=SAXSparameters.SAXSparameters()
                val.getFromXML(subelement)
            elif tag=="model":
                #print "TAG Model"
                val=model.getModelFromXML(subelement)
                    
            else:
                #tag exist in dataset class -> get attrib and text
                attrib=subelement.attrib
                #print attrib
                text=subelement.text
                val=text
                # attrib has 'datatype' ?
                if 'datatype' in attrib:
                    datatype=attrib['datatype']
                    #check datatype 
                    val=xmltools.convertText(text,datatype)
                
                #IMPORTANT set value of attibute
            setattr(dt,tag,val)
    try:
        if len(dt.q)>len(dt.i):
            #print "lecture error"
            dt.q=dt.q[:len(dt.i)]
        if len(dt.q)<len(dt.i):
            #print "lecture error"
            dt.i=dt.i[:len(dt.q)]
    except:
        #print("error when reading %s"%str(subelement.tag))
        pass
            
    return dt

def getDatasetFromXMLFile(filename):
    element=ElementTree.parse(filename)
    root=element.getroot()
    return getDatasetFromXML(root)

def saveDatasetOnXMLFile(filename,dataset):
    xmltools.data2XMLfile(dataset,filename,name="data")

def getQIEFromXMLFile(filename,dataname):
    dt=getDataDictFromXMLFile(filename)
    #return dt[dataname].getQIE()
    return dt[dataname].q, dt[dataname].i, dt[dataname].error

def saveDataDictOnXMLFile(filename, data_dict):
    root_element = ElementTree.Element("dataset",name=filename)
    for key,value in list(data_dict.items()):
        #print("---",key,value.checked)
        #value=data_dict[key]
        el=value._xml()
        #print ElementTree.tostring(el)
        root_element.append(el)
    tree=ElementTree.ElementTree()
    tree._setroot(root_element)
    tree.write(filename)
        
def getDataDictFromXMLFile(filename):
    element=ElementTree.parse(filename)
    root=element.getroot()
    data_dict={}
    for subelement in list(root):
        #get all element in root element
        tag=subelement.tag #get tag ie <data> or <q>
        #print tag
        if tag=='data':
            dataset_name=subelement.attrib["name"]
            #new dataset
            data_dict[dataset_name]=getDatasetFromXML(subelement)
            #print list(subelement)
    return data_dict
                
def saveDataDictRaw(filename,data_dict):
    '''
    save dataset in old format
    '''
    import pickle
    f=open(filename,mode='w')
    newdict={}
    for key, datas in list(data_dict.items()):
        newdict[key]=datas._copy()
    pickle.dump(newdict,f)
    f.close()
        
def getDataDictRaw(filename):
    '''
    get dataset from old format
    '''
    import pickle
    f=open(filename,mode='r')
    data_dict=pickle.load(f)
    f.close()
    return data_dict

def getListOfChecked(data_dict):
    '''
    return the list of the checked elements
    '''
    l=[]
    for key,value in list(data_dict.items()):
        if value.checked:
            l.append(key)
    return l

def dataset_merging(dataset0,dataset1,cut1,cut2):
    #check if curve are in the good order
    if dataset0.q[0]>dataset1.q[0]:
        datasettemp=dataset1
        dataset1=dataset0
        dataset0=datasettemp
    if cut1>cut2:
        cuttemp=cut2
        cut2=cut1
        cut1=cuttemp

    q0=dataset0.q
    i0=dataset0.i
    e0=dataset0.error

    q1 = dataset1.q
    i1 = dataset1.i
    e1 = dataset1.error

    q_f,i_f,e_f=merging(q0, i0, e0, q1, i1, e1, cut1, cut2)
    d=dataset(q=q_f,i=i_f,error=e_f)
    return d

def merging(q0, i0, e0, q1, i1, e1, cut1, cut2):
    '''
    merging two iq curve
    mean between cut1 & cut2
    '''
    # print("start of merge : from %6.4f"%(q0[0]))
    # cut1=dataArray[j+1][5]
    # print("cut1 %6.4f"%cut1)
    # cut2=dataArray[j][6]
    # print("cut2 %6.4f"%cut2)
    # first datas
    # print("first data")
    # print("end of merge : from %6.4f"%(dataArray[j][1][-1]))
    # 1 part from actual qmin to cut1
    q_final, i_final, e_final = clipQmax(q0, i0, e0, cut1)

    # 2 part from next qmin large to actual qmax merging
    # to be merge with next datas
    q2, i2, e2 = clipQmin(q0, i0, e0, cut1)
    q2, i2, e2 = clipQmax(q2, i2, e2, cut2)

    # 3 part from second qmin large to
    q3, i3, e3 = clipQmin(q1, i1, e1, cut1)
    q3, i3, e3 = clipQmax(q3, i3, e3, cut2)
    # MEAN
    newf = interpolate.interp1d(q2, i2, kind='linear', bounds_error=False, fill_value='extrapolate')
    newfe = interpolate.interp1d(q2, e2, kind='linear', bounds_error=False, fill_value='extrapolate')
    newi = newf(q3)
    newe = newfe(q3)

    q_final = numpy.concatenate((q_final, q3))
    i_final = numpy.concatenate((i_final, (newi + i3) / 2))
    e_final = numpy.concatenate((e_final, (newe + e3)))

    # 4 part
    q4, i4, e4 = clipQmin(q1, i1, e1, cut2)
    q_final = numpy.concatenate((q_final, q4))
    i_final = numpy.concatenate((i_final, i4))
    e_final = numpy.concatenate((e_final, e4))
    return q_final, i_final, e_final


def clipQmin(q, i, error, qmin):
    # clip q min
    if error is not None:
        error = numpy.repeat(error, q >= qmin)
    i = numpy.repeat(i, q >= qmin)
    q = numpy.repeat(q, q >= qmin)
    return q, i, error


def clipQmax(q, i, error, qmax):
    # clip q max
    if error is not None:
        error = numpy.repeat(error, q <= qmax)
    i = numpy.repeat(i, q <= qmax)
    q = numpy.repeat(q, q <= qmax)
    return q, i, error

# ----------------------------------------------------------------------------------
if __name__== '__main__':
    '''
    q=arange(0,10,0.1)
    i=q**2
    dt=dataset('test',q,i)
    x=dt._getxml()
    print ElementTree.tostring(x)
    print "save in file :", "test_xml_from_python.xml"
    dt._data2XMLfile("c:\\python26\\test_xml_from_python.xml")
    '''
    dt=getDataDictFromXMLFile("c:\\python26\\test1.xml")
