#!/usr/bin/env python3
"""
Debug Transformer User Tower

Debug version of TransformerUserTower with detailed print statements at each layer
to help identify issues with transformer recommendation inference.
"""

import tensorflow as tf
import numpy as np
from typing import Optional, List, Dict
from src.models.transformer.transformer_user_tower import (
    TransformerEncoder,
    ConcatFusion,
    AttentionPooling
)


class DebugTransformerUserTower(tf.keras.Model):
    """Debug version of transformer user tower with detailed layer-by-layer output printing."""

    def __init__(self,
                 embedding_dim: int = 128,
                 transformer_layers: int = 2,
                 transformer_heads: int = 4,
                 transformer_ff_dim: int = 256,
                 hidden_dims: List[int] = [128, 64],
                 dropout_rate: float = 0.2,
                 max_sequence_length: Optional[int] = None,
                 name: str = "debug_transformer_user_tower",
                 compatibility_mode: bool = False):
        super().__init__(name=name)

        self.embedding_dim = embedding_dim
        self.max_sequence_length = max_sequence_length

        print(f"\n🔧 INITIALIZING DEBUG TRANSFORMER USER TOWER")
        print(f"   Embedding dim: {embedding_dim}")
        print(f"   Transformer layers: {transformer_layers}")
        print(f"   Transformer heads: {transformer_heads}")
        print(f"   Max sequence length: {max_sequence_length}")

        # Demographic embeddings (balanced dimensions to preserve information flow equality)
        self.age_embedding = tf.keras.layers.Embedding(6, 8, name="age_embedding")
        self.income_embedding = tf.keras.layers.Embedding(5, 8, name="income_embedding")
        self.gender_embedding = tf.keras.layers.Embedding(2, 8, name="gender_embedding")
        self.profession_embedding = tf.keras.layers.Embedding(8, 8, name="profession_embedding")
        self.location_embedding = tf.keras.layers.Embedding(3, 8, name="location_embedding")
        self.education_embedding = tf.keras.layers.Embedding(5, 8, name="education_embedding")
        self.marital_embedding = tf.keras.layers.Embedding(4, 8, name="marital_embedding")

        # Transformer encoder for interaction sequences
        self.transformer_encoder = TransformerEncoder(
            num_layers=transformer_layers,
            num_heads=transformer_heads,
            model_dim=embedding_dim,
            ff_dim=transformer_ff_dim,
            dropout_rate=dropout_rate,
            name="interaction_transformer"
        )

        # Attention-based sequence aggregation
        self.sequence_pooling = AttentionPooling(embedding_dim, name="sequence_pooling")

        # Concatenation fusion for demographic and behavioral features with cold-start support
        self.concat_fusion = ConcatFusion(embedding_dim, name="concat_fusion")

        # Dense layers for combining features
        dense_layer_list = []
        for i, dim in enumerate(hidden_dims):
            dense_layer_list.extend([
                tf.keras.layers.Dense(
                    dim,
                    activation="relu",
                    name=f"user_dense_{i}"
                ),
                tf.keras.layers.Dropout(dropout_rate, name=f"user_dropout_{i}")
            ])

        self.dense_layers = tf.keras.Sequential(dense_layer_list, name="user_dense_stack")

        # Output layer without L2 regularization
        self.output_layer = tf.keras.layers.Dense(
            embedding_dim,
            activation=None,
            name="user_output"
        )

        # Sequence length processing (compatibility mode for old weights)
        if compatibility_mode:
            # Old architecture: simple embedding lookup
            self.sequence_length_embedding = tf.keras.layers.Embedding(
                max_sequence_length or 200, 4, name="sequence_length_embedding"
            )
        else:
            # New architecture: sequential dense processing
            self.sequence_length_processor = tf.keras.Sequential([
                tf.keras.layers.Dense(embedding_dim // 8, activation='relu', name="seq_len_dense_1"),
                tf.keras.layers.Dense(embedding_dim // 16, activation='relu', name="seq_len_dense_2")
            ], name="sequence_length_processor")

        self.compatibility_mode = compatibility_mode

        # Add layer normalization for regularization
        self.demographic_layer_norm = tf.keras.layers.LayerNormalization(
            name="demographic_layer_norm"
        )
        self.history_layer_norm = tf.keras.layers.LayerNormalization(
            name="history_layer_norm"
        )

    def call(self, inputs, training=None):
        """
        Forward pass of the debug transformer user tower with detailed prints.
        """
        print(f"\n🎯 DEBUG TRANSFORMER USER TOWER FORWARD PASS")
        print(f"   Training mode: {training}")

        # Extract inputs
        age = inputs["age"]
        gender = inputs["gender"]
        income = inputs["income"]
        profession = inputs["profession"]
        location = inputs["location"]
        education = inputs["education_level"]
        marital_status = inputs["marital_status"]
        item_history = inputs["item_history_embeddings"]
        attention_mask = inputs.get("attention_masks")
        sequence_lengths = inputs.get("sequence_lengths")

        print(f"\n📋 INPUT ANALYSIS:")
        print(f"   Age: {age.numpy()}")
        print(f"   Gender: {gender.numpy()}")
        print(f"   Income: {income.numpy()}")
        print(f"   Profession: {profession.numpy()}")
        print(f"   Location: {location.numpy()}")
        print(f"   Education: {education.numpy()}")
        print(f"   Marital status: {marital_status.numpy()}")
        print(f"   Item history shape: {item_history.shape}")
        print(f"   Attention mask shape: {attention_mask.shape if attention_mask is not None else 'None'}")
        print(f"   Sequence lengths: {sequence_lengths.numpy() if sequence_lengths is not None else 'None'}")

        # Check if item history contains meaningful data
        history_sum = tf.reduce_sum(tf.abs(item_history))
        print(f"   Item history total magnitude: {history_sum.numpy():.4f}")

        # Process demographics through embeddings
        print(f"\n🏷️ DEMOGRAPHIC EMBEDDINGS:")
        age_emb = self.age_embedding(age)
        print(f"   Age embedding shape: {age_emb.shape}, values: {age_emb.numpy().flatten()[:5]}...")

        income_emb = self.income_embedding(income)
        print(f"   Income embedding shape: {income_emb.shape}, values: {income_emb.numpy().flatten()[:5]}...")

        gender_emb = self.gender_embedding(gender)
        print(f"   Gender embedding shape: {gender_emb.shape}, values: {gender_emb.numpy().flatten()[:5]}...")

        profession_emb = self.profession_embedding(profession)
        print(f"   Profession embedding shape: {profession_emb.shape}, values: {profession_emb.numpy().flatten()[:5]}...")

        location_emb = self.location_embedding(location)
        print(f"   Location embedding shape: {location_emb.shape}, values: {location_emb.numpy().flatten()[:5]}...")

        education_emb = self.education_embedding(education)
        print(f"   Education embedding shape: {education_emb.shape}, values: {education_emb.numpy().flatten()[:5]}...")

        marital_emb = self.marital_embedding(marital_status)
        print(f"   Marital embedding shape: {marital_emb.shape}, values: {marital_emb.numpy().flatten()[:5]}...")

        # Process interaction sequences through transformer with proper masking
        print(f"\n🔄 ATTENTION MASK PROCESSING:")
        if attention_mask is None:
            # Create mask from item history: 1 for valid tokens, 0 for padding
            attention_mask = tf.reduce_any(tf.not_equal(item_history, 0.0), axis=-1)
            attention_mask = tf.cast(attention_mask, tf.float32)
            print(f"   Generated attention mask from item history")

        print(f"   Attention mask shape: {attention_mask.shape}")
        print(f"   Valid tokens count: {tf.reduce_sum(attention_mask, axis=1).numpy()}")
        print(f"   Sample attention mask: {attention_mask.numpy()[0][:10]}...")

        # Apply transformer encoder
        print(f"\n🤖 TRANSFORMER ENCODER:")
        print(f"   Input item history stats:")
        print(f"     - Shape: {item_history.shape}")
        print(f"     - Min: {tf.reduce_min(item_history).numpy():.6f}")
        print(f"     - Max: {tf.reduce_max(item_history).numpy():.6f}")
        print(f"     - Mean: {tf.reduce_mean(item_history).numpy():.6f}")
        print(f"     - Std: {tf.math.reduce_std(item_history).numpy():.6f}")

        transformed_history = self.transformer_encoder(
            item_history,
            attention_mask=attention_mask,
            training=training
        )

        print(f"   Transformer output stats:")
        print(f"     - Shape: {transformed_history.shape}")
        print(f"     - Min: {tf.reduce_min(transformed_history).numpy():.6f}")
        print(f"     - Max: {tf.reduce_max(transformed_history).numpy():.6f}")
        print(f"     - Mean: {tf.reduce_mean(transformed_history).numpy():.6f}")
        print(f"     - Std: {tf.math.reduce_std(transformed_history).numpy():.6f}")

        # Use attention-based pooling for better sequence aggregation
        print(f"\n🎯 ATTENTION POOLING:")
        history_aggregated = self.sequence_pooling(transformed_history, mask=attention_mask)
        print(f"   Pooled history stats:")
        print(f"     - Shape: {history_aggregated.shape}")
        print(f"     - Min: {tf.reduce_min(history_aggregated).numpy():.6f}")
        print(f"     - Max: {tf.reduce_max(history_aggregated).numpy():.6f}")
        print(f"     - Mean: {tf.reduce_mean(history_aggregated).numpy():.6f}")
        print(f"     - Std: {tf.math.reduce_std(history_aggregated).numpy():.6f}")
        print(f"     - First 5 values: {history_aggregated.numpy()[0][:5]}")

        # Apply layer normalization to history
        history_aggregated = self.history_layer_norm(history_aggregated)
        print(f"   After layer norm: mean={tf.reduce_mean(history_aggregated).numpy():.6f}, std={tf.math.reduce_std(history_aggregated).numpy():.6f}")

        # Add sequence length as a continuous feature
        print(f"\n📏 SEQUENCE LENGTH PROCESSING:")
        if self.compatibility_mode:
            # Old architecture: use embedding lookup
            if sequence_lengths is not None:
                clipped_lengths = tf.clip_by_value(
                    tf.cast(sequence_lengths, tf.int32),
                    0, (self.max_sequence_length or 200) - 1
                )
                seq_len_emb = self.sequence_length_embedding(clipped_lengths)
                print(f"   Using provided sequence lengths: {clipped_lengths.numpy()}")
            else:
                actual_lengths = tf.reduce_sum(tf.cast(attention_mask, tf.float32), axis=1)
                clipped_lengths = tf.clip_by_value(
                    tf.cast(actual_lengths, tf.int32),
                    0, (self.max_sequence_length or 200) - 1
                )
                seq_len_emb = self.sequence_length_embedding(clipped_lengths)
                print(f"   Calculated sequence lengths: {clipped_lengths.numpy()}")
        else:
            # New architecture: use dense processing
            if sequence_lengths is not None:
                seq_len_emb = self.sequence_length_processor(
                    tf.cast(sequence_lengths, tf.float32)[:, tf.newaxis]
                )
                print(f"   Using provided sequence lengths: {sequence_lengths.numpy()}")
            else:
                actual_lengths = tf.reduce_sum(tf.cast(attention_mask, tf.float32), axis=1)
                seq_len_emb = self.sequence_length_processor(actual_lengths[:, tf.newaxis])
                print(f"   Calculated sequence lengths: {actual_lengths.numpy()}")

        print(f"   Sequence length embedding shape: {seq_len_emb.shape}")
        print(f"   Sequence length embedding values: {seq_len_emb.numpy()[0][:5]}...")

        # Combine demographic features and apply layer normalization
        print(f"\n👥 DEMOGRAPHIC COMBINATION:")
        demographics_combined = tf.concat([
            age_emb,
            income_emb,
            gender_emb,
            profession_emb,
            location_emb,
            education_emb,
            marital_emb
        ], axis=-1)

        print(f"   Combined demographics shape: {demographics_combined.shape}")
        print(f"   Demographics stats before norm: mean={tf.reduce_mean(demographics_combined).numpy():.6f}, std={tf.math.reduce_std(demographics_combined).numpy():.6f}")

        demographics_combined = self.demographic_layer_norm(demographics_combined)
        print(f"   Demographics stats after norm: mean={tf.reduce_mean(demographics_combined).numpy():.6f}, std={tf.math.reduce_std(demographics_combined).numpy():.6f}")

        # Use concatenation fusion with cold-start fallback
        print(f"\n🔗 CONCATENATION FUSION:")
        print(f"   Input demographics shape: {demographics_combined.shape}")
        print(f"   Input behavioral shape: {history_aggregated.shape}")

        fused_features = self.concat_fusion(
            demographic_features=demographics_combined,
            behavioral_features=history_aggregated,
            sequence_lengths=sequence_lengths,
            training=training
        )

        print(f"   Fused features stats:")
        print(f"     - Shape: {fused_features.shape}")
        print(f"     - Min: {tf.reduce_min(fused_features).numpy():.6f}")
        print(f"     - Max: {tf.reduce_max(fused_features).numpy():.6f}")
        print(f"     - Mean: {tf.reduce_mean(fused_features).numpy():.6f}")
        print(f"     - Std: {tf.math.reduce_std(fused_features).numpy():.6f}")

        # Combine fused features with sequence length embedding
        print(f"\n🔧 FINAL COMBINATION:")
        combined = tf.concat([
            fused_features,
            seq_len_emb
        ], axis=-1)

        print(f"   Combined features shape: {combined.shape}")
        print(f"   Combined features stats: mean={tf.reduce_mean(combined).numpy():.6f}, std={tf.math.reduce_std(combined).numpy():.6f}")

        # Pass through dense layers
        print(f"\n🧠 DENSE LAYERS:")
        x = self.dense_layers(combined, training=training)
        print(f"   After dense layers:")
        print(f"     - Shape: {x.shape}")
        print(f"     - Min: {tf.reduce_min(x).numpy():.6f}")
        print(f"     - Max: {tf.reduce_max(x).numpy():.6f}")
        print(f"     - Mean: {tf.reduce_mean(x).numpy():.6f}")
        print(f"     - Std: {tf.math.reduce_std(x).numpy():.6f}")

        # Final output
        output = self.output_layer(x)
        print(f"\n📤 OUTPUT LAYER:")
        print(f"   Raw output stats:")
        print(f"     - Shape: {output.shape}")
        print(f"     - Min: {tf.reduce_min(output).numpy():.6f}")
        print(f"     - Max: {tf.reduce_max(output).numpy():.6f}")
        print(f"     - Mean: {tf.reduce_mean(output).numpy():.6f}")
        print(f"     - Std: {tf.math.reduce_std(output).numpy():.6f}")

        # L2 normalize for consistent similarity calculations
        normalized_embedding = tf.nn.l2_normalize(output, axis=-1)

        print(f"   Normalized output stats:")
        print(f"     - Shape: {normalized_embedding.shape}")
        print(f"     - L2 norm: {tf.norm(normalized_embedding, axis=-1).numpy()}")
        print(f"     - Min: {tf.reduce_min(normalized_embedding).numpy():.6f}")
        print(f"     - Max: {tf.reduce_max(normalized_embedding).numpy():.6f}")
        print(f"     - Mean: {tf.reduce_mean(normalized_embedding).numpy():.6f}")
        print(f"     - First 5 values: {normalized_embedding.numpy()[0][:5]}")

        print(f"\n✅ DEBUG USER TOWER FORWARD PASS COMPLETE")

        return normalized_embedding