"""Build a self-contained Colab notebook and a checksummed data archive."""
import json
from pathlib import Path
import zipfile

from tiny_json_llm.data import sha256, write_json

ROOT = Path(__file__).resolve().parent
PREPARED = ROOT / "prepared/home_commands_v1"
manifest = json.loads((PREPARED / "manifest.json").read_text(encoding="utf-8"))
bundle = ROOT / "home_commands_training_bundle.zip"
if not bundle.exists():
    with zipfile.ZipFile(bundle, "w", compression=zipfile.ZIP_DEFLATED, compresslevel=6) as archive:
        for path in sorted(PREPARED.rglob("*")):
            if path.is_file():
                archive.write(path, str(path.relative_to(ROOT)).replace("\\", "/"))
bundle_hash = sha256(bundle)
sources = {str(path.relative_to(ROOT)).replace("\\", "/"): path.read_text(encoding="utf-8")
           for path in sorted((ROOT / "tiny_json_llm").glob("*.py"))}
sources["requirements.txt"] = (ROOT / "requirements.txt").read_text(encoding="utf-8")
cells = []


def markdown(text):
    cells.append({"cell_type": "markdown", "metadata": {}, "source": text.splitlines(keepends=True)})


def code(text):
    cells.append({"cell_type": "code", "metadata": {}, "source": text.splitlines(keepends=True),
                  "execution_count": None, "outputs": []})


a = manifest["architecture"]
markdown(f"""# Tiny home-command JSON model — A100 training and recovery

Account: **your Google account**. Select **Runtime → Change runtime type → A100 GPU**.

Randomly initialized Llama-style decoder: **{a['parameters']:,} parameters**, **{a['num_hidden_layers']} layers**, width **{a['hidden_size']}**, vocabulary **{manifest['tokenizer']['vocabulary_size']}**.

One pass over **{manifest['splits']['train']['prediction_tokens']:,} training prediction tokens** gives **{a['unique_train_tokens_per_parameter']:.3f} tokens/parameter**. No pretrained weights, external training corpus, or repeated epochs. Tokenizer fitting uses training rows only.

The source has {manifest['source_rows']:,} rows; {manifest['duplicates_removed']:,} duplicate input rows were removed. Core phrase families keep casing, punctuation and common politeness variants together in the 80/10/10 split. This is highly templated English and Romanized Hindi data. The Chinchilla ratio is a sizing heuristic extrapolated to this task and model size, not a proof of compute optimality.

**Recovery:** rerun all cells after reconnecting. The bundle and run live in Google Drive. Training checks data/settings hashes and restores model, optimizer, scheduler, random states and exact next-example position. Two checksummed checkpoint generations are retained, every 100 updates or 120 seconds. An interrupted partial file is ignored.

Colab's lifetime, compute-unit and GPU limits still apply. Progress messages do not override idle timeouts. The supplied colab_keep_alive.sh uses Google's official CLI keep-alive on Linux/macOS. No automatic purchase or public publication.
""")

code("""from pathlib import Path
import os, sys, json, shutil, subprocess, zipfile
WORK = Path('/content/tiny_json_workspace')
WORK.mkdir(parents=True, exist_ok=True)
os.chdir(WORK)
SOURCE_FILES = """ + repr(sources) + """
for relative, content in SOURCE_FILES.items():
    destination = WORK / relative
    destination.parent.mkdir(parents=True, exist_ok=True)
    destination.write_text(content, encoding='utf-8')
subprocess.check_call([sys.executable, '-m', 'pip', 'install', '-q', '-r', 'requirements.txt'])
print('Training code and pinned tokenizer/transformers dependencies are ready.')
""")

markdown("## Persistent storage\nAuthorize the standard Colab Drive mount if prompted. Keep this project folder when resuming.")
code("""from google.colab import drive
drive.mount('/content/drive')
DRIVE_ROOT = Path('/content/drive/MyDrive/TinyJSONHomeCommands/v1_10822f6617d2')
DRIVE_ROOT.mkdir(parents=True, exist_ok=True)
RUN = DRIVE_ROOT / 'run'
print('Persistent run:', RUN)
""")

markdown("## Dataset snapshot\nFirst run: choose home_commands_training_bundle.zip. Future runs use the saved Drive copy.")
code("EXPECTED_BUNDLE_SHA256 = " + repr(bundle_hash) + """
BUNDLE_NAME = 'home_commands_training_bundle.zip'
PERSISTENT_BUNDLE = DRIVE_ROOT / BUNDLE_NAME
from tiny_json_llm.data import sha256
if not PERSISTENT_BUNDLE.exists():
    from google.colab import files
    uploaded = files.upload()
    if BUNDLE_NAME not in uploaded:
        raise ValueError('Choose ' + BUNDLE_NAME)
    local_bundle = WORK / BUNDLE_NAME
    if sha256(local_bundle) != EXPECTED_BUNDLE_SHA256:
        raise ValueError('Uploaded bundle checksum mismatch')
    temporary = PERSISTENT_BUNDLE.with_suffix('.zip.tmp')
    shutil.copy2(local_bundle, temporary)
    temporary.replace(PERSISTENT_BUNDLE)
    del uploaded
if sha256(PERSISTENT_BUNDLE) != EXPECTED_BUNDLE_SHA256:
    raise ValueError('Drive bundle checksum mismatch')
local_bundle = WORK / BUNDLE_NAME
if not local_bundle.exists() or sha256(local_bundle) != EXPECTED_BUNDLE_SHA256:
    shutil.copy2(PERSISTENT_BUNDLE, local_bundle)
with zipfile.ZipFile(local_bundle) as archive:
    for member in archive.infolist():
        destination = (WORK / member.filename).resolve()
        if not destination.is_relative_to(WORK.resolve()):
            raise ValueError('Unsafe archive path')
    archive.extractall(WORK)
PREPARED = WORK / 'prepared/home_commands_v1'
print(json.dumps(json.loads((PREPARED / 'manifest.json').read_text())['architecture'], indent=2))
""")

markdown("## Train or resume\nRequires an actual A100; BF16 compute with FP32 optimizer weights. Resume automatically preserves batch size, schedule and training position.")
code("""import torch
assert torch.cuda.is_available(), 'Select A100 GPU in Runtime settings.'
assert 'A100' in torch.cuda.get_device_name(0), torch.cuda.get_device_name(0)
print('GPU:', torch.cuda.get_device_name(0))
subprocess.run([sys.executable, '-u', '-m', 'tiny_json_llm.train', '--prepared', str(PREPARED),
                '--out', str(RUN), '--resume'], check=True)
""")

markdown("## Held-out evaluation\nAll test examples use the training-observed JSON command vocabulary. A separate stratified diagnostic checks unconstrained generation. JSON validity is separate from semantic accuracy.")
code("""EVAL = RUN / 'evaluation'
subprocess.run([sys.executable, '-u', '-m', 'tiny_json_llm.evaluate', '--model', str(RUN / 'best_model'),
                '--prepared', str(PREPARED), '--out', str(EVAL), '--device', 'cuda'], check=True)
metrics = json.loads((EVAL / 'metrics.json').read_text())
print(json.dumps({k: v for k, v in metrics['constrained_generation'].items() if k not in ('per_command','errors')}, indent=2))
""")

markdown("## Use the model\nProduces JSON only; it does not control devices. Unknown devices, negation and out-of-distribution requests have no validated reject behavior.")
code("""from tiny_json_llm.infer import JsonCommandModel
predictor = JsonCommandModel(RUN / 'best_model')
for text in ['Please turn on the fan', 'pankha band karo', 'check the water pump status']:
    print(text, '=>', json.dumps(predictor.predict(text), ensure_ascii=False))
del predictor
torch.cuda.empty_cache()
""")

markdown("## Export results\nCheckpoints and the dataset bundle stay in Drive. Download includes best-validation and final one-pass weights, inference code, evaluation and dataset manifest.")
code("""run_summary = json.loads((RUN / 'run.json').read_text())
model_card = '''# Tiny Home Commands JSON

Trained from scratch on home-commands-3.4M.jsonl.
Purpose: translate English and Romanized Hindi/Hinglish home commands into activity/subject/action JSON.
Architecture, tokenizer, split counts, source SHA-256 and data hashes: see manifest.json.
Training: one unique pass; A100 BF16; AdamW, LR 0.001, cosine schedule, 5% warmup, effective batch 512, gradient clip 1.
Usage: from tiny_json_llm.infer import JsonCommandModel
JsonCommandModel('best_model').predict('turn on the fan')
Evaluation: evaluation/metrics.json and predictions. Constrained decoding permits 30 training-observed commands; JSON validity does not imply semantic correctness.
No general chat, unseen devices/actions, native-script Hindi or reliable negation/rejection behavior is claimed. No device actions are executed.
The highly templated corpus and tiny model fall outside the Chinchilla experiments; the 20-token/parameter heuristic is not a compute-optimality proof.
Original user dataset contributions: Copyright (c) 2026 sraivante, Apache License 2.0. Third-party attribution and terms remain unchanged. Verify source provenance before redistribution.
No pretrained base weights or external training corpus. Private artifacts; no public publication performed.
'''
(RUN / 'README.md').write_text(model_card, encoding='utf-8')
delivery = DRIVE_ROOT / 'tiny_home_json_model_v1.zip'
with zipfile.ZipFile(delivery, 'w', zipfile.ZIP_DEFLATED, compresslevel=6) as archive:
    for name in ['best_model', 'final_model', 'evaluation']:
        for path in (RUN / name).rglob('*'):
            if path.is_file():
                archive.write(path, str(path.relative_to(RUN)))
    for name in ['run.json', 'train_log.jsonl', 'README.md']:
        archive.write(RUN / name, name)
    for relative in SOURCE_FILES:
        archive.write(WORK / relative, relative)
    archive.write(PREPARED / 'manifest.json', 'dataset/manifest.json')
print('Model package:', delivery, 'SHA256:', sha256(delivery))
print('Dataset package:', PERSISTENT_BUNDLE, 'SHA256:', EXPECTED_BUNDLE_SHA256)
print('Resume checkpoints:', RUN / 'checkpoints')
from google.colab import files
files.download(str(delivery))
""")

notebook = {"nbformat": 4, "nbformat_minor": 5,
            "metadata": {"colab": {"name": "Tiny_Home_JSON_A100.ipynb", "provenance": []},
                         "kernelspec": {"name": "python3", "display_name": "Python 3"},
                         "language_info": {"name": "python"}, "accelerator": "GPU", "gpuClass": "premium"},
            "cells": cells}
for i, cell in enumerate(cells):
    cell["id"] = f"tinyjson-{i:02d}"
write_json(ROOT / "Tiny_Home_JSON_A100.ipynb", notebook)
write_json(ROOT / "transfer_manifest.json", {"bundle": bundle.name, "sha256": bundle_hash,
            "bytes": bundle.stat().st_size, "notebook": "Tiny_Home_JSON_A100.ipynb"})
print(json.dumps({"bundle_bytes": bundle.stat().st_size, "sha256": bundle_hash, "cells": len(cells)}))