src/phase4.py

#!/usr/bin/env python3
"""
Phase 4 — Variant Assignment / Variant Modeling
Compliments Reference DB Pipeline

Input: Phase 3 outputs + Phase 2 output
Output: Variant tables and mappings

This phase transforms the Phase 3 product-group representation into a
structured product variant model, where products belonging to the same
reference product group can have different purchasable variants such as
size, package quantity, formulation, flavour, fat level, etc.

CRITICAL: This phase does NOT use nutrition data.
CRITICAL: This phase does NOT use any external taxonomy.
CRITICAL: This phase is 100% deterministic (no LLM).
"""

import json
import hashlib
import sys
from datetime import datetime, timezone
from pathlib import Path

import pandas as pd
import numpy as np

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------
BASE_DIR = Path(__file__).resolve().parent.parent
PHASE3_DIR = BASE_DIR.parent / "phase3"
PHASE2_DIR = BASE_DIR.parent / "phase2"
OUTPUT_DIR = BASE_DIR / "outputs"
VALIDATION_DIR = BASE_DIR / "validation"
STATISTICS_DIR = BASE_DIR / "statistics"
AUDIT_DIR = BASE_DIR / "audit"

VERSION = "1.0.0"
TIMESTAMP = datetime.now(timezone.utc).isoformat()


def log(msg: str) -> None:
    print(f"[Phase4] {msg}")


# ---------------------------------------------------------------------------
# 1. VARIANT DEFINITION
# ---------------------------------------------------------------------------
# A variant is a purchasable representation of the same reference product
# that differs in attributes such as:
#   - Size (500 g, 1 kg, 2 kg)
#   - Package quantity (1 × 500 g, 2 × 500 g)
#   - Pack configuration (single, multipack, case)
#
# Variants within the same group share the same identity attributes:
#   - is_organic, is_gluten_free, is_naturally_simple, etc.
#   - product_line
#   - flavour
#   - formulation
#   - fat_level
#
# Variants differ only in size/package attributes.


# ---------------------------------------------------------------------------
# 2. VARIANT KEY GENERATION
# ---------------------------------------------------------------------------

def build_variant_key(row: dict) -> str:
    """
    Build a deterministic variant key from normalized variant attributes.
    
    The variant key captures size/package information that distinguishes
    variants within the same reference product group.
    
    Structure: amount|unit|qty|count
    """
    va_str = row.get("variant_attributes", "{}")
    if isinstance(va_str, str):
        try:
            va = json.loads(va_str)
        except (json.JSONDecodeError, TypeError):
            va = {}
    elif isinstance(va_str, dict):
        va = va_str
    else:
        va = {}
    
    # Extract variant-distinguishing attributes
    amount = va.get("amount")
    unit = va.get("unit", "")
    qty = va.get("qty", 1)
    count = va.get("count")
    multiplier = va.get("multiplier", 1)
    raw = va.get("raw", "")
    
    # Build key parts
    parts = []
    
    if count is not None:
        # Count-based variant (e.g., "20 per pack")
        parts.append(f"count{int(count)}")
    elif amount is not None:
        # Size-based variant
        parts.append(f"amt{amount}")
        if unit:
            parts.append(f"unit{unit}")
        if qty and qty > 1:
            parts.append(f"qty{int(qty)}")
        if multiplier and multiplier > 1:
            parts.append(f"mult{int(multiplier)}")
    elif raw:
        # Fallback to raw string
        parts.append(f"raw{raw.strip().lower()}")
    else:
        # No size information
        parts.append("nosize")
    
    return "|".join(parts)


def generate_variant_id(group_id: str, variant_key: str) -> str:
    """
    Generate deterministic variant_id from group_id + variant_key.
    
    Uses UUID5 (namespace-based) for reproducibility.
    """
    import uuid
    namespace = uuid.NAMESPACE_DNS
    name = f"{group_id}||{variant_key}"
    return str(uuid.uuid5(namespace, name))


# ---------------------------------------------------------------------------
# 3. VARIANT ASSIGNMENT
# ---------------------------------------------------------------------------

def assign_variants(df: pd.DataFrame) -> pd.DataFrame:
    """
    Assign variant_ids to products within each group.
    
    Products in the same group with the same variant_key get the same variant_id.
    Products in the same group with different variant_keys get different variant_ids.
    """
    log("Building variant keys ...")
    df["variant_key"] = df.apply(lambda row: build_variant_key(row.to_dict()), axis=1)
    
    log("Generating variant IDs ...")
    df["variant_id"] = df.apply(
        lambda row: generate_variant_id(row["group_id"], row["variant_key"]),
        axis=1
    )
    
    n_variants = df["variant_id"].nunique()
    log(f"Total unique variants: {n_variants}")
    
    return df


# ---------------------------------------------------------------------------
# 4. OUTPUT GENERATION
# ---------------------------------------------------------------------------

def build_variant_table(df: pd.DataFrame) -> pd.DataFrame:
    """
    Build the main variant table: one row per unique variant.
    """
    log("Building variant table ...")
    
    # Group by variant_id and aggregate
    variant_rows = []
    for vid, grp in df.groupby("variant_id"):
        row = grp.iloc[0]
        
        # Get all sizes in this variant
        sizes = grp["size"].dropna().unique().tolist()
        
        variant_rows.append({
            "variant_id": vid,
            "group_id": row["group_id"],
            "group_name": row["group_name"],
            "brand": row["brand"],
            "product_line": row["product_line"],
            "product_domain": row["product_domain"],
            "taxonomy": row["taxonomy"],
            "core_title": row["core_title"],
            "variant_key": row["variant_key"],
            "size": sizes[0] if len(sizes) == 1 else "; ".join(str(s) for s in sizes),
            "size_amount": row.get("size_amount"),
            "size_unit": row.get("size_unit"),
            "size_qty": row.get("size_qty"),
            "flavour": row.get("flavour"),
            "formulation": row.get("formulation"),
            "fat_level": row.get("fat_level"),
            "fat_percentage": row.get("fat_percentage"),
            "product_count": len(grp),
            "unique_upcs": grp["upc"].nunique(),
            "source": row["source"],
            "source_url": row["source_url"],
        })
    
    variant_df = pd.DataFrame(variant_rows)
    log(f"  Variant table: {len(variant_df)} variants")
    return variant_df


def build_variant_mapping(df: pd.DataFrame) -> pd.DataFrame:
    """
    Build the product → variant mapping table.
    """
    log("Building variant mapping ...")
    
    mapping = df[[
        "external_id", "upc", "group_id", "variant_id",
        "core_title", "original_title", "brand", "size",
        "variant_key", "source", "source_url"
    ]].copy()
    
    log(f"  Variant mapping: {len(mapping)} products")
    return mapping


def build_variant_summary(df: pd.DataFrame) -> pd.DataFrame:
    """
    Build the group-level variant summary.
    """
    log("Building variant summary ...")
    
    summary_rows = []
    for gid, grp in df.groupby("group_id"):
        variant_ids = grp["variant_id"].unique().tolist()
        row = grp.iloc[0]
        summary_rows.append({
            "group_id": gid,
            "group_name": row["group_name"],
            "brand": row["brand"],
            "product_domain": row["product_domain"],
            "taxonomy": row["taxonomy"],
            "variant_count": len(variant_ids),
            "variant_ids": ", ".join(variant_ids),
            "product_count": len(grp),
        })
    
    summary_df = pd.DataFrame(summary_rows)
    log(f"  Variant summary: {len(summary_df)} groups")
    return summary_df


# ---------------------------------------------------------------------------
# 5. VALIDATION
# ---------------------------------------------------------------------------

def validate_phase4(df: pd.DataFrame, variant_table: pd.DataFrame,
                    variant_mapping: pd.DataFrame, variant_summary: pd.DataFrame) -> dict:
    """Comprehensive Phase 4 validation."""
    checks = {}
    
    # Rule 1: Every product has exactly one variant
    checks["rule_1_product_has_one_variant"] = {
        "description": "Every product has exactly one variant_id",
        "products": len(df),
        "unique_variant_ids_per_product": int(df.groupby("external_id")["variant_id"].nunique().max()),
        "pass": int(df.groupby("external_id")["variant_id"].nunique().max()) == 1,
    }
    
    # Rule 2: Every variant belongs to exactly one group
    checks["rule_2_variant_belongs_to_one_group"] = {
        "description": "Every variant belongs to exactly one group_id",
        "variants": len(variant_table),
        "unique_groups_per_variant": int(variant_table.groupby("variant_id")["group_id"].nunique().max()),
        "pass": int(variant_table.groupby("variant_id")["group_id"].nunique().max()) == 1,
    }
    
    # Rule 3: No product maps to multiple variants
    checks["rule_3_no_product_multiple_variants"] = {
        "description": "No product maps to multiple variant_ids",
        "products_with_multiple_variants": int((df.groupby("external_id")["variant_id"].nunique() > 1).sum()),
        "pass": int((df.groupby("external_id")["variant_id"].nunique() > 1).sum()) == 0,
    }
    
    # Rule 4: Variant count matches
    checks["rule_4_variant_count"] = {
        "description": "Variant count in mapping equals variant table",
        "mapping_variants": int(variant_mapping["variant_id"].nunique()),
        "table_variants": len(variant_table),
        "pass": int(variant_mapping["variant_id"].nunique()) == len(variant_table),
    }
    
    # Rule 5: All products have group_id
    checks["rule_5_all_products_have_group"] = {
        "description": "All products have group_id",
        "null_group_ids": int(df["group_id"].isna().sum()),
        "pass": int(df["group_id"].isna().sum()) == 0,
    }
    
    # Rule 6: All products have variant_id
    checks["rule_6_all_products_have_variant"] = {
        "description": "All products have variant_id",
        "null_variant_ids": int(df["variant_id"].isna().sum()),
        "pass": int(df["variant_id"].isna().sum()) == 0,
    }
    
    # Overall
    all_pass = all(c.get("pass", True) for c in checks.values())
    checks["overall"] = {"result": "PASS" if all_pass else "FAIL"}
    
    return checks


# ---------------------------------------------------------------------------
# 6. STATISTICS
# ---------------------------------------------------------------------------

def build_statistics(df: pd.DataFrame, variant_table: pd.DataFrame,
                     variant_mapping: pd.DataFrame, variant_summary: pd.DataFrame) -> dict:
    """Build Phase 4 statistics."""
    
    # Variant count distribution
    variant_counts = variant_summary["variant_count"].value_counts().sort_index()
    
    return {
        "version": VERSION,
        "timestamp": TIMESTAMP,
        "input": {
            "source": "phase3_product_group_mapping.csv + phase2_output.parquet",
            "products": len(df),
            "groups": df["group_id"].nunique(),
        },
        "output": {
            "variants": len(variant_table),
            "products_mapped": len(variant_mapping),
            "groups": len(variant_summary),
        },
        "variant_distribution": {
            "groups_with_1_variant": int((variant_summary["variant_count"] == 1).sum()),
            "groups_with_2_variants": int((variant_summary["variant_count"] == 2).sum()),
            "groups_with_3_variants": int((variant_summary["variant_count"] == 3).sum()),
            "groups_with_4plus_variants": int((variant_summary["variant_count"] >= 4).sum()),
            "max_variants_per_group": int(variant_summary["variant_count"].max()),
            "avg_variants_per_group": float(variant_summary["variant_count"].mean()),
        },
        "variant_count_distribution": {str(k): int(v) for k, v in variant_counts.items()},
        "missing_size": {
            "products_with_missing_size": int(df["size"].isna().sum()),
            "products_with_no_size_variant": int((df["variant_key"] == "nosize").sum()),
        },
    }


# ---------------------------------------------------------------------------
# 7. PEANUT BUTTER EXAMPLE
# ---------------------------------------------------------------------------

def generate_peanut_butter_example(df: pd.DataFrame) -> pd.DataFrame:
    """Generate Peanut Butter variant example for audit."""
    log("Generating Peanut Butter example ...")
    
    # Find peanut butter products
    pb = df[df["original_title"].str.contains("peanut butter", case=False, na=False)].copy()
    
    # Select relevant columns
    example = pb[[
        "external_id", "upc", "group_id", "variant_id",
        "original_title", "core_title", "brand", "size",
        "variant_key", "variant_attributes", "product_line",
        "identity_hash", "fat_level", "product_domain", "taxonomy"
    ]].copy()
    
    # Sort by group_id and size
    example = example.sort_values(["group_id", "size"]).reset_index(drop=True)
    
    log(f"  Peanut Butter example: {len(example)} products in {example['group_id'].nunique()} groups")
    return example


# ---------------------------------------------------------------------------
# MAIN
# ---------------------------------------------------------------------------

def main():
    log("Starting Phase 4 (v1.0.0 — Variant Assignment)")
    
    # Ensure output dirs exist
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    VALIDATION_DIR.mkdir(parents=True, exist_ok=True)
    STATISTICS_DIR.mkdir(parents=True, exist_ok=True)
    AUDIT_DIR.mkdir(parents=True, exist_ok=True)
    
    # ------------------------------------------------------------------
    # 1. LOAD PHASE 3 OUTPUTS
    # ------------------------------------------------------------------
    log("Loading Phase 3 outputs ...")
    mapping_path = PHASE3_DIR / "outputs" / "product_group_mapping.csv"
    mapping = pd.read_csv(mapping_path)
    log(f"  product_group_mapping: {len(mapping)} rows")
    
    # ------------------------------------------------------------------
    # 2. LOAD PHASE 2 OUTPUT (for variant-relevant columns)
    # ------------------------------------------------------------------
    log("Loading Phase 2 output ...")
    p2_path = PHASE2_DIR / "outputs" / "phase2_output.parquet"
    p2 = pd.read_parquet(p2_path)
    log(f"  phase2_output: {len(p2)} rows, {len(p2.columns)} columns")
    
    # ------------------------------------------------------------------
    # 3. MERGE PHASE 3 MAPPING WITH PHASE 2 VARIANT DATA
    # ------------------------------------------------------------------
    log("Merging Phase 3 mapping with Phase 2 variant data ...")
    variant_cols = ["external_id", "variant_attributes", "flavour", "formulation",
                    "fat_level", "fat_percentage", "size_amount", "size_unit", "size_qty"]
    df = mapping.merge(p2[variant_cols], on="external_id", how="left")
    log(f"  Merged: {len(df)} rows, {len(df.columns)} columns")
    
    # ------------------------------------------------------------------
    # 4. ASSIGN VARIANTS
    # ------------------------------------------------------------------
    df = assign_variants(df)
    
    # ------------------------------------------------------------------
    # 5. GENERATE OUTPUTS
    # ------------------------------------------------------------------
    variant_table = build_variant_table(df)
    variant_mapping = build_variant_mapping(df)
    variant_summary = build_variant_summary(df)
    
    # ------------------------------------------------------------------
    # 6. VALIDATION
    # ------------------------------------------------------------------
    log("Running validation ...")
    validation = validate_phase4(df, variant_table, variant_mapping, variant_summary)
    
    # ------------------------------------------------------------------
    # 7. STATISTICS
    # ------------------------------------------------------------------
    log("Building statistics ...")
    statistics = build_statistics(df, variant_table, variant_mapping, variant_summary)
    
    # ------------------------------------------------------------------
    # 8. PEANUT BUTTER EXAMPLE
    # ------------------------------------------------------------------
    pb_example = generate_peanut_butter_example(df)
    
    # ------------------------------------------------------------------
    # 9. SAVE OUTPUTS
    # ------------------------------------------------------------------
    log("Saving outputs ...")
    
    # Main outputs (Parquet)
    variant_table.to_parquet(OUTPUT_DIR / "reference_product_variants.parquet", index=False)
    log(f"  Saved reference_product_variants.parquet ({len(variant_table)} rows)")
    
    variant_mapping.to_parquet(OUTPUT_DIR / "product_variant_mapping.parquet", index=False)
    log(f"  Saved product_variant_mapping.parquet ({len(variant_mapping)} rows)")
    
    variant_summary.to_parquet(OUTPUT_DIR / "reference_product_variant_summary.parquet", index=False)
    log(f"  Saved reference_product_variant_summary.parquet ({len(variant_summary)} rows)")
    
    # CSV versions for inspection
    variant_table.to_csv(OUTPUT_DIR / "reference_product_variants.csv", index=False)
    variant_mapping.to_csv(OUTPUT_DIR / "product_variant_mapping.csv", index=False)
    variant_summary.to_csv(OUTPUT_DIR / "reference_product_variant_summary.csv", index=False)
    
    # Validation
    with open(VALIDATION_DIR / "phase4_validation.json", "w") as f:
        json.dump(validation, f, indent=2, default=str)
    log("  Saved phase4_validation.json")
    
    # Statistics
    with open(STATISTICS_DIR / "phase4_statistics.json", "w") as f:
        json.dump(statistics, f, indent=2, default=str)
    log("  Saved phase4_statistics.json")
    
    # Peanut Butter example
    pb_example.to_parquet(AUDIT_DIR / "phase4_peanut_butter_example.parquet", index=False)
    log(f"  Saved phase4_peanut_butter_example.parquet ({len(pb_example)} rows)")
    
    # ------------------------------------------------------------------
    # SUMMARY
    # ------------------------------------------------------------------
    log("")
    log("=== PHASE 4 COMPLETE (v1.0.0) ===")
    log(f"Input:  {len(df)} products, {df['group_id'].nunique()} groups")
    log(f"Output: {len(variant_table)} variants, {len(variant_mapping)} products mapped")
    log(f"Validation: {validation['overall']['result']}")
    log("========================")
    
    return variant_table, variant_mapping, variant_summary, validation, statistics


if __name__ == "__main__":
    main()