#!/usr/bin/env python3
import argparse
import json
import os
import sys

import torch
from transformers import GPT2LMHeadModel

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from model import GPT, GPTConfig


TRANSPOSED = (
    'attn.c_attn.weight',
    'attn.c_proj.weight',
    'mlp.c_fc.weight',
    'mlp.c_proj.weight',
)


def copy_tensor(dst, src, name):
    value = src.t() if any(name.endswith(suffix) for suffix in TRANSPOSED) else src
    if dst.shape != value.shape:
        raise ValueError(f'shape mismatch for {name}: dst={tuple(dst.shape)} src={tuple(value.shape)}')
    dst.copy_(value)


def copy_hf_block(model, hf_state, src_layer, dst_layer):
    own_state = model.state_dict()
    prefix_hf = f'transformer.h.{src_layer}.'
    prefix_own = f'transformer.h.{dst_layer}.'
    keys = [
        'ln_1.weight',
        'ln_1.bias',
        'attn.c_attn.weight',
        'attn.c_attn.bias',
        'attn.c_proj.weight',
        'attn.c_proj.bias',
        'ln_2.weight',
        'ln_2.bias',
        'mlp.c_fc.weight',
        'mlp.c_fc.bias',
        'mlp.c_proj.weight',
        'mlp.c_proj.bias',
    ]
    for suffix in keys:
        own_key = prefix_own + suffix
        hf_key = prefix_hf + suffix
        copy_tensor(own_state[own_key], hf_state[hf_key], own_key)


def zero_residual_outputs(model, layer_ids):
    with torch.no_grad():
        for i in layer_ids:
            block = model.transformer.h[i]
            block.attn.c_proj.weight.zero_()
            block.mlp.c_proj.weight.zero_()
            if block.attn.c_proj.bias is not None:
                block.attn.c_proj.bias.zero_()
            if block.mlp.c_proj.bias is not None:
                block.mlp.c_proj.bias.zero_()


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--out_dir', default='out-init-distilgpt2-plus2')
    parser.add_argument('--model_name', default='distilgpt2')
    parser.add_argument('--dropout', type=float, default=0.0)
    args = parser.parse_args()

    os.makedirs(args.out_dir, exist_ok=True)

    hf = GPT2LMHeadModel.from_pretrained(args.model_name)
    hf.eval()
    hf_state = hf.state_dict()
    hf_cfg = hf.config
    if hf_cfg.vocab_size != 50257 or hf_cfg.n_layer != 6 or hf_cfg.n_head != 12 or hf_cfg.n_embd != 768:
        raise ValueError(
            'Expected distilgpt2 shape 6L/12H/768d/vocab50257, got '
            f'{hf_cfg.n_layer}L/{hf_cfg.n_head}H/{hf_cfg.n_embd}d/vocab{hf_cfg.vocab_size}'
        )

    model_args = dict(
        block_size=1024,
        vocab_size=50257,
        n_layer=8,
        n_head=12,
        n_embd=768,
        dropout=args.dropout,
        bias=True,
    )
    torch.manual_seed(1337)
    model = GPT(GPTConfig(**model_args))

    with torch.no_grad():
        own_state = model.state_dict()
        own_state['transformer.wte.weight'].copy_(hf_state['transformer.wte.weight'])
        own_state['transformer.wpe.weight'].copy_(hf_state['transformer.wpe.weight'])
        own_state['transformer.ln_f.weight'].copy_(hf_state['transformer.ln_f.weight'])
        own_state['transformer.ln_f.bias'].copy_(hf_state['transformer.ln_f.bias'])
        own_state['lm_head.weight'].copy_(hf_state['transformer.wte.weight'])
        for i in range(6):
            copy_hf_block(model, hf_state, i, i)
        zero_residual_outputs(model, [6, 7])

    total_params = sum(p.numel() for p in model.parameters())
    if total_params >= 100_000_000:
        raise ValueError(f'model has {total_params:,} params, exceeds 100M limit')

    checkpoint = {
        'model': model.state_dict(),
        'optimizer': None,
        'model_args': model_args,
        'iter_num': 0,
        'best_val_loss': 1e9,
        'config': {
            'source_model': args.model_name,
            'architecture_note': 'distilgpt2 6 pretrained layers + 2 zero-output residual layers',
            'total_params': total_params,
        },
    }
    torch.save(checkpoint, os.path.join(args.out_dir, 'checkpoint.pt'))
    with open(os.path.join(args.out_dir, 'config.json'), 'w', encoding='utf-8') as f:
        json.dump(model_args, f, indent=2)
    with open(os.path.join(args.out_dir, 'init_info.json'), 'w', encoding='utf-8') as f:
        json.dump(checkpoint['config'], f, indent=2)
    print(f'wrote {args.out_dir}/checkpoint.pt')
    print(f'total params: {total_params:,}')
    print('tokenizer/vocab: GPT-2 ids, vocab_size=50257')


if __name__ == '__main__':
    main()