#!/usr/bin/env python3
"""Train BERT multi-task model with explicit 90/10 train/test split on real case data."""

from __future__ import annotations

import json
import os
import random
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

# NOTE: transformers must import before torch on Windows
from transformers import AutoTokenizer, get_linear_schedule_with_warmup
import torch
from torch.utils.data import DataLoader
from tqdm import tqdm

from app.services.bert_multi_task_model import ChineseRobertaMultiTask
from train_bert import LaborCaseDataset, collate_batch, train_epoch, evaluate

DATASET = Path(__file__).resolve().parent.parent.parent / "backend" / "data" / "training_dataset.json"
OUTPUT = Path(__file__).resolve().parent.parent / "models" / "chinese_roberta_labor_extractor_v2"
MODEL_NAME = "hfl/chinese-roberta-wwm-ext"
EPOCHS = 10
BATCH_SIZE = 4
LR = 2e-5
MAX_LENGTH = 512
PATIENCE = 3
TEST_RATIO = 0.1  # 10% for test

random.seed(42)
torch.manual_seed(42)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Device: {device}")

# Load data
with open(DATASET, encoding="utf-8") as f:
    data = json.load(f)

cases = data["dataset"]
random.shuffle(cases)

n_test = max(1, int(len(cases) * TEST_RATIO))
n_train = len(cases) - n_test

train_cases = cases[:n_train]
test_cases = cases[n_train:]

print(f"Total: {len(cases)}, Train: {n_train}, Test: {n_test}")

# Tokenizer
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)

# Datasets
train_ds = LaborCaseDataset(train_cases, tokenizer, MAX_LENGTH, for_training=True)
test_ds = LaborCaseDataset(test_cases, tokenizer, MAX_LENGTH, for_training=True)

train_loader = DataLoader(train_ds, batch_size=BATCH_SIZE, shuffle=True,
                          collate_fn=collate_batch, num_workers=0)
test_loader = DataLoader(test_ds, batch_size=BATCH_SIZE, shuffle=False,
                         collate_fn=collate_batch, num_workers=0)

# Model
model = ChineseRobertaMultiTask(model_name=MODEL_NAME)
model.to(device)
print(f"Model params: {sum(p.numel() for p in model.parameters()):,}")

# Optimizer
no_decay = ["bias", "LayerNorm.weight"]
opt_params = [
    {"params": [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)],
     "weight_decay": 0.01},
    {"params": [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)],
     "weight_decay": 0.0},
]
optimizer = torch.optim.AdamW(opt_params, lr=LR)
total_steps = len(train_loader) * EPOCHS
scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=int(total_steps * 0.1),
                                            num_training_steps=total_steps)

# Training
best_loss = float("inf")
patience_cnt = 0

for epoch in range(1, EPOCHS + 1):
    train_loss = train_epoch(model, train_loader, optimizer, scheduler, device, epoch)
    dev_metrics = evaluate(model, test_loader, device)

    print(f"Epoch {epoch:>2}: train_loss={train_loss:.4f}, "
          f"test_loss={dev_metrics['loss']:.4f}, "
          f"test_cause_acc={dev_metrics['cause_accuracy']:.3f}")

    if dev_metrics["loss"] < best_loss:
        best_loss = dev_metrics["loss"]
        patience_cnt = 0
        os.makedirs(OUTPUT, exist_ok=True)
        model.save_pretrained(str(OUTPUT))
        tokenizer.save_pretrained(str(OUTPUT))
        print(f"  -> Saved to {OUTPUT}")
    else:
        patience_cnt += 1
        if patience_cnt >= PATIENCE:
            print(f"Early stop at epoch {epoch}")
            break

print(f"\nDone. Best test loss: {best_loss:.4f}")
print(f"Model: {OUTPUT}")