"""Portable MSE fine-tuning from the RNASeek pretrained base or a regression checkpoint.

Paths are repository-relative. This is a new fine-tuning run, not an automatic
resume of historical optimizer state or hyperparameter-search experiments.
"""
import argparse
import json
from pathlib import Path

from validate_release import DEFAULTS, REGRESSION_BASE, ROOT, load_regression_checkpoint, local_path


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("task", choices=DEFAULTS)
    parser.add_argument("--checkpoint", default=REGRESSION_BASE,
                        help="Repository-relative pretrained base or fine-tuned regression checkpoint (default: %(default)s)")
    parser.add_argument("--tokenizer")
    parser.add_argument("--train-data", required=True)
    parser.add_argument("--valid-data", required=True)
    parser.add_argument("--output", required=True)
    parser.add_argument("--max-steps", type=int, default=1000)
    parser.add_argument("--batch-size", type=int, default=4)
    parser.add_argument("--learning-rate", type=float, default=1e-6)
    parser.add_argument("--limit", type=int, default=0)
    parser.add_argument("--skip-model-save", action="store_true")
    args = parser.parse_args()
    output = local_path(args.output)
    if output.exists():
        parser.error("Output already exists; use a new directory to protect prior runs")
    if args.max_steps < 1 or args.batch_size < 1 or args.limit < 0:
        parser.error("max-steps/batch-size must be positive and limit nonnegative")

    import numpy as np
    import pandas as pd
    import torch
    from torch import nn
    from datasets import Dataset
    from transformers import Trainer, TrainingArguments, DataCollatorWithPadding, set_seed
    from safetensors.torch import save_file
    set_seed(42)
    backbone, head, tokenizer = load_regression_checkpoint(args.checkpoint, args.tokenizer, initialize_head=True)

    class Regressor(nn.Module):
        def __init__(self):
            super().__init__()
            self.backbone, self.regression_head = backbone, head
        def forward(self, input_ids, attention_mask, labels=None):
            hidden = self.backbone(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state
            positions = torch.arange(hidden.shape[1], device=hidden.device)[None, :].expand(hidden.shape[:2])
            last = positions.masked_fill(~attention_mask.bool(), -1).max(dim=1).values
            logits = self.regression_head.net(hidden[torch.arange(hidden.shape[0], device=hidden.device), last]).flatten()
            result = {"logits": logits}
            if labels is not None:
                result["loss"] = nn.functional.mse_loss(logits.float(), labels.float())
            return result

    def dataset(path, split):
        path = local_path(path)
        if args.task == "efficiency":
            mapping = json.loads(path.read_text())
            texts, labels = list(mapping), list(mapping.values())
        else:
            frame = pd.read_csv(path, sep="\t")
            if "split" in frame:
                allowed = ["train"] if split == "train" else ["valid", "val", "validation"]
                frame = frame[frame["split"].astype(str).str.lower().isin(allowed)]
            if frame.empty:
                raise ValueError("The supplied table contains no records for the requested split")
            texts = frame["text"].tolist() if "text" in frame else ("<s>" + frame.seq + "~$predict_stability\n</s>").tolist()
            labels = frame["label" if "label" in frame else "score"].tolist()
        if args.limit:
            texts, labels = texts[:args.limit], labels[:args.limit]
        encoded = tokenizer(texts, add_special_tokens=args.task == "efficiency", truncation=True, max_length=512)
        return Dataset.from_dict({"input_ids": encoded["input_ids"], "attention_mask": encoded["attention_mask"], "labels": labels})
    train, valid = dataset(args.train_data, "train"), dataset(args.valid_data, "valid")
    output.mkdir(parents=True, exist_ok=False)
    model = Regressor()
    trainer = Trainer(model=model, processing_class=tokenizer, train_dataset=train, eval_dataset=valid,
                      data_collator=DataCollatorWithPadding(tokenizer, pad_to_multiple_of=8),
                      args=TrainingArguments(output_dir=str(output), overwrite_output_dir=False,
                          logging_dir=str(output / "tensorboard"), report_to=["tensorboard"], logging_steps=1,
                          max_steps=args.max_steps, per_device_train_batch_size=args.batch_size,
                          per_device_eval_batch_size=args.batch_size, learning_rate=args.learning_rate,
                          bf16=torch.cuda.is_available(), optim="adamw_bnb_8bit" if torch.cuda.is_available() else "adamw_torch",
                          save_strategy="no", eval_strategy="no", seed=42))
    result = trainer.train()
    metrics = {**result.metrics, **trainer.evaluate(), "task": args.task, "smoke_test": bool(args.limit)}
    if not np.isfinite([v for v in metrics.values() if isinstance(v, float)]).all():
        raise ValueError("Non-finite training/evaluation metric")
    (output / "metrics.json").write_text(json.dumps(metrics, indent=2) + "\n")
    if not args.skip_model_save:
        final = output / "model"
        final.mkdir()
        save_file({k: v.detach().cpu().contiguous() for k, v in model.state_dict().items()}, str(final / "model.safetensors"))
        backbone.config.save_pretrained(final)
        tokenizer.save_pretrained(final)
    print(json.dumps(metrics, indent=2))


if __name__ == "__main__":
    main()