import argparse
from os import makedirs
from os.path import dirname, exists, join
from typing import List

import torch
import math

import torch.nn as nn
import torch.nn.functional as F
from einops.layers.torch import Rearrange

"""
Encoder should take a wave file, then return an embedding
"""


# 16 kHz / 160 in CPC's conv stack. `downsample` halves it once per block, so
# the rate the model sees is a choice: 50 Hz is the original VAP's, 25 Hz halves
# it again to meet Mimi on a shared grid.
CPC_CONV_HZ = 100.0
CPC_RATES = (100.0, 50.0, 25.0, 12.5)


class CPC(nn.Module):
    def __init__(self, frame_hz=25.0):
        super().__init__()
        frame_hz = float(frame_hz)
        if frame_hz not in CPC_RATES:
            raise ValueError(
                f"cpc frame_hz must be one of {CPC_RATES} (the conv stack runs at "
                f"{CPC_CONV_HZ:g} Hz and `downsample` halves it), got {frame_hz:g}"
            )
        self.sample_rate = 16000
        self.frame_hz = frame_hz
        self.encoder = load_CPC(True)
        self.in_dim = self.encoder.gEncoder.conv4.out_channels
        self.out_dim = 256

        halvings = round(math.log2(CPC_CONV_HZ / frame_hz))
        self.downsample = nn.Sequential(
            *[
                get_cnn_layer(
                    in_dim=self.in_dim if block == 0 else self.out_dim,
                    out_dim=self.out_dim,
                    kernel=[3],
                    stride=[2],
                    dilation=[1],
                    activation="GELU",
                )
                for block in range(halvings)
            ]
        )
        if not halvings:   # 100 Hz: still bridge the encoder width onto out_dim
            self.downsample = nn.Sequential(nn.Linear(self.in_dim, self.out_dim))

        self.freeze()

    def forward(self, waveform):
        if waveform.ndim < 3:
            waveform = waveform.unsqueeze(1)
        z = self.encoder.gEncoder(waveform)
        z = z.permute(0, 2, 1)
        z = self.encoder.gAR(z)
        z = self.downsample(z)
        return z

    def freeze(self):
        for p in self.encoder.parameters():
            p.requires_grad_(False)


"""

#############################################################
#############################################################
WARNING - ATTENTION - HEAR YEE HEAAAR YEEEE
#############################################################
#############################################################

Most of the code in this file are scaled down (and heavily copied) versions of

https://github.com/facebookresearch/CPC_audio

Please checkout their codebase if you are interested in CPC networks
----------------------------------------------------------------------------

torch.hub downloads to: `$HOME/.cache/torch/hub/checkpoints/`
Explicit checkpoint path saved manually in "assets/" see CHECKPOINTS below.
"""


def repo_root():
    """
    Returns the absolute path to the git repository
    """
    root = dirname(__file__)
    root = dirname(root)
    return root


CHECKPOINTS = {
    "cpc": join(repo_root(), "assets/checkpoints/cpc/60k_epoch4-d0f474de.pt")
}
NAMES = list(CHECKPOINTS.keys())


class ChannelNorm(nn.Module):
    """
    Most of the code in this file are scaled down (and heavily copied) versions of
        https://github.com/facebookresearch/CPC_audio
    """

    def __init__(self, numFeatures, epsilon=1e-05, affine=True):

        super(ChannelNorm, self).__init__()
        if affine:
            self.weight = nn.parameter.Parameter(torch.Tensor(1, numFeatures, 1))
            self.bias = nn.parameter.Parameter(torch.Tensor(1, numFeatures, 1))
        else:
            self.weight = None
            self.bias = None
        self.epsilon = epsilon
        self.p = 0
        self.affine = affine
        self.reset_parameters()

    def reset_parameters(self):
        if self.affine:
            torch.nn.init.ones_(self.weight)
            torch.nn.init.zeros_(self.bias)

    def forward(self, x):

        cumMean = x.mean(dim=1, keepdim=True)
        cumVar = x.var(dim=1, keepdim=True)
        x = (x - cumMean) * torch.rsqrt(cumVar + self.epsilon)

        if self.weight is not None:
            x = x * self.weight + self.bias
        return x


class CPCEncoder(nn.Module):
    """
    Most of the code in this file are scaled down (and heavily copied) versions of
        https://github.com/facebookresearch/CPC_audio
    """

    def __init__(self, sizeHidden=512, normMode="layerNorm"):
        super(CPCEncoder, self).__init__()
        normLayer = ChannelNorm
        self.dimEncoded = sizeHidden
        self.conv0 = nn.Conv1d(1, sizeHidden, 10, stride=5, padding=3)
        self.batchNorm0 = normLayer(sizeHidden)
        self.conv1 = nn.Conv1d(sizeHidden, sizeHidden, 8, stride=4, padding=2)
        self.batchNorm1 = normLayer(sizeHidden)
        self.conv2 = nn.Conv1d(sizeHidden, sizeHidden, 4, stride=2, padding=1)
        self.batchNorm2 = normLayer(sizeHidden)
        self.conv3 = nn.Conv1d(sizeHidden, sizeHidden, 4, stride=2, padding=1)
        self.batchNorm3 = normLayer(sizeHidden)
        self.conv4 = nn.Conv1d(sizeHidden, sizeHidden, 4, stride=2, padding=1)
        self.batchNorm4 = normLayer(sizeHidden)
        self.DOWNSAMPLING = 160

    def getDimOutput(self):
        return self.conv4.out_channels

    def forward(self, x):
        x = F.relu(self.batchNorm0(self.conv0(x)))
        x = F.relu(self.batchNorm1(self.conv1(x)))
        x = F.relu(self.batchNorm2(self.conv2(x)))
        x = F.relu(self.batchNorm3(self.conv3(x)))
        x = F.relu(self.batchNorm4(self.conv4(x)))
        return x


class CPCAR(nn.Module):
    """
    Most of the code in this file are scaled down (and heavily copied) versions of
        https://github.com/facebookresearch/CPC_audio
    """

    def __init__(
        self, dimEncoded, dimOutput, keepHidden, nLevelsGRU, mode="GRU", reverse=False
    ):

        super(CPCAR, self).__init__()
        self.RESIDUAL_STD = 0.1

        if mode == "LSTM":
            self.baseNet = nn.LSTM(
                dimEncoded, dimOutput, num_layers=nLevelsGRU, batch_first=True
            )
        elif mode == "RNN":
            self.baseNet = nn.RNN(
                dimEncoded, dimOutput, num_layers=nLevelsGRU, batch_first=True
            )
        else:
            self.baseNet = nn.GRU(
                dimEncoded, dimOutput, num_layers=nLevelsGRU, batch_first=True
            )

        self.hidden = None
        self.keepHidden = keepHidden
        self.reverse = reverse

    def getDimOutput(self):
        return self.baseNet.hidden_size

    def forward(self, x):

        if self.reverse:
            x = torch.flip(x, [1])
        try:
            self.baseNet.flatten_parameters()
        except RuntimeError:
            pass
        x, h = self.baseNet(x, self.hidden)
        if self.keepHidden:
            if isinstance(h, tuple):
                self.hidden = tuple(x.detach() for x in h)
            else:
                self.hidden = h.detach()

        # For better modularity, a sequence's order should be preserved
        # by each module
        if self.reverse:
            x = torch.flip(x, [1])
        return x


class CPCModel(nn.Module):
    """
    Most of the code in this file are scaled down (and heavily copied) versions of
        https://github.com/facebookresearch/CPC_audio
    """

    def __init__(self, encoder, AR):
        super(CPCModel, self).__init__()
        self.gEncoder = encoder
        self.gAR = AR

    def forward(self, batchData, label):
        encodedData = self.gEncoder(batchData).permute(0, 2, 1)
        cFeature = self.gAR(encodedData)
        return cFeature, encodedData, label


def load_CPC(load_state_dict=True):
    """
    Contrast predictive learning model for audio data
    pretrained: if True, load a model trained on libri-light 60k
    (https://arxiv.org/abs/1912.07875)
    **kwargs : see cpc/cpc_default_config to get the list of possible arguments

    Most of the code in this file are scaled down (and heavily copied) versions of
        https://github.com/facebookresearch/CPC_audio
    """

    def loadArgs(args, locArgs, forbiddenAttr=None):
        for k, v in vars(locArgs).items():
            if forbiddenAttr is not None:
                if k not in forbiddenAttr:
                    setattr(args, k, v)
            else:
                setattr(args, k, v)

    def get_default_cpc_config():
        parser = argparse.ArgumentParser()

        # Run parameters
        group = parser.add_argument_group(
            "Architecture configuration",
            description="The arguments defining the model's architecture.",
        )
        group.add_argument(
            "--hiddenEncoder",
            type=int,
            default=256,
            help="Hidden dimension of the encoder network.",
        )
        group.add_argument(
            "--hiddenGar",
            type=int,
            default=256,
            help="Hidden dimension of the auto-regressive network",
        )
        group.add_argument(
            "--nPredicts", type=int, default=12, help="Number of steps to predict."
        )
        group.add_argument(
            "--negativeSamplingExt",
            type=int,
            default=128,
            help="Number of negative samples to take.",
        )
        group.add_argument("--learningRate", type=float, default=2e-4)
        group.add_argument(
            "--schedulerStep",
            type=int,
            default=-1,
            help="Step of the learning rate scheduler: at each "
            "step the learning rate is divided by 2. Default: "
            "no scheduler.",
        )
        group.add_argument(
            "--schedulerRamp",
            type=int,
            default=None,
            help="Enable a warm up phase for the learning rate: "
            "adds a linear ramp of the given size.",
        )
        group.add_argument(
            "--beta1",
            type=float,
            default=0.9,
            help="Value of beta1 for the Adam optimizer",
        )
        group.add_argument(
            "--beta2",
            type=float,
            default=0.999,
            help="Value of beta2 for the Adam optimizer",
        )
        group.add_argument(
            "--epsilon",
            type=float,
            default=1e-08,
            help="Value of epsilon for the Adam optimizer",
        )
        group.add_argument(
            "--sizeWindow",
            type=int,
            default=20480,
            help="Number of frames to consider at each batch.",
        )
        group.add_argument(
            "--nEpoch", type=int, default=200, help="Number of epoch to run"
        )
        group.add_argument(
            "--samplingType",
            type=str,
            default="samespeaker",
            choices=["samespeaker", "uniform", "samesequence", "sequential"],
            help="How to sample the negative examples in the CPC loss.",
        )
        group.add_argument(
            "--nLevelsPhone",
            type=int,
            default=1,
            help="(Supervised mode only). Number of layers in "
            "the phone classification network.",
        )
        group.add_argument(
            "--cpc_mode",
            type=str,
            default=None,
            choices=["reverse", "none"],
            help="Some variations on CPC.",
        )
        group.add_argument(
            "--encoder_type",
            type=str,
            choices=["cpc", "mfcc", "lfb"],
            default="cpc",
            help="Replace the encoder network by mfcc features or learned filter banks",
        )
        group.add_argument(
            "--normMode",
            type=str,
            default="layerNorm",
            choices=["instanceNorm", "ID", "layerNorm", "batchNorm"],
            help="Type of normalization to use in the encoder "
            "network (default is layerNorm).",
        )
        group.add_argument(
            "--onEncoder",
            action="store_true",
            help="(Supervised mode only) Perform the "
            "classification on the encoder's output.",
        )
        group.add_argument(
            "--random_seed", type=int, default=None, help="Set a specific random seed."
        )
        group.add_argument(
            "--speakerEmbedding",
            type=int,
            default=0,
            help="(Depreciated) Feed the prediction network with "
            "speaker embeddings along with the usual sequence.",
        )
        group.add_argument(
            "--arMode",
            default="LSTM",
            choices=["GRU", "LSTM", "RNN", "no_ar", "transformer"],
            help="Architecture to use for the auto-regressive "
            "network (default is lstm).",
        )
        group.add_argument(
            "--nLevelsGRU",
            type=int,
            default=1,
            help="Number of layers in the autoregressive network.",
        )
        group.add_argument(
            "--rnnMode",
            type=str,
            default="transformer",
            choices=[
                "transformer",
                "RNN",
                "LSTM",
                "linear",
                "ffd",
                "conv4",
                "conv8",
                "conv12",
            ],
            help="Architecture to use for the prediction network",
        )
        group.add_argument(
            "--dropout",
            action="store_true",
            help="Add a dropout layer at the output of the prediction network.",
        )
        group.add_argument(
            "--abspos",
            action="store_true",
            help="If the prediction network is a transformer, "
            "active to use absolute coordinates.",
        )
        return parser.parse_args([])

    # from cpc.model import CPCModel as cpcmodel
    # from cpc.cpc_default_config import get_default_cpc_config
    # from cpc.feature_loader import getEncoder, getAR, loadArgs
    # from cpc.feature_loader import loadArgs

    locArgs = get_default_cpc_config()
    if exists(CHECKPOINTS["cpc"]):
        checkpoint = torch.load(
            CHECKPOINTS["cpc"], map_location="cpu", weights_only=True
        )
    else:
        checkpoint_url = "https://dl.fbaipublicfiles.com/librilight/CPC_checkpoints/60k_epoch4-d0f474de.pt"
        checkpoint = torch.hub.load_state_dict_from_url(
            checkpoint_url, progress=False, map_location="cpu"
        )
        makedirs(dirname(CHECKPOINTS["cpc"]), exist_ok=True)
        torch.save(checkpoint, CHECKPOINTS["cpc"])
    loadArgs(locArgs, argparse.Namespace(**checkpoint["config"]))
    # encoderNet = getEncoder(locArgs)
    encoderNet = CPCEncoder(locArgs.hiddenEncoder, locArgs.normMode)
    # arNet = getAR(locArgs)
    arNet = CPCAR(
        locArgs.hiddenEncoder,
        locArgs.hiddenGar,
        locArgs.samplingType == "sequential",
        locArgs.nLevelsGRU,
        mode=locArgs.arMode,
        reverse=locArgs.cpc_mode == "reverse",
    )
    # model = cpcmodel(encoderNet, arNet)
    model = CPCModel(encoderNet, arNet)

    # always load pretrained
    if load_state_dict:
        model.load_state_dict(checkpoint["weights"], strict=False)
    model.name = "cpc"
    return model


class LayerNorm(nn.Module):
    def __init__(self, dim: int, rearrange_outputs: bool = True) -> None:
        super().__init__()
        self.ln = nn.LayerNorm(dim)
        self.in_rearrange = Rearrange("b d t -> b t d")
        if rearrange_outputs:
            self.out_rearrange = Rearrange("b t d -> b d t")
        else:
            self.out_rearrange = nn.Identity()

    def __repr__(self):
        return str(self.ln)

    def forward(self, x):
        return self.out_rearrange(self.ln(self.in_rearrange(x)))


class CConv1d(nn.Conv1d):
    """source: https://github.com/pytorch/pytorch/issues/1333"""

    def __init__(
        self,
        in_channels,
        out_channels,
        kernel_size,
        stride=1,
        dilation=1,
        groups=1,
        padding_value=0,
        bias=True,
        **kwargs,
    ):
        super().__init__(
            in_channels,
            out_channels,
            kernel_size=kernel_size,
            stride=stride,
            dilation=dilation,
            groups=groups,
            bias=bias,
            **kwargs,
        )

        ks = kernel_size if isinstance(kernel_size, int) else kernel_size[0]
        pad_dim1_pre = ks - 1
        pad_dim1_post = 0
        if dilation > 0:
            pad_dim1_pre *= dilation
        pad = (pad_dim1_pre, pad_dim1_post)
        self.pad = nn.ConstantPad1d(padding=pad, value=padding_value)

    def debug_weights(self, type="sum"):
        w = 1.0
        if type == "mean":
            w = 1.0 / self.kernel_size[0]

        elif type == "range":
            k = self.kernel_size[0]
            w = torch.arange(1, k + 1).float().pow(2)
            w = w.repeat(self.out_channels, self.in_channels, 1)
            self.weight.data = w
            if self.bias:
                self.bias.data = self.bias.data.fill_(0.0)
            return None

        self.weight.data = self.weight.data.fill_(w)
        if self.bias:
            self.bias.data = self.bias.data.fill_(0.0)

    def forward(self, input):
        return super().forward(self.pad(input))


def get_cnn_layer(
    in_dim: int,
    out_dim: int,
    kernel: List[int] = [3],
    stride: List[int] = [2],
    dilation: List[int] = [1],
    padding: List[int] = [0],
    activation: str = "GELU",
):
    layers = []
    layers.append(Rearrange("b t d -> b d t"))
    for k, s, d, p in zip(kernel, stride, dilation, padding):
        layers.append(
            CConv1d(in_dim, out_dim, kernel_size=k, stride=s, dilation=d, padding=p)
        )
        layers.append(LayerNorm(out_dim))
        if activation is not None:
            layers.append(getattr(torch.nn, activation)())
    layers.append(Rearrange("b d t -> b t d"))
    return nn.Sequential(*layers)