#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Listening experiment — Day 2: Optimized stories.

Subjects listen to 32 short narratives, shuffled within 3 groups
(1-12, 13-24, 25-32) for counterbalancing. Audio is pre-loaded
and normalized to 24 kHz before playback.
"""

import os
import wave
import numpy as np
from pydub import AudioSegment
from psychopy import logging

from exp_common import (
    HARDWARE_DEFAULTS,
    show_dialog,
    open_serial_port,
    setup_window,
    setup_logging,
    instruction_show,
    cross_show,
    sound_play,
    preload_sounds,
)


def convert_audio_sample_rates(directory, target_rate=24000):
    """Convert all WAV files in directory to target_rate Hz (in-place)."""
    for filename in os.listdir(directory):
        if filename.endswith('.wav'):
            path = os.path.join(directory, filename)
            try:
                with wave.open(path, 'rb') as wf:
                    current_rate = wf.getframerate()
                if current_rate != target_rate:
                    print(f'Converting {filename} from {current_rate} Hz to {target_rate} Hz')
                    audio = AudioSegment.from_wav(path)
                    audio = audio.set_frame_rate(target_rate)
                    audio.export(path, format='wav')
                else:
                    print(f'{filename} is already at {target_rate} Hz')
            except Exception as e:
                print(f'Error processing {filename}: {e}')


# ── Hardware config (override per site) ──────────────────────────────
config = {**HARDWARE_DEFAULTS, 'serial_port': 'COM3'}

# ── Data paths ───────────────────────────────────────────────────────
data_path = '../data_experiments'
audio_dir = os.path.join(data_path, 'audio', 'optimized_stories_pat_6684')
log_dir = '../logs'

# ── Normalize audio sample rates ────────────────────────────────────
convert_audio_sample_rates(audio_dir, target_rate=24000)

# ── Dialog box ───────────────────────────────────────────────────────
exp_info = show_dialog(
    exp_name='Listening - Optimized Stories',
    exp_info={
        'Enter subject ID': 'subject_1',
        'ITI duration (s)': '2',
        'Baseline duration (s)': '30',
    },
    field_order=['Enter subject ID', 'ITI duration (s)', 'Baseline duration (s)'],
)

# ── Serial port ──────────────────────────────────────────────────────
mmbts, serial_port = open_serial_port(config['serial_port'])

# ── Window ───────────────────────────────────────────────────────────
win = setup_window(config)
frame_rate = config['frame_rate']

# ── Logging ──────────────────────────────────────────────────────────
subject_id = str(exp_info['Enter subject ID'])
time_string, log_file = setup_logging(subject_id, log_dir)

# ── Parameters (convert to frames) ──────────────────────────────────
tITI_frames = int(float(exp_info['ITI duration (s)']) * frame_rate)
tbaseline_frames = int(float(exp_info['Baseline duration (s)']) * frame_rate)

logging.data('Starting the experiment. Time and date: ' + time_string)
logging.data('Subject ID: ' + subject_id)
logging.data('ITI duration: ' + exp_info['ITI duration (s)'])
logging.data('Baseline duration: ' + exp_info['Baseline duration (s)'])

# ── Randomize story order (fixed seed for reproducibility) ───────────
np.random.seed(1)
filenames_1 = [str(i) for i in range(1, 13)]
filenames_2 = [str(i) for i in range(13, 25)]
filenames_3 = [str(i) for i in range(25, 33)]
np.random.shuffle(filenames_1)
np.random.shuffle(filenames_2)
np.random.shuffle(filenames_3)
filenames = np.concatenate([filenames_1, filenames_2, filenames_3])
print(filenames)
logging.data('file order: ' + '\n'.join(filenames))

# ── Pre-load audio ──────────────────────────────────────────────────
sound_stims = preload_sounds(filenames, audio_dir, '6684_story_{name}.wav')

# ── Experiment ───────────────────────────────────────────────────────
DESCRIPTION_TEXT = (
    "Bonjour !\n\nBienvenue dans cette experience. \n\n"
    "Vous allez écouter des petites histoires."
    "\n\nPendant l'experience, bougez le moins possible et fixez la croix."
    "\n\nPressez 'Espace' pour continuer"
)

instruction_show(win, 'description', description_text=DESCRIPTION_TEXT)
instruction_show(win, 'baseline', int(2 * frame_rate))
cross_show(win, tbaseline_frames, config=config)
instruction_show(win, 'start')

for isent, filename in enumerate(filenames):
    if serial_port:
        mmbts.write(bytes(bytearray([isent % 254 + 1])))  # trigger codes 1-255
    sound_play(win, sound_stims[filename], filename, frame_rate, config=config)
    cross_show(win, tITI_frames, photodiode=0, config=config)
    if (isent + 1) % 2 == 0:
        instruction_show(win, 'rest')
        instruction_show(win, 'baseline', int(2 * frame_rate))
        cross_show(win, tITI_frames, config=config)

instruction_show(win, 'baseline', int(2 * frame_rate))
cross_show(win, tbaseline_frames, config=config)
instruction_show(win, 'end', int(2 * frame_rate))

# ── Cleanup ──────────────────────────────────────────────────────────
if serial_port:
    mmbts.close()

win.close()
from psychopy import core
core.quit()
