import torch
import numpy as np

from src.config import SR, CODE2IDX


def get_idx(duration, n_secs=5, sr=SR, random_chunk=True):
    num_frames = np.ceil(sr * duration)
    chunk_idx = (n_secs*sr)
    DEFAULT_OFFSET = 10
    start = np.random.randint(DEFAULT_OFFSET, num_frames-chunk_idx) if random_chunk else DEFAULT_OFFSET
    return start, start+chunk_idx

def to_square(arr):
    """Convert (almost square) array to a square array by padding/truncating."""
    rows, cols = arr.shape
    
    if cols < rows:
        pad_width = ((0, 0), (0, rows - cols))
        return np.pad(arr, pad_width, mode='constant')
    else:
        return arr[:, :rows]

def to_tensor(data):
    return [torch.FloatTensor(x) for x in data]

def one_hot(idx):
    y = torch.zeros(len(CODE2IDX.keys()))
    y[idx] = 1.
    return y