from model import Wav2VecModel
from dataset import S2IDataset, collate_fn
import requests
requests.packages.urllib3.disable_warnings()
import gradio as gr
import torch
import torch.nn as nn
import torchaudio
import torch.nn.functional as F
import pytorch_lightning as pl

from pytorch_lightning import Trainer
from pytorch_lightning.callbacks import ModelCheckpoint
from pytorch_lightning.loggers import WandbLogger

# SEED
SEED=100
pl.utilities.seed.seed_everything(SEED)
torch.manual_seed(SEED)

import os
os.environ['WANDB_MODE'] = 'online'
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"  
os.environ["CUDA_VISIBLE_DEVICES"]="1"

class LightningModel(pl.LightningModule):
    def __init__(self,):
        super().__init__()
        self.model = Wav2VecModel()

    def forward(self, x):
        return self.model(x)

    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=1e-5)
        return [optimizer]

    def loss_fn(self, prediction, targets):
        return nn.CrossEntropyLoss()(prediction, targets)

    def training_step(self, batch, batch_idx):
        x, y = batch
        y = y.view(-1)

        logits = self(x)        
        probs = F.softmax(logits, dim=1)
        loss = self.loss_fn(logits, y)

        winners = logits.argmax(dim=1)
        corrects = (winners == y)
        acc = corrects.sum().float()/float(logits.size(0))

        self.log('train/loss', loss, on_step=False, on_epoch=True, prog_bar=True)
        self.log('train/acc', acc, on_step=False, on_epoch=True, prog_bar=True)
        torch.cuda.empty_cache()
        return {
            'loss':loss, 
            'acc':acc
            }

    def validation_step(self, batch, batch_idx):
        x, y = batch
        y = y.view(-1)

        logits = self(x)
        loss = self.loss_fn(logits, y)

        winners = logits.argmax(dim=1)
        corrects = (winners == y)
        acc = corrects.sum().float() / float( logits.size(0))

        self.log('val/loss' , loss, on_step=False, on_epoch=True, prog_bar=True)
        self.log('val/acc',acc, on_step=False, on_epoch=True, prog_bar=True)

        return {'val_loss':loss, 
                'val_acc':acc,
                }

    def test_step(self, batch, batch_idx):
        x, y = batch
        y = y.view(-1)

        logits = self(x)
        loss = self.loss_fn(logits, y)

        winners = logits.argmax(dim=1)
        corrects = (winners == y)
        acc = corrects.sum().float() / float( logits.size(0))

        self.log('val/loss' , loss, on_step=False, on_epoch=True, prog_bar=True)
        self.log('val/acc',acc, on_step=False, on_epoch=True, prog_bar=True)

        return {'val_loss':loss, 
                'val_acc':acc,
                }

    def predict(self, wav):
        self.eval()
        with torch.no_grad():
            output = self.forward(wav)
            predicted_class = torch.argmax(output, dim=1)
        return predicted_class


print(torch.cuda.mem_get_info())

model = LightningModel()

run_name = "wav2vec"

checkpoint_path = "./wav2vec-epoch=epoch=4.ckpt.ckpt"
checkpoint = torch.load(checkpoint_path)
model.load_state_dict(checkpoint['state_dict'])
trainer = Trainer(
    gpus=1
)

#trainer.fit(model, train_dataloader=trainloader, val_dataloaders=valloader)
#trainer.test(model,dataloaders=testloader,verbose=True)

#with torch.no_grad():
 #   y_hat = model(wav_tensor)

def trabscribe(audio):
    wav_tensor,_ = audio
    wav_tensor = resmaple(wav_tensor)
    #model = model.to('cuda') 
    y_hat = model.predict(wav_tensor)
    labels = {0:"branch_address : enquiry about bank branch location",
    1:"activate_card : enquiry about activating card products",
    2:"past_transactions : enquiry about past transactions in a specific time period",
    3:"dispatch_status : enquiry about the dispatch status of card products",
    4:"outstanding_balance : enquiry about outstanding balance on card products",
    5:"card_issue : report about an issue with using card products",
    6:"ifsc_code : enquiry about IFSC code of bank branch",
    7:"generate_pin : enquiry about changing or generating a new pin for their card product",
    8:"unauthorised_transaction : report about an unauthorised or fraudulent transaction",
    9:"loan_query : enquiry about different kinds of loans",
    10:"balance_enquiry : enquiry about bank account balance",
    11:"change_limit : enquiry about changing the limit for card products",
    12:"block : enquiry about blocking card or banking product",
    13:"lost : report about losing a card product}
    return labels[y_hat]

print(y_hat)
get_intent = gr.Interface(fn = transcribe,
                       gr.Audio(source="microphone"), outputs="text").launch()