import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.optim import Adam
from torch.utils.data import DataLoader, random_split

import pytorch_lightning as pl
from pytorch_lightning.core.lightning import LightningModule
from pytorch_lightning import Trainer, seed_everything
from pytorch_lightning.loggers import WandbLogger
from pytorch_lightning.callbacks import EarlyStopping

import pandas as pd
import numpy as np
import transformers
from transformers import AutoModelForSequenceClassification, AutoTokenizer
from FastAutoAugment.read_data import *
from FastAutoAugment.classification_models.TestClassifier import *
from FastAutoAugment.classification_models.BertBasedClassifier import *
from FastAutoAugment.classification_models.RandomDistilBertClassifier import *

import pickle
import wandb

def evaluate_dataset_on_model(wandb_name, checkpoint, dataset, model_type='BertBasedClassifier'):
    # seed_everything(42)
    # print('Inside evaluate ')
    wandb_logger = WandbLogger(name=wandb_name, project='autoaugment')

    if model_type == 'BertBasedClassifier':
        model = BertBasedClassifier.load_from_checkpoint(checkpoint)
    elif model_type == 'RandomDistilBertClassifier':
        model = RandomDistilBertClassifier.load_from_checkpoint(checkpoint)
    
    dataloader = DataLoader(dataset, batch_size=64, num_workers=3)

    trainer = pl.Trainer(deterministic=True, 
        logger=wandb_logger, 
        distributed_backend='dp',
        gpus=[2,3])
    result_dict = trainer.test(model = model, test_dataloaders = dataloader)
    return result_dict

if __name__ == "__main__":
    pass
    # model_name = 'distilbert-base-uncased'
    # full_val_data = pickle.load(open('data/yahoo_answers_v1/yahoo_answers_full_val.pkl', 'rb'))
    # val_dataset = create_dataset(
    #     full_val_data['X'], full_val_data['y'], model_name, 256)
    
    # print('Final accuracy ', evaluate_dataset_on_model(
    #     wandb_name = 'test_full_yahoo_answers_model_no_aug', 
    #     checkpoint = 'checkpoints/full_yahoo_answers_classifier_baseline/lightning_logs/version_7/checkpoints/epoch=6.ckpt',
    #     dataset = val_dataset
    #     )['val_accuracy']
    # )