import audformat
import pandas as pd
# segment the data
from inaSpeechSegmenter import Segmenter
from audformat.utils import to_filewise_index
from audformat import segmented_index
import warnings

class Ina_segmenter:

    def __init__(self):
        # initialize the VAD model
        self.seg = Segmenter()
        warnings.simplefilter(action='ignore', category=FutureWarning)

    def get_segmentation(self, file):
        print(f'segmenting {file[0]}')
        segmentation = self.seg(file[0])
        files, starts, ends = [], [], []
        number = 0
        for entry in segmentation:
            kind = entry[0]
            start = entry[1]
            end = entry[2]
            if kind == 'female' or kind == 'male':
                files.append(file[0])
                starts.append(start)
                ends.append(end)
                number += 1
        print(f'found {number} segments')        
        seg_index = segmented_index(files, starts, ends)
        return seg_index

    def segment_dataframe(self, df):
        dfs = []
        for file, values in df.iterrows():
            index = self.get_segmentation(file)
            dfs.append(
                pd.DataFrame(
                    values.to_dict(),
                    index,
                )
            )
        return audformat.utils.concat(dfs)

def main():
    files = pd.Series(['test_wavs/test_for_segment_2.wav'])
    df_sample = pd.DataFrame(index = files)
    df_sample['target'] = 'anger'
    df_sample.index = audformat.utils.to_segmented_index(df_sample.index, allow_nat=False)
    df_sample.head(10)
    segmenter = Ina_segmenter()
    df_seg = segmenter.segment_dataframe(df_sample)
    def calc_dur(x):
        starts = x[1]
        ends = x[2]
        return (ends - starts).total_seconds()
    df_seg['duration'] = df_seg.index.to_series().map(lambda x:calc_dur(x)) 
    print(df_seg.head())

if __name__ == '__main__':
    main()