import matplotlib.pyplot as plt
import numpy as np
from sklearn.metrics import r2_score
import csv  # useful for reading csv files
import os   # run operating system commands to set current working directory

# change direct to program's folder
os.chdir(os.path.dirname(os.path.realpath(__file__)))

# declare two lists for the temperature data and cricket chirp rate data
temps=[]
chirps=[]

with open("Crickets.csv", "r") as f:
   csv_reader= csv.reader(f, delimiter=",")
   next(csv_reader)     # skip the first row with the headers
   for row in csv_reader:
      temps.append(float(row[0]))    # cast speed data as a float
      chirps.append(float(row[1]))     # cast heart beat data as a float

# Calculate regression line
#m, b = np.polyfit(temps, chirps, 1)
the_fit = np.polynomial.Polynomial.fit(temps, chirps, deg=[1], domain=[], rcond=None, full=False, w=None, window=None, symbol='x')
print(the_fit)
b=the_fit(0)
m= the_fit(1)-the_fit(0)
print(m)
print(b)

# Generate points for the regression line
x_line = np.linspace(min(temps), max(temps), 100)
y_line = m * x_line + b

# Calculate R-squared
y_predicted = np.multiply(m,temps) +b
r2 = r2_score(chirps, y_predicted)

# Create the scatter plot
plt.scatter(temps, chirps, label="Data points")

# Plot the regression line
plt.plot(x_line, y_line, color='red', label=f"Regression line (R² = {r2:.4f})")
#plt.plot(x_line, y_line, color='red')
# Add labels and title
plt.xlabel("Temperature (F)")
plt.ylabel("Chirp Rate (Hz)")
plt.title("Crickets")

# Add legend
plt.legend()

# Show the plot
plt.show()