import argparse
import json
import numpy as np
from common import ROOT, inside, save_json
from weights import load


def mean_window(directory, metadata, text, signal, window):
    total, count = None, 0
    for file in metadata['files']:
        if file['text'] != text or file['signal'] != signal:
            continue
        start, end = file['token_start'], file['token_end']
        left, right = max(start, window[0]), min(end, window[1])
        if right <= left:
            continue
        array = np.load(directory / file['name'], mmap_mode='r')[0, left-start:right-start]
        value = array.sum(0, dtype=np.float64)
        total = value if total is None else total + value
        count += right-left
    if count != window[1]-window[0]:
        raise ValueError(f'Incomplete {signal} window for text {text}: {count}')
    return (total/count).astype(np.float32)


def branch_logps(directory, metadata, labels):
    result = []
    for row in labels:
        values = []
        for file in metadata['files']:
            if file['text'] != row['text'] or file['signal'] != 'token_logprobs':
                continue
            array = np.load(directory / file['name'])
            start = max(0, row['branch_start']-1-file['token_start'])
            end = min(len(array), row.get('branch_end',file['token_start']+len(array)+1)-1-file['token_start'])
            if end > start:
                values.extend(array[start:end].tolist())
        if not values:
            raise ValueError('Missing branch scores')
        result.append(dict(row, mean_logp=float(np.mean(values)), logp_sum=float(np.sum(values)), scored_tokens=len(values)))
    return result


def main():
    p = argparse.ArgumentParser()
    p.add_argument('--probe', required=True)
    p.add_argument('--labels', default=str(ROOT/'configs/matched_probe_001_labels.json'))
    p.add_argument('--out', required=True)
    p.add_argument('--model', default=str(ROOT/'models/original'))
    p.add_argument('--logp-only', action='store_true')
    p.add_argument('--protect-productive', action='store_true')
    args = p.parse_args()
    directory, out = inside(args.probe), inside(args.out)
    metadata = json.loads((directory/'meta.json').read_text())
    if not metadata.get('complete'):
        raise ValueError('Probe is incomplete')
    labels = json.loads(inside(args.labels).read_text())
    out.mkdir(parents=True, exist_ok=True)
    save_json(out/'branch_logps.json', branch_logps(directory, metadata, labels))
    if args.logp_only:
        return
    arrays, ranking = {}, []
    discovery = [r for r in labels if r['split']=='discovery']
    pairs = sorted({r['pair'] for r in discovery})
    for layer in range(24):
        hidden, mlp = {}, {}
        for row in discovery:
            key = (row['pair'], row['condition'])
            hidden[key] = mean_window(directory, metadata, row['text'], f'hidden_{layer+1}', row['mean_window'])
            mlp[key] = mean_window(directory, metadata, row['text'], f'mlp_act_{layer}', row['mean_window'])
        hp = np.stack([hidden[(pair,'productive')] for pair in pairs])
        hr = np.stack([hidden[(pair,'repetitive')] for pair in pairs])
        dh = hr-hp
        dm = np.stack([mlp[(pair,'repetitive')]-mlp[(pair,'productive')] for pair in pairs])
        direction = dh.mean(0)
        if args.protect_productive:
            _, singular, basis = np.linalg.svd(hp,full_matrices=False)
            basis = basis[singular>singular[0]*1e-6]
            direction = direction-basis.T@(basis@direction)
        magnitude = np.linalg.norm(direction)
        direction = direction/magnitude
        weight = load(args.model, f'model.language_model.layers.{layer}.mlp.down_proj.weight').float().numpy()
        projection = direction @ weight
        scores = dm * projection
        mean = scores.mean(0)
        columns = np.argsort(mean)[-64:][::-1].copy()
        arrays[f'u_{layer}'] = direction
        arrays[f'columns_{layer}'] = columns
        arrays[f'scores_{layer}'] = mean
        arrays[f'pair_delta_mlp_{layer}'] = dm
        arrays[f'productive_mlp_{layer}'] = np.stack([mlp[(pair,'productive')] for pair in pairs])
        arrays[f'repetitive_mlp_{layer}'] = np.stack([mlp[(pair,'repetitive')] for pair in pairs])
        arrays[f'pair_delta_hidden_{layer}'] = dh
        arrays[f'productive_hidden_{layer}'] = hp
        arrays[f'repetitive_hidden_{layer}'] = hr
        ranking.append({'layer':layer,'direction_norm':float(magnitude),'top64_writer_contrast':float(mean[columns].sum()),
                        'top64_fraction_positive_contrast':float(mean[columns].sum()/np.maximum(mean,0).sum()),
                        'direction_pair_cosines':(dh@direction/np.linalg.norm(dh,axis=1)).tolist(),
                        'top_neurons':[{'neuron':int(j),'contrast':float(mean[j]),'positive_pairs':int((scores[:,j]>0).sum())} for j in columns[:12]],
                        'top64_consistent_pairs_mean':float((scores[:,columns]>0).mean()),
                        'productive_protection':args.protect_productive,
                        'max_productive_projection':float(np.max(np.abs(hp@direction)))})
    np.savez(out/'directions_and_writers.npz', **arrays)
    save_json(out/'layer_ranking.json', sorted(ranking,key=lambda x:x['top64_writer_contrast'],reverse=True))
    print(json.dumps(sorted(ranking,key=lambda x:x['top64_writer_contrast'],reverse=True)[:5],indent=2))


if __name__ == '__main__':
    main()