# Copyright 2021 Huawei Technologies Co., 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.
# ============================================================================
"""THOR"""
import numpy as np
from mindspore.ops import functional as F, composite as C, operations as P
from mindspore.common.initializer import initializer
from mindspore.common.parameter import Parameter, ParameterTuple
from mindspore.common.tensor import Tensor
import mindspore.nn as nn
import mindspore.common.dtype as mstype
from mindspore._checkparam import Validator
from mindspore.nn.optim.optimizer import Optimizer
from mindspore.parallel._utils import _get_device_num, _get_gradients_mean
from mindspore import context
from mindspore.context import ParallelMode
from mindspore.nn.layer import Dense_Thor, Conv2d_Thor, Embedding_Thor
from mindspore.nn.wrap import DistributedGradReducer
from mindspore.train.train_thor.convert_utils import ConvertNetUntils
from mindspore.parallel._auto_parallel_context import auto_parallel_context

# Enumerates types of Layer
Other = -1
Conv = 1
FC = 2
Embedding = 3
LayerNorm = 4
BatchNorm = 5

_momentum_opt = C.MultitypeFuncGraph("momentum_opt")

op_add = P.AddN()
apply_decay = C.MultitypeFuncGraph("apply_decay")

@apply_decay.register("Number", "Bool", "Tensor", "Tensor")
def _tensor_apply_decay(weight_decay, if_apply, weight, gradient):
    """Get grad with weight_decay."""
    if if_apply:
        return op_add((weight * weight_decay, gradient))
    return gradient

@_momentum_opt.register("Function", "Tensor", "Tensor", "Tensor", "Tensor", "Tensor")
def _tensor_run_opt_ext(opt, momentum, learning_rate, gradient, weight, moment):
    """Apply momentum optimizer to the weight parameter using Tensor."""
    success = True
    success = F.depend(success, opt(weight, moment, learning_rate, gradient, momentum))
    return success

C0 = 16

def caculate_device_shape(matrix_dim, channel, is_A):
    ll = (0)
    if is_A:
        if channel // C0 == 0:
            matrix_dim = (matrix_dim / channel) * C0
        ll = (int(matrix_dim // C0), int(matrix_dim // C0), C0, C0), int(matrix_dim)
    else:
        ll = (int(matrix_dim // C0), int(matrix_dim // C0), C0, C0), int(matrix_dim)
    return ll

def caculate_matmul_shape(matrix_A_dim, matrix_G_dim, split_dim):
    """get matmul shape"""
    split_dimA = split_dim
    split_dimG = split_dim
    if matrix_A_dim % split_dim == 0:
        batch_w = matrix_A_dim // split_dim
    else:
        if matrix_A_dim < split_dim:
            batch_w = 1
            split_dimA = matrix_A_dim
        else:
            batch_w = matrix_A_dim // split_dim + 1

    if matrix_G_dim % split_dim == 0:
        batch_h = matrix_G_dim // split_dim
    else:
        if matrix_G_dim < split_dim:
            batch_h = 1
            split_dimG = matrix_G_dim
        else:
            batch_h = matrix_G_dim // split_dim + 1
    matrix_A_shape = (batch_h, batch_w, split_dimA, split_dimA)
    matrix_G_shape = (batch_h, split_dimG, split_dimG)
    return matrix_A_shape, matrix_G_shape

def find_net_layertype_recur(net, layertype_map):
    """get net layer type recursively."""
    cells = net.name_cells()
    for name in cells:
        subcell = cells[name]
        print("thor subcell name: ", name)
        if subcell == net:
            continue
        elif isinstance(subcell, Conv2d_Thor):
            layertype_map.append(Conv)
        elif isinstance(subcell, Dense_Thor):
            layertype_map.append(FC)
        elif isinstance(subcell, Embedding_Thor):
            layertype_map.append(Embedding)
        elif isinstance(subcell, nn.LayerNorm):
            layertype_map.append(LayerNorm)
        elif isinstance(subcell, nn.BatchNorm2d):
            layertype_map.append(BatchNorm)
        elif isinstance(subcell, (nn.Conv2d, nn.Dense, nn.Embedding, nn.Conv2dTranspose, nn.Conv1d, nn.Conv1dTranspose,
                                  nn.BatchNorm1d, nn.GroupNorm, nn.GlobalBatchNorm)):
            layertype_map.append(Other)
        else:
            find_net_layertype_recur(subcell, layertype_map)

def get_net_layertype_mask(net):
    layertype_map = []
    find_net_layertype_recur(net, layertype_map)
    return layertype_map

def get_layer_counter(layer_type, layer_counter, params, idx):
    """get layer counter"""
    if layer_type in [Conv, FC, LayerNorm, BatchNorm]:
        if layer_type in [LayerNorm, BatchNorm]:
            if "beta" in params[idx].name.lower():
                layer_counter = layer_counter + 1
        else:
            if "bias" in params[idx].name.lower():
                layer_counter = layer_counter + 1
            else:
                if idx < len(params) - 1 and "bias" not in params[idx + 1].name.lower():
                    layer_counter = layer_counter + 1
    else:
        layer_counter = layer_counter + 1
    return layer_counter

def THOR(net, learning_rate, damping, momentum, weight_decay=0.0, loss_scale=1.0, batch_size=32,
         use_nesterov=False, decay_filter=lambda x: x.name not in [], split_indices=None):
    context.set_context(max_call_depth=10000)
    ConvertNetUntils().convert_to_thor_net(net)
    if context.get_context("device_target") == "Ascend":
        return THOR_Ascend(net, learning_rate, damping, momentum, weight_decay, loss_scale, batch_size, decay_filter,
                           split_indices=split_indices)
    return THOR_GPU(net, learning_rate, damping, momentum, weight_decay, loss_scale, batch_size,
                    use_nesterov, decay_filter, split_indices=split_indices)

class THOR_GPU(Optimizer):
    """
    THOR_GPU
    """
    def __init__(self, net, learning_rate, damping, momentum, weight_decay=0.0, loss_scale=1.0, batch_size=32,
                 use_nesterov=False, decay_filter=lambda x: x.name not in [], split_indices=None):
        params = filter(lambda x: x.requires_grad, net.get_parameters())
        super(THOR_GPU, self).__init__(learning_rate, params, weight_decay, loss_scale)
        Validator.check_value_type("momentum", momentum, [float], self.cls_name)
        if isinstance(momentum, float) and momentum < 0.0:
            raise ValueError("momentum should be at least 0.0, but got momentum {}".format(momentum))
        self.momentum = Parameter(Tensor(momentum, mstype.float32), name="momentum")
        self.params = self.parameters
        self.use_nesterov = Validator.check_bool(use_nesterov)
        self.moments = self.params.clone(prefix="moments", init='zeros')
        self.hyper_map = C.HyperMap()
        self.opt = P.ApplyMomentum(use_nesterov=self.use_nesterov)
        self.net = net
        self.matrix_A_cov = ParameterTuple(filter(lambda x: 'matrix_A' in x.name, net.get_parameters()))
        self.matrix_G_cov = ParameterTuple(filter(lambda x: 'matrix_G' in x.name, net.get_parameters()))
        self.A_normalizer = ParameterTuple(filter(lambda x: 'A_normalizer' in x.name, net.get_parameters()))
        self.G_normalizer = ParameterTuple(filter(lambda x: 'G_normalizer' in x.name, net.get_parameters()))
        self.transpose = P.Transpose()
        self.shape = P.Shape()
        self.reshape = P.Reshape()
        self.matmul = P.MatMul()
        self.assign = P.Assign()
        self.mul = P.Mul()
        self.damping = damping
        self.gather = P.GatherV2()
        self.one = Tensor(1, mstype.int32)
        self.batch_size = Tensor(batch_size, mstype.float32)
        self.loss_scale = Tensor(1 / (loss_scale * loss_scale), mstype.float32)
        self.batch_size_scale = Tensor(batch_size * batch_size, mstype.float32)
        self.feature_map = Tensor(1.0, mstype.float32)
        self.axis = 0
        self.cov_step = Parameter(initializer(0, [1], mstype.int32), name="cov_step", requires_grad=False)
        self.cast = P.Cast()
        self.sqrt = P.Sqrt()
        self.eye = P.Eye()
        split_dim = 128
        self.embedding_cholesky = P.CholeskyTrsm()
        self.cholesky = P.CholeskyTrsm(split_dim=split_dim)
        self.vector_matmul = P.BatchMatMul(transpose_a=True)
        self.reduce_sum = P.ReduceSum(keep_dims=False)
        self.inv = P.Reciprocal()
        self.square = P.Square()
        self.expand = P.ExpandDims()
        self.thor = True

        self.matrix_A = ()
        self.matrix_G = ()
        self.matrix_A_shape = ()
        self.thor_layer_count = 0
        self.conv_layer_count = 0
        self.weight_fim_idx_map = ()
        self.weight_conv_idx_map = ()
        self.weight_layerType_idx_map = ()
        layer_type_map = get_net_layertype_mask(net)

        layer_counter = 0
        for idx in range(len(self.params)):
            layer_type = layer_type_map[layer_counter]
            weight = self.params[idx]
            weight_shape = self.shape(weight)
            if layer_type in [Conv, FC] and "bias" not in self.params[idx].name.lower():
                in_channels = weight_shape[1]
                out_channels = weight_shape[0]
                matrix_A_dim = in_channels
                if layer_type == Conv:
                    matrix_A_dim = in_channels * weight_shape[2] * weight_shape[3]
                matrix_G_dim = out_channels
                matrix_A_shape, matrix_G_shape = caculate_matmul_shape(matrix_A_dim, matrix_G_dim, split_dim)
                matrix_A_inv = Parameter(np.zeros(matrix_A_shape).astype(np.float32),
                                         name='matrix_A_inv_' + str(self.thor_layer_count), requires_grad=False)
                matrix_G_inv = Parameter(np.zeros(matrix_G_shape).astype(np.float32),
                                         name="matrix_G_inv_" + str(self.thor_layer_count), requires_grad=False)
                self.matrix_A = self.matrix_A + (matrix_A_inv,)
                self.matrix_G = self.matrix_G + (matrix_G_inv,)
                self.matrix_A_shape = self.matrix_A_shape + (matrix_A_shape,)
            elif layer_type == Embedding:
                vocab_size = weight_shape[0]
                embedding_size = weight_shape[1]
                matrix_A_inv = Parameter(Tensor(np.zeros([vocab_size]).astype(np.float32)),
                                         name='matrix_A_inv_' + str(self.thor_layer_count), requires_grad=False)
                matrix_G_inv = Parameter(Tensor(np.zeros([embedding_size, embedding_size]).astype(np.float32)),
                                         name="matrix_G_inv_" + str(self.thor_layer_count), requires_grad=False)
                self.matrix_A = self.matrix_A + (matrix_A_inv,)
                self.matrix_G = self.matrix_G + (matrix_G_inv,)
                self.matrix_A_shape = self.matrix_A_shape + ((vocab_size,),)

            if layer_type in [Conv, FC, Embedding] and "bias" not in self.params[idx].name.lower():
                self.weight_fim_idx_map = self.weight_fim_idx_map + (self.thor_layer_count,)
                self.weight_layerType_idx_map = self.weight_layerType_idx_map + (layer_type,)
                self.thor_layer_count = self.thor_layer_count + 1
                if layer_type == Conv:
                    self.weight_conv_idx_map = self.weight_conv_idx_map + (self.conv_layer_count,)
                    self.conv_layer_count = self.conv_layer_count + 1
                else:
                    self.weight_conv_idx_map = self.weight_conv_idx_map + (-1,)
            else:
                self.weight_fim_idx_map = self.weight_fim_idx_map + (-1,)
                self.weight_conv_idx_map = self.weight_conv_idx_map + (-1,)
                if layer_type == LayerNorm:
                    self.weight_layerType_idx_map = self.weight_layerType_idx_map + (LayerNorm,)
                else:
                    self.weight_layerType_idx_map = self.weight_layerType_idx_map + (Other,)
                # bert.cls1.output_bias: not a network layer, only a trainable param
            if "output_bias" not in self.params[idx].name.lower():
                layer_counter = get_layer_counter(layer_type, layer_counter, self.params, idx)

        self.matrix_A = ParameterTuple(self.matrix_A)
        self.matrix_G = ParameterTuple(self.matrix_G)
        self.weight_decay = weight_decay
        self.decay_flags = tuple(decay_filter(x) for x in self.parameters)
        self.update_gradient = P.UpdateThorGradient(split_dim=split_dim)

        self.parallel_mode = context.get_auto_parallel_context("parallel_mode")
        self.is_distributed = (self.parallel_mode != ParallelMode.STAND_ALONE)
        if self.is_distributed:
            mean = _get_gradients_mean()
            degree = _get_device_num()
            if self.conv_layer_count > 0:
                if not split_indices:
                    self.split_indices = split_indices
                else:
                    self.split_indices = [len(self.matrix_A) - 1]
                auto_parallel_context().set_all_reduce_fusion_split_indices(self.split_indices, "hccl_world_groupsum2")
                auto_parallel_context().set_all_reduce_fusion_split_indices(self.split_indices, "hccl_world_groupsum4")
                auto_parallel_context().set_all_reduce_fusion_split_indices(self.split_indices, "hccl_world_groupsum6")
                auto_parallel_context().set_all_reduce_fusion_split_indices(self.split_indices, "hccl_world_groupsum8")
                self.grad_reducer_Amax = DistributedGradReducer(self.matrix_A, mean, degree, fusion_type=2)
                self.grad_reducer_Gmax = DistributedGradReducer(self.matrix_A, mean, degree, fusion_type=4)
                self.grad_reducer_A = DistributedGradReducer(self.matrix_A, mean, degree, fusion_type=6)
                self.grad_reducer_G = DistributedGradReducer(self.matrix_A, mean, degree, fusion_type=8)
            else:
                if not split_indices:
                    self.split_indices = split_indices
                else:
                    self.split_indices = [len(self.params) - 1]
                auto_parallel_context().set_all_reduce_fusion_split_indices(self.split_indices, "hccl_world_groupsum3")
                self.grad_reducer_g = DistributedGradReducer(self.params, mean, degree, fusion_type=3)

    def _get_Ainv_Ginv_list(self, gradients, damping_step, matrix_a_allreduce, matrix_g_allreduce):
        """get matrixA inverse list and matrix G inverse list"""
        for i in range(len(self.params)):
            thor_layer_count = self.weight_fim_idx_map[i]
            conv_layer_count = self.weight_conv_idx_map[i]
            layer_type = self.weight_layerType_idx_map[i]
            if layer_type in [Conv, FC, Embedding]:
                g = gradients[i]
                matrix_A = self.matrix_A_cov[thor_layer_count]
                matrix_G = self.matrix_G_cov[thor_layer_count]
                matrix_A = F.depend(matrix_A, g)
                matrix_G = F.depend(matrix_G, g)
                dampingA = damping_step
                dampingG = damping_step
                feature_map = self.feature_map
                if layer_type == Conv:
                    A_normalizer = self.A_normalizer[conv_layer_count]
                    G_normalizer = self.G_normalizer[conv_layer_count]
                    A_normalizer = F.depend(A_normalizer, g)
                    G_normalizer = F.depend(G_normalizer, g)
                    dampingA = self.mul(damping_step, 1.0 / A_normalizer)
                    dampingG = self.mul(damping_step, 1.0 / G_normalizer)
                    feature_map = self.sqrt(1.0 / A_normalizer)
                A_shape = self.shape(matrix_A)
                A_eye = self.eye(A_shape[0], A_shape[0], mstype.float32)
                dampingA = self.sqrt(dampingA)
                dampingG = self.sqrt(dampingG)
                G_shape = self.shape(matrix_G)
                G_eye = self.eye(G_shape[0], G_shape[1], mstype.float32)
                matrix_G = self.mul(matrix_G, self.loss_scale)
                matrix_G = self.mul(matrix_G, self.batch_size_scale)
                matrix_G = matrix_G + dampingG * G_eye
                if layer_type == Embedding:
                    A_eye = P.OnesLike()(matrix_A)
                    matrix_A = self.mul(matrix_A, 1.0 / self.batch_size)
                    matrix_A = matrix_A + dampingA * A_eye
                    matrix_A = self.inv(matrix_A)
                    matrix_G = self.embedding_cholesky(matrix_G)
                    matrix_G = self.matmul(matrix_G, matrix_G)
                else:
                    matrix_A = matrix_A + dampingA * A_eye
                    matrix_A = self.cholesky(matrix_A)
                    matrix_A = self.vector_matmul(matrix_A, matrix_A)
                    matrix_A = P.BroadcastTo(self.matrix_A_shape[thor_layer_count])(matrix_A)
                    matrix_G = self.cholesky(matrix_G)
                    matrix_G = self.vector_matmul(matrix_G, matrix_G)
                matrix_A = self.mul(matrix_A, feature_map)
                matrix_G = self.mul(matrix_G, feature_map)
                matrix_a_allreduce = matrix_a_allreduce + (matrix_A,)
                matrix_g_allreduce = matrix_g_allreduce + (matrix_G,)
        return matrix_a_allreduce, matrix_g_allreduce

    def construct(self, gradients):
        params = self.params
        moments = self.moments
        gradients = self.scale_grad(gradients)
        damping_step = self.gather(self.damping, self.cov_step, self.axis)
        damping_step = self.cast(damping_step, mstype.float32)
        new_grads = ()
        if self.thor:
            matrix_Ainv_list = ()
            matrix_Ginv_list = ()
            matrix_A_allreduce, matrix_G_allreduce = self._get_Ainv_Ginv_list(gradients, damping_step,
                                                                              matrix_Ainv_list, matrix_Ginv_list)
            if self.is_distributed and self.conv_layer_count > 0:
                matrix_A_allreduce = self.grad_reducer_A(matrix_A_allreduce)
                matrix_G_allreduce = self.grad_reducer_G(matrix_G_allreduce)

            for i in range(len(self.params)):
                g = gradients[i]
                thor_layer_count = self.weight_fim_idx_map[i]
                conv_layer_count = self.weight_conv_idx_map[i]
                layer_type = self.weight_layerType_idx_map[i]
                if layer_type in [Conv, FC]:
                    g_shape = self.shape(g)
                    g = self.reshape(g, (g_shape[0], -1))
                    matrix_A = matrix_A_allreduce[thor_layer_count]
                    matrix_G = matrix_G_allreduce[thor_layer_count]
                    g = self.update_gradient(matrix_G, g, matrix_A)
                    fake_A = self.assign(self.matrix_A[thor_layer_count], matrix_A)
                    fake_G = self.assign(self.matrix_G[thor_layer_count], matrix_G)
                    g = F.depend(g, fake_A)
                    g = F.depend(g, fake_G)
                    if conv_layer_count != -1:
                        g = self.reshape(g, g_shape)
                elif layer_type == Embedding:
                    matrix_A = matrix_A_allreduce[thor_layer_count]
                    matrix_G = matrix_G_allreduce[thor_layer_count]
                    fake_A = self.assign(self.matrix_A[thor_layer_count], matrix_A)
                    fake_G = self.assign(self.matrix_G[thor_layer_count], matrix_G)
                    g = F.depend(g, fake_A)
                    g = F.depend(g, fake_G)
                    temp_a = self.expand(matrix_A, 1)
                    g = self.mul(temp_a, g)
                    g = self.matmul(g, matrix_G)
                elif layer_type == LayerNorm:
                    damping = self.sqrt(damping_step)
                    normalizer = self.batch_size
                    normalizer = self.cast(normalizer, mstype.float32)
                    fim_cov = self.square(g)
                    fim_cov = self.mul(fim_cov, 1.0 / normalizer)
                    fim_cov = fim_cov + damping
                    fim_inv = self.inv(fim_cov)
                    g = self.mul(fim_inv, g)
                new_grads = new_grads + (g,)
        else:
            for j in range(len(self.params)):
                g = gradients[j]
                thor_layer_count = self.weight_fim_idx_map[j]
                conv_layer_count = self.weight_conv_idx_map[j]
                layer_type = self.weight_layerType_idx_map[j]
                if layer_type in [Conv, FC]:
                    g_shape = self.shape(g)
                    g = self.reshape(g, (g_shape[0], -1))
                    matrix_A = self.matrix_A[thor_layer_count]
                    matrix_G = self.matrix_G[thor_layer_count]
                    g = self.update_gradient(matrix_G, g, matrix_A)
                    if conv_layer_count != -1:
                        g = self.reshape(g, g_shape)
                elif layer_type == Embedding:
                    matrix_A = self.matrix_A[thor_layer_count]
                    matrix_G = self.matrix_G[thor_layer_count]
                    g = gradients[j]
                    temp_a = self.expand(matrix_A, 1)
                    g = self.mul(temp_a, g)
                    g = self.matmul(g, matrix_G)
                elif layer_type == LayerNorm:
                    damping = self.sqrt(damping_step)
                    normalizer = self.batch_size
                    normalizer = self.cast(normalizer, mstype.float32)
                    fim_cov = self.square(g)
                    fim_cov = self.mul(fim_cov, 1.0 / normalizer)
                    fim_cov = fim_cov + damping
                    fim_inv = self.inv(fim_cov)
                    g = self.mul(fim_inv, g)
                new_grads = new_grads + (g,)
        gradients = new_grads

        if self.is_distributed and self.conv_layer_count == 0:
            gradients = self.grad_reducer_g(gradients)
        self.cov_step = self.cov_step + self.one
        if self.weight_decay > 0:
            gradients = self.hyper_map(F.partial(apply_decay, self.weight_decay), self.decay_flags, params, gradients)
        lr = self.get_lr()
        success = self.hyper_map(F.partial(_momentum_opt, self.opt, self.momentum, lr), gradients, params, moments)
        return success

class THOR_Ascend(Optimizer):
    """THOR"""

    def __init__(self, net, learning_rate, damping, momentum, weight_decay=0.0, loss_scale=1.0, batch_size=32,
                 decay_filter=lambda x: x.name not in [], split_indices=None):
        params = filter(lambda x: x.requires_grad, net.get_parameters())
        super(THOR_Ascend, self).__init__(learning_rate, params, weight_decay, loss_scale)
        if isinstance(momentum, float) and momentum < 0.0:
            raise ValueError("momentum should be at least 0.0, but got momentum {}".format(momentum))
        self.momentum = Parameter(Tensor(momentum, mstype.float32), name="momentum")
        self.params = self.parameters
        self.moments = self.params.clone(prefix="moments", init='zeros')
        self.hyper_map = C.HyperMap()
        self.opt = P.ApplyMomentum()
        self.net = net
        self.matrix_A_cov = ParameterTuple(filter(lambda x: 'matrix_A' in x.name, net.get_parameters()))
        self.matrix_G_cov = ParameterTuple(filter(lambda x: 'matrix_G' in x.name, net.get_parameters()))
        self.A_normalizer = ParameterTuple(filter(lambda x: 'A_normalizer' in x.name, net.get_parameters()))
        self.G_normalizer = ParameterTuple(filter(lambda x: 'G_normalizer' in x.name, net.get_parameters()))
        self.cube_matmul_left = P.CusMatMulCubeFraczLeftCast()
        self.cube_matmul_left_fc = P.CusMatMulCubeDenseLeft()
        self.cube_matmul_right_fc = P.CusMatMulCubeDenseRight()
        self.cube_matmul_right_mul = P.CusMatMulCubeFraczRightMul()
        self.transpose = P.Transpose()
        self.shape = P.Shape()
        self.reshape = P.Reshape()
        self.mul = P.Mul()

        self.C0 = 16
        self.matrix_A_dim = ()
        self.padA_flag = ()
        self.device_shape_pad_flag = ()
        self.diag_block_dim = 128
        self.matrix_A = ()
        self.matrix_G = ()
        print("matrix_A_cov len is", len(self.matrix_A_cov))
        self.thor_layer_count = 0
        self.conv_layer_count = 0
        self.weight_fim_idx_map = ()
        self.weight_conv_idx_map = ()
        self.weight_layerType_idx_map = ()
        layer_type_map = get_net_layertype_mask(net)
        layer_counter = 0
        for idx in range(len(self.params)):
            layer_type = layer_type_map[layer_counter]
            weight = self.params[idx]
            weight_shape = self.shape(weight)
            if layer_type == Conv and "bias" not in self.params[idx].name.lower():
                in_channels = weight_shape[1]
                out_channels = weight_shape[0]
                matrix_A_dim = in_channels * weight_shape[2] * weight_shape[3]
                matrix_G_dim = out_channels
                matrix_A_device_shape, matrix_A_device_dim = caculate_device_shape(matrix_A_dim, in_channels, True)
                matrix_G_device_shape, matrix_G_device_dim = caculate_device_shape(matrix_G_dim, in_channels, False)
                matrix_A_inv = Parameter(
                    Tensor(np.reshape(np.identity(matrix_A_device_dim).astype(np.float16), matrix_A_device_shape)),
                    name='matrix_A_inv_' + str(self.thor_layer_count), requires_grad=False)
                matrix_G_inv = Parameter(
                    Tensor(np.reshape(np.identity(matrix_G_device_dim).astype(np.float16), matrix_G_device_shape)),
                    name="matrix_G_inv_" + str(self.thor_layer_count), requires_grad=False)
                self.matrix_A = self.matrix_A + (matrix_A_inv,)
                self.matrix_G = self.matrix_G + (matrix_G_inv,)
                self.matrix_A_dim = self.matrix_A_dim + (matrix_A_dim,)
                padA_flag = False
                if (matrix_A_dim // self.diag_block_dim) * self.diag_block_dim != matrix_A_dim \
                    and matrix_A_dim > self.diag_block_dim:
                    padA_flag = True
                self.padA_flag = self.padA_flag + (padA_flag,)
                device_shape_pad_flag = False
                if matrix_A_dim != matrix_A_device_dim:
                    device_shape_pad_flag = True
                self.device_shape_pad_flag = self.device_shape_pad_flag + (device_shape_pad_flag,)
            elif layer_type == FC and "bias" not in self.params[idx].name.lower():
                out_channels = weight_shape[0]
                if out_channels == 1001:
                    fc_matrix_A = Parameter(Tensor(np.zeros([128, 128, 16, 16]).astype(np.float16)),
                                            name='matrix_A_inv_' + str(self.thor_layer_count),
                                            requires_grad=False)
                    fc_matrix_G = Parameter(Tensor(np.zeros([63, 63, 16, 16]).astype(np.float16)),
                                            name="matrix_G_inv_" + str(self.thor_layer_count),
                                            requires_grad=False)
                    self.matrix_A = self.matrix_A + (fc_matrix_A,)
                    self.matrix_G = self.matrix_G + (fc_matrix_G,)

            if layer_type in [Conv, FC, Embedding] and "bias" not in self.params[idx].name.lower():
                self.weight_fim_idx_map = self.weight_fim_idx_map + (self.thor_layer_count,)
                self.weight_layerType_idx_map = self.weight_layerType_idx_map + (layer_type,)
                self.thor_layer_count = self.thor_layer_count + 1
                if layer_type == Conv:
                    self.weight_conv_idx_map = self.weight_conv_idx_map + (self.conv_layer_count,)
                    self.conv_layer_count = self.conv_layer_count + 1
                else:
                    self.weight_conv_idx_map = self.weight_conv_idx_map + (-1,)
            else:
                self.weight_fim_idx_map = self.weight_fim_idx_map + (-1,)
                self.weight_conv_idx_map = self.weight_conv_idx_map + (-1,)
                if layer_type == LayerNorm:
                    self.weight_layerType_idx_map = self.weight_layerType_idx_map + (LayerNorm,)
                else:
                    self.weight_layerType_idx_map = self.weight_layerType_idx_map + (Other,)
            # bert.cls1.output_bias: not a network layer, only a trainable param
            if "output_bias" not in self.params[idx].name.lower():
                layer_counter = get_layer_counter(layer_type, layer_counter, self.params, idx)

        self.matrix_A = ParameterTuple(self.matrix_A)
        self.matrix_G = ParameterTuple(self.matrix_G)
        self.matrix_max_inv = ()
        for i in range(len(self.matrix_A)):
            self.matrix_max_inv = self.matrix_max_inv + (
                Parameter(initializer(1, [1], mstype.float32), name="matrix_max" + str(i), requires_grad=False),)
        self.log = P.Log()
        self.exp = P.Exp()
        self.sqrt = P.Sqrt()
        self.matrix_max_inv = ParameterTuple(self.matrix_max_inv)
        self.assign = P.Assign()
        self.cast = P.Cast()
        self.thor = True
        self.weight_decay = weight_decay * loss_scale
        self.decay_flags = tuple(decay_filter(x) for x in self.parameters)
        self.damping = damping
        self.gather = P.GatherV2()
        self.one = Tensor(1, mstype.int32)
        self.batch_size = Tensor(batch_size, mstype.float32)
        self.loss_scale = Tensor(1 / (loss_scale * loss_scale), mstype.float32)
        self.batch_size_scale = Tensor(batch_size * batch_size, mstype.float32)
        self.axis = 0
        self.cov_step = Parameter(initializer(0, [1], mstype.int32), name="cov_step", requires_grad=False)
        self.cast = P.Cast()
        self.eye = P.Eye()
        self.cholesky = P.CusCholeskyTrsm()
        self.vector_matmul = P.CusBatchMatMul()
        self.fused_abs_max2 = P.CusFusedAbsMax1()
        self.matrix_combine = P.CusMatrixCombine()
        self.slice = P.Slice()
        self.expand = P.ExpandDims()
        self.reduce_sum = P.ReduceSum(keep_dims=False)
        self.square = P.Square()
        self.inv = P.Inv()
        self.matmul = P.MatMul()

        self.parallel_mode = context.get_auto_parallel_context("parallel_mode")
        self.is_distributed = (self.parallel_mode != ParallelMode.STAND_ALONE)
        if self.is_distributed:
            mean = _get_gradients_mean()
            degree = _get_device_num()
            if self.conv_layer_count > 0:
                if not split_indices:
                    self.split_indices = split_indices
                else:
                    self.split_indices = [len(self.matrix_A) - 1]
                auto_parallel_context().set_all_reduce_fusion_split_indices(self.split_indices, "hccl_world_groupsum2")
                auto_parallel_context().set_all_reduce_fusion_split_indices(self.split_indices, "hccl_world_groupsum4")
                auto_parallel_context().set_all_reduce_fusion_split_indices(self.split_indices, "hccl_world_groupsum6")
                auto_parallel_context().set_all_reduce_fusion_split_indices(self.split_indices, "hccl_world_groupsum8")
                self.grad_reducer_Amax = DistributedGradReducer(self.matrix_A, mean, degree, fusion_type=2)
                self.grad_reducer_Gmax = DistributedGradReducer(self.matrix_A, mean, degree, fusion_type=4)
                self.grad_reducer_A = DistributedGradReducer(self.matrix_A, mean, degree, fusion_type=6)
                self.grad_reducer_G = DistributedGradReducer(self.matrix_A, mean, degree, fusion_type=8)
            else:
                if not split_indices:
                    self.split_indices = split_indices
                else:
                    self.split_indices = [len(self.params) - 1]
                auto_parallel_context().set_all_reduce_fusion_split_indices(self.split_indices, "hccl_world_groupsum3")
                self.grad_reducer_g = DistributedGradReducer(self.params, mean, degree, fusion_type=3)

    def _get_Ainv_Ginv_Amax_Gmax_list(self, gradients, damping_step, matrix_a_allreduce, matrix_g_allreduce,
                                      matrix_a_max_allreduce, matrix_g_max_allreduce):
        """get matrixA inverse list, matrixG inverse list, matrixA_max list, matrixG_max list"""
        for i in range(len(self.params)):
            thor_layer_count = self.weight_fim_idx_map[i]
            conv_layer_count = self.weight_conv_idx_map[i]
            layer_type = self.weight_layerType_idx_map[i]
            if layer_type in [Conv, FC, Embedding]:
                g = gradients[i]
                matrix_A = self.matrix_A_cov[thor_layer_count]
                matrix_G = self.matrix_G_cov[thor_layer_count]
                matrix_A = F.depend(matrix_A, g)
                matrix_G = F.depend(matrix_G, g)
                A_shape = self.shape(matrix_A)
                A_eye = self.eye(A_shape[0], A_shape[0], mstype.float32)
                G_shape = self.shape(matrix_G)
                G_eye = self.eye(G_shape[0], G_shape[0], mstype.float32)
                if layer_type == Conv:
                    A_normalizer = self.A_normalizer[conv_layer_count]
                    G_normalizer = self.G_normalizer[conv_layer_count]
                    A_normalizer = F.depend(A_normalizer, g)
                    G_normalizer = F.depend(G_normalizer, g)
                    dampingA = self.mul(damping_step, self.batch_size / A_normalizer)
                    dampingG = self.mul(damping_step, self.batch_size / G_normalizer)
                    dampingA = self.sqrt(dampingA)
                    matrix_A = matrix_A + dampingA * A_eye
                    matrix_A_inv = self.cholesky(matrix_A)
                    matrix_A_inv = self.vector_matmul(matrix_A_inv, matrix_A_inv)
                    A_max = P.CusFusedAbsMax1([self.matrix_A_dim[conv_layer_count],
                                               self.matrix_A_dim[conv_layer_count]])(matrix_A_inv)
                    A_max = self.fused_abs_max2(A_max)
                    matrix_A_inv = self.matrix_combine(matrix_A_inv)
                    if self.padA_flag[conv_layer_count]:
                        matrix_A_inv = self.slice(matrix_A_inv, (0, 0), (self.matrix_A_dim[conv_layer_count],
                                                                         self.matrix_A_dim[conv_layer_count]))
                    if self.device_shape_pad_flag[conv_layer_count]:
                        weight = self.params[i]
                        weight_shape = self.shape(weight)
                        kernel_hw = weight_shape[2] * weight_shape[3]
                        in_channels = weight_shape[1]
                        matrix_A_inv = self.reshape(matrix_A_inv, (kernel_hw, in_channels, kernel_hw, in_channels))
                        matrix_A_inv = P.Pad(((0, 0), (0, self.C0 - in_channels), (0, 0),
                                              (0, self.C0 - in_channels)))(matrix_A_inv)
                    matrix_A_inv_shape = self.shape(self.matrix_A[thor_layer_count])
                    matrix_A_device_temp_shape = (matrix_A_inv_shape[0], matrix_A_inv_shape[2],
                                                  matrix_A_inv_shape[1], matrix_A_inv_shape[3])
                    matrix_A_inv = self.reshape(matrix_A_inv, matrix_A_device_temp_shape)
                    matrix_A_inv = self.transpose(matrix_A_inv, (2, 0, 1, 3))

                    dampingG = self.sqrt(dampingG)
                    matrix_G = self.mul(matrix_G, self.loss_scale)
                    matrix_G = self.mul(matrix_G, self.batch_size_scale)
                    matrix_G = matrix_G + dampingG * G_eye
                    matrix_G_inv = self.cholesky(matrix_G)
                    matrix_G_inv = self.vector_matmul(matrix_G_inv, matrix_G_inv)
                    G_max = self.fused_abs_max2(matrix_G_inv)
                    G_max = self.fused_abs_max2(G_max)
                    matrix_G_inv = self.matrix_combine(matrix_G_inv)
                    matrix_G_inv_shape = self.shape(self.matrix_G[thor_layer_count])
                    matrix_G_device_temp_shape = (matrix_G_inv_shape[0], matrix_G_inv_shape[2],
                                                  matrix_G_inv_shape[1], matrix_G_inv_shape[3])
                    matrix_G_inv = self.reshape(matrix_G_inv, matrix_G_device_temp_shape)
                    matrix_G_inv = self.transpose(matrix_G_inv, (2, 0, 1, 3))

                    A_max = F.depend(A_max, g)
                    G_max = F.depend(G_max, g)
                    matrix_a_allreduce = matrix_a_allreduce + (matrix_A_inv,)
                    matrix_g_allreduce = matrix_g_allreduce + (matrix_G_inv,)
                    matrix_a_max_allreduce = matrix_a_max_allreduce + (A_max,)
                    matrix_g_max_allreduce = matrix_g_max_allreduce + (G_max,)
                elif layer_type == FC:
                    damping = self.sqrt(damping_step)
                    matrix_A = matrix_A + damping * A_eye
                    matrix_A_inv = self.cholesky(matrix_A)
                    matrix_A_inv = self.vector_matmul(matrix_A_inv, matrix_A_inv)
                    weight_shape = self.shape(self.params[i])
                    out_channels = weight_shape[0]
                    if out_channels == 2:
                        matrix_A_inv = self.matrix_combine(matrix_A_inv)
                        matrix_G_inv = G_eye
                    else:
                        matrix_G = self.mul(matrix_G, self.loss_scale)
                        matrix_G = self.mul(matrix_G, self.batch_size_scale)
                        matrix_G = matrix_G + damping * G_eye
                        matrix_G_inv = self.cholesky(matrix_G)
                        matrix_G_inv = self.vector_matmul(matrix_G_inv, matrix_G_inv)
                        if out_channels == 1001:
                            matrix_A_inv_max = self.fused_abs_max2(matrix_A_inv)
                            A_max = self.fused_abs_max2(matrix_A_inv_max)
                            matrix_A_inv = self.matrix_combine(matrix_A_inv)
                            matrix_A_inv_shape = self.shape(matrix_A_inv)
                            matrix_A_inv = self.reshape(matrix_A_inv,
                                                        (matrix_A_inv_shape[0] / 16, 16,
                                                         matrix_A_inv_shape[0] / 16, 16))
                            matrix_A_inv = self.transpose(matrix_A_inv, (2, 0, 1, 3))
                            matrix_G_inv_max = P.CusFusedAbsMax1([1001, 1001])(matrix_G_inv)
                            G_max = self.fused_abs_max2(matrix_G_inv_max)
                            matrix_G_inv = self.matrix_combine(matrix_G_inv)
                            matrix_G_inv = self.slice(matrix_G_inv, (0, 0), (1001, 1001))
                            matrix_G_inv = P.Pad(((0, 7), (0, 7)))(matrix_G_inv)
                            matrix_G_inv_shape = self.shape(matrix_G_inv)
                            matrix_G_inv = self.reshape(matrix_G_inv,
                                                        (matrix_G_inv_shape[0] / 16, 16,
                                                         matrix_G_inv_shape[0] / 16, 16))
                            matrix_G_inv = self.transpose(matrix_G_inv, (2, 0, 1, 3))
                            A_max = F.depend(A_max, g)
                            G_max = F.depend(G_max, g)
                            matrix_a_max_allreduce = matrix_a_max_allreduce + (A_max,)
                            matrix_g_max_allreduce = matrix_g_max_allreduce + (G_max,)
                        else:
                            matrix_A_inv = self.matrix_combine(matrix_A_inv)
                            matrix_G_inv = self.matrix_combine(matrix_G_inv)
                    matrix_a_allreduce = matrix_a_allreduce + (matrix_A_inv,)
                    matrix_g_allreduce = matrix_g_allreduce + (matrix_G_inv,)
                elif layer_type == Embedding:
                    damping = self.sqrt(damping_step)
                    A_eye = P.OnesLike()(matrix_A)
                    matrix_A = self.mul(matrix_A, 1.0 / self.batch_size)
                    matrix_A = matrix_A + damping * A_eye
                    matrix_A_inv = self.inv(matrix_A)
                    matrix_G = self.mul(matrix_G, self.loss_scale)
                    matrix_G = self.mul(matrix_G, self.batch_size_scale)
                    matrix_G = matrix_G + damping * G_eye
                    matrix_G_inv = self.cholesky(matrix_G)
                    matrix_G_inv = self.vector_matmul(matrix_G_inv, matrix_G_inv)
                    matrix_G_inv = self.matrix_combine(matrix_G_inv)
                    matrix_a_allreduce = matrix_a_allreduce + (matrix_A_inv,)
                    matrix_g_allreduce = matrix_g_allreduce + (matrix_G_inv,)
        return matrix_a_allreduce, matrix_g_allreduce, matrix_a_max_allreduce, matrix_g_max_allreduce

    def _process_layernorm(self, damping_step, gradient):
        """process layernorm layer for thor"""
        damping = self.sqrt(damping_step)
        normalizer = self.cast(self.batch_size, mstype.float32)
        fim_cov = self.square(gradient)
        fim_cov = self.mul(fim_cov, 1.0 / normalizer)
        fim_cov = fim_cov + damping
        fim_inv = self.inv(fim_cov)
        gradient = self.mul(fim_inv, gradient)
        return gradient

    def _get_second_gradients(self, new_grads, damping_step, gradients):
        """get second gradients for thor"""
        params_len = len(self.params)
        for i in range(params_len):
            g = gradients[i]
            thor_layer_count = self.weight_fim_idx_map[i]
            layer_type = self.weight_layerType_idx_map[i]
            if self.conv_layer_count > 0:
                matrix_A = self.matrix_A[thor_layer_count]
                matrix_G = self.matrix_G[thor_layer_count]
                matrix_max = self.matrix_max_inv[thor_layer_count]
                if layer_type == FC:
                    g = self.cube_matmul_left_fc(matrix_G, g)
                    g = self.cube_matmul_right_fc(g, matrix_A, matrix_max)
                elif layer_type == Conv:
                    g = self.cube_matmul_left(matrix_G, g)
                    g = self.cube_matmul_right_mul(g, matrix_A, matrix_max)
            else:
                if layer_type == Embedding:
                    temp_a_ori = self.matrix_A_cov[thor_layer_count]
                    temp_g = self.matrix_G_cov[thor_layer_count]
                    temp_a = self.expand(temp_a_ori, 1)
                    g = self.mul(temp_a, g)
                    temp_g = self.cast(temp_g, mstype.float16)
                    g = self.cast(g, mstype.float16)
                    g = self.matmul(g, temp_g)
                    g = self.cast(g, mstype.float32)
                elif layer_type == FC:
                    temp_a = self.matrix_A_cov[thor_layer_count]
                    temp_g = self.matrix_G_cov[thor_layer_count]
                    temp_a = self.cast(temp_a, mstype.float16)
                    temp_g = self.cast(temp_g, mstype.float16)
                    g = self.cast(g, mstype.float16)
                    g = self.matmul(temp_g, g)
                    g = self.matmul(g, temp_a)
                    g = self.cast(g, mstype.float32)
                elif layer_type == LayerNorm:
                    g = self._process_layernorm(damping_step, g)
            new_grads = new_grads + (g,)
        return new_grads

    def construct(self, gradients):
        params = self.params
        moments = self.moments
        damping_step = self.gather(self.damping, self.cov_step, self.axis)
        damping_step = self.cast(damping_step, mstype.float32)
        if self.thor:
            matrix_A_allreduce = ()
            matrix_G_allreduce = ()
            matrix_A_max_allreduce = ()
            matrix_G_max_allreduce = ()
            matrix_A_allreduce, matrix_G_allreduce, matrix_A_max_allreduce, matrix_G_max_allreduce = \
                self._get_Ainv_Ginv_Amax_Gmax_list(gradients, damping_step, matrix_A_allreduce, matrix_G_allreduce,
                                                   matrix_A_max_allreduce, matrix_G_max_allreduce)
            if self.is_distributed and self.conv_layer_count > 0:
                matrix_A_allreduce = self.grad_reducer_A(matrix_A_allreduce)
                matrix_G_allreduce = self.grad_reducer_G(matrix_G_allreduce)
                matrix_A_max_allreduce = self.grad_reducer_Amax(matrix_A_max_allreduce)
                matrix_G_max_allreduce = self.grad_reducer_Gmax(matrix_G_max_allreduce)

            new_grads = ()
            for i in range(len(self.params)):
                g = gradients[i]
                thor_layer_count = self.weight_fim_idx_map[i]
                conv_layer_count = self.weight_conv_idx_map[i]
                layer_type = self.weight_layerType_idx_map[i]
                if self.conv_layer_count > 0:
                    temp_a = matrix_A_allreduce[thor_layer_count]
                    temp_g = matrix_G_allreduce[thor_layer_count]
                    matrix_A_inv_max = self.log(matrix_A_max_allreduce[thor_layer_count])
                    matrix_A_inv_max = self.mul(matrix_A_inv_max, -1)
                    matrix_A_inv_max = self.exp(matrix_A_inv_max)
                    temp_a = self.mul(temp_a, matrix_A_inv_max)
                    matrix_G_inv_max = self.log(matrix_G_max_allreduce[thor_layer_count])
                    matrix_G_inv_max = self.mul(matrix_G_inv_max, -1)
                    matrix_G_inv_max = self.exp(matrix_G_inv_max)
                    temp_g = self.mul(temp_g, matrix_G_inv_max)
                    temp_max = self.mul(matrix_A_max_allreduce[thor_layer_count],
                                        matrix_G_max_allreduce[thor_layer_count])
                    temp_a = self.cast(temp_a, mstype.float16)
                    temp_g = self.cast(temp_g, mstype.float16)
                    if layer_type == FC:
                        g = self.cube_matmul_left_fc(temp_g, g)
                        g = self.cube_matmul_right_fc(g, temp_a, temp_max)
                    elif layer_type == Conv:
                        A_normalizer = self.A_normalizer[conv_layer_count]
                        A_normalizer = F.depend(A_normalizer, g)
                        temp_max = self.mul(temp_max, self.batch_size / A_normalizer)
                        g = self.cube_matmul_left(temp_g, g)
                        g = self.cube_matmul_right_mul(g, temp_a, temp_max)
                    fake_A = self.assign(self.matrix_A[thor_layer_count], temp_a)
                    fake_G = self.assign(self.matrix_G[thor_layer_count], temp_g)
                    fake_max = self.assign(self.matrix_max_inv[thor_layer_count], temp_max)
                    g = F.depend(g, fake_A)
                    g = F.depend(g, fake_G)
                    g = F.depend(g, fake_max)
                else:
                    if layer_type == Embedding:
                        temp_a_ori = matrix_A_allreduce[thor_layer_count]
                        temp_g = matrix_G_allreduce[thor_layer_count]
                        fake_A = self.assign(self.matrix_A_cov[thor_layer_count], temp_a_ori)
                        fake_G = self.assign(self.matrix_G_cov[thor_layer_count], temp_g)
                        g = F.depend(g, fake_A)
                        g = F.depend(g, fake_G)
                        temp_a = self.expand(temp_a_ori, 1)
                        g = self.mul(temp_a, g)
                        temp_g = self.cast(temp_g, mstype.float16)
                        g = self.cast(g, mstype.float16)
                        g = self.matmul(g, temp_g)
                        g = self.cast(g, mstype.float32)
                    elif layer_type == FC:
                        temp_a = matrix_A_allreduce[thor_layer_count]
                        temp_g = matrix_G_allreduce[thor_layer_count]
                        fake_A = self.assign(self.matrix_A_cov[thor_layer_count], temp_a)
                        fake_G = self.assign(self.matrix_G_cov[thor_layer_count], temp_g)
                        g = F.depend(g, fake_A)
                        g = F.depend(g, fake_G)
                        temp_a = self.cast(temp_a, mstype.float16)
                        temp_g = self.cast(temp_g, mstype.float16)
                        g = self.cast(g, mstype.float16)
                        g = self.matmul(temp_g, g)
                        g = self.matmul(g, temp_a)
                        g = self.cast(g, mstype.float32)
                    elif layer_type == LayerNorm:
                        g = self._process_layernorm(damping_step, g)
                new_grads = new_grads + (g,)
            gradients = new_grads
        else:
            new_grads = ()
            gradients = self._get_second_gradients(new_grads, damping_step, gradients)

        if self.is_distributed and self.conv_layer_count == 0:
            gradients = self.grad_reducer_g(gradients)
        self.cov_step = self.cov_step + self.one
        if self.weight_decay > 0:
            gradients = self.hyper_map(F.partial(apply_decay, self.weight_decay), self.decay_flags, params, gradients)
        gradients = self.scale_grad(gradients)
        lr = self.get_lr()
        success = self.hyper_map(F.partial(_momentum_opt, self.opt, self.momentum, lr), gradients, params, moments)
        return success


# linflat.py - FlatSystem subclass for linear systems
# RMM, 10 November 2012
#
# This file defines a FlatSystem class for a linear system.
#
# Copyright (c) 2012 by California Institute of Technology
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#
# 1. Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#
# 3. Neither the name of the California Institute of Technology nor
#    the names of its contributors may be used to endorse or promote
#    products derived from this software without specific prior
#    written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
# FOR A PARTICULAR PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL CALTECH
# OR THE CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF
# USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT
# OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
# SUCH DAMAGE.

import numpy as np
import control
from .flatsys import FlatSystem
from ..iosys import LinearIOSystem

class LinearFlatSystem(FlatSystem, LinearIOSystem):
    """Base class for a linear, differentially flat system.

    This class is used to create a differentially flat system representation
    from a linear system.

    Parameters
    ----------
    linsys : StateSpace
        LTI StateSpace system to be converted
    inputs : int, list of str or None, optional
        Description of the system inputs.  This can be given as an integer
        count or as a list of strings that name the individual signals.
        If an integer count is specified, the names of the signal will be
        of the form `s[i]` (where `s` is one of `u`, `y`, or `x`).  If
        this parameter is not given or given as `None`, the relevant
        quantity will be determined when possible based on other
        information provided to functions using the system.
    outputs : int, list of str or None, optional
        Description of the system outputs.  Same format as `inputs`.
    states : int, list of str, or None, optional
        Description of the system states.  Same format as `inputs`.
    dt : None, True or float, optional
        System timebase.  None (default) indicates continuous
        time, True indicates discrete time with undefined sampling
        time, positive number is discrete time with specified
        sampling time.
    params : dict, optional
        Parameter values for the systems.  Passed to the evaluation
        functions for the system as default values, overriding internal
        defaults.
    name : string, optional
        System name (used for specifying signals)

    """

    def __init__(self, linsys, inputs=None, outputs=None, states=None,
                 name=None):
        """Define a flat system from a SISO LTI system.

        Given a reachable, single-input/single-output, linear time-invariant
        system, create a differentially flat system representation.

        """
        # Make sure we can handle the system
        if (not control.isctime(linsys)):
            raise control.ControlNotImplemented(
                "requires continuous time, linear control system")
        elif (not control.issiso(linsys)):
            raise control.ControlNotImplemented(
                "only single input, single output systems are supported")

        # Initialize the object as a LinearIO system
        LinearIOSystem.__init__(
            self, linsys, inputs=inputs, outputs=outputs, states=states,
            name=name)

        # Find the transformation to chain of integrators form
        # Note: store all array as ndarray, not matrix
        zsys, Tr = control.reachable_form(linsys)
        Tr = np.array(Tr[::-1, ::])     # flip rows

        # Extract the information that we need
        self.F = np.array(zsys.A[0, ::-1])      # input function coeffs
        self.T = Tr                             # state space transformation
        self.Tinv = np.linalg.inv(Tr)           # compute inverse once

        # Compute the flat output variable z = C x
        Cfz = np.zeros(np.shape(linsys.C)); Cfz[0, 0] = 1
        self.Cf = Cfz @ Tr

    # Compute the flat flag from the state (and input)
    def forward(self, x, u):
        """Compute the flat flag given the states and input.

        See :func:`control.flatsys.FlatSystem.forward` for more info.

        """
        x = np.reshape(x, (-1, 1))
        u = np.reshape(u, (1, -1))
        zflag = [np.zeros(self.nstates + 1)]
        zflag[0][0] = self.Cf @ x
        H = self.Cf                     # initial state transformation
        for i in range(1, self.nstates + 1):
            zflag[0][i] = H @ (self.A @ x + self.B @ u)
            H = H @ self.A       # derivative for next iteration
        return zflag

    # Compute state and input from flat flag
    def reverse(self, zflag):
        """Compute the states and input given the flat flag.

        See :func:`control.flatsys.FlatSystem.reverse` for more info.

        """
        z = zflag[0][0:-1]
        x = self.Tinv @ z
        u = zflag[0][-1] - self.F @ z
        return np.reshape(x, self.nstates), np.reshape(u, self.ninputs)


import os
import time
import logging
import torch
import torch.nn.functional as F
import torch.backends.cudnn as cudnn
import numpy as np
import nibabel as nib
import scipy.misc
import pdb

cudnn.benchmark = True

path = os.path.dirname(__file__)

# dice socre is equal to f1 score
def dice_score(o, t, eps=1e-8):
    num = 2 * (o * t).sum() + eps  #
    den = o.sum() + t.sum() + eps  # eps
    # print(o.sum(),t.sum(),num,den)
    print('All_voxels:240*240*155 | numerator:{} | denominator:{} | pred_voxels:{} | GT_voxels:{}'.format(int(num),
                                                                                int(den),
                                                                                o.sum(),
                                                                                int(t.sum())))
    return num / den

def softmax_output_dice(output, target):
    ret = []

    # whole
    o = output > 0;
    t = target > 0  # ce
    ret += dice_score(o, t),
    # core
    o = (output == 1) | (output == 3)
    t = (target == 1) | (target == 4)
    ret += dice_score(o, t),
    # active
    o = (output == 3);
    t = (target == 4)
    ret += dice_score(o, t),

    return ret

keys = 'whole', 'core', 'enhancing', 'loss'

def validate_softmax(
        valid_loader,
        model,
        cfg='',
        savepath='',  # when in validation set, you must specify the path to save the 'nii' segmentation results here
        names=None,  # The names of the patients orderly!
        scoring=True,  # If true, print the dice score.
        verbose=False,
        use_TTA=False,  # Test time augmentation, False as default!
        save_format=None,  # ['nii','npy'], use 'nii' as default. Its purpose is for submission.
        snapshot=False,  # for visualization. Default false. It is recommended to generate the visualized figures.
        postprocess=False,  # Defualt False, when use postprocess, the score of dice_ET would be changed.
        cpu_only=False):
    assert cfg is not None
    H, W, T = 240, 240, 155
    model.eval()
    runtimes = []
    vals = AverageMeter()
    for i, data in enumerate(valid_loader):
        target_cpu = data[1][0, :H, :W,
                     :T].numpy() if scoring else None  # when validing, make sure that argument 'scoring' must be false, else it raise a error!

        if cpu_only == False:
            data = [t.cuda(non_blocking=True) for t in data]
        x, target = data[:2]
        # pdb.set_trace()
        # compute output
        if not use_TTA:

            # torch.cuda.synchronize()
            start_time = time.time()
            logit = model(x)
            # torch.cuda.synchronize()
            elapsed_time = time.time() - start_time
            runtimes.append(elapsed_time)

            output = F.softmax(logit, dim=1)
        else:
            logit = F.softmax(model(x), 1)  # 000
            logit += F.softmax(model(x.flip(dims=(2,))).flip(dims=(2,)), 1)
            logit += F.softmax(model(x.flip(dims=(3,))).flip(dims=(3,)), 1)
            logit += F.softmax(model(x.flip(dims=(4,))).flip(dims=(4,)), 1)
            logit += F.softmax(model(x.flip(dims=(2, 3))).flip(dims=(2, 3)), 1)
            logit += F.softmax(model(x.flip(dims=(2, 4))).flip(dims=(2, 4)), 1)
            logit += F.softmax(model(x.flip(dims=(3, 4))).flip(dims=(3, 4)), 1)
            logit += F.softmax(model(x.flip(dims=(2, 3, 4))).flip(dims=(2, 3, 4)), 1)
            output = logit / 8.0  # mean
        # pdb.set_trace()
        output = output[:, :, :H, :W, :T].cpu().numpy()

        ############
        # todo 0->1
        output = output.argmax(1)  # (channels,height,width,depth)

        if postprocess == True:
            for j in range(3):
                ET_voxels = (output[j, :, :, :] == 3).sum()
                if ET_voxels < 500:
                    # zhuyi zai output hou jia 'J'
                    output[j][np.where(output[j, :, :, :] == 3)] = 1

                # todo

                TC_voxels = (output[j, :, :, :] == 1).sum()
                if TC_voxels < 100:
                    output[j][np.where(output[j, :, :, :] == 2)] = 1

        msg = 'Subject {}/{}, '.format(i + 1, len(valid_loader))
        name = str(i)
        if names:
            # todo
            name = [names[3 * i], names[3 * i + 1], names[3 * i + 2]]
            # todo
            msg += '{:>20}, {:>20},{:>20},'.format(name[0], name[1], name[2])

        if savepath:
            # .npy for farthur model ensemble
            # .nii for directly model submission
            assert save_format in ['npy', 'nii']
            if save_format == 'npy':
                np.save(os.path.join(savepath, name + '_preds'), output)
            if save_format == 'nii':
                for i in range(3):
                    oname = os.path.join(savepath, 'submission', name[i] + '.nii.gz')
                    seg_img = np.zeros(shape=(H, W, T), dtype=np.uint8)

                    seg_img[np.where(output[i] == 1)] = 1
                    seg_img[np.where(output[i] == 2)] = 2
                    seg_img[np.where(output[i] == 3)] = 4
                    if verbose:
                        # todo
                        # print('1:',np.sum(seg_img==1),' | 2:',np.sum(seg_img==2),' | 4:',np.sum(seg_img==4))
                        # print('WT:',np.sum((seg_img==1)|(seg_img==2)|(seg_img==4)),' | TC:',np.sum((seg_img==1)|(seg_img==4)),' | ET:',np.sum(seg_img==4))
                        # todo
                        logging.info(
                            ('1:', np.sum(seg_img == 1), ' | 2:', np.sum(seg_img == 2), ' | 4:', np.sum(seg_img == 4)))
                        logging.info(('WT:', np.sum((seg_img == 1) | (seg_img == 2) | (seg_img == 4)), ' | TC:',
                                      np.sum((seg_img == 1) | (seg_img == 4)), ' | ET:', np.sum(seg_img == 4)))
                        nib.save(nib.Nifti1Image(seg_img, None), oname)
                logging.info(msg)
                if snapshot:
                    """ --- grey figure---"""
                    # Snapshot_img = np.zeros(shape=(H,W,T),dtype=np.uint8)
                    # Snapshot_img[np.where(output[1,:,:,:]==1)] = 64
                    # Snapshot_img[np.where(output[2,:,:,:]==1)] = 160
                    # Snapshot_img[np.where(output[3,:,:,:]==1)] = 255
                    """ --- colorful figure--- """
                    Snapshot_img = np.zeros(shape=(H, W, 3, T), dtype=np.uint8)
                    Snapshot_img[:, :, 0, :][np.where(output == 1)] = 255
                    Snapshot_img[:, :, 1, :][np.where(output == 2)] = 255
                    Snapshot_img[:, :, 2, :][np.where(output == 3)] = 255

                    for frame in range(T):
                        os.makedirs(os.path.join(savepath, 'snapshot', name), exist_ok=True)
                        scipy.misc.imsave(os.path.join(savepath, 'snapshot', name, str(frame) + '.png'),
                                          Snapshot_img[:, :, :, frame])

        if scoring:
            scores = softmax_output_dice(output, target_cpu)
            vals.update(np.array(scores))
            msg += ', '.join(['{}: {:.4f}'.format(k, v) for k, v in zip(keys, scores)])

            if snapshot:
                # red: (255,0,0) green:(0,255,0) blue:(0,0,255) 1 for NCR & NET, 2 for ED, 4 for ET, and 0 for everything else.
                gap_width = 2  # boundary width = 2
                Snapshot_img = np.zeros(shape=(H, W * 2 + gap_width, 3, T), dtype=np.uint8)
                Snapshot_img[:, W:W + gap_width, :] = 255  # white boundary

                empty_fig = np.zeros(shape=(H, W, T), dtype=np.uint8)
                empty_fig[np.where(output == 1)] = 255
                Snapshot_img[:, :W, 0, :] = empty_fig
                empty_fig = np.zeros(shape=(H, W, T), dtype=np.uint8)
                empty_fig[np.where(target_cpu == 1)] = 255
                Snapshot_img[:, W + gap_width:, 0, :] = empty_fig

                empty_fig = np.zeros(shape=(H, W, T), dtype=np.uint8)
                empty_fig[np.where(output == 2)] = 255
                Snapshot_img[:, :W, 1, :] = empty_fig
                empty_fig = np.zeros(shape=(H, W, T), dtype=np.uint8)
                empty_fig[np.where(target_cpu == 2)] = 255
                Snapshot_img[:, W + gap_width:, 1, :] = empty_fig

                empty_fig = np.zeros(shape=(H, W, T), dtype=np.uint8)
                empty_fig[np.where(output == 3)] = 255
                Snapshot_img[:, :W, 2, :] = empty_fig
                empty_fig = np.zeros(shape=(H, W, T), dtype=np.uint8)
                empty_fig[np.where(target_cpu == 4)] = 255
                Snapshot_img[:, W + gap_width:, 2, :] = empty_fig

                for frame in range(T):
                    os.makedirs(os.path.join('snapshot', cfg, name), exist_ok=True)
                    scipy.misc.imsave(os.path.join('snapshot', cfg, name, str(frame) + '.png'),
                                      Snapshot_img[:, :, :, frame])

    if scoring:
        msg = 'Average scores:'
        msg += ', '.join(['{}: {:.4f}'.format(k, v) for k, v in zip(keys, vals.avg)])
        logging.info(msg)
    if len(runtimes):
        computational_runtime(runtimes)

    model.train()
    return vals.avg

def computational_runtime(runtimes):
    # remove the maximal value and minimal value
    runtimes = np.array(runtimes)
    maxvalue = np.max(runtimes)
    minvalue = np.min(runtimes)
    nums = runtimes.shape[0] - 2
    meanTime = (np.sum(runtimes) - maxvalue - minvalue) / nums
    fps = 1 / meanTime
    print('mean runtime:', meanTime, 'fps:', fps)

class AverageMeter(object):
    """Computes and stores the average and current value"""

    def __init__(self):
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count


#!/usr/bin/env python
# -*- coding: utf-8 -*-

# Copyright (c) 2015 IAS / CNRS / Univ. Paris-Sud
# BSD License - see attached LICENSE file
# Author: Alexandre Boucaud <alexandre.boucaud@ias.u-psud.fr>

"""
PyPHER - Python-based PSF Homogenization kERnels
================================================

Compute the homogenization kernel between two PSFs

Usage:
  pypher psf_source psf_target output
         [-s ANGLE_SOURCE] [-t ANGLE_TARGET] [-r REG_FACT]
  pypher (-h | --help)

Example:
  pypher psf_a.fits psf_b.fits kernel_a_to_b.fits -r 1.e-5
"""
from __future__ import absolute_import, print_function, division

import os
import sys
import logging
import logging.handlers
import argparse
import numpy as np

from scipy.ndimage import rotate, zoom

from . import fitsutils as fits
from .parser import ThrowingArgumentParser, ArgumentParserError

__version__ = '0.6.4'

def parse_args():
    """Argument parser for the command line interface of `pypher`"""
    parser = ThrowingArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
        prog='pypher',
        description="Compute the homogenization kernel between two PSFs")

    parser.add_argument('psf_source', type=str,
                        help="FITS file of PSF image with highest resolution")

    parser.add_argument('psf_target', type=str,
                        help="FITS file of PSF image with lowest resolution")

    parser.add_argument('output', type=str,
                        help="File name for the output kernel")

    parser.add_argument('-s', '--angle_source', type=float, default=0.0,
                        help="Rotation angle to apply to `psf_source` (deg)")

    parser.add_argument('-t', '--angle_target', type=float, default=0.0,
                        help="Rotation angle to apply to `psf_target` (deg)")

    parser.add_argument('-r', '--reg_fact', type=float, default=1.e-4,
                        help="Regularisation parameter for the Wiener filter")

    return parser.parse_args()

################
# IMAGE METHODS
################

def format_kernel_header(fits_file, args, pixel_scale):
    """
    Write the input parameters of pypher as comments in the header

    The kernel header therefore contains the name of the PSF files
    it has been created from.
    The pixel scale of the kernel is also written as a dedicated
    kernel key.

    Parameters
    ----------
    fits_file: str
        Path to the FITS kernel image
    args: `argparse.Namespace`
        Container for the parsed values
    pixel_scale: float
        Pixel scale of the kernel

    """
    fits.clear_comments(fits_file)

    pypher_comments = [
        '=' * 50, '',
        'File written with PyPHER',
        '------------------------', '',
        'Kernel from PSF', '',
        '=> {0}'.format(os.path.basename(args.psf_source)), '',
        'to PSF', '',
        '=> {0}'.format(os.path.basename(args.psf_target)), '',
        'using a regularisation parameter '
        'R = {0:1.1e}'.format(args.reg_fact), '',
        '=' * 50
    ]
    fits.add_comments(fits_file, pypher_comments)

    fits.write_pixelscale(fits_file, pixel_scale)

def imrotate(image, angle, interp_order=1):
    """
    Rotate an image from North to East given an angle in degrees

    Parameters
    ----------
    image : `numpy.ndarray`
        Input data array
    angle : float
        Angle in degrees
    interp_order : int, optional
        Spline interpolation order [0, 5] (default 1: linear)

    Returns
    -------
    output : `numpy.ndarray`
        Rotated data array

    """
    return rotate(image, -1.0 * angle,
                  order=interp_order, reshape=False, prefilter=False)

def imresample(image, source_pscale, target_pscale, interp_order=1):
    """
    Resample data array from one pixel scale to another

    The resampling ensures the parity of the image is conserved
    to preserve the centering.

    Parameters
    ----------
    image : `numpy.ndarray`
        Input data array
    source_pscale : float
        Pixel scale of ``image`` in arcseconds
    target_pscale : float
        Pixel scale of output array in arcseconds
    interp_order : int, optional
        Spline interpolation order [0, 5] (default 1: linear)

    Returns
    -------
    output : `numpy.ndarray`
        Resampled data array

    """
    old_size = image.shape[0]
    new_size_raw = old_size * source_pscale / target_pscale
    new_size = int(np.ceil(new_size_raw))

    if new_size > 10000:
        raise MemoryError("The resampling will yield a too large image. "
                          "Please resize the input PSF image.")

    # Chech for parity
    if (old_size - new_size) % 2 == 1:
        new_size += 1

    ratio = new_size / old_size

    return zoom(image, ratio, order=interp_order) / ratio**2

def trim(image, shape):
    """
    Trim image to a given shape

    Parameters
    ----------
    image: 2D `numpy.ndarray`
        Input image
    shape: tuple of int
        Desired output shape of the image

    Returns
    -------
    new_image: 2D `numpy.ndarray`
        Input image trimmed

    """
    shape = np.asarray(shape, dtype=int)
    imshape = np.asarray(image.shape, dtype=int)

    if np.alltrue(imshape == shape):
        return image

    if np.any(shape <= 0):
        raise ValueError("TRIM: null or negative shape given")

    dshape = imshape - shape
    if np.any(dshape < 0):
        raise ValueError("TRIM: target size bigger than source one")

    if np.any(dshape % 2 != 0):
        raise ValueError("TRIM: source and target shapes "
                         "have different parity")

    idx, idy = np.indices(shape)
    offx, offy = dshape // 2

    return image[idx + offx, idy + offy]

def zero_pad(image, shape, position='corner'):
    """
    Extends image to a certain size with zeros

    Parameters
    ----------
    image: real 2d `numpy.ndarray`
        Input image
    shape: tuple of int
        Desired output shape of the image
    position : str, optional
        The position of the input image in the output one:
            * 'corner'
                top-left corner (default)
            * 'center'
                centered

    Returns
    -------
    padded_img: real `numpy.ndarray`
        The zero-padded image

    """
    shape = np.asarray(shape, dtype=int)
    imshape = np.asarray(image.shape, dtype=int)

    if np.alltrue(imshape == shape):
        return image

    if np.any(shape <= 0):
        raise ValueError("ZERO_PAD: null or negative shape given")

    dshape = shape - imshape
    if np.any(dshape < 0):
        raise ValueError("ZERO_PAD: target size smaller than source one")

    pad_img = np.zeros(shape, dtype=image.dtype)

    idx, idy = np.indices(imshape)

    if position == 'center':
        if np.any(dshape % 2 != 0):
            raise ValueError("ZERO_PAD: source and target shapes "
                             "have different parity.")
        offx, offy = dshape // 2
    else:
        offx, offy = (0, 0)

    pad_img[idx + offx, idy + offy] = image

    return pad_img

##########
# FOURIER
##########

def udft2(image):
    """Unitary fft2"""
    norm = np.sqrt(image.size)
    return np.fft.fft2(image) / norm

def uidft2(image):
    """Unitary ifft2"""
    norm = np.sqrt(image.size)
    return np.fft.ifft2(image) * norm

def psf2otf(psf, shape):
    """
    Convert point-spread function to optical transfer function.

    Compute the Fast Fourier Transform (FFT) of the point-spread
    function (PSF) array and creates the optical transfer function (OTF)
    array that is not influenced by the PSF off-centering.
    By default, the OTF array is the same size as the PSF array.

    To ensure that the OTF is not altered due to PSF off-centering, PSF2OTF
    post-pads the PSF array (down or to the right) with zeros to match
    dimensions specified in OUTSIZE, then circularly shifts the values of
    the PSF array up (or to the left) until the central pixel reaches (1,1)
    position.

    Parameters
    ----------
    psf : `numpy.ndarray`
        PSF array
    shape : int
        Output shape of the OTF array

    Returns
    -------
    otf : `numpy.ndarray`
        OTF array

    Notes
    -----
    Adapted from MATLAB psf2otf function

    """
    if np.all(psf == 0):
        return np.zeros_like(psf)

    inshape = psf.shape
    # Pad the PSF to outsize
    psf = zero_pad(psf, shape, position='corner')

    # Circularly shift OTF so that the 'center' of the PSF is
    # [0,0] element of the array
    for axis, axis_size in enumerate(inshape):
        psf = np.roll(psf, -int(axis_size / 2), axis=axis)

    # Compute the OTF
    otf = np.fft.fft2(psf)

    # Estimate the rough number of operations involved in the FFT
    # and discard the PSF imaginary part if within roundoff error
    # roundoff error  = machine epsilon = sys.float_info.epsilon
    # or np.finfo().eps
    n_ops = np.sum(psf.size * np.log2(psf.shape))
    otf = np.real_if_close(otf, tol=n_ops)

    return otf

################
# DECONVOLUTION
################

LAPLACIAN = np.array([[ 0, -1,  0],
                      [-1,  4, -1],
                      [ 0, -1,  0]])

def deconv_wiener(psf, reg_fact):
    r"""
    Create a Wiener filter using a PSF image

    The signal is $\ell_2$ penalized by a 2D Laplacian operator that
    serves as a high-pass filter for the regularization process.
    The key to the process is to use optical transfer functions (OTF)
    instead of simple Fourier transform, since it ensures the phase
    of the psf is adequately placed.

    Parameters
    ----------
    psf: `numpy.ndarray`
        PSF array
    reg_fact: float
        Regularisation parameter for the Wiener filter

    Returns
    -------
    wiener: complex `numpy.ndarray`
        Fourier space Wiener filter

    """
    # Optical transfer functions
    trans_func = psf2otf(psf, psf.shape)
    reg_op = psf2otf(LAPLACIAN, psf.shape)

    wiener = np.conj(trans_func) / (np.abs(trans_func)**2 +
                                    reg_fact * np.abs(reg_op)**2)

    return wiener

def homogenization_kernel(psf_target, psf_source, reg_fact=1e-4, clip=True):
    r"""
    Compute the homogenization kernel to match two PSFs

    The deconvolution step is done using a Wiener filter with $\ell_2$
    penalization.
    The output is given both in Fourier and in the image domain to serve
    different purposes.

    Parameters
    ----------
    psf_target: `numpy.ndarray`
        2D array
    psf_source: `numpy.ndarray`
        2D array
    reg_fact: float, optional
        Regularisation parameter for the Wiener filter
    clip: bool, optional
        If `True`, enforces the non-amplification of the noise
        (default `True`)

    Returns
    -------
    kernel_image: `numpy.ndarray`
        2D deconvolved image
    kernel_fourier: `numpy.ndarray`
        2D discrete Fourier transform of deconvolved image

    """
    wiener = deconv_wiener(psf_source, reg_fact)

    kernel_fourier = wiener * udft2(psf_target)
    kernel_image = np.real(uidft2(kernel_fourier))

    if clip:
        kernel_image.clip(-1, 1)

    return kernel_image, kernel_fourier

########
# DEBUG
########

def setup_logger(log_filename='pypher.log'):  # pragma: no cover
    """
    Set up and return a logger

    The logger records the time, modulename, method and message

    Parameters
    ----------
    log_filename: str
        Name of the output logfile

    """
    # create logger
    logger = logging.getLogger('logger')
    logger.setLevel(logging.DEBUG)
    # Add the log message handler to the logger
    handler = logging.handlers.RotatingFileHandler(log_filename)
    # create formatter
    formatter = logging.Formatter('%(asctime)s - '
                                  '%(module)s - '
                                  '%(levelname)s - '
                                  '%(message)s')
    handler.setFormatter(formatter)
    # add handler to logger
    logger.addHandler(handler)

    return logger

#######
# MAIN
#######

def main():  # pragma: no cover
    """Main script for pypher"""
    try:
        args = parse_args()
    except ArgumentParserError:
        print(__doc__)
        sys.exit()

    kernel_basename, _ = os.path.splitext(args.output)
    kernel_fits = kernel_basename + '.fits'

    logname = '%s.log' % kernel_basename
    if os.path.exists(logname):
        os.remove(logname)
    log = setup_logger(logname)

    # Load images (NaNs are set to 0)
    psf_source = fits.getdata(args.psf_source)
    psf_target = fits.getdata(args.psf_target)

    log.info('Source PSF loaded: %s', args.psf_source)
    log.info('Target PSF loaded: %s', args.psf_target)

    # Set NaNs to 0.0
    psf_source = np.nan_to_num(psf_source)
    psf_target = np.nan_to_num(psf_target)

    # Retrieve the pixel scale of each image
    pixscale_source = fits.get_pixscale(args.psf_source)
    pixscale_target = fits.get_pixscale(args.psf_target)

    log.info('Source PSF pixel scale: %.2f arcsec', pixscale_source)
    log.info('Target PSF pixel scale: %.2f arcsec', pixscale_target)

    # Rotate images (if necessary)
    if args.angle_source != 0.0:
        psf_source = imrotate(psf_source, args.angle_source)
    if args.angle_target != 0.0:
        psf_target = imrotate(psf_target, args.angle_target)

    log.info('Source PSF rotated by %.2f degrees', args.angle_source)
    log.info('Target PSF rotated by %.2f degrees', args.angle_target)

    # Normalize the PSFs
    psf_source /= psf_source.sum()
    psf_target /= psf_target.sum()

    # Resample high resolution image to the low one
    if pixscale_source != pixscale_target:
        try:
            psf_source = imresample(psf_source,
                                    pixscale_source,
                                    pixscale_target)
        except MemoryError:
            log.error('- COMPUTATION ABORTED -')
            log.error('The size of the resampled PSF would have '
                      'exceeded 10K x 10K')
            log.error('Please resize your image and try again')

            print('Issue during the resampling step - see pypher.log')
            sys.exit()

        log.info('Source PSF resampled to the target pixel scale')

    # check the new size of the source vs. the target
    if psf_source.shape > psf_target.shape:
        psf_source = trim(psf_source, psf_target.shape)
    else:
        psf_source = zero_pad(psf_source, psf_target.shape, position='center')

    kernel, _ = homogenization_kernel(psf_target, psf_source,
                                      reg_fact=args.reg_fact)

    log.info('Kernel computed using Wiener filtering and a regularisation '
             'parameter r = %.2e', args.reg_fact)

    # Write kernel to FITS file
    fits.writeto(kernel_fits, data=kernel)
    format_kernel_header(kernel_fits, args, pixscale_target)

    log.info('Kernel saved in %s', kernel_fits)

    print("pypher: Output kernel saved to %s" % kernel_fits)

if __name__ == '__main__':
    main()