import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader

import numpy as np

import random
import math
import time
import sys, os
sys.path.append(os.path.join(sys.path[0],'baseline'))

from datasets.base_dataset import BaseDataset

from baseline.baseline_attention import Attention
from baseline.baseline_decoder import Decoder
from baseline.baseline_encoder import Encoder
from baseline.baseline_model import Seq2Seq
from datasets.vocab import VocabEntry

SEED = 1000

random.seed(SEED)
np.random.seed(SEED)
torch.manual_seed(SEED)
torch.cuda.manual_seed(SEED)
torch.backends.cudnn.deterministic = True

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

BATCH_SIZE = 4

dataset = BaseDataset(['kushman.json'])

eqFlattened = [item for sublist in dataset.equations for item in sublist]
quFlattened = [item for sublist in dataset.questions for item in sublist]
eqFlattened.append('<unk>')
quFlattened.append('<unk>')

INPUT_DIM = len(set(quFlattened))
OUTPUT_DIM = len(set(eqFlattened))

src_vocab = VocabEntry(dict([(v, i) for i, v in enumerate(set(quFlattened))]))
trg_vocab = VocabEntry(dict([(v, i) for i, v in enumerate(set(eqFlattened))]))

train_size = int(len(dataset) * (3./10))
valid_size = int(len(dataset) * (1./10))
test_size = len(dataset) - train_size - valid_size

dataset.equations = src_vocab.to_input_tensor(dataset.equations, device)
dataset.questions = trg_vocab.to_input_tensor(dataset.questions, device)
train_data = dataset
#train_data, valid_data, test_data = torch.utils.data.random_split(dataset, [train_size, valid_size, test_size])
#print(train_data)
#train_data = src_vocab.to_input_tensor(train_data, device)
#valid_data = src_vocab.to_input_tensor(valid_data, device)
#test_data = src_vocab.to_input_tensor(test_data, device)



train_iterator = DataLoader(train_data, batch_size=BATCH_SIZE, collate_fn=collate_fn_padd)
'''
valid_iterator = DataLoader(valid_data, batch_size=BATCH_SIZE)
test_iterator = DataLoader(test_data, batch_size=BATCH_SIZE)
'''

ENC_EMB_DIM = 256
DEC_EMB_DIM = 256
ENC_HID_DIM = 512
DEC_HID_DIM = 512
ENC_DROPOUT = 0.5
DEC_DROPOUT = 0.5

attn = Attention(ENC_HID_DIM, DEC_HID_DIM)
enc = Encoder(INPUT_DIM, ENC_EMB_DIM, ENC_HID_DIM, DEC_HID_DIM, ENC_DROPOUT)
dec = Decoder(OUTPUT_DIM, DEC_EMB_DIM, ENC_HID_DIM, DEC_HID_DIM, DEC_DROPOUT, attn)

model = Seq2Seq(enc, dec, device).to(device)

def init_weights(m):
    for name, param in m.named_parameters():
        if 'weight' in name:
            nn.init.normal_(param.data, mean=0, std=0.01)
        else:
            nn.init.constant_(param.data, 0)
            
model.apply(init_weights)

def count_parameters(model):
    return sum(p.numel() for p in model.parameters() if p.requires_grad)

print(f'The model has {count_parameters(model):,} trainable parameters')

optimizer = optim.Adam(model.parameters())

#TRG_PAD_IDX = TRG.vocab.stoi[TRG.pad_token]

criterion = nn.CrossEntropyLoss() #ignore_index = TRG_PAD_IDX

def train(model, iterator, optimizer, criterion, clip):
    
    model.train()
    
    epoch_loss = 0
    
    for i, batch in enumerate(iterator):

        print(batch)
        src = batch[0] #src_vocab.to_input_tensor(batch[0], device).permute(1, 0)
        trg = batch[1] #trg_vocab.to_input_tensor(batch[1], device).permute(1, 0)
        print("PRINTING SRC AND TRG")
        print(src.shape)
        print(trg.shape)
        optimizer.zero_grad()
        
        output = model(src, trg)
        print("OUTPUT OF MODEL")
        print(output)
        print(output.shape)
        #trg = [trg len, batch size]
        #output = [trg len, batch size, output dim]
        
        output_dim = output.shape[-1]
        
        output = output[1:].view(-1, output_dim)
        trg = trg[1:].view(-1)
        
        #trg = [(trg len - 1) * batch size]
        #output = [(trg len - 1) * batch size, output dim]
        
        loss = criterion(output, trg)
        print("AFTER MODIFICATIONS")
        print(output)
        print(loss)
        
        loss.backward()
        
        torch.nn.utils.clip_grad_norm_(model.parameters(), clip)
        
        optimizer.step()
        
        epoch_loss += loss.item()
        
    return epoch_loss / len(iterator)

def evaluate(model, iterator, criterion):
    
    model.eval()
    
    epoch_loss = 0
    
    with torch.no_grad():
    
        for i, batch in enumerate(iterator):

            src = batch[0]
            trg = batch[1]

            output = model(src, trg, 0) #turn off teacher forcing

            #trg = [trg len, batch size] 
            #output = [trg len, batch size, output dim]

            output_dim = output.shape[-1]
            
            output = output[1:].view(-1, output_dim)
            trg = trg[1:].view(-1)

            #trg = [(trg len - 1) * batch size]
            #output = [(trg len - 1) * batch size, output dim]

            loss = criterion(output, trg)

            epoch_loss += loss.item()
        
    return epoch_loss / len(iterator)

def epoch_time(start_time, end_time):
    elapsed_time = end_time - start_time
    elapsed_mins = int(elapsed_time / 60)
    elapsed_secs = int(elapsed_time - (elapsed_mins * 60))
    return elapsed_mins, elapsed_secs

N_EPOCHS = 10
CLIP = 1

best_valid_loss = float('inf')

for epoch in range(N_EPOCHS):
    print("Epoch: %i" % epoch)
    
    start_time = time.time()
    
    train_loss = train(model, train_iterator, optimizer, criterion, CLIP)
    valid_loss = evaluate(model, valid_iterator, criterion)
    
    end_time = time.time()
    
    epoch_mins, epoch_secs = epoch_time(start_time, end_time)
    
    if valid_loss < best_valid_loss:
        best_valid_loss = valid_loss
        torch.save(model.state_dict(), 'tut3-model.pt')
    
    print(f'Epoch: {epoch+1:02} | Time: {epoch_mins}m {epoch_secs}s')
    print(f'\tTrain Loss: {train_loss:.3f} | Train PPL: {math.exp(train_loss):7.3f}')
    print(f'\t Val. Loss: {valid_loss:.3f} |  Val. PPL: {math.exp(valid_loss):7.3f}')

model.load_state_dict(torch.load('tut3-model.pt'))

test_loss = evaluate(model, test_iterator, criterion)

print(f'| Test Loss: {test_loss:.3f} | Test PPL: {math.exp(test_loss):7.3f} |')
