"""Contains classes for transformer architecture within CrabNet."""
from os.path import join, dirname

import numpy as np
import pandas as pd

import torch
from torch import nn
from collections import OrderedDict

RNG_SEED = 42
torch.manual_seed(RNG_SEED)
np.random.seed(RNG_SEED)
data_type_torch = torch.float32

# %%
class ResidualNetwork(nn.Module):
    """
    Feed forward Residual Neural Network as seen in Roost.

    https://doi.org/10.1038/s41467-020-19964-7
    """

    def __init__(self, input_dim, output_dim, hidden_layer_dims, bias=False):
        """Instantiate a ResidualNetwork model.

        Parameters
        ----------
        input_dim : int
            Input dimensions for the Residual Network, specified in SubCrab() model class, by default 512
        output_dim : int
            Output dimensions for Residual Network, by default 3
        hidden_layer_dims : list(int)
            Hidden layer architecture for the Residual Network, by default [1024, 512, 256, 128]
        bias : bool
            Whether to bias the linear network, by default False
        """
        super(ResidualNetwork, self).__init__()
        dims = [input_dim] + hidden_layer_dims
        self.fcs = nn.ModuleList(
            [nn.Linear(dims[i], dims[i + 1]) for i in range(len(dims) - 1)]
        )
        self.res_fcs = nn.ModuleList(
            [
                nn.Linear(dims[i], dims[i + 1], bias=bias)
                if (dims[i] != dims[i + 1])
                else nn.Identity()
                for i in range(len(dims) - 1)
            ]
        )
        self.acts = nn.ModuleList([nn.LeakyReLU() for _ in range(len(dims) - 1)])
        self.fc_out = nn.Linear(dims[-1], output_dim)

    def forward(self, fea):
        """Propagate Residual Network weights forward.

        Parameters
        ----------
        fea : torch.tensor (n_dim)
            Tensor output of self attention block
        Returns
        -------
        fc_out
            The output of the Residual Network
        """
        for fc, res_fc, act in zip(self.fcs, self.res_fcs, self.acts):
            fea = act(fc(fea)) + res_fc(fea)
        return self.fc_out(fea)

    def __repr__(self):
        """Return the class name."""
        return f"{self.__class__.__name__}"

class TransferNetwork(nn.Module):
    """Learn extended representations of materials during transfer learning.

    This network was designed to have little impact on predictions during
    training and enhance learning with the inclusion of extended features.
    """

    def __init__(self, input_dims, output_dims):
        """Instantiate a TransferNetwork to learn extended representations.

        Parameters
        ----------
        input_dims : int
            Dimensions of input layer

        output_dims : int
            Dimensions of output layer
        """
        super().__init__()
        self.layers = nn.Sequential(
            OrderedDict(
                [
                    ("fc1", nn.Linear(input_dims, 512)),
                    ("leakyrelu1", nn.LeakyReLU()),
                    ("fc2", nn.Linear(512, output_dims)),
                    ("leakyrelu2", nn.LeakyReLU()),
                ]
            )
        )

    def forward(self, x):
        """Perform a forward pass of the TransferNetwork.

        Parameters
        ----------
        x : _type_
            _description_

        Returns
        -------
        _type_
            _description_
        """
        x = self.layers(x)
        return x

class Embedder(nn.Module):
    """Perform composition-based embeddings of elemental features."""

    def __init__(
        self,
        d_model: int,
        compute_device: str = None,
        elem_prop: str = "mat2vec",
    ):
        """Embed elemental features, similar to CBFV.

        Parameters
        ----------
        d_model : int
            Row dimenions of elemental emeddings, by default 512
        compute_device : str
            Name of device which the model will be run on
        elem_prop : str
            Which elemental feature vector to use. Possible values are "jarvis",
            "magpie", "mat2vec", "oliynyk", "onehot", "ptable", and "random_200", by
            default "mat2vec"
        """
        super().__init__()
        self.d_model = d_model
        self.compute_device = compute_device

        elem_dir = join(dirname(__file__), "data", "element_properties")
        # # Choose what element information the model receives
        mat2vec = join(elem_dir, elem_prop + ".csv")  # element embedding
        # mat2vec = f'{elem_dir}/onehot.csv'  # onehot encoding (atomic number)
        # mat2vec = f'{elem_dir}/random_200.csv'  # random vec for elements

        cbfv = pd.read_csv(mat2vec, index_col=0).values
        feat_size = cbfv.shape[-1]
        self.fc_mat2vec = nn.Linear(feat_size, d_model).to(self.compute_device)
        zeros = np.zeros((1, feat_size))
        cat_array = np.concatenate([zeros, cbfv])
        cat_array = torch.as_tensor(cat_array, dtype=data_type_torch)
        # NOTE: Parameters within nn.Embedding
        self.cbfv = nn.Embedding.from_pretrained(cat_array).to(
            self.compute_device, dtype=data_type_torch
        )

    def forward(self, src):
        """Compute forward call for embedder class to perform elemental embeddings.

        Parameters
        ----------
        src : torch.tensor
            Tensor containing element numbers corresponding to elements in compound

        Returns
        -------
        torch.tensor
            Tensor containing elemental embeddings for compounds, reduced to d_model dimensions
        """
        mat2vec_emb = self.cbfv(src)
        x_emb = self.fc_mat2vec(mat2vec_emb)
        return x_emb

# %%
class FractionalEncoder(nn.Module):
    """Encode element fractional amount using a "fractional encoding".

    This is inspired by the positional encoder discussed by Vaswani.
    https://arxiv.org/abs/1706.03762
    """

    def __init__(self, d_model, resolution=100, log10=False, compute_device=None):
        """Instantiate the FractionalEncoder.

        Parameters
        ----------
        d_model : int
            Model size, see paper, by default 512
        resolution : int
            Number of discretizations for the fractional prevalence encoding, by default 100
        log10 : bool
            Whether to apply a log operation to fraction prevalence encoding, by default False
        compute_device : str
            The compute device to store and run the FractionalEncoder class
        """
        super().__init__()
        self.d_model = d_model // 2
        self.resolution = resolution
        self.log10 = log10
        self.compute_device = compute_device

        x = torch.linspace(
            0, self.resolution - 1, self.resolution, requires_grad=False
        ).view(self.resolution, 1)
        fraction = (
            torch.linspace(0, self.d_model - 1, self.d_model, requires_grad=False)
            .view(1, self.d_model)
            .repeat(self.resolution, 1)
        )

        pe = torch.zeros(self.resolution, self.d_model)
        pe[:, 0::2] = torch.sin(x / torch.pow(50, 2 * fraction[:, 0::2] / self.d_model))
        pe[:, 1::2] = torch.cos(x / torch.pow(50, 2 * fraction[:, 1::2] / self.d_model))
        pe = self.register_buffer("pe", pe)

    def forward(self, x):
        """Perform the forward pass of the fractional encoding.

        Parameters
        ----------
        x : torch.tensor
            Tensor of linear spaced values based on defined resolution

        Returns
        -------
        out
            Sinusoidal expansions of elemental fractions
        """
        x = x.clone()
        if self.log10:
            x = 0.0025 * (torch.log2(x)) ** 2
            x[x > 1] = 1
            # x = 1 - x  # for sinusoidal encoding at x=0
        x[x < 1 / self.resolution] = 1 / self.resolution
        frac_idx = torch.round(x * (self.resolution)).to(dtype=torch.long) - 1
        out = self.pe[frac_idx]

        return out

# %%
class Encoder(nn.Module):
    """Create elemental descriptor matrix via element embeddings and frac. encodings.

    See the CrabNet paper for further details:
    https://www.nature.com/articles/s41524-021-00545-1
    """

    def __init__(
        self,
        d_model,
        N,
        heads,
        extend_features=None,
        fractional=True,
        attention=True,
        compute_device=None,
        pe_resolution=5000,
        ple_resolution=5000,
        elem_prop="mat2vec",
        emb_scaler=1.0,
        pos_scaler=1.0,
        pos_scaler_log=1.0,
        dim_feedforward=2048,
        dropout=0.1,
    ):
        """Instantiate the Encoder class to create elemental descriptor matrix (EDM).

        Parameters
        ----------
        d_model : _type_
            _description_
        N : int, optional
            Number of encoder layers, by default 3
        heads : int, optional
            Number of attention heads to use, by default 4
        extend_features : Optional[List[str]]
            Additional features to grab from columns of the other DataFrames (e.g. state
            variables such as temperature or applied load), by default None
        fractional : bool, optional
            Whether to weight each element by its fractional contribution, by default True.
        attention : bool, optional
            Whether to perform self attention, by default True
        pe_resolution : int, optional
            Number of discretizations for the prevalence encoding, by default 5000
        ple_resolution : int, optional
            Number of discretizations for the prevalence log encoding, by default 5000
        elem_prop : str, optional
            Which elemental feature vector to use. Possible values are "jarvis",
            "magpie", "mat2vec", "oliynyk", "onehot", "ptable", and "random_200", by
            default "mat2vec"
        emb_scaler : float, optional
            _description_, by default 1.0
        pos_scaler : float, optional
            Scaling factor applied to fractional encoder, by default 1.0
        pos_scaler_log : float, optional
            Scaling factor applied to log fractional encoder, by default 1.0
        dim_feedforward : int, optional
            Dimenions of the feed forward network following transformer, by default 2048
        dropout : float, optional
            Percent dropout in the feed forward network following the transformer, by default 0.1
        """
        super().__init__()
        self.d_model = d_model
        self.N = N
        self.heads = heads
        self.extend_features = extend_features
        self.fractional = fractional
        self.attention = attention
        self.compute_device = compute_device
        self.pe_resolution = pe_resolution
        self.ple_resolution = ple_resolution
        self.elem_prop = elem_prop
        self.embed = Embedder(d_model=self.d_model, compute_device=self.compute_device)
        self.prevalence_encoder = FractionalEncoder(
            self.d_model, resolution=pe_resolution, log10=False
        )
        self.prevalence_log_encoder = FractionalEncoder(
            self.d_model, resolution=ple_resolution, log10=True
        )

        self.emb_scaler = nn.parameter.Parameter(torch.tensor([emb_scaler]))
        self.pos_scaler = nn.parameter.Parameter(torch.tensor([pos_scaler]))
        self.pos_scaler_log = nn.parameter.Parameter(torch.tensor([pos_scaler_log]))

        if self.attention:
            encoder_layer = nn.TransformerEncoderLayer(
                self.d_model,
                nhead=self.heads,
                dim_feedforward=dim_feedforward,
                dropout=dropout,
            )
            self.transformer_encoder = nn.TransformerEncoder(
                encoder_layer, num_layers=self.N
            )

    def forward(self, src, frac, extra_features=None):
        """Compute the forward pass for encoding the elemental descriptor matrix.

        Parameters
        ----------
        src : torch.tensor
            Tensor containing integers corresponding to elements in compound
        frac : torch.tensor
            Tensor containing the fractions of each element in compound
        extra_features : bool, optional
            Whether to append extra features after encoding, by default None

        Returns
        -------
        torch.tensor
            Tensor containing flattened transformer representations of compounds
            concatenated with extended features.
        """
        x = self.embed(src) * self.emb_scaler  # * 2 ** self.emb_scaler

        pe = torch.zeros_like(x)
        ple = torch.zeros_like(x)
        pe_scaler = self.pos_scaler
        ple_scaler = self.pos_scaler_log
        pe[:, :, : self.d_model // 2] = self.prevalence_encoder(frac) * pe_scaler
        ple[:, :, self.d_model // 2 :] = self.prevalence_log_encoder(frac) * ple_scaler

        mask = frac.unsqueeze(dim=-1)
        mask = torch.matmul(mask, mask.transpose(-2, -1))
        mask[mask != 0] = 1
        src_mask = mask[:, 0] != 1

        if self.attention:
            x_src = x + pe + ple
            x_src = x_src.transpose(0, 1)
            x = self.transformer_encoder(x_src, src_key_padding_mask=src_mask)
            x = x.transpose(0, 1)

        if self.fractional:
            x = x * frac.unsqueeze(2).repeat(1, 1, self.d_model)

        hmask = mask[:, :, 0:1].repeat(1, 1, self.d_model)
        if mask is not None:
            x = x.masked_fill(hmask == 0, 0)

        if self.extend_features is not None:
            n_elements = x.shape[1]
            X_extra = extra_features.repeat(1, 1, n_elements).permute([1, 2, 0])
            x = torch.concat((x, X_extra), axis=2)

        return x

# %%
class SubCrab(nn.Module):
    """SubCrab model class which implements the transformer architecture."""

    def __init__(
        self,
        out_dims=3,
        d_model=512,
        extend_features=None,
        d_extend=0,
        N=3,
        heads=4,
        fractional=False,
        attention=True,
        compute_device=None,
        out_hidden=[1024, 512, 256, 128],
        pe_resolution=5000,
        ple_resolution=5000,
        elem_prop="mat2vec",
        bias=False,
        emb_scaler=1.0,
        pos_scaler=1.0,
        pos_scaler_log=1.0,
        dim_feedforward=2048,
        dropout=0.1,
    ):
        """Instantiate a SubCrab class to be used within CrabNet.

        Parameters
        ----------
        out_dims : int, optional
            Output dimensions for Residual Network, by default 3
        d_model : int, optional
            Model size. See paper, by default 512
        extend_features : _type_, optional
            Additional features to grab from columns of the other DataFrames (e.g. state
            variables such as temperature or applied load), by default None
        d_extend : int, optional
            Number of extended features, by default 0
        N : int, optional
            Number of attention layers, by default 3
        heads : int, optional
            Number of attention heads, by default 4
        frac : bool, optional
            Whether to multiply `x` by the fractional amounts for each element, by default False
        attn : bool, optional
            Whether to perform self attention, by default True
        compute_device : _type_, optional
            Computing device to run model on, by default None
        out_hidden : list(int), optional
            Architecture of hidden layers in the Residual Network, by default [1024, 512, 256, 128]
        pe_resolution : int, optional
            Number of discretizations for the prevalence encoding, by default 5000
        ple_resolution : int, optional
            Number of discretizations for the prevalence log encoding, by default 5000
        elem_prop : str, optional
            Which elemental feature vector to use. Possible values are "jarvis", "magpie",
            "mat2vec", "oliynyk", "onehot", "ptable", and "random_200", by default "mat2vec"
        bias : bool, optional
            Whether to bias the Residual Network, by default False
        emb_scaler : float, optional
            Float value by which to scale the elemental embeddings, by default 1.0
        pos_scaler : float, optional
            Float value by which to scale the fractional encodings, by default 1.0
        pos_scaler_log : float, optional
            Float value by which to scale the log fractional encodings, by default 1.0
        dim_feedforward : int, optional
            Dimenions of the feed forward network following transformer, by default 2048
        dropout : float, optional
            Percent dropout in the feed forward network following the transformer, by default 0.1
        """
        super().__init__()
        self.avg = True
        self.out_dims = out_dims
        self.d_model = d_model
        self.extend_features = extend_features
        self.d_extend = d_extend
        self.N = N
        self.heads = heads
        self.fractional = fractional
        self.attention = attention
        self.compute_device = compute_device
        self.bias = bias
        self.encoder = Encoder(
            d_model=self.d_model,
            N=self.N,
            heads=self.heads,
            attention=self.attention,
            compute_device=self.compute_device,
            pe_resolution=pe_resolution,
            ple_resolution=ple_resolution,
            elem_prop=elem_prop,
            emb_scaler=emb_scaler,
            pos_scaler=pos_scaler,
            pos_scaler_log=pos_scaler_log,
            dim_feedforward=dim_feedforward,
            dropout=dropout,
        )

        self.out_hidden = out_hidden
        self.output_nn = ResidualNetwork(
            self.d_model + self.d_extend,
            self.out_dims,
            self.out_hidden,
            self.bias,
        )

    def forward(self, src, frac, extra_features=None):
        """Compute forward pass of the SubCrab model class (i.e. transformer).

        Parameters
        ----------
        src : torch.tensor
            Tensor containing element numbers corresponding to elements in compound
        frac : torch.tensor
            Tensor containing fractional amounts of each element in compound
        extra_features : bool, optional
            Whether to append extra features after encoding, by default None

        Returns
        -------
        torch.tensor
            Model output containing predicted value and uncertainty for that value
        """
        output = self.encoder(src, frac, extra_features)
        # output = self.transfer_nn(output)

        # average the "element contribution", mask so you only average "elements" (i.e.
        # not padded zero values)
        elem_pad_mask = (src == 0).unsqueeze(-1).repeat(1, 1, self.out_dims)
        output = self.output_nn(output)  # simple linear
        if self.avg:
            output = output.masked_fill(elem_pad_mask, 0)
            output = output.sum(dim=1) / (~elem_pad_mask).sum(dim=1)
            output, logits = output.chunk(2, dim=-1)
            probability = torch.ones_like(output)
            probability[:, : logits.shape[-1]] = torch.sigmoid(logits)
            output = output * probability

        return output

# %%
if __name__ == "__main__":
    model = SubCrab()


"""
Trade exit triggers

"""

import numpy as np

class TradeTargets():
    """
    Calculate profit targets, stop loss levels etc.

    """
    @classmethod
    def exit_and_stop_targets(
            cls, prices, params, trade_price_dict):
        """
        Calculate exit and stop targets.

        Parameters
        ----------
        prices : DataFrame
            The OHLC data.
        exit_amount : Float
            The dollar exit amount. The default is $1000.00.
        stop_amount : Float
            The dollar stop amount. The default is $500.00.
        position_size : Int, optional
            The number of units to trade. The default is based on equity.
        trade_price_dict : Dict
            Dictionary of trade entry/high/low/close series.

        Returns
        -------
        prices : DataFrame
            The OHLC data.

        """

        if params['exit_type'] is not None:
            prices = cls._exit_targets(
                prices=prices,
                exit_amount=params['exit_amount'],
                trade_price_dict=trade_price_dict,
                params=params)

        if params['stop_type'] is not None:
            prices = cls._stop_targets(
                prices=prices,
                stop_amount=params['stop_amount'],
                trade_price_dict=trade_price_dict,
                params=params)

        return prices

    @classmethod
    def _exit_targets(
            cls, prices, exit_amount, trade_price_dict, params):
        """
        Create 4 series of exit targets

        Parameters
        ----------
        prices : DataFrame
            The OHLC data.
        exit_amount : Float
            The dollar exit amount. The default is $1000.00.
        position_size : Int, optional
            The number of units to trade. The default is based on equity.
        trade_price_dict : Dict
            Dictionary of trade entry/high/low/close series.

        Returns
        -------
        prices : DataFrame
            The OHLC data..

        """
        # Generate profit targets / trailing stops
        prices['exit_profit_target'], prices['exit_initial_dollar_loss'], \
            prices['exit_trailing_close'], \
                prices['exit_trailing_high_low'] = cls._pnl_targets(
                    prices=prices, dollar_amount=exit_amount,
                    trade_price_dict=trade_price_dict,
                    params=params)

        return prices

    @classmethod
    def _stop_targets(
            cls, prices, stop_amount, trade_price_dict, params):
        """
        Create 4 series of stop targets

        Parameters
        ----------
        prices : DataFrame
            The OHLC data.
        stop_amount : Float
            The dollar stop amount. The default is $500.00.
        position_size : Int, optional
            The number of units to trade. The default is based on equity.
        trade_price_dict : Dict
            Dictionary of trade entry/high/low/close series.

        Returns
        -------
        prices : DataFrame
            The OHLC data.

        """
        # Generate profit targets / trailing stops
        prices['stop_profit_target'], prices['stop_initial_dollar_loss'], \
            prices['stop_trailing_close'], \
                prices['stop_trailing_high_low'] = cls._pnl_targets(
                    prices=prices, dollar_amount=stop_amount,
                    trade_price_dict=trade_price_dict,
                    params=params)

        return prices

    @staticmethod
    def _pnl_targets(
            prices, dollar_amount, trade_price_dict, params):
        """
        Create profit and loss stop and exit points

        Parameters
        ----------
        prices : DataFrame
            The OHLC data.
        exit_amount : Float
            The dollar exit amount. The default is $1000.00.
        position_size : Int
            The number of units of the chosen ticker to trade.
        trade_number : Series
            Array of trade numbers.
        end_of_day_position : Series
            The number of units of position held at the end of day.
        trade_entry_price : Series
            The entry price for each trade.
        trade_high_price : Series
            The high price for each trade.
        trade_low_price : Series
            The low price for each trade.
        trade_close_high_price : Series
            The highest closing price for each trade.
        trade_close_low_price : Series
            The lowest closing price for each trade.

        Returns
        -------
        profit_target : Series
            The exit levels for each trade based on a dollar loss from the
            entry level.
        initial_dollar_loss : Series
            The exit levels for each trade based on a profit target
        trailing_close : Series
            The exit levels for each trade based on the trailing close.
        trailing_high_low : Series
            The exit levels for each trade based on the trailing high / low.

        """

        # Create empty arrays to store the values
        trade_target = np.array([0.0]*len(prices))
        profit_target = np.array([0.0]*len(prices))
        initial_dollar_loss = np.array([0.0]*len(prices))
        trailing_close = np.array([0.0]*len(prices))
        trailing_high_low = np.array([0.0]*len(prices))

        end_of_day_position = prices['raw_end_of_day_position']
        position_size = prices['position_size']

        # For each row in the data
        for row in range(1, len(prices)):

            # Calculate the trade target (distance that price has to change) by
            # dividing the dollar amount by the number of units making up the
            # position size
            if position_size[row] == 0:
                trade_target[row] = 0
            else:
                trade_target[row] = np.round(
                    (dollar_amount / params['contract_point_value'])
                    / position_size[row], 2)

            # If there is a trade on
            if prices['raw_trade_number'][row] != 0:

                # If there is a long position
                if end_of_day_position[row] > 0:

                    # Set the profit target to the trade entry price plus the
                    # trade target
                    profit_target[row] = (
                        trade_price_dict['trade_entry_price'][row]
                        + trade_target[row])

                    # Set the initial dollar loss target to the trade entry
                    # price minus the trade target
                    initial_dollar_loss[row] = (
                        trade_price_dict['trade_entry_price'][row]
                        - trade_target[row])

                    # Set the trailing close target to the closing high price
                    # of the trade minus the trade target
                    trailing_close[row] = (
                        trade_price_dict['trade_close_high_price'][row]
                        - trade_target[row])

                    # Set the trailing high/low target to the high price
                    # of the trade minus the trade target
                    trailing_high_low[row] = (
                        trade_price_dict['trade_high_price'][row]
                        - trade_target[row])

                # If there is a short position
                else:
                    # Set the profit target to the trade entry price minus the
                    # trade target
                    profit_target[row] = (
                        trade_price_dict['trade_entry_price'][row]
                        - trade_target[row])

                    # Set the initial dollar loss target to the trade entry
                    # price plus the trade target
                    initial_dollar_loss[row] = (
                        trade_price_dict['trade_entry_price'][row]
                        + trade_target[row])

                    # Set the trailing close target to the closing low price
                    # of the trade plus the trade target
                    trailing_close[row] = (
                        trade_price_dict['trade_close_low_price'][row]
                        + trade_target[row])

                    # Set the trailing high/low target to the low price
                    # of the trade minus the trade target
                    trailing_high_low[row] = (
                        trade_price_dict['trade_low_price'][row]
                        + trade_target[row])

        return profit_target, initial_dollar_loss, trailing_close, \
            trailing_high_low


import numpy as np

from .libfft import FFT
from .pencil import Pencil
from .pencil import Subcomm

class Transform(object):
    """Class for performing any parallel transform, forward or backward

    Parameters
    ----------
    xfftn : list of serial transform objects
    transfer : list of global redistribution objects
    pencil : list of two pencil objects
        The two pencils represent the input and final output configuration of
        the distributed global arrays

    """
    def __init__(self, xfftn, transfer, pencil):
        assert len(xfftn) == len(transfer) + 1 and len(pencil) == 2
        self._xfftn = tuple(xfftn)
        self._transfer = tuple(transfer)
        self._pencil = tuple(pencil)

    @property
    def input_array(self):
        """Return input array of Transform"""
        return self._xfftn[0].input_array

    @property
    def output_array(self):
        """Return output array of Transform"""
        return self._xfftn[-1].output_array

    @property
    def input_pencil(self):
        """Return input pencil of Transform"""
        return self._pencil[0]

    @property
    def output_pencil(self):
        """Return output pencil of Transform"""
        return self._pencil[1]

    def __call__(self, input_array=None, output_array=None, **kw):
        """Compute transform

        Parameters
        ----------
        input_array : array, optional
        output_array : array, optional
        kw : dict
            parameters to serial transforms
            Note in particular that the keyword 'normalize'=True/False can be
            used to turn normalization on or off. Default is to enable
            normalization for forward transforms and disable it for backward.

        Note
        ----
        If input_array/output_array are not given, then use predefined arrays
        as planned with serial transform object _xfftn.

        """
        if input_array is not None:
            self.input_array[...] = input_array

        for i in range(len(self._transfer)):
            self._xfftn[i](**kw)
            arrayA = self._xfftn[i].output_array
            arrayB = self._xfftn[i+1].input_array
            self._transfer[i](arrayA, arrayB)
        self._xfftn[-1](**kw)

        if output_array is not None:
            output_array[...] = self.output_array
            return output_array
        else:
            return self.output_array

class PFFT(object):
    """Base class for parallel FFT transforms

    Parameters
    ----------
    comm : MPI communicator
    shape : sequence of ints, optional
        shape of input array planned for
    axes : None, int, sequence of ints or sequence of sequence of ints, optional
        axes to transform over.

        - None -> All axes are transformed
        - int -> Just one axis to transform over
        - sequence of ints -> e.g., (0, 1, 2) or (0, 2, 1)
        - sequence of sequence of ints -> e.g., ((0,), (1,)) or ((0,), (1, 2))
          For seq. of seq. of ints all but the last transformed sequence
          may be longer than 1. This corresponds to collapsing axes, where
          serial FFTs are performed for all collapsed axes in one single call
    dtype : np.dtype, optional
        Type of input array
    grid : sequence of ints, optional
        Define processor grid sizes. Non positive values act as wildcards to
        allow MPI compute optimal decompositions. The sequence is padded with
        ones to match the global transform dimension.
        Use ``(-1,)`` to get a slab decomposition on the first axis.
        Use ``(1, -1)`` to get a slab decomposition  on the second axis.
        Use ``(P, Q)`` or ``(P, Q, 1)`` to get a 3D transform with 2D-pencil
        decomposition on a PxQ processor grid with the last axis non distributed.
        Use ``(P, 1, Q)`` to get a 3D transform with 2D-pencil decomposition on
        a PxQ processor grid with the second to last axis non distributed.
    padding : bool, number or sequence of numbers, optional
        If False, then no padding. If number, then apply this number as padding
        factor for all axes. If sequence of numbers, then each number gives the
        padding for each axis. Must be same length as axes.
    collapse : bool, optional
        If True try to collapse several serial transforms into one
    backend : str, optional
        Choose backend for serial transforms (``fftw``, ``pyfftw``, ``numpy``,
        ``scipy``, ``mkl_fft``). Default is ``fftw``
    transforms : None or dict, optional
        Dictionary of axes to serial transforms (forward and backward) along
        those axes. For example::

            {(0, 1): (dctn, idctn), (2, 3): (dstn, idstn)}

        If missing the default is to use rfftn/irfftn for real input arrays and
        fftn/ifftn for complex input arrays. Real-to-real transforms can be
        configured using this dictionary and real-to-real transforms from the
        :mod:`.fftw.xfftn` module. See Examples.

    Other Parameters
    ----------------
    darray : DistArray object, optional
        Create PFFT using information contained in ``darray``, neglecting most
        optional Parameters above
    slab : bool or int, optional
        DEPRECATED. If True then distribute only one axis of the global array.

    Methods
    -------
    forward(input_array=None, output_array=None, **kw)
        Parallel forward transform. The method is an instance of the
        :class:`.Transform` class. See :meth:`.Transform.__call__`

        Parameters
        ----------
        input_array : array, optional
        output_array : array, optional
        kw : dict
            parameters to serial transforms

        Returns
        -------
        output_array : array

    backward(input_array=None, output_array=None, **kw)
        Parallel backward transform. The method is an instance of the
        :class:`.Transform` class. See :meth:`.Transform.__call__`

        Parameters
        ----------
        input_array : array, optional
        output_array : array, optional
        kw : dict
            parameters to serial transforms

        Returns
        -------
        output_array : array

    Examples
    --------
    >>> import numpy as np
    >>> from mpi4py import MPI
    >>> from mpi4py_fft import PFFT, newDistArray
    >>> N = np.array([12, 14, 15], dtype=int)
    >>> fft = PFFT(MPI.COMM_WORLD, N, axes=(0, 1, 2))
    >>> u = newDistArray(fft, False)
    >>> u[:] = np.random.random(u.shape).astype(u.dtype)
    >>> u_hat = fft.forward(u)
    >>> uj = np.zeros_like(u)
    >>> uj = fft.backward(u_hat, uj)
    >>> assert np.allclose(uj, u)

    Now configure with real-to-real discrete cosine transform type 3

    >>> from mpi4py_fft.fftw import rfftn, irfftn, dctn, idctn
    >>> import functools
    >>> dct = functools.partial(dctn, type=3)
    >>> idct = functools.partial(idctn, type=3)
    >>> transforms = {(1, 2): (dct, idct)}
    >>> r2c = PFFT(MPI.COMM_WORLD, N, axes=((0,), (1, 2)), transforms=transforms)
    >>> u = newDistArray(r2c, False)
    >>> u[:] = np.random.random(u.shape).astype(u.dtype)
    >>> u_hat = r2c.forward(u)
    >>> uj = np.zeros_like(u)
    >>> uj = r2c.backward(u_hat, uj)
    >>> assert np.allclose(uj, u)

    """
    def __init__(self, comm, shape=None, axes=None, dtype=float, grid=None,
                 padding=False, collapse=False, backend='fftw',
                 transforms=None, darray=None, **kw):
        # pylint: disable=too-many-locals
        # pylint: disable=too-many-branches
        # pylint: disable=too-many-statements

        if shape is None:
            assert darray is not None
            shape = darray.pencil.shape

        if axes is not None:
            axes = list(axes) if np.ndim(axes) else [axes]
        else:
            axes = list(range(len(shape)))
            if darray is not None:
                # Make sure aligned axis of darray is transformed first
                axes = list(np.roll(axes, len(shape)-1-darray.alignment))

        for i, ax in enumerate(axes):
            if isinstance(ax, (int, np.integer)):
                if ax < 0:
                    ax += len(shape)
                axes[i] = (ax,)
            else:
                assert isinstance(ax, (tuple, list))
                ax = list(ax)
                for j, a in enumerate(ax):
                    assert isinstance(a, int)
                    if a < 0:
                        a += len(shape)
                        ax[j] = a
                axes[i] = ax
            assert min(axes[i]) >= 0
            assert max(axes[i]) < len(shape)
            assert 0 < len(axes[i]) <= len(shape)
            assert sorted(axes[i]) == sorted(set(axes[i]))

        self.axes = axes
        shape = list(shape)

        if darray is None:
            dtype = np.dtype(dtype)
            assert dtype.char in 'fdgFDG'

            if padding is not False:
                assert len(padding) == len(shape)
                for ax in axes:
                    if len(ax) == 1 and padding[ax[0]] > 1.0+1e-6:
                        old = float(shape[ax[0]])
                        shape[ax[0]] = int(np.floor(shape[ax[0]]*padding[ax[0]]))
                        padding[ax[0]] = shape[ax[0]] / old

            self._input_shape = tuple(shape)
            assert len(shape) > 0
            assert min(shape) > 0

            slab = kw.pop('slab', False)

            if grid is not None:
                assert not isinstance(comm, Subcomm)
                assert slab is False
                grid = tuple(grid)
                assert len(grid) <= len(shape)
                dims = list(grid) + [1] * (len(shape) - len(grid))
                comm = Subcomm(comm, dims)

            if isinstance(comm, Subcomm):
                assert slab is False
                assert len(comm) == len(shape)
                assert np.all([comm[ax].Get_size() == 1 for ax in axes[-1]])
                self.subcomm = comm
            else:
                if slab is False or slab is None:
                    dims = [0] * len(shape)
                    for ax in axes[-1]:
                        dims[ax] = 1
                else: #pragma: no cover
                    if slab is True:
                        axis = (axes[-1][-1] + 1) % len(shape)
                    else:
                        axis = slab
                        if axis < 0:
                            axis = axis + len(shape)
                        assert 0 <= axis < len(shape)
                    dims = [1] * len(shape)
                    dims[axis] = comm.Get_size()

                self.subcomm = Subcomm(comm, dims)
        else:
            dtype = darray.dtype
            self.subcomm = darray.subcomm
            self._input_shape = tuple(shape)
            commsizes = darray.commsizes
            assert np.all([commsizes[ax] == 1 for ax in axes[-1]]), "Set keyword axes such that axes to transform first are aligned"

        self.collapse = collapse
        if collapse is True:
            groups = [[]]
            for ax in reversed(axes):
                if np.all([self.subcomm[axis].Get_size() == 1 for axis in ax]):
                    [groups[0].insert(0, axis) for axis in reversed(ax)]
                else:
                    groups.insert(0, ax)
            axes = groups

        self.axes = tuple(map(tuple, axes))
        self.xfftn = []
        self.transfer = []
        self.pencil = [None, None]

        axes = self.axes[-1]
        pencil = Pencil(self.subcomm, shape, axes[-1])
        xfftn = FFT(pencil.subshape, axes, dtype, padding, backend=backend,
                    transforms=transforms, **kw)
        self.xfftn.append(xfftn)
        self.pencil[0] = pencilA = pencil
        if not shape[axes[-1]] == xfftn.forward.output_array.shape[axes[-1]]:
            dtype = xfftn.forward.output_array.dtype
            shape[axes[-1]] = xfftn.forward.output_array.shape[axes[-1]]
            pencilA = Pencil(self.subcomm, shape, axes[-1])

        for axes in reversed(self.axes[:-1]):
            pencilB = pencilA.pencil(axes[-1])
            transAB = pencilA.transfer(pencilB, dtype)
            xfftn = FFT(pencilB.subshape, axes, dtype, padding, backend=backend,
                        transforms=transforms, **kw)
            self.xfftn.append(xfftn)
            self.transfer.append(transAB)
            pencilA = pencilB
            if not shape[axes[-1]] == xfftn.forward.output_array.shape[axes[-1]]:
                dtype = xfftn.forward.output_array.dtype
                shape[axes[-1]] = xfftn.forward.output_array.shape[axes[-1]]
                pencilA = Pencil(pencilB.subcomm, shape, axes[-1])

        self.pencil[1] = pencilA
        self._output_shape = tuple(shape)

        self.forward = Transform(
            [o.forward for o in self.xfftn],
            [o.forward for o in self.transfer],
            self.pencil)
        self.backward = Transform(
            [o.backward for o in self.xfftn[::-1]],
            [o.backward for o in self.transfer[::-1]],
            self.pencil[::-1])

    def destroy(self):
        if isinstance(self.subcomm, Subcomm):
            self.subcomm.destroy()
        for trans in self.transfer:
            trans.destroy()

    def shape(self, forward_output=True):
        """The local (to each processor) shape of data

        Parameters
        ----------
        forward_output : bool, optional
            Return shape of output array (spectral space) if True, else return
            shape of input array (physical space)
        """
        if forward_output is not True:
            return self.forward.input_pencil.subshape
        return self.backward.input_pencil.subshape

    def local_slice(self, forward_output=True):
        """The local view into the global data

        Parameters
        ----------
        forward_output : bool, optional
            Return local slices of output array (spectral space) if True, else
            return local slices of input array (physical space)

        """
        if forward_output is not True:
            ip = self.forward.input_pencil
            s = [slice(start, start+shape) for start, shape in zip(ip.substart,
                                                                   ip.subshape)]
        else:
            ip = self.backward.input_pencil
            s = [slice(start, start+shape) for start, shape in zip(ip.substart,
                                                                   ip.subshape)]
        return tuple(s)

    def global_shape(self, forward_output=False):
        """Return global shape of associated tensors

        Parameters
        ----------
        forward_output : bool, optional
            If True then return global shape of spectral space, i.e., the input
            to a backward transfer. If False then return shape of physical
            space, i.e., the input to a forward transfer.
        """
        if forward_output:
            return self._output_shape
        return self._input_shape

    @property
    def dimensions(self):
        """The number of dimensions for transformed arrays"""
        return len(self.forward.input_array.shape)

    def dtype(self, forward_output=False):
        """The type of transformed arrays

        Parameters
        ----------
        forward_output : bool, optional
            If True then return dtype of an array that is the result of a
            forward transform. Otherwise, return the dtype of an array that
            is input to a forward transform.
        """
        if forward_output:
            return self.forward.output_array.dtype
        return self.forward.input_array.dtype