# -- Code Cell --
import torchaudio
import torchaudio.transforms as T
import os

# -- Code Cell --
import pandas as pd
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from torchvision import models, transforms
from PIL import Image
from torchvision.transforms import functional as TF
import torch.nn.functional as F

# -- Code Cell --
train = pd.read_csv('train.csv')
train.head()

# -- Code Cell --
spectograme = []
for i in range(len(train)):
    audio = os.path.join("./audio",f"{train['sampleID'][i]}.wav")
    waveform, sr = torchaudio.load(audio)
    mel_spec = T.MelSpectrogram(
    sample_rate=sr,
    n_mels=128,
    n_fft=1024,
    hop_length=512
    )(waveform)
    mel_spec_db = T.AmplitudeToDB()(mel_spec)
    spectograme.append(mel_spec_db)
    # print(audio)
spectograme

# -- Code Cell --
import torch
import torch.nn.functional as F

def spec_to_image(spec):
    if spec.dim() == 2:
        spec = spec.unsqueeze(0)  
    spec = F.interpolate(spec.unsqueeze(0), size=(224, 224), mode='bilinear').squeeze(0)
    spec = spec.repeat(3, 1, 1)
    from torchvision import transforms
    normalize = transforms.Normalize(
        mean=[0.485, 0.456, 0.406],
        std=[0.229, 0.224, 0.225]
    )
    spec = normalize(spec)
    return spec

# -- Code Cell --
import pandas as pd

# -- Code Cell --
train = pd.read_csv('train.csv')
train.head()

# -- Code Cell --
class_map = {'cat':0, "pig":1, "hen":2, "water_drops":3,"rain":4,"dog":5,"sheep":6,"cow":7,"chirping_birds":8}

# -- Code Cell --
class AudioDataset(Dataset):
    def __init__(self, path):
        self.df = pd.read_csv(path)
    def __len__(self):
        return len(self.df)
    def __getitem__(self, idx):
        row = self.df.iloc[idx]
        spec = torch.tensor(spectograme[idx], dtype=torch.float32)
        img = spec_to_image(spec)
        label = class_map[row['class']]
        return img, label

# -- Code Cell --
train_ds = AudioDataset("./train.csv")
train_dataloader = DataLoader(train_ds, batch_size=64, shuffle=True, num_workers=0)

# -- Code Cell --
train['class'].unique()

# -- Code Cell --
train['class'].unique()

# -- Code Cell --
device = torch.device('cuda')

# -- Code Cell --
weights = models.ResNet50_Weights.IMAGENET1K_V2
model = models.resnet50(weights=weights)
model.fc = nn.Linear(2048,9)

for param in model.parameters():
    param.requires_grad = False
for param in model.fc.parameters():
    param.requires_grad = True
model = model.to(device)
optimizer= optim.Adam(model.fc.parameters(), lr=1e-3)
criterion = nn.CrossEntropyLoss()

# -- Code Cell --
for epoch in range(40):
    total_loss = 0
    for spectogram, label in train_dataloader:
        spectogram = spectogram.to(device)
        label = label.to(device)
        optimizer.zero_grad()
        output = model(spectogram)
        loss = criterion(output,label)
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
    print(f"Epoch {epoch}: loss={total_loss/len(train_dataloader):.4f}")
    if epoch == 20:
        for param in model.parameters():
            param.requires_grad = True
        optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)

# -- Code Cell --
test = pd.read_csv('test.csv')

# -- Code Cell --
spectograme_test = []
for i in range(len(test)):
    audio = os.path.join("./audio",f"{test['sampleID'][i]}.wav")
    waveform, sr = torchaudio.load(audio)
    mel_spec = T.MelSpectrogram(
    sample_rate=sr,
    n_mels=128,
    n_fft=1024,
    hop_length=512
    )(waveform)
    mel_spec_db = T.AmplitudeToDB()(mel_spec)
    spectograme_test.append(mel_spec_db)
    # print(audio)
spectograme_test

# -- Code Cell --
class AudioDatasetTest(Dataset):
    def __init__(self, path):
        self.df = pd.read_csv(path)
    def __len__(self):
        return len(self.df)
    def __getitem__(self, idx):
        spec = torch.tensor(spectograme_test[idx], dtype=torch.float32)
        img = spec_to_image(spec)
        return img

# -- Code Cell --
test_ds = AudioDatasetTest("./test.csv")
test_loader = DataLoader(test_ds,batch_size=32, num_workers=0, shuffle=False)

# -- Code Cell --
preds= []
for spectogram in test_loader:
    with torch.no_grad():
        spectogram = spectogram.to(device)
        pred = model(spectogram)
        _,pred = torch.max(pred, 1)
        preds.append(pred.cpu())
        preds_final = torch.cat(preds)
preds_final

# -- Code Cell --
preds_final = preds_final.numpy()

# -- Code Cell --
test = pd.read_csv('test.csv')
test.head()

# -- Code Cell --


# -- Code Cell --
class_map = {'cat':0, "pig":1, "hen":2, "water_drops":3,"rain":4,"dog":5,"sheep":6,"cow":7,"chirping_birds":8}

# -- Code Cell --
rev = {}
for label,value in class_map.items():
    rev[value] = label
rev

# -- Code Cell --
preds_final_rev = []
for i in preds_final:
    preds_final_rev.append(rev[i])
preds_final_rev

# -- Code Cell --
subs = pd.DataFrame({
    "sampleID":test['sampleID'],
    "class":preds_final_rev
}).to_csv("subs.csv",index=False)

# -- Code Cell --
