import numpy as np

from ._cython import build_sumtree_from_array
from ._cython import get_prefix_sum_idx
from ._cython import strided_sum
from ._cython import update_sumtree

def _temporarily_enable_update(func):
    def wrapper(self, *args, **kwargs):
        self._enable_writes(True)
        try:
            output = func(self, *args, **kwargs)
        finally:
            self._enable_writes(False)
        return output

    return wrapper

class SumTreeArray(np.ndarray):
    """
    A subclass of an ``numpy.ndarray`` that maintains an internal sumtree
    for fast Categorical distribution sampling and fast sum operations,
    at the expense of slower write operations. Because ``SumTreeArray``
    is a subclass of ``numpy.ndarray``, it can be used just like a
    ``numpy.ndarray``. It also comes with three new methods: ``sample``,
    ``get_prefix_sum_id``, and ``sumtree``.

    All elements of ``SumTreeArray`` must be non-negative (explanation below).

    Parameters
    ----------
    shape_or_array : int, tuple of ints, array-like
        When an int or tuple of ints is provided, a ``SumTreeArray`` is created
        from an array of zeros of the given shape. When array-like, a
        ``SumTreeArray`` is created from the given array.
    dtype : data-type, optional
        When provided creates a ``SumTreeArray`` of the specified dtype. If not
        provided, and ``shape_or_array`` is a shape, then an array of zeros with
        dtype ``float`` is created.  Defaults to ``None``. All integer and floating
        point dtypes are supported. Because of the non-negativity requirement,
        complex dtypes are not supported.

    Notes
    -----
    Because the elements of a ``SumTreeArray`` represent an unnormalized
    Categorical probability distribution, we require that all elements of a
    ``SumTreeArray`` be non-negative.

    For the sake of the integrity of the sumtree, the memory of the array is
    carefully guarded. Elements of ``SumTreeArray`` should only be updated through
    the ``SumTreeArray`` API.  This ensures that the underlying sumtree correctly
    models the underlying array. For example, when creating a ``SumTreeArray``
    from ``another_array``, the new ``SumTreeArray`` uses a copy of ``another_array``.
    Thus, changes to ``another_array`` do not affect the ``SumTreeArray``. As another
    example, using ``SumTreeArray.view(np.ndarray)`` will return an object
    with read-only access to the underlying array ``SumTreeArray``.

    Similarly, when retrieving elements of a ``SumTreeArray`` through indexing,
    the returned object is always an ``np.ndarray`` that is a copy or read-only
    view of the underlying array in ``SumTreeArray``, and thus the integrity of the
    sumtree is protected. Another reason for doing this is that we always assume
    that we do not want to re-compute a new sumtree on top of the returned object,
    which could be unnecessarily expensive.

    References
    ----------
    [1] NumPy https://numpy.org

    Examples
    --------
    >>> SumTreeArray(4,dtype='int32')
    SumTreeArray([0, 0, 0, 0], dtype=int32)
    >>> SumTreeArray((2,2),dtype='int32')
    SumTreeArray([[0, 0],
                  [0, 0]], dtype=int32)
    >>> sum_tree = SumTreeArray(np.array([1,2,3,4],dtype='float32'))
    >>> sum_tree
    SumTreeArray([1., 2., 3., 4.], dtype=float32)
    >>> # set and get just like an ndarray
    >>> sum_tree[:2] = [2,1]
    >>> sum_tree
    SumTreeArray([2., 1., 3., 4.], dtype=float32)
    """

    def __new__(self, shape_or_array, dtype=None):

        if isinstance(shape_or_array, SumTreeArray):
            if dtype is None:
                return shape_or_array
            else:
                return SumTreeArray(shape_or_array.view(np.ndarray), dtype)

        elif isinstance(shape_or_array, np.ndarray):
            if shape_or_array.size <= 1:
                raise ValueError(
                    "input to SumTreeArray must have shape with at least 2 elements"
                )
            assert shape_or_array.size > 1
            dtype = shape_or_array.dtype if dtype is None else dtype
            array = shape_or_array.astype(dtype, copy=True)  # strictly copies
            return array.view(SumTreeArray)

        else:
            dtype = float if dtype is None else dtype
            array = np.zeros(shape_or_array, dtype=dtype)
            if array.size <= 1:
                raise ValueError(
                    "input to SumTreeArray must have shape with at least 2 elements"
                )
            return array.view(SumTreeArray)

    @_temporarily_enable_update
    def __array_finalize__(self, array):

        if not np.shares_memory(array, self):
            # note that we would have ended up here without overriding copy(...)
            raise NotImplementedError("input array and self must share memory")

        if isinstance(self.base, SumTreeArray) and self.base.dtype == self.dtype:
            # inherit the same base and sum tree
            if array.size != self.size:
                # we should never end up here
                raise NotImplementedError(
                    "input array and base SumTreeArray must have same number of elements"
                )
            self._flat_base = self.base._flat_base
            self._indices = self.base._indices.reshape(self.shape)
            self._sumtree = self.base._sumtree
        else:
            # initialize
            self._flat_base = array.view(np.ndarray).ravel()
            self._indices = np.arange(array.size, dtype=np.intp).reshape(array.shape)
            self._sumtree = np.zeros_like(self._flat_base)
            # sumtree needs to be initialize
            self._rebuild_sumtree()

    # When a transformation is applied to a SumTreeArray object, it is assumed that
    # we do not want a new SumTreeArray object (which could result in a large
    # number of unwanted prefix sum tree updates)...and thus the transformation
    # is applied to the underlying array object, and an NDArray is returned.
    # The exception to this rule is in-place operators, such as +=
    def __array_prepare__(self, out_arr, context=None):
        return out_arr.view(np.ndarray)

    def __array_wrap__(self, out_arr, context=None):
        return out_arr.view(np.ndarray)

    @_temporarily_enable_update
    def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
        inputs = list(inputs)
        for i, x in enumerate(inputs):
            if isinstance(x, SumTreeArray):
                inputs[i] = x.view(np.ndarray)
        inputs = tuple(inputs)
        if "out" in kwargs and kwargs["out"] is not None:
            out = kwargs["out"]
            if len(out) == 1 and out[0] is self:
                # this is an in-place update on self
                # proceed with update and then rebuild sumtree
                kwargs["out"] = (self.view(np.ndarray),)
                _ = super(SumTreeArray, self).__array_ufunc__(
                    ufunc, method, *inputs, **kwargs
                )
                self._rebuild_sumtree()
                return self
            else:
                raise NotImplementedError
        else:
            return super(SumTreeArray, self).__array_ufunc__(
                ufunc, method, *inputs, **kwargs
            )

    def _enable_writes(self, val):
        self.setflags(write=val)

    def __setitem__(self, idx, val):
        self.put(idx, val)

    def __getitem__(self, idx):
        return self.view(np.ndarray)[idx]

    def _rebuild_sumtree(self):
        build_sumtree_from_array(self._flat_base, self._sumtree)

    def astype(self, *args, **kwargs):
        return self.view(np.ndarray).astype(*args, **kwargs)

    def choose(self, *args, **kwargs):
        return self.view(np.ndarray).choose(*args, **kwargs)

    def copy(self, *args, **kwargs):
        return SumTreeArray(self.view(np.ndarray))

    def diagonal(self, *args, **kwargs):
        return self.view(np.ndarray).diagonal(*args, **kwargs)

    def dot(self, *args, **kwargs):
        return self.view(np.ndarray).dot(*args, **kwargs)

    @_temporarily_enable_update
    def fill(self, *args, **kwargs):
        super(SumTreeArray, self).fill(*args, **kwargs)
        self._rebuild_sumtree()

    def flatten(self, *args, **kwargs):
        return self.view(np.ndarray).flatten(*args, **kwargs)

    @property
    def imag(self):
        return self.view(np.ndarray).imag

    def mean(self, *args, **kwargs):
        output = self.sum(*args, **kwargs)
        m = np.array(output).size
        n = self.size
        assert n % m == 0
        return output / float(n / m)

    def newbyteorder(self, *args, **kwargs):
        return self.view(np.ndarray).newbyteorder(*args, **kwargs)

    @_temporarily_enable_update
    def partition(self, *args, **kwargs):
        super(SumTreeArray, self).partition(*args, **kwargs)
        self._rebuild_sumtree()

    def put(self, indices, values):
        # TODO: there's probably a better way of building an index iterator
        indices = np.ascontiguousarray(self._indices[indices]).ravel()
        values = np.ascontiguousarray(values, dtype=self._flat_base.dtype).ravel()
        update_sumtree(indices, values, self._flat_base, self._sumtree)

    @property
    def real(self):
        return self.view(np.ndarray).real

    def repeat(self, *args, **kwargs):
        return self.view(np.ndarray).repeat(*args, **kwargs)

    def round(self, *args, **kwargs):
        return self.view(np.ndarray).round(*args, **kwargs)

    @_temporarily_enable_update
    def sort(self, *args, **kwargs):
        super(SumTreeArray, self).sort(*args, **kwargs)
        self._rebuild_sumtree()

    def squeeze(self):
        return self.reshape(self.view(np.ndarray).squeeze().shape)

    def sumtree(self):
        """
        Returns a copy of the sumtree.

        Returns
        -------
        sumtree : ndarray
            A copy of the (flat) sumtree object maintained by ``SumTreeArray``.

        Examples
        --------
        >>> sum_tree = SumTreeArray(np.array([1,2,3,4],dtype='float32'))
        >>> sum_tree
        SumTreeArray([1., 2., 3., 4.], dtype=float32)
        >>> sum_tree.sumtree()
        array([ 0., 10.,  3.,  7.], dtype=float32)
        >>> sum_tree_from_2d_array = SumTreeArray(np.array([[1,2],[3,4]],dtype='int32'))
        >>> sum_tree_from_2d_array
        SumTreeArray([[1., 2.],
                      [3., 4.]], dtype=float32)
        >>> sum_tree_from_2d_array.sumtree()
        array([ 0., 10.,  3.,  7.], dtype=float32)
        """
        return np.copy(self._sumtree)

    def swapaxes(self, *args, **kwargs):
        return self.view(np.ndarray).swapaxes(*args, **kwargs)

    @property
    def T(self):
        return self.view(np.ndarray).T

    def take(self, *args, **kwargs):
        return self.view(np.ndarray).take(*args, **kwargs)

    def trace(self, *args, **kwargs):
        return self.view(np.ndarray).trace(*args, **kwargs)

    def transpose(self, *args, **kwargs):
        return self.view(np.ndarray).transpose(*args, **kwargs)

    def get_prefix_sum_id(self, prefix_sum, flatten_indices=False):
        """
        Returns an array of indices of the same shape is the input array
        ``prefix_sum`` where each element ``i`` in the output is ``j``
        such that ``self.ravel()[:j+1] < prefix_sum[i]``.  In other words,
        the output array returns the index of the largest prefix sum of
        self that is less than the provided input.

        Parameters
        ----------
        prefix_sum : array-like
            An n-dimensional array of prefix sums.  Does not need to be
            of the same shape as ``self``.
        flatten_indices : bool, optional
            Defaults to ``False``.  When ``True``, returns the indices
            as if ``self`` is a 1d array (i.e. ``self.ravel()``)

        Returns
        -------
        prefix_sum_indices : ndarray, tuple(ndarray)
            A new array holding the result is returned containing the
            indices of the supplied prefix sums. The returned array
            has the same shape as the input array. When supplied with
            ``flatten_indices=False`` and ``self`` is an n-d array
            where ``n>1``, then a tuple of ``n`` arrays is returned
            where each element of the tuple corresponds to each dimension
            of the ``SumTreeArray``.

        Notes
        -----
        Arithmetic is modular when using integer types, and no error is
        raised on overflow.

        Examples
        --------
        >>> sum_tree = SumTreeArray(np.array([1,2,3,4],dtype='float32'))
        >>> sum_tree
        SumTreeArray([1., 2., 3., 4.], dtype=float32)
        >>> sum_tree.get_prefix_sum_id(np.arange(10))
        array([0, 1, 1, 2, 2, 2, 3, 3, 3, 3], dtype=int32)
        >>> sum_tree.get_prefix_sum_id(np.arange(10).reshape(2,5))
        array([[0, 1, 1, 2, 2],
               [2, 3, 3, 3, 3]], dtype=int32)
        >>> sum_tree_from_2d_array = SumTreeArray(np.array([[1,2],[3,4]],dtype='int32'))
        >>> sum_tree_from_2d_array
        SumTreeArray([[1., 2.],
                       [3., 4.]], dtype=float32)
        >>> # output is "flattened" and thus will be the same as ``sum_tree.get_prefix_sum_id`` above
        >>> sum_tree_from_2d_array.get_prefix_sum_id(np.arange(10))
        array([0, 1, 1, 2, 2, 2, 3, 3, 3, 3], dtype=int32)
        >>> sum_tree_from_2d_array.get_prefix_sum_id(np.arange(10).reshape(2,5))
        array([[0, 1, 1, 2, 2],
               [2, 3, 3, 3, 3]], dtype=int32)
        """
        # ensure prefix sum is the correct type and is contiguous
        prefix_sum = np.ascontiguousarray(prefix_sum, dtype=self.dtype)
        prefix_sum_flat = prefix_sum.ravel()
        # init return array
        flat_idx = np.zeros(prefix_sum.size, dtype=np.intp)
        # get ids
        get_prefix_sum_idx(flat_idx, prefix_sum_flat, self._flat_base, self._sumtree)
        # output shape should be same as prefix_sum shape
        if prefix_sum.ndim > 1:
            output = flat_idx.reshape(prefix_sum.shape)
        else:
            output = flat_idx
        if self.ndim <= 1 or flatten_indices:
            return output
        else:
            return np.unravel_index(output, self.shape)

    def sample(self, nsamples=1, flatten_indices=False):
        """
        Return a sample of indices, where the probability an index being
        sampled is equal to the value at that index divided by the
        sum of the array, i.e. ``probs = self/self.sum()``.

        Parameters
        ----------
        nsamples : int
            Number of samples to return
        flatten_indices : bool, optional
            Defaults to ``False``.  When ``True``, returns the indices
            as if ``self`` is a 1d array (i.e. ``self.ravel()``)

        Returns
        -------
        sample_of_indices : ndarray, tuple(ndarray)
            A new array holding the result is returned containing the
            sampled indices.  If the underlying ``SumTreeArray`` array
            is n-dimensional, where n>1, and ``flatten_indices=False``,
            then a tuple of ``n`` arrays is returned where each element of
            the tuple corresponds to each dimension of the ``SumTreeArray``
            array.

        Notes
        -----
        Arithmetic is modular when using integer types, and no error is
        raised on overflow.

        Examples
        --------
        >>> sum_tree = SumTreeArray(np.array([1,2,3,4],dtype='float32'))
        >>> sum_tree
        SumTreeArray([1., 2., 3., 4.], dtype=float32)
        >>> sum_tree.sample(10)
        array([2, 3, 3, 3, 3, 1, 2, 2, 2, 0], dtype=int32)
        >>> # probability of being sampled
        >>> sum_tree / sum_tree.sum()
        array([0.1, 0.2, 0.3, 0.4], dtype=float32)
        >>> # sampled proportions
        >>> (sum_tree.sample(1000)[None] == np.arange(4)[:,None]).mean(axis=1)
        array([0.10057, 0.19919, 0.29983, 0.40041])
        >>> # sampling from a 2-d array
        >>> sum_tree_from_2d_array = SumTreeArray(np.array([[1,2,3],[4,5,6]],dtype='int32'))
        >>> sum_tree_from_2d_array
        SumTreeArray([[1, 2, 3],
                      [4, 5, 6]], dtype=int32)
        >>> sum_tree_from_2d_array.sample(4)
        (array([1, 1, 1, 0]), array([1, 1, 2, 2]))
        >>> sum_tree_from_2d_array.sample(4,flatten_indices=True)
        array([4, 4, 5, 2], dtype=int32)
        """
        if self.sum() == 0:
            raise ValueError("array must have at least 1 positive value")
        # sample priority values in the cumulative sum
        vals = (self.sum() * np.random.rand(nsamples)).astype(self.dtype)
        return self.get_prefix_sum_id(vals, flatten_indices)

    def _parse_axis_arg(self, axis):
        if axis is None:
            return None
        else:
            axes = np.array(axis).reshape(-1)
            if len(set(axes)) < len(axes):
                raise IndexError(
                    "invalid axis argument: %s contains duplicates" % str(axis)
                )
            return np.arange(self.ndim)[axes]

    def sum(self, axis=None, keepdims=False):
        """
        Functions the same as ``numpy.sum(...)`` except that some ``sum``
        operations can be significantly faster for large arrays
        when the sum operation is performed over C-contiguous
        ranges of indices.  For example, for a 2d array, supplying
        ``axis=1`` to the ``sum`` function would be a sum operation over
        C-contiguous ranges of indices, and thus benefit from sum-tree
        speedups.  However, ``axis=0`` is not C-contiguous sum operation,
        and would thus revert to the standard ``numpy.sum`` method.

        Parameters
        ----------
        axis : None or int or tuple of ints, optional
            Axis or axes along which a sum is performed. The default,
            ``axis=None``, will sum all of the elements of the input array.
            If axis is negative it counts from the last to the first axis.
            If axis is a tuple of ints, a sum is performed on all of the
            axes specified in the tuple instead of a single axis or all the
            axes as before.
        keepdims : bool, optional
            If this is set to True, the axes which are reduced are left in
            the result as dimensions with size one. With this option, the
            result will broadcast correctly against the input array. Defaults
            to ``False``.

        Returns
        -------
        sum_along_axis : ndarray
            An array with the same shape as self, with the specified axis
            or axes removed, unless ``keepdims=True`` in which case the
            specified axis or axes are not removed, but will have size ``1``.
            If axis is None, a scalar is returned.

        Notes
        -----
        Arithmetic is modular when using integer types, and no error is raised on overflow.
        SumTreeArray does not use the improved precision techniques that NumPy uses when
        summing along the fast axis (C-contiguous), instead relying on the standard floating
        point precision arithmetic provided by Python.

        References
        ----------
        [1] NumPy.  For more information on ``numpy.sum``, please see
        https://numpy.org/doc/stable/reference/generated/numpy.sum.html

        Examples
        --------
        >>> x = SumTreeArray(np.ones((1000,1000)))
        >>> # lightning fast
        >>> x.sum()
        1000000.0
        >>> # very fast
        >>> x.sum(axis=1)
        array([1000., ..., 1000.], dtype='float64')
        >>> # slightly slower than NumPy
        >>> x.sum(axis=0)
        array([1000., ..., 1000.], dtype='float64')
        >>> # maintain dims with keepdims=True
        >>> x.sum(axis=1,keepdims=True).shape
        (1000, 1)
        >>> x.sum(axis=1).shape
        (1000,)
        """
        axes = self._parse_axis_arg(axis)
        if axes is None:
            if keepdims:
                return self._sumtree[1:2].reshape([1] * self.ndim)
            else:
                return self._sumtree[1]
        else:
            if axes.min() == self.ndim - len(axes):
                # strides are contiguous along leaves of sumtree
                stride = int(np.prod(np.array(self.shape)[axes]))
                output = strided_sum(self._flat_base, self._sumtree, stride)
                if keepdims:
                    return output.reshape(list(output.shape) + [1] * len(axes))
                else:
                    return output
            else:
                return super(SumTreeArray, self).sum(axis=axis, keepdims=keepdims)


'''
    @Author: fxm
    @Date: Dec 27, 2020.
    @Title: Learn class.
'''

import logging
import os
import numpy as np
from collections import deque
from random import shuffle

import numpy as np
from tqdm import tqdm

from Framework.Mcts import MCTS
from Framework.Net import dotdict
from SelfPlay import SelfPlay

log = logging.getLogger(__name__)

args = dotdict({
    'numIters': 100,  # 神经网络训练迭代次数
    'numEps': 100,  # 自学习过程每次比赛场数
    'tempThreshold': 15,  # 搜索树的阈值
    'updateThreshold': 0.5,  # 如果新模型赢了超过updateThreshold比例的比赛，就接受新模型
    'maxlenOfQueue': 200000,  # 训练数据上限
    'numMCTSSims': 25,  # 搜索树采样数量
    'arenaCompare': 2,  # 新旧模型对比需要的比赛次数
    'cpuct': 1,  # 搜索树参数
    'checkpoint': './Model/',  # 模型保存路径
    'load_model': False,  # 是否加载模型
    'numItersFortable': 20,
})

class Learn():
    '''
        初始化
        参数设置:
        game：游戏对象
        net：网络对象
        pnet：竞争对手的网络对象
        mcts：蒙特卡洛树对象
        args：其他参数
        table：保存的训练数据，每次训练都使用这里面的数据
    '''

    def __init__(self, game, net, args):
        self.game = game
        self.net = net
        self.pnet = self.net.__class__(self.game)
        self.args = args
        self.mcts = MCTS(self.game, self.net, self.args)
        self.table = []
        self.firstrain = True

    '''
        process过程算法
        将蒙特卡洛树搜索得到的
        棋盘(即状态)、得到的action概率向量、值
        传入神经网络进行学习
    '''

    def process(self):
        # 结果保存表
        table = []
        board = self.game.initBoard()
        # 由于是学习过程，谁先手无所谓
        self.curcolor = 1
        step = 0

        # 只要比赛不结束就一直进行
        while True:
            step += 1
            # 当前玩家所使用棋盘对象
            canonicalBoard = self.game.getCanonicalForm(board, self.curcolor)
            # 如果step大于阈值，则不再计算准确的概率向量，而是将概率最大的设为1，其他为0
            temp = int(step < self.args.tempThreshold)

            # 获得当前的概率向量
            prob = self.mcts.getActionProb(canonicalBoard, temp=temp)
            sym = self.game.getSymmetries(canonicalBoard, prob)
            for b, p in sym:
                table.append([b, self.curcolor, p, None])

            # 执行动作后获得下一状态，并交换掌棋权，然后继续执行上述过程
            action = np.random.choice(len(prob), p=prob)
            board, self.curcolor = self.game.getNextState(board, self.curcolor, action)
            ended = self.game.getGameEnded(board, self.curcolor)

            # 如果结束就返回结果，是一个三元组
            # 最终的结果包括了搜索轨迹中所有的状态对应的
            # 状态本身，动作action概率向量，以及最有希望得到的值
            if ended != -2:
                return [(x[0], x[2], ended * ((-1) ** (x[1] != self.curcolor))) for x in table]

    '''
        学习过程
        AC模块采用神经网络学习
        process模块采用蒙特卡洛树搜索过程
    '''

    def learn(self):
        for i in range(1, self.args.numIters + 1):
            log.info(f'Starting Iter #{i} ...')
            deq = deque([], maxlen=self.args.maxlenOfQueue)

            # 在自学习过程中，执行process算法过程
            for _ in tqdm(range(self.args.numEps), desc="Self Play"):
                self.mcts = MCTS(self.game, self.net, self.args)
                deq += self.process()

            # 将训练结果保存到table中
            self.table.append(deq)

            if len(self.table) > self.args.numItersFortable:
                log.warning(f"Removing the oldest entry in data. len(table) = {len(self.table)}")
                self.table.pop(0)

            # 训练前需要将数据打乱
            data = []
            for e in self.table:
                data.extend(e)
            shuffle(data)

            # 训练新的网络前需要保存旧的网络
            self.net.saveCheckpoint(folder=self.args.checkpoint, filename='temp.pth.tar')
            self.pnet.loadCheckpoint(folder=self.args.checkpoint, filename='temp.pth.tar')
            pmcts = MCTS(self.game, self.pnet, self.args)

            self.net.train(data)
            mcts = MCTS(self.game, self.net, self.args)

            # 如果是第一次训练，那么不需要比较网络
            if self.firstrain == True:
                self.net.saveCheckpoint(folder=self.args.checkpoint, filename='temp.pth.tar')
                self.net.saveCheckpoint(folder=self.args.checkpoint, filename='best.pth.tar')
                self.firstrain = False
                continue

            log.info('与旧的网络对比中：')
            arena = SelfPlay(lambda x: np.argmax(pmcts.getActionProb(x, temp=0)),
                             lambda x: np.argmax(mcts.getActionProb(x, temp=0)), self.game)
            pwins, wins, draws = arena.playGames(self.args.arenaCompare)

            log.info('胜/负 : %d / %d ; 平局 : %d' % (wins, pwins, draws))
            if pwins + wins == 0 or float(wins) / (pwins + wins) <= self.args.updateThreshold:
                log.info('不接受新的模型')
                self.net.loadCheckpoint(folder=self.args.checkpoint, filename='temp.pth.tar')
            else:
                log.info('接受新的模型')
                self.net.saveCheckpoint(folder=self.args.checkpoint, filename='best.pth.tar')


"""Methods for decorrelating simulation results."""

import copy
import math

import numpy as np
from pymbar import timeseries

from origamipy import datatypes
from origamipy import utility

NUM_STAPLES_TAG = 'numstaples'

class DecorrelatedOutputs:
    _datatypes = ['enes', 'ops', 'staples', 'staplestates']
    _trjtypes = ['trj', 'vcf', 'ores', 'states']
#    _trjtypes = ['trj', 'vcf', 'states']

    def __init__(self, sim_collections, all_conditions):
        self.all_conditions = all_conditions
        self._sim_collections = sim_collections
        self._decor_masks = []
        self._num_decorrelated_steps = 0
        self._datatype_to_decors = {}
        self._trjtype_to_decors = {}

    def get_concatenated_datatype(self, tag):
        concat = []
        for data in self._datatype_to_decors[tag]:
            concat.append(datatypes.OutputData.concatenate(data))

        return datatypes.OutputData.concatenate(concat)

    def get_num_steps_per_condition(self):
        steps = []
        for rep_to_data in self._datatype_to_decors['enes']:
            steps.append(sum([s.steps for s in rep_to_data]))

        return steps

    def get_concatenated_series(self, tag):
        for reps_data in self._datatype_to_decors.values():
            if tag in reps_data[0][0].tags:
                concat = []
                for data in reps_data:
                    reps = []
                    for series in data:
                        reps.append(series[tag])

                    concat.append(np.concatenate(reps))

                return np.concatenate(concat)

        else:
            raise Exception

    @property
    def all_series_tags(self):
        tags = []
        for decors in self._datatype_to_decors.values():
            datatype = decors[0][0]
            for tag in datatype.tags:
                if tag == 'step':
                    continue

                tags.append(tag)

        return tags

    def perform_decorrelation(self, skip):
        print('Performing decorrelations')
        print('State,   configs, t0, g,   Neff')
        for sim_collection in self._sim_collections:
            self._decor_masks.append([])
            for rep in sim_collection._reps:
                mask = self._construct_decorrelation_mask(sim_collection, rep,
                        skip)
                self._decor_masks[-1].append(mask)

    def _construct_decorrelation_mask(self, sim_collection, rep, skip):
        enes = sim_collection.reps_energies[rep]
        ops = sim_collection.reps_order_params[rep]
        num_staples = sim_collection.reps_staples[rep]
        steps = enes.steps
        rpots = utility.calc_reduced_potentials(enes, ops, num_staples,
                                                sim_collection.conditions)
        start_i, g, Neff = timeseries.detectEquilibration(rpots, nskip=skip)
        template = '{:<8} {:<8} {:<3} {:<4.1f} {:<.1f}'
        print(template.format(sim_collection.conditions.fileformat, steps,
                start_i, g, Neff))
        indices = (timeseries.subsampleCorrelatedData(rpots[start_i:], g=skip*g))
        return [i + start_i for i in indices]

    def read_decors_from_files(self, data_only=False):
        for datatype in self._datatypes:
            self._datatype_to_decors[datatype] = []
            for sim_collection in self._sim_collections:
                reps_series = sim_collection.get_decor_reps_data(datatype)
                self._datatype_to_decors[datatype].append(reps_series)

        if not data_only:
            for trjtype in self._trjtypes:
                self._trjtype_to_decors[trjtype] = []
                for sim_collection in self._sim_collections:
                    reps_series = sim_collection.get_decor_reps_trj(trjtype)
                    self._trjtype_to_decors[trjtype].append(reps_series)

    def apply_masks(self):
        # The mask numbering is different than the rep number
        for datatype in self._datatypes:
            self._datatype_to_decors[datatype] = []
            for i, sim_collection in enumerate(self._sim_collections):
                self._datatype_to_decors[datatype].append([])
                reps_to_data = sim_collection.get_reps_data(datatype)
                for j, rep in enumerate(sim_collection._reps):
                    data = reps_to_data[rep]
                    data.apply_mask(self._decor_masks[i][j])
                    self._datatype_to_decors[datatype][i].append(data)

        for trjtype in self._trjtypes:
            self._trjtype_to_decors[trjtype] = []
            for i, sim_collection in enumerate(self._sim_collections):
                self._trjtype_to_decors[trjtype].append([])
                reps_to_trjs = sim_collection.get_reps_trj(trjtype)
                for j, rep in enumerate(sim_collection._reps):
                    trjs = reps_to_trjs[rep]
                    filebase = sim_collection.decor_filebase_template.format(
                            sim_collection.filebase, rep,
                            sim_collection.conditions.fileformat)
                    filename = '{}.{}'.format(filebase, trjtype)
                    decor_trj = self._apply_mask_to_trjs(self._decor_masks[i][j],
                            trjs, filename)
                    self._trjtype_to_decors[trjtype][i].append(decor_trj)

    def _apply_mask_to_trjs(self, mask, trjs, filename):
        out_file = open(filename, 'w')
        step_i = 0
        mask_i = 0
        for trj in trjs:
            for step in trj:
                step_included = step_i == mask[mask_i]
                if step_included:
                    out_file.write(step)
                    mask_i += 1
                    if mask_i == len(mask):
                        return

                step_i += 1
            trj.close()

        out_file.close()

    def write_decors_to_files(self):
        for datatype in self._datatypes:
            for i, sim_collection in enumerate(self._sim_collections):
                for j, rep in enumerate(sim_collection._reps):
                    filebase = sim_collection.decor_filebase_template.format(
                            sim_collection.filebase, rep,
                            sim_collection.conditions.fileformat)
                    if datatype == 'enes':
                        self._datatype_to_decors[datatype][i][j].to_file(filebase,
                                float(sim_collection.conditions.temp))
                    else:
                        self._datatype_to_decors[datatype][i][j].to_file(filebase)

    def filter_collections(self, filter_tag, value):
        filtered_count = 0
        for i, sim_collection in enumerate(self._sim_collections):
            for j, rep in enumerate(sim_collection._reps):

                # Create mask
                selected_op = self._datatype_to_decors['ops'][i][j][filter_tag]
                mask = selected_op == value
                filtered_count += mask.sum()

                # Apply mask
                for datatype in self._datatypes:
                    data = self._datatype_to_decors[datatype][i][j]._data
                    reduced_data = []
                    for series in data:
                        reduced_data.append(series[mask])

                    reduced_data = np.array(reduced_data)
                    self._datatype_to_decors[datatype][i][j]._data = reduced_data

        return filtered_count

class SimpleDecorrelatedOutputs:
    _datatypes = ['enes', 'ops', 'staples', 'staplestates']
    _trjtypes = ['trj', 'vcf', 'ores', 'states']

    def __init__(self, sim_collections, all_conditions):
        self.all_conditions = all_conditions
        self._sim_collections = sim_collections
        self._decor_masks = []
        self._num_decorrelated_steps = 0
        self._datatype_to_decors = {}
        self._trjtype_to_decors = {}

    def get_concatenated_datatype(self, tag):
        return datatypes.OutputData.concatenate(self._datatype_to_decors[tag])

    def get_num_steps_per_condition(self):
        steps = []
        for data in self._datatype_to_decors['enes']:
            steps.append(data.steps)

        return steps

    def get_concatenated_series(self, tag):
        for data in self._datatype_to_decors.values():
            if tag in data[0].tags:
                concat = []
                for series in data:
                    concat.append(series[tag])

                return np.concatenate(concat)

        else:
            raise Exception

    @property
    def all_series_tags(self):
        tags = []
        for decors in self._datatype_to_decors.values():
            datatype = decors[0]
            for tag in datatype.tags:
                if tag == 'step':
                    continue

                tags.append(tag)

        return tags

    def perform_decorrelation(self, skip):
        print('Performing decorrelations')
        print('State,   configs, t0, g,   Neff')
        for sim_collection in self._sim_collections:
            mask = self._construct_decorrelation_mask(sim_collection, skip)
            self._decor_masks.append(mask)

    def _construct_decorrelation_mask(self, sim_collection, skip):
        enes = sim_collection.get_data('enes')
        ops = sim_collection.get_data('ops')
        num_staples = sim_collection.get_data('staples')
        steps = enes.steps
        rpots = utility.calc_reduced_potentials(enes, ops, num_staples,
                                                sim_collection.conditions)
        start_i, g, Neff = timeseries.detectEquilibration(rpots, nskip=skip)
        template = '{:<8} {:<8} {:<3} {:<4.1f} {:<.1f}'
        print(template.format(sim_collection.conditions.fileformat, steps,
                start_i, g, Neff))
        indices = (timeseries.subsampleCorrelatedData(rpots[start_i:], g=skip*g))
        return [i + start_i for i in indices]

    def read_decors_from_files(self, data_only=False):
        for datatype in self._datatypes:
            self._datatype_to_decors[datatype] = []
            for sim_collection in self._sim_collections:
                series = sim_collection.get_data(datatype)
                self._datatype_to_decors[datatype].append(series)

        if not data_only:
            for trjtype in self._trjtypes:
                self._trjtype_to_decors[trjtype] = []
                for sim_collection in self._sim_collections:
                    series = sim_collection.get_trj(trjtype)
                    self._trjtype_to_decors[trjtype].append(series)

    def apply_masks(self):
        # The mask numbering is different than the rep number
        for datatype in self._datatypes:
            self._datatype_to_decors[datatype] = []
            for i, sim_collection in enumerate(self._sim_collections):
                data = sim_collection.get_data(datatype)
                data.apply_mask(self._decor_masks[i])
                self._datatype_to_decors[datatype].append(data)

        for trjtype in self._trjtypes:
            self._trjtype_to_decors[trjtype] = []
            for i, sim_collection in enumerate(self._sim_collections):
                trj = sim_collection.get_trj(trjtype)
                filebase = '{}_decor'.format(sim_collection.filebase)
                filename = '{}.{}'.format(filebase, trjtype)
                decor_trj = self._apply_mask_to_trj(self._decor_masks[i],
                        trj, filename)
                self._trjtype_to_decors[trjtype].append(decor_trj)

    def _apply_mask_to_trj(self, mask, trj, filename):
        out_file = open(filename, 'w')
        step_i = 0
        mask_i = 0
        for step in trj:
            step_included = step_i == mask[mask_i]
            if step_included:
                out_file.write(step)
                mask_i += 1
                if mask_i == len(mask):
                    return

            step_i += 1
        trj.close()

        out_file.close()

    def write_decors_to_files(self):
        for datatype in self._datatypes:
            for i, sim_collection in enumerate(self._sim_collections):
                filebase = '{}_decor'.format(sim_collection.filebase)
                if datatype == 'enes':
                    self._datatype_to_decors[datatype][i].to_file(filebase,
                            float(sim_collection.conditions.temp))
                else:
                    self._datatype_to_decors[datatype][i].to_file(filebase)

    def filter_collections(self, filter_tag, value):
        filtered_count = 0
        for i, sim_collection in enumerate(self._sim_collections):

            # Create mask
            selected_op = self._datatype_to_decors['ops'][i][filter_tag]
            mask = selected_op == value
            filtered_count += mask.sum()

            # Apply mask
            for datatype in self._datatypes:
                data = self._datatype_to_decors[datatype][i]._data
                reduced_data = []
                for series in data:
                    reduced_data.append(series[mask])

                reduced_data = np.array(reduced_data)
                self._datatype_to_decors[datatype][i]._data = reduced_data

        return filtered_count


"""

Written by JT Fuchs in July 2015
Based off pySALT redution routine specsens.py by S. Crawford
And reading the darned IRAF documentation

flux_calibration.py performs flux calibration on a 1D and wavelength-calibrated spectrum

To run file:
python flux_calibration.py spec_list --flux_list listflux.txt --stan_list liststandards.txt --extinct False
python flux_calibration.py GD1212.ms.fits --usemaster True

:INPUTS:
    spec_list: either single *.fits file or text file containing list of files to flux calibrate.

:OPTIONS:
    --flux_list: string, file containing standard star fluxes. These are typically m*.dat.

    --stan_list: string, file with list of 1D standard star spectra

    --usemaster: boolean, Option to use master response function instead of single star observation. Default: False

    --extinct: boolean, Option to extinction correct spectra. Default: True

:OUTPUTS:
        flux calibrated files (_flux is added to the filename). User will be prompted if file will overwrite existing file.

        sensitivity_params.txt:  File is updated everytime spec_sens.py is run. Contains information used in the flux calibration. Columns are: input observed spectrum, date/time program was run, observed standard spectrum used for calibration, flux calibration file (m*dat), pixel regions excluded in fit, order of polynomial to flux standard, width in Angstroms used for rebinning, output spectrum filename

        sens_fits_DATE.txt: File for diagnostics. Columns are: wavelength, observed flux, polynomial fit, and residuals for each standard listed above. There are extra zeros at the bottom of some columns.

Each list should have the names of the stars, with blue and red exposures next to each other.
The ordering of the standard star flux files should match the order of the standard star list.
Example:

liststandard:
wtfb.LTT3218_930_blue.ms.fits
wtfb.LTT3218_930_red.ms.fits
wnb.GD50_930_blue.ms.fits
wnb.GD50_930_red.ms.fits

listflux:
mltt3218.dat
mgd50.dat

spec_list:
wnb.WD0122p0030_930_blue.ms.fits
wnb.WD0122p0030_930_red.ms.fits
wnb.WD0235p069_930_blue.ms.fits
wnb.WD0235p069_930_red.ms.fits

#####

Counting variables are fruits and vegetables.

"""

import os
import sys
import numpy as np
#import pyfits as fits
import astropy.io.fits as fits
import spectools as st
import datetime
from glob import glob
import matplotlib.pyplot as plt
from scipy.interpolate import UnivariateSpline
import argparse

#=============================================
#To help with command line interpretation
def str2bool(v):
    if v.lower() in ('yes','true','t','y','1'):
        return True
    if v.lower() in ('no','false','f','n','0'):
        return False
    else:
        raise argparse.ArgumentTypeError('Boolean value expected.')

#=============================================
#These functions are to help with excluding regions from the sensitivity function
def find_nearest(array,value):
    idx = (np.abs(array-value)).argmin()
    return array[idx]

def onclick(event):
    global ix,iy
    ix, iy = event.xdata,event.ydata
    global coords
    ax.axvline(x=ix,color='k',linewidth='3')
    fig.canvas.draw()
    coords.append((ix,iy))

#=============================================

def flux_calibrate_now(stdlist,fluxlist,speclist,extinct_correct=False,masterresp=False):
    if extinct_correct:
        extinctflag = 0
    else:
        extinctflag = -1
    if masterresp: #Use the master response function
        #Read in master response function and use that.
        cwd = os.getcwd()
        os.chdir('/afs/cas.unc.edu/depts/physics_astronomy/clemens/students/group/standards/response_curves/')
        standards = sorted(glob('*resp*.npy'))

        master_response_blue_in = np.load(standards[0])
        master_response_blue_in_pol = np.poly1d(master_response_blue_in)
        master_response_blue_out = np.load(standards[1])
        master_response_blue_out_pol = np.poly1d(master_response_blue_out)
        master_response_red_in = np.load(standards[2])
        master_response_red_in_pol = np.poly1d(master_response_red_in)
        master_response_red_out = np.load(standards[3])
        master_response_red_out_pol = np.poly1d(master_response_red_out)

        os.chdir(cwd)

        airstd = np.ones([4])
        #airstd[0] = 1.1

        #For saving files correctly
        stdflux = np.array(['mmaster.dat'])
        #standards = np.array([masterlist])
        allexcluded = [[None] for i in range(len(standards))]
        orderused = np.zeros([len(standards)])
        size = 0.

        #Find shift for each night
        #For blue setup: use mean of 4530-4590
        #for red setup: use mean of 6090-6190
        try:
            flux_tonight_list = np.genfromtxt('response_curves.txt',dtype=str)
            print 'Found response_curves.txt file.'
            print flux_tonight_list
            if flux_tonight_list.size == 1:
                flux_tonight_list = np.array([flux_tonight_list])
            for x in flux_tonight_list:
                #print x
                if 'blue' in x.lower():
                    wave_tonight, sens_tonight = np.genfromtxt(x,unpack=True)
                    blue_low_index = np.min(np.where(wave_tonight > 4530.))
                    blue_high_index = np.min(np.where(wave_tonight > 4590.))
                    blue_mean_tonight = np.mean(sens_tonight[blue_low_index:blue_high_index])
                elif 'red' in x.lower():
                    wave_tonight, sens_tonight = np.genfromtxt(x,unpack=True)
                    red_low_index = np.min(np.where(wave_tonight > 6090.))
                    red_high_index = np.min(np.where(wave_tonight > 6190.))
                    red_mean_tonight = np.mean(sens_tonight[red_low_index:red_high_index])
        except:
            print 'No response_curves.txt file found.'
            blue_mean_tonight = None
            red_mean_tonight = None
            flux_tonight_list = ['None','None']

    else: #Use the standard star fluxes in the typical manner
        #Read in each standard star spectrum
        standards = np.genfromtxt(stdlist,dtype=str)
        if standards.size ==1:
            standards = np.array([standards])
        stdflux = np.genfromtxt(fluxlist,dtype=str)
        if stdflux.size == 1:
            stdflux = np.array([stdflux]) #save stdflux explicitly as an array so you can index if only 1 element
        #Check that the files are set up correctly to avoid mixing standards.
        #This checks that the files in liststandard have similar characters to those in listflux and the correct order. But might break if flux file doesn't match. E.G. mcd32d9927.dat is often called CD-32_9927 in our system.
        '''
        onion = 0
        for stanspec in standards:
            quickcheck = stdflux[onion//2].lower()[1:-4] in stanspec.lower()
            if not quickcheck:
                print 'Check your standard star and flux files. They are mixed up.'
                sys.exit()
            onion += 1
        '''
        orderused = np.zeros([len(standards)])
        senspolys = []
        airstd = np.zeros([len(standards)])
        allexcluded = [[None] for i in range(len(standards))]

        #Calculating the sensitivity function of each standard star
        cucumber = 0
        for stdspecfile in standards:
            print stdspecfile
            #Read in the observed spectrum of the standard star
            obs_spectra,airmass,exptime,dispersion = st.readspectrum(stdspecfile) #obs_spectra is an object containing opfarr,farr,sky,sigma,warr
            airstd[cucumber] = airmass
            #plt.clf()
            #plt.plot(obs_spectra.warr,obs_spectra.opfarr)
            #plt.show()

            #Do the extinction correction
            if extinct_correct:
                print 'Extinction correcting spectra.'
                plt.clf()
                plt.plot(obs_spectra.warr,obs_spectra.opfarr)
                obs_spectra.opfarr = st.extinction_correction(obs_spectra.warr,obs_spectra.opfarr,airmass)
                plt.plot(obs_spectra.warr,obs_spectra.opfarr)
                #plt.show()

            #Change to the standard star directory
            cwd = os.getcwd()
            os.chdir('/afs/cas.unc.edu/depts/physics_astronomy/clemens/students/group/standards')

            #read in the standard file
            placeholder = cucumber // 2
            stdfile = stdflux[placeholder]
            std_spectra = st.readstandard(stdfile)
            os.chdir(cwd)
            #plt.clf()
            #plt.plot(std_spectra.warr,std_spectra.magarr,'.')
            #plt.show()
            #Only keep the part of the standard file that overlaps with observation.
            lowwv = np.where(std_spectra.warr >= np.min(obs_spectra.warr))
            lowwv = np.asarray(lowwv)
            highwv = np.where(std_spectra.warr <= np.max(obs_spectra.warr))
            highwv = np.asarray(highwv)
            index = np.intersect1d(lowwv,highwv)

            std_spectra.warr = std_spectra.warr[index]
            std_spectra.magarr = std_spectra.magarr[index]
            std_spectra.wbin = std_spectra.wbin[index]

            #Convert from AB mag to fnu, then to fwave (ergs/s/cm2/A)
            stdzp = 3.68e-20 #The absolute flux per unit frequency at an AB mag of zero
            std_spectra.magarr = st.magtoflux(std_spectra.magarr,stdzp)
            std_spectra.magarr = st.fnutofwave(std_spectra.warr, std_spectra.magarr)

            #plt.clf()
            #plt.plot(std_spectra.warr,std_spectra.magarr,'.')
            #plt.show()
            #np.savetxt('hz4_stan.txt',np.transpose([std_spectra.warr,std_spectra.magarr]))
            #exit()

            #We want to rebin the observed spectrum to match with the bins in the standard file. This makes summing up counts significantly easier.
            #Set the new binning here.
            print 'Starting to rebin: ',stdspecfile
            low = np.rint(np.min(obs_spectra.warr)) #Rounds to nearest integer
            high = np.rint(np.max(obs_spectra.warr))
            size = 0.05 #size in Angstroms you want each bin

            num = (high - low) / size + 1. #number of bins. Must add one to get correct number.
            wavenew = np.linspace(low,high,num=num) #wavelength of each new bin

            #Now do the rebinning using Ian Crossfield's rebinning package
            binflux = st.resamplespec(wavenew,obs_spectra.warr,obs_spectra.opfarr,200.) #200 is the oversampling factor
            print 'Done rebinning. Now summing the spectrum into new bins to match', stdfile
            #plt.clf()
            #plt.plot(obs_spectra.warr,obs_spectra.opfarr)
            #plt.plot(wavenew,binflux)
            #plt.show()

            #Now sum the rebinned spectra into the same bins as the standard star file
            counts = st.sum_std(std_spectra.warr,std_spectra.wbin,wavenew,binflux)
            #plt.clf()
            #plt.plot(std_spectra.warr,std_spectra.magarr)
            #plt.plot(obs_spectra.warr,obs_spectra.opfarr,'b')
            #plt.plot(std_spectra.warr,counts,'g+')
            #plt.show()

            #Calculate the sensitivity function
            sens_function = st.sensfunc(counts,std_spectra.magarr,exptime,std_spectra.wbin,airmass)
            #plt.clf()
            #plt.plot(std_spectra.warr,sens_function)
            #plt.show()
            #sys.exit()
            #Fit a low order polynomial to this function so that it is smooth.
            #The sensitivity function is in units of 2.5 * log10[counts/sec/Ang / ergs/cm2/sec/Ang]
            #Choose regions to not include in fit, first by checking if a mask file exists, and if not the prompt for user interaction.
            if 'blue' in stdspecfile.lower():
                std_mask = stdfile[0:-4] + '_blue_maskasdf.dat'
            if 'red' in stdspecfile.lower():
                std_mask = stdfile[0:-4] + '_red_maskasdf.dat'
            std_mask2 = glob(std_mask)
            if len(std_mask2) == 1.:
                print 'Found mask file.\n'
                mask = np.ones(len(std_spectra.warr))
                excluded_wave = np.genfromtxt(std_mask) #Read in wavelengths to exclude
                #print excluded_wave
                #print type(excluded_wave)
                #Find index of each wavelength
                excluded = []
                for x in excluded_wave:
                    #print x
                    #print np.where(std_spectra.warr == find_nearest(std_spectra.warr,x))
                    pix_val = np.where(std_spectra.warr == find_nearest(std_spectra.warr,x))
                    excluded.append(pix_val[0][0])
                #print excluded
                lettuce = 0
                while lettuce < len(excluded):
                    mask[excluded[lettuce]:excluded[lettuce+1]+1] = 0
                    lettuce += 2
                excluded =  np.array(excluded).tolist()
                allexcluded[cucumber] = excluded
                indices = np.where(mask !=0.)
                lambdasfit = std_spectra.warr[indices]
                fluxesfit = sens_function[indices]
            else:
                print 'No mask found. User interaction required.\n'

                global ax, fig, coords
                coords = []
                plt.clf()
                fig = plt.figure(1)
                ax = fig.add_subplot(111)
                ax.plot(std_spectra.warr,sens_function)
                cid = fig.canvas.mpl_connect('button_press_event',onclick)
                print 'Please click on both sides of regions you want to exclude. Then close the plot.'
                plt.title('Click both sides of regions you want to exclude. Then close the plot.')
                plt.show(1)

                #Mask our the regions you don't want to fit
                #We need make sure left to right clicking and right to left clicking both work.
                mask = np.ones(len(std_spectra.warr))
                excluded = np.zeros(len(coords))
                lettuce = 0
                if len(coords) > 0:
                    while lettuce < len(coords):
                        x1 = np.where(std_spectra.warr == (find_nearest(std_spectra.warr,coords[lettuce][0])))
                        excluded[lettuce] = np.asarray(x1)
                        lettuce += 1
                        x2 = np.where(std_spectra.warr == (find_nearest(std_spectra.warr,coords[lettuce][0])))
                        if x2 < x1:
                            x1,x2 = x2,x1
                        mask[x1[0][0]:x2[0][0]+1] = 0 #have to add 1 here to the second index so that we exclude through that index. Most important for when we need to exclude the last point of the array.
                        excluded[lettuce-1] = np.asarray(x1)
                        excluded[lettuce] = np.asarray(x2)
                        lettuce += 1

                excluded =  np.array(excluded).tolist()
                allexcluded[cucumber] = excluded
                indices = np.where(mask !=0.)
                lambdasfit = std_spectra.warr[indices]
                fluxesfit = sens_function[indices]

                #Save masked wavelengths
                lambdasnotfit = std_spectra.warr[excluded]
                #print lambdasnotfit
                #print stdfile
                if 'blue' in stdspecfile.lower():
                    std_mask_name = stdfile[0:-4] + '_blue_mask.dat'
                if 'red' in stdspecfile.lower():
                    std_mask_name = stdfile[0:-4] + '_red_mask.dat'
                np.savetxt(std_mask_name,np.transpose(np.array(lambdasnotfit)))
                #exit()

            ##Move back to directory with observed spectra
            #os.chdir(cwd)

            #Make sure they are finite
            ind1 = np.isfinite(lambdasfit) & np.isfinite(fluxesfit)
            lambdasfit = lambdasfit[ind1]
            fluxesfit = fluxesfit[ind1]

            print 'Fitting the sensitivity funtion now.'
            order = 4
            repeat = 'yes'
            while repeat == 'yes':
                p = np.polyfit(lambdasfit,fluxesfit,order)
                f = np.poly1d(p)
                smooth_sens = f(lambdasfit)
                residual = fluxesfit - smooth_sens
                plt.close()
                plt.ion()
                g, (ax1,ax2) = plt.subplots(2,sharex=True)
                ax1.plot(lambdasfit,fluxesfit,'b+')
                ax1.plot(lambdasfit,smooth_sens,'r',linewidth=2.0)
                ax1.set_ylabel('Sensitivity Function')
                ax2.plot(lambdasfit,residual,'k+')
                ax2.set_ylabel('Residuals')
                ax1.set_title('Current polynomial order: %s' % order)
                g.subplots_adjust(hspace=0)
                plt.setp([a.get_xticklabels() for a in g.axes[:-1]],visible=False)
                plt.show()
                plt.ioff()
                #Save this sensitivity curve
                '''
                try:
                    temp_file = fits.open(stdspecfile)
                    ADCstat = temp_file[0].header['ADCSTAT']
                except:
                    ADCstat = 'none'
                    pass
                if 'blue' in stdspecfile.lower():
                    resp_name = 'senscurve_' + stdfile[1:-4] + '_' + str(np.round(airstd[cucumber],decimals=3))  + '_' + ADCstat  + '_' + cwd[60:70] + '_blue.txt'
                elif 'red' in stdspecfile.lower():
                    resp_name = 'senscurve_' + stdfile[1:-4] + '_' + str(np.round(airstd[cucumber],decimals=3))  + '_' + ADCstat  + '_' + cwd[60:70] + '_red.txt'
                print resp_name
                #exit()
                np.savetxt(resp_name,np.transpose([lambdasfit,fluxesfit]))
                '''
                repeat = raw_input('Do you want to try again (yes/no)? ')
                if repeat == 'yes':
                    order = raw_input('New order for polynomial: ')

            orderused[cucumber] = order
            senspolys.append(f)

            #Save arrays for diagnostic plots
            if cucumber == 0:
                bigarray = np.zeros([len(lambdasfit),4.*len(standards)])
                artichoke = 0
            bigarray[0:len(lambdasfit),artichoke] = lambdasfit
            bigarray[0:len(fluxesfit),artichoke+1] = fluxesfit
            bigarray[0:len(smooth_sens),artichoke+2] = smooth_sens
            bigarray[0:len(residual),artichoke+3] = residual
            artichoke += 4

            cucumber += 1

        #Save fit and residuals into text file for diagnostic plotting later.
        #Need to save lambdasfit,fluxesfit,smooth_sens,residual for each standard
        #List of standards is found as standards
        now = datetime.datetime.now().strftime("%Y-%m-%dT%H:%M")
        with open('sens_fits_' + now + '.txt','a') as handle:
            header = str(standards) + '\n Set of four columns correspond to wavelength, observed flux, polynomial fit, \n and residuals for each standard listed above. \n You will probably need to strip zeros from the bottoms of some columns.'
            np.savetxt(handle,bigarray,fmt='%f',header = header)

    #Outline for next steps:
    #Read in both red and blue files
    #compute airmass and compare to airstd
    #choose best standard and flux calibrate both blue and red
    #save files and write to sensitivity_params.txt

    if speclist[-4:] == 'fits':
        specfile = np.array([speclist])
    else:
        specfile = np.genfromtxt(speclist,dtype=str)
        if specfile.size ==1:
            specfile = np.array([specfile])

    length = len(specfile)
    airwd = np.zeros([length])
    bean = 0
    #if length == 1:
    #    redfile = False
    #else:
    #    redfile = True

    avocado = 0
    while avocado < length:
        #Read in the blue and red spectra we want to flux calibrate. Save the airmass
        WD_spectra1,airmass1,exptime1,dispersion1 = st.readspectrum(specfile[avocado])
        if (len(specfile) >= 1) and (avocado+1 < length):
            if 'red' in specfile[avocado+1]:
                redfile = True
            else:
                redfile = False
        else:
            redfile = False
        if redfile:
            WD_spectra2,airmass2,exptime2,dispersion2 = st.readspectrum(specfile[avocado+1])

        #Extinction correct WD
        if extinct_correct:
            print 'Extinction correcting spectra.'
            #plt.clf()
            #plt.plot(WD_spectra1.warr,WD_spectra1.opfarr)
            WD_spectra1.opfarr = st.extinction_correction(WD_spectra1.warr,WD_spectra1.opfarr,airmass1)
            WD_spectra1.farr = st.extinction_correction(WD_spectra1.warr,WD_spectra1.farr,airmass1)
            WD_spectra1.sky = st.extinction_correction(WD_spectra1.warr,WD_spectra1.sky,airmass1)
            WD_spectra1.sigma = st.extinction_correction(WD_spectra1.warr,WD_spectra1.sigma,airmass1)
            #plt.plot(WD_spectra1.warr,WD_spectra1.opfarr)
            #plt.show()

            if redfile:
                #plt.clf()
                #plt.plot(WD_spectra2.warr,WD_spectra2.opfarr)
                WD_spectra2.opfarr = st.extinction_correction(WD_spectra2.warr,WD_spectra2.opfarr,airmass2)
                WD_spectra2.farr = st.extinction_correction(WD_spectra2.warr,WD_spectra2.farr,airmass2)
                WD_spectra2.sky = st.extinction_correction(WD_spectra2.warr,WD_spectra2.sky,airmass2)
                WD_spectra2.sigma = st.extinction_correction(WD_spectra2.warr,WD_spectra2.sigma,airmass2)

                #zaplt.plot(WD_spectra2.warr,WD_spectra2.opfarr)
                #plt.show()
        airwd[avocado] = airmass1
        if redfile:
            airwd[avocado+1] = airmass2
        #Compare the airmasses to determine the best standard star
        tomato = 0
        while tomato < len(airstd):
            if redfile:
                diff = np.absolute(np.mean([airwd[avocado],airwd[avocado+1]]) - np.mean([airstd[tomato],airstd[tomato+1]]))
            else:
                diff = np.absolute(airwd[avocado] - airstd[tomato])
            if tomato == 0:
                difference = diff
                choice = tomato
            if diff < difference:
                difference = diff
                choice = tomato
            tomato += 2

        #To get the flux calibration, perform the following
        #Flux = counts / (Exptime * dispersion * 10**(sens/2.5))
        #Get the sensitivity function at the correct wavelength spacing
        if masterresp:
            header_temp = st.readheader(specfile[avocado])
            ADCstatus = header_temp['ADCSTAT']
            if ADCstatus == 'IN':
                sens_wave1_unscale = master_response_blue_in_pol(WD_spectra1.warr)
                blue_low_index = np.min(np.where(WD_spectra1.warr > 4530.))
                blue_high_index = np.min(np.where(WD_spectra1.warr > 4590.))
                blue_mean_stan = np.mean(sens_wave1_unscale[blue_low_index:blue_high_index])
                if blue_mean_tonight == None:
                    sens_wave1 = sens_wave1_unscale
                else:
                    sens_wave1 = sens_wave1_unscale + (blue_mean_tonight - blue_mean_stan)
                choice = 0
            else:
                sens_wave1_unscale = master_response_blue_out_pol(WD_spectra1.warr)
                blue_low_index = np.min(np.where(WD_spectra1.warr > 4530.))
                blue_high_index = np.min(np.where(WD_spectra1.warr > 4590.))
                blue_mean_stan = np.mean(sens_wave1_unscale[blue_low_index:blue_high_index])
                if blue_mean_tonight == None:
                    sens_wave1 = sens_wave1_unscale
                else:
                    sens_wave1 = sens_wave1_unscale + (blue_mean_tonight - blue_mean_stan)
                choice = 1
            if redfile:
                header_temp = st.readheader(specfile[avocado+1])
                ADCstatus = header_temp['ADCSTAT']
                if ADCstatus == 'IN':
                    sens_wave2_unscale = master_response_red_in_pol(WD_spectra2.warr)
                    red_low_index = np.min(np.where(WD_spectra2.warr > 6090.))
                    red_high_index = np.min(np.where(WD_spectra2.warr > 6190.))
                    red_mean_stan = np.mean(sens_wave2_unscale[red_low_index:red_high_index])
                    if red_mean_tonight == None:
                        sens_wave2 = sens_wave2_unscale
                    else:
                        sens_wave2 = sens_wave2_unscale + (red_mean_tonight - red_mean_stan)
                    choice2 = 2
                else:
                    sens_wave2_unscale = master_response_red_out_pol(WD_spectra2.warr)
                    red_low_index = np.min(np.where(WD_spectra2.warr > 6090.))
                    red_high_index = np.min(np.where(WD_spectra2.warr > 6190.))
                    red_mean_stan = np.mean(sens_wave2_unscale[red_low_index:red_high_index])
                    if red_mean_tonight == None:
                        sens_wave2 = sens_wave2_unscale
                    else:
                        sens_wave2 = sens_wave2_unscale + (red_mean_tonight - red_mean_stan)
                    choice2 = 3
        else:
            sens_wave1 = senspolys[choice](WD_spectra1.warr)
            if redfile:
                sens_wave2 = senspolys[choice+1](WD_spectra2.warr)

        #Perform the flux calibration. We do this on the optimal extraction, non-variance weighted aperture, the sky spectrum, and the sigma spectrum.
        print 'Doing the final flux calibration.'
        #np.savetxt('response_g60-54_extinction_2016-03-17.txt',np.transpose([WD_spectra1.warr,(exptime1 * dispersion1 * 10.**(sens_wave1/2.5))]))#,WD_spectra2.warr,(exptime2 * dispersion2 * 10.**(sens_wave2/2.5))]))
        #exit()
        star_opflux1 = st.cal_spec(WD_spectra1.opfarr,sens_wave1,exptime1,dispersion1)
        star_flux1 = st.cal_spec(WD_spectra1.farr,sens_wave1,exptime1,dispersion1)
        sky_flux1 = st.cal_spec(WD_spectra1.sky,sens_wave1,exptime1,dispersion1)
        sigma_flux1 = st.cal_spec(WD_spectra1.sigma,sens_wave1,exptime1,dispersion1)

        if redfile:
            star_opflux2 = st.cal_spec(WD_spectra2.opfarr,sens_wave2,exptime2,dispersion2)
            star_flux2 = st.cal_spec(WD_spectra2.farr,sens_wave2,exptime2,dispersion2)
            sky_flux2 = st.cal_spec(WD_spectra2.sky,sens_wave2,exptime2,dispersion2)
            sigma_flux2 = st.cal_spec(WD_spectra2.sigma,sens_wave2,exptime2,dispersion2)

        #plt.clf()
        #plt.plot(WD_spectra.warr,star_opflux)
        #plt.show()

        #Save final spectra if using master response
        if masterresp:
            if avocado == 0:
                diagnostic_array = np.zeros([len(WD_spectra1.warr),2*length])
            diagnostic_array[0:len(WD_spectra1.warr),bean] = WD_spectra1.warr
            bean += 1
            diagnostic_array[0:len(star_opflux1),bean] = star_opflux1
            bean += 1
            if redfile:
                diagnostic_array[0:len(WD_spectra2.warr),bean] = WD_spectra2.warr
                bean += 1
                diagnostic_array[0:len(star_opflux2),bean] = star_opflux2
                bean += 1
        #if avocado == (length -1 ) or (redfile == True and avocado == (length-2)):
        #    print 'Saveing diagnostic file.'
        #    now = datetime.datetime.now().strftime("%Y-%m-%dT%H:%M")
        #    with open('flux_fits_' + now + '.txt','a') as handle:
        #        header = str(specfile) + '\n Each star is formatted as wavelength, flux'
        #        np.savetxt(handle,diagnostic_array,fmt='%.10e',header=header)

        print 'Saving the final spectrum.'

        #Save the flux-calibrated spectrum and update the header
        header1 = st.readheader(specfile[avocado])
        header1.set('EX-FLAG',extinctflag) #Extiction correction? 0=yes, -1=no
        header1.set('CA-FLAG',0) #Calibrated to flux scale? 0=yes, -1=no
        header1.set('BUNIT','erg/cm2/s/A') #physical units of the array value
        header1.set('STANDARD',str(standards[choice]),'Flux standard used') #flux standard used for flux-calibration
        if masterresp:
            header1.set('STDOFF',str(flux_tonight_list[0]),'Night offset used')

        if redfile:
            header2 = st.readheader(specfile[avocado+1])
            header2.set('EX-FLAG',extinctflag) #Extiction correction? 0=yes, -1=no
            header2.set('CA-FLAG',0) #Calibrated to flux scale? 0=yes, -1=no
            header2.set('BUNIT','erg/cm2/s/A') #physical units of the array value
            if masterresp:
                header2.set('STANDARD',str(standards[choice2]),'Flux standard used') #flux standard used for flux-calibration
                header1.set('STDOFF',str(flux_tonight_list[1]),'Night offset used')
            else:
                header2.set('STANDARD',str(standards[choice+1]),'Flux standard used') #flux standard used for flux-calibration

        #Set up size of new fits image
        Ni = 4. #Number of extensions
        Nx1 = len(star_flux1)
        if redfile:
            Nx2 = len(star_flux2)
        Ny = 1. #All 1D spectra

        data1 = np.empty(shape = (Ni,Ny,Nx1))
        data1[0,:,:] = star_opflux1
        data1[1,:,:] = star_flux1
        data1[2,:,:] = sky_flux1
        data1[3,:,:] = sigma_flux1

        if redfile:
            data2 = np.empty(shape = (Ni,Ny,Nx2))
            data2[0,:,:] = star_opflux2
            data2[1,:,:] = star_flux2
            data2[2,:,:] = sky_flux2
            data2[3,:,:] = sigma_flux2

        #Add '_flux' to the end of the filename
        loc1 = specfile[avocado].find('.ms.fits')
        if masterresp:
            newname1 = specfile[avocado][0:loc1] + '_flux_' + stdflux[0][1:-4]  + '.ms.fits'
        else:
            newname1 = specfile[avocado][0:loc1] + '_flux_' + stdflux[choice//2][1:-4]  + '.ms.fits'
        clob = False
        mylist = [True for f in os.listdir('.') if f == newname1]
        exists = bool(mylist)

        if exists:
            print 'File %s already exists.' % newname1
            nextstep = raw_input('Do you want to overwrite or designate a new name (overwrite/new)? ')
            if nextstep == 'overwrite':
                clob = True
                exists = False
            elif nextstep == 'new':
                newname1 = raw_input('New file name: ')
                exists = False
            else:
                exists = False
        print 'Saving: ', newname1
        newim1 = fits.PrimaryHDU(data=data1,header=header1)
        newim1.writeto(newname1,clobber=clob)

        if redfile:
            loc2 = specfile[avocado+1].find('.ms.fits')
            if masterresp:
                newname2 = specfile[avocado+1][0:loc2] + '_flux_' + stdflux[0][1:-4] + '.ms.fits'
            else:
                newname2 = specfile[avocado+1][0:loc2] + '_flux_' + stdflux[choice//2][1:-4] + '.ms.fits'
            clob = False
            mylist = [True for f in os.listdir('.') if f == newname2]
            exists = bool(mylist)

            if exists:
                print 'File %s already exists.' % newname2
                nextstep = raw_input('Do you want to overwrite or designate a new name (overwrite/new)? ')
                if nextstep == 'overwrite':
                    clob = True
                    exists = False
                elif nextstep == 'new':
                    newname2 = raw_input('New file name: ')
                    exists = False
                else:
                    exists = False

            newim2 = fits.PrimaryHDU(data=data2,header=header2)
            newim2.writeto(newname2,clobber=clob)
            print 'Saving: ', newname2

        #Finally, save all the used parameters into a file for future reference.
        # specfile,current date, stdspecfile,stdfile,order,size,newname
        f = open('sensitivity_params.txt','a')
        now = datetime.datetime.now().strftime("%Y-%m-%dT%H:%M")
        if masterresp:
            newinfo1 = specfile[avocado] + '\t' + now + '\t' + standards[choice] + '\t' + stdflux[0] + '\t' + str(allexcluded[choice]) + '\t' + str(orderused[choice]) + '\t' + str(size) + '\t' + newname1
        else:
            newinfo1 = specfile[avocado] + '\t' + now + '\t' + standards[choice] + '\t' + stdflux[choice//2] + '\t' + str(allexcluded[choice]) + '\t' + str(orderused[choice]) + '\t' + str(size) + '\t' + newname1
        if redfile:
            if masterresp:
                newinfo2 = specfile[avocado+1] + '\t' + now + '\t' + standards[choice2] + '\t' + stdflux[0] + '\t' + str(allexcluded[choice+1]) + '\t' + str(orderused[choice+1]) + '\t' + str(size) + '\t' + newname2
            else:
                newinfo2 = specfile[avocado+1] + '\t' + now + '\t' + standards[choice+1] + '\t' + stdflux[choice//2] + '\t' + str(allexcluded[choice+1]) + '\t' + str(orderused[choice+1]) + '\t' + str(size) + '\t' + newname2
            f.write(newinfo1 + "\n" + newinfo2 + "\n")
        else:
            f.write(newinfo1 + "\n")
        f.close()

        if redfile:
            avocado += 2
        else:
            avocado += 1

    print 'Done flux calibrating the spectra.'

#Run from command line
if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('spec_list')
    parser.add_argument('--flux_list',default=None)
    parser.add_argument('--stan_list',default=None)
    parser.add_argument('--usemaster',type=str2bool,nargs='?',const=False,default=False,help='Activate nice mode.')
    parser.add_argument('--extinct',type=str2bool,nargs='?',const=True,default=True,help='Activate nice mode.')
    args = parser.parse_args()
    #print args.stand_list
    flux_calibrate_now(args.stan_list,args.flux_list,args.spec_list,extinct_correct=args.extinct,masterresp=args.usemaster)


import re
# import sys
import os
import pandas as pd
import numpy as np
from scipy import signal
from scipy import ndimage
from sklearn import decomposition
import matplotlib.pyplot as plt
plt.rcParams["font.family"] = "IPAexGothic"

class EMG:
    """
    筋電位を扱うクラス

    Attributes
    ----------
    data : pandas.DataFrame
        筋電位データ

    fs : float
        サンプリング周波数

    H : pandas.DataFrame
        筋シナジーの重み

    W : pandas.DataFrame
        筋シナジーの時間変化

    muscle_colors : str | list = 'tab:blue'
        筋肉に対応する色

    begin_time : float
        対象範囲の開始時間

    end_time : float
        対象範囲の終了時間

    Examples
    --------
    >>> from measurexp.EMG import EMG
    >>> emg = EMG()
    >>> emg.read('EMG.csv')
    >>> emg.set_time(0, 30)
    >>> emg.prep()
    >>> emg.calc_synergy()
    >>> emg.plot_synergy()
    >>> emg.set_colors('../muscle_colors.csv')
    >>> emg.plot_synergy_weights()
    """

    def __init__(self):
        # 筋電位データ
        self.data: pd.DataFrame = None
        # フィルター後のデータ
        self.rms: pd.DataFrame = None
        # 正規化後のデータ
        self.norm: pd.DataFrame = None
        # サンプリング周波数
        self.fs: float = None
        self.H: pd.DataFrame = None
        self.W: pd.DataFrame = None
        self.begin_time: float = None
        self.end_time: float = None
        self.muscles_color: str | list = 'tab:blue'
        self.taskname: str = 'Unnamed'

    def set_colors(self, colors) -> 'EMG':
        """筋肉に対応する色を設定します。

        Parameters
        ----------
        colors : str | list
            筋肉に対応する色のリストまたは文字列

        Returns
        -------
        self : EMG
        """
        if type(colors) in (str, list):
            if os.path.isfile(colors):
                self.muscles_color = pd.read_csv(
                    colors, header=None).values.reshape(-1).tolist()
                return self
            self.muscles_color = colors
        else:
            print('筋肉に対応する色を指定してください。')
        return self

    def _col2muscles_name(self, col_muscle: str):
        match = re.match(r'^(.+):', col_muscle)
        if match is not None:
            muscles_name = match.group(1)
        return muscles_name

    def read(self, filename: str) -> 'EMG':
        """データファイル (*.csv) を読み込みます。

        Parameters
        ----------
        filename : str
            データファイル (*.csv)

        Returns
        -------
        self : EMG
        """
        try:
            # self.data = pd.read_csv(filename,
            # encoding='Shift-JIS', header=116)
            self.data = pd.read_csv(filename)
        except UnicodeDecodeError as err:
            print(err)
            print('デコードできません。正しい文字コードを指定してください。')
        except FileNotFoundError as err:
            print(err)
            print('筋電位データのファイルを指定してください。')

        # インデックスの設定
        self.data.set_index(['Time [s]'], inplace=True)
        self.begin_time, self.end_time = \
            self.data.index[0], self.data.index[-1]
        self.fs = self.data.shape[0] / (self.end_time - self.begin_time)
        return self

    def name(self, taskname: str) -> 'EMG':
        """タスク名を設定します。

        Parameters
        ----------
        taskname : str
            タスク名

        Return
        ------
        self : EMG
        """
        self.taskname = taskname
        return self

    def prep(self, period: float = 0.2, n: int = 5, Fc: np.ndarray = np.array([5, 500])) -> 'EMG':
        """データの下処理を行います。

        Parameters
        ----------
        period : float = 0.1
            RMS するウィンドウの範囲 (秒)

        n : int = 5
            ローパスフィルターの次数

        Fc : int = 50
            ローパスフィルターの遮断周波数

        Returns
        -------
        self : EMG
        """
        # b, a = signal.butter(n, Fc / self.fs * 2, 'band', analog=True)
        x = self.data.loc[self.begin_time:self.end_time, :].copy()
        # for m in range(x.shape[1]):
            # x.iloc[:, m] = signal.filtfilt(b, a, x.iloc[:, m])

        x[:] = np.power(x - x.mean(), 2)
        for _ in x:
            x[_][:] = ndimage.gaussian_filter1d(x[_], self.fs * period)
        self.rms = np.sqrt(x)
        return self

    def resample(self, fs: int) -> 'EMG':
        df = self.rms.copy()
        self.rms = pd.DataFrame(
            signal.resample(df.to_numpy(), self.end_time * fs),
            index=(np.arange(self.end_time * fs) / fs),
            columns=df.columns
        )
        df = self.data.copy()
        self.data = pd.DataFrame(
            signal.resample(df.to_numpy(), self.end_time * fs),
            index=(np.arange(self.end_time * fs) / fs),
            columns=df.columns
        )
        self.fs = fs

        return self

    def set_muscles(self, muscles) -> 'EMG':
        self.rms = self.rms.loc[:, muscles]
        self.data = self.data.loc[:, muscles]
        return self

    def _vaf(self, X: np.ndarray, W: np.ndarray, H: np.ndarray):
        VAF = 1 - np.power(X - W.dot(H), 2).sum() / np.power(X, 2).sum()
        return VAF

    def _calc_synergy(self, X: np.ndarray, max_vaf: float):
        VAFs = []
        for n_components in range(1, X.shape[1]):
            model = decomposition.NMF(
                n_components=n_components,
                init='nndsvda',
                random_state=0,
                tol=1e-1
            )
            W = model.fit_transform(X)
            H = model.components_
            VAFs.append(self._vaf(X, W, H))
            if VAFs[-1] > max_vaf:
                break

        return (VAFs, W, H)

    def set_time(self, begin_time: float, end_time: float) -> 'EMG':
        """処理対象時間を設定します。

        Parameters
        ----------
        begin_time : float
            処理範囲の開始時間

        end_time : float
            処理範囲の終了時間

        Returns
        -------
        self : EMG
        """
        self.begin_time, self.end_time = begin_time, end_time
        return self

    def calc_synergy(self, max_vaf: float = 0.9, norm: bool = True) -> 'EMG':
        """筋シナジーを計算します。

        Parameters
        ----------
        max_vaf : float = 0.9
            最大 VAF

        norm : bool = True
            正規化した筋シナジーで計算する

        Returns
        -------
        self : EMG
        """
        if self.norm is None:
            print('正規化されたデータが存在しません。正規化する必要があります。')
            return EMG

        X = self.norm.values if norm else self.rms.values

        self.VAFs, W, H = self._calc_synergy(X, max_vaf)
        self.W = pd.DataFrame(
            W,
            columns=[f'筋シナジー {_+1}' for _ in range(W.shape[1])],
            index=self.rms.index
        )
        self.H = pd.DataFrame(
            H,
            columns=self.rms.columns,
            index=[f'筋シナジー {_+1}' for _ in range(W.shape[1])]
        )
        self.n_synergy = self.H.shape[0]
        return self

    def plot_synergy(self, **kwargs) -> 'EMG':
        """筋シナジーの時間変化を表示します。

        Parameters
        ----------
        kwargs : Any

        Returns
        -------
        self : EMG
        """
        fig, ax = plt.subplots()

        plt_kwargs = {'ax': ax}
        if kwargs is not None:
            plt_kwargs.update(kwargs)
        self.W.plot(**plt_kwargs)
        ax.set_xlim(self.begin_time, self.end_time)
        ax.yaxis.set_major_formatter(plt.ScalarFormatter(useMathText=True))
        ax.ticklabel_format(style='sci', axis='y', scilimits=(0, 0))
        plt.show()
        return self

    def plot_synergy_weights(self) -> 'EMG':
        """筋シナジーを表示します。

        Returns
        -------
        self : EMG
        """

        fig, ax = plt.subplots(
            self.n_synergy, figsize=(6.4, 1.6 * self.n_synergy))
        if self.n_synergy == 1:
            ax = [ax]
        for _ in range(self.n_synergy):
            self.H.iloc[_, :].plot.bar(ax=ax[_], color=self.muscles_color)
            if _ == self.n_synergy - 1:
                break
            ax[_].axes.xaxis.set_visible(False)
        fig.tight_layout()
        plt.show()
        return self


"""
This file is part of the OpenProtein project.

For license information, please see the LICENSE file in the root directory.
"""

import math
import os
from datetime import datetime
import torch
import torch.utils.data
import h5py
import PeptideBuilder
import Bio.PDB
import numpy as np
import pnerf.pnerf as pnerf

AA_ID_DICT = {'A': 1, 'C': 2, 'D': 3, 'E': 4, 'F': 5, 'G': 6, 'H': 7, 'I': 8, 'K': 9,
              'L': 10, 'M': 11, 'N': 12, 'P': 13, 'Q': 14, 'R': 15, 'S': 16, 'T': 17,
              'V': 18, 'W': 19, 'Y': 20}

def contruct_dataloader_from_disk(filename, minibatch_size):
    return torch.utils.data.DataLoader(H5PytorchDataset(filename),
                                       batch_size=minibatch_size,
                                       shuffle=True,
                                       collate_fn=merge_samples_to_minibatch)

class H5PytorchDataset(torch.utils.data.Dataset):
    def __init__(self, filename):
        super(H5PytorchDataset, self).__init__()

        self.h5pyfile = h5py.File(filename, 'r')
        self.num_proteins, self.max_sequence_len = self.h5pyfile['primary'].shape

    def __getitem__(self, index):
        mask = torch.Tensor(self.h5pyfile['mask'][index, :]).type(dtype=torch.bool)
        prim = torch.masked_select(
            torch.Tensor(self.h5pyfile['primary'][index, :]).type(dtype=torch.long),
            mask)
        tertiary = torch.Tensor(self.h5pyfile['tertiary'][index][:int(mask.sum())])# max length x 9
        return prim, tertiary, mask

    def __len__(self):
        return self.num_proteins

def merge_samples_to_minibatch(samples):
    samples_list = []
    for sample in samples:
        samples_list.append(sample)
    # sort according to length of aa sequence
    samples_list.sort(key=lambda x: len(x[0]), reverse=True)
    return zip(*samples_list)

def set_experiment_id(data_set_identifier, learning_rate, minibatch_size):
    output_string = datetime.now().strftime('%Y-%m-%d_%H_%M_%S')
    output_string += "-" + str(os.getpid())
    output_string += "-" + data_set_identifier
    output_string += "-LR" + str(learning_rate).replace(".", "_")
    output_string += "-MB" + str(minibatch_size)
    globals().__setitem__("experiment_id", output_string)

def get_experiment_id():
    return globals().get("experiment_id")

def write_out(*args, end='\n'):
    output_string = datetime.now().strftime('%Y-%m-%d %H:%M:%S') \
                    + ": " + str.join(" ", [str(a) for a in args]) + end
    if globals().get("experiment_id") is not None:
        with open("output/" + globals().get("experiment_id") + ".txt", "a+") as output_file:
            output_file.write(output_string)
            output_file.flush()
    print(output_string, end="")

def write_model_to_disk(model):
    path = "output/models/" + globals().get("experiment_id") + ".model"
    torch.save(model, path)
    return path

def write_prediction_data_to_disk(prediction_data):
    filepath = "output/predictions/" + globals().get("experiment_id") + ".txt"
    output_file = open(filepath, 'w')
    output_file.write(prediction_data)
    output_file.close()

def draw_plot(fig, plt, validation_dataset_size, sample_num, train_loss_values,
              validation_loss_values):
    def draw_with_vars():
        ax = fig.gca()
        ax2 = ax.twinx()
        plt.grid(True)
        plt.title("Training progress (" + str(validation_dataset_size)
                  + " samples in validation set)")
        train_loss_plot, = ax.plot(sample_num, train_loss_values)
        ax.set_ylabel('Train Negative log likelihood')
        ax.yaxis.labelpad = 0
        validation_loss_plot, = ax2.plot(sample_num, validation_loss_values, color='black')
        ax2.set_ylabel('Validation loss')
        ax2.set_ylim(bottom=0)
        plt.legend([train_loss_plot, validation_loss_plot],
                   ['Train loss on last batch', 'Validation loss'])
        ax.set_xlabel('Minibatches processed (=network updates)', color='black')

    return draw_with_vars

def draw_ramachandran_plot(fig, plt, phi, psi):
    def draw_with_vars():
        ax = fig.gca()
        plt.grid(True)
        plt.title("Ramachandran plot")
        train_loss_plot, = ax.plot(phi, psi)
        ax.set_ylabel('Psi')
        ax.yaxis.labelpad = 0
        plt.legend([train_loss_plot],
                   ['Phi psi'])
        ax.set_xlabel('Phi', color='black')

    return draw_with_vars

def write_result_summary(accuracy):
    output_string = globals().get("experiment_id") + ": " + str(accuracy) + "\n"
    with open("output/result_summary.txt", "a+") as output_file:
        output_file.write(output_string)
        output_file.flush()
    print(output_string, end="")

def calculate_dihedral_angles_over_minibatch(atomic_coords_padded, batch_sizes, use_gpu):
    angles = []
    atomic_coords = atomic_coords_padded.transpose(0, 1)
    for idx, _ in enumerate(batch_sizes):
        angles.append(calculate_dihedral_angles(atomic_coords[idx][:batch_sizes[idx]], use_gpu))
    return torch.nn.utils.rnn.pad_packed_sequence(
        torch.nn.utils.rnn.pack_sequence(angles))

def protein_id_to_str(protein_id_list):
    _aa_dict_inverse = {v: k for k, v in AA_ID_DICT.items()}
    aa_list = []
    for protein_id in protein_id_list:
        aa_symbol = _aa_dict_inverse[int(protein_id)]
        aa_list.append(aa_symbol)
    return aa_list

def calculate_dihedral_angles(atomic_coords, use_gpu):
    assert int(atomic_coords.shape[1]) == 9
    atomic_coords = atomic_coords.contiguous().view(-1, 3)

    zero_tensor = torch.tensor(0.0)
    if use_gpu:
        zero_tensor = zero_tensor.cuda()

    dihedral_list = [zero_tensor, zero_tensor]
    dihedral_list.extend(compute_dihedral_list(atomic_coords))
    dihedral_list.append(zero_tensor)
    angles = torch.tensor(dihedral_list).view(-1, 3)
    return angles

def compute_dihedral_list(atomic_coords):
    # atomic_coords is -1 x 3
    ba = atomic_coords[1:] - atomic_coords[:-1]
    ba /= ba.norm(dim=1).unsqueeze(1)
    ba_neg = -1 * ba

    n1_vec = torch.cross(ba[:-2], ba_neg[1:-1], dim=1)
    n2_vec = torch.cross(ba_neg[1:-1], ba[2:], dim=1)
    n1_vec /= n1_vec.norm(dim=1).unsqueeze(1)
    n2_vec /= n2_vec.norm(dim=1).unsqueeze(1)

    m1_vec = torch.cross(n1_vec, ba_neg[1:-1], dim=1)

    x_value = torch.sum(n1_vec * n2_vec, dim=1)
    y_value = torch.sum(m1_vec * n2_vec, dim=1)

    return torch.atan2(y_value, x_value)

def get_structure_from_angles(aa_list_encoded, angles):
    aa_list = protein_id_to_str(aa_list_encoded)
    omega_list = angles[1:, 0]
    phi_list = angles[1:, 1]
    psi_list = angles[:-1, 2]
    assert len(aa_list) == len(phi_list) + 1 == len(psi_list) + 1 == len(omega_list) + 1
    structure = PeptideBuilder.make_structure(aa_list,
                                              list(map(lambda x: math.degrees(x), phi_list)),
                                              list(map(lambda x: math.degrees(x), psi_list)),
                                              list(map(lambda x: math.degrees(x), omega_list)))
    return structure

def write_to_pdb(structure, prot_id):
    out = Bio.PDB.PDBIO()
    out.set_structure(structure)
    out.save("output/protein_" + str(prot_id) + ".pdb")

def calc_pairwise_distances(chain_a, chain_b, use_gpu):
    distance_matrix = torch.Tensor(chain_a.size()[0], chain_b.size()[0]).type(torch.float)
    # add small epsilon to avoid boundary issues
    epsilon = 10 ** (-4) * torch.ones(chain_a.size(0), chain_b.size(0))
    if use_gpu:
        distance_matrix = distance_matrix.cuda()
        epsilon = epsilon.cuda()

    for idx, row in enumerate(chain_a.split(1)):
        distance_matrix[idx] = torch.sum((row.expand_as(chain_b) - chain_b) ** 2, 1).view(1, -1)

    return torch.sqrt(distance_matrix + epsilon)

def calc_drmsd(chain_a, chain_b, use_gpu=False):
    assert len(chain_a) == len(chain_b)
    distance_matrix_a = calc_pairwise_distances(chain_a, chain_a, use_gpu)
    distance_matrix_b = calc_pairwise_distances(chain_b, chain_b, use_gpu)
    return torch.norm(distance_matrix_a - distance_matrix_b, 2) \
           / math.sqrt((len(chain_a) * (len(chain_a) - 1)))

# method for translating a point cloud to its center of mass
def transpose_atoms_to_center_of_mass(atoms_matrix):
    # calculate com by summing x, y and z respectively
    # and dividing by the number of points
    center_of_mass = np.matrix([[atoms_matrix[0, :].sum() / atoms_matrix.shape[1]],
                                [atoms_matrix[1, :].sum() / atoms_matrix.shape[1]],
                                [atoms_matrix[2, :].sum() / atoms_matrix.shape[1]]])
    # translate points to com and return
    return atoms_matrix - center_of_mass

def calc_rmsd(chain_a, chain_b):
    # move to center of mass
    chain_a_value = chain_a.cpu().numpy().transpose()
    chain_b_value = chain_b.cpu().numpy().transpose()
    X = transpose_atoms_to_center_of_mass(chain_a_value)
    Y = transpose_atoms_to_center_of_mass(chain_b_value)

    R = Y * X.transpose()
    # extract the singular values
    _, S, _ = np.linalg.svd(R)
    # compute RMSD using the formular
    E0 = sum(list(np.linalg.norm(x) ** 2 for x in X.transpose())
             + list(np.linalg.norm(x) ** 2 for x in Y.transpose()))
    TraceS = sum(S)
    RMSD = np.sqrt((1 / len(X.transpose())) * (E0 - 2 * TraceS))
    return RMSD

def calc_angular_difference(values_1, values_2):
    values_1 = values_1.transpose(0, 1).contiguous()
    values_2 = values_2.transpose(0, 1).contiguous()
    acc = 0
    for idx, _ in enumerate(values_1):
        assert values_1[idx].shape[1] == 3
        assert values_2[idx].shape[1] == 3
        a1_element = values_1[idx].view(-1, 1)
        a2_element = values_2[idx].view(-1, 1)
        acc += torch.sqrt(torch.mean(
            torch.min(torch.abs(a2_element - a1_element),
                      2 * math.pi - torch.abs(a2_element - a1_element)
                      ) ** 2))
    return acc / values_1.shape[0]

def structures_to_backbone_atoms_padded(structures):
    backbone_atoms_list = []
    for structure in structures:
        backbone_atoms_list.append(structure_to_backbone_atoms(structure))
    backbone_atoms_padded, batch_sizes_backbone = torch.nn.utils.rnn.pad_packed_sequence(
        torch.nn.utils.rnn.pack_sequence(backbone_atoms_list))
    return backbone_atoms_padded, batch_sizes_backbone

def structure_to_backbone_atoms(structure):
    predicted_coords = []
    for res in structure.get_residues():
        predicted_coords.append(torch.Tensor(res["N"].get_coord()))
        predicted_coords.append(torch.Tensor(res["CA"].get_coord()))
        predicted_coords.append(torch.Tensor(res["C"].get_coord()))
    return torch.stack(predicted_coords).view(-1, 9)

def get_backbone_positions_from_angular_prediction(angular_emissions, batch_sizes, use_gpu):
    # angular_emissions -1 x minibatch size x 3 (omega, phi, psi)
    points = pnerf.dihedral_to_point(angular_emissions, use_gpu)
    coordinates = pnerf.point_to_coordinate(points, use_gpu) / 100  # devide by 100 to angstrom unit
    return coordinates.transpose(0, 1).contiguous()\
               .view(len(batch_sizes), -1, 9).transpose(0, 1), batch_sizes

def calc_avg_drmsd_over_minibatch(backbone_atoms_padded, actual_coords_padded, batch_sizes):
    backbone_atoms_list = list(
        [backbone_atoms_padded[:batch_sizes[i], i] for i in range(int(backbone_atoms_padded
                                                                      .size(1)))])
    actual_coords_list = list(
        [actual_coords_padded[:batch_sizes[i], i] for i in range(int(actual_coords_padded
                                                                     .size(1)))])
    drmsd_avg = 0
    for idx, backbone_atoms in enumerate(backbone_atoms_list):
        actual_coords = actual_coords_list[idx].transpose(0, 1).contiguous().view(-1, 3)
        drmsd_avg += calc_drmsd(backbone_atoms.transpose(0, 1).contiguous().view(-1, 3),
                                actual_coords)
    return drmsd_avg / len(backbone_atoms_list)

def encode_primary_string(primary):
    return list([AA_ID_DICT[aa] for aa in primary])

def initial_pos_from_aa_string(batch_aa_string):
    structures = []
    for aa_string in batch_aa_string:
        structure = get_structure_from_angles(aa_string,
                                              np.repeat([-120], len(aa_string) - 1),
                                              np.repeat([140], len(aa_string) - 1),
                                              np.repeat([-370], len(aa_string) - 1))
        structures.append(structure)
    return structures

def pass_messages(aa_features, message_transformation, use_gpu):
    # aa_features (#aa, #features) - each row represents the amino acid type
    # (embedding) and the positions of the backbone atoms
    # message_transformation: (-1 * 2 * feature_size) -> (-1 * output message size)
    feature_size = aa_features.size(1)
    aa_count = aa_features.size(0)
    eye = torch.eye(aa_count, dtype=torch.uint8).view(-1).expand(2, feature_size, -1)\
        .transpose(1, 2).transpose(0, 1)
    eye_inverted = torch.ones(eye.size(), dtype=torch.uint8) - eye
    if use_gpu:
        eye_inverted = eye_inverted.cuda()
    features_repeated = aa_features.repeat((aa_count, 1)).view((aa_count, aa_count, feature_size))
    # (aa_count^2 - aa_count) x 2 x aa_features     (all pairs except for reflexive connections)
    aa_messages = torch.stack((features_repeated.transpose(0, 1), features_repeated))\
        .transpose(0, 1).transpose(1, 2).view(-1, 2, feature_size)
    aa_msg_pairs = torch.masked_select(aa_messages, eye_inverted).view(-1, 2, feature_size)
    transformed = message_transformation(aa_msg_pairs).view(aa_count, aa_count - 1, -1)
    transformed_sum = transformed.sum(dim=1)  # aa_count x output message size
    return transformed_sum

def load_model_from_disk(path, force_cpu=True):
    if force_cpu:
        # load model with map_location set to storage (main mem)
        model = torch.load(path, map_location=lambda storage, loc: storage)
        # flattern parameters in memory
        model.flatten_parameters()
        # update internal state accordingly
        model.use_gpu = False
    else:
        # load model using default map_location
        model = torch.load(path)
        model.flatten_parameters()
    return model