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


def matched_direction(productive, repetitive):
    a, b = np.asarray(productive), np.asarray(repetitive)
    if a.shape != b.shape or a.ndim != 2:
        raise ValueError('Нужны попарные средние активации одинаковой формы [contexts, width]')
    direction = (b - a).mean(0)
    norm = np.linalg.norm(direction)
    if not np.isfinite(norm) or norm == 0:
        raise ValueError('Направление не определено')
    return direction / norm


def writer_scores(weight, direction, productive_mlp, repetitive_mlp):
    coefficients = np.asarray(direction) @ weight.float().numpy()
    deltas = np.asarray(repetitive_mlp) - np.asarray(productive_mlp)
    values = deltas * coefficients
    return values.mean(0), (values > 0).mean(0)


if __name__ == '__main__':
    p = argparse.ArgumentParser()
    p.add_argument('--model', required=True)
    p.add_argument('--tensor', required=True)
    p.add_argument('--productive', required=True)
    p.add_argument('--repetitive', required=True)
    p.add_argument('--productive-mlp', required=True)
    p.add_argument('--repetitive-mlp', required=True)
    p.add_argument('--columns', type=int, required=True)
    p.add_argument('--alpha', type=float, required=True)
    p.add_argument('--out', required=True)
    p.add_argument('--hypothesis', required=True)
    a = p.parse_args()
    permit_compute()
    direction = matched_direction(np.load(inside(a.productive)), np.load(inside(a.repetitive)))
    weight = load(a.model, a.tensor)
    score, consistency = writer_scores(weight, direction, np.load(inside(a.productive_mlp)), np.load(inside(a.repetitive_mlp)))
    if not 1 <= a.columns < weight.shape[1]:
        raise ValueError('Неверное количество столбцов')
    columns = np.argsort(score)[-a.columns:]
    edited = project_writers(weight, direction, columns, a.alpha)
    record = save_edit(a.model, a.out, a.tensor, edited, a.hypothesis)
    save_json(inside(a.out) / (a.tensor + '.analysis.json'), {'columns': columns.tolist(), 'alpha': a.alpha,
                                                          'scores': score[columns].tolist(),
                                                          'consistency': consistency[columns].tolist(),
                                                          'input_files': [a.productive, a.repetitive, a.productive_mlp, a.repetitive_mlp]})
    np.save(inside(a.out) / (a.tensor + '.direction.npy'), direction)
    print(json.dumps(record, ensure_ascii=False))