"""Actual causal-LM training, random initialization by default.
--resume MODEL starts from model weights with a NEW optimizer, not an exact resume.
Exports FP16 safetensors; training itself uses FP32. No pretrained weights used.
"""
import argparse
import json
import math
import random
import time
from pathlib import Path
import numpy as np
import torch
from transformers import GPT2Config, GPT2LMHeadModel, GPT2TokenizerFast, set_seed

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

def main():
    p = argparse.ArgumentParser()
    p.add_argument('--steps', type=int, default=300)
    p.add_argument('--batch-size', type=int, default=2)
    p.add_argument('--seq-len', type=int, default=128)
    p.add_argument('--lr', type=float, default=0.0005)
    p.add_argument('--seed', type=int, default=42)
    p.add_argument('--threads', type=int, default=2)
    p.add_argument('--device', default='cuda' if torch.cuda.is_available() else 'cpu')
    p.add_argument('--resume', default=None)
    p.add_argument('--output', default=str(ROOT/'model'))
    args = p.parse_args()
    assert 1 < args.seq_len <= 512 and args.batch_size >= 2 and args.batch_size % 2 == 0
    assert args.steps > 0
    set_seed(args.seed)
    torch.set_num_threads(args.threads)
    output = Path(args.output)
    output.mkdir(parents=True, exist_ok=True)
    tok = GPT2TokenizerFast.from_pretrained(args.resume or ROOT/'model')
    if args.resume:
        model = GPT2LMHeadModel.from_pretrained(args.resume, torch_dtype=torch.float32)
    else:
        config = GPT2Config(vocab_size=len(tok), n_positions=512, n_ctx=512,
            n_embd=512, n_layer=12, n_head=8, n_inner=2048,
            activation_function='gelu_new', resid_pdrop=.1, embd_pdrop=.1,
            attn_pdrop=.1, bos_token_id=tok.bos_token_id,
            eos_token_id=tok.eos_token_id, pad_token_id=tok.pad_token_id,
            use_cache=False)
        model = GPT2LMHeadModel(config)
    model.to(args.device)
    parameters = sum(x.numel() for x in model.parameters())
    print(f'Parameters: {parameters:,}; device: {args.device}', flush=True)
    data = {f'{split}_{lang}': np.load(ROOT/'.cache'/f'{split}_{lang}.npy')
            for split in ['train', 'valid'] for lang in ['ru', 'en']}
    def batch(split, rng, language=None):
        rows = []
        for i in range(args.batch_size):
            lang = language or ('ru' if i % 2 == 0 else 'en')
            a = data[f'{split}_{lang}']
            start = rng.randrange(len(a) - args.seq_len + 1)
            rows.append(a[start:start+args.seq_len].copy())
        return torch.tensor(np.stack(rows), dtype=torch.long, device=args.device)
    @torch.no_grad()
    def evaluate():
        model.eval()
        result = {}
        for lang in ['ru','en']:
            rng = random.Random(12345)
            losses = []
            for _ in range(8):
                x = batch('valid', rng, lang)
                losses.append(model(x, labels=x, use_cache=False).loss.item())
            result[lang] = sum(losses) / len(losses)
        model.train()
        return result
    before = evaluate()
    print('Validation before:', before, flush=True)
    optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr,
        betas=(.9,.95), weight_decay=.1, foreach=False)
    rng = random.Random(args.seed)
    started = time.monotonic()
    history = []
    log_path = output/'training_log.jsonl'
    with log_path.open('w') as log:
        for step in range(1, args.steps+1):
            warmup = min(20, max(1,args.steps//10))
            factor = min(1.,step/warmup) * (.1 + .9 * .5 * (1 + math.cos(math.pi*(step-1)/args.steps)))
            for group in optimizer.param_groups:
                group['lr'] = args.lr * factor
            x = batch('train', rng)
            optimizer.zero_grad(set_to_none=True)
            loss = model(x, labels=x, use_cache=False).loss
            if not torch.isfinite(loss):
                raise RuntimeError('Non-finite loss; refusing to export')
            loss.backward()
            grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(),1.)
            optimizer.step()
            row = {'step':step,'loss':loss.item(),'lr':optimizer.param_groups[0]['lr'],
                'grad_norm':float(grad_norm),'elapsed_seconds':round(time.monotonic()-started,2)}
            log.write(json.dumps(row)+'\n')
            log.flush()
            history.append(row)
            if step == 1 or step % 10 == 0:
                print(json.dumps(row), flush=True)
    elapsed = time.monotonic()-started
    after = evaluate()
    report = {'name':'Galactic 2.5 Mini','developer_label':'SpaceModel',
        'status':'experimental short-run checkpoint; not a capable bilingual assistant',
        'initialization':'random' if not args.resume else f'weights from {args.resume}; new optimizer',
        'parameters':parameters, 'steps':args.steps,
        'input_token_presentations':args.steps*args.batch_size*args.seq_len,
        'prediction_targets':args.steps*args.batch_size*(args.seq_len-1),
        'sampling':'equal RU/EN sequences, random windows with replacement',
        'batch_size':args.batch_size,'sequence_length':args.seq_len,
        'max_context':512,'seed':args.seed,'device':args.device,
        'training_seconds':round(elapsed,2), 'optimizer':'AdamW',
        'learning_rate_peak':args.lr,'betas':[.9,.95],'weight_decay':.1,
        'gradient_clip_norm':1.,'training_dtype':'float32','export_dtype':'float16',
        'validation_before_nats':before,'validation_after_nats':after,
        'validation_method':'8 fixed random batches per language from held-out final 5% of each book; not a benchmark',
        'corpus_tokens':{k:len(v) for k,v in data.items()},
        'torch_version':torch.__version__,
        'optimizer_state_included':False}
    # Free optimizer states before FP16 export.
    optimizer.zero_grad(set_to_none=True)
    del optimizer
    model.eval().half()
    model.config.use_cache = True
    model.config.galactic_model_name = 'Galactic 2.5 Mini'
    model.config.galactic_developer = 'SpaceModel'
    model.save_pretrained(output, safe_serialization=True)
    tok.save_pretrained(output)
    (output/'training_report.json').write_text(json.dumps(report,ensure_ascii=False,indent=2),encoding='utf-8')
    print('Validation after:',after,flush=True)
    print('SAVED:', output, flush=True)

if __name__ == '__main__':
    main()