import torch
import torch.nn as nn

from torchseq.models.encoder import SequenceEncoder
from torchseq.models.decoder import SequenceDecoder
from torchseq.models.bottleneck import PoolingBottleneck
from torchseq.models.modular_bottleneck import ModularBottleneck
from torchseq.utils.logging import Logger
from torchseq.utils.functions import cos_sim
from torchseq.models.vq_code_predictor import VQCodePredictor


class BottleneckAutoencoderModel(nn.Module):
    def __init__(self, config, src_field="s1"):
        super().__init__()
        self.config = config

        # TEMP: deprecation warning
        if self.config.bottleneck.get("num_similar_heads", None) is not None:
            print('num_similar_heads is deprecated! Use "splice_head_offset" instead')

        self.src_field = src_field

        self.seq_encoder = SequenceEncoder(config)

        if config.bottleneck.get("modular", False):
            self.bottleneck = ModularBottleneck(config)
        else:
            self.bottleneck = PoolingBottleneck(config)

        self.seq_decoder = SequenceDecoder(config, embeddings=self.seq_encoder.embeddings)

        if self.config.bottleneck.get("split_encoder", False):
            self.seq_encoder_2 = SequenceEncoder(config)
            self.bottleneck_2 = PoolingBottleneck(config)

            self.split_projection_1 = nn.utils.weight_norm(
                nn.Linear(config.encoder.embedding_dim, config.decoder.embedding_dim // 2, bias=False)
            )
            self.split_projection_2 = nn.utils.weight_norm(
                nn.Linear(config.encoder.embedding_dim, config.decoder.embedding_dim // 2, bias=False)
            )

        if self.config.bottleneck.get("code_predictor", None) is not None:
            pred_config = self.config.bottleneck.code_predictor
            if self.config.bottleneck.get("quantizer_transitions", False):
                vq_transitions = self.bottleneck.quantizer._transitions
            else:
                vq_transitions = None

            if self.config.bottleneck.get("modular", False):
                # HACK!
                vq_embeddings = self.bottleneck.module_list[1].quantizer._embedding
            else:
                vq_embeddings = self.bottleneck.quantizer._embedding

            self.code_predictor = VQCodePredictor(pred_config, transitions=vq_transitions, embeddings=vq_embeddings)

    def forward(self, batch, output, memory=None, tgt_field=None):
        if memory is None:
            memory = {}

        # First pass? Construct the encoding
        if "encoding" not in memory:
            encoding, memory = self.seq_encoder(
                batch[self.src_field],
                batch[self.src_field + "_len"],
                memory,
                include_position=self.config.encoder.get("position_embeddings", True),
            )

            encoding_pooled, memory = self.bottleneck(
                encoding, memory, batch["_global_step"], head_mask=batch.get("head_mask", None)
            )

            raw_encoding_pooled = memory["encoding_pooled"]

            if self.config.bottleneck.get(
                "code_predictor", None
            ) is not None and self.config.bottleneck.code_predictor.get("infer_codes", False):
                pred_codes = self.code_predictor.infer(
                    raw_encoding_pooled.max(dim=1).values.detach(), batch, outputs_to_block=memory.get("vq_codes")
                )
                batch["forced_codes"] = pred_codes

            if self.config.bottleneck.get("split_encoder", False):
                encoding2, memory = self.seq_encoder(batch[self.src_field], batch[self.src_field + "_len"], memory)
                encoding_pooled2, memory = self.bottleneck(
                    encoding2,
                    memory,
                    batch["_global_step"],
                    forced_codes=batch.get("forced_codes", None),
                    head_mask=batch.get("head_mask", None),
                )

                # TODO: Instead of 2x full size encoders + down projection, change to 2x half size encoders
                if self.config.encoder.embedding_dim != self.config.decoder.embedding_dim:
                    encoding_pooled = torch.cat(
                        [self.split_projection_1(encoding_pooled), self.split_projection_2(encoding_pooled2)], -1
                    )
                else:
                    encoding_pooled = torch.cat([encoding_pooled, encoding_pooled2], -1)

            num_heads = self.config.bottleneck.get("num_heads", self.config.encdec.get("num_heads", 1))

            if self.config.bottleneck.get("num_similar_heads", None) is not None:
                # print('num_similar_heads is deprecated! Use "splice_head_offset" instead')
                splice_head_offset = self.config.bottleneck.num_similar_heads
            else:
                splice_head_offset = self.config.bottleneck.get("splice_head_offset", num_heads)

            sep_splice_ix = self.config.bottleneck.embedding_dim // num_heads * splice_head_offset

            if "template" in batch:
                template_memory = {}
                template_encoding, template_memory = self.seq_encoder(
                    batch["template"],
                    batch["template_len"],
                    template_memory,
                    include_position=self.config.encoder.get("template_position_embeddings", True),
                )

                if "forced_templ_encoding" in batch:
                    # print(batch["forced_templ_encoding"].shape, template_encoding)
                    template_encoding_pooled = batch["forced_templ_encoding"].unsqueeze(1)
                    template_memory["encoding_pooled"] = batch["forced_templ_encoding"].unsqueeze(1)
                else:
                    template_encoding_pooled, template_memory = self.bottleneck(
                        template_encoding,
                        template_memory,
                        batch["_global_step"],
                        forced_codes=batch.get("forced_codes", None),
                        head_mask=batch.get("head_mask", None),
                    )

                if "loss" in memory:
                    memory["loss"] += template_memory["loss"]

                if "vq_codes" in template_memory:
                    memory["vq_codes"] = template_memory["vq_codes"]

                if self.config.bottleneck.get("split_encoder", False):
                    template_encoding2, memory = self.seq_encoder(
                        batch["template"], batch["template_len"], template_memory
                    )
                    template_encoding_pooled2, memory = self.bottleneck(
                        template_encoding2, template_memory, batch["_global_step"]
                    )

                    template_encoding_pooled_joint = torch.cat(
                        [
                            self.split_projection_1(template_encoding_pooled),
                            self.split_projection_2(template_encoding_pooled2),
                        ],
                        -1,
                    )

                    template_encoding_pooled = template_encoding_pooled_joint

                if self.config.bottleneck.get("use_templ_encoding", False):
                    if self.config.bottleneck.get("invert_templ", False):
                        raise Exception("invert_templ is no longer supported")
                        encoding_pooled[:, :, :sep_splice_ix] = template_encoding_pooled[:, :, :sep_splice_ix]
                        raw_encoding_pooled[:, :, :sep_splice_ix] = template_memory["encoding_pooled"][
                            :, :, :sep_splice_ix
                        ]
                    else:
                        encoding_pooled = torch.cat(
                            [
                                encoding_pooled[:, :, :sep_splice_ix],
                                template_encoding_pooled[:, :1, sep_splice_ix:].expand(
                                    -1, encoding_pooled.shape[1], -1
                                ),
                            ],
                            dim=2,
                        )
                        templ_raw_enc = template_memory["encoding_pooled"][:, :1, sep_splice_ix:].expand(
                            -1, raw_encoding_pooled.shape[1], -1
                        )
                        raw_encoding_pooled = torch.cat(
                            [raw_encoding_pooled[:, :, :sep_splice_ix], templ_raw_enc], dim=2
                        )

                if self.config.bottleneck.get("separation_loss_weight", 0) > 0:
                    if "loss" not in memory:
                        memory["loss"] = 0

                    # TODO: there must be a cleaner way of flipping the tensor ranges here...
                    if self.config.bottleneck.get("cos_separation_loss", True):
                        diff1 = cos_sim(
                            encoding_pooled[:, :, sep_splice_ix:], template_encoding_pooled[:, :, sep_splice_ix:]
                        )
                        diff2 = cos_sim(
                            encoding_pooled[:, :, :sep_splice_ix], template_encoding_pooled[:, :, :sep_splice_ix]
                        )
                        if self.config.bottleneck.get("flip_separation_loss", False):
                            similarity_loss = 1 - 1 * diff1.mean(dim=-1).mean(dim=-1)
                            dissimilarity_loss = 1 + 1 * diff2.mean(dim=-1).mean(dim=-1)
                        else:
                            similarity_loss = 1 - 1 * diff2.mean(dim=-1).mean(dim=-1)
                            dissimilarity_loss = 1 + 1 * diff1.mean(dim=-1).mean(dim=-1)
                    else:
                        if self.config.bottleneck.get("flip_separation_loss", False):
                            diff1 = (
                                encoding_pooled[:, :, sep_splice_ix:] - template_encoding_pooled[:, :, sep_splice_ix:]
                            )
                            similarity_loss = (diff1 ** 2).mean(dim=-1).mean(dim=-1)

                            diff2 = (
                                encoding_pooled[:, :, :sep_splice_ix] - template_encoding_pooled[:, :, :sep_splice_ix]
                            )
                            dissimilarity_loss = torch.log(1 + 1 / (1e-2 + (diff2) ** 2)).mean(dim=-1).mean(dim=-1)
                        else:
                            diff1 = (
                                encoding_pooled[:, :, :sep_splice_ix] - template_encoding_pooled[:, :, :sep_splice_ix]
                            )
                            similarity_loss = (diff1 ** 2).mean(dim=-1).mean(dim=-1)

                            diff2 = (
                                encoding_pooled[:, :, sep_splice_ix:] - template_encoding_pooled[:, :, sep_splice_ix:]
                            )
                            dissimilarity_loss = torch.log(1 + 1 / (1e-2 + (diff2) ** 2)).mean(dim=-1).mean(dim=-1)

                        similarity_loss *= 10

                    separation_loss = (similarity_loss + dissimilarity_loss) * self.config.bottleneck.get(
                        "separation_loss_weight", 0
                    )

                    if self.training:
                        Logger().log_scalar("bottleneck/sim_loss", similarity_loss.mean(), batch["_global_step"])
                        Logger().log_scalar("bottleneck/dissim_loss", dissimilarity_loss.mean(), batch["_global_step"])

                    memory["loss"] += separation_loss

            memory["encoding"] = encoding_pooled
            memory["encoding_mask"] = None

            # Note: this returns the encodings *before* bottleneck!
            if sep_splice_ix > 0:
                memory["sep_encoding_1"] = (
                    raw_encoding_pooled[:, :, :sep_splice_ix].max(dim=1, keepdim=True).values.detach()
                )
                memory["sep_encoding_1_after_bottleneck"] = (
                    encoding_pooled[:, :, :sep_splice_ix].max(dim=1, keepdim=True).values.detach()
                )
            memory["sep_encoding_2"] = (
                raw_encoding_pooled[:, :, sep_splice_ix:].max(dim=1, keepdim=True).values.detach()
            )
            memory["sep_encoding_2_after_bottleneck"] = (
                encoding_pooled[:, :, sep_splice_ix:].max(dim=1, keepdim=True).values.detach()
            )

        # Fwd pass through decoder block
        logits, memory = self.seq_decoder(output, memory)

        return logits, memory

import torch
import torch.nn as nn

from torchseq.models.encoder import SequenceEncoder
from torchseq.models.decoder import SequenceDecoder
from torchseq.models.bottleneck import PoolingBottleneck
from torchseq.models.modular_bottleneck import ModularBottleneck
from torchseq.utils.logging import Logger
from torchseq.utils.functions import cos_sim
from torchseq.models.vq_code_predictor import VQCodePredictor


class BottleneckAutoencoderModel(nn.Module):
    def __init__(self, config, src_field="s1"):
        super().__init__()
        self.config = config

        # TEMP: deprecation warning
        if self.config.bottleneck.get("num_similar_heads", None) is not None:
            print('num_similar_heads is deprecated! Use "splice_head_offset" instead')

        self.src_field = src_field

        self.seq_encoder = SequenceEncoder(config)

        if config.bottleneck.get("modular", False):
            self.bottleneck = ModularBottleneck(config)
        else:
            self.bottleneck = PoolingBottleneck(config)

        self.seq_decoder = SequenceDecoder(config, embeddings=self.seq_encoder.embeddings)

        if self.config.bottleneck.get("split_encoder", False):
            self.seq_encoder_2 = SequenceEncoder(config)
            self.bottleneck_2 = PoolingBottleneck(config)

            self.split_projection_1 = nn.utils.weight_norm(
                nn.Linear(config.encoder.embedding_dim, config.decoder.embedding_dim // 2, bias=False)
            )
            self.split_projection_2 = nn.utils.weight_norm(
                nn.Linear(config.encoder.embedding_dim, config.decoder.embedding_dim // 2, bias=False)
            )

        if self.config.bottleneck.get("code_predictor", None) is not None:
            pred_config = self.config.bottleneck.code_predictor
            if self.config.bottleneck.get("quantizer_transitions", False):
                vq_transitions = self.bottleneck.quantizer._transitions
            else:
                vq_transitions = None

            if self.config.bottleneck.get("modular", False):
                # HACK!
                vq_embeddings = self.bottleneck.module_list[1].quantizer._embedding
            else:
                vq_embeddings = self.bottleneck.quantizer._embedding

            self.code_predictor = VQCodePredictor(pred_config, transitions=vq_transitions, embeddings=vq_embeddings)

    def forward(self, batch, output, memory=None, tgt_field=None):
        if memory is None:
            memory = {}

        # First pass? Construct the encoding
        if "encoding" not in memory:
            encoding, memory = self.seq_encoder(
                batch[self.src_field],
                batch[self.src_field + "_len"],
                memory,
                include_position=self.config.encoder.get("position_embeddings", True),
            )

            encoding_pooled, memory = self.bottleneck(
                encoding, memory, batch["_global_step"], head_mask=batch.get("head_mask", None)
            )

            raw_encoding_pooled = memory["encoding_pooled"]

            if self.config.bottleneck.get(
                "code_predictor", None
            ) is not None and self.config.bottleneck.code_predictor.get("infer_codes", False):
                pred_codes = self.code_predictor.infer(
                    raw_encoding_pooled.max(dim=1).values.detach(), batch, outputs_to_block=memory.get("vq_codes")
                )
                batch["forced_codes"] = pred_codes

            if self.config.bottleneck.get("split_encoder", False):
                encoding2, memory = self.seq_encoder(batch[self.src_field], batch[self.src_field + "_len"], memory)
                encoding_pooled2, memory = self.bottleneck(
                    encoding2,
                    memory,
                    batch["_global_step"],
                    forced_codes=batch.get("forced_codes", None),
                    head_mask=batch.get("head_mask", None),
                )

                # TODO: Instead of 2x full size encoders + down projection, change to 2x half size encoders
                if self.config.encoder.embedding_dim != self.config.decoder.embedding_dim:
                    encoding_pooled = torch.cat(
                        [self.split_projection_1(encoding_pooled), self.split_projection_2(encoding_pooled2)], -1
                    )
                else:
                    encoding_pooled = torch.cat([encoding_pooled, encoding_pooled2], -1)

            num_heads = self.config.bottleneck.get("num_heads", self.config.encdec.get("num_heads", 1))

            if self.config.bottleneck.get("num_similar_heads", None) is not None:
                # print('num_similar_heads is deprecated! Use "splice_head_offset" instead')
                splice_head_offset = self.config.bottleneck.num_similar_heads
            else:
                splice_head_offset = self.config.bottleneck.get("splice_head_offset", num_heads)

            sep_splice_ix = self.config.bottleneck.embedding_dim // num_heads * splice_head_offset

            if "template" in batch:
                template_memory = {}
                template_encoding, template_memory = self.seq_encoder(
                    batch["template"],
                    batch["template_len"],
                    template_memory,
                    include_position=self.config.encoder.get("template_position_embeddings", True),
                )

                if "forced_templ_encoding" in batch:
                    # print(batch["forced_templ_encoding"].shape, template_encoding)
                    template_encoding_pooled = batch["forced_templ_encoding"].unsqueeze(1)
                    template_memory["encoding_pooled"] = batch["forced_templ_encoding"].unsqueeze(1)
                else:
                    template_encoding_pooled, template_memory = self.bottleneck(
                        template_encoding,
                        template_memory,
                        batch["_global_step"],
                        forced_codes=batch.get("forced_codes", None),
                        head_mask=batch.get("head_mask", None),
                    )

                if "loss" in memory:
                    memory["loss"] += template_memory["loss"]

                if "vq_codes" in template_memory:
                    memory["vq_codes"] = template_memory["vq_codes"]

                if self.config.bottleneck.get("split_encoder", False):
                    template_encoding2, memory = self.seq_encoder(
                        batch["template"], batch["template_len"], template_memory
                    )
                    template_encoding_pooled2, memory = self.bottleneck(
                        template_encoding2, template_memory, batch["_global_step"]
                    )

                    template_encoding_pooled_joint = torch.cat(
                        [
                            self.split_projection_1(template_encoding_pooled),
                            self.split_projection_2(template_encoding_pooled2),
                        ],
                        -1,
                    )

                    template_encoding_pooled = template_encoding_pooled_joint

                if self.config.bottleneck.get("use_templ_encoding", False):
                    if self.config.bottleneck.get("invert_templ", False):
                        raise Exception("invert_templ is no longer supported")
                        encoding_pooled[:, :, :sep_splice_ix] = template_encoding_pooled[:, :, :sep_splice_ix]
                        raw_encoding_pooled[:, :, :sep_splice_ix] = template_memory["encoding_pooled"][
                            :, :, :sep_splice_ix
                        ]
                    else:
                        encoding_pooled = torch.cat(
                            [
                                encoding_pooled[:, :, :sep_splice_ix],
                                template_encoding_pooled[:, :1, sep_splice_ix:].expand(
                                    -1, encoding_pooled.shape[1], -1
                                ),
                            ],
                            dim=2,
                        )
                        templ_raw_enc = template_memory["encoding_pooled"][:, :1, sep_splice_ix:].expand(
                            -1, raw_encoding_pooled.shape[1], -1
                        )
                        raw_encoding_pooled = torch.cat(
                            [raw_encoding_pooled[:, :, :sep_splice_ix], templ_raw_enc], dim=2
                        )

                if self.config.bottleneck.get("separation_loss_weight", 0) > 0:
                    if "loss" not in memory:
                        memory["loss"] = 0

                    # TODO: there must be a cleaner way of flipping the tensor ranges here...
                    if self.config.bottleneck.get("cos_separation_loss", True):
                        diff1 = cos_sim(
                            encoding_pooled[:, :, sep_splice_ix:], template_encoding_pooled[:, :, sep_splice_ix:]
                        )
                        diff2 = cos_sim(
                            encoding_pooled[:, :, :sep_splice_ix], template_encoding_pooled[:, :, :sep_splice_ix]
                        )
                        if self.config.bottleneck.get("flip_separation_loss", False):
                            similarity_loss = 1 - 1 * diff1.mean(dim=-1).mean(dim=-1)
                            dissimilarity_loss = 1 + 1 * diff2.mean(dim=-1).mean(dim=-1)
                        else:
                            similarity_loss = 1 - 1 * diff2.mean(dim=-1).mean(dim=-1)
                            dissimilarity_loss = 1 + 1 * diff1.mean(dim=-1).mean(dim=-1)
                    else:
                        if self.config.bottleneck.get("flip_separation_loss", False):
                            diff1 = (
                                encoding_pooled[:, :, sep_splice_ix:] - template_encoding_pooled[:, :, sep_splice_ix:]
                            )
                            similarity_loss = (diff1 ** 2).mean(dim=-1).mean(dim=-1)

                            diff2 = (
                                encoding_pooled[:, :, :sep_splice_ix] - template_encoding_pooled[:, :, :sep_splice_ix]
                            )
                            dissimilarity_loss = torch.log(1 + 1 / (1e-2 + (diff2) ** 2)).mean(dim=-1).mean(dim=-1)
                        else:
                            diff1 = (
                                encoding_pooled[:, :, :sep_splice_ix] - template_encoding_pooled[:, :, :sep_splice_ix]
                            )
                            similarity_loss = (diff1 ** 2).mean(dim=-1).mean(dim=-1)

                            diff2 = (
                                encoding_pooled[:, :, sep_splice_ix:] - template_encoding_pooled[:, :, sep_splice_ix:]
                            )
                            dissimilarity_loss = torch.log(1 + 1 / (1e-2 + (diff2) ** 2)).mean(dim=-1).mean(dim=-1)

                        similarity_loss *= 10

                    separation_loss = (similarity_loss + dissimilarity_loss) * self.config.bottleneck.get(
                        "separation_loss_weight", 0
                    )

                    if self.training:
                        Logger().log_scalar("bottleneck/sim_loss", similarity_loss.mean(), batch["_global_step"])
                        Logger().log_scalar("bottleneck/dissim_loss", dissimilarity_loss.mean(), batch["_global_step"])

                    memory["loss"] += separation_loss

            memory["encoding"] = encoding_pooled
            memory["encoding_mask"] = None

            # Note: this returns the encodings *before* bottleneck!
            if sep_splice_ix > 0:
                memory["sep_encoding_1"] = (
                    raw_encoding_pooled[:, :, :sep_splice_ix].max(dim=1, keepdim=True).values.detach()
                )
                memory["sep_encoding_1_after_bottleneck"] = (
                    encoding_pooled[:, :, :sep_splice_ix].max(dim=1, keepdim=True).values.detach()
                )
            memory["sep_encoding_2"] = (
                raw_encoding_pooled[:, :, sep_splice_ix:].max(dim=1, keepdim=True).values.detach()
            )
            memory["sep_encoding_2_after_bottleneck"] = (
                encoding_pooled[:, :, sep_splice_ix:].max(dim=1, keepdim=True).values.detach()
            )

        # Fwd pass through decoder block
        logits, memory = self.seq_decoder(output, memory)

        return logits, memory

# TEMP: deprecation warning # HACK! # First pass? Construct the encoding # TODO: Instead of 2x full size encoders + down projection, change to 2x half size encoders # print('num_similar_heads is deprecated! Use "splice_head_offset" instead') # print(batch["forced_templ_encoding"].shape, template_encoding) # TODO: there must be a cleaner way of flipping the tensor ranges here... # Note: this returns the encodings *before* bottleneck! # Fwd pass through decoder block