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

# %%
train.head()

# %%
train['answer'] = train['answer'].str.replace("|", "")
train['answer']

# %%
from scipy.io import wavfile
import scipy.signal as signal
import matplotlib.pyplot as plt
import numpy as np

# %%
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

def get_waveform(wav_path):
    sr, data = wavfile.read(wav_path)                                       # Reading the data from the given path
    plt.plot(data)                                                          # Plotting the waveforms using matplotlib
    plt.show()
    

# %%
get_waveform("./train/0.wav")

# %%
get_spectrogram("./train/0.wav")

# %%
len(list("967153686617"))

# %%
from scipy.ndimage import gaussian_filter1d         
import torch                          
import torch.nn as nn                 
import torch.optim as optim           
from torch.utils.data import Dataset, DataLoader 

# %%
import os

# %%
def get_dist(vect1, vect2):
    v1, v2 = np.array(vect1), np.array(vect2)
    return 1-(v1 @ v2)/(np.linalg.norm(v1)*np.linalg.norm(v2) + 1e-10)

# %%
def pad_segment(segment, target_len):
    segment = np.array(segment)
    if len(segment) < target_len:
        pad =np.zeros((target_len-len(segment), segment.shape[1]))
        segment = np.vstack([segment,pad])
    else:
        segment = segment[:target_len]
    return segment

# %%
def preproc_audio(wav_path, target_len=50):
    sr, data_init = wavfile.read(wav_path)
    f, t, Sxx = signal.spectrogram(data_init, sr, nperseg=1024, noverlap=500)
    Sxx = 10 * np.log10(Sxx + 1e-10)

    data = Sxx.T
    noise_base = np.mean(data[0:10], axis=0)

    dists = []
    for i in range(len(data)):
        dists.append(get_dist(data[i], noise_base))

    dists = np.array(dists)
    dists_smooth = gaussian_filter1d(dists.astype(float), sigma=4)
    thresh = dists_smooth.min() + 0.1

    curated_data = []
    temp = []

    for i in range(len(dists)):
        if dists_smooth[i] > thresh:
            temp.append(Sxx.T[i])
        else:
            if temp and len(temp) > 10:
                padd = pad_segment(temp, target_len)
                curated_data.append(padd)
            temp = []

    if temp and len(temp) > 10:
        padd = pad_segment(temp, target_len)
        curated_data.append(padd)

    return np.array(curated_data)

# %%
train_df = pd.read_csv('train_data.csv')
train_df['answer'] = train_df['answer'].str.replace("|", "")
all_segments = []
all_labels = []
skipped = 0
for _,row in train_df.iterrows():
    wav_path = f"./train/{row['datapointID']}.wav"
    segments = preproc_audio(wav_path)
    answer = str(row['answer'])
    digits = [int(c) for c in answer if c.isdigit()]
    if len(segments) != len(digits):
        skipped += 1
        continue
    for seg, d in zip(segments, digits):
        all_segments.append(seg)
        all_labels.append(d)
        

# %%
all_segments = np.array(all_segments) 
all_labels = np.array(all_labels) 

# %%
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]

# %%
dataset = SegmentDataset(all_segments, all_labels)
loader = DataLoader(dataset, batch_size=64, shuffle=True)

# %%
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                  

# %%
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)

# %%
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}%")

# %%
test_df = pd.read_csv('test_data.csv')

# %%
model.eval()  
predictions = []
 
for _, row in test_df.iterrows():
    wav_path = f"./test/{row['datapointID']}.wav"
    segments = preproc_audio(wav_path)
 
    if len(segments) == 0:
        predictions.append("")
        continue
 
    
    segs_tensor = torch.tensor(segments, dtype=torch.float32).unsqueeze(1).to(device)
 
    with torch.no_grad():                    
        logits = model(segs_tensor)          
        preds = logits.argmax(dim=1)         
 
    digit_str = "".join([str(p.item()) for p in preds])
    predictions.append(digit_str)

# %%
with open("submission.csv", "w") as f:
    f.write("subtaskID,datapointID,answer\n")
    for i, pred in enumerate(predictions):
        f.write(f"1,{i},|{pred}|\n")


