"""
GPU worker for separator-attention and random-target intervention experiments.
For each random sequence, at each sorted output position:
  1. Check if layer 0 max attention is on the separator token.
     If yes, intervene with standard method (ub=60) and record result.
  2. Intervene by boosting a random unsorted number's attention and record result.
Collects per-number success data across many trials.
"""
import argparse
import json
import os
import sys
import time
import numpy as np
import torch

sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'grid-run'))
from model_analysis import GPT, GPTConfig, GPTIntervention

INTENSITIES = [2.0, 6.0, 10.0]
UB_STANDARD = 60


def remap_state_dict(sd):
    new = {}
    for k, v in sd.items():
        nk = k
        for i in range(10):
            nk = nk.replace(f'transformer.h.{i}.attn.', f'transformer.h.{i}.c_attn.')
            nk = nk.replace(f'transformer.h.{i}.mlp.', f'transformer.h.{i}.c_fc.')
        new[nk] = v
    return new


def load_model(ckpt_path, device):
    ckpt = torch.load(ckpt_path, map_location='cpu')
    mc = ckpt['model_config']
    config = GPTConfig(block_size=mc['block_size'], vocab_size=mc['vocab_size'] - 1,
                       with_layer_norm=mc.get('use_final_LN', True))
    model = GPT(config)
    sd = remap_state_dict(ckpt['model_state_dict'])
    wpe_max = config.block_size * 4 + 1
    if 'transformer.wpe.weight' in sd and sd['transformer.wpe.weight'].shape[0] > wpe_max:
        sd['transformer.wpe.weight'] = sd['transformer.wpe.weight'][:wpe_max]
    for k in [k for k in sd if k.endswith('.c_attn.bias') and 'c_attn.c_attn' not in k]:
        del sd[k]
    if 'lm_head.weight' in sd:
        del sd['lm_head.weight']
    model.load_state_dict(sd, strict=False)
    model.to(device)
    model.eval()
    return model, config


def get_batch(vocab_size, block_size, device):
    x = torch.randperm(vocab_size)[:block_size]
    vals, _ = torch.sort(x)
    return torch.cat((x, torch.tensor([vocab_size]), vals), dim=0).unsqueeze(0).to(device)


def try_standard_intervention(model, idx, config, location, intensity):
    try:
        im = GPTIntervention(model, idx)
        im.intervent_attention(
            attention_layer_num=0, location=location,
            unsorted_lb=UB_STANDARD, unsorted_ub=UB_STANDARD,
            unsorted_lb_num=0, unsorted_ub_num=1,
            unsorted_intensity_inc=intensity,
            sorted_lb=0, sorted_num=0, sorted_intensity_inc=0.0)
        g, n = im.check_if_still_works()
        im.revert_attention(0)
        return g == n
    except:
        return None


def try_random_intervention(model, idx, config, location, intensity):
    vs = config.vocab_size
    try:
        im = GPTIntervention(model, idx)
        im.intervent_attention(
            attention_layer_num=0, location=location,
            unsorted_lb=0, unsorted_ub=vs,
            unsorted_lb_num=0, unsorted_ub_num=1,
            unsorted_intensity_inc=intensity,
            sorted_lb=0, sorted_num=0, sorted_intensity_inc=0.0)
        g, n = im.check_if_still_works()
        im.revert_attention(0)
        return g == n
    except:
        pass
    try:
        im = GPTIntervention(model, idx)
        im.intervent_attention(
            attention_layer_num=0, location=location,
            unsorted_lb=vs, unsorted_ub=0,
            unsorted_lb_num=1, unsorted_ub_num=0,
            unsorted_intensity_inc=intensity,
            sorted_lb=0, sorted_num=0, sorted_intensity_inc=0.0)
        g, n = im.check_if_still_works()
        im.revert_attention(0)
        return g == n
    except:
        return None


def run_trials(model, config, device, n_trials):
    bs = config.block_size
    vs = config.vocab_size
    sep_pos = bs

    sep_records = []
    rand_records = []

    for trial in range(n_trials):
        idx = get_batch(vs, bs, device)

        with torch.no_grad():
            logits, _ = model(idx)

        attn_layer0 = model.transformer.h[0].c_attn.attn

        for p in range(bs - 1):
            sorted_loc = bs + 1 + p
            number_val = idx[0, sorted_loc].item()
            next_num = idx[0, sorted_loc + 1].item()

            attn_row = attn_layer0[sorted_loc, :sorted_loc + 1]
            max_attn_pos = attn_row.argmax().item()
            attends_to_sep = (max_attn_pos == sep_pos)

            for intensity in INTENSITIES:
                if attends_to_sep:
                    result = try_standard_intervention(model, idx, config, sorted_loc, intensity)
                    if result is not None:
                        sep_records.append((number_val, intensity, int(result)))

                result_rand = try_random_intervention(model, idx, config, sorted_loc, intensity)
                if result_rand is not None:
                    rand_records.append((number_val, intensity, int(result_rand)))

        if (trial + 1) % 200 == 0:
            print(f"  Trial {trial+1}/{n_trials}: sep={len(sep_records)}, rand={len(rand_records)}",
                  flush=True)

    return np.array(sep_records, dtype=np.int32), np.array(rand_records, dtype=np.int32)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--ckpt', required=True)
    parser.add_argument('--gpu', type=int, required=True)
    parser.add_argument('--trials', type=int, default=1000)
    parser.add_argument('--out', required=True)
    args = parser.parse_args()

    os.environ['CUDA_VISIBLE_DEVICES'] = str(args.gpu)
    device = 'cuda'

    print(f"GPU {args.gpu}: loading model...", flush=True)
    t0 = time.time()
    model, config = load_model(args.ckpt, device)
    print(f"  Loaded in {time.time()-t0:.1f}s", flush=True)

    print(f"GPU {args.gpu}: running {args.trials} trials...", flush=True)
    t0 = time.time()
    sep_data, rand_data = run_trials(model, config, device, args.trials)
    elapsed = time.time() - t0

    os.makedirs(os.path.dirname(args.out), exist_ok=True)
    np.savez(args.out, sep_data=sep_data, rand_data=rand_data)
    print(f"GPU {args.gpu}: done in {elapsed:.0f}s, "
          f"sep={len(sep_data)} rand={len(rand_data)}", flush=True)


if __name__ == '__main__':
    main()