"""Small LoRA routing exercise. QLoRA path needs a compatible NVIDIA CUDA setup."""
import argparse
import json
import random
from pathlib import Path

TRAIN = [
    ('Where is my parcel?', 'shipping'), ('My order has not arrived.', 'shipping'),
    ('Can you track the delivery?', 'shipping'), ('Which carrier has my package?', 'shipping'),
    ('I want my money back.', 'refund'), ('Please refund this purchase.', 'refund'),
    ('The item was returned; when is the refund?', 'refund'), ('Cancel and reimburse my order.', 'refund'),
    ('I cannot log in.', 'account'), ('Please reset my password.', 'account'),
    ('My account is locked.', 'account'), ('How do I change my email address?', 'account'),
]
HELD_OUT = [
    ('Has my package been dispatched?', 'shipping'), ('Can I get a reimbursement?', 'refund'),
    ('I forgot my sign-in password.', 'account'),
]


def prompt(text):
    return f'Request: {text}\nIntent:'


def main():
    import torch
    from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
    from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
    parser = argparse.ArgumentParser()
    parser.add_argument('--model', default='distilbert/distilgpt2')
    parser.add_argument('--steps', type=int, default=80)
    parser.add_argument('--device', choices=['cpu', 'mps', 'cuda'], default='cpu')
    parser.add_argument('--qlora', action='store_true')
    parser.add_argument('--output', type=Path, default=Path('routing-adapter'))
    args = parser.parse_args()
    if args.steps < 1:
        parser.error('--steps must be positive')
    if args.qlora and (args.device != 'cuda' or not torch.cuda.is_available()):
        parser.error('This QLoRA example requires --device cuda and a compatible CUDA installation.')
    random.seed(7)
    torch.manual_seed(7)
    torch.set_num_threads(4)
    tokenizer = AutoTokenizer.from_pretrained(args.model)
    tokenizer.pad_token = tokenizer.eos_token
    kwargs = {'dtype': torch.float32}
    if args.qlora:
        compute_dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
        kwargs = {'quantization_config': BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type='nf4', bnb_4bit_use_double_quant=True, bnb_4bit_compute_dtype=compute_dtype), 'device_map': {'': 0}}
    model = AutoModelForCausalLM.from_pretrained(args.model, **kwargs)
    if args.qlora:
        model = prepare_model_for_kbit_training(model, use_gradient_checkpointing=False)
    else:
        model = model.to(args.device)
    model.config.use_cache = False
    config = LoraConfig(r=8, lora_alpha=16, lora_dropout=0.05, target_modules='all-linear', bias='none', task_type='CAUSAL_LM')
    model = get_peft_model(model, config)
    model.print_trainable_parameters()
    frozen_before = {name: parameter.detach().cpu().clone() for name, parameter in model.named_parameters() if not parameter.requires_grad} if not args.qlora else {}

    def classify(text):
        model.eval()
        inputs = tokenizer(prompt(text), return_tensors='pt').to(args.device)
        with torch.no_grad():
            output = model.generate(**inputs, max_new_tokens=5, do_sample=False, pad_token_id=tokenizer.eos_token_id)
        generated = tokenizer.decode(output[0, inputs['input_ids'].shape[1]:], skip_special_tokens=True).strip()
        return generated

    before = [{'request': text, 'expected': label, 'prediction': classify(text)} for text, label in HELD_OUT]
    optimizer = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], lr=5e-4)
    losses = []
    model.train()
    for step in range(args.steps):
        text, label = TRAIN[step % len(TRAIN)]
        prefix = tokenizer.encode(prompt(text), add_special_tokens=False)
        suffix = tokenizer.encode(' ' + label, add_special_tokens=False) + [tokenizer.eos_token_id]
        ids = torch.tensor([prefix + suffix], device=args.device)
        labels = ids.clone()
        labels[:, :len(prefix)] = -100
        optimizer.zero_grad()
        loss = model(input_ids=ids, attention_mask=torch.ones_like(ids), labels=labels).loss
        loss.backward()
        optimizer.step()
        losses.append(float(loss.detach().cpu()))
    after = [{'request': text, 'expected': label, 'prediction': classify(text)} for text, label in HELD_OUT]
    frozen_unchanged = all(torch.equal(parameter.detach().cpu(), frozen_before[name]) for name, parameter in model.named_parameters() if name in frozen_before)
    args.output.mkdir(parents=True, exist_ok=True)
    model.save_pretrained(args.output)
    tokenizer.save_pretrained(args.output)
    report = {'base_model': args.model, 'steps': args.steps, 'qlora': args.qlora, 'frozen_base_unchanged': frozen_unchanged if frozen_before else 'not checked for quantized tensors', 'first_loss': losses[0], 'last_loss': losses[-1], 'before': before, 'after': after}
    (args.output / 'training-report.json').write_text(json.dumps(report, indent=2))
    print(json.dumps(report, indent=2))


if __name__ == '__main__':
    main()