#!/usr/bin/python3
# -*- coding: utf-8 -*-

"""
FlexCalc : S(Q,hw)に対して、Q,hwを使用した計算を行う
"""
from __future__ import print_function

import Manyo
import Manyo.Utsusemi as mu
import os
import sys
import time
import numpy
import importlib


class FlexCalc(object):
    ##################################################
    def __init__(self, DAT, EX1="I", EX2="Err", PATH=""):
        """
        Constructor

        @param DAT   (ElementContainerMatrix)
        @param EX1   (String) Intensity計算式
        @param EX2   (String) Error計算式
        @param PATH  (String) 計算用関数ファイルの保存場所
        """
        self.CommentHead = "FlexCalc >>> "

        # DATはElementContainerMatrixか
        # if type(DAT)!=type(Manyo.ElementContainerMatrix()):
        if not isinstance(DAT, Manyo.ElementContainerMatrix):
            msg = self.CommentHead + "Error: DAT is not ElementContainerMatrix."
            mu.UtsusemiError(msg)
            raise UserWarning("Invalid Arguments")
        """
        PROCESS = DAT.PutHeader().PutStringVector("DATAPROCESSED")
        if not "TOF TO ENERGY CONVERSION" in PROCESS:
            print self.CommentHead+"Error: You must execute TofToEnergyTransfer function on this data"
            raise "InvalidArgument"
        """
        CDP = mu.UtsusemiCheckDataProcess()
        if not CDP.CheckProcess(DAT, "TOF TO ENERGY", False):
            mu.UtsusemiError(
                self.CommentHead + "Error: You must execute TofToEnergyTransfer function on this data")
            raise UserWarning("Invalid Arguments")

        self.DAT = DAT

        # DATのデータの情報取得
        if DAT.PutHeader().CheckKey("SAMPLETYPE") == 1:
            self.SAMPLETYPE = DAT.PutHeader().PutString("SAMPLETYPE")
        else:
            self.SAMPLETYPE = None

        if (PATH != "") and os.path.exists(PATH):
            self.PATH = PATH
        else:
            if "UTSUSEMI_USR_DIR" in os.environ:
                self.PATH = os.path.join(
                    os.environ["UTSUSEMI_USR_DIR"], "ana", "tmp")
            elif "HOME" in os.environ:
                self.PATH = os.path.join(os.environ["HOME"], "ana", "tmp")
            else:
                self.PATH = os.path.join("ana", "tmp")
            if not os.path.exists(self.PATH):
                self.PATH = ""

        self.filename = "TempFlexCalcFile.py"

        # Import the temporal file to calc at first to be reloaded later.
        self.EX1 = "I"
        self.EX2 = "Err"
        try:
            self._MakeTempFile()
        except:
            raise UserWarning("FlexCalc Error")

        sys.path.insert(0, self.PATH)
        try:
            self.TF = __import__(self.filename.split(".")[0])
        except:
            raise UserWarning("FlexFuncImportError: file = {}".format(self.filename))

        # 計算式は入っているか
        if (EX1 != "I") and (not isinstance(EX1, str)):
            raise UserWarning("Error: EX1 is not String.")

        if (EX2 != "Err") and (not isinstance(EX2, str)):
            raise UserWarning("Error: EX2 is not String.")

        # 計算式や関数ファイルの名前の登録
        self.EX1 = EX1
        self.EX2 = EX2

        # 計算式が空でなければ関数ファイル作成
        if self.EX1 != "I":
            self._MakeTempFile()

    def _MakeTempFile(self):
        """
        関数ファイルの作成
        これを読み込んで実際の計算に用いる。
        @retval None
        """
        # データがPowderなら、引数は|Q|、PowderでなければQx,Qy,Qzで関数を作成
        if self.SAMPLETYPE == "Powder":
            mu.UtsusemiMessage(self.CommentHead + "SAMPELTYPE is Powder.")
            SS = "from numpy import *\n"
            SS += "def func(hw,I,Err,Q):\n"
            SS += "    I={}\n".format(self.EX1)
            SS += "    Err={}\n".format(self.EX2)
            SS += "    return (I,Err)\n"
        else:
            mu.UtsusemiMessage(self.CommentHead + "SAMPELTYPE is not Powder.")
            SS = "from numpy import *\n"
            SS += "def func(hw,I,Err,Qx,Qy,Qz):\n"
            SS += "    I={}\n".format(self.EX1)
            SS += "    Err={}\n".format(self.EX2)
            SS += "    return (I,Err)\n"

        # 関数ファイルのオープンと書き込み
        try:
            fi = open(os.path.join(self.PATH, self.filename), "w")
        except PermissionError:
            mu.UtsusemiError(self.CommentHead + "Cannot open file:" + os.path.join(self.PATH, self.filename))
            raise UserWarning("ErrorOnOpenFile")

        fi.write(SS)
        fi.close()

    def DoFuncSelf(self):
        """
        計算実行
        @retval None
        """
        self._dofunc(self.DAT)

    def DoFunc(self):
        """
        オリジナルデータをコピーして計算の実行
        @retval new ElementContainerMatrix
        """
        DAT = Manyo.ElementContainerMatrix(self.DAT)
        self._dofunc(DAT)
        return DAT

    def _dofunc(self, DAT):
        """
        計算実行
        関数ファイルを読み込み関数登録して、実行する
        @param DAT (ElementContainerMatrix)
        @retval None
        """

        # 関数ファイルの再インポート
        try:
            self.TF = importlib.reload(self.TF)
        except:
            raise UserWarning("FlexFuncImportError: failed to reload")

        # 入射エネルギーの入手とki算出
        Ei = self.DAT.PutHeader().PutDouble("Ei")
        Ki = numpy.sqrt(Ei / 2.072)

        # 関数の実行
        if self.SAMPLETYPE == "Powder":
            for i in range(DAT.PutTableSize()):
                eca = DAT(i)
                mu.UtsusemiMessage(self.CommentHead + "PSD no={}".format(i))
                for j in range(eca.PutTableSize()):
                    pangle_vec = eca(j).PutHeader(
                    ).PutDoubleVector("PixelPolarAngle")

                    xx_org = eca(j).PutXList()
                    xx = []
                    for k in range(len(xx_org) - 1):
                        xx.append((xx_org[k] + xx_org[k + 1]) / 2.0)
                    xx = numpy.array(xx)
                    yy = numpy.array(eca(j).PutYList())
                    ee = numpy.array(eca(j).PutEList())
                    qq = self._calcQarray(Ei, Ki, pangle_vec[0], xx)
                    xkey = eca(j).PutXKey()
                    ykey = eca(j).PutYKey()
                    ekey = eca(j).PutEKey()

                    (res_yy, res_ee) = self.TF.func(xx, yy, ee, qq)
                    eca(j).Replace(ykey, res_yy.tolist())
                    eca(j).Replace(ekey, res_ee.tolist())
                    eca(j).SetKeys(xkey, ykey, ekey)
            # print xx
            # print qq
        else:
            for i in range(DAT.PutTableSize()):
                eca = DAT(i)
                mu.UtsusemiMessage(self.CommentHead + "PSD no={}".format(i))
                for j in range(eca.PutTableSize()):
                    pangle_vec = eca(j).PutHeader(
                    ).PutDoubleVector("PixelPolarAngle")
                    aangle_vec = eca(j).PutHeader(
                    ).PutDoubleVector("PixelAzimAngle")
                    xx_org = eca(j).PutXList()
                    xx = []
                    for k in range(len(xx_org) - 1):
                        xx.append((xx_org[k] + xx_org[k + 1]) / 2.0)
                    xx = numpy.array(xx)
                    yy = numpy.array(eca(j).PutYList())
                    ee = numpy.array(eca(j).PutEList())
                    (qx, qy, qz) = self._calcQxyzarray(
                        Ei, Ki, pangle_vec[0], aangle_vec[0], xx)
                    xkey = eca(j).PutXKey()
                    ykey = eca(j).PutYKey()
                    ekey = eca(j).PutEKey()

                    (res_yy, res_ee) = self.TF.func(xx, yy, ee, qx, qy, qz)
                    eca(j).Replace(ykey, res_yy.tolist())
                    eca(j).Replace(ekey, res_ee.tolist())
                    eca(j).SetKeys(xkey, ykey, ekey)

    def _calcQarray(self, Ei, Ki, pangle, hw):
        """
        |Q|の計算
        @param Ei     (Double) 入射エネルギー
        @param Ki     (Double) 入射波数ベクトル
        @param pangle (Double) PixelのPolar Angle
        @param hw     (numpy.array) エネルギー遷移情報
        @retval numpy.array of MomentTransfer
        """
        Kf = numpy.sqrt((Ei - hw) / 2.072)
        return numpy.sqrt(Ki * Ki + Kf * Kf - 2. * Ki * Kf * numpy.cos(pangle * numpy.pi / 180.0))

    def _calcQxyzarray(self, Ei, Ki, pangle, aangle, hw):
        """
        (Qx, Qy, Qz)の計算
        @param Ei     (Double) 入射エネルギー
        @param Ki     (Double) 入射波数ベクトル
        @param pangle (Double) PixelのPolar Angle
        @param aangle (Double) Azimuthal Angle
        @param hw     (numpy.array) エネルギー遷移情報
        @retval tupples of numpy.array of MomentTransfer
        """
        Kf = numpy.sqrt((Ei - hw) / 2.072)
        qx = Kf * numpy.sin(pangle * numpy.pi / 180.0) * numpy.cos(aangle * numpy.pi / 180.0)
        qy = Kf * numpy.sin(pangle * numpy.pi / 180.0) * numpy.sin(aangle * numpy.pi / 180.0)
        qz = Kf * numpy.cos(pangle * numpy.pi / 180.0)
        return (qx, qy, qz)

    def finish(self):
        """
        終了処理
        """
        if sys.platform == 'linux':
            com = "rm -rf " + os.path.join(self.PATH, self.filename)
            os.system(com)
