# -*- coding: utf-8 -*-
# +
# Importing stock libraries
import numpy as np
import pandas as pd
import torch
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader, RandomSampler, SequentialSampler
import json
from typing import List

# Importing the T5 modules from huggingface/transformers
from transformers import T5ForConditionalGeneration, T5Tokenizer, Adafactor

# Import os for env varibles via Beaker
import os

# WandB – Import the wandb library
import wandb
import logging

from torch import cuda

from split.utils import write_items

from optparse import OptionParser

device = 'cuda' if cuda.is_available() else 'cpu'

logger = logging.getLogger("t5-comet")
logging.basicConfig(level=logging.DEBUG)

# logger.info for allenai beaker verification
logger.info(device)
logger.info(torch.cuda.device_count())

from mosaic.infra.modeling import train, validate, beam_generations
from mosaic.datasets.KGDataset import KGDataset

DEBUG = False
NUM_INST = 100

def read_jsonl_lines(input_file: str) -> List[dict]:
    with open(input_file) as f:
        lines = f.readlines()
        return [json.loads(l.strip()) for l in lines]


# +
wandb.init(project="t5_comet_atomic")
config = wandb.config
config.TRAIN_BATCH_SIZE = int(os.environ.get("TRAIN_BATCH_SIZE", 2))
config.VALID_BATCH_SIZE = int(os.environ.get("VALID_BATCH_SIZE", 2))
config.TRAIN_EPOCHS = int(os.environ.get("TRAIN_EPOCHS", 1))
config.VAL_EPOCHS = int(os.environ.get("VAL_EPOCHS", 1))
config.LEARNING_RATE = float(os.environ.get("LEARNING_RATE", "1e-5"))
config.SEED = int(os.environ.get("SEED", 42))
config.IN_LEN = int(os.environ.get("IN_LEN", 16))
config.OUT_LEN = int(os.environ.get("OUT_LEN", 100))
config.SUMMARY_LEN = 100 # Used for t5
config.OUT_DIR = os.environ.get("OUT_DIR", "/home/pouramini/atomic_models/myt5")
config.DO_TRAIN = os.environ.get("DO_TRAIN", "False") == "True"
config.DO_PRED = os.environ.get("DO_PRED", "True") == "True"
config.PRED_FILE = str(os.environ.get("PRED_FILE", "/tmp/atomic/atomic_pred.tsv"))
config.TOP_K = int(os.environ.get("TOP_K", 40))
config.PRED_BATCH = 64
config.TOKENIZER = os.environ.get('TOKENIZER', "t5")

torch.manual_seed(config.SEED)  # pytorch random seed
np.random.seed(config.SEED)  # numpy random seed
torch.backends.cudnn.deterministic = True

model_name = "t5-small" if 'T5_MODEL' not in os.environ else os.environ['T5_MODEL']

try:
    tokenizer = T5Tokenizer.from_pretrained(model_name)
except:
    tokenizer = T5Tokenizer.from_pretrained(config.TOKENIZER)

tokenizer.add_special_tokens({
'eos_token': '[EOS]',
'additional_special_tokens': [
    'LocationOfAction',
    'HinderedBy',
    'HasFirstSubevent',
    'NotHasProperty',
    'NotHasA',
    'HasA',
    'AtLocation',
    'NotCapableOf',
    'CausesDesire',
    'HasPainCharacter',
    'NotDesires',
    'MadeUpOf',
    'InstanceOf',
    'SymbolOf',
    'xReason',
    'isAfter',
    'HasPrerequisite',
    'UsedFor',
    'MadeOf',
    'MotivatedByGoal',
    'Causes',
    'oEffect',
    'CreatedBy',
    'ReceivesAction',
    'NotMadeOf',
    'xWant',
    'PartOf',
    'DesireOf',
    'HasPainIntensity',
    'xAttr',
    'DefinedAs',
    'oReact',
    'xIntent',
    'HasSubevent',
    'oWant',
    'HasProperty',
    'IsA',
    'HasSubEvent',
    'LocatedNear',
    'Desires',
    'isFilledBy',
    'isBefore',
    'InheritsFrom',
    'xNeed',
    'xEffect',
    'xReact',
    'HasLastSubevent',
    'RelatedTo',
    'CapableOf',
    'NotIsA',
    'ObjectUse',
    '[GEN]'
]
})
tokenizer.add_special_tokens({'pad_token': '[PAD]'})
# -

DEBUT = True
train_dataset = pd.read_csv(
os.environ.get('TRAIN_DATA_PATH', "/tmp/atomic/atomic_train.tsv"),
encoding='latin-1', sep="\t",header=None,names=['head_event','relation','tail_event'])
if DEBUG:
    train_dataset = train_dataset.head(NUM_INST)
    

#train_dataset

#train_dataset = train_dataset[['head_event','relation','tail_event']]
train_dataset.head_event = train_dataset.head_event + ' ' + train_dataset.relation + ' [GEN]'
train_dataset.tail_event = train_dataset.tail_event + ' [EOS]'
logger.info(train_dataset.head())
logger.info(train_dataset.tail_event)

# +
val_dataset = pd.read_csv(os.environ.get('DEV_DATA_PATH', "/tmp/atomic/atomic_dev.tsv"), encoding='latin-1', sep="\t",header=None,names=['head_event','relation','tail_event'])
if DEBUG:
    val_dataset = val_dataset.head(NUM_INST)
val_dataset = val_dataset[['head_event','relation','tail_event']]
val_dataset.head_event = val_dataset.head_event + ' ' + val_dataset.relation + ' [GEN]'
val_dataset.tail_event = val_dataset.tail_event + ' [EOS]'
logger.info(val_dataset.tail_event)
logger.info(val_dataset.head())

test_dataset = pd.read_csv(os.environ.get('TEST_DATA_PATH', "/tmp/atomic/atomic_test.tsv"), encoding='latin-1', sep="\t",header=None,names=['head_event','relation','tail_event'])
if DEBUG:
    test_dataset = test_dataset.head(NUM_INST)
test_dataset = test_dataset[['head_event','relation','tail_event']]
test_dataset.head_event = test_dataset.head_event + ' ' + test_dataset.relation \
                      + ' [GEN]'
test_dataset.tail_event = test_dataset.tail_event + ' [EOS]'
logger.info(test_dataset.tail_event)
logger.info(test_dataset.head())

val_dataset_mini = pd.read_csv(os.environ.get('DEV_DATA_PATH', "/tmp/atomic/atomic_dev.tsv"), encoding='latin-1', sep="\t",header=None,names=['head_event','relation','tail_event'])
if DEBUG:
    val_dataset_mini = val_dataset_mini.head(5)
val_dataset_mini = val_dataset_mini.sample(n=min(int(val_dataset_mini.size / 3), 100),
                                       random_state=config.SEED)
val_dataset_mini = val_dataset_mini[['head_event','relation','tail_event']]
val_dataset_mini.head_event = val_dataset_mini.head_event + ' ' + val_dataset_mini.relation + ' [GEN]'
val_dataset_mini.tail_event = val_dataset_mini.tail_event + ' [EOS]'
logger.info(val_dataset_mini.tail_event)
logger.info(val_dataset_mini.head())

logger.info("TRAIN Dataset tuple count: {}".format(train_dataset.shape))
logger.info("DEV Dataset tuple_count: {}".format(val_dataset.shape))
logger.info("DEV MINI Dataset tuple_count: {}".format(val_dataset_mini.shape))

# +
training_set = KGDataset(train_dataset, tokenizer, config.OUT_LEN, 
                         config.SUMMARY_LEN, model="t5")
val_set = KGDataset(val_dataset, tokenizer, config.IN_LEN, config.OUT_LEN - config.IN_LEN, model="t5", is_eval=True)
val_set_mini = KGDataset(val_dataset.head(2000), tokenizer, config.IN_LEN,  config.OUT_LEN - config.IN_LEN, model="t5", is_eval=True)
test_set = KGDataset(test_dataset, tokenizer, config.IN_LEN,  config.OUT_LEN - config.IN_LEN, model="t5", is_eval=True)

train_params = {
'batch_size': config.TRAIN_BATCH_SIZE,
'shuffle': True,
'num_workers': 0
}

val_params = {
'batch_size': 1,
'shuffle': False,
'num_workers': 0
}

# -

training_set.summ_len

# +
training_loader = DataLoader(training_set, **train_params, drop_last=True)
val_loader = DataLoader(val_set, **val_params, drop_last=True)
test_loader = DataLoader(test_set, **val_params, drop_last=True)
val_loader_mini = DataLoader(val_set_mini, **val_params, drop_last=True)

logging.info("Loading model from {}".format(model_name))
model = T5ForConditionalGeneration.from_pretrained(model_name, return_dict=True)
logging.info("Move model to device {}".format(device))
model = model.to(device)
model.resize_token_embeddings(len(tokenizer))

optimizer = torch.optim.Adam(params=model.parameters(), lr=config.LEARNING_RATE)
#optimizer = Adafactor(model.parameters(),lr=1e-3,
#                      eps=(1e-30, 1e-3),
#                      clip_threshold=1.0,
#                      decay_rate=-0.8,
#                      beta1=None,
#                      weight_decay=0.0,
#                      relative_step=False,
#                      scale_parameter=False,
#                      warmup_init=False)
wandb.watch(model, log="all")
# -



if True: #config.DO_TRAIN:
    logger.info('Initiating Fine-Tuning for the model on our dataset')

    for epoch in range(config.TRAIN_EPOCHS):
        train(epoch, tokenizer, model, device, training_loader, optimizer, val_loader_mini, model_class="t5",save_dir="/home/pouramini/atomic_models/myt5")
        model.save_pretrained('{}/checkpoint_{}'.format(config.OUT_DIR, epoch))
        tokenizer.save_pretrained('{}/checkpoint_{}'.format(config.OUT_DIR, epoch))
    model.save_pretrained(config.OUT_DIR)

# +
if config.DO_PRED:

    if config.PRED_FILE.endswith("jsonl"):
        records = read_jsonl_lines(config.PRED_FILE)
        pred_dataset = pd.DataFrame.from_records(records)
        pred_dataset = pred_dataset.rename(columns={"head": "head_event", "tails": "tail_event"})
        pred_dataset = pred_dataset.explode('tail_event')
    else:
        pred_dataset = pd.read_csv(config.PRED_FILE, encoding='latin-1', sep="\t")

    if DEBUG:
        pred_dataset = pred_dataset.head(NUM_INST)

    pred_dataset = pred_dataset.drop_duplicates(['head_event', 'relation'], ignore_index=True)

    pred_dataset.head_event = pred_dataset.head_event + ' ' + pred_dataset.relation + ' [GEN]'
    pred_dataset.tail_event = pred_dataset.tail_event + ' [EOS]'
    logger.info(pred_dataset.tail_event)
    logger.info(pred_dataset.head())

    pred_set = KGDataset(pred_dataset, tokenizer, config.IN_LEN, config.OUT_LEN - config.IN_LEN, model="t5", is_eval=True)
    pred_loader = DataLoader(pred_set, **val_params, drop_last=False)

    pred_generations = beam_generations(tokenizer, model, device, pred_loader, top_k=config.TOP_K)
    write_items(os.path.join(config.OUT_DIR, "pred_generations.jsonl"),
                [json.dumps(r) for r in pred_generations])

# Resave the model to keep generations and model associated
model.save_pretrained(config.OUT_DIR)
tokenizer.save_pretrained(config.OUT_DIR)
