"""Script 03: Generate weak supervision labels for ACSA, AND fit the MetaEncoder.

v2 upgrade: labeling now depends on metadata (features, price, category),
so that a text-only model CANNOT perfectly replicate the labels.
"""
import argparse
import sys
from pathlib import Path

import pandas as pd

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from src.utils import setup_logging
from src import config as cfg
from src.aspect_dict import (
    get_aspect_dict, extend_aspect_dict_from_metadata, save_aspect_dict,
)
from src.weak_labeling import (
    WeakLabeler, label_distribution_summary, save_audit_sample,
    compute_category_median_prices,
)
from src.meta_encoder import fit_and_save


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--no_extension", action="store_true",
                        help="Skip extending aspect dict from metadata")
    parser.add_argument("--top_k", type=int, default=15,
                        help="Top-K new tokens to add per aspect from metadata mining")
    args = parser.parse_args()

    setup_logging()

    # 1. Aspect dictionary
    aspect_dict = get_aspect_dict()
    if not args.no_extension:
        proc = pd.read_parquet(cfg.PROCESSED_PATH)
        meta_blobs = []
        for col in ("features_text", "categories_text", "product_title"):
            if col in proc.columns:
                meta_blobs.extend(proc[col].dropna().astype(str).tolist())
        meta_blobs = list({m for m in meta_blobs if m})[:50000]
        print(f"Mining {len(meta_blobs)} metadata blobs to extend aspect dict...")
        aspect_dict = extend_aspect_dict_from_metadata(meta_blobs, aspect_dict, top_k=args.top_k)
    save_aspect_dict(aspect_dict)
    print(f"Aspect dict saved to {cfg.ASPECT_DICT_PATH}")
    for k, v in aspect_dict.items():
        print(f"  {k}: {len(v)} keywords  (sample: {v[:5]}...)")

    # 2. Compute category median prices from the full processed data
    #    (needed for price-aware VALUE labeling)
    proc = pd.read_parquet(cfg.PROCESSED_PATH)
    cat_median_prices = compute_category_median_prices(proc)
    print(f"\nComputed median prices for {len(cat_median_prices)-1} categories "
          f"(global median=${cat_median_prices.get('__global__', 0):.1f})")

    # 3. Weak labeling each split — now with meta-dependent logic
    labeler = WeakLabeler(aspect_dict, category_median_prices=cat_median_prices)
    for split_name, path in [("train", cfg.TRAIN_PATH),
                             ("val", cfg.VAL_PATH),
                             ("test", cfg.TEST_PATH)]:
        df = pd.read_parquet(path)
        print(f"\nLabeling {split_name} ({len(df)} rows)...")
        labeled = labeler.label_dataframe(df)
        labeled.to_parquet(path, index=False)
        summary = label_distribution_summary(labeled)
        print(f"\n{split_name.upper()} label distribution:")
        print(summary.to_string(index=False))

    # 4. Audit sample from train
    train_labeled = pd.read_parquet(cfg.TRAIN_PATH)
    save_audit_sample(train_labeled)
    train_labeled.to_parquet(cfg.LABELED_PATH, index=False)
    print(f"\nAudit sample at {cfg.REPORT_DIR / 'weak_label_audit_sample.csv'}")

    # 5. Fit MetaEncoder on the TRAIN split only and save
    print(f"\nFitting MetaEncoder on train ({len(train_labeled)} rows)...")
    enc = fit_and_save(train_labeled)
    print(f"MetaEncoder saved to {cfg.META_ENCODER_PATH}")
    print(f"  total_dim={enc.total_dim} (tfidf={enc.tfidf_dim} + numeric={enc.num_dim})")


if __name__ == "__main__":
    main()