# Copyright (c) 2023-2024, NVIDIA CORPORATION.
# 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 warnings
import torch

import pylibcugraphops

import e3nn.o3 as o3
from e3nn.o3 import Irreps
from e3nn.o3._reduce import _wigner_nj

from typing import Iterator, List

import pylibcugraphops.operators as ops


def _transpose_irrep_to_m_last(tensor: torch.Tensor, irreps: Irreps):
    new_tensor = torch.empty_like(tensor)
    offset = 0
    for irr in irreps:
        num_channels_per_m = irr.mul
        m = irr.ir.dim
        tmp = tensor[:, offset : offset + irr.dim].clone()
        if m > 1:
            tmp = tmp.reshape(-1, m, num_channels_per_m)
            tmp = torch.transpose(tmp, 2, 1).reshape(-1, irr.dim)
        new_tensor[:, offset : offset + irr.dim] = tmp
        offset += irr.dim

    return new_tensor


def _transpose_irrep_to_channels_last(tensor: torch.Tensor, irreps: Irreps):
    new_tensor = torch.empty_like(tensor)
    offset = 0
    for irr in irreps:
        num_channels_per_m = irr.mul
        m = irr.ir.dim
        tmp = tensor[:, offset : offset + irr.dim].clone()
        if m > 1:
            tmp = tmp.reshape(-1, num_channels_per_m, m)
            tmp = torch.transpose(tmp, 2, 1).reshape(-1, irr.dim)
        new_tensor[:, offset : offset + irr.dim] = tmp
        offset += irr.dim

    return new_tensor


def _irreps_to_segment_info(irreps: Irreps):
    info = []
    offset = 0
    is_sh = True
    for irr in irreps:
        info.append([offset, irr.ir.dim, irr.mul, -1])
        offset += irr.dim
        is_sh = is_sh and (irr.mul == 1)

    if is_sh:
        return None
    return torch.IntTensor(info)


class _transpose_irrep_to_channels_last_autograd(torch.autograd.Function):
    @staticmethod
    def forward(ctx, tensor: torch.Tensor, segment_info: torch.Tensor):
        segment_info = segment_info.detach().contiguous()
        ctx.save_for_backward(segment_info)
        tensor_transpose = torch.empty_like(tensor)
        stream = torch.cuda.current_stream().cuda_stream

        if tensor.numel() > 0:
            ops.segmented_transpose(
                tensor_transpose.detach().contiguous(),
                tensor.detach().contiguous(),
                segment_info,
                True,
                stream,
            )
        return tensor_transpose

    @staticmethod
    def backward(ctx, grad_out: torch.Tensor):
        (segment_info,) = ctx.saved_tensors
        grad_tensor = _transpose_irrep_to_m_last_autograd.apply(grad_out, segment_info)
        return grad_tensor, None


class _transpose_irrep_to_m_last_autograd(torch.autograd.Function):
    @staticmethod
    def forward(ctx, tensor: torch.Tensor, segment_info: torch.Tensor):
        segment_info = segment_info.detach().contiguous()
        ctx.save_for_backward(segment_info)
        tensor_transpose = torch.empty_like(tensor)
        stream = torch.cuda.current_stream().cuda_stream

        if tensor.numel() > 0:
            ops.segmented_transpose(
                tensor_transpose.detach().contiguous(),
                tensor.detach().contiguous(),
                segment_info,
                False,
                stream,
            )
        return tensor_transpose

    @staticmethod
    def backward(ctx, grad_out: torch.Tensor):
        (segment_info,) = ctx.saved_tensors
        grad_tensor = _transpose_irrep_to_channels_last_autograd.apply(
            grad_out, segment_info
        )
        return grad_tensor, None


def transpose_irrep_to_channels_last(
    tensor: torch.Tensor, segment_info: torch.Tensor
) -> torch.Tensor:
    if segment_info is None:
        return tensor
    return _transpose_irrep_to_channels_last_autograd.apply(tensor, segment_info)


def transpose_irrep_to_m_last(
    tensor: torch.Tensor, segment_info: torch.Tensor
) -> torch.Tensor:
    if segment_info is None:
        return tensor
    return _transpose_irrep_to_m_last_autograd.apply(tensor, segment_info)


def sort_tp_info_for_operand(
    path_offsets_and_dims: torch.Tensor, path_cg_values: torch.Tensor, op_idx: int
):
    # idx = 0 -> in1
    # idx = 1 -> in2
    # idx = 2 -> weight
    # idx = 3 -> out
    key = path_offsets_and_dims[:, op_idx]
    perm = key.sort(descending=False, stable=True)[1]
    return path_offsets_and_dims[perm, :], path_cg_values[perm]


def path_offsets_to_csr(
    path_offsets_and_dims: torch.Tensor, op_idx: int, op_stride: int
):
    diff = torch.diff(
        path_offsets_and_dims[:, op_idx],
        n=1,
        append=path_offsets_and_dims.new_tensor([op_stride]),
    )
    path_csr_offsets = diff.nonzero().to(dtype=torch.int32).squeeze(dim=1)
    path_csr_offsets = torch.cat(
        [
            path_csr_offsets.new_tensor([0]),
            path_csr_offsets + 1,
        ]
    )
    return path_csr_offsets


def make_tp_info(
    op_idx: int,
    op_stride: int,
    path_offsets: torch.Tensor,
    path_cg_values: torch.Tensor,
):
    path_offsets = path_offsets.clone().detach()
    path_cg_values = path_cg_values.clone().detach()
    path_offsets, path_cg_values = sort_tp_info_for_operand(
        path_offsets, path_cg_values, op_idx
    )
    path_csr_offsets = path_offsets_to_csr(path_offsets, op_idx, op_stride)

    path_cg_values = path_cg_values
    path_offsets = path_offsets
    path_csr_offsets = path_csr_offsets

    return (
        path_csr_offsets,
        path_offsets,
        path_cg_values,
    )


def get_tp_infos_from_e3nn_instructions(
    instructions: Iterator[o3.Instruction],
    irreps_in1: o3.Irreps,
    irreps_in2: o3.Irreps,
    irreps_out: o3.Irreps,
):
    connection_mode = instructions[0].connection_mode

    path_offsets_and_dims = []
    path_cg_values = []

    slices_a = irreps_in1.slices()
    slices_b = irreps_in2.slices()
    slices_o = irreps_out.slices()

    off_weight = 0

    for inst in instructions:
        error_msg = "Only one connection mode at a time supported"
        error_msg += f"got {connection_mode} and {inst.connection_mode} respectively"
        if not inst.connection_mode == connection_mode:
            raise ValueError(error_msg)
        ir_idx_a = inst.i_in1
        ir_idx_b = inst.i_in2
        ir_idx_o = inst.i_out

        off_a = slices_a[ir_idx_a].start
        off_b = slices_b[ir_idx_b].start
        off_o = slices_o[ir_idx_o].start
        path_weight = inst.path_weight

        dim_u = inst.path_shape[0]
        dim_v = inst.path_shape[1] if len(inst.path_shape) > 1 else 0
        dim_w = inst.path_shape[2] if len(inst.path_shape) > 2 else 0

        if connection_mode == "uvw":
            dim_weight = dim_u * dim_v * dim_w
        elif connection_mode in {"uvu", "uvv"}:
            dim_weight = dim_u * dim_v
        elif connection_mode == "uuw":
            dim_w = dim_v
            dim_v = dim_u
            dim_weight = dim_u * dim_w
        else:
            dim_weight = 0

        # dim_a, dim_b, dim_o here actually are the number of channels
        dim_a = irreps_in1[ir_idx_a].mul
        dim_b = irreps_in2[ir_idx_b].mul
        dim_o = irreps_out[ir_idx_o].mul

        l_a = irreps_in1[ir_idx_a].ir.l
        l_b = irreps_in2[ir_idx_b].ir.l
        l_o = irreps_out[ir_idx_o].ir.l

        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            cg = o3.wigner_3j(l_a, l_b, l_o)

        for m_a in range(2 * l_a + 1):
            for m_b in range(2 * l_b + 1):
                for m_o in range(2 * l_o + 1):
                    cg_val = cg[m_a, m_b, m_o]
                    if cg_val != 0:
                        path_cg_values.append(cg_val * path_weight)
                        path_offsets_and_dims.append(
                            (
                                off_a + m_a * dim_a,
                                off_b + m_b * dim_b,
                                off_weight,
                                off_o + m_o * dim_o,
                                dim_u,
                                dim_v,
                                dim_w,
                                -1,  # place-holder / padding
                            )
                        )

        if inst.has_weight:
            off_weight += dim_weight

    strides = (
        irreps_in1.dim,
        irreps_in2.dim,
        off_weight,
        irreps_out.dim,
    )

    if connection_mode == "uvw" and all([ir.mul == 1 for ir in irreps_in1]):
        connection_mode = "1vw"
        pass
    elif connection_mode == "uvw" and all([ir.mul == 1 for ir in irreps_in2]):
        connection_mode = "u1w"
        pass
    connection_mode = getattr(pylibcugraphops.ConnectionMode, connection_mode.upper())

    path_cg_values = torch.FloatTensor(path_cg_values)
    path_offsets_and_dims = torch.IntTensor(path_offsets_and_dims)

    tp_info_fwd = make_tp_info(3, strides[3], path_offsets_and_dims, path_cg_values)
    tp_info_bwd_dgrad_in1 = make_tp_info(
        0, strides[0], path_offsets_and_dims, path_cg_values
    )
    tp_info_bwd_dgrad_in2 = make_tp_info(
        1, strides[1], path_offsets_and_dims, path_cg_values
    )
    tp_info_bwd_dgrad_w = make_tp_info(
        2, strides[2], path_offsets_and_dims, path_cg_values
    )

    return (
        tp_info_fwd,
        tp_info_bwd_dgrad_in1,
        tp_info_bwd_dgrad_in2,
        tp_info_bwd_dgrad_w,
        connection_mode,
        strides,
    )


def get_linear_infos_from_e3nn_instructions(
    instructions: Iterator[o3.Instruction],
    irreps_in: o3.Irreps,
    irreps_out: o3.Irreps,
):
    connection_mode = "uuvv"

    path_offsets_and_dims = []
    path_cg_values = []

    slices_i = irreps_in.slices()
    slices_o = irreps_out.slices()

    off_weight = 0

    for inst in instructions:
        ir_idx_i = inst.i_in
        ir_idx_o = inst.i_out

        off_i = slices_i[ir_idx_i].start
        off_o = slices_o[ir_idx_o].start
        path_weight = inst.path_weight

        dim_u = inst.path_shape[0]
        dim_v = inst.path_shape[1]

        ir_i = irreps_in[ir_idx_i].ir
        ir_o = irreps_out[ir_idx_o].ir
        l_i = ir_i.l
        l_o = ir_o.l
        assert l_i == l_o

        for m in range(2 * l_i + 1):
            path_cg_values.append(path_weight)
            path_offsets_and_dims.append(
                (
                    off_i + m * dim_u,
                    off_weight,
                    -1,
                    off_o + m * dim_v,
                    dim_u,
                    dim_v,
                    -1,
                    -1,  # place-holder / padding
                )
            )

        off_weight += dim_u * dim_v

    strides = (
        irreps_in.dim,
        off_weight,
        -1,
        irreps_out.dim,
    )

    connection_mode = getattr(pylibcugraphops.ConnectionMode, connection_mode.upper())

    path_cg_values = torch.FloatTensor(path_cg_values)
    path_offsets_and_dims = torch.IntTensor(path_offsets_and_dims)

    tp_info_fwd = make_tp_info(3, strides[3], path_offsets_and_dims, path_cg_values)
    tp_info_bwd_dgrad_in = make_tp_info(
        0, strides[0], path_offsets_and_dims, path_cg_values
    )
    tp_info_bwd_dgrad_w = make_tp_info(
        1, strides[1], path_offsets_and_dims, path_cg_values
    )

    return (
        tp_info_fwd,
        tp_info_bwd_dgrad_in,
        tp_info_bwd_dgrad_w,
        connection_mode,
        strides,
    )


# taken from https://github.com/ACEsuit/mace/blob/main/mace/tools/cg.py
def u_matrix_real(
    irreps_in: o3.Irreps,
    irreps_out: o3.Irreps,
    correlation: int,
    normalization: str = "component",
    filter_ir_mid=None,
    dtype=None,
):
    irrepss = [o3.Irreps(irreps_in)] * correlation
    if correlation == 4:
        filter_ir_mid = [
            (0, 1),
            (1, -1),
            (2, 1),
            (3, -1),
            (4, 1),
            (5, -1),
            (6, 1),
            (7, -1),
            (8, 1),
            (9, -1),
            (10, 1),
            (11, -1),
        ]
    wigners = _wigner_nj(
        *irrepss, normalization=normalization, filter_ir_mid=filter_ir_mid, dtype=dtype
    )
    current_ir = wigners[0][0]
    out = []
    stack = torch.tensor([])

    for ir, _, base_o3 in wigners:
        if ir in irreps_out and ir == current_ir:
            stack = torch.cat((stack, base_o3.squeeze().unsqueeze(-1)), dim=-1)
            last_ir = current_ir
        elif ir in irreps_out and ir != current_ir:
            if len(stack) != 0:
                out += [last_ir, stack]
            stack = base_o3.squeeze().unsqueeze(-1)
            current_ir, last_ir = ir, ir
        else:
            current_ir = ir
    out += [last_ir, stack]
    return out


def get_sym_cont_u_matrices(irreps_in, irreps_out, correlation):
    u_matrices = []
    coupling_irreps = o3.Irreps([irrep.ir for irrep in irreps_in])
    for ir in irreps_out:
        u_list = []
        for nu in range(0, correlation):
            u_matrix = u_matrix_real(
                irreps_in=coupling_irreps,
                irreps_out=ir,
                correlation=nu + 1,
            )[-1]
            if ir.ir.dim == 1:
                u_matrix = u_matrix.unsqueeze(0)
            u_list.append(u_matrix)
        u_matrices.append(u_list)

    return u_matrices


def init_sym_cont_weights(u_matrices, num_elements, num_features, correlation):
    weights = []
    for u_list in u_matrices:
        w_list = []
        for nu in range(0, correlation):
            num_params = u_list[nu].size(-1)
            w = torch.randn((num_elements, num_params, num_features)) / num_params
            w_list.append(w)
        weights.append(w_list)
    return weights


def _uniquify_index_tuple(tup):
    nnz, idx = torch.unique(torch.vstack(tup), dim=1, return_inverse=True)
    nnz_tup = tuple([t.squeeze(0) for t in torch.tensor_split(nnz, nnz.size(0))])
    return nnz_tup, nnz, idx


def _get_nonzero_indices(w):
    tmp = w.sum(dim=-1).sum(dim=-1).nonzero(as_tuple=True)
    return tmp, torch.vstack(tmp)


def contract_sym_cont_weights(
    weights: List[List[torch.Tensor]],
    u_matrices: List[List[torch.Tensor]],
    correlation: int,
):
    contracted_weights = []
    for nu in range(0, correlation):
        w_list = []
        for idx in range(len(u_matrices)):
            u = u_matrices[idx][nu]
            w = weights[idx][nu]
            u_shape = u.shape
            u = u.view(-1, u_shape[-1])
            tmp = torch.einsum("xb,abc->xca", u, w)
            tmp = tmp.view(*u_shape[:-1], tmp.size(-2), tmp.size(-1))
            w_list.append(tmp)
        tmp = torch.cat(w_list, dim=0)
        contracted_weights.append(tmp)

    return contracted_weights


def get_sym_contraction_path_offsets(
    weights, u_matrices, correlation, num_features, num_elements
):
    weight_indices = [None] * correlation
    tp_infos = [None] * correlation
    strides = [None] * correlation

    # contract weights as this makes selecting the indices easier
    with torch.no_grad():
        contracted_weights = contract_sym_cont_weights(weights, u_matrices, correlation)
    w_max = contracted_weights[-1]
    nnz_tup, _ = _get_nonzero_indices(w_max)

    for nu in range(correlation - 1, -1, -1):
        nnz_tup_next, _, rev_idx = _uniquify_index_tuple(nnz_tup[:-1])
        tmp = []

        if nu == correlation - 1:
            for n in range(nnz_tup[0].size(0)):
                tmp.append(
                    (
                        nnz_tup[-1][n].item() * num_features,
                        0,
                        n * num_features * num_elements,
                        rev_idx[n].item() * num_features,
                        num_features,
                        num_elements,
                        -1,
                        -1,
                    )
                )

            strides[nu] = (
                w_max.size(1) * num_features,
                num_elements,
                nnz_tup[0].size(0) * num_features * num_elements,
                nnz_tup_next[0].size(0) * num_features,
            )

            offset_and_dims = torch.IntTensor(tmp)
            cg_values = torch.ones(len(tmp), dtype=torch.float)

            tp_info_fwd = make_tp_info(3, strides[nu][3], offset_and_dims, cg_values)
            tp_info_bwd_dgrad_x = make_tp_info(
                0, strides[nu][0], offset_and_dims, cg_values
            )
            tp_info_bwd_dgrad_y = make_tp_info(
                1, strides[nu][1], offset_and_dims, cg_values
            )
            tp_info_bwd_dgrad_w = make_tp_info(
                2, strides[nu][2], offset_and_dims, cg_values
            )

            tp_infos[nu] = (
                tp_info_fwd,
                tp_info_bwd_dgrad_x,
                tp_info_bwd_dgrad_y,
                tp_info_bwd_dgrad_w,
            )

        else:
            for n in range(nnz_tup[0].size(0)):
                tmp.append(
                    (
                        nnz_tup[-1][n].item() * num_features,
                        n * num_features,
                        -1,
                        rev_idx[n].item() * num_features,
                        num_features,
                        -1,
                        -1,
                        -1,
                    )
                )

            strides[nu] = (
                w_max.size(1) * num_features,
                nnz_tup[0].size(0) * num_features,
                -1,
                nnz_tup_next[0].size(0) * num_features,
            )

            offset_and_dims = torch.IntTensor(tmp)
            cg_values = torch.ones(len(tmp), dtype=torch.float)

            tp_info_fwd = make_tp_info(3, strides[nu][3], offset_and_dims, cg_values)
            tp_info_bwd_dgrad_x = make_tp_info(
                0, strides[nu][0], offset_and_dims, cg_values
            )
            tp_info_bwd_dgrad_w = make_tp_info(
                1, strides[nu][1], offset_and_dims, cg_values
            )

            tp_infos[nu] = (tp_info_fwd, tp_info_bwd_dgrad_x, tp_info_bwd_dgrad_w)

        weight_indices[nu] = nnz_tup
        nnz_tup = nnz_tup_next

    return tp_infos, weight_indices, strides
