"""Actual low-memory CPU LoRA experiment. Frozen base, adapters in last two blocks."""
import os
os.environ.setdefault('TOKENIZERS_PARALLELISM', 'false')
os.environ.setdefault('HF_HUB_DISABLE_XET', '1')
import time, json, hashlib, random, resource, sys
from pathlib import Path
import torch
import torch.nn.functional as F
from transformers import AutoModelForCausalLM, AutoTokenizer, set_seed
from peft import LoraConfig, get_peft_model, get_peft_model_state_dict

ROOT=Path(__file__).resolve().parents[1]
if (ROOT/'adapter_model.safetensors').exists() and os.environ.get('TERRA_OVERWRITE_ADAPTER') != '1':
    raise SystemExit('Adapter already exists. Work in a COPY and set TERRA_OVERWRITE_ADAPTER=1 to retrain.')
CFG=json.loads((ROOT/'terra_config.json').read_text())
BASE=CFG['base_model']; REV=CFG['base_revision']
torch.set_num_threads(2)
torch.set_num_interop_threads(1)
set_seed(42)
START=time.time()

def log(event, **fields):
    record={'event':event,'elapsed_s':round(time.time()-START,2),**fields}
    print(json.dumps(record,ensure_ascii=False),flush=True)
    with (ROOT/'training_events.jsonl').open('a') as f: f.write(json.dumps(record,ensure_ascii=False)+'\n')

def digest_state(model):
    h=hashlib.sha256()
    for key,v in sorted(get_peft_model_state_dict(model).items()):
        h.update(key.encode()); h.update(v.detach().cpu().float().contiguous().numpy().tobytes())
    return h.hexdigest()

log('loading',base_model=BASE,revision=REV,device='cpu',dtype='bfloat16')
tok=AutoTokenizer.from_pretrained(BASE,revision=REV,trust_remote_code=False)
model=AutoModelForCausalLM.from_pretrained(BASE,revision=REV,torch_dtype=torch.bfloat16,
    low_cpu_mem_usage=True,trust_remote_code=False,attn_implementation='sdpa')
base_count=sum(p.numel() for p in model.parameters())
model.config.use_cache=False
last=model.config.num_hidden_layers-1
modules=[f'model.layers.{i}.self_attn.{name}' for i in [last-1,last] for name in ['q_proj','v_proj','o_proj']]
model=get_peft_model(model,LoraConfig(task_type='CAUSAL_LM',r=8,lora_alpha=16,lora_dropout=0.0,
    target_modules=modules,bias='none',base_model_name_or_path=BASE,revision=REV))
# No input-gradient or checkpointing hook: frozen early blocks need no autograd graph.
trainable=[p for p in model.parameters() if p.requires_grad]
assert trainable and all('lora_' in n for n,p in model.named_parameters() if p.requires_grad)
initial=digest_state(model)
log('loaded',base_parameters=base_count,trainable_parameters=sum(p.numel() for p in trainable),
    target_modules=modules,rss_peak_mib=round(resource.getrusage(resource.RUSAGE_SELF).ru_maxrss/1024,1))

def load_rows(name):
    rows=[json.loads(l) for l in (ROOT/'data'/name).read_text().splitlines() if l.strip()]
    result=[]
    for row in rows:
        msgs=[{'role':'system','content':CFG['system_prompt']}]+row['messages']
        pre=tok.apply_chat_template(msgs[:-1],tokenize=True,add_generation_prompt=True)
        ids=tok.apply_chat_template(msgs,tokenize=True,add_generation_prompt=False)
        assert ids[:len(pre)]==pre and len(ids)<=256
        result.append({'ids':torch.tensor([ids],dtype=torch.long),'prompt_len':len(pre),'language':row['language']})
    return result
train=load_rows('train.example.jsonl'); valid=load_rows('eval.example.jsonl')

# Precompute output-token positions. Compute logits only for supervised answer
# positions, avoiding full-sequence 151k-vocabulary logits and their gradients.
# This is the standard next-token cross-entropy objective, not an approximation.
def loss_for(row):
    base=model.get_base_model()
    ids=row['ids']
    hidden=base.model(input_ids=ids,attention_mask=torch.ones_like(ids),use_cache=False).last_hidden_state
    start=row['prompt_len']-1
    answer_hidden=hidden[0,start:-1,:]
    labels=ids[0,start+1:]
    logits=base.lm_head(answer_hidden)
    loss=F.cross_entropy(logits.float(),labels)
    return loss,labels.numel()

@torch.no_grad()
def evaluate():
    model.eval(); totals={k:[0.,0] for k in ['ru','en']}
    for i,row in enumerate(valid):
        value,n=loss_for(row); totals[row['language']][0]+=float(value)*n;totals[row['language']][1]+=n
        log('eval_item',item=i+1,loss=round(float(value),6))
    results={k:v[0]/v[1] for k,v in totals.items()}
    results['overall']=sum(v[0] for v in totals.values())/sum(v[1] for v in totals.values())
    return results

before=evaluate();log('eval_before',loss=before)
optimizer=torch.optim.AdamW(trainable,lr=5e-4,weight_decay=0.0)
order=list(range(len(train)));random.Random(42).shuffle(order)
losses=[]
for step,index in enumerate(order,1):
    model.train();optimizer.zero_grad(set_to_none=True)
    t=time.time();loss,n=loss_for(train[index]);assert torch.isfinite(loss)
    loss.backward()
    norm=torch.nn.utils.clip_grad_norm_(trainable,1.0)
    assert torch.isfinite(norm)
    optimizer.step();losses.append(float(loss.detach()))
    log('train_step',step=step,total_steps=len(order),example=index,language=train[index]['language'],
        loss=round(losses[-1],6),grad_norm=round(float(norm),6),answer_tokens=n,
        step_s=round(time.time()-t,2),rss_peak_mib=round(resource.getrusage(resource.RUSAGE_SELF).ru_maxrss/1024,1))
    if step%8==0:
        model.save_pretrained(ROOT,safe_serialization=True,save_embedding_layers=False)
        log('checkpoint',step=step)
after=evaluate();log('eval_after',loss=after)
final=digest_state(model);assert final!=initial
model.save_pretrained(ROOT,safe_serialization=True,save_embedding_layers=False)
old='You are Qwen, created by Alibaba Cloud. You are a helpful assistant.'
assert old in tok.chat_template
tok.chat_template=tok.chat_template.replace(old,CFG['system_prompt'].replace('\\','\\\\').replace("'","\\'").replace('\n','\\n'))
tok.save_pretrained(ROOT)
manifest={'name':'Terra 2.6','publisher':'SpaceModel','base_model':BASE,'base_revision':REV,
    'method':'LoRA SFT on CPU; adapters in final two attention blocks',
    'demo_only':True,'completed':True,'base_parameters':base_count,
    'trainable_parameters':sum(p.numel() for p in trainable),'target_modules':modules,
    'steps':len(order),'train_dialogs':len(train),'eval_dialogs':len(valid),
    'seed':42,'epochs':1,'learning_rate':5e-4,'batch_size':1,'gradient_accumulation':1,
    'rank':8,'alpha':16,'dtype_base':'bfloat16','dtype_adapter':'float32',
    'optimizer':'AdamW','weight_decay':0.0,'gradient_clip_norm':1.0,
    'train_loss_mean':sum(losses)/len(losses),'validation_loss_before':before,'validation_loss_after':after,
    'initial_adapter_sha256':initial,'final_adapter_sha256':final,'adapter_weights_changed':True,
    'data_sha256':{name:hashlib.sha256((ROOT/'data'/name).read_bytes()).hexdigest() for name in ['train.example.jsonl','eval.example.jsonl']},
    'wall_time_seconds':round(time.time()-START,2),'rss_peak_mib':round(resource.getrusage(resource.RUSAGE_SELF).ru_maxrss/1024,1),
    'torch_version':torch.__version__,'python_version':sys.version,
    'note':'Tiny illustrative run; no independent quality improvement claim. Adapter requires the original base weights.'}
(ROOT/'terra_training.json').write_text(json.dumps(manifest,ensure_ascii=False,indent=2)+'\n')
log('complete',adapter_changed=True,seconds=manifest['wall_time_seconds'])