"""
train_context.py

Train the HV Context model on a conversation dataset.

If you have your own turns (list of (topic, user, assistant)),
edit the load_turns() function to read from your source.

Run:
    python train_context.py
    python train_context.py --output ./my_context_model
"""

import argparse
import time
from pathlib import Path

from hv_context import (
    HVContext, generate_conversation, D, SEED,
)

OUTPUT_DIR = "context_model"


def load_turns():
    """Replace this with your own data loader.

    Returns a list of (topic, user_text, assistant_text) tuples.
    The topic label is only used for evaluation -- it is not
    consumed by the fit() method.
    """
    return generate_conversation(n_per_topic=10, seed=SEED)


def main():
    parser = argparse.ArgumentParser(description="Train HV Context.")
    parser.add_argument("--output", type=str, default=OUTPUT_DIR)
    parser.add_argument("--dim", type=int, default=D)
    args = parser.parse_args()

    print(f"Loading turns...")
    turns = load_turns()
    print(f"  {len(turns)} turns, {len(set(t for t, _, _ in turns))} topics")

    print(f"Training...")
    t0 = time.perf_counter()
    model = HVContext(D=args.dim, seed=SEED)
    model.fit(turns)
    dt = time.perf_counter() - t0

    print(f"  vocabulary: {len(model.vocab_index)} words")
    print(f"  codebook:   {model.K.size / 8 / 1024:.1f} KB")
    print(f"  bank:       {model.bank_int8.nbytes / 1024:.1f} KB")
    print(f"  time:       {dt*1000:.0f} ms")

    print(f"Saving to {args.output}/")
    model.save(args.output)
    print("Done.")


if __name__ == "__main__":
    main()