import numpy as np
import csv
import matplotlib.pyplot as plt
import pandas as pd
from scipy.fft import rfft, rfftfreq

def read_csv(file_path):
    timestamps = []
    raw_values = []
    voltage_values = []
    with open('eeg_data.csv', mode = 'r') as f:
        reader = csv.reader(f)
        next(reader)
        for row in reader:
            if len(row) == 3:
                timestamps.append(float(row[0]))
                raw_values.append(float(row[1]))
                voltage_values.append(float(row[2]))
    return timestamps, raw_values, voltage_values

def fourier_transform(voltage_values, samples, T):
    signal = np.array(voltage_values)
    fited_signal = signal * np.hamming(samples)

    xf = rfftfreq(samples, T)
    yf = rfft(fited_signal)

    return xf, yf

def compute_band_powers(xf, yf, eeg_bands):
    powers = {}
    energy = np.abs(yf) ** 2
    for band_name, (f_low, f_high) in eeg_bands.items():
        band_power = 0
        for i in range(len(xf)):
            if f_low <= xf[i] <= f_high:
                band_power += energy[i]
        powers[band_name] = band_power
    return powers

def plot_eeg_bands_heatmap(band_powers):
    plt.figure(figsize=(6, 1.5))
    band_labels = list(band_powers.keys())
    power_values = list(band_powers.values())

    plt.imshow([power_values], cmap='viridis', aspect='auto')
    plt.xticks(ticks=np.arange(len(band_labels)), labels=band_labels)
    plt.yticks([])
    plt.colorbar(label="Putere relativa")
    plt.title("Putere pe benzi EEG")
    plt.tight_layout()
    plt.show()    
def main():
    T = 1.0 / 250.0
    timestamps = []
    raw_values = []
    voltage_values = []

    timestamps, raw_values, voltage_values = read_csv('eeg_data.csv')
    samples = len(voltage_values)

    xf, yf = fourier_transform(voltage_values, samples, T)
    
    eeg_bands = {
        "Delta": (0.5, 4),
        "Theta": (4, 8),
        "Alpha": (8, 12),
        "Beta": (12, 35),
        "Gamma": (35, 50)
    }

    band_powers = compute_band_powers(xf, yf, eeg_bands)
    magnitude = np.abs(yf)

    print("\nPutere pe benzi EEG:")
    for band, power in band_powers.items():
        print(f"{band}: {power:.2f}")

    plot_eeg_bands_heatmap(band_powers)

if __name__ == "__main__":
    main()