from __future__ import print_function

import Manyo
import os


"""
Base Functions for SAS
File.py
- Save1dDataToText
- Save1dDataToSrlz
- ManyoDeserializationEC
- ManyoDeserializationECM
- ManyoSerialization
- SaveQzToText
- SaveIqToText
- SaveRPMTToText
- SaveQSpaceToText
"""
#########################################
def Save1dDataToText(dat, dirpath = "/home/mlfdev/users", filename="SAS104000.dat", prec=6):
    """
    Save 1d data to text file(s).
    @param dat       (ElementContainer or list of I(q)) 1d data
    @param dirpath   (string)           directory where files are saved
    @param filename  (string)           filename
    @param prec      (int)              precision
    """
    if not os.path.exists(dirpath):
        raise UserWarning("%s dose not exist!!"%dirpath)
    if isinstance(dat, list):
        for ec in dat:
            xKey = ec.PutXKey()
            try:
                ec.SetKeys("Q", "Intensity", "Error")
            except:
                raise UserWarning("Data is not I(q).")
            bank_name = ec.PutHeader().PutString("BANKNAME")
            fn_split = os.path.splitext(filename)
            new_filename = "%s-%s%s"%(fn_split[0], bank_name, fn_split[1])
            print("Save %s to %s"%(bank_name, os.path.join(dirpath, new_filename)))
            ec.SaveTextFile(os.path.join(dirpath, new_filename), prec)
            ec.SetKeys(xKey, "Intensity", "Error")
    else:
        try:
            dat.SaveTextFile(os.path.join(dirpath, filename), prec)
        except:
            raise UserWarning("Unexpected data format")

#########################################
def Save1dDataToSrlz(dat, dirpath = "/home/mlfdev/users", filename="SAS104000.srlz"):
    """
    Save 1d data to text file(s).
    @param dat       (ElementContainer or list of I(q)) 1d data
    @param dirpath   (string)           directory where files are saved
    @param filename  (string)           filename
    """
    if not os.path.exists(dirpath):
        raise UserWarning("%s dose not exist!!"%dirpath)
    if isinstance(dat, list):
        for ec in dat:
            xKey = ec.PutXKey()
            try:
                ec.SetKeys("Q", "Intensity", "Error")
            except:
                raise UserWarning("Data is not I(q).")
            bank_name = ec.PutHeader().PutString("BANKNAME")
            fn_split = os.path.splitext(filename)
            new_filename = "%s-%s%s"%(fn_split[0], bank_name, fn_split[1])
            print("Save %s to %s"%(bank_name, os.path.join(dirpath, new_filename)))
            write = Manyo.WriteSerializationFileBinary(os.path.join(dirpath, new_filename))
            write.Save(ec)
            del write
            ec.SetKeys(xKey, "Intensity", "Error")
    else:
        try:
            write = Manyo.WriteSerializationFileBinary(os.path.join(dirpath, filename))
            write.Save(dat)
            del write
        except:
            raise UserWarning("Unexpected data format")
#########################################
def SaveQzToText(dat, dirpath = "/home/mlfdev/users", filename="SAS104000_Qz.dat"):
    """
    Save Qz to a text file.
    @param dat          Qz ElementContainerArray
    @param dirpath      directory where files are saved
    @param filename     filename
    """
    if not isinstance( dat, Manyo.ElementContainerArray ):
        raise UserWarning("dat is not ElementContainerArray")
    fn = open(os.path.join(dirpath, filename), "w")
    qxs = dat.PutHeaderPointer().PutDoubleVector("Qx")
    num_of_ec = dat.PutSize()
    ##[inamura 170118]-->
    fn.write("##Qx, Qy, Intensity, Error\n")
    qxsize = qxs.size()
    qy = dat(0).ReduceColumnList("Qy")
    qysize = len(qy)
    fn.write("## XRANGE=%f,%f,%f\n"%( (qxs[0]+qxs[1])/2.0, (qxs[qxsize-2]+qxs[qxsize-1])/2.0, qxs[1]-qxs[0] ) )
    fn.write("## YRANGE=%f,%f,%f\n"%( qy[0],qy[qysize-1],qy[1]-qy[0] ) )
    fn.write("## XLABEL=Qx [1/Ang.]\n")
    fn.write("## YLABEL=Qy [1/Ang.]\n")

    out_ymatrix=[]
    out_ematrix=[]
    for i in range(num_of_ec):
        ys = dat(i).PutYList()
        es = dat(i).PutEList()
        out_ymatrix.append( ys )
        out_ematrix.append( es )

    for i in range( num_of_ec ):
        for j in range( qysize ):
            fn.write( "%f, %f, %e, %e\n"%( (qxs[j]+qxs[j+1])/2.0, qy[i], out_ymatrix[j][i], out_ematrix[j][i] ) )

    fn.close()
    """
    for i in range(num_of_ec):
        xs = dat(i).ReduceColumnList("Qy")
        ys = dat(i).PutYList()
        es = dat(i).PutEList()
        for j in range(len(ys)):
            fn.write("%f, %f, %e, %e\n"%( (qxs[i]+qxs[i+1])*0.5, xs[j], ys[j], es[j]))
    """
    ##<--[inamura 170118]
#########################################
def SaveIqToText(dat, dirpath = "/home/mlfdev/users", filename="SAS104000_Iq.dat", prec=6):
    """
    Save I(q) to a text file.
    @param dat         I(q) data which is list of ElementContainers
    @param dirpath      directory where files are saved
    @param filename     filename
    @param prec         precision
    """
    xKey = dat[0].PutXKey()
    try:
        dat[0].SetKeys("Q", "Intensity", "Error")
    except:
        raise UserWarning("Data is not I(q).")
    if not os.path.exists(dirpath):
        raise UserWarning("%s dose not exist!!"%dirpath)
    dat[0].SaveTextFile(os.path.join(dirpath, filename), prec)
    dat[0].SetKeys(xKey, "Intensity", "Error")

#########################################
def ManyoDeserializationEC(dirpath = ".", filename=""):
    """
    Deserialize an EC file by Manyo-Lib.
    @param dirpath              directory where files are saved
    @param filename             filename
    @retval EC                 ElementContainer
    """
    filepath = os.path.join(dirpath, filename)
    if not os.path.exists(filepath):
        raise UserWarning("%s dose not exist!!"%filepath)
    ec = Manyo.ElementContainer()
    read = Manyo.ReadSerializationFileBinary(filepath)
    read.Load(ec)
    del read
    return ec

#########################################
def ManyoDeserializationECM(dirpath = ".", filename=""):
    """
    Deserialize an ECM file by Manyo-Lib.
    @param dirpath              directory where files are saved
    @param filename             filename
    @retval DAT                 ElementContainerMatrix
    """
    filepath = os.path.join(dirpath, filename)
    if not os.path.exists(filepath):
        raise UserWarning("%s dose not exist!!"%filepath)
    ecm = Manyo.ElementContainerMatrix()
    read = Manyo.ReadSerializationFileBinary(filepath)
    read.Load(ecm)
    del read
    return ecm

#########################################
def ManyoSerialization(dat, dirpath = "/home/mlfdev/users", filename="SAS104000.srlz"):
    """
    Serialize data by Manyo-Lib.
    @param dat                  data container
    @param dirpath              directory where files are saved
    @param filepath             filename
    """
    if not os.path.exists(dirpath):
        raise UserWarning("%s dose not exist!!"%dirpath)
    filepath = os.path.join(dirpath, filename)
    write = Manyo.WriteSerializationFileBinary(filepath)
    write.Save(dat)
    del write

#########################################
def SaveRPMTToText(dat, dirpath="/home/mlfdev/users", filename="SAS000001_RPMT.txt", lambdaRange="0.7 7.5"):
    """
    Save RPMT data to text
    @param dat             ElementContainerMatrix
    @param dirpath         directory where files are saved
    @param lambdaRange
    @param filepath        file name
    """
    import math
    if not os.path.exists(dirpath):
        raise UserWarning("%s dose not exist!!"%dirpath)
    filepath = os.path.join(dirpath, filename)

    lambdaRange_list=lambdaRange.strip().split(" ")
    if len(lambdaRange_list)!=2:
        raise UserWarning("Lambda Range is invalid")
    lamb_min = float(lambdaRange_list[0])
    lamb_max = float(lambdaRange_list[1])

    existed_xpos_list = []
    existed_ypos_list = []
    for i in range(dat.PutSize()):
        if dat(i).PutHeaderPointer().PutInt4("MASKED")==0:
            xpos = dat(i).PutHeaderPointer().PutInt4("XPOS")
            existed_xpos_list.append( xpos )
            if len(existed_ypos_list)==0:
                for j in range(dat(i).PutSize()):
                    ypos = dat(i,j).PutHeaderPointer().PutInt4("YPOS")
                    existed_ypos_list.append( ypos )

    existed_xpos_list.sort()
    existed_ypos_list.sort()

    fo = open( filepath, "w" )
    if fo==None:
        raise UserWarning("Failed to open file = "+filepath)

    SIH = Manyo.SearchInHeader()
    SIH.SetTarget(dat)

    for xpos in existed_xpos_list:
        SIH.SearchArray( "XPOS", xpos )
        iv = SIH.PutResultIndex()
        if iv.size()!=0:
            eca=dat(iv[0])
            int_list = []
            err_list = []
            for i in range( len(existed_ypos_list) ):
                int_list.append(0.0)
                err_list.append(0.0)

            for i in range(eca.PutSize()):
                ypos = eca(i).PutHeaderPointer().PutInt4("YPOS")
                ii = eca(i).PutYList()
                ee = eca(i).PutEList()
                LL = eca(i).Put('Lamb')
                for j in range( LL.size() ):
                    if LL[j]>=lamb_min and LL[j]<lamb_max:
                        int_list[ypos] = int_list[ypos]+ii[j]
                        err_list[ypos] = err_list[ypos]+( ee[j]*ee[j] )

            for ypos in range( len(existed_ypos_list) ):
                conts = "%d, %d, %f, %f\n"%(xpos,ypos,int_list[ypos],math.sqrt( err_list[ypos] ))
                fo.write( conts )

    fo.close()

#########################################
def SaveQSpaceToText(dat, dirpath="/home/mlfdev/users", filename="SAS000001_QSpace.txt", QxRange="1.0 1.2", QyRange="-0.1 0.1", QzRange="-0.4 -0.2", cutoff=0.01):
    """
    Save Q space data to text
    @param dat      (ElementContainerMatrix) Q space ECM
    @param dirpath  (string) Directory where files are saved
    @param filename (string) File name
    @param QxRange  (string) Minimum and maximum qx separated with space
    @param QyRange  (string) Minimum and maximum qy separated with space
    @param QzRange  (string) Minimum and maximum qz separated with space
    @param cutoff   (float) Cutoff value for intensity. If it is negative, it is ignored.
    """
    if not os.path.exists(dirpath):
        raise UserWarning("{} dose not exist!!".format(dirpath))
    filepath = os.path.join(dirpath, filename)

    qx_min, qx_max = list(map(float, QxRange.strip().split(" ")))
    qy_min, qy_max = list(map(float, QyRange.strip().split(" ")))
    qz_min, qz_max = list(map(float, QzRange.strip().split(" ")))

    ecm_size = dat.PutSize()
    eca_size = dat(0).PutSize()
    with open(filepath, "wb") as fp:
        for i in range(ecm_size):
            if (dat(i).PutHeaderPointer().PutInt4("MASKED"))==1:
                continue
            for j in range(dat(i).PutSize()):
                if (dat(i, j).PutHeaderPointer().PutInt4("MASKED"))==1:
                    continue
                qx_l = dat(i, j).ReduceColumnList("Qx")
                qy_l = dat(i, j).ReduceColumnList("Qy")
                qz_l = dat(i, j).ReduceColumnList("Qz")
                ys = dat(i, j).PutYList()
                es = dat(i, j).PutEList()
                for k in range(len(qx_l)):
                    if cutoff < 0.0:
                        if ( qx_min <= qx_l[k] <= qx_max ) and ( qy_min <= qy_l[k] <= qy_max ) and ( qz_min <= qz_l[k] <= qz_max ):
                            fp.write("{0:g}, {1:g}, {2:g}, {3:g}, {4:g}\n".format(qx_l[k], qy_l[k], qz_l[k], ys[k], es[k]))
                    else:
                        if (ys[k] >= cutoff) and ( qx_min <= qx_l[k] <= qx_max ) and ( qy_min <= qy_l[k] <= qy_max ) and ( qz_min <= qz_l[k] <= qz_max ):
                            fp.write("{0:g}, {1:g}, {2:g}, {3:g}, {4:g}\n".format(qx_l[k], qy_l[k], qz_l[k], ys[k], es[k]))
                del qx_l, qy_l, qz_l, ys
    del dat
