# coding=utf-8
# 2021.03 modified (1) notears_linear(def) to Notears(class)
# 2021.03 added    (1) logging; 
#                  (2) BaseLearner
# 2021.03 deleted  (1) __main__
# Huawei Technologies Co., Ltd. 
# 
# Copyright (C) 2021. Huawei Technologies Co., Ltd. All rights reserved.
# 
# Copyright (c) Xun Zheng (https://github.com/xunzheng/notears)
#
# 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 logging
import numpy as np
import scipy.optimize as sopt
from scipy.special import expit as sigmoid

from castle.common import BaseLearner, Tensor
from castle.common.consts import NOTEARS_VALID_PARAMS
from castle.common.validator import check_args_value


class Notears(BaseLearner):
    """
    Notears Algorithm.
    A gradient-based algorithm for linear data models (typically with least-squares loss).

    Parameters
    ----------
    lambda1: float 
        l1 penalty parameter
    loss_type: str 
        l2, logistic, poisson
    max_iter: int 
        max num of dual ascent steps
    h_tol: float 
        exit if |h(w_est)| <= htol
    rho_max: float 
        exit if rho >= rho_max
    w_threshold: float 
        drop edge if |weight| < threshold

    Attributes
    ----------
    causal_matrix : numpy.ndarray
        Learned causal structure matrix

    References
    ----------
    https://arxiv.org/abs/1803.01422
    
    Examples
    --------
    >>> from castle.algorithms import Notears
    >>> from castle.datasets import load_dataset
    >>> from castle.common import GraphDAG
    >>> from castle.metrics import MetricsDAG
    >>> X, true_dag, _ = load_dataset('IID_Test')
    >>> n = Notears()
    >>> n.learn(X)
    >>> GraphDAG(n.causal_matrix, true_dag)
    >>> met = MetricsDAG(n.causal_matrix, true_dag)
    >>> print(met.metrics)
    """

    @check_args_value(NOTEARS_VALID_PARAMS)
    def __init__(self, lambda1=0.1, 
                 loss_type='l2', 
                 max_iter=100, 
                 h_tol=1e-8, 
                 rho_max=1e+16, 
                 w_threshold=0.3):

        super().__init__()

        self.lambda1 = lambda1
        self.loss_type = loss_type
        self.max_iter = max_iter
        self.h_tol = h_tol
        self.rho_max = rho_max
        self.w_threshold = w_threshold

    def learn(self, data, columns=None, **kwargs):
        """
        Set up and run the Notears algorithm.

        Parameters
        ----------
        data: castle.Tensor or numpy.ndarray
            The castle.Tensor or numpy.ndarray format data you want to learn.
        columns : Index or array-like
            Column labels to use for resulting tensor. Will default to
            RangeIndex (0, 1, 2, ..., n) if no column labels are provided.
        """
        X = Tensor(data, columns=columns)

        W_est = self.notears_linear(X, lambda1=self.lambda1, 
                                    loss_type=self.loss_type, 
                                    max_iter=self.max_iter, 
                                    h_tol=self.h_tol, 
                                    rho_max=self.rho_max)
        causal_matrix = (abs(W_est) > self.w_threshold).astype(int)
        self.weight_causal_matrix = Tensor(W_est,
                                           index=X.columns,
                                           columns=X.columns)
        self.causal_matrix = Tensor(causal_matrix, index=X.columns,
                                    columns=X.columns)

    def notears_linear(self, X, lambda1, loss_type, max_iter, h_tol, 
                       rho_max):
        """
        Solve min_W L(W; X) + lambda1 ‖W‖_1 s.t. h(W) = 0 using 
        augmented Lagrangian.

        Parameters
        ----------
        X: np.ndarray 
            n*d sample matrix

        Return
        ------
        W_est: np.ndarray
            d*d estimated DAG
        """
        def _loss(W):
            """Evaluate value and gradient of loss."""
            M = X @ W
            if loss_type == 'l2':
                R = X - M
                loss = 0.5 / X.shape[0] * (R ** 2).sum()
                G_loss = - 1.0 / X.shape[0] * X.T @ R
            elif loss_type == 'logistic':
                loss = 1.0 / X.shape[0] * (np.logaddexp(0, M) - X * M).sum()
                G_loss = 1.0 / X.shape[0] * X.T @ (sigmoid(M) - X)
            elif loss_type == 'poisson':
                S = np.exp(M)
                loss = 1.0 / X.shape[0] * (S - X * M).sum()
                G_loss = 1.0 / X.shape[0] * X.T @ (S - X)
            else:
                raise ValueError('unknown loss type')
            return loss, G_loss

        def _h(W):
            """
            Evaluate value and gradient of acyclicity constraint.
            """
            #     E = slin.expm(W * W)  # (Zheng et al. 2018)
            #     h = np.trace(E) - d
            M = np.eye(d) + W * W / d  # (Yu et al. 2019)
            E = np.linalg.matrix_power(M, d - 1)
            h = (E.T * M).sum() - d
            G_h = E.T * W * 2
            return h, G_h

        def _adj(w):
            """
            Convert doubled variables ([2 d^2] array) back to original 
            variables ([d, d] matrix).
            """
            return (w[:d * d] - w[d * d:]).reshape([d, d])

        def _func(w):
            """
            Evaluate value and gradient of augmented Lagrangian for 
            doubled variables ([2 d^2] array).
            """
            W = _adj(w)
            loss, G_loss = _loss(W)
            h, G_h = _h(W)
            obj = loss + 0.5 * rho * h * h + alpha * h + lambda1 * w.sum()
            G_smooth = G_loss + (rho * h + alpha) * G_h
            g_obj = np.concatenate((G_smooth + lambda1, - G_smooth + lambda1), 
                                   axis=None)
            return obj, g_obj

        n, d = X.shape
        # double w_est into (w_pos, w_neg)
        w_est, rho, alpha, h = np.zeros(2 * d * d), 1.0, 0.0, np.inf
        bnds = [(0, 0) if i == j else (0, None) for _ in range(2) 
                for i in range(d) for j in range(d)]
        if loss_type == 'l2':
            X = X - np.mean(X, axis=0, keepdims=True)
        
        logging.info('[start]: n={}, d={}, iter_={}, h_={}, rho_={}'.format( \
                    n, d, max_iter, h_tol, rho_max))

        for i in range(max_iter):
            w_new, h_new = None, None
            while rho < rho_max:
                sol = sopt.minimize(_func, w_est, method='L-BFGS-B', 
                                    jac=True, bounds=bnds)
                w_new = sol.x
                h_new, _ = _h(_adj(w_new))
                
                logging.info(
                    '[iter {}] h={:.3e}, loss={:.3f}, rho={:.1e}'.format( \
                    i, h_new, _func(w_est)[0], rho))
                
                if h_new > 0.25 * h:
                    rho *= 10
                else:
                    break
            w_est, h = w_new, h_new
            alpha += rho * h

            if h <= h_tol or rho >= rho_max:
                break

        W_est = _adj(w_est)

        logging.info('FINISHED')

        return W_est
