import torch
import argparse

from load_model import load_model
from transformers import GPT2TokenizerFast
import torch.nn.functional as F
import sampling
import numpy as np


def main():
    parser = argparse.ArgumentParser(description="Generate some samples")
    parser.add_argument("--model_path", default="exp_local/openwebtext/x", type=str)
    parser.add_argument("--visualize", default=False, type=bool)
    parser.add_argument("--dataset", default="wikitext103", type=str)
    parser.add_argument("--batch_size", type=int, default=4)
    parser.add_argument("--steps", type=int, default=1024)
    args = parser.parse_args()

    
    device = torch.device('cuda')
    model, graph, noise = load_model(args.model_path, device)
    tokenizer = GPT2TokenizerFast.from_pretrained('gpt2')




    loop_max = 100
    jumps = []

    for loop in range(loop_max):
        print(loop)

        sampling_fn = sampling.get_pc_sampler(
            graph, noise, (args.batch_size, 128), 'euler', args.steps, args.visualize, device=device
        )

        samples, switch_hist = sampling_fn(model)
        jumps.append(switch_hist)
        print('Average Jumps:', np.array(jumps).mean())
        print('Std Jumps:', np.array(jumps).std()/np.sqrt(0*np.array(jumps)+1).sum())

if __name__=="__main__":
    main()
