from dataclasses import dataclass
import torch
from tqdm.auto import trange
import typing as tp
from einops import rearrange
from torch import nn

from .autoencoders import AudioAutoencoder
from .conditioners import MultiConditioner, create_multi_conditioner_from_conditioning_config
from .factory import create_pretransform_from_config
from .lm_backbone import AudioLMBackbone, ContinuousTransformerAudioLMBackbone
from .pretransforms import Pretransform, AutoencoderPretransform, PretrainedDACPretransform, AudiocraftCompressionPretransform
from .utils import multinomial, sample_top_k, sample_top_p

from .codebook_patterns import (
    CodebooksPatternProvider,
    DelayedPatternProvider,
    MusicLMPattern,
    ParallelPatternProvider,
    UnrolledPatternProvider
)

# Copied and modified from https://github.com/facebookresearch/audiocraft/blob/main/audiocraft/models/lm.py under MIT license
# License can be found in LICENSES/LICENSE_META.txt

@dataclass
class LMOutput:
    # The logits are already re-aligned with the input codes
    # hence no extra shift is required, e.g. when computing CE
    logits: torch.Tensor  # [B, K, T, card]
    mask: torch.Tensor  # [B, K, T]

# Wrapper for a multi-codebook language model
# Handles patterns and quantizer heads
class AudioLanguageModel(nn.Module):
    def __init__(
            self, 
            pattern_provider: CodebooksPatternProvider, 
            backbone: AudioLMBackbone,
            num_quantizers: int,
            codebook_size: int
        ):
        super().__init__()

        self.pattern_provider = pattern_provider
        self.backbone = backbone
        self.num_quantizers = num_quantizers
        self.codebook_size = codebook_size

        self.masked_token_id = codebook_size

        # Per-quantizer embedders
        # Add one for the mask embed
        self.embeds = nn.ModuleList([nn.Embedding(codebook_size + 1, backbone.embed_dim) for _ in range(num_quantizers)])

        # Per-quantizer output heads
        self.quantizer_heads = nn.ModuleList([
            nn.Linear(backbone.embed_dim, codebook_size) for _ in range(num_quantizers)
        ])

    def forward(self,
            sequence: torch.Tensor, #[batch, seq_len, 
            prepend_cond=None, #[batch, seq, channels]
            prepend_cond_mask=None,
            cross_attn_cond=None, #[batch, seq, channels],
            **kwargs
        ):

        batch, num_quantizers, seq_len = sequence.shape

        assert num_quantizers == self.num_quantizers, "Number of quantizers in sequence must match number of quantizers in model"

        backbone_input = sum([self.embeds[i](sequence[:, i]) for i in range(num_quantizers)]) # [batch, seq_len, embed_dim]

        dtype = next(self.parameters()).dtype

        if cross_attn_cond is not None:
            cross_attn_cond = cross_attn_cond.to(dtype)

        if prepend_cond is not None:
            prepend_cond = prepend_cond.to(dtype)

            if prepend_cond_mask is not None:
                prepend_cond_mask = prepend_cond_mask.to(dtype)
            
        backbone_input = backbone_input.to(dtype)

        output = self.backbone(
            backbone_input,
            cross_attn_cond=cross_attn_cond,
            prepend_cond=prepend_cond,
            prepend_cond_mask=prepend_cond_mask,
            **kwargs
        ) # [batch, seq_len, embed_dim]

        # Run output through quantizer heads
        logits = torch.stack([self.quantizer_heads[i](output) for i in range(num_quantizers)], dim=1) # [batch, num_quantizers, seq_len, codebook_size]

        return logits
    
    def compute_logits(
            self, 
            codes, #[batch, num_quantizers, seq_len]
            **kwargs):
        """
        Compute logits for a batch of codes, optionally conditioning on cross-attention and prepend conditioning
        Handles translation between input sequence and pattern-shifted sequence
        Only used during training
        """
        
        batch, _, seq_len = codes.shape

        pattern = self.pattern_provider.get_pattern(seq_len)

        # Apply the token pattern to the codes, shifting the codes as needed and masking out invalid steps
        shifted_codes, _, _ = pattern.build_pattern_sequence(
            codes,
            self.masked_token_id,
            keep_only_valid_steps=True
        )

        # Run the model to get logits for each quantizer [batch, num_quantizers, seq_len, codebook_size]
        logits = self(shifted_codes, **kwargs)

        # Rearrange logits to prepare to revert pattern
        logits = rearrange(logits, "b n s c -> b c n s")

        # Revert sequence logits back to original sequence length, removing masked steps
        logits, _, logits_mask = pattern.revert_pattern_logits(
            logits, float('nan'), keep_only_valid_steps=True
        )

        logits = rearrange(logits, "b c n t -> b n t c")

        logits_mask = logits_mask[None, :, :].expand(batch, -1, -1) # [batch, num_quantizers, seq_len]

        return LMOutput(logits=logits, mask=logits_mask)

# Conditioning and generation wrapper for a multi-codebook language model
# Handles conditioning, CFG, generation, and encoding/decoding
class AudioLanguageModelWrapper(nn.Module):
    def __init__(
            self, 
            pretransform: Pretransform,
            lm: AudioLanguageModel,
            sample_rate: int,
            min_input_length: int,
            conditioner: MultiConditioner = None,
            cross_attn_cond_ids: tp.List[str] = [],
            prepend_cond_ids: tp.List[str] = [],
            global_cond_ids: tp.List[str] = []
        ):
        super().__init__()
        
        assert pretransform.is_discrete, "Pretransform must be discrete"
        self.pretransform = pretransform

        self.pretransform.requires_grad_(False)
        self.pretransform.eval()

        if isinstance(self.pretransform, AutoencoderPretransform):
            self.num_quantizers = self.pretransform.model.bottleneck.num_quantizers
            self.codebook_size = self.pretransform.model.bottleneck.codebook_size
        elif isinstance(self.pretransform, PretrainedDACPretransform):
            self.num_quantizers = self.pretransform.model.num_quantizers
            self.codebook_size = self.pretransform.model.codebook_size
        elif isinstance(self.pretransform, AudiocraftCompressionPretransform):
            self.num_quantizers = self.pretransform.num_quantizers
            self.codebook_size = self.pretransform.codebook_size
        else:
            raise NotImplementedError(f"Unrecognized pretransform type {type(self.pretransform)}")

        self.conditioner = conditioner

        self.lm = lm

        self.sample_rate = sample_rate
        self.min_input_length = min_input_length

        self.cross_attn_cond_ids = cross_attn_cond_ids
        self.prepend_cond_ids = prepend_cond_ids
        self.global_cond_ids = global_cond_ids
    
    def get_conditioning_inputs(self, cond: tp.Dict[str, tp.Any], negative=False):
        cross_attention_input = None
        prepend_cond = None
        prepend_cond_mask = None
        global_cond = None

        if len(self.cross_attn_cond_ids) > 0:
            # Concatenate all cross-attention inputs over the sequence dimension
            # Assumes that the cross-attention inputs are of shape (batch, seq, channels)
            cross_attention_input = torch.cat([cond[key][0] for key in self.cross_attn_cond_ids], dim=1)

        if len(self.prepend_cond_ids) > 0:
            # Concatenate all prepend conditioning inputs over the sequence dimension
            # Assumes that the prepend conditioning inputs are of shape (batch, seq, channels)
            prepend_cond = torch.cat([cond[key][0] for key in self.prepend_cond_ids], dim=1)
            prepend_cond_mask = torch.cat([cond[key][1] for key in self.prepend_cond_ids], dim=1)

        if len(self.global_cond_ids) > 0:
            # Concatenate all global conditioning inputs over the channel dimension
            # Assumes that the global conditioning inputs are of shape (batch, channels)
            global_cond = torch.cat([cond[key][0] for key in self.global_cond_ids], dim=-1)
            if len(global_cond.shape) == 3:
                global_cond = global_cond.squeeze(1)

        if negative:
            return {
                "negative_cross_attn_cond": cross_attention_input,
                "negative_prepend_cond": prepend_cond,
                "negative_prepend_cond_mask": prepend_cond_mask,
                "negative_global_cond": global_cond
            }
        else:
            return {
                "cross_attn_cond": cross_attention_input,
                "prepend_cond": prepend_cond,
                "prepend_cond_mask": prepend_cond_mask,
                "global_cond": global_cond
            }
        
    def compute_logits(
            self, 
            codes, 
            condition_tensors=None, 
            cfg_dropout_prob=0.0,
            **kwargs
        ):
        """
        Compute logits for a batch of codes, and translates from conditioning inputs to model inputs
        Handles CFG dropout
        """

        if condition_tensors is None:
            condition_tensors = {}

        conditioning_inputs = self.get_conditioning_inputs(condition_tensors)

        cross_attn_cond = conditioning_inputs["cross_attn_cond"]
        prepend_cond = conditioning_inputs["prepend_cond"]
        prepend_cond_mask = conditioning_inputs["prepend_cond_mask"]
        global_cond = conditioning_inputs["global_cond"]

        if cfg_dropout_prob > 0.0:
            if cross_attn_cond is not None:
                null_embed = torch.zeros_like(cross_attn_cond, device=cross_attn_cond.device)
                dropout_mask = torch.bernoulli(torch.full((cross_attn_cond.shape[0], 1, 1), cfg_dropout_prob, device=cross_attn_cond.device)).to(torch.bool)
                cross_attn_cond = torch.where(dropout_mask, null_embed, cross_attn_cond)
        
            if prepend_cond is not None:
                null_embed = torch.zeros_like(prepend_cond, device=prepend_cond.device)
                dropout_mask = torch.bernoulli(torch.full((prepend_cond.shape[0], 1, 1), cfg_dropout_prob, device=prepend_cond.device)).to(torch.bool)
                prepend_cond = torch.where(dropout_mask, null_embed, prepend_cond)

            if global_cond is not None:
                null_embed = torch.zeros_like(global_cond, device=global_cond.device)
                dropout_mask = torch.bernoulli(torch.full((global_cond.shape[0], 1), cfg_dropout_prob, device=global_cond.device)).to(torch.bool)
                global_cond = torch.where(dropout_mask, null_embed, global_cond)

        return self.lm.compute_logits(codes, cross_attn_cond=cross_attn_cond, prepend_cond=prepend_cond, prepend_cond_mask=prepend_cond_mask, global_cond=global_cond, **kwargs)
    
    def _sample_next_token(
            self, 
            sequence, #[batch, num_quantizers, seq_len]
            conditioning_tensors=None, 
            cross_attn_use_cfg=True,
            prepend_use_cfg=True,
            global_use_cfg=True,
            cfg_scale=1.0,
            top_k=250,
            top_p=0.0,
            temp=1.0,
            **kwargs
        ):
        """
        Sample the next token for a batch of codes, and translates from conditioning inputs to model inputs
        Handles CFG inference
        """

        if conditioning_tensors is None:
            conditioning_tensors = {}

        conditioning_inputs = self.get_conditioning_inputs(conditioning_tensors)

        cross_attn_cond = conditioning_inputs["cross_attn_cond"]
        prepend_cond = conditioning_inputs["prepend_cond"]
        prepend_cond_mask = conditioning_inputs["prepend_cond_mask"]
        global_cond = conditioning_inputs["global_cond"]

        if cfg_scale != 1.0:
            
            # Batch size is doubled to account for negative samples
            sequence = torch.cat([sequence, sequence], dim=0)

            if cross_attn_cond is not None and cross_attn_use_cfg:
                null_embed = torch.zeros_like(cross_attn_cond, device=cross_attn_cond.device)

                cross_attn_cond = torch.cat([cross_attn_cond, null_embed], dim=0)
            
            if prepend_cond is not None and prepend_use_cfg:
                null_embed = torch.zeros_like(prepend_cond, device=prepend_cond.device)

                prepend_cond = torch.cat([prepend_cond, null_embed], dim=0) 

                if prepend_cond_mask is not None:
                    prepend_cond_mask = torch.cat([prepend_cond_mask, prepend_cond_mask], dim=0)

            if global_cond is not None and global_use_cfg:
                null_embed = torch.zeros_like(global_cond, device=global_cond.device)

                global_cond = torch.cat([global_cond, null_embed], dim=0)

        logits = self.lm(sequence, cross_attn_cond=cross_attn_cond, prepend_cond=prepend_cond, prepend_cond_mask=prepend_cond_mask, global_cond=global_cond, **kwargs)

        if cfg_scale != 1.0:
            cond_logits, uncond_logits = logits.chunk(2, dim=0)

            logits = uncond_logits + (cond_logits - uncond_logits) * cfg_scale

        logits = rearrange(logits, "b n s c -> b n c s") # [batch, num_quantizers, codebook_size, seq_len]
        
        # Grab the logits for the last step
        logits = logits[:, :, :, -1] # [batch, num_quantizers, codebook_size]

        # Apply top-k or top-p sampling

        if temp > 0:
            probs = torch.softmax(logits / temp, dim=-1)

            if top_p > 0.0:
                next_token = sample_top_p(probs, p=top_p)
            elif top_k > 0:
                next_token = sample_top_k(probs, k=top_k)
            else:
                next_token = multinomial(probs, num_samples=1)

        else:
            next_token = torch.argmax(logits, dim=-1, keepdim=True) # [batch, num_quantizers, 1]

        return next_token

    @torch.no_grad()
    def generate(
        self,
        max_gen_len: int = 256,
        batch_size: tp.Optional[int] = None,
        init_data: tp.Optional[torch.Tensor] = None,
        conditioning: tp.Optional[tp.Dict[str, tp.Any]] = None,
        conditioning_tensors: tp.Optional[tp.Dict[str, tp.Any]] = None,
        callback: tp.Optional[tp.Callable[[int, int], None]] = None,
        use_cache: bool = True,
        cfg_scale: float = 1.0,
        **kwargs
    ):
        device = next(self.parameters()).device

        if conditioning_tensors is None and conditioning is not None:
            # Convert conditioning inputs to conditioning tensors
            conditioning_tensors = self.conditioner(conditioning, device)

        # Check that batch size is consistent across inputs
        possible_batch_sizes = []

        if batch_size is not None:
            possible_batch_sizes.append(batch_size)
        elif init_data is not None:
            possible_batch_sizes.append(init_data.shape[0])
        elif conditioning_tensors is not None:
            # Assume that the first conditioning tensor has the batch dimension
            possible_batch_sizes.append(conditioning_tensors[list(conditioning_tensors.keys())[0]][0].shape[0])
        else:
            possible_batch_sizes.append(1)

        assert [x == possible_batch_sizes[0] for x in possible_batch_sizes], "Batch size must be consistent across inputs"

        batch_size = possible_batch_sizes[0]
        
        if init_data is None:
            # Initialize with zeros
            assert batch_size > 0
            init_data = torch.zeros((batch_size, self.num_quantizers, 0), device=device, dtype=torch.long)

        batch_size, num_quantizers, seq_len = init_data.shape

        start_offset = seq_len
        assert start_offset < max_gen_len, "init data longer than max gen length"

        pattern = self.lm.pattern_provider.get_pattern(max_gen_len)

        unknown_token = -1

        # Initialize the generated codes with the init data, padded with unknown tokens
        gen_codes = torch.full((batch_size, num_quantizers, max_gen_len), unknown_token, device=device, dtype=torch.long)
        gen_codes[:, :, :start_offset] = init_data # [batch, num_quantizers, max_gen_len]

        gen_sequence, _, mask = pattern.build_pattern_sequence(gen_codes, self.lm.masked_token_id) # [batch, num_quantizers, gen_sequence_len]

        start_offset_sequence = pattern.get_first_step_with_timesteps(start_offset)
        assert start_offset_sequence is not None

        # Generation
        prev_offset = 0
        gen_sequence_len = gen_sequence.shape[-1]

        # Reset generation cache
        if use_cache and self.lm.backbone.use_generation_cache:
            self.lm.backbone.reset_generation_cache(max_gen_len, batch_size if cfg_scale == 1.0 else batch_size * 2)

        for offset in trange(start_offset_sequence, gen_sequence_len):

            # Get the full sequence up to the current offset
            curr_sequence = gen_sequence[..., prev_offset:offset]

            next_token = self._sample_next_token(
                curr_sequence,
                conditioning_tensors=conditioning_tensors,
                use_cache=use_cache,
                cfg_scale=cfg_scale,
                **kwargs
            )

            valid_mask = mask[..., offset:offset+1].expand(batch_size, -1, -1)
            next_token[~valid_mask] = self.lm.masked_token_id

            # Update the generated sequence with the next token
            gen_sequence[..., offset:offset+1] = torch.where(
                gen_sequence[..., offset:offset+1] == unknown_token,
                next_token, 
                gen_sequence[..., offset:offset+1]
            )

            if use_cache and self.lm.backbone.use_generation_cache:
                # Only update the offset if caching is being used
                prev_offset = offset

                self.lm.backbone.update_generation_cache(offset)

            if callback is not None:
                # Callback to report progress
                # Pass in the offset relative to the start of the sequence, and the length of the current sequence
                callback(1 + offset - start_offset_sequence, gen_sequence_len - start_offset_sequence)

        assert not (gen_sequence == unknown_token).any(), "Unknown tokens in generated sequence"

        out_codes, _, out_mask = pattern.revert_pattern_sequence(gen_sequence, special_token=unknown_token)

        # sanity checks over the returned codes and corresponding masks
        assert (out_codes[..., :max_gen_len] != unknown_token).all()
        assert (out_mask[..., :max_gen_len] == 1).all()

        #out_codes = out_codes[..., 0:max_gen_len]

        return out_codes
    

    def generate_audio(
        self,
        **kwargs
    ):
        """
        Generate audio from a batch of codes
        """

        codes = self.generate(**kwargs)

        audio = self.pretransform.decode_tokens(codes)

        return audio


def create_audio_lm_from_config(config):
    model_config = config.get('model', None)
    assert model_config is not None, 'model config must be specified in config'

    sample_rate = config.get('sample_rate', None)
    assert sample_rate is not None, "Must specify sample_rate in config"
    
    lm_config = model_config.get('lm', None)
    assert lm_config is not None, 'lm config must be specified in model config'

    codebook_pattern = lm_config.get("codebook_pattern", "delay")

    pattern_providers = {
        'parallel': ParallelPatternProvider,
        'delay': DelayedPatternProvider,
        'unroll': UnrolledPatternProvider,
        'musiclm': MusicLMPattern,
    }

    pretransform_config = model_config.get("pretransform", None)
    
    pretransform = create_pretransform_from_config(pretransform_config, sample_rate)

    assert pretransform.is_discrete, "Pretransform must be discrete"

    min_input_length = pretransform.downsampling_ratio

    pattern_provider = pattern_providers[codebook_pattern](n_q=pretransform.num_quantizers)

    conditioning_config = model_config.get('conditioning', None)

    conditioner = None
    if conditioning_config is not None:
        conditioner = create_multi_conditioner_from_conditioning_config(conditioning_config, pretransform=pretransform)

    cross_attn_cond_ids = lm_config.get('cross_attention_cond_ids', [])
    prepend_cond_ids = lm_config.get('prepend_cond_ids', [])
    global_cond_ids = lm_config.get('global_cond_ids', [])

    lm_type = lm_config.get("type", None)
    lm_model_config = lm_config.get("config", None)

    assert lm_type is not None, "Must specify lm type in lm config"
    assert lm_model_config is not None, "Must specify lm model config in lm config"

    if lm_type == "continuous_transformer":
        backbone = ContinuousTransformerAudioLMBackbone(**lm_model_config)
    else:
        raise NotImplementedError(f"Unrecognized lm type {lm_type}")

    lm = AudioLanguageModel(
        pattern_provider=pattern_provider,
        backbone=backbone,
        num_quantizers=pretransform.num_quantizers,
        codebook_size=pretransform.codebook_size
    )

    model = AudioLanguageModelWrapper(
        pretransform=pretransform,
        lm=lm,
        conditioner=conditioner,
        sample_rate=sample_rate,
        min_input_length=min_input_length,
        cross_attn_cond_ids=cross_attn_cond_ids,
        prepend_cond_ids=prepend_cond_ids,
        global_cond_ids=global_cond_ids
    )

    return model