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

# Load data from the CSV file, skipping the first row (header) and handling missing or bad values
data = np.genfromtxt('data2.csv', delimiter=',', skip_header=1, invalid_raise=False, missing_values='NaN')

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

# Extract columns
diff_amplifier_measured = data[:, 2]
wavegen_measured = data[:, 1]
real_wavegen_gain = data[:, 4]

# Define polynomial degree
degree = 1

# Calculate polynomial fits
coefficients_measured = np.polyfit(wavegen_measured, diff_amplifier_measured, degree)

# Generate polynomial functions using the coefficients
poly_func_measured = np.poly1d(coefficients_measured)

# Calculate R^2 value
r_squared_measured = r2_score(diff_amplifier_measured, poly_func_measured(wavegen_measured))

# Calculate absolute percentage error for polynomial fits (handling zero values)
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

# Plot original data and fitted curves for "Wavegen Measured vs Diff Amplifier Measured"
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()

os.system('cls' if os.name == 'nt' else 'clear')

# Print coefficients, R^2 value, and average absolute percentage error for "Wavegen Measured vs Diff Amplifier Measured"
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, "%")

plt.close()

# Plot "Wavegen Measured vs Real Wavegen Gain"
plt.figure(figsize=(8, 6))
plt.scatter(wavegen_measured, real_wavegen_gain, color='blue', label='Amplifier Gain', s=1)
plt.xlabel('Input (mV)')
plt.ylabel('Amplifier Gain')
plt.title('Amplifier Gain vs Input')
plt.legend()
plt.grid(True)
plt.savefig("images/Amplifier Gain vs Input")
plt.show()