"""Publish a trained DINNNER checkpoint and its model code to the Hugging Face Hub.

The checkpoint is checked against the model code (weights/dinnner_core.py by
default) first, then both are uploaded with a short model card. The repo is
private unless --public is given: keep it private until the paper is
published. Run it after `hf auth login` (or with an HF_TOKEN environment
variable that has write access):

    python scripts/upload_weights.py \
        --checkpoint path/to/checkpoint.pth \
        --repo faridan/dinnner-pubmedbert-bionlp13pc

Leave out --checkpoint to update only the model code and the model card.
"""

import argparse
import os
import sys
import tempfile

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import torch  # noqa: E402
from huggingface_hub import HfApi  # noqa: E402

from dinnner.config import MODELS  # noqa: E402
from dinnner.model import load_model  # noqa: E402

CARD = """---
license: {license}
language: en
library_name: pytorch
pipeline_tag: token-classification
tags: [biomedical, named-entity-recognition, ner, bionlp13pc, dinnner]
datasets: [bionlp13pc]
base_model: {encoder}
---

# DINNNER + PubMedBERT (BioNLP13PC)

Weights for **DINNNER** (Distribution-Informed Neural Network for Named Entity
Recognition), trained on the BioNLP 2013 Pathway Curation corpus.

The paper describing DINNNER is under review. Details of the method and
results will be added here once it is published.

* Encoder: `{encoder}`
* Labels: {labels}
* Max subword length: {max_length}

Try it online: {demo}

The checkpoint is a plain PyTorch `state_dict` for the `DINNNER` class in
`dinnner_core.py`.
"""


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--checkpoint", help="local .pth file (omit to update only the code and card)")
    ap.add_argument("--code", help="model code to upload (default: the model's local_code_path)")
    ap.add_argument("--repo", required=True, help="e.g. your-username/dinnner-pubmedbert-bionlp13pc")
    ap.add_argument("--model", default="pubmedbert-bionlp13pc", choices=list(MODELS))
    ap.add_argument("--demo", default="https://huggingface.co/spaces/faridan/dinnner")
    ap.add_argument("--license", default="other", help="licence id for the card, e.g. mit or cc-by-4.0")
    ap.add_argument("--public", action="store_true", help="make the repo public (after publication)")
    args = ap.parse_args()

    spec = MODELS[args.model]
    code = args.code or spec.local_code_path
    print("Checking the checkpoint against the model code...")
    local = {"local_code_path": code, **({"local_path": args.checkpoint} if args.checkpoint else {})}
    load_model(spec.__class__(**{**spec.__dict__, **local}), torch.device("cpu"))
    print("OK")

    api = HfApi()
    api.create_repo(args.repo, repo_type="model", private=not args.public, exist_ok=True)
    api.update_repo_settings(args.repo, private=not args.public)
    if args.checkpoint:
        print(f"Uploading {args.checkpoint} as {spec.weights_file} ...")
        api.upload_file(path_or_fileobj=args.checkpoint, path_in_repo=spec.weights_file, repo_id=args.repo,
                        commit_message="Add DINNNER checkpoint")
    api.upload_file(path_or_fileobj=code, path_in_repo=spec.code_file, repo_id=args.repo,
                    commit_message="Add DINNNER model code")

    card = CARD.format(license=args.license, encoder=spec.encoder, max_length=spec.max_length,
                       demo=args.demo, labels=", ".join(spec.id2label.values()))
    with tempfile.NamedTemporaryFile("w", suffix=".md", delete=False) as f:
        f.write(card)
    api.upload_file(path_or_fileobj=f.name, path_in_repo="README.md", repo_id=args.repo,
                    commit_message="Add model card")
    print(f"Done: https://huggingface.co/{args.repo}")


if __name__ == "__main__":
    main()