#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Copyright 1999-2020 Alibaba Group Holding Ltd.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import numpy as np
from numpy.linalg import LinAlgError

from ... import opcodes as OperandDef
from ...serialize import KeyField, StringField
from ...core import ExecutableTuple
from ..array_utils import device, as_same_device
from ..datasource import tensor as astensor
from ..operands import TensorHasInput, TensorOperandMixin
from ..core import TensorOrder
from .core import SFQR, TSQR


class TensorQR(TensorHasInput, TensorOperandMixin):
    _op_type_ = OperandDef.QR

    _input = KeyField('input')
    _method = StringField('method')

    def __init__(self, method=None, dtype=None, **kw):
        super().__init__(_method=method, _dtype=dtype, **kw)

    @property
    def method(self):
        return self._method

    @property
    def output_limit(self):
        return 2

    def _set_inputs(self, inputs):
        super()._set_inputs(inputs)
        self._input = self._inputs[0]

    def __call__(self, a):
        a = astensor(a)

        if a.ndim != 2:
            raise LinAlgError(f'{a.ndim}-dimensional tensor given. '
                              'Tensor must be two-dimensional')

        tiny_q, tiny_r = np.linalg.qr(np.ones((1, 1), dtype=a.dtype))

        x, y = a.shape
        q_shape, r_shape = (a.shape, (y, y)) if x > y else ((x, x), a.shape)
        q, r = self.new_tensors([a],
                                kws=[{'side': 'q', 'dtype': tiny_q.dtype,
                                      'shape': q_shape, 'order': TensorOrder.C_ORDER},
                                     {'side': 'r', 'dtype': tiny_r.dtype,
                                      'shape': r_shape, 'order': TensorOrder.C_ORDER}])
        return ExecutableTuple([q, r])

    @classmethod
    def tile(cls, op):
        q, r = op.outputs
        q_dtype, r_dtype = q.dtype, r.dtype
        q_shape, r_shape = q.shape, r.shape
        in_tensor = op.input
        if in_tensor.chunk_shape == (1, 1):
            in_chunk = in_tensor.chunks[0]
            chunk_op = op.copy().reset_key()
            qr_chunks = chunk_op.new_chunks([in_chunk], kws=[
                {'side': 'q', 'shape': q_shape, 'index': in_chunk.index},
                {'side': 'r', 'shape': r_shape, 'index': in_chunk.index}
            ])
            q_chunk, r_chunk = qr_chunks

            new_op = op.copy()
            kws = [
                {'chunks': [q_chunk], 'nsplits': ((q_shape[0],), (q_shape[1],)),
                 'dtype': q_dtype, 'shape': q_shape, 'order': q.order},
                {'chunks': [r_chunk], 'nsplits': ((r_shape[0],), (r_shape[1],)),
                 'dtype': r_dtype, 'shape': r_shape, 'order': r.order}
            ]
            return new_op.new_tensors(op.inputs, kws=kws)
        elif op.method == 'tsqr':
            return TSQR.tile(op)
        elif op.method == 'sfqr':
            return SFQR.tile(op)
        else:
            raise NotImplementedError('Only tsqr method supported for now')

    @classmethod
    def execute(cls, ctx, op):
        (a,), device_id, xp = as_same_device(
            [ctx[c.key] for c in op.inputs], device=op.device, ret_extra=True)

        with device(device_id):
            q, r = xp.linalg.qr(a)
            qc, rc = op.outputs
            ctx[qc.key] = q
            ctx[rc.key] = r


def qr(a, method='tsqr'):
    """
    Compute the qr factorization of a matrix.

    Factor the matrix `a` as *qr*, where `q` is orthonormal and `r` is
    upper-triangular.

    Parameters
    ----------
    a : array_like, shape (M, N)
        Matrix to be factored.
    method: {'tsqr', 'sfqr'}, optional
        method to calculate qr factorization, tsqr as default

        TSQR is presented in:

            A. Benson, D. Gleich, and J. Demmel.
            Direct QR factorizations for tall-and-skinny matrices in
            MapReduce architectures.
            IEEE International Conference on Big Data, 2013.
            http://arxiv.org/abs/1301.1071

        FSQR is a QR decomposition for fat and short matrix:
            A = [A1, A2, A3, ...], A1 may be decomposed as A1 = Q1 * R1,
            for A = Q * R, Q = Q1, R = [R1, R2, R3, ...] where A2 = Q1 * R2, A3 = Q1 * R3, ...

    Returns
    -------
    q : Tensor of float or complex, optional
        A matrix with orthonormal columns. When mode = 'complete' the
        result is an orthogonal/unitary matrix depending on whether or not
        a is real/complex. The determinant may be either +/- 1 in that
        case.
    r : Tensor of float or complex, optional
        The upper-triangular matrix.

    Raises
    ------
    LinAlgError
        If factoring fails.

    Notes
    -----
    For more information on the qr factorization, see for example:
    http://en.wikipedia.org/wiki/QR_factorization

    Examples
    --------
    >>> import mars.tensor as mt

    >>> a = mt.random.randn(9, 6)
    >>> q, r = mt.linalg.qr(a)
    >>> mt.allclose(a, mt.dot(q, r)).execute()  # a does equal qr
    True

    """
    op = TensorQR(method=method)
    return op(a)
