import numpy as np
import pandas as pd
from scipy.optimize import curve_fit
import matplotlib.pyplot as plt

# Load your data from a .dat file
data = pd.read_csv('Viscosity_Water.dat', delimiter=r'\s+')                # pandas assumes that the first non-empty line = header

# Extract the temperature and viscosity data
x_data = data['Temperature'].values                                        # x_data - array with all temperature values
y_data = data['Viscosity'].values                                          # y_data - array with all viscosity values

# Define the model function
def model_func(x, a, b, c):
    y = a * (x / b - 1) ** c
    return y

# Perform the non-linear curve fitting
initial_guess = [1, 200, 1]                                                                             # Initial guess for a, b, c
params, covariance = curve_fit(f=model_func, xdata=x_data, ydata=y_data, p0=initial_guess)              # curve_fit() is part of scipy.optimize
 
# Extract the optimized parameters
a_opt, b_opt, c_opt = params
print(f"Optimized parameters: a = {a_opt}, b = {b_opt}, c = {c_opt}")

# Calculate the fitted values
y_fitted = model_func(x_data, a_opt, b_opt, c_opt)

# Calculate residuals
residuals = y_data - y_fitted
np.savetxt("residuals.txt", residuals)

# Compute the correlation coefficient
correlation_matrix = np.corrcoef(y_data, y_fitted)
correlation_coefficient = correlation_matrix[0, 1]
print(f"Correlation coefficient: {correlation_coefficient}")

# Plot the data and the fitted curve
plt.figure(figsize=(10, 6))
plt.scatter(x_data, y_data, label='Exp.', color='blue')
plt.plot(x_data, y_fitted, label='Speedy and Angel equation', color='red')
plt.xlabel('Temperature (K)')
plt.ylabel('Viscosity (mPas)')
plt.title('Water Viscosity')
plt.legend()
plt.savefig('water_viscosity.png', dpi=300, bbox_inches='tight')
plt.show()