import argparse
import json
from pathlib import Path

from transformers import (
    AutoModelForSequenceClassification,
    DataCollatorWithPadding,
)

from curve_logger import CurveLoggerCallback
from trainer_utils import (
    LABELS,
    compute_binary_metrics,
    load_or_prepare_imdb,
    make_trainer,
    make_training_arguments,
)


def parse_args():
    parser = argparse.ArgumentParser(description="Trainer basic usage")
    parser.add_argument("--checkpoint", default="bert-base-uncased")
    parser.add_argument("--output-dir", default="./results/page_01_basic")
    parser.add_argument("--data-dir", default="./data/page_01_imdb")
    parser.add_argument("--max-length", type=int, default=512)
    parser.add_argument("--epochs", type=float, default=3)
    parser.add_argument("--max-train-samples", type=int, default=None)
    parser.add_argument("--max-eval-samples", type=int, default=None)
    return parser.parse_args()


def main():
    args = parse_args()
    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    train_ds, eval_ds, tokenizer = load_or_prepare_imdb(
        checkpoint=args.checkpoint,
        data_dir=args.data_dir,
        max_length=args.max_length,
        padding="max_length",
        max_train_samples=args.max_train_samples,
        max_eval_samples=args.max_eval_samples,
    )
    model = AutoModelForSequenceClassification.from_pretrained(
        args.checkpoint,
        num_labels=2,
        **LABELS,
    )

    training_args = make_training_arguments(
        output_dir=str(output_dir),
        eval_strategy="epoch",
        learning_rate=2e-5,
        per_device_train_batch_size=16,
        per_device_eval_batch_size=16,
        num_train_epochs=args.epochs,
        weight_decay=0.01,
        save_strategy="epoch",
        load_best_model_at_end=True,
        metric_for_best_model="eval_f1",
        greater_is_better=True,
        report_to="none",
    )

    trainer = make_trainer(
        tokenizer=tokenizer,
        model=model,
        args=training_args,
        train_dataset=train_ds,
        eval_dataset=eval_ds,
        data_collator=DataCollatorWithPadding(tokenizer=tokenizer),
        compute_metrics=compute_binary_metrics,
        callbacks=[CurveLoggerCallback(output_dir, stage="basic")],
    )

    trainer.train()
    metrics = trainer.evaluate()
    predictions = trainer.predict(eval_ds.select(range(min(16, len(eval_ds)))))
    trainer.save_model(str(output_dir / "final_model"))
    tokenizer.save_pretrained(str(output_dir / "final_model"))

    summary = {
        "metrics": metrics,
        "prediction_shape": list(predictions.predictions.shape),
    }
    (output_dir / "summary.json").write_text(
        json.dumps(summary, ensure_ascii=False, indent=2),
        encoding="utf-8",
    )


if __name__ == "__main__":
    main()