import Manyo
import Manyo.Utsusemi as mu
import Manyo.SAS as ms

import utsusemi.ana.Reduction.BaseCommands as BC
import sys
import os

DAT = None

def UBSliceBase(DAT, outfile, UBmat, vaxes, vtitles, vunits, sliceInfo):
    """UBSlice command line version
    Defines:
    UBmat = [ 0.0, 0.0, 0.0, \
              0.0, 0.0, 0.0, \
              0.0, 0.0, 0.0  ]
    vaxes = [[1.0, 0.0, 0.0], \
             [0.0, 1.0, 0.0], \
             [0.0, 0.0, 1.0] ]
    vtitles = ["H", "K", "L"]
    vunits = ["rlu", "rlu", "rlu"]
    sliceInfo = [[0.0, 1.0, "X", 0.1], \
                 [0.0, 1.0, "Y", 0.1], \
                 [0.0, 1.0, "T", 0.1]]

    @param DAT (ElementContainerMatrix) ElementContainerMatrix of data
    @param outifle (str) path to text file for output
    @param UBmat
    @param vaxes
    @param vtitles
    @param vunits
    @param sliceInfo
    @retval None
    """
    # Preparing data process
    TQ = ms.TreatQhklWithUB()
    TQ.SetTarget(DAT)

    # Projection
    if TQ.SetUBmatrix(UBmat):
        TQ.SetViewAxes( vaxes[0], vaxes[1], vaxes[2] )
        TQ.SetViewAxesLabels(vtitles[0], vtitles[1], vtitles[2], vunits[0], vunits[1], vunits[2])
        if TQ.Projection():
            pass
        else:
            msg = "ERROR: Faled projection."
            print(msg)
            raise UserWarning(msg)
    else:
        msg = "ERROR: given UBmatrix is invalid."
        print(msg)
        raise UserWarning(msg)

    # Slicing
    indSliceX = -1
    indSliceY = -1
    indSliceT = -1
    for i in range(3):
        ax_type = sliceInfo[i][2]
        if ax_type.upper() == "X":
            indSliceX = i
        elif ax_type.upper() == "Y":
            indSliceY = i
        elif ax_type.upper() == "T":
            indSliceT = i
        else:
            msg = "ERROR: given axis type in slice info is invalid ({})".format(ax_type)
            print(msg)
            raise UserWarning(msg)
    if indSliceX == -1 or indSliceY == -1 or indSliceT == -1:
        msg = "ERROR: given axis type in slice info is invalid."
        print(msg)
        raise UserWarning(msg)

    TQ.SetSliceAxes(sliceInfo[0][2], sliceInfo[1][2], sliceInfo[2][2])

    x_range = [float(sliceInfo[indSliceX][0]), float(sliceInfo[indSliceX][1]), float(sliceInfo[indSliceX][3])]
    y_range = [float(sliceInfo[indSliceY][0]), float(sliceInfo[indSliceY][1]), float(sliceInfo[indSliceY][3])]
    t_range = [float(sliceInfo[indSliceT][0]), float(sliceInfo[indSliceT][1]), float(sliceInfo[indSliceT][3])]

    sliced_dat = Manyo.ElementContainerArray()
    TQ.Slice(sliced_dat, x_range, y_range, t_range)

    # Save sliced data as text
    keys = sliced_dat.PutHeader().PutStringVector(TQ.SAS_TREATQHK_KEYS)

    xKey = keys[0]
    yKey = keys[1]

    # output the sliced data as text
    Yvect = sliced_dat(0).PutX()
    Ylist = []
    for i in range(Yvect.size() - 1):
        Ylist.append((Yvect[i] + Yvect[i + 1]) / 2.0)

    with open(outfile, "w") as fo:
        fo.write("# {}, {}, Intensity, Error \n".format(xKey, yKey))
        for i in range(sliced_dat.PutSize()):
            Xval = sliced_dat(i).PutHeader().PutDouble(xKey)
            Ilist = sliced_dat(i).PutYList()
            Elist = sliced_dat(i).PutEList()
            for j in range(len(Ilist)):
                a_line = "{:16.8f}, {:16.8f}, {:16.8f}, {:16.8f}\n".format(Xval, Ylist[j], Ilist[j], Elist[j])
                fo.write(a_line)


def UBSlice(dat, xtalparam, outfile=""):
    """UBSlice command line version

    @param dat (str or ElementContainerMatrix) "path/to/dat.mdb" or ElementContainerMatrix of data
    @param xtalparam (str) path to XtalParam
    @param outifle (str) path to text file for output
    @retval None
    """
    if isinstance(dat, str):
        if os.path.exists(dat):
            dir_name = os.path.dirname(dat)
            filename = os.path.basename(dat)
            dat = BC.LoadDataFromManyoBinary(dir_name, filename)
            if outfile == "":
                outfile = os.path.join(dir_name, filename + ".txt")
            elif os.path.dirname(outfile) == "":
                outfile = os.path.join(dir_name, outfile)
        else:
            msg = "ERROR: file not found {}".format(dat)
            print(msg)
            raise UserWarning(msg)

    elif isinstance(dat, Manyo.ElementContainerMatrix):
        pass
    else:
        msg = "ERROR: argument dat must be path/to/data.mdb or ElementContainerMatrix"
        print(msg)
        raise UserWarning(msg)

    if not os.path.exists(xtalparam):
        msg = "ERROR: argument xtalparam not found."
        print(msg)
        raise UserWarning(msg)

    # Define UB matrix
    UBmat = [ 0.0, 0.0, 0.0, \
              0.0, 0.0, 0.0, \
              0.0, 0.0, 0.0  ]
    vaxes = [[1.0, 0.0, 0.0], \
             [0.0, 1.0, 0.0], \
             [0.0, 0.0, 1.0] ]
    vtitles = ["H", "K", "L"]
    vunits = ["rlu", "rlu", "rlu"]

    sliceInfo = [[0.0, 1.0, "X", 0.1], \
                 [0.0, 1.0, "Y", 0.1], \
                 [0.0, 1.0, "T", 0.1]]

    vcp=mu.UtsusemiSqeCalcXtalParams()
    flags = vcp.TOXML_UBMATRIX + vcp.TOXML_SLICINGINFO + vcp.TOXML_PROJECTION
    if vcp.LoadFile(xtalparam, flags):
        ubmat_v = vcp.PutUBMatrix()
        qaxes_v = vcp.PutViewAxes()
        qaxes_labels_v=vcp.PutViewAxesLabels()
        qaxes_units_v=vcp.PutViewAxesUnits()

        for i in range(9):
            UBmat[i] = ubmat_v[i]

        for i in range(3):
            ax = []
            for j in range(3):
                vaxes[i][j] = float(qaxes_v[(i * 4) + j])
            vtitles[i] = qaxes_labels_v[i]
            vunits[i] = qaxes_units_v[i]

        slice_type_v = vcp.PutAxType()

        for i in range(3):
            slice_v = vcp.PutAxRange(i)
            sliceInfo[i][0] = float(slice_v[0])
            sliceInfo[i][1] = float(slice_v[1])
            sliceInfo[i][2] = str(slice_type_v[i]).upper()
            sliceInfo[i][3] = float(slice_v[2])
    else:
        msg = "ERROR: "
        print(msg)
        raise UserWarning(msg)

    UBSliceBase(dat, outfile, UBmat, vaxes, vtitles, vunits, sliceInfo)



if __name__ == '__main__':
    # Load data file
    dir_path = "nakajima/250425UBslicerCUI/"
    filename = "213620_steady_HL.mdb"

    # Load XtalParam file
    xtalparam_file = None
    #xtalparam_file = "nakajima/250425UBslicerCUI/XtalParam_CFGO_omg_29.xml"
    xtalparam_file = "test.xml"

    dat_file = os.path.join(dir_path, filename)
    UBSlice(dat_file, xtalparam_file, "./output_UBSlice.txt")

    DAT = BC.LoadDataFromManyoBinary(dir_path, filename)
    UBSlice(DAT, xtalparam_file, "output_UBSlice2.txt")
