"""
temporal.py
-----------
Create plots for model's prediction over time for steps 1, 4, 7.
Two plots for each continua, containing two facets for english type.

- Black line: start of vowel
- Red line: end of non padded audio
"""

import torch
import torch.nn as nn
import torchaudio
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from i_dataset import StimuliSounds
from ii_structure import SimpleLSTM

# Configuration (same as training)
# -------------------------------------------
CSV_FILE = "../../dataset.csv"
AUDIO_DIR = "../..sound/"
SAMPLES = 6394

INPUT_SIZE = 128
HIDDEN_SIZE = 64
LEARNING_RATE = 0.01
MODEL_PATH = "model.pt"

mel_transform = nn.Sequential(
    torchaudio.transforms.MelSpectrogram(
        sample_rate=16000,
        n_fft=1024,
        win_length=320,
        hop_length=44,              # 2.75 ms per frame
        window_fn=torch.hamming_window,
        n_mels=128
    ),
    torchaudio.transforms.AmplitudeToDB()
)

# Load model and dataset
# -------------------------------------------
model = SimpleLSTM(input_size=INPUT_SIZE, hidden_size=HIDDEN_SIZE, learning_rate=LEARNING_RATE)
model.load_state_dict(torch.load(MODEL_PATH))

dataset = StimuliSounds(CSV_FILE, AUDIO_DIR, SAMPLES, mel_transform)
csv_data = pd.read_csv(CSV_FILE)

# Parse filename
# -------------------------------------------
def parse_filename(fname):
    parts = fname.split('_')
    continuum = int(parts[0])           # Cruː or Cɔt
    accent = parts[1].split('-')[0]     # asm or amr
    vot = int(parts[2])                 # vot step
    accent_map = {'amr': 'AmE', 'asm': 'AsE'}
    continuum_map = {1: 'Cɔt', 2: 'Cruː'}
    return vot, accent_map[accent], continuum_map[continuum], continuum

# VOT endpoints (ms)    - dictionary
# -------------------------------------------
vot_endpoints_1 = {     # Cɔt
    1: 27.25,   # ms
    4: 49.75,
    7: 72.25
}

vot_endpoints_2 = {     # Cruː
    1: 57.125,
    4: 73.938,
    7: 90.75
}

# Frame duration
frame_duration_ms = 44 / 16000 * 1000   # 2.75 ms

# Sound endpoints before padding (ms)   - nested dictionary
# -------------------------------------------
sound_endpoints = {
    1: {    # Cɔt
        'AmE': {1: 252.375, 4: 274.875, 7: 297.375},
        'AsE':    {1: 231.563, 4: 254.063, 7: 276.562}
    },
    2: {    # Cruː
        'AmE': {1: 366.0, 4: 382.812, 7: 399.625},
        'AsE':    {1: 255.312, 4: 272.125, 7: 288.938}
    }
}

# Extract temporal outputs (steps: 1,4,7)
# -------------------------------------------
target_steps = {1, 4, 7}
temporal_records = []   # each frame: {vot, accent, continuum_label, continuum_num, output}

with torch.no_grad():
    for fname in csv_data.iloc[:, 0].unique():
        idx = csv_data.index[csv_data.iloc[:, 0] == fname].tolist()
        mel, _ = dataset[idx[0]]
        mel = mel.unsqueeze(0)                     # (1, time, 128)

        lstm_out, (_, _) = model.lstm(mel)          # lstm_out (1, time, 64)
        fc_out = model.fc(lstm_out.squeeze(0))      # (time, 1)
        outputs = fc_out.squeeze(-1).tolist()       # list outputs

        vot, accent, cont_label, cont_num = parse_filename(fname)
        if vot not in target_steps:
            continue

        temporal_records.append({
            'vot': vot,
            'accent': accent,
            'continuum_label': cont_label,
            'continuum_number': cont_num,
            'outputs': outputs
        })

# Plotting
# -------------------------------------------
def plot_continuum(cont_label, cont_num, records, vot_end, sound_end):

    fig, axes = plt.subplots(1, 2, figsize=(12, 5), sharey=True)

    for ax, accent in zip(axes, ['AmE', 'AsE']):    # pairs 1st axis with AmE, 2nd with AsE

        # Time axis
        n_frames = 146
        time_axis = [i * frame_duration_ms for i in range(n_frames)]    # 0, 2.75, 5.5, 8.25, ..., 2.75x145

        # frame lines
        for t in time_axis:
            ax.axvline(x=t, color='lightgrey', linestyle='-', linewidth=0.3, alpha=0.8)

        # curves & vot endpoints
        for step in [1, 4, 7]:

            # Get the single curve for this VOT step (exactly one stimulus per condition)
            curve = [r['outputs'] for r in records
                      if r['accent'] == accent
                      and r['continuum_number'] == cont_num  
                      and r['vot'] == step][0]
                    # [0] so its not nested list 
                    # (there is only on r in records which satisfies these condition)

            ax.plot(time_axis, curve, linewidth=2, label=f'Step {step}')

            # Black dashed
            if step in vot_end:
                ax.axvline(x=vot_end[step], color='black', linestyle='--', linewidth=1.5)

            # Red dash‑dot
            if cont_num in sound_end and accent in sound_end[cont_num]:
                if step in sound_end[cont_num][accent]:
                    sound_ms = sound_end[cont_num][accent][step]
                    ax.axvline(x=sound_ms, color='red', linestyle='-.', linewidth=1.2)

        ax.set_title(f'{accent}')
        ax.set_xlabel('Time (ms)')
        ax.set_ylabel('Predicted rating')
        ax.axhline(0, color='grey', linestyle=':', alpha=0.5)
        ax.legend(fontsize=8)
        ax.set_xlim(0, max(time_axis) + 5)

    fig.suptitle(f'Continuum: {cont_label}', fontsize=14)
    plt.tight_layout()
    return fig


fig1 = plot_continuum("Cɔt", 1, temporal_records, vot_endpoints_1, sound_endpoints)
plt.savefig("temporal_Cɔt.png", dpi=300, bbox_inches='tight')
fig2 = plot_continuum("Cruː", 2, temporal_records, vot_endpoints_2, sound_endpoints)
plt.savefig("temporal_Cruː.png", dpi=300, bbox_inches='tight')
plt.show()