"""Load, evaluate, or smoke-train a locally shipped encoder benchmark checkpoint.

Use the benchmark Conda environment. Paths are relative to the repository root.
Data must contain seq_model and label columns (the archived prediction tables do).
"""
import argparse
import importlib.util
import json
from pathlib import Path

ROOT = Path(__file__).resolve().parent

def local_path(value):
    return (ROOT / Path(value).expanduser()).resolve()

def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--checkpoint", required=True)
    parser.add_argument("--data", required=True)
    parser.add_argument("--output", required=True)
    parser.add_argument("--device", default="cuda:0")
    parser.add_argument("--limit", type=int, default=0)
    parser.add_argument("--batch-size", type=int, default=8)
    parser.add_argument("--smoke-train", action="store_true", help="One gradient update on a fresh in-memory copy; no weights are saved")
    args = parser.parse_args()
    checkpoint, data, output = map(local_path, [args.checkpoint, args.data, args.output])
    if output.exists():
        parser.error("Output exists; choose a new directory")
    if args.limit < 0 or args.batch_size < 1:
        parser.error("limit must be nonnegative and batch-size positive")
    source = ROOT / "regression_stability_functionalviral/benchmark_shipped_encoders/benchmark_shipped_encoders.py"
    spec = importlib.util.spec_from_file_location("rnaseek_benchmark", source)
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    import numpy as np
    import pandas as pd
    import torch
    from transformers import AutoConfig, AutoTokenizer
    from scipy.stats import pearsonr, spearmanr
    from sklearn.metrics import r2_score
    config = AutoConfig.from_pretrained(checkpoint, trust_remote_code=True, local_files_only=True)
    tokenizer = AutoTokenizer.from_pretrained(checkpoint / "tokenizer" if (checkpoint / "tokenizer").is_dir() else checkpoint,
                                               trust_remote_code=True, local_files_only=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token or tokenizer.unk_token
    model = module.EncoderRegressor(str(checkpoint), config, 0.1).to(args.device).eval()
    frame = pd.read_csv(data, sep="\t")
    if args.limit:
        frame = frame.head(args.limit)
    maximum = module.infer_max_length(config, None)
    predictions = []
    first_batch = None
    for start in range(0, len(frame), args.batch_size):
        rows = frame.iloc[start:start+args.batch_size]
        batch = tokenizer(rows.seq_model.tolist(), return_tensors="pt", padding=True, truncation=True, max_length=maximum)
        batch = {k: batch[k].to(args.device) for k in ("input_ids", "attention_mask")}
        if first_batch is None:
            first_batch = batch | {"labels": torch.tensor(rows.label.tolist(), device=args.device, dtype=torch.float32)}
        # no_grad keeps lazily cached rotary tensors usable by a subsequent
        # training update on this same model instance.
        with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16, enabled=str(args.device).startswith("cuda")):
            predictions.extend(model(**batch)["logits"].float().cpu().tolist())
    if not np.isfinite(predictions).all():
        raise ValueError("Non-finite benchmark prediction")
    metrics = {"n": len(frame), "full_validation": not bool(args.limit),
               "r2": float(r2_score(frame.label, predictions)),
               "pearson_r": float(pearsonr(frame.label, predictions).statistic),
               "spearman_r": float(spearmanr(frame.label, predictions).statistic)}
    if args.smoke_train:
        model.train()
        optimizer = torch.optim.AdamW(model.parameters(), lr=1e-6)
        with torch.autocast("cuda", dtype=torch.bfloat16, enabled=str(args.device).startswith("cuda")):
            loss = model(**first_batch)["loss"]
        loss.backward()
        norm = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        if not torch.isfinite(loss) or not torch.isfinite(norm):
            raise ValueError("Non-finite benchmark loss or gradient")
        optimizer.step()
        metrics.update(smoke_train_loss=float(loss.detach()), gradient_norm=float(norm))
    output.mkdir(parents=True, exist_ok=False)
    (output / "metrics.json").write_text(json.dumps(metrics, indent=2) + "\n")
    pd.DataFrame({"label": frame.label, "prediction": predictions}).to_csv(output / "predictions.tsv", sep="\t", index=False)
    print(json.dumps(metrics, indent=2), flush=True)

if __name__ == "__main__":
    main()