# SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES.
# SPDX-FileCopyrightText: All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# 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 math
import warnings

import torch
import torch.fft
import torch.nn as nn
import torch.nn.functional as F
from torch.cuda import amp
from torch.utils.checkpoint import checkpoint
from torch_harmonics import *  # noqa

from fcnv2_activations import ComplexReLU  # noqa
from fcnv2_contractions import (
    compl_contract2d_fwd,
    compl_contract2d_fwd_c,
    compl_contract_fwd,
    compl_contract_fwd_c,
    compl_mul2d_fwd,
    compl_mul2d_fwd_c,
    compl_muladd2d_fwd,
    compl_muladd2d_fwd_c,
    contract_tt,
)


def _no_grad_trunc_normal_(tensor, mean, std, a, b):
    # Cut & paste from PyTorch official master until it's in a few official releases - RW
    # Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf
    def norm_cdf(x):
        # Computes standard normal cumulative distribution function
        return (1.0 + math.erf(x / math.sqrt(2.0))) / 2.0

    if (mean < a - 2 * std) or (mean > b + 2 * std):
        warnings.warn(
            "mean is more than 2 std from [a, b] in nn.init.trunc_normal_. "
            "The distribution of values may be incorrect.",
            stacklevel=2,
        )

    with torch.no_grad():
        # Values are generated by using a truncated uniform distribution and
        # then using the inverse CDF for the normal distribution.
        # Get upper and lower cdf values
        l = norm_cdf((a - mean) / std)  # noqa
        u = norm_cdf((b - mean) / std)

        # Uniformly fill tensor with values from [l, u], then translate to
        # [2l-1, 2u-1].
        tensor.uniform_(2 * l - 1, 2 * u - 1)

        # Use inverse cdf transform for normal distribution to get truncated
        # standard normal
        tensor.erfinv_()

        # Transform to proper mean, std
        tensor.mul_(std * math.sqrt(2.0))
        tensor.add_(mean)

        # Clamp to ensure it's in the proper range
        tensor.clamp_(min=a, max=b)
        return tensor


def trunc_normal_(tensor, mean=0.0, std=1.0, a=-2.0, b=2.0):
    r"""Fills the input Tensor with values drawn from a truncated
    normal distribution. The values are effectively drawn from the
    normal distribution :math:`\mathcal{N}(\text{mean}, \text{std}^2)`
    with values outside :math:`[a, b]` redrawn until they are within
    the bounds. The method used for generating the random values works
    best when :math:`a \leq \text{mean} \leq b`.
    Args:
    tensor: an n-dimensional `torch.Tensor`
    mean: the mean of the normal distribution
    std: the standard deviation of the normal distribution
    a: the minimum cutoff value
    b: the maximum cutoff value
    Examples:
    >>> w = torch.empty(3, 5)
    >>> nn.init.trunc_normal_(w)
    """
    return _no_grad_trunc_normal_(tensor, mean, std, a, b)


@torch.jit.script
def drop_path(
    x: torch.Tensor, drop_prob: float = 0.0, training: bool = False
) -> torch.Tensor:
    """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
    This is the same as the DropConnect impl I created for EfficientNet, etc networks, however,
    the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper...
    See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for
    changing the layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use
    'survival rate' as the argument.
    """
    if drop_prob == 0.0 or not training:
        return x
    keep_prob = 1.0 - drop_prob
    shape = (x.shape[0],) + (1,) * (
        x.ndim - 1
    )  # work with diff dim tensors, not just 2d ConvNets
    random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device)
    random_tensor.floor_()  # binarize
    output = x.div(keep_prob) * random_tensor
    return output


class DropPath(nn.Module):
    """Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks)."""

    def __init__(self, drop_prob=None):
        super(DropPath, self).__init__()
        self.drop_prob = drop_prob

    def forward(self, x):
        return drop_path(x, self.drop_prob, self.training)


class PatchEmbed(nn.Module):
    def __init__(
        self, img_size=(224, 224), patch_size=(16, 16), in_chans=3, embed_dim=768
    ):
        super(PatchEmbed, self).__init__()
        num_patches = (img_size[1] // patch_size[1]) * (img_size[0] // patch_size[0])
        self.img_size = img_size
        self.patch_size = patch_size
        self.num_patches = num_patches
        self.proj = nn.Conv2d(
            in_chans, embed_dim, kernel_size=patch_size, stride=patch_size
        )

    def forward(self, x):
        # gather input
        B, C, H, W = x.shape
        assert (  # noqa
            H == self.img_size[0] and W == self.img_size[1]
        ), f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
        # new: B, C, H*W
        x = self.proj(x).flatten(2)
        return x


class MLP(nn.Module):
    def __init__(
        self,
        in_features,
        hidden_features=None,
        out_features=None,
        act_layer=nn.GELU,
        output_bias=True,
        drop_rate=0.0,
        checkpointing=False,
    ):
        super(MLP, self).__init__()
        self.checkpointing = checkpointing
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features

        fc1 = nn.Conv2d(in_features, hidden_features, 1, bias=True)
        act = act_layer()
        fc2 = nn.Conv2d(hidden_features, out_features, 1, bias=output_bias)
        if drop_rate > 0.0:
            drop = nn.Dropout(drop_rate)
            self.fwd = nn.Sequential(fc1, act, drop, fc2, drop)
        else:
            self.fwd = nn.Sequential(fc1, act, fc2)

    @torch.jit.ignore
    def checkpoint_forward(self, x):
        return checkpoint(self.fwd, x)

    def forward(self, x):
        if self.checkpointing:
            return self.checkpoint_forward(x)
        else:
            return self.fwd(x)


class RealFFT2(nn.Module):
    """
    Helper routine to wrap FFT similarly to the SHT
    """

    def __init__(self, nlat, nlon, lmax=None, mmax=None):
        super(RealFFT2, self).__init__()

        self.nlat = nlat
        self.nlon = nlon
        self.lmax = lmax or self.nlat
        self.mmax = mmax or self.nlon // 2 + 1
        self.num_batches = 1

        assert self.lmax % 2 == 0  # noqa

    def forward(self, x):
        # do batched FFT
        xs = torch.split(x, x.shape[1] // self.num_batches, dim=1)

        ys = []
        for xt in xs:
            yt = torch.fft.rfft2(xt, dim=(-2, -1), norm="ortho")
            ys.append(
                torch.cat(
                    (
                        yt[..., : math.ceil(self.lmax / 2), : self.mmax],
                        yt[..., -math.floor(self.lmax / 2) :, : self.mmax],
                    ),
                    dim=-2,
                )
            )

        # connect
        y = torch.cat(ys, dim=1).contiguous()

        # y = torch.fft.rfft2(x, dim=(-2, -1), norm="ortho")
        # y = torch.cat((y[..., :math.ceil(self.lmax/2), :self.mmax], y[..., -math.floor(self.lmax/2):, :self.mmax]), dim=-2)
        return y


class InverseRealFFT2(nn.Module):
    """
    Helper routine to wrap FFT similarly to the SHT
    """

    def __init__(self, nlat, nlon, lmax=None, mmax=None):
        super(InverseRealFFT2, self).__init__()

        self.nlat = nlat
        self.nlon = nlon
        self.lmax = lmax or self.nlat
        self.mmax = mmax or self.nlon // 2 + 1
        self.num_batches = 1

    def forward(self, x):
        # do batched FFT
        xs = torch.split(x, x.shape[1] // self.num_batches, dim=1)

        ys = []
        for xt in xs:
            ys.append(
                torch.fft.irfft2(
                    xt, dim=(-2, -1), s=(self.nlat, self.nlon), norm="ortho"
                )
            )
        out = torch.cat(ys, dim=1).contiguous()

        # out = torch.fft.irfft2(x, dim=(-2, -1), s=(self.nlat, self.nlon), norm="ortho")
        return out


class SpectralConv2d(nn.Module):
    """
    Spectral Convolution as utilized in
    """

    def __init__(
        self,
        forward_transform,
        inverse_transform,
        hidden_size,
        sparsity_threshold=0.0,
        hard_thresholding_fraction=1,
        use_complex_kernels=False,
        compression=None,
        rank=0,
        bias=False,
    ):
        super(SpectralConv2d, self).__init__()

        self.hidden_size = hidden_size
        self.sparsity_threshold = sparsity_threshold
        self.hard_thresholding_fraction = hard_thresholding_fraction
        self.scale = 1 / hidden_size**2
        self.contract_handle = (
            compl_contract2d_fwd_c if use_complex_kernels else compl_contract2d_fwd
        )

        self.forward_transform = forward_transform
        self.inverse_transform = inverse_transform

        self.output_dims = (self.inverse_transform.nlat, self.inverse_transform.nlon)
        modes_lat = self.inverse_transform.lmax
        modes_lon = self.inverse_transform.mmax
        self.modes_lat = int(modes_lat * self.hard_thresholding_fraction)
        self.modes_lon = int(modes_lon * self.hard_thresholding_fraction)

        # new simple linear layer
        self.w = nn.Parameter(
            self.scale
            * torch.randn(
                self.hidden_size, self.hidden_size, self.modes_lat, self.modes_lon, 2
            )
        )
        # optional bias
        if bias:
            self.b = nn.Parameter(
                self.scale * torch.randn(1, self.hidden_size, *self.output_dims)
            )

    def forward(self, x):
        dtype = x.dtype
        # x = x.float()
        B, C, H, W = x.shape

        with amp.autocast(enabled=False):
            x = x.to(torch.float32)
            x = self.forward_transform(x)
            x = torch.view_as_real(x)
            x = x.to(dtype)

        # do spectral conv
        modes = torch.zeros(x.shape, device=x.device)

        # modes[:, :, :self.modes_lat,  :self.modes_lon, :] = self.contract_handle(x[:, :, :self.modes_lat,  :self.modes_lon, :], self.wh)
        # modes[:, :, -self.modes_lat:, :self.modes_lon, :] = self.contract_handle(x[:, :, -self.modes_lat:, :self.modes_lon, :], self.wl)
        modes = self.contract_handle(x, self.w)

        # finalize
        x = F.softshrink(modes, lambd=self.sparsity_threshold)
        x = torch.view_as_complex(x)

        with amp.autocast(enabled=False):
            x = x.to(torch.float32)
            x = torch.view_as_complex(x)
            x = self.inverse_transform(x)
            x = x.to(dtype)

        if hasattr(self, "b"):
            x = x + self.b

        return x


class SpectralConvS2(nn.Module):
    """
    Spectral Convolution as utilized in
    """

    def __init__(
        self,
        forward_transform,
        inverse_transform,
        hidden_size,
        sparsity_threshold=0.0,
        use_complex_kernels=False,
        compression=None,
        rank=128,
        bias=False,
    ):
        super(SpectralConvS2, self).__init__()

        self.hidden_size = hidden_size
        self.sparsity_threshold = sparsity_threshold
        self.scale = 0.02

        self.forward_transform = forward_transform
        self.inverse_transform = inverse_transform

        self.modes_lat = self.forward_transform.lmax
        self.modes_lon = self.forward_transform.mmax

        assert self.inverse_transform.lmax == self.modes_lat  # noqa
        assert self.inverse_transform.mmax == self.modes_lon  # noqa

        # remember the lower triangular indices
        ii, jj = torch.tril_indices(self.modes_lat, self.modes_lon)
        self.register_buffer("ii", ii)
        self.register_buffer("jj", jj)

        if compression == "tt":
            self.rank = rank
            # tensortrain coefficients
            g1 = nn.Parameter(self.scale * torch.randn(self.hidden_size, self.rank, 2))
            g2 = nn.Parameter(
                self.scale * torch.randn(self.rank, self.hidden_size, self.rank, 2)
            )
            g3 = nn.Parameter(self.scale * torch.randn(self.rank, len(ii), 2))
            self.w = nn.ParameterList([g1, g2, g3])

            self.contract_handle = (
                contract_tt  # if use_complex_kernels else raise(NotImplementedError)
            )
        else:
            self.w = nn.Parameter(
                self.scale * torch.randn(self.hidden_size, self.hidden_size, len(ii), 2)
            )
            self.contract_handle = (
                compl_contract_fwd_c if use_complex_kernels else compl_contract_fwd
            )

        if bias:
            self.b = nn.Parameter(
                self.scale * torch.randn(1, self.hidden_size, *self.output_dims)
            )

    def forward(self, x):
        dtype = x.dtype
        # x = x.float()
        B, C, H, W = x.shape

        with amp.autocast(enabled=False):
            x = x.to(torch.float32)
            x = self.forward_transform(x)
            x = torch.view_as_real(x)
            x = x.to(dtype)

        # Populate the sparse spectral grid without an in-place write. The
        # latter breaks autograd under multi-process DDP on some HIP builds.
        spectral_height, spectral_width = x.shape[2:4]
        contracted = self.contract_handle(
            x[:, :, self.ii, self.jj, :], self.w
        )
        spectral_indices = self.ii * spectral_width + self.jj
        modes = torch.zeros_like(x).reshape(
            B, C, spectral_height * spectral_width, 2
        ).index_copy(
            2, spectral_indices, contracted
        )
        modes = modes.view_as(x)

        # finalize
        x = F.softshrink(modes, lambd=self.sparsity_threshold)

        with amp.autocast(enabled=False):
            x = x.to(torch.float32)
            x = torch.view_as_complex(x)
            x = self.inverse_transform(x)
            x = x.to(dtype)

        if hasattr(self, "b"):
            x = x + self.b

        return x


class SpectralAttention2d(nn.Module):
    """
    2d Spectral Attention layer
    """

    def __init__(
        self,
        forward_transform,
        inverse_transform,
        embed_dim,
        sparsity_threshold=0.0,
        hidden_size_factor=2,
        use_complex_network=True,
        use_complex_kernels=False,
        complex_activation="real",
        bias=False,
        spectral_layers=1,
        drop_rate=0.0,
    ):
        super(SpectralAttention2d, self).__init__()

        self.embed_dim = embed_dim
        self.sparsity_threshold = sparsity_threshold
        self.hidden_size = int(hidden_size_factor * self.embed_dim)
        self.scale = 0.02
        self.spectral_layers = spectral_layers
        self.mul_add_handle = (
            compl_muladd2d_fwd_c if use_complex_kernels else compl_muladd2d_fwd
        )
        self.mul_handle = compl_mul2d_fwd_c if use_complex_kernels else compl_mul2d_fwd

        self.modes_lat = forward_transform.lmax
        self.modes_lon = forward_transform.mmax

        # only storing the forward handle to be able to call it
        self.forward_transform = forward_transform.forward
        self.inverse_transform = inverse_transform.forward

        assert inverse_transform.lmax == self.modes_lat  # noqa
        assert inverse_transform.mmax == self.modes_lon  # noqa

        # weights
        w = [self.scale * torch.randn(self.embed_dim, self.hidden_size, 2)]
        # w = [self.scale * torch.randn(self.embed_dim + 2*self.embed_freqs, self.hidden_size, 2)]
        # w = [self.scale * torch.randn(self.embed_dim + 4*self.embed_freqs, self.hidden_size, 2)]
        for l in range(1, self.spectral_layers):
            w.append(self.scale * torch.randn(self.hidden_size, self.hidden_size, 2))
        self.w = nn.ParameterList(w)

        if bias:
            self.b = nn.ParameterList(
                [
                    self.scale * torch.randn(self.hidden_size, 1, 2)
                    for _ in range(self.spectral_layers)
                ]
            )

        self.wout = nn.Parameter(
            self.scale * torch.randn(self.hidden_size, self.embed_dim, 2)
        )

        self.drop = nn.Dropout(drop_rate) if drop_rate > 0.0 else nn.Identity()

        self.activation = ComplexReLU(
            mode=complex_activation, bias_shape=(self.hidden_size, 1, 1)
        )

    def forward_mlp(self, xr):
        for l in range(self.spectral_layers):
            if hasattr(self, "b"):
                xr = self.mul_add_handle(
                    xr, self.w[l].to(xr.dtype), self.b[l].to(xr.dtype)
                )
            else:
                xr = self.mul_handle(xr, self.w[l].to(xr.dtype))
            xr = torch.view_as_complex(xr)
            xr = self.activation(xr)
            xr = self.drop(xr)
            xr = torch.view_as_real(xr)

        xr = self.mul_handle(xr, self.wout)

        return xr

    def forward(self, x):
        dtype = x.dtype
        # x = x.to(torch.float32)

        # FWD transform
        with amp.autocast(enabled=False):
            x = x.to(torch.float32)
            x = self.forward_transform(x)
            x = torch.view_as_real(x)

        # MLP
        x = self.forward_mlp(x)

        # BWD transform
        with amp.autocast(enabled=False):
            x = torch.view_as_complex(x)
            x = self.inverse_transform(x)
            x = x.to(dtype)

        return x


class SpectralAttentionS2(nn.Module):
    """
    geometrical Spectral Attention layer
    """

    def __init__(
        self,
        forward_transform,
        inverse_transform,
        embed_dim,
        sparsity_threshold=0.0,
        hidden_size_factor=2,
        use_complex_network=True,
        use_complex_kernels=False,
        complex_activation="real",
        bias=False,
        spectral_layers=1,
        drop_rate=0.0,
    ):
        super(SpectralAttentionS2, self).__init__()

        self.embed_dim = embed_dim
        self.sparsity_threshold = sparsity_threshold
        self.hidden_size = int(hidden_size_factor * self.embed_dim)
        self.scale = 0.02
        # self.mul_add_handle = compl_muladd1d_fwd_c if use_complex_kernels else compl_muladd1d_fwd
        self.mul_add_handle = (
            compl_muladd2d_fwd_c if use_complex_kernels else compl_muladd2d_fwd
        )
        # self.mul_handle = compl_mul1d_fwd_c if use_complex_kernels else compl_mul1d_fwd
        self.mul_handle = compl_mul2d_fwd_c if use_complex_kernels else compl_mul2d_fwd
        self.spectral_layers = spectral_layers

        self.modes_lat = forward_transform.lmax
        self.modes_lon = forward_transform.mmax

        # only storing the forward handle to be able to call it
        self.forward_transform = forward_transform.forward
        self.inverse_transform = inverse_transform.forward

        assert inverse_transform.lmax == self.modes_lat  # noqa
        assert inverse_transform.mmax == self.modes_lon  # noqa

        # weights
        w = [self.scale * torch.randn(self.embed_dim, self.hidden_size, 2)]
        # w = [self.scale * torch.randn(self.embed_dim + 4*self.embed_freqs, self.hidden_size, 2)]
        for l in range(1, self.spectral_layers):
            w.append(self.scale * torch.randn(self.hidden_size, self.hidden_size, 2))
        self.w = nn.ParameterList(w)

        if bias:
            self.b = nn.ParameterList(
                [
                    self.scale * torch.randn(2 * self.hidden_size, 1, 1, 2)
                    for _ in range(self.spectral_layers)
                ]
            )

        self.wout = nn.Parameter(
            self.scale * torch.randn(self.hidden_size, self.embed_dim, 2)
        )

        self.drop = nn.Dropout(drop_rate) if drop_rate > 0.0 else nn.Identity()

        self.activation = ComplexReLU(
            mode=complex_activation, bias_shape=(self.hidden_size, 1, 1)
        )

    def forward_mlp(self, xr):
        for l in range(self.spectral_layers):
            if hasattr(self, "b"):
                xr = self.mul_add_handle(
                    xr, self.w[l].to(xr.dtype), self.b[l].to(xr.dtype)
                )
            else:
                xr = self.mul_handle(xr, self.w[l].to(xr.dtype))
            xr = torch.view_as_complex(xr)
            xr = self.activation(xr)
            xr = self.drop(xr)
            xr = torch.view_as_real(xr)

        # final MLP
        xr = self.mul_handle(xr, self.wout)

        return xr

    def forward(self, x):
        dtype = x.dtype
        # x = x.to(torch.float32)

        # FWD transform
        with amp.autocast(enabled=False):
            x = x.to(torch.float32)
            x = self.forward_transform(x)
            x = torch.view_as_real(x)

        # MLP
        x = self.forward_mlp(x)

        # BWD transform
        with amp.autocast(enabled=False):
            x = torch.view_as_complex(x)
            x = self.inverse_transform(x)
            x = x.to(dtype)

        return x