CPTSD vs PTSD vs Anxiety Py

import mne
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import json
import csv
from scipy.signal import welch
from mne.preprocessing import ICA, create_eog_epochs, create_ecg_epochs
from mne.viz import plot_topomap
from mne.channels import make_standard_montage
from mne import create_info, EvokedArray

# ---- CONFIG ----
BANDS = {
    "delta": (0.5, 4),
    "theta": (4, 8),
    "alpha": (8, 12),
    "smr": (12, 15),
    "beta": (15, 25),
    "hi_beta": (25, 30),
    "gamma": (30, 45)
}

CHANNELS_10_20 = ['Fp1', 'Fp2', 'F3', 'F4', 'C3', 'C4', 'P3', 'P4',
                  'O1', 'O2', 'F7', 'F8', 'T3', 'T4', 'T5', 'T6',
                  'Fz', 'Cz', 'Pz']

# ---- FEATURE EXTRACTION ----
def compute_bandpower(data, sfreq, band):
    fmin, fmax = band
    freqs, psd = welch(data, sfreq, nperseg=1024)
    idx_band = np.logical_and(freqs >= fmin, freqs <= fmax)
    return np.sum(psd[:, idx_band], axis=1)

def extract_features(raw):
    sfreq = raw.info['sfreq']
    data = raw.get_data(picks=CHANNELS_10_20)
    features = {}
    for i, ch_name in enumerate(CHANNELS_10_20):
        ch_data = data[i, :]
        features[ch_name] = {}
        for band_name, band_range in BANDS.items():
            power = compute_bandpower(ch_data[np.newaxis, :], sfreq, band_range)
            features[ch_name][band_name] = float(power[0])
    return features

# ---- ARTIFACT REJECTION ----
def reject_artifacts(raw):
    raw.set_eeg_reference('average', projection=True)
    ica = ICA(n_components=15, random_state=97, max_iter='auto')
    ica.fit(raw)

    # Eye blinks
    eog_epochs = create_eog_epochs(raw, reject_by_annotation=True)
    eog_inds, _ = ica.find_bads_eog(eog_epochs)
    ica.exclude.extend(eog_inds)

    # ECG artifacts
    ecg_epochs = create_ecg_epochs(raw, reject_by_annotation=True)
    ecg_inds, _ = ica.find_bads_ecg(ecg_epochs)
    ica.exclude.extend(ecg_inds)

    # EMG noise (high frequency)
    freqs, psd = welch(raw.get_data(), raw.info['sfreq'])
    emg_noise = psd.mean(axis=1).max() > 1e-6  # crude threshold
    if emg_noise:
        print("⚠️ High EMG noise detected")

    raw_clean = ica.apply(raw.copy())
    return raw_clean

# ---- CLASSIFICATION ----
def classify_conditions_advanced(features):
    rationale = {}
    results = {}

    beta_fz = features['Fz']['beta']
    hi_beta_fz = features['Fz']['hi_beta']
    alpha_occ = np.mean([features['O1']['alpha'], features['O2']['alpha']])
    theta_cz = features['Cz']['theta']
    theta_coherence_fz_pz = abs(features['Fz']['theta'] - features['Pz']['theta'])
    gamma_post = np.mean([features['O1']['gamma'], features['O2']['gamma']])
    alpha_asymmetry = features['F3']['alpha'] - features['F4']['alpha']

    results['PTSD'] = int(hi_beta_fz > 10 and beta_fz > 12 and alpha_occ < 6)
    rationale['PTSD'] = {
        "hi_beta_fz": hi_beta_fz,
        "beta_fz": beta_fz,
        "alpha_occipital_mean": alpha_occ,
        "criteria_met": results['PTSD']
    }

    results['cPTSD'] = int(
        results['PTSD'] and
        theta_cz > 8 and
        theta_coherence_fz_pz < 2.0 and
        gamma_post < 3.5
    )
    rationale['cPTSD'] = {
        "inherits_PTSD": results['PTSD'],
        "theta_cz": theta_cz,
        "theta_coherence_fz_pz": theta_coherence_fz_pz,
        "posterior_gamma_mean": gamma_post,
        "criteria_met": results['cPTSD']
    }

    beta_frontal = np.mean([features['F3']['beta'], features['F4']['beta']])
    alpha_frontal = np.mean([features['F3']['alpha'], features['F4']['alpha']])

    results['Anxiety (Primary)'] = int(beta_frontal > 12 and alpha_frontal > 6 and not results['PTSD'])
    results['Anxiety (Symptom)'] = int(beta_frontal > 10 and results['PTSD'])

    rationale['Anxiety (Primary)'] = {
        "beta_frontal": beta_frontal,
        "alpha_frontal": alpha_frontal,
        "criteria_met": results['Anxiety (Primary)']
    }

    rationale['Anxiety (Symptom)'] = {
        "beta_frontal": beta_frontal,
        "PTSD_flag": results['PTSD'],
        "criteria_met": results['Anxiety (Symptom)']
    }

    return results, rationale

# ---- ARCHETYPE MAPPING ----
def map_archetypes(features):
    alpha_frontal = np.mean([features['F3']['alpha'], features['F4']['alpha']])
    alpha_occip = np.mean([features['O1']['alpha'], features['O2']['alpha']])
    theta_frontal = np.mean([features['F3']['theta'], features['F4']['theta']])
    gamma_cortex = np.mean([features[site]['gamma'] for site in ['Cz', 'Pz', 'Fz']])
    delta_occip = np.mean([features['O1']['delta'], features['O2']['delta']])
    beta_fz = features['Fz']['beta']

    archetypes = []

    if alpha_frontal > 10 and gamma_cortex > 5:
        archetypes.append("The Sage")
    if beta_fz > 12 and alpha_occip < 6:
        archetypes.append("The Hero")
    if alpha_occip > 10 and delta_occip < 5 and theta_frontal < 6:
        archetypes.append("The Caregiver")
    if alpha_frontal > 12 and beta_fz < 6 and gamma_cortex < 3:
        archetypes.append("The Innocent")
    if theta_frontal > 8 and gamma_cortex > 7:
        archetypes.append("The Explorer")
    if alpha_frontal < 6 and theta_frontal > 8:
        archetypes.append("The Orphan")
    if gamma_cortex > 8:
        archetypes.append("The Magician")

    if not archetypes:
        archetypes.append("Mixed or transitional")

    return archetypes

# ---- TOPOGRAPHIC MAPPING ----
def generate_topomaps(features, output_prefix='topomap'):
    montage = make_standard_montage("standard_1020")
    ch_names = list(features.keys())
    ch_types = ["eeg"] * len(ch_names)
    info = create_info(ch_names=ch_names, sfreq=250, ch_types=ch_types)
    info.set_montage(montage)

    bands = list(next(iter(features.values())).keys())
    for band in bands:
        band_values = [features[ch][band] for ch in ch_names]
        evoked = EvokedArray(np.expand_dims(band_values, axis=1), info)
        fig, ax = plt.subplots()
        plot_topomap(evoked.data[:, 0], evoked.info, axes=ax, show=False)
        ax.set_title(f'{band.capitalize()} Band Topomap')
        plt.savefig(f'{output_prefix}_{band}.png', dpi=150)
        plt.close()

# ---- NEUROFEEDBACK PROTOCOL ----
def suggest_protocol(classifications):
    protocol = {}
    if classifications.get("PTSD", 0) > 0.8:
        protocol = {"Site": "Fz, Pz", "Reward": "12–15 Hz (SMR)", "Inhibit": "Hi-Beta, Theta", "Type": "Stabilization"}
    elif classifications.get("Anxiety (Primary)", 0) > 0.8:
        protocol = {"Site": "F3, F4", "Reward": "Low Beta", "Inhibit": "Hi-Beta", "Type": "Calming"}
    elif classifications.get("cPTSD", 0) > 0.8:
        protocol = {"Site": "O1, Cz", "Reward": "Alpha", "Inhibit": "Delta, Theta", "Type": "Alpha-Theta"}
    else:
        protocol = {"Site": "Cz", "Reward": "12–15 Hz", "Inhibit": "Theta", "Type": "Default"}
    return protocol

# ---- EXPORT ----
def export_research_data(subject_id, features, classifications, rationale):
    flat_features = {'subject_id': subject_id}
    for ch, bands in features.items():
        for band, val in bands.items():
            flat_features[f"{ch}_{band}"] = val

    with open(f"{subject_id}_features.csv", "w", newline="") as csvfile:
        writer = csv.DictWriter(csvfile, fieldnames=flat_features.keys())
        writer.writeheader()
        writer.writerow(flat_features)

    out_json = {
        "subject_id": subject_id,
        "features": features,
        "classification": classifications,
        "rationale_trace": rationale
    }

    with open(f"{subject_id}_research_export.json", "w") as f:
        json.dump(out_json, f, indent=4)

# ---- MAIN PIPELINE ----
def run_qeeg_research_mode(edf_path, subject_id="subject001"):
    raw = mne.io.read_raw_edf(edf_path, preload=True)
    raw.pick_channels(CHANNELS_10_20)
    raw.filter(1., 45., fir_design='firwin')

    raw_clean = reject_artifacts(raw)
    features = extract_features(raw_clean)
    classifications, rationale = classify_conditions_advanced(features)
    protocol = suggest_protocol(classifications)
    archetypes = map_archetypes(features)

    output = {
        "subject_id": subject_id,
        "features": features,
        "diagnostic_confidence": classifications,
        "protocol_suggestion": protocol,
        "archetype_mapping": archetypes,
        "diagnostic_rationale": rationale
    }

    with open(f"{subject_id}_full_output.json", "w") as f:
        json.dump(output, f, indent=4)

    export_research_data(subject_id, features, classifications, rationale)
    generate_topomaps(features, output_prefix=f"{subject_id}_topomap")
    print(f"🧠 QEEG Analysis complete for {subject_id}")
    return output

# Example usage:
# run_qeeg_research_mode("/mnt/data/subject1.edf", subject_id="anon_001")