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

Debug version of TransformerItemTower 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 List, Dict


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

    def __init__(self,
                 item_vocab_size: int,
                 category_vocab_size: int,
                 category_code_vocab_size: int,
                 brand_vocab_size: int,
                 embedding_dim: int = 128,
                 hidden_dims: List[int] = [256, 128],
                 dropout_rate: float = 0.2,
                 price_mean: float = 0.0,
                 price_std: float = 1.0,
                 name: str = "debug_transformer_item_tower"):
        super().__init__(name=name)

        self.embedding_dim = embedding_dim

        print(f"\n🔧 INITIALIZING DEBUG TRANSFORMER ITEM TOWER")
        print(f"   Item vocab size: {item_vocab_size}")
        print(f"   Category vocab size: {category_vocab_size}")
        print(f"   Category code vocab size: {category_code_vocab_size}")
        print(f"   Brand vocab size: {brand_vocab_size}")
        print(f"   Embedding dim: {embedding_dim}")
        print(f"   Price normalization: mean={price_mean:.4f}, std={price_std:.4f}")

        # Embedding dimensions chosen for efficiency
        self.product_embedding_dim = 56
        self.category_embedding_dim = 16
        self.brand_embedding_dim = 16
        self.price_embedding_dim = 16

        # Embedding layers
        self.item_embedding = tf.keras.layers.Embedding(
            item_vocab_size, self.product_embedding_dim, name="item_embedding"
        )
        self.category_embedding = tf.keras.layers.Embedding(
            category_vocab_size, self.category_embedding_dim, name="category_embedding"
        )
        self.category_code_embedding = tf.keras.layers.Embedding(
            category_code_vocab_size, self.category_embedding_dim, name="category_code_embedding"
        )
        self.brand_embedding = tf.keras.layers.Embedding(
            brand_vocab_size, self.brand_embedding_dim, name="brand_embedding"
        )

        # Price scaling using explicit mean/std
        self.price_mean = tf.constant(price_mean, dtype=tf.float32, name="price_mean")
        self.price_std = tf.constant(price_std, dtype=tf.float32, name="price_std")
        self.price_mlp = tf.keras.Sequential([
            tf.keras.layers.Dense(32, activation="relu", name="price_dense1"),
            tf.keras.layers.Dropout(dropout_rate / 2, name="price_dropout"),
            tf.keras.layers.Dense(self.price_embedding_dim, activation=None, name="price_dense2")
        ], name="price_mlp")

        # Dense hidden layers
        dense_layers = []
        for i, dim in enumerate(hidden_dims):
            dense_layers.extend([
                tf.keras.layers.Dense(dim, activation="relu", name=f"item_dense_{i}"),
                tf.keras.layers.Dropout(dropout_rate, name=f"item_dropout_{i}")
            ])
        self.dense_stack = tf.keras.Sequential(dense_layers, name="item_dense_stack")

        # Output projection
        self.output_layer = tf.keras.layers.Dense(
            embedding_dim, activation=None, name="item_output"
        )

    def _preprocess_price(self, price: tf.Tensor) -> tf.Tensor:
        """Log transform → explicit standardization → MLP → price embedding with debug prints."""
        print(f"\n💰 PRICE PREPROCESSING:")
        print(f"   Raw price: {price.numpy()}")

        log_price = tf.math.log1p(price)
        print(f"   Log(1+price): {log_price.numpy()}")

        # Explicit standardization using precomputed mean/std
        norm_price = (log_price - self.price_mean) / (self.price_std + 1e-6)
        print(f"   Normalized price: {norm_price.numpy()}")
        print(f"   Using mean={self.price_mean.numpy():.4f}, std={self.price_std.numpy():.4f}")

        price_input = tf.expand_dims(norm_price, -1)
        price_embedding = self.price_mlp(price_input)
        print(f"   Price embedding shape: {price_embedding.shape}")
        print(f"   Price embedding values: {price_embedding.numpy()[0][:5]}...")

        return price_embedding

    def call(self, inputs: Dict[str, tf.Tensor], training: bool = None) -> tf.Tensor:
        """Forward pass of debug item tower with detailed prints."""

        print(f"\n🏪 DEBUG TRANSFORMER ITEM TOWER FORWARD PASS")
        print(f"   Training mode: {training}")

        print(f"\n📋 INPUT ANALYSIS:")
        print(f"   Product ID: {inputs['product_id'].numpy()}")
        print(f"   Category ID: {inputs['category_id'].numpy()}")
        print(f"   Category Code ID: {inputs.get('category_code_id', inputs['category_id']).numpy()}")
        print(f"   Brand ID: {inputs['brand_id'].numpy()}")
        print(f"   Price: {inputs['price'].numpy()}")

        # Generate embeddings with debug info
        print(f"\n🏷️ ITEM EMBEDDINGS:")

        item_emb = self.item_embedding(inputs["product_id"])
        print(f"   Product embedding shape: {item_emb.shape}, values: {item_emb.numpy()[0][:5]}...")

        cat_emb = self.category_embedding(inputs["category_id"])
        print(f"   Category embedding shape: {cat_emb.shape}, values: {cat_emb.numpy()[0][:5]}...")

        cat_code_emb = self.category_code_embedding(
            inputs.get("category_code_id", inputs["category_id"])
        )
        print(f"   Category code embedding shape: {cat_code_emb.shape}, values: {cat_code_emb.numpy()[0][:5]}...")

        brand_emb = self.brand_embedding(inputs["brand_id"])
        print(f"   Brand embedding shape: {brand_emb.shape}, values: {brand_emb.numpy()[0][:5]}...")

        price_emb = self._preprocess_price(inputs["price"])

        # Concatenate: 56 + 16 + 16 + 16 + 16 = 120D
        print(f"\n🔗 FEATURE CONCATENATION:")
        combined = tf.concat([item_emb, cat_emb, cat_code_emb, brand_emb, price_emb], axis=-1)
        print(f"   Combined features shape: {combined.shape}")
        print(f"   Expected: 56 + 16 + 16 + 16 + 16 = 120 dimensions")
        print(f"   Combined stats:")
        print(f"     - Min: {tf.reduce_min(combined).numpy():.6f}")
        print(f"     - Max: {tf.reduce_max(combined).numpy():.6f}")
        print(f"     - Mean: {tf.reduce_mean(combined).numpy():.6f}")
        print(f"     - Std: {tf.math.reduce_std(combined).numpy():.6f}")

        # Dense stack → output layer
        print(f"\n🧠 DENSE PROCESSING:")
        x = self.dense_stack(combined, training=training)
        print(f"   After dense stack:")
        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}")

        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}")

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

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

        return normalized_output