#!/usr/bin/python
# -*- coding: utf-8 -*-
"""
U2IFJupyter.py
uGao u2dplot I/F on Jupyter and on the dark mode
"""
from __future__ import print_function
import os
import Manyo
import Manyo.MLF as mm
import numpy as np
import datetime

from uGao.u2dplotBase import U2Params, D2Chart, DataConverter, format_int_range
from uGao.uGaoColorMap import COLOR_MAPS
from uGao.uGaoUtil import IFEvtProp, MASKVALUE

import matplotlib
ISJUPYTER = False
if matplotlib.get_backend() == "inline":
    # print("U2IFJupyter : backend = inline")
    ISJUPYTER = True
    import matplotlib.pyplot as plt
else:
    # print("U2IFJupyter : backend is not inline")
    from matplotlib.figure import Figure
    from matplotlib.backends.backend_agg import FigureCanvasAgg as FigCanvas

#########################################
# U2IFJupyter
#########################################
class U2IFJupyter(object):
    """
    U2IFJupyter class

    This enables to plot on Jupyter and to save plot file using u2dplot of uGao

    import uGao.U2IFJupyter as mp
    p = mp.U2IFJupyter(sizeX=10.0, sizeY=5.0)
    p.SetKeys("index", "index", "Intensity")
    p.SetData(dat)
    p.SetRange(False, 0.0, 40.0)
    p.SetTitles("Test in Dark", "")
    p.SetLabels("QA", "hw")
    p.SaveFile("test_InDark.png")
    # On Jupyter
    p.Show()
    """

    #########################################
    def __init__(self, data=None, orderNo=1, sizeX=7.0, sizeY=7.0, dpi=100):
        """
        Constructor
        @param data(list of numpy arrays or ElementContainerMatrix or -Array)
        @param orderNo(int) Plot number
        @param sizeX (float) size of X axis direction
        @param sizeY (float) size of Y axis direction
        @param dpi (float) dot per inch
        @retval None
        """
        if ISJUPYTER:
            # print("U2IFJupyter : ISJUPYTER is True")
            self.plt = plt
            self.fig = plt.figure(figsize=(sizeX, sizeY))
            self.canvas = self.fig.canvas
        else:
            # print("U2IFJupyter : ISJUPYTER is False")
            self.fig = Figure((sizeX, sizeY), dpi)
            self.canvas = FigCanvas(self.fig)
        self.order = orderNo
        self.params = U2Params()
        self.sliceMode = False
        self.ifi = IFEvtProp(orderNo)
        self.ifi.AddProperty('params', lambda: self.params)
        self.ifi.AddProperty('slice_mode', lambda: self.sliceMode)

        self.d2 = None

        self.ndata = None
        self.data = data
        self.keys = (None, None, None)
        self.intRange = None
        self.converter = DataConverter()

    ###################################################
    def Show(self):
        """
        """
        if ISJUPYTER:
            if self.ndata is None:
                if not self.SetData():
                    msg = "U2IFJupyter::SaveFile >> data is invalid."
                    raise UserWarning(msg)

            if self.d2 is not None:
                del self.d2

            self.d2 = D2Chart(self, self.canvas, self.ndata, True)

            self.plt.show()
        else:
            pass

    #########################################
    def SetKeys(self, xkey, ykey, intkey=None, rkey=None):
        """キー文字列設定処理

        Parameters:
            xkey (str or None): X Key
            ykey (str or None): Y Key
            intkey (str or None): Int Key
            rkey (str or None): Range Key
        """
        if xkey is not None and not isinstance(xkey, str):
            raise PlotException('Common', 'C025', ('xkey',))
        if ykey is not None and not isinstance(ykey, str):
            raise PlotException('Common', 'C025', ('ykey',))
        if intkey is not None and not isinstance(intkey, str):
            raise PlotException('Common', 'C025', ('intkey',))
        if rkey is not None and not isinstance(rkey, str):
            raise PlotException('Common', 'C025', ('rkey',))

        self.keys = (xkey, ykey, intkey, rkey)

    #########################################
    def SetIntRange(self, range0, range1):
        """強度積算範囲設定処理

        self.intRangeに積算範囲を設定する。
        引数のどちらかがNoneの時は全範囲になるようにself.intRangeをNoneにする。

        Parameters:
            range0 (float or None): 積算範囲始値
            range1 (float or None): 積算範囲終値
        """
        self.intRange = format_int_range(range0, range1)

    #########################################
    def SetData(self, data=None):
        """
        Set Data
        @param data(EleentContainerArray, -Matrix or list of arrays)
        @retval None
        """
        if data is None:
            if self.data is None:
                return False
            data = self.data
        else:
            if self.data is not None:
                del self.data
            self.data = data

        if isinstance(data, Manyo.ElementContainerMatrix) or isinstance(data, Manyo.ElementContainerArray):
            self.ndata = self.converter.ConvertData(data, self.keys, self.intRange)
        elif data is not None:
            self.ndata = data
        else:
            msg = "U2IFJupyter::SetData >> ChangeData argument is invalid."
            raise UserWarning(msg)
        return True

    #########################################
    def ReadData(self, filename):
        """
        """
        print("ReadData not work yet")

    #########################################
    def SetTitles(self, title, comment=None):
        """
        Set title strings
        @param title(string)
        @retval None
        """
        self.params.title = str(title)
        if comment is not None:
            self.params.comment = comment

    #########################################
    def SetTransPosition(self, trans):
        """
        軸の転置指定
        @param  trans (bool) True means axes transpose
        @retval 無し
        """
        # 引数の型チェック
        if isinstance(trans, bool):
            if self.params.transPosition != trans:
                self.params.transPosition = trans
        else:
            msg = "U2IFJupyter::SetTransPosition >> argument is invalid."
            raise UserWarning(msg)

    #########################################
    def SetXScale(self, autoFlag, r_min=None, r_max=None):
        """
        Set Xaxis range
        @param autoFlag(bool) True means auto scale mode
        @param r_min(float) minimum value of the range (enabled when autoFlag is False) / None (autoscale)
        @param r_max(float) maximum value of the range (enabled when autoFlag is False) / None (autoscale)
        @retval None
        """
        self.params.xscale.autoRange = autoFlag
        if not autoFlag:
            if r_min is not None:
                r_min = float(r_min)
            if r_max is not None:
                r_max = float(r_max)
            self.params.xscale.range = (r_min, r_max)

    #########################################
    def SetYScale(self, autoFlag, r_min=None, r_max=None):
        """
        Set Yaxis range
        @param autoFlag(bool) True means auto scale mode
        @param r_min(float) minimum value of the range (enabled when autoFlag is False) / None (autoscale)
        @param r_max(float) maximum value of the range (enabled when autoFlag is False) / None (autoscale)
        @retval None
        """
        self.params.yscale.autoRange = autoFlag
        if not autoFlag:
            if r_min is not None:
                r_min = float(r_min)
            if r_max is not None:
                r_max = float(r_max)
            self.params.yscale.range = (r_min, r_max)

    #########################################
    def SetUnits(self, xunit=None, yunit=None):
        """
        軸の単位(unit)を指定
        @param  xunit (string or None) X軸の単位
        @param  yunit (string oe None) Y軸の単位
        @retval 無し
        """
        if xunit is not None:
            self.params.xscale.unit = xunit
        if yunit is not None:
            self.params.yscale.unit = yunit

    #########################################
    def SetLabels(self, xlabel, ylabel):
        """
        軸のラベルを指定
        @param  xlabel (string or None) X軸のラベル
        @param  ylabel (string or None) Y軸のラベル
        @retval 無し
        """
        if xlabel is not None:
            self.params.xscale.label = xlabel
        if ylabel is not None:
            self.params.yscale.label = ylabel

    #########################################
    def SetHGrid(self, grid, color=None, interval=None):
        """
        水平グリッドの指定
        @param  grid (bool) True:On / False:Off
        @param  color (int) 0:White, 1:Yellow 2:Red, 3:Black, 4:Blue, 5:Green, 6:Cyan, 7:Purple
        @param  interval (float) 間隔
        @retval 無し
        """
        # 引数の型チェック
        if isinstance(grid, bool):
            self.params.hgrid = grid
            # カラーが指定されていれば
            if color is not None:
                self.params.gridColorH = color
            # 間隔が指定されていれば
            if interval is not None:
                self.params.hgridInterval = interval

    #########################################
    def SetVGrid(self, grid, color=None, interval=None):
        """
        垂直グリッドの指定
        @param  grid (bool) True:On / False:Off
        @param  color (int) 0:White, 1:Yellow 2:Red, 3:Black, 4:Blue, 5:Green, 6:Cyan, 7:Purple
        @param  interval (float) 間隔
        @retval 無し
        """
        # 引数の型チェック
        if isinstance(grid, bool):
            self.params.vgrid = grid
            # カラーが指定されていれば
            if color is not None:
                self.params.gridColorV = color
            # 間隔が指定されていれば
            if interval is not None:
                self.params.vgridInterval = interval

    #########################################
    def SetColorMap(self, cmap):
        """
        Set color map
        @param cmap(int or string) Color map No. or name ['Default', 'autumn', 'winter', 'gray', 'haxby', 'terrain', 'turbo']
        @retval None
        """
        if isinstance(cmap, int):
            if cmap >= 0 and cmap < len(COLOR_MAPS):
                self.params.colorMap = cmap
                return
        elif isinstance(cmap, str):
            for i, a_map in enumerate(COLOR_MAPS):
                if cmap == a_map:
                    self.params.colorMap = i
                    return

        msg = "U2IFJupyter::SetColorMap >> invalid argument. Acceptable : "
        for a_name in COLOR_MAPS:
            msg += (a_name + ", ")
        msg = msg[:-2]
        raise UserWarning(msg)

    #########################################
    def SetSmooth(self, flag, times=1):
        """
        Set smoothing parameters
        @param flag(int) if flag<=0, Smoothing is off
        @param times(int) the number of times
        @retval None
        """
        if float(flag) <= 0:
            self.params.smoothing = [False, 1, 1]
        else:
            self.params.smoothing = [True, int(flag), int(times)]

    #########################################
    def SetSmoothing(self, isSmooth=False, windowVal=1, timesVal=1):
        """
        Set Smoothiing
        @param  isSmooth (bool) Smooth on or not
        @param windowVal (int) Window Size for smoothing
        @param timesVal (int) The number of times
        @retval None
        """
        if isinstance(isSmooth, bool):
            if isSmooth:
                self.params.smoothing = [True, int(windowVal), int(timesVal)]
            else:
                self.params.smoothing = [False, 1, 1]

    #########################################
    def SetBankSeparator(self, flag=False):
        """
        Bank Separatorの指定
        @param flag   Showing bank separators or not
        @retval None
        """
        # 引数の型チェック
        if isinstance(flag, bool):
            self.params.bankseparator = flag
        else:
            msg = "U2IFJupyter::SetBankSeparator >> invalid argument."
            raise UserWarning(msg)

    #########################################
    def SetIntRange(self, autoFlag, r_min=None, r_max=None):
        """
        Set Zaxis(Intensity) range
        @param autoFlag(bool) True means auto scale mode
        @param r_min(float) minimum value of the range (enabled when autoFlag is False) / None (autoscale)
        @param r_max(float) maximum value of the range (enabled when autoFlag is False) / None (autoscale)
        @retval None
        """
        self.SetRange(autoFlag, r_min, r_max)

    #########################################
    def SetRange(self, autoFlag, r_min=None, r_max=None):
        """
        Set Zaxis(Intensity) range
        @param autoFlag(bool) True means auto scale mode
        @param r_min(float) minimum value of the range (enabled when autoFlag is False) / None (autoscale)
        @param r_max(float) maximum value of the range (enabled when autoFlag is False) / None (autoscale)
        @retval None
        """
        self.params.autoRange = autoFlag
        if autoFlag:
            self.params.range = None
        else:
            if r_min is not None:
                r_min = float(r_min)
            if r_max is not None:
                r_max = float(r_max)
            self.params.range = (r_min, r_max)

    #########################################
    def SetLog(self, zlog):
        """
        Set log scale on Z axis
        @param zlog(bool) True means plot with log scale
        @retval None
        """
        if isinstance(zlog, bool):
            self.params.logScale = zlog
        else:
            msg = "U2IFJupyter::SetLog >> invalid argument."
            raise UserWarning(msg)

    #########################################
    def SetFont(self, key, fontSize=10, fontName=None):
        """
        Set font size and font
        @param key(str) target
        @param fontSize(int) font size
        @param fontName(str) font name
        @retval None
        """
        if key in self.params.fontDict.keys():
            if isinstance(fontSize, int):
                self.params.fontDict[key][0] = int(fontSize)
            else:
                msg = "U2IFJupyter::SetFont >> invalid argument."
                raise UserWarning(msg)
            if fontName is not None:
                if isinstance(fontName, str) and fontName in ["Serif", "Sans"]:
                    self.params.fontDict[key][1] = fontName
                else:
                    msg = "U2IFJupyter::SetFont >> invalid argument."
                    raise UserWarning(msg)
            else:
                pass
        else:
            msg = "U2IFJupyter::SetFont >> invalid key argument. key must be {}".format(self.params.fontDict.keys())
            raise UserWarning(msg)

    #########################################
    def SetFigureSize(self, px, py):
        """
        Set Size of Figure
        @param px(float) size of X direction
        @param py(float) size of Y direction
        @retval None
        """
        if ISJUPYTER:
            return
        self.fig = Figure((px, py), 100)
        self.canvas = FigCanvas(self.fig)

    #########################################
    def SaveFile(self, filename):
        """
        Save figure file
        @param filename(string) The path of a plot figure file.
        @retval None
        """
        if self.ndata is None:
            if not self.SetData():
                msg = "U2IFJupyter::SaveFile >> data is invalid."
                raise UserWarning(msg)

        if self.d2 is not None:
            del self.d2

        self.d2 = D2Chart(self, self.canvas, self.ndata, True)
        if os.path.isdir(os.path.dirname(filename)):
            pass
        else:
            filename = os.path.join(os.getcwd(), filename)

        self.fig.savefig(filename)
    #########################################

    def SaveAsTextFromPlot(self, filename):
        """
        データファイルをテキスト形式で保存処理
        This function is copied from u2dplot::D2Frame::SaveAsTxt.
        @param  filename (string) ファイルのパス
        @retval 無し
        """
        import codecs
        root, ext = os.path.splitext(filename)
        if ext == "":
            filename = root + ".txt"
        elif ext != ".txt":
            raise PlotException('Common', 'C043', ("txt",))
        try:
            # テキストファイルを開く
            fw = codecs.open(filename, 'w', 'utf-8')
        except:
            raise PlotException('Common', 'C018', (filename, ))
        # グラフのタイトルをコメントとして書き込む
        fw.write("# Title: " + self.params.title + os.linesep)
        # グラフのコメントを書き込む
        lines = self.params.comment.split('\n')
        if len(lines) > 0:
            fw.write("# Comment: " + lines[0] + os.linesep)
        if len(lines) > 1:
            fw.write("#          " + lines[1] + os.linesep)
        if len(lines) > 2:
            fw.write("#          " + lines[2] + os.linesep)

        # 行タイトル
        fw.write("#       X           Y         Int" + os.linesep)
        # データを展開
        if self.ndata is None:
            if not self.SetData():
                msg = "U2IFJupyter::SaveAsText >> data is invalid."
                raise UserWarning(msg)
        Z, X, Y = self.ndata
        # テキストファイルに書き込み
        for i in range(len(Z[0])):
            for j in range(len(Z)):
                xx = (X[j][i] + X[j][i + 1]) / 2.0
                yy = (Y[j][i] + Y[j + 1][i]) / 2.0
                value = Z[j][i]
                if value >= MASKVALUE:
                    value *= 1.1
                line = "%.5e %.5e %.5e" % (xx, yy, value)
                fw.write(line + os.linesep)
            fw.write(os.linesep)
        fw.close()
    #########################################

    def SaveAsText(self, filename, dQ=0.0, isIgnoreMask=False, maskinfo=[False, 0.0]):
        """
        Save data as a text file (format and arguments are completely same as M2PlotPlusCui::SaveAsText)
        @param filename (str) path/to/file
        @param dQ (float) binning of Q at elastic position which is used for inelastic powder data
        @param isIgnoreMask (bool) Whether points witch intensity is mask value are ignore ( not saved in file ) or not
        @param maskinfo (list of bool and float) [True, <val>] is to replace mask value to <val>
        @retval None
        """
        if len(maskinfo) != 2:
            raise UserWarning(
                "M2PlotPlusCui::SaveAsText >> maskinfo is invalid")
        filename = str(filename)
        saveTxt = U2IFJupyterTextFileIO(self)
        mask_val = 0.0
        try:
            mask_val = float(maskinfo[1])
        except:
            UserWarning("Invalid argument maskinfo [<bool>, <float>]")
        saveTxt.Save(filename, dQ, isIgnoreMask, [maskinfo[0], mask_val])

    #########################################

    def OnXYlimChanged(self, axes):
        """XY軸変更コールバック関数

        Parameters:
            axes (matplotlib.axes.Axes)
        """
        params = self.ifi.GetProperty('params')
        if params.transPosition:
            params.xscale.range = list(axes.get_ylim())
            params.yscale.range = list(axes.get_xlim())
        else:
            params.xscale.range = list(axes.get_xlim())
            params.yscale.range = list(axes.get_ylim())
        self.ifi.NotifyEvent(self, 'change_displayed_params')


class U2IFJupyterTextFileIO(object):
    """
    M2PlotやM2PlotPlus用のデータ保存用テキストファイルのIOクラス
    """
    ##########################################################

    def __init__(self, parent):
        # 各種初期化
        self.parent = parent
        self.XRANGE = "XRANGE"
        self.YRANGE = "YRANGE"
        self.XKEY = "XKEY"
        self.YKEY = "YKEY"
        self.XUNIT = "XUNIT"
        self.YUNIT = "YUNIT"
        self.MASKVALUE = "MASKVALUE"
        self.XLABEL = "XLABEL"
        self.YLABEL = "YLABEL"
        self.TITLE = "Title"
        self.COMMENT = "Comment"
        self.HEADER = "## M2Plot+ Output (%s)\n"
        self.DATATITLE = "## X,Y,Intensity,Error"

    ##########################################################
    def Save(self, filename, dQ=0.0, isIgnoreMask=False, maskinfo=[False, 0.0]):
        """
        テキスト保存用メソッド
        @param filename (str) 保存先のフルパス
        @param dQ (float) 非弾性散乱粉末データ用の運動量軸のbin幅
        @param isIgnoreMask (bool) 強度がMASKVALUEだった時に、その点を書き出すかどうか
        @param maskinfo (list) MASKVALUEの値をどうするかを決める [<bool>,<float>] bool : Trueのとき float の値を使用、bool : Falseのとき本来のMASKVALUEを使用
        @retval None
        """
        self.params = self.parent.ifi.GetProperty('params')
        # fd = open( filename,'w' )
        try:
            fd = open(filename, 'w')
        except:
            msg = "SaveAsText : failed to open file = %s" % (filename)
            raise UserWarning(msg)

        # プロットされたbinデータを取り出し
        xbin = self.parent.ndata[1][0]
        ybin = self.parent.ndata[2].T[0]

        # キーの取り出し : self.keys = (xkey, ykey, intkey, rkey)
        _ykey = self.parent.keys[1]
        _xkey = self.parent.keys[0]

        # クリップ範囲の初期化
        clip_region = [0.0, 0.0, 0.0, 0.0]

        # データは非弾性散乱粉末データか？
        if self.parent.converter.isInelasticPowder:
            # キーはデフォルト
            _xkey = "Q"
            _ykey = "hw"
            # スライス関数の取り出しとスライス領域のクリア
            clip = self.parent.converter.arraySlicer
            clip.ClearClipRegion()

            # dQの設定（Q範囲はデータ内最大範囲）
            clip.SetQbin(dQ, -1.0, -1.0)

            # ybin （Energy軸）で区切ってスライス設定
            for i in range(len(ybin) - 1):
                clip_region[0] = -1.0
                clip_region[1] = ybin[i]
                clip_region[2] = -1.0
                clip_region[3] = ybin[i + 1]
                clip.SetClipRegion(*clip_region)

            # スライスの実行＆結果を取り出す
            ret_eca = clip.CutAlongX(True)

        else:
            # DetectMapモードなら
            # if self.parent.isDetMapMode:
            # データがElementContainerMatrixなら
            if isinstance(self.parent.data, Manyo.ElementContainerMatrix):
                data = self.parent.converter._eca
            else:
                data = self.parent.data

            # 非弾性散乱粉末データのスライスデータの構造に合わせてArrayの構造を変換する
            ret_eca = Manyo.ElementContainerArray()
            num_of_y = len(data(0).PutXList()) - 1

            II = []
            EE = []
            for i in range(data.PutSize()):
                II.append(data(i).PutYList())
                EE.append(data(i).PutEList())

            for i in range(num_of_y):
                ii = []
                ee = []
                for j in range(len(II)):
                    ii.append(II[j][i])
                    ee.append(EE[j][i])
                ec = Manyo.ElementContainer()
                ec.Add("X", xbin.tolist())
                ec.Add("I", ii)
                ec.Add("E", ee)
                ec.SetKeys("X", "I", "E")
                ret_eca.Add(ec)

        XX = ret_eca(0).PutXList()
        YY = np.array(ybin)

        # テキストファイル書き出し開始
        t_now = datetime.datetime.now()
        fd.write(self.HEADER % (t_now.strftime('%Y-%m-%d %H:%M:%S')))
        # 領域情報
        # 非弾性粉末か？
        if self.parent.converter.isInelasticPowder:
            xrange_text = "## %s=%g,%g,%g \n" % (self.XRANGE, (XX[0] + XX[1]) / 2.0, (XX[-1] + XX[-2]) / 2.0, dQ)
        else:
            xrange_text = "## %s=%g,%g,%g \n" % (self.XRANGE, (XX[0] + XX[1]) / 2.0, (XX[-1] + XX[-2]) / 2.0, XX[1] - XX[0])

        yrange_text = "## %s=%g,%g,%g \n" % (self.YRANGE, (YY[0] + YY[1]) / 2.0, (YY[-1] + YY[-2]) / 2.0, (ybin[1] - ybin[0]))
        fd.write(xrange_text)
        fd.write(yrange_text)

        # ラベル情報
        fd.write("## %s=%s\n" % (self.XLABEL, self.params.xscale.label))
        fd.write("## %s=%s\n" % (self.YLABEL, self.params.yscale.label))

        # UNIT情報
        fd.write("## %s=%s\n" % (self.XUNIT, self.params.xscale.unit))
        fd.write("## %s=%s\n" % (self.YUNIT, self.params.yscale.unit))

        # マスク情報
        if not isIgnoreMask:
            if maskinfo[0]:
                fd.write("## %s=%g\n" % (self.MASKVALUE, maskinfo[1]))

        # 軸のキー名
        fd.write("## %s=%s\n" % (self.XKEY, _xkey))
        fd.write("## %s=%s\n" % (self.YKEY, _ykey))

        # タイトル・コメント
        p_title = self.params.title
        p_comm = self.params.comment
        fd.write("# %s=%s\n" % (self.TITLE, p_title))
        fd.write("# %s=%s\n" % (self.COMMENT, p_comm))

        # データ
        fd.write(self.DATATITLE + " \n")
        for i in range(ret_eca.PutSize()):
            yy = (ybin[i] + ybin[i + 1]) / 2.0
            if yy < 1e-13 and yy > -1e-13:
                yy = 0.0
            ec = ret_eca(i)
            XX = ec.PutXList()
            YY = ec.PutYList()
            EE = ec.PutEList()
            for j in range(len(YY)):
                xx = (XX[j] + XX[j + 1]) / 2.0
                ii = YY[j]
                ee = EE[j]
                if ii < mm.MLF_MASKVALUE:
                    fd.write("%g, %g, %g, %g\n" % (xx, yy, ii, ee))
                else:
                    if isIgnoreMask:
                        continue
                    if maskinfo[0]:
                        fd.write("%g, %g, %g, %g\n" %
                                 (xx, yy, maskinfo[1], 0.0))
                    else:
                        fd.write("%g, %g, %g, %g\n" % (xx, yy, ii, ee))
        # ファイル閉
        fd.close()

if __name__ == '__main__':
    import Cmm
    ecm = Cmm.LoadDataFromDump("./", "Load373.dmp")
    p = U2IFJupyter(ecm, 1)

    p.SetKeys(None, None, None)
    p.SetFigureSize(5.0, 5.0)
    p.SetRange(False, 0, 40.0)
    p.SetTitles("U2IFJupyter Test", "YAHOOO!")
    p.SetLabels("|Q| (1/Ang.)", "Energy (meV)")
    p.SetColorMap("haxby")
    p.SaveFile(os.path.join(os.getcwd(), "U2IFJupyter.png"))
