"""Quick run: intensity up to 30 for the 1M checkpoint, layer 0, ub=60."""
import os
import sys
import numpy as np
import torch
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

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

CKPT = os.path.join(os.path.dirname(os.path.abspath(__file__)),
    'sortgpt_k16_methfixed_mlp1_L2_N256_E64_pos0_fln1_wd0p0_lr0p03_dseed1337_iseed1337__final.pt')
PLOT_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'outputs',
    'plots_V256_B16_LR3e-2_MI1000000_E64_H1_L2_ds1337_is1337_ckpt1000000')

INTENSITIES = [1.0, 2.0, 4.0, 6.0, 8.0, 10.0, 14.0, 18.0, 22.0, 26.0, 30.0]
UB = 60
MIN_VALID = 200
GPU = 0


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(device):
    ckpt = torch.load(CKPT, 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(vs, bs, device):
    x = torch.randperm(vs)[:bs]
    vals, _ = torch.sort(x)
    return torch.cat((x, torch.tensor([vs]), vals), dim=0).unsqueeze(0).to(device)


def main():
    os.environ['CUDA_VISIBLE_DEVICES'] = str(GPU)
    device = 'cuda'
    print("Loading model...", flush=True)
    model, config = load_model(device)
    bs = config.block_size
    vs = config.vocab_size
    location = bs + 5

    for layer in [0, 1]:
        print(f"\nLayer {layer}:", flush=True)
        rates, counts = [], []
        for intens in INTENSITIES:
            attempts, rounds = [], 0
            while len(attempts) < MIN_VALID and rounds < 3000:
                rounds += 1
                idx = get_batch(vs, bs, device)
                try:
                    im = GPTIntervention(model, idx)
                    im.intervent_attention(
                        attention_layer_num=layer, location=location,
                        unsorted_lb=UB, unsorted_ub=UB,
                        unsorted_lb_num=0, unsorted_ub_num=1,
                        unsorted_intensity_inc=intens,
                        sorted_lb=0, sorted_num=0, sorted_intensity_inc=0.0)
                    g, n = im.check_if_still_works()
                    attempts.append(g == n)
                    im.revert_attention(layer)
                except:
                    continue
            rate = sum(attempts) / len(attempts) if attempts else 0.0
            rates.append(rate)
            counts.append(len(attempts))
            print(f"  intensity={intens:5.1f}: success={rate:.4f} (n={len(attempts)})", flush=True)

        intensities = np.array(INTENSITIES)
        rates = np.array(rates)

        plt.figure(figsize=(5.5, 3.8))
        plt.plot(intensities, rates, marker='o', linewidth=1.5, markersize=5, color='#e6850e')
        plt.xlabel('Intervention Intensity', fontsize=10)
        plt.ylabel('Success Probability', fontsize=10)
        plt.title(f'Robustness to Attention Intervention (Layer {layer})  [ub={UB}]\n'
                  f'V=256  B=16  lr=0.03  iters=1000000  dseed=1337  iseed=1337',
                  fontsize=10, fontweight='bold')
        plt.grid(True, alpha=0.3)
        plt.xticks(intensities[::2], fontsize=9)
        plt.yticks(fontsize=9)
        plt.ylim(0, 1.05)
        plt.tight_layout()
        out = os.path.join(PLOT_DIR, f'intensity_layer{layer}_ub60_high.png')
        plt.savefig(out, dpi=300, bbox_inches='tight')
        plt.close()
        print(f"  Saved: {out}", flush=True)

    print("\nDone!", flush=True)


if __name__ == '__main__':
    main()