import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
import torch.backends.cudnn as cudnn
import torch.utils
from torch.utils.data import Dataset, DataLoader

import os
import argparse

# from ..utils import progress_bar
import time
import random
import numpy as np
import pickle
import hashlib
import io
import torch.utils.data
from tqdm import tqdm
from transformers import AutoModelForCausalLM, DataCollatorForLanguageModeling, AutoTokenizer, LlamaForCausalLM
from datasets import load_dataset
from functools import partial
import copy
import wandb

def tokenize(dp, tokenizer):
    inputs = tokenizer(
            dp['text'], 
            # return_tensors="pt", 
            max_length=128, 
            truncation=True, 
            padding=False
        )["input_ids"]
    
    inputs=inputs[:128]

    return {'input_ids': inputs, 'labels': copy.deepcopy(inputs)}

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='PyTorch CIFAR10 Training')
    parser.add_argument('--lr', default=2e-5, type=float, help='learning rate')
    parser.add_argument('--batch-size', default=64, type=int, help='batch size')
    parser.add_argument('--model-ckpt', default=None, type=str, help='model checkpoint')
    parser.add_argument('--save', default=None, type=str, help='model checkpoint save dir')
    parser.add_argument('--epoch', default=1, type=int, help='number of epochs')
    parser.add_argument('--save_interval', default=5, type=int, help='model checkpoint saving interval')
    parser.add_argument('--pseudo_random', type=int, default=1234, help='pseudo random seed for all')
    args = parser.parse_args()

    if args.pseudo_random is not None:
        os.environ['PYTHONHASHSEED'] = '0'
        os.environ['TF_DETERMINISTIC_OPS'] = '1'
        random.seed(args.pseudo_random + 1)
        np.random.seed(args.pseudo_random + 1)
        torch.manual_seed(args.pseudo_random)
        torch.cuda.manual_seed(args.pseudo_random)
        torch.cuda.manual_seed_all(args.pseudo_random)
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False
        print(f'set seed to {args.pseudo_random}')

    wandb.init(
        project='InfoScore',
        name='finetune-smol',
        config=args
    )

    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    best_acc = 0  # best test accuracy
    batch_size = args.batch_size
    # Data
    print('==> Preparing data..')
    raw_texts = load_dataset("tatsu-lab/alpaca", split='train')
    ds = raw_texts.map(lambda x: {'text': x['instruction']+x['input']+x['output']})

    model=AutoModelForCausalLM.from_pretrained('HuggingFaceTB/SmolLM-135M', attn_implementation="flash_attention_2", torch_dtype=torch.bfloat16).to(device)
    tokenizer=AutoTokenizer.from_pretrained('HuggingFaceTB/SmolLM-135M')
    tokenizer.pad_token = tokenizer.eos_token
    ds = ds.map(lambda x: tokenize(x, tokenizer)).remove_columns('instruction').remove_columns('input').remove_columns('output').remove_columns('text').remove_columns('labels')
    ds=ds.map(lambda x, idx: {'index': idx}, with_indices=True)
    print(ds[0])
    train_data = torch.utils.data.Subset(ds, list(range(40000)))
    test_data = torch.utils.data.Subset(ds, list(range(40000, 52002)))

    # texts = torch.utils.data.Subset(raw_texts, list(range(40000, 52002)))
    train_loader = DataLoader(
        train_data,
        shuffle=False,
        collate_fn=DataCollatorForLanguageModeling(tokenizer, mlm=False, pad_to_multiple_of=8, return_tensors="pt"),
        num_workers=8,
        batch_size=batch_size)
    
    test_loader = DataLoader(
        test_data,
        shuffle=False,
        collate_fn=DataCollatorForLanguageModeling(tokenizer, mlm=False, pad_to_multiple_of=8, return_tensors="pt"),
        num_workers=8,
        batch_size=batch_size)

    optimizer = optim.SGD(model.parameters(), lr=args.lr,
                        momentum=0.9, weight_decay=5e-4)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=200)

    # Training

    def train(epoch):
        print('\nEpoch: %d' % epoch)
        st_time = time.time()
        model.train()
        train_loss = 0
        # print(next(iter(trainloader)))
        for batch_idx, batch in enumerate(tqdm(train_loader)):
            optimizer.zero_grad()

            batch = batch.to(device)
            input_ids=batch['input_ids']
            labels=batch['labels']
            attn_mask=batch['attention_mask']

            res_model = model(input_ids, labels=labels, attention_mask=attn_mask)
            loss = res_model.loss
            loss.backward()
            optimizer.step()
            train_loss += loss.item()

        duration=time.time()-st_time
        print('Epoch: %d | Train Loss: %.3f | Time: %ds' % (epoch, train_loss/(batch_idx+1), duration), flush=True)
        model.push_to_hub('pxyyy/SmolLM-135M-epoch1', use_temp_dir=True)
        tokenizer.push_to_hub('pxyyy/SmolLM-135M-epoch1', use_temp_dir=True)

        return train_loss/(batch_idx+1)
    
    def test(epoch):
        model.eval()
        test_loss = 0
        with torch.no_grad():
            for batch_idx, batch in enumerate(tqdm(test_loader)):
                batch = batch.to(device)
                input_ids=batch['input_ids']
                labels=batch['labels']
                attn_mask=batch['attention_mask']
                outputs = model(input_ids, labels=labels, attention_mask=attn_mask)
                loss = outputs.loss
                test_loss += loss.item()

        print('Epoch: %d | Test Loss: %.3f ' % (epoch, test_loss/(batch_idx+1)), flush=True)
        
        # Save checkpoint.
        if epoch % args.save_interval == 0 and args.save is not None:
            print('Saving..')
            if not os.path.isdir(args.save):
                os.mkdir(args.save)
            torch.save(model.state_dict(), f'{args.save}/ckpt-{epoch}.pth')
        return test_loss/(batch_idx+1)

    for epoch in range(1, args.epoch+1):
        train_loss = train(epoch)
        test_loss = test(epoch)
        scheduler.step()
        wandb.log({'train/train_loss': train_loss, 'eval/test_loss': test_loss})