import json
from pathlib import Path
import torch
from transformers import AutoTokenizer
from safetensors import safe_open
from common import ROOT, permit_compute, append
from weights import headers, load as read_weight, save_edit as write_edit

permit_compute()
torch.set_grad_enabled(False)
MODEL = ROOT / 'models/original'
CURRENT = ROOT / 'edits/current'
NAMES = sorted(headers(MODEL))
tok = AutoTokenizer.from_pretrained(MODEL, local_files_only=True)


def load_original(name):
    return read_weight(MODEL, name).float()


def load(name):
    if name not in NAMES:
        raise KeyError(name)
    path = CURRENT / (name + '.safetensors')
    if path.exists():
        with safe_open(str(path), framework='pt', device='cpu') as f:
            return f.get_tensor(name).float()
    return load_original(name)


def save_edit(name, tensor, hypothesis):
    return write_edit(MODEL, CURRENT, name, tensor, hypothesis)


def edited():
    return sorted(p.name.removesuffix('.safetensors') for p in CURRENT.glob('*.safetensors'))


def drop_edit(name):
    permit_compute()
    if name not in NAMES:
        raise KeyError(name)
    for suffix in ('.safetensors', '.json', '.analysis.json', '.direction.npy'):
        path = CURRENT / (name + suffix)
        if path.exists():
            path.unlink()
    append(ROOT / 'logs/experiments.jsonl', {'action': 'drop_edit', 'tensor': name})


print('CPU kernel готов: NAMES, load, load_original, save_edit(name,tensor,hypothesis), drop_edit, edited, tok. Градиенты отключены. Модель не запущена.')