import numpy as np # linear algebra
import pandas as pd
import os

import librosa as lr
import torch
import torch.nn as nn
import pytorch_lightning as pl
import gradio

# HYPERPARAMS
EPOCHS = 200
BATCH_SIZE = 32
NUM_OF_CLASSES = 14

class MFCC_CNN(pl.LightningModule):
    def __init__(self, num_of_classes):
        super(MFCC_CNN, self).__init__()

        self.example_input_array = torch.Tensor(32, 1, 50, 94)
        self.train_loss_output = []
        self.train_acc_output = []
        self.val_acc_output = []
        self.val_loss_output = []

        self.number_of_classes = num_of_classes

        self.conv_1 = nn.Sequential(
            nn.Conv2d(in_channels = 1, 
                      out_channels = 64, 
                      kernel_size =3,
                      padding = 1,
                      stride = 1),
            nn.BatchNorm2d(64),
            nn.LeakyReLU(),
            nn.MaxPool2d(kernel_size=2),
            nn.Dropout(0.1)
        )

        self.conv_2 = nn.Sequential(
            nn.Conv2d(in_channels = 64, 
                      out_channels = 128, 
                      kernel_size = 3,
                      padding = 1,
                      stride = 1),
            nn.BatchNorm2d(128),
            nn.LeakyReLU(),
            nn.MaxPool2d(kernel_size=2),
            nn.Dropout(0.1)
        )

        self.conv_3 = nn.Sequential(
            nn.Conv2d(in_channels = 128, 
                      out_channels = 256, 
                      kernel_size = 3,
                      padding = 1,
                      stride = 1),
            nn.BatchNorm2d(256),
            nn.LeakyReLU(),
            nn.MaxPool2d(kernel_size=2),
            nn.Dropout(0.1)
        )

        self.conv_4 = nn.Sequential(
            nn.Conv2d(in_channels = 256, 
                      out_channels = 512,
                      kernel_size = 3,
                      padding = 1,
                      stride = 1),
            nn.BatchNorm2d(512),
            nn.LeakyReLU(),
            nn.MaxPool2d(kernel_size=2)
        )

        self.conv_5 = nn.Sequential(
            nn.Conv2d(in_channels = 512, 
                      out_channels = 512, 
                      kernel_size = 2,
                      padding = 0,
                      stride = 1),
            nn.BatchNorm2d(512),
            nn.LeakyReLU(),
            nn.MaxPool2d(kernel_size=2)
        )

        self.drop = nn.Dropout(0.1)
        self.lin_1 = nn.Linear(1024, 128)
        self.lin_2 = nn.Linear(128, 64)
        self.lin_3 = nn.Linear(64, num_of_classes)

        self.relu = nn.ReLU()
        self.softmax = nn.Softmax()
           
    def forward(self, x):
        out = self.conv_1(x)
        out = self.conv_2(out)
        out = self.conv_3(out)
        out = self.conv_4(out)
        out = self.conv_5(out)

        out = torch.flatten(out, start_dim=1)

        out = self.drop(self.lin_1(self.relu(out)))
        out = self.drop(self.lin_2(self.relu(out)))
        out = self.drop(self.lin_3(self.relu(out)))
        
        out = self.softmax(out)

        return out

    def loss_fn(self, out, target):
        return nn.CrossEntropyLoss()(input=out.view(-1, self.number_of_classes),
                                     target=target)

    def configure_optimizers(self):
        LR=1e-3
        optimizer = torch.optim.Adam(self.parameters(), lr=LR, weight_decay=1e-3)
        scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer,
                                                               mode='min',
                                                               factor=0.5,
                                                               patience=5,
                                                               verbose=True)

        return {
                'optimizer': optimizer,
                "lr_scheduler": {
                    "scheduler": scheduler,
                    "monitor": "val_loss",
                    'interval': 'epoch',
                    'frequency': 1
                },
            }

    def training_step(self, batch, batch_idx):
        mfcc, lable = batch
        mfcc = mfcc.view(-1, 1, 50, 94)
        lable = lable.view(-1, self.number_of_classes)
        
        out = self(mfcc)
        
        loss = self.loss_fn(out=out, target=lable)

        lable = torch.argmax(lable,dim=1)
        predictions = torch.argmax(out,dim=1)
        accuracy = torch.sum(lable==predictions)/float(len(lable))

        self.train_acc_output.append(accuracy.detach().numpy())
        self.train_loss_output.append(loss.detach().numpy())
        #wandb.log({'train_accuracy_step': accuracy, 'train_loss_step':loss})\

        self.log('train_accuracy', accuracy, prog_bar=True, on_epoch=True, on_step=False)
        self.log('train_loss', loss, prog_bar=True, on_epoch=True, on_step=False)
        return loss
    
    def validation_step(self, batch, batch_idx):
        mfcc, lable = batch
        mfcc = mfcc.view(-1, 1, 50, 94)
        lable = lable.view(-1, self.number_of_classes)
        
        out = self(mfcc)
        
        loss = self.loss_fn(out=out, target=lable)
        
        lable = torch.argmax(lable,dim=1)
        predictions = torch.argmax(out,dim=1)
        accuracy = torch.sum(lable==predictions)/float(len(lable))

        self.val_acc_output.append(accuracy.detach().numpy())
        self.val_loss_output.append(loss.detach().numpy())
        #wandb.log({'val_accuracy_step': accuracy, 'val_loss_step':loss})
        
        self.log('val_accuracy', accuracy, prog_bar=True, on_epoch=True)
        self.log('val_loss', loss, prog_bar=True, on_epoch=True)

        return loss
    
    def on_train_epoch_end(self):
        train_loss_epoch = self.train_loss_output
        train_acc_epoch = self.train_acc_output

        #wandb.log({'train_loss_epoch':np.mean(train_loss_epoch), 
        #       'train_acc_epoch':np.mean(train_acc_epoch)})

        self.train_loss_output.clear()
        self.train_acc_output.clear()

    def on_validation_epoch_end(self):
        val_loss_epoch = self.val_loss_output
        val_acc_epoch = self.val_acc_output

        #wandb.log({'val_loss_epoch':np.mean(val_loss_epoch), 
        #       'val_acc_epoch':np.mean(val_acc_epoch)})

        self.val_acc_output.clear()
        self.val_loss_output.clear()