# structure.py
import torch
import torch.nn as nn
import lightning as L
from torch.optim.lr_scheduler import ReduceLROnPlateau

class SimpleLSTM(L.LightningModule):
    def __init__(self, input_size=128, hidden_size=64, learning_rate=0.01):
        super().__init__()
        # initialise lr, lstm unit, fc layer, loss func
        self.learning_rate = learning_rate
        self.lstm = nn.LSTM(input_size, hidden_size, batch_first=True)  # (batch, frames, mels)
        self.fc = nn.Linear(hidden_size, 1) # weight matrix (64, 1), bias (1,)
        self.loss_fn = nn.MSELoss()

    def forward(self, x):
        # pass input x (batch, frames, mels) through lstm
        lstm_out, (h_n, c_n) = self.lstm(x)   # h_n: (num_layers, batch, 64), lstm_out: (batch, frames, 64)
        return self.fc(h_n[-1]) # (batch, 64) > (batch, 1)

    def training_step(self, batch, batch_idx):
        x, y = batch    # batch comes from dataloader
        y = y.float().unsqueeze(1)  # make sure y is float32 & (batch,) > (batch, 1)
        pred = self.forward(x)      # forward pass
        loss = self.loss_fn(pred, y)
        # record loss for lightning to track:
        self.log('train_loss', loss, on_step=True, on_epoch=True, prog_bar=True)
        return loss # return loss for backprop

    def validation_step(self, batch, batch_idx):
        x, y = batch
        y = y.float().unsqueeze(1)
        pred = self.forward(x)
        loss = self.loss_fn(pred, y)
        self.log('val_loss', loss, on_step=True, on_epoch=True, prog_bar=True)
        return loss

    def configure_optimizers(self):
        # pass all trainable parameters and learning rate to Adam
        optimizer = torch.optim.Adam(self.parameters(), lr=self.learning_rate)
        # reduce learning rate by a factor of .5 when val_loss stops improving for 10 epochs
        scheduler = ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=10)
        return {
            'optimizer': optimizer,
            'lr_scheduler': {
                'scheduler': scheduler,
                'monitor': 'val_loss'
            }
        }