# -*- coding: utf-8 -*-
"""
Convert BDF file to MAT format.

This script converts a BDF (Biosemi Data Format) EEG file to .mat format,
matching the structure of eeg_data.mat.

Structure of eeg_data.mat:
- data: (n_samples, n_channels) float64
- chn: (1, n_channels) object - channel names
- sample_rate: (1, 1) int64
- node_number: (1, 1) int64
- t: (n_samples, 1) float64 - time vector in seconds
- electrode_name: (1, n_channels) object - electrode names (10-20 system)
- electrode_xyz: (n_channels, 3) float64 - electrode 3D coordinates
- electrode_coord_system: (1,) <U21
- meta: (1, 1) structured - metadata
"""

import numpy as np
import scipy.io
import mne
from datetime import datetime


def get_standard_electrode_coords():
    """
    Standard 10-20 system electrode coordinates.
    Returns a dictionary mapping electrode names to x, y, z coordinates.
    Coordinates are approximate spherical projections.
    """
    coords = {
        'FP1': (-0.0293, 0.0903, -0.0033),
        'FP2': (0.0293, 0.0903, -0.0033),
        'FPZ': (0.0, 0.0903, -0.0033),
        'AF7': (-0.0658, 0.0734, -0.0224),
        'AF3': (-0.0350, 0.0812, -0.0183),
        'AF4': (0.0350, 0.0812, -0.0183),
        'AF8': (0.0658, 0.0734, -0.0224),
        'F7': (-0.0815, 0.0467, -0.0336),
        'F5': (-0.0667, 0.0503, -0.0351),
        'F3': (-0.0489, 0.0560, -0.0370),
        'F1': (-0.0254, 0.0584, -0.0384),
        'FZ': (0.0, 0.0584, -0.0384),
        'F2': (0.0254, 0.0584, -0.0384),
        'F4': (0.0489, 0.0560, -0.0370),
        'F6': (0.0667, 0.0503, -0.0351),
        'F8': (0.0815, 0.0467, -0.0336),
        'FT7': (-0.0880, 0.0229, -0.0397),
        'FC5': (-0.0699, 0.0317, -0.0402),
        'FC3': (-0.0514, 0.0362, -0.0411),
        'FC1': (-0.0268, 0.0383, -0.0419),
        'FCZ': (0.0, 0.0383, -0.0419),
        'FC2': (0.0268, 0.0383, -0.0419),
        'FC4': (0.0514, 0.0362, -0.0411),
        'FC6': (0.0699, 0.0317, -0.0402),
        'FT8': (0.0880, 0.0229, -0.0397),
        'T7': (-0.0958, 0.0, -0.0411),
        'T8': (0.0958, 0.0, -0.0411),
        'C5': (-0.0739, 0.0, -0.0425),
        'C3': (-0.0544, 0.0, -0.0436),
        'C1': (-0.0283, 0.0, -0.0444),
        'CZ': (0.0, 0.0, -0.0444),
        'C2': (0.0283, 0.0, -0.0444),
        'C4': (0.0544, 0.0, -0.0436),
        'C6': (0.0739, 0.0, -0.0425),
        'TP7': (-0.0880, -0.0229, -0.0397),
        'CP5': (-0.0699, -0.0317, -0.0402),
        'CP3': (-0.0514, -0.0362, -0.0411),
        'CP1': (-0.0268, -0.0383, -0.0419),
        'CPZ': (0.0, -0.0383, -0.0419),
        'CP2': (0.0268, -0.0383, -0.0419),
        'CP4': (0.0514, -0.0362, -0.0411),
        'CP6': (0.0699, -0.0317, -0.0402),
        'TP8': (0.0880, -0.0229, -0.0397),
        'P7': (-0.0815, -0.0467, -0.0336),
        'P5': (-0.0667, -0.0503, -0.0351),
        'P3': (-0.0489, -0.0560, -0.0370),
        'P1': (-0.0254, -0.0584, -0.0384),
        'PZ': (0.0, -0.0584, -0.0384),
        'P2': (0.0254, -0.0584, -0.0384),
        'P4': (0.0489, -0.0560, -0.0370),
        'P6': (0.0667, -0.0503, -0.0351),
        'P8': (0.0815, -0.0467, -0.0336),
        'PO7': (-0.0658, -0.0734, -0.0224),
        'PO5': (-0.0503, -0.0744, -0.0258),
        'PO3': (-0.0350, -0.0812, -0.0183),
        'POZ': (0.0, -0.0829, -0.0172),
        'PO4': (0.0350, -0.0812, -0.0183),
        'PO6': (0.0503, -0.0744, -0.0258),
        'PO8': (0.0658, -0.0734, -0.0224),
        'O1': (-0.0293, -0.0903, -0.0033),
        'OZ': (0.0, -0.0903, -0.0033),
        'O2': (0.0293, -0.0903, -0.0033),
        'CB1': (-0.0618, -0.0380, -0.0387),
        'CB2': (0.0618, -0.0380, -0.0387),
        'A1': (-0.0958, 0.0, 0.0),
        'A2': (0.0958, 0.0, 0.0),
    }
    return coords


def bdf_to_mat(bdf_path, output_path, subject_id='unknown', session_id='unknown'):
    """
    Convert BDF file to MAT format matching eeg_data.mat structure.

    Parameters
    ----------
    bdf_path : str
        Path to the input BDF file.
    output_path : str
        Path to the output MAT file.
    subject_id : str, optional
        Subject identifier. Default is 'unknown'.
    session_id : str, optional
        Session identifier. Default is 'unknown'.
    """
    print(f'Loading BDF file: {bdf_path}')
    raw = mne.io.read_raw_bdf(bdf_path, preload=True, verbose=False)

    # Get basic info
    ch_names = raw.ch_names
    n_channels = len(ch_names)
    sfreq = int(raw.info['sfreq'])
    data = raw.get_data()

    # BDF data shape: (n_channels, n_samples)
    # Convert to eeg_data.mat format: (n_samples, n_channels)
    data = data.T
    n_samples = data.shape[0]

    # Create time vector (in seconds)
    t = np.arange(n_samples) / sfreq
    t = t.reshape(-1, 1)

    # Create channel names array (matching eeg_data.mat structure)
    chn = np.array([[name] for name in ch_names], dtype=object)

    # Create electrode names (same as channel names for BDF)
    electrode_name = np.array([[name] for name in ch_names], dtype=object)

    # Get electrode coordinates
    standard_coords = get_standard_electrode_coords()
    electrode_xyz = np.zeros((n_channels, 3))
    for i, name in enumerate(ch_names):
        if name in standard_coords:
            electrode_xyz[i] = standard_coords[name]
        else:
            print(f'Warning: No standard coordinate for electrode {name}')

    # Create metadata structure
    start_time_str = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
    meta = np.array([(subject_id, session_id, 'CMS/DRL', start_time_str)],
                    dtype=[('subject_id', 'O'), ('session_id', 'O'),
                           ('ref', 'O'), ('start_time', 'O')])

    # Create the EEG structure (matching eeg_data.mat format)
    eeg_struct = np.array([(data, chn, [[sfreq]], [[n_channels]],
                           t, electrode_name, electrode_xyz,
                           'buzsaki', meta)],
                          dtype=[('data', 'O'), ('chn', 'O'),
                                 ('sample_rate', 'O'), ('node_number', 'O'),
                                 ('t', 'O'), ('electrode_name', 'O'),
                                 ('electrode_xyz', 'O'), ('electrode_coord_system', 'O'),
                                 ('meta', 'O')])

    # Save to MAT file
    print(f'Saving to: {output_path}')
    scipy.io.savemat(output_path, {'eeg': eeg_struct}, do_compression=True)

    print(f'\nConversion complete!')
    print(f'  Channels: {n_channels}')
    print(f'  Samples: {n_samples}')
    print(f'  Duration: {n_samples / sfreq:.2f} seconds')
    print(f'  Sample rate: {sfreq} Hz')
    print(f'  Data shape: {data.shape}')


def main():
    # File paths
    bdf_path = r'D:\Ivey\Code_New_Proj\Debug_Depression\algorithm_version_0521_v0\0515-18.bdf'
    output_path = r'D:\Ivey\Code_New_Proj\Debug_Depression\algorithm_version_0521_v0\0515-18.mat'

    # Convert
    bdf_to_mat(bdf_path, output_path, subject_id='lvpeng', session_id='01')

    # Verify the output
    print('\n=== Verification ===')
    mat_data = scipy.io.loadmat(output_path)
    eeg = mat_data['eeg'][0, 0]

    print(f'Output file keys: {list(mat_data.keys())}')
    print(f'eeg.data shape: {eeg["data"].shape}')
    print(f'eeg.chn shape: {eeg["chn"].shape}')
    print(f'eeg.sample_rate: {eeg["sample_rate"][0, 0]}')
    print(f'eeg.t shape: {eeg["t"].shape}')
    print(f'eeg.electrode_name: {eeg["electrode_name"].shape}')
    print(f'eeg.electrode_xyz shape: {eeg["electrode_xyz"].shape}')
    print(f'eeg.electrode_coord_system: {eeg["electrode_coord_system"][0]}')


if __name__ == '__main__':
    main()