"""Gemma 4 text backbone + LoRA + a multi-label head; no generation/modalities."""
import json
import os
from pathlib import Path
from types import SimpleNamespace
os.environ.setdefault('HF_HOME',str(Path(__file__).resolve().parent/'.cache/huggingface'))
import torch
from accelerate import init_empty_weights
from huggingface_hub import hf_hub_download
from peft import LoraConfig, get_peft_model, PeftModel, TaskType
from safetensors import safe_open
from safetensors.torch import save_file, load_file
from transformers import Gemma4TextConfig, Gemma4TextModel, AutoTokenizer

BASE='google/gemma-4-E2B'
REVISION='d29ff6b45f081a49ee2733a859c9c9c2d95d1a6f'


def load_backbone(device='cuda'):
    full=json.loads(Path(hf_hub_download(BASE,'config.json',revision=REVISION,token=False)).read_text())
    config=Gemma4TextConfig(**full['text_config'])
    config.use_cache=False
    config._name_or_path=BASE
    config._attn_implementation='eager'
    with init_empty_weights():
        backbone=Gemma4TextModel(config)
    path=hf_hub_download(BASE,'model.safetensors',revision=REVISION,token=False)
    prefix='model.language_model.'
    # Match Transformers' documented redundant KV keys for shared layers.
    # Keep strict loading for every parameter used by the text backbone.
    ignored=tuple(backbone._keys_to_ignore_on_load_unexpected)
    with safe_open(path,framework='pt',device='cpu') as f:
        state={k[len(prefix):]:f.get_tensor(k) for k in f.keys()
               if k.startswith(prefix) and not any(k[len(prefix):].startswith(p+'.') for p in ignored)}
    report=backbone.load_state_dict(state,strict=True,assign=True)
    del state
    backbone=backbone.to(device=device,dtype=torch.bfloat16)
    backbone.name_or_path=BASE
    return backbone


class GemmaClassifier(torch.nn.Module):
    def __init__(self,backbone,labels):
        super().__init__()
        self.backbone=backbone
        self.head=torch.nn.Linear(backbone.config.hidden_size,len(labels)).to(next(backbone.parameters()).device)
        self.config=SimpleNamespace(id2label=dict(enumerate(labels)),num_labels=len(labels))

    def forward(self,input_ids,attention_mask,labels=None,**kwargs):
        device=self.head.weight.device
        input_ids=input_ids.to(device);attention_mask=attention_mask.to(device)
        if labels is not None: labels=labels.to(device)
        hidden=self.backbone(input_ids=input_ids,attention_mask=attention_mask,use_cache=False).last_hidden_state
        # Right-padding; final non-padding token sees the preceding causal context.
        index=attention_mask.sum(dim=1)-1
        pooled=hidden[torch.arange(len(index),device=hidden.device),index]
        logits=self.head(pooled.float())
        loss=None if labels is None else torch.nn.functional.binary_cross_entropy_with_logits(logits,labels.float())
        return SimpleNamespace(logits=logits,loss=loss)

    def save(self,path,tokenizer,options):
        path=Path(path);path.mkdir(parents=True,exist_ok=True)
        self.backbone.save_pretrained(path)
        save_file({k:v.detach().cpu().contiguous() for k,v in self.head.state_dict().items()},str(path/'classifier_head.safetensors'))
        tokenizer.save_pretrained(path)
        (path/'classifier_config.json').write_text(json.dumps(dict(base_model=BASE,base_revision=REVISION,
          task=path.parents[1].name,
          labels=list(self.config.id2label.values()),pooling='last_non_padding_token',precision='BF16 backbone; FP32 head and trainable LoRA',
          architecture='Gemma4TextModel+LoRA+Linear',text_only=True,generation_enabled=False,lora_rank=8),indent=2)+'\n')
        (path/'thresholds.json').write_text(json.dumps(options,indent=2)+'\n')


def new_classifier(labels,device='cuda'):
    backbone=load_backbone(device)
    backbone=get_peft_model(backbone,LoraConfig(task_type=TaskType.FEATURE_EXTRACTION,revision=REVISION,r=8,lora_alpha=16,lora_dropout=.05,
                            target_modules=['q_proj','v_proj'],bias='none'))
    return GemmaClassifier(backbone,labels)


def load_classifier(path,device='cuda',trainable=False):
    path=Path(path);config=json.loads((path/'classifier_config.json').read_text())
    assert config['base_revision']==REVISION
    backbone=PeftModel.from_pretrained(load_backbone(device),path,is_trainable=trainable)
    model=GemmaClassifier(backbone,config['labels'])
    model.config.repo_id=f'mts-the-architect/{config["task"]}-gemma4-e2b-bf16-lora'
    model.head.load_state_dict(load_file(str(path/'classifier_head.safetensors'),device=device))
    tokenizer=AutoTokenizer.from_pretrained(path);tokenizer.padding_side='right'
    return model,tokenizer,json.loads((path/'thresholds.json').read_text())


if __name__=='__main__':
    torch.set_num_threads(8)
    tokenizer=AutoTokenizer.from_pretrained(BASE,revision=REVISION);tokenizer.padding_side='right'
    model=new_classifier(['care','equality'])
    batch=tokenizer(['Everyone deserves the same rights.','We should protect vulnerable people.'],return_tensors='pt',padding=True)
    batch={k:v.cuda() for k,v in batch.items()}
    loss=model(**batch,labels=torch.tensor([[0.,1.],[1.,0.]],device='cuda')).loss
    loss.backward()
    print('GPU forward/backward passed',float(loss.detach()),'total parameters',sum(p.numel() for p in model.parameters()),
          'trainable',sum(p.numel() for p in model.parameters() if p.requires_grad),flush=True)