# Copyright (c) Kyutai, all rights reserved.
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
#
# Adapted from
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

import logging
import typing as tp
from collections import defaultdict
from dataclasses import dataclass, field
from itertools import chain

import torch
from torch import nn

from ..modules.transformer import create_sin_embedding

logger = logging.getLogger(__name__)
TextCondition = tp.Optional[str]  # a text condition can be a string or None (if doesn't exist)
ConditionTensors = dict[str, 'ConditionType']


class ConditionType(tp.NamedTuple):
    """Return type for a conditioner: both a condition tensor, and a mask indicating valid positions.
    """
    condition: torch.Tensor
    mask: torch.Tensor


@dataclass(frozen=True)
class TensorCondition:
    """Looks quite similar to ConditionType, but represents the input to TensorConditioners.
    `tensor` should be [B | 1, T, D], and `mask` should be `[B | 1, T]`.
    """
    tensor: torch.Tensor
    mask: torch.Tensor

    @staticmethod
    def from_tensor(tensor: torch.Tensor):
        B, T, _ = tensor.shape
        mask = torch.ones(B, T, dtype=torch.bool, device=tensor.device)
        return TensorCondition(tensor, mask)

    @staticmethod
    def cat(conditions: tp.Sequence['TensorCondition']) -> 'TensorCondition':
        assert conditions, "Cannot cat empty list."
        ref_tensor = conditions[0].tensor
        B, _, D = ref_tensor.shape
        assert B == 1
        B = len(conditions)
        T = max(condition.tensor.shape[1] for condition in conditions)
        mask = torch.zeros(B, T, dtype=torch.bool, device=ref_tensor.device)
        tensor = torch.zeros(B, T, D, dtype=ref_tensor.dtype, device=ref_tensor.device)
        for b, condition in enumerate(conditions):
            tensor[b, :condition.tensor.shape[1], :] = condition.tensor[0]
            mask[b, :condition.mask.shape[1]] = condition.mask[0]
        return TensorCondition(tensor, mask)


@dataclass
class ConditionAttributes:
    """Standard class for representing the set of potential inputs to the conditioners.
    Typically, `audiocraft.data.audio_dataset.SegmentInfo` will convert
    to this class to make conditioning agnostic to the type of dataset.

    There are two kinds of conditionings: text (or None), or raw torch tensors (with a mask).

    """
    text: tp.Dict[str, tp.Optional[str]] = field(default_factory=dict)
    tensor: tp.Dict[str, TensorCondition] = field(default_factory=dict)

    @property
    def text_attributes(self) -> tp.Iterable[str]:
        return self.text.keys()

    @property
    def tensor_attributes(self) -> tp.Iterable[str]:
        return self.text.keys()

    @staticmethod
    def condition_types() -> tp.FrozenSet[str]:
        return frozenset(["text", "tensor"])

    def copy(self) -> 'ConditionAttributes':
        return ConditionAttributes(dict(self.text), dict(self.tensor))


Prepared = tp.TypeVar('Prepared')  # represents the prepared condition input type.


class BaseConditioner(nn.Module, tp.Generic[Prepared]):
    """Base model for all conditioner modules.

    Args:
        dim (int): internal dim of the model.
        output_dim (int): Output dim of the conditioner.
        force_linear (bool, optional): Force linear projection even when `dim == output_dim`.
        pad_empty (bool): if True, conditionings of 0 length will be padded to have length 1.
        output_bias (bool): if True, the output projection will have a bias.
        learn_padding (bool): if True, the padding value will be learnt, zero otherwise.
    """

    def __init__(self, dim: int,
                 output_dim: int,
                 device: tp.Union[torch.device, str],
                 force_linear: bool = True,
                 pad_empty: bool = True,
                 output_bias: bool = False,
                 learn_padding: bool = True):
        super().__init__()
        self.dim = dim
        self.output_dim = output_dim
        self.pad_empty = pad_empty
        self.device = device
        self.output_proj: nn.Module
        if force_linear or dim != output_dim:
            self.output_proj = nn.Linear(dim, output_dim, bias=output_bias, device=device)
            assert not output_bias
        else:
            self.output_proj = nn.Identity()
        self.learnt_padding: tp.Optional[torch.Tensor]
        if learn_padding:
            self.learnt_padding = nn.Parameter(
                torch.randn(1, 1, output_dim, device=device), requires_grad=True)
            self.learnt_padding.data *= 0.2
        else:
            self.learnt_padding = None

    def prepare(self, *args, **kwargs) -> Prepared:
        """Should be any part of the processing that will lead to a synchronization
        point, e.g. BPE tokenization with transfer to the GPU.

        The returned value will be saved and return later when calling forward().
        """
        raise NotImplementedError()

    def _get_condition(self, inputs: Prepared) -> ConditionType:
        """Gets input that should be used as conditioning (e.g, genre, description or a waveform).
        Outputs a ConditionType, after the input data was embedded as a dense vector.

        Returns:
            ConditionType:
                - A tensor of size [B, T, dim] where B is the batch size, T is the length of the
                  output embedding and `dim` is the internal dimension of the embedding.
                - And a mask indicating where the padding tokens, of shape `[B, T]`.
        """
        raise NotImplementedError()

    def forward(self, inputs: Prepared) -> ConditionType:
        cond, mask = self._get_condition(inputs)
        B, T, C = cond.shape
        if T == 0 and self.pad_empty:
            cond = torch.zeros(B, T, C, device=cond.device, dtype=cond.dtype)
            mask = torch.zeros_like(cond[..., 0], dtype=torch.bool)

        cond = self.output_proj(cond)

        maskf = mask.float()[..., None]
        if self.learnt_padding is not None:
            cond = cond * maskf + self.learnt_padding * (1 - maskf)
        else:
            cond = cond * maskf
        return ConditionType(cond, mask)


class _BaseTextConditioner(BaseConditioner[Prepared]):
    pass


class _BaseTensorConditioner(BaseConditioner[Prepared]):
    pass


def dropout_tensor(condition: TensorCondition) -> TensorCondition:
    """Utility function for nullifying a WavCondition object.
    """
    return TensorCondition(
        tensor=torch.zeros_like(condition.tensor),
        mask=torch.zeros_like(condition.mask))


def dropout_condition_(sample: ConditionAttributes, condition_type: str, condition: str) -> None:
    """Utility function for nullifying an attribute inside a ConditionAttributes object.
    Works in-place.
    """
    valid_conditions = ConditionAttributes.condition_types()
    if condition_type not in valid_conditions:
        raise ValueError(
            "dropout_condition got an unexpected condition type!"
            f" expected one of {valid_conditions} but got '{condition_type}'")

    if condition not in getattr(sample, condition_type):
        raise ValueError(
            "dropout_condition received an unexpected condition!"
            f" expected tensor={sample.tensor.keys()} and text={sample.text.keys()}"
            f" but got '{condition}' of type '{condition_type}'!"
        )

    if condition_type == 'tensor':
        tensor_condition = sample.tensor[condition]
        sample.tensor[condition] = dropout_tensor(tensor_condition)
    elif condition_type == 'text':
        sample.text[condition] = None
    else:
        assert False


def dropout_all_conditions(attributes: tp.Sequence[ConditionAttributes]) -> list[ConditionAttributes]:
    """
    Args:
        attributes (list[ConditionAttributes]): All conditions attributes.
    Returns:
        list[ConditionAttributes]: Same with all conditions dropped.
    """
    attributes = [attribute.copy() for attribute in attributes]
    for condition_type in ConditionAttributes.condition_types():
        for attribute in attributes:
            for condition in getattr(attribute, condition_type):
                dropout_condition_(attribute, condition_type, condition)
    return attributes


class ConditionProvider(nn.Module):
    """Prepare and provide conditions given all the supported conditioners.

    Args:
        conditioners (dict): Dictionary of conditioners.
        device (torch.device or str, optional): Device for conditioners and output condition types.
    """

    def __init__(self, conditioners: tp.Dict[str, BaseConditioner], device: tp.Union[torch.device, str] = "cpu"):
        super().__init__()
        self.device = device
        self.conditioners = nn.ModuleDict(conditioners).to(device)

    @property
    def text_conditions(self):
        return [k for k, v in self.conditioners.items() if isinstance(v, _BaseTextConditioner)]

    @property
    def tensor_conditions(self):
        return [k for k, v in self.conditioners.items() if isinstance(v, _BaseTensorConditioner)]

    def _collate_text(self, samples: tp.Sequence[ConditionAttributes]) -> tp.Dict[str, tp.List[tp.Optional[str]]]:
        """Given a list of ConditionAttributes objects, compile a dictionary where the keys
        are the attributes and the values are the aggregated input per attribute.
        For example:
        Input:
        [
            ConditionAttributes(text={"genre": "Rock", "description": "A rock song with a guitar solo"}, wav=...),
            ConditionAttributes(text={"genre": "Hip-hop", "description": "A hip-hop verse"}, wav=...),
        ]
        Output:
        {
            "genre": ["Rock", "Hip-hop"],
            "description": ["A rock song with a guitar solo", "A hip-hop verse"]
        }

        Args:
            samples (list of ConditionAttributes): List of ConditionAttributes samples.
        Returns:
            dict[str, list[str, optional]]: A dictionary mapping an attribute name to text batch.
        """
        out: tp.Dict[str, tp.List[tp.Optional[str]]] = defaultdict(list)
        texts = [x.text for x in samples]
        for text in texts:
            for condition in self.text_conditions:
                out[condition].append(text[condition])
        return out

    def _collate_tensors(self, samples: tp.Sequence[ConditionAttributes]) -> tp.Dict[str, TensorCondition]:
        """For each tensor attribute, collate the tensor from individual batch items.

        Args:
            samples (list of ConditionAttributes): List of ConditionAttributes samples.
        Returns:
            dict[str, TensorCondition]: A dictionary mapping an attribute name to tensor.
        """
        per_attribute = defaultdict(list)
        out: tp.Dict[str, TensorCondition] = {}
        for sample in samples:
            for attribute in self.tensor_conditions:
                per_attribute[attribute].append(sample.tensor[attribute])

        # stack all tensors to a single tensor
        for attribute in self.tensor_conditions:
            out[attribute] = TensorCondition.cat(per_attribute[attribute])

        return out

    def prepare(self, inputs: tp.Sequence[ConditionAttributes]) -> tp.Dict[str, tp.Any]:
        """Match attributes/tensors with existing conditioners in self, and call `prepare` for each one.
        This should be called before starting any real GPU work to avoid synchronization points.
        This will return a dict matching conditioner names to their arbitrary prepared representations.

        Args:
            inputs (list[ConditionAttributes]): List of ConditionAttributes objects containing
                text and tensors conditions.
        """
        assert all([isinstance(x, ConditionAttributes) for x in inputs]), (
            "Got unexpected types input for conditioner! should be tp.List[ConditionAttributes]",
            f" but types were {set([type(x) for x in inputs])}"
        )

        output = {}
        text = self._collate_text(inputs)
        tensors = self._collate_tensors(inputs)

        assert set(text.keys() | tensors.keys()).issubset(set(self.conditioners.keys())), (
            f"Got an unexpected attribute! Expected {self.conditioners.keys()}, ",
            f"got {text.keys(), tensors.keys()}"
        )

        missing_inputs = set(self.conditioners.keys()) - (set(text.keys()) | set(tensors.keys()))
        if missing_inputs:
            raise RuntimeError(f"Some conditioners did not receive an input: {missing_inputs}")
        for attribute, batch in chain(text.items(), tensors.items()):
            conditioner = self.conditioners[attribute]
            assert isinstance(conditioner, BaseConditioner)
            output[attribute] = conditioner.prepare(batch)
        return output

    def forward(self, prepared: tp.Dict[str, tp.Any]) -> tp.Dict[str, ConditionType]:
        """Compute pairs of `(embedding, mask)` using the configured conditioners and the prepared representations.
        The output is for example:
        {
            "genre": (torch.Tensor([B, 1, D_genre]), torch.Tensor([B, 1])),
            "description": (torch.Tensor([B, T_desc, D_desc]), torch.Tensor([B, T_desc])),
            ...
        }

        Args:
            prepared (dict): Dict of prepared representations as returned by `prepare()`.
        """
        output = {}
        for name, inputs in prepared.items():
            condition, mask = self.conditioners[name](inputs)
            output[name] = ConditionType(condition, mask)
        return output

    def prepare_and_provide(self, inputs: tp.Sequence[ConditionAttributes]):
        """See .prepare() and .forward()."""
        prepared = self.prepare(inputs)
        return self(prepared)


class ConditionFuser(nn.Module):
    """Condition fuser handles the logic to combine the different conditions
    to the actual model input.

    Args:
        fuse2cond (tp.Dict[str, str]): A dictionary that says how to fuse
            each condition. For example:
            {
                "prepend": ["description"],
                "sum": ["genre", "bpm"],
                "cross": ["description"],
            }
        cross_attention_pos_emb (bool, optional): Use positional embeddings in cross attention.
        cross_attention_pos_emb_scale (int): Scale for positional embeddings in cross attention if used.
    """
    FUSING_METHODS = ["sum", "prepend", "cross"]

    def __init__(self, fuse2cond: tp.Dict[str, tp.List[str]], cross_attention_pos_emb: bool = False,
                 cross_attention_pos_emb_scale: float = 1.0):
        super().__init__()
        assert all(
            [k in self.FUSING_METHODS for k in fuse2cond.keys()]
        ), f"Got invalid fuse method, allowed methods: {self.FUSING_METHODS}"
        self.cross_attention_pos_emb = cross_attention_pos_emb
        self.cross_attention_pos_emb_scale = cross_attention_pos_emb_scale
        self.fuse2cond: tp.Dict[str, tp.List[str]] = fuse2cond
        self.cond2fuse: tp.Dict[str, str] = {}
        for fuse_method, conditions in fuse2cond.items():
            for condition in conditions:
                self.cond2fuse[condition] = fuse_method
                if fuse_method not in ['cross', 'sum']:
                    raise RuntimeError("only `sum` and `cross` conditionings are supported "
                                       f"for now, got {fuse_method}.")

    @property
    def has_conditions(self) -> bool:
        return bool(self.cond2fuse)

    @property
    def has_prepend(self) -> bool:
        """Is there a conditioning that needs to be prepending to the Transformer sequence."""
        return bool(self.fuse2cond['prepend'])

    def get_cross(self, conditions: ConditionTensors) -> torch.Tensor | None:
        """Return the tensor to be provided for the cross attention."""
        cross = None
        for name in self.fuse2cond['cross']:
            cond, _ = conditions[name]
            if cross is None:
                cross = cond
            else:
                cross = torch.cat([cross, cond], dim=1)

        if self.cross_attention_pos_emb and cross is not None:
            positions = torch.arange(
                cross.shape[1],
                device=cross.device
            ).view(1, -1, 1)
            pos_emb = create_sin_embedding(positions, cross.shape[-1]).to(cross.dtype)
            cross = cross + self.cross_attention_pos_emb_scale * pos_emb
        return cross

    def get_sum(self, conditions: ConditionTensors) -> torch.Tensor | None:
        """Return the tensor to be provided as an extra sum offset shared for each step."""
        sum = None
        for name in self.fuse2cond['sum']:
            cond, _ = conditions[name]
            assert cond.shape[1] == 1, cond.shape
            if sum is None:
                sum = cond
            else:
                sum = sum + cond
        return sum

    def get_prepend(self, conditions: ConditionTensors) -> torch.Tensor | None:
        """Return the tensor to be prepended to the transformer."""
        prepend = None
        for name in self.fuse2cond['prepend']:
            cond, _ = conditions[name]
            if prepend is None:
                prepend = cond
            else:
                prepend = torch.cat([cond, prepend], dim=1)
        if prepend is not None:
            sum = self.get_sum(conditions)
            if sum is not None:
                prepend = prepend + sum
        return prepend
