# -- Code Cell --
import pandas as pd
train = pd.read_csv('train_data.csv')
test = pd.read_csv('test_data.csv')

# -- Code Cell --
train.head()

# -- Code Cell --
from scipy.io import wavfile
import scipy.signal as signal
import matplotlib.pyplot as plt
import os
import numpy as np

# -- Code Cell --
def get_spectrogram(wav_path):
    sr, data = wavfile.read(wav_path)                                       # Reading the data from the given path
    _, _, Sxx = signal.spectrogram(data, sr, nperseg = 1024, noverlap = 500)# Computing the standard linear spectrogram
    Sxx = 10 * np.log10(Sxx + 1e-10)                                        # Transforming the spectrogram in logarithmic scale
    return Sxx

# -- Code Cell --
def get_dist(v1,v2):
    v1,v2=np.array(v1),np.array(v2)
    return 1-(v1@v2)/(np.linalg.norm(v1)*np.linalg.norm(v2))

# -- Code Cell --
get_spectrogram("./train/2.wav").shape

# -- Code Cell --
def pad_segment(segment,target):
    segment = np.array(segment)
    if len(segment)>target:
        segment = segment[:target]
    else:
        pad = np.zeros((target-len(segment),segment.shape[1])) #acelasi numar de frecvente segment.shape[1]
        segment = np.vstack([segment,pad])
    return segment

# -- Code Cell --
def preproc_data(wav, target=50):
    spec = get_spectrogram(wav)
    data = spec.T
    noise_base = np.mean(data[:10],axis=0)
    dists = []
    for i in range(len(data)):
        dists.append(get_dist(data[i], noise_base))
    dists = np.array(dists)
    thresh = dists.min()+0.1 #deci daca e mai mare ca thres atunci NU e noise <=> distanta mai mare de noise
    good_data = []
    temp = []
    for i in range(len(dists)):
        if dists[i]>thresh:
            temp.append(data[i])
        else:
            if temp and len(temp)>10:
                padd =pad_segment(temp, target)
                good_data.append(padd)
            temp = []
    if temp and len(temp)>10:
        padd=pad_segment(temp,target)
        good_data.append(padd)
    return np.array(good_data)

# -- Code Cell --
train['answer']=train['answer'].str.replace("|","")
all_labels = []
all_segments =[]
for i in range(len(train)):
    wav_path = f"./train/{train['datapointID'][i]}.wav"
    segments = preproc_data(wav_path)
    answer = str(train['answer'][i])
    digits = [int(d) for d in answer if d.isdigit()]
    for seg,d in zip(segments,digits):
        all_labels.append(d)
        all_segments.append(seg)

# -- Code Cell --
all_segments = np.array(all_segments) 
all_labels = np.array(all_labels) 

# -- Code Cell --
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset

# -- Code Cell --
class SegmentDataset(Dataset):
    def __init__(self, segments, labels):
        self.segments = torch.tensor(segments, dtype=torch.float32)
        self.labels = torch.tensor(labels, dtype=torch.long)
 
    def __len__(self):
        return len(self.labels)
 
    def __getitem__(self, idx):
        return self.segments[idx].unsqueeze(0), self.labels[idx]

# -- Code Cell --
dataset = SegmentDataset(all_segments, all_labels)
loader = DataLoader(dataset, batch_size=64, shuffle=True)

# -- Code Cell --
class DigitCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(1, 32, 3, padding=1),   
            nn.BatchNorm2d(32),                
            nn.ReLU(),                         
            nn.MaxPool2d(2),                   
            nn.Conv2d(32, 64, 3, padding=1),   
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.MaxPool2d(2),                   
            nn.Conv2d(64, 128, 3, padding=1),  
            nn.BatchNorm2d(128),
            nn.ReLU(),
            nn.AdaptiveAvgPool2d((4, 4)),     
           
        )
        self.classifier = nn.Sequential(
            nn.Flatten(),                      
            nn.Linear(128 * 4 * 4, 256),       
            nn.ReLU(),
            nn.Dropout(0.3),                   
            nn.Linear(256, 10),               
        )
    def forward(self, x):
        x = self.features(x)      
        x = self.classifier(x)    
        return x                  

# -- Code Cell --
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = DigitCNN().to(device)          
criterion = nn.CrossEntropyLoss()      
optimizer = optim.Adam(model.parameters(), lr=1e-3)  
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=15, gamma=0.5)

# -- Code Cell --
for epoch in range(40):
    model.train()       
    total_loss = 0
    correct = 0
    total = 0
 
    for x, y in loader:
        x, y = x.to(device), y.to(device)  
 
        out = model(x)                  
        loss = criterion(out, y)        
 
        optimizer.zero_grad()           
        loss.backward()                
        optimizer.step()                
 
        total_loss += loss.item()
        correct += (out.argmax(1) == y).sum().item() 
        total += y.size(0)                             
 
    scheduler.step()  
    acc = correct / total * 100
    print(f"Epoch {epoch:2d} | Loss: {total_loss/len(loader):.4f} | Acc: {acc:.1f}%")

# -- Code Cell --
