import argparse
import hashlib
import json
import numpy as np
import torch
from common import ROOT, inside, permit_compute, save_json
from weights import load, save_edit, export, project_writers

permit_compute()
p = argparse.ArgumentParser()
p.add_argument('--analysis', required=True)
p.add_argument('--layers', required=True)
p.add_argument('--alpha', type=float, required=True)
p.add_argument('--topk', type=int, default=64)
p.add_argument('--name', required=True)
p.add_argument('--control', choices=['targeted','orthogonal','random_columns'], default='targeted')
p.add_argument('--seed', type=int, default=1701)
p.add_argument('--hypothesis', required=True)
p.add_argument('--protect-coefficients', action='store_true')
p.add_argument('--match-protected-norm', action='store_true')
a = p.parse_args()
if not a.name or any(c not in 'abcdefghijklmnopqrstuvwxyz0123456789_-' for c in a.name):
    raise ValueError('Используй простое уникальное имя кандидата')
directory, model = ROOT/'edits'/a.name, ROOT/'models'/a.name
if directory.exists() or model.exists():
    raise ValueError('Имя уже занято')
z = np.load(inside(a.analysis)/'directions_and_writers.npz')
records = []
for layer in map(int,a.layers.split(',')):
    if not 0 <= layer < 24:
        raise ValueError('Недопустимый слой')
    name = f'model.language_model.layers.{layer}.mlp.down_proj.weight'
    weight = load(ROOT/'models/original', name).float()
    u = torch.from_numpy(z[f'u_{layer}'])
    scores = z[f'scores_{layer}']
    k = a.topk or weight.shape[1]
    if not 1 <= k <= weight.shape[1]:
        raise ValueError('Недопустимое число столбцов')
    columns = np.argsort(scores)[-k:].copy()
    targeted = project_writers(weight, u, columns, a.alpha)
    candidate = targeted
    rng = np.random.default_rng(a.seed+layer)
    if a.control == 'orthogonal':
        v = torch.tensor(rng.normal(size=u.numel()),dtype=torch.float32)
        v = v-u*torch.dot(u,v)
        v = v/v.norm()
        delta = weight-project_writers(weight,v,columns,1)
        candidate = weight-delta*((weight-targeted).norm()/delta.norm())
    elif a.control == 'random_columns':
        selected = rng.choice(weight.shape[1],size=k,replace=False)
        delta = weight-project_writers(weight,u,selected,1)
        candidate = weight-delta*((weight-targeted).norm()/delta.norm())
        columns = selected
    protection = None
    if a.protect_coefficients:
        if a.control != 'targeted':
            raise ValueError('Coefficient protection currently requires targeted control')
        productive = z[f'productive_mlp_{layer}'][:,columns].astype(np.float64)
        _, singular, basis = np.linalg.svd(productive,full_matrices=False)
        basis = basis[singular>singular[0]*1e-6]
        coefficients = (u@weight).numpy()[columns].astype(np.float64)
        projected = coefficients-basis.T@(basis@coefficients)
        retained = float(np.linalg.norm(projected)/np.linalg.norm(coefficients))
        if a.match_protected_norm:
            if retained < 1e-8:
                raise ValueError('Protected coefficient vector vanishes')
            projected = projected/retained
        vector = torch.zeros(weight.shape[1],dtype=torch.float32)
        vector[columns] = torch.from_numpy(projected.astype(np.float32))
        candidate = weight-a.alpha*torch.outer(u,vector)
        protection = {'constraint':'coefficient vector orthogonal to original productive MLP means on selected columns',
                      'retained_coefficient_norm_fraction_before_matching':retained,
                      'norm_matched':a.match_protected_norm,
                      'max_constraint_residual':float(np.max(np.abs(productive@projected)))}
    record = save_edit(ROOT/'models/original',directory,name,candidate,a.hypothesis)
    record.update({'layer':layer,'alpha':a.alpha,'control':a.control,'columns':columns.tolist(),'coefficient_protection':protection})
    records.append(record)
save_json(directory/'candidate.json',{'origin':'independent_sol','analysis':a.analysis,'hypothesis':a.hypothesis,
                                    'analysis_sha256':hashlib.sha256((inside(a.analysis)/'directions_and_writers.npz').read_bytes()).hexdigest(),
                                    'control':a.control,'records':records,'status':'unconfirmed'})
export(ROOT/'models/original',directory,model)
print(json.dumps({'candidate':a.name,'model':str(model),'changed_entries':sum(r['changed_entries'] for r in records)},ensure_ascii=False))