# coding=utf-8
# Copyright 2024 The Google Research Authors.
#
# 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.

"""Model Architecture."""

import functools
from typing import Any, Callable

from flax import linen as nn
from flax import struct
from jax import numpy as jnp

@struct.dataclass
class TransformerConfig:
    """Global hyperparameters used to minimize obnoxious kwarg plumbing."""
    vocab_size: int = 1
    dtype: Any = jnp.float32
    emb_dim: int = 512
    num_heads: int = 8
    num_layers: int = 6
    qkv_dim: int = 512
    mlp_dim: int = 2048
    seq_len: int = 2048  # Maximum sequence length
    dropout_rate: float = 0.1
    attention_dropout_rate: float = 0.1
    deterministic: bool = False
    num_latent_slots: int = 0  # K continuous latent thought slots
    # When False, the fed-back latent thought vectors are NOT injected into the
    # slot positions (they keep their placeholder token embedding). This turns
    # the K slots into static, parallel per-stage readouts with NO carried
    # recurrent state -- the "stagewise supervision, no recurrence" control.
    inject_latents: bool = True

    
class TransformerBlock(nn.Module):
    config: Any = None

    def setup(self):
        self.vocab_size = self.config.vocab_size
        self.emb_dim = self.config.emb_dim
        self.num_layers = self.config.num_layers

    @nn.compact
    def __call__(self, inputs, causal_mask_inputs, training=True):
        """
        Transformer Block call function.

        Args:
            inputs: Input tensor.
            causal_mask_inputs: Causal mask for the inputs.
            training: Whether the model is in training mode.

        Returns:
            Transformed tensor after self-attention and MLP layers.
        """
        
        x = inputs + nn.SelfAttention(
            num_heads=self.config.num_heads, dtype=self.config.dtype,
            qkv_features=self.config.qkv_dim,
            kernel_init=nn.initializers.xavier_uniform(),
            bias_init=nn.initializers.normal(stddev=1e-6), 
            use_bias=False, broadcast_dropout=False,
            dropout_rate=self.config.attention_dropout_rate, normalize_qk=True,
            deterministic=self.config.deterministic)(inputs, causal_mask_inputs)

        def mlp(x):
            """
            Multi-Layer Perceptron function.

            Args:
                x: Input tensor.

            Returns:
                Transformed tensor after applying MLP layers.
            """
            dense_with_init = functools.partial(
                nn.Dense,
                kernel_init=nn.initializers.xavier_uniform(),
                bias_init=nn.initializers.normal(stddev=1e-6)
                )
            x = dense_with_init(features=self.config.mlp_dim)(x)
            x = nn.gelu(x)
            x = dense_with_init(features=self.config.emb_dim)(x)
            x = nn.Dropout(rate=self.config.dropout_rate,
                        deterministic=self.config.deterministic)(x)
            return x

        x = x + mlp(x)
        return x


class TransformerLMHeadModel(nn.Module):
    config: Any = None

    def setup(self):
        self.vocab_size = self.config.vocab_size
        self.emb_dim = self.config.emb_dim
        self.num_layers = self.config.num_layers

    @nn.compact
    def __call__(self, inputs, latent_values=None, latent_positions=None,
                 latent_active=None, training=True):
        """
        Transformer LM Head call function.

        Args:
            inputs: Input token ids (batch, seq).
            latent_values: Optional (batch, K, emb_dim) continuous thought
                vectors (raw last-layer hiddens fed back, Coconut/ATC style).
            latent_positions: Optional (batch, K) int positions of the latent
                slots in the sequence (per-example, after the clue block).
            latent_active: Optional (batch, K) bool; slot j uses the projected
                latent vector when True, otherwise keeps the placeholder
                token embedding.
            training: Whether the model is in training mode.

        Returns:
            (logits, hidden): LM logits and final (post-LayerNorm) hidden
            states, the latter used to build the next continuous thought.
        """
        batch_size, seq_size = inputs.shape

        causal_mask_x = nn.make_causal_mask(inputs, dtype=self.config.dtype)

        # Embed the input tensor using a learnable embedding matrix.
        embed_with_init = functools.partial(
            nn.Embed, embedding_init=nn.initializers.normal(stddev=0.02))
        token_embeddings = embed_with_init(
            num_embeddings=self.config.vocab_size,
            features=self.config.emb_dim,
        )(inputs)

        # Check the shape of the embedded tensor.
        assert token_embeddings.shape == (batch_size, seq_size,
                                      self.config.emb_dim)

        # Continuous latent thoughts: project fed-back hidden states and
        # scatter them into the latent slot positions, replacing the
        # placeholder token embedding (position embeddings still added below).
        # Skipped when inject_latents is False (no-recurrence control: the slots
        # stay as static placeholders and the candidate heads become parallel
        # per-stage readouts with no carried state).
        if latent_values is not None and self.config.inject_latents:
            proj = nn.Dense(features=self.config.emb_dim,
                            kernel_init=nn.initializers.xavier_uniform(),
                            name="latent_proj_in")(latent_values)
            proj = nn.gelu(proj)
            proj = nn.Dense(features=self.config.emb_dim,
                            kernel_init=nn.initializers.xavier_uniform(),
                            name="latent_proj_out")(proj)
            bidx = jnp.arange(batch_size)[:, None]
            cur = token_embeddings[bidx, latent_positions]
            new = jnp.where(latent_active[..., None],
                            proj.astype(cur.dtype), cur)
            token_embeddings = token_embeddings.at[
                bidx, latent_positions].set(new)

        # Initialize the positional embedding variable.
        pos_embedding_variable = self.variable(
            "params",
            "position_embeddings",
            jnp.zeros,
            (self.config.seq_len, self.config.emb_dim),
        )

        # Slice the positional embedding array to the correct sequence length.
        pos_embeddings = pos_embedding_variable.value[:seq_size, :]

        # Check the shape of the positional embedding array.
        output_tuple = (pos_embeddings.shape, token_embeddings.shape[1:])
        assert pos_embeddings.shape == token_embeddings.shape[1:], output_tuple

        # Add the positional embeddings to the token embeddings.
        x = token_embeddings + pos_embeddings[None, :, :]

        # Apply dropout to the input.
        x = nn.Dropout(rate=self.config.dropout_rate,
                    deterministic=self.config.deterministic)(x)

        # Apply the Transformer layers. remat (gradient checkpointing) keeps
        # the multi-pass latent recurrence within GPU memory under full BPTT.
        RematBlock = nn.remat(TransformerBlock)
        for i in range(self.num_layers):
            x = RematBlock(config=self.config)(
                    x, causal_mask_x, training=training)
      
            self.sow('intermediates', 'feature_' + str(i), x)

        # Apply the final layer normalization.
        x = nn.LayerNorm()(x)

        # Apply the LM head.
        logits = nn.Dense(features=self.config.vocab_size,
                        kernel_init=nn.initializers.xavier_uniform(),
                        bias_init=nn.initializers.normal(stddev=1e-6),
                        use_bias=False)(x)

        # Check the shape of the output tensor.
        assert logits.shape == (batch_size, seq_size, self.config.vocab_size)

        # ---- Auxiliary multi-candidate head ----
        # Each latent slot reads its (post-LayerNorm) hidden and predicts the
        # full 81x9 candidate grid for its reasoning stage. Trained with BCE
        # (independent per-digit sigmoids = candidate-set membership), NOT
        # softmax, so multiple digits can be "on" at intermediate stages.
        cand_logits = None
        if latent_positions is not None:
            bidx = jnp.arange(batch_size)[:, None]
            slot_hidden = x[bidx, latent_positions]          # (bs, K, emb)
            h = nn.Dense(features=self.config.emb_dim,
                         kernel_init=nn.initializers.xavier_uniform(),
                         name="cand_head_in")(slot_hidden)
            h = nn.gelu(h)
            cand_logits = nn.Dense(features=81 * 9,
                                   kernel_init=nn.initializers.xavier_uniform(),
                                   name="cand_head_out")(h)   # (bs, K, 729)
            cand_logits = cand_logits.reshape(
                batch_size, -1, 81, 9)                        # (bs, K, 81, 9)

        return logits, x, cand_logits