import argparse
import os
import sys
import torch

sys.path.insert(0, os.path.dirname(__file__))
from modeling import load_model, encode_text, pad_batch, TAXONOMY, JIGSAW_LABELS

SAMPLE = "You are an idiot and I will find you. My email is jane.doe@example.com."
VERIFICATION = {
    "tax": [
        0.3051,
        0.0119,
        0.1066,
        0.0103,
        0.4735,
        0.4706,
        0.4947,
        0.0363,
        0.4683,
        0.0714,
        0.4361,
        0.0968,
        0.3022,
        0.0607,
        0.0827,
        0.038,
        0.0001,
        0.5119,
        0.5415,
    ],
    "jig": [0.6895, 0.025, 0.2517, 0.0263, 0.4001, 0.244],
    "tol": 0.02,
}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--weights", default=os.environ.get("FELA_MODERATION_WEIGHTS", "."))
    args = ap.parse_args()
    model = load_model(args.weights)
    ids, _ = encode_text(SAMPLE, model.cfg.max_len)
    input_ids, mask = pad_batch([ids], model.cfg.max_len)
    with torch.no_grad():
        out = model(input_ids, mask, task="both")
    tax = torch.sigmoid(out["taxonomy"][0]).tolist()
    jig = torch.sigmoid(out["jigsaw"][0]).tolist()
    pii_shape = tuple(out["pii"].shape)
    exp_tax = (1, len(TAXONOMY))
    exp_jig = (1, len(JIGSAW_LABELS))
    if tuple(out["taxonomy"].shape) != exp_tax or tuple(out["jigsaw"].shape) != exp_jig:
        print(
            f"FAILURE: head shapes {tuple(out['taxonomy'].shape)}, {tuple(out['jigsaw'].shape)} != expected {exp_tax}, {exp_jig}"
        )
        sys.exit(1)
    if pii_shape[0] != 1 or pii_shape[2] != model.cfg.n_pii_tags:
        print(
            f"FAILURE: pii shape {pii_shape} unexpected (n_pii_tags {model.cfg.n_pii_tags})"
        )
        sys.exit(1)
    print(
        f"Shapes verified OK: taxonomy {tuple(out['taxonomy'].shape)}, jigsaw {tuple(out['jigsaw'].shape)}, pii {pii_shape}"
    )
    print("Captured taxonomy probs: " + str([round(p, 4) for p in tax]))
    print("Captured jigsaw probs: " + str([round(p, 4) for p in jig]))
    tol = VERIFICATION["tol"]
    diffs = [abs(a - b) for a, b in zip(tax, VERIFICATION["tax"])]
    diffs += [abs(a - b) for a, b in zip(jig, VERIFICATION["jig"])]
    if max(diffs) > tol:
        print(
            f"FAILURE: probs vs verification max abs diff {max(diffs):.4f} > tol {tol}"
        )
        sys.exit(1)
    print(f"Verification check OK (max abs diff {max(diffs):.4f} <= tol {tol})")


if __name__ == "__main__":
    main()