import numpy as np
import matplotlib.pyplot as plt
from sklearn.metrics import r2_score
from sklearn.linear_model import LinearRegression
import os

# Load data from CSV file, handling missing or malformed data
data = np.genfromtxt('data.csv', delimiter=',', skip_header=1, invalid_raise=False, missing_values='NaN')

# Remove any rows containing NaN values
data = data[~np.isnan(data).any(axis=1)]

# Data extraction from the cleaned dataset
wavegen_measured = data[:, 1]  # Input voltage in mV
diff_amplifier_measured = data[:, 2]  # Amplifier output in volts
real_wavegen_gain = data[:, 4]  # Real gain from the dataset

# Linear regression for Wavegen Measured vs Diff Amplifier Measured
coefficients_measured = np.polyfit(wavegen_measured, diff_amplifier_measured, 1)
poly_func_measured = np.poly1d(coefficients_measured)
r_squared_measured = r2_score(diff_amplifier_measured, poly_func_measured(wavegen_measured))

# Calculate the absolute percentage error for the linear fit, avoiding division by zero
non_zero_indices = np.nonzero(diff_amplifier_measured)
absolute_percentage_error_measured = np.mean(np.abs((diff_amplifier_measured[non_zero_indices] - poly_func_measured(wavegen_measured[non_zero_indices])) / diff_amplifier_measured[non_zero_indices])) * 100

# Plotting the linear fit results
plt.figure(figsize=(8, 6))
plt.scatter(wavegen_measured, diff_amplifier_measured, color='green', label='Amplifier Output', s=1)
plt.plot(wavegen_measured, poly_func_measured(wavegen_measured), color='red', label='Linear Fit')
plt.xlabel('Input (mV)')
plt.ylabel('Amplifier Output (V)')
plt.title('Amplifier Output vs Input')
plt.legend()
plt.grid(True)
plt.savefig("images/Amplifier Output vs Input")
plt.show()

# Clear the console after plotting
os.system('cls' if os.name == 'nt' else 'clear')

# Output the results of the linear regression analysis
print("\nPolynomial Fit Coefficients for Wavegen Measured vs Diff Amplifier Measured:", coefficients_measured)
print("R^2 Value for Wavegen Measured vs Diff Amplifier Measured:", r_squared_measured)
print("Average Absolute Percentage Error for Wavegen Measured vs Diff Amplifier Measured:", absolute_percentage_error_measured, "%")

# Close the current plot
plt.close()

# Segment the data and fit multiple linear regression models
set_size = 80
slopes = []
for i in range(len(data) // set_size):
    set_data = data[i * set_size: (i + 1) * set_size]
    x_values = set_data[:, 1]  # Input values
    y_values = set_data[:, 2]  # Output values
    model = LinearRegression().fit(x_values.reshape(-1, 1), y_values)
    slopes.append(model.coef_[0])
    print("Slope of linear fitting function for set", i+1, ":", model.coef_[0])

# Calculate and print the average slope
average_slope = np.mean(slopes)
print("\nAverage slope (mV/V):", average_slope * 1000)
print("Standard Deviation of slopes:", np.std(slopes) * 1000)
