import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from scipy import stats
from sklearn.linear_model import LinearRegression
import statsmodels.api as sm
from scipy import optimize

# Data from the table
x = np.array([0.05, 0.15, 0.25, 0.35, 0.45, 0.55, 0.65, 0.75, 0.85, 0.95, 
              1.05, 1.15, 1.25, 1.35, 1.45, 1.55, 1.65, 1.75, 1.85, 1.95])
y = np.array([5.942, 5.469, 5.872, 5.182, 5.196, 5.249, 5.136, 5.226, 5.081, 5.224,
              4.735, 4.595, 5.154, 5.284, 5.345, 5.146, 5.409, 5.650, 6.026, 5.535])

# Part 1: Linear Regression Analysis
print("Part 1: Linear Regression Analysis")
print("---------------------------------")

# Create a scatter plot
plt.figure(figsize=(10, 6))
plt.scatter(x, y)
plt.title('Scatter plot of Y vs X')
plt.xlabel('X')
plt.ylabel('Y')
plt.grid(True)
plt.savefig('scatter_plot.png')
plt.close()

# Simple linear regression
X = x.reshape(-1, 1)  # Reshape for sklearn
model = LinearRegression().fit(X, y)
y_pred_linear = model.predict(X)

print(f"Linear Regression Equation: Y = {model.intercept_:.4f} + {model.coef_[0]:.4f}X")
print(f"R² = {model.score(X, y):.4f}")

# Residual analysis
residuals = y - y_pred_linear
residuals_squared = residuals**2
SSE = np.sum(residuals_squared)
MSE = SSE / (len(x) - 2)  # n-2 degrees of freedom for simple linear regression

plt.figure(figsize=(12, 8))

# Plot 1: Residuals vs Fitted values
plt.subplot(2, 2, 1)
plt.scatter(y_pred_linear, residuals)
plt.axhline(y=0, color='r', linestyle='-')
plt.title('Residuals vs Fitted Values')
plt.xlabel('Fitted Values')
plt.ylabel('Residuals')
plt.grid(True)

# Plot 2: Histogram of residuals
plt.subplot(2, 2, 2)
plt.hist(residuals, bins=10, edgecolor='black')
plt.title('Histogram of Residuals')
plt.xlabel('Residuals')
plt.ylabel('Frequency')
plt.grid(True)

# Plot 3: Q-Q plot
plt.subplot(2, 2, 3)
stats.probplot(residuals, dist="norm", plot=plt)
plt.title('Q-Q Plot')
plt.grid(True)

# Plot 4: Residuals vs X
plt.subplot(2, 2, 4)
plt.scatter(x, residuals)
plt.axhline(y=0, color='r', linestyle='-')
plt.title('Residuals vs X')
plt.xlabel('X')
plt.ylabel('Residuals')
plt.grid(True)

plt.tight_layout()
plt.savefig('residual_analysis.png')
plt.close()

print(f"Sum of Squared Residuals (SSE): {SSE:.4f}")
print(f"Mean Squared Error (MSE): {MSE:.4f}")
print(f"Root Mean Squared Error (RMSE): {np.sqrt(MSE):.4f}")

# True model: Y = 5 + (X-1)² + ε, where ε ~ N(0, 0.625)
def true_model(x):
    return 5 + (x - 1)**2

x_range = np.linspace(0, 2, 100)
y_true = true_model(x_range)

# Quadratic regression
X_quad = np.column_stack((x, x**2))
X_quad = sm.add_constant(X_quad)
quad_model = sm.OLS(y, X_quad).fit()
print("\nQuadratic Model Summary:")
print(quad_model.summary().tables[1])

# Generate predictions for the quadratic model
x_for_pred = np.linspace(min(x), max(x), 100)
X_pred = np.column_stack((np.ones(100), x_for_pred, x_for_pred**2))
y_pred_quad = quad_model.predict(X_pred)

# Plot original data, linear fit, and quadratic fit
plt.figure(figsize=(12, 7))
plt.scatter(x, y, label='Original Data')
plt.plot(x_range, y_true, 'r-', label='True Model: Y = 5 + (X-1)²')
plt.plot(x, y_pred_linear, 'g--', label=f'Linear Model: Y = {model.intercept_:.4f} + {model.coef_[0]:.4f}X')
plt.plot(x_for_pred, y_pred_quad, 'b-', 
         label=f'Quadratic Model: Y = {quad_model.params[0]:.4f} + {quad_model.params[1]:.4f}X + {quad_model.params[2]:.4f}X²')
plt.title('Data with True Model and Fitted Models')
plt.xlabel('X')
plt.ylabel('Y')
plt.legend()
plt.grid(True)
plt.savefig('model_comparison.png')
plt.close()

# Residual analysis for quadratic model
quad_residuals = y - quad_model.predict(X_quad)
quad_SSE = np.sum(quad_residuals**2)
quad_MSE = quad_SSE / (len(x) - 3)  # n-3 degrees of freedom for quadratic

print(f"\nQuadratic model equation: Y = {quad_model.params[0]:.4f} + {quad_model.params[1]:.4f}X + {quad_model.params[2]:.4f}X²")
print(f"Quadratic R²: {quad_model.rsquared:.4f}")
print(f"Quadratic SSE: {quad_SSE:.4f}")
print(f"Quadratic MSE: {quad_MSE:.4f}")
print(f"Quadratic RMSE: {np.sqrt(quad_MSE):.4f}")

# Try a different parameterization that better matches the true model: Y = a + b(X-1)²
X_param = ((x - 1)**2).reshape(-1, 1)
X_param = sm.add_constant(X_param)
param_model = sm.OLS(y, X_param).fit()

print("\nParameterized Model: Y = a + b(X-1)²")
print(param_model.summary().tables[1])
print(f"Equation: Y = {param_model.params[0]:.4f} + {param_model.params[1]:.4f}(X-1)²")
print(f"Compare with true model: Y = 5 + (X-1)² + ε, where ε ~ N(0, 0.625)")

# Part 2: Box-Cox Transformation
print("\n\nPart 2: Box-Cox Transformation")
print("---------------------------------")

# Define Box-Cox transformation
def boxcox_transform(y, lmbda):
    if lmbda == 0:
        return np.log(y)
    else:
        return (y**lmbda - 1) / lmbda

# Objective function to minimize (negative log-likelihood)
def neg_log_likelihood(lmbda, data):
    n = len(data)
    y_transformed = boxcox_transform(data, lmbda)
    s2 = np.var(y_transformed, ddof=1)
    return -(-n/2 * np.log(s2) + (lmbda - 1) * np.sum(np.log(data)))

# Find optimal lambda
# First, check that all y values are positive (required for Box-Cox)
if np.min(y) <= 0:
    print("Warning: Box-Cox transformation requires all y values to be positive!")
else:
    result = optimize.minimize_scalar(neg_log_likelihood, args=(y,), bounds=(-2, 2), method='bounded')
    lambda_opt = result.x
    print(f"Optimal λ for Box-Cox transformation: {lambda_opt:.4f}")

    # Apply the Box-Cox transformation with optimal lambda
    y_transformed = boxcox_transform(y, lambda_opt)
    
    # Fit linear model on transformed data
    model_bc = LinearRegression().fit(X, y_transformed)
    y_pred_bc = model_bc.predict(X)
    
    print(f"Linear regression on Box-Cox transformed data:")
    print(f"Equation (transformed scale): Y' = {model_bc.intercept_:.4f} + {model_bc.coef_[0]:.4f}X")
    print(f"R² of transformed model: {model_bc.score(X, y_transformed):.4f}")
    
    # Comparison plots
    plt.figure(figsize=(14, 10))
    
    # Original data
    plt.subplot(2, 2, 1)
    plt.scatter(x, y)
    plt.title('Original Data')
    plt.xlabel('X')
    plt.ylabel('Y')
    plt.grid(True)
    
    # Transformed data
    plt.subplot(2, 2, 2)
    plt.scatter(x, y_transformed)
    plt.plot(x, y_pred_bc, 'r-')
    plt.title(f'Box-Cox Transformed Data (λ = {lambda_opt:.4f})')
    plt.xlabel('X')
    plt.ylabel(f'Y(λ = {lambda_opt:.4f})')
    plt.grid(True)
    
    # Residuals of original linear model
    plt.subplot(2, 2, 3)
    plt.scatter(x, residuals)
    plt.axhline(y=0, color='r', linestyle='-')
    plt.title('Residuals of Linear Model on Original Data')
    plt.xlabel('X')
    plt.ylabel('Residuals')
    plt.grid(True)
    
    # Residuals of Box-Cox transformed model
    bc_residuals = y_transformed - y_pred_bc
    plt.subplot(2, 2, 4)
    plt.scatter(x, bc_residuals)
    plt.axhline(y=0, color='r', linestyle='-')
    plt.title('Residuals of Linear Model on Box-Cox Transformed Data')
    plt.xlabel('X')
    plt.ylabel('Residuals')
    plt.grid(True)
    
    plt.tight_layout()
    plt.savefig('boxcox_analysis.png')
    plt.close()
    
    # Check residuals of the Box-Cox transformed model
    bc_SSE = np.sum(bc_residuals**2)
    bc_MSE = bc_SSE / (len(x) - 2)
    
    print(f"\nBox-Cox transformed model SSE: {bc_SSE:.4f}")
    print(f"Box-Cox transformed model MSE: {bc_MSE:.4f}")
    print(f"Box-Cox transformed model RMSE: {np.sqrt(bc_MSE):.4f}")
    
    # Compare with quadratic model
    print("\nModel Comparison:")
    print(f"Linear Model R²: {model.score(X, y):.4f}")
    print(f"Quadratic Model R²: {quad_model.rsquared:.4f}")
    print(f"Box-Cox Transformed Model R²: {model_bc.score(X, y_transformed):.4f}")
    
    # Insights about Box-Cox transformation
    print("\nInsights about Box-Cox transformation:")
    if abs(lambda_opt - 1.0) < 0.1:
        print("The optimal λ is close to 1, suggesting that no transformation is needed.")
    elif abs(lambda_opt - 0.5) < 0.1:
        print("The optimal λ is close to 0.5, suggesting a square root transformation might be appropriate.")
    elif abs(lambda_opt - 0.0) < 0.1:
        print("The optimal λ is close to 0, suggesting a log transformation might be appropriate.")
    elif lambda_opt < 0:
        print(f"The negative λ value ({lambda_opt:.4f}) suggests a more complex relationship.")
    else:
        print(f"The λ value of {lambda_opt:.4f} suggests that a power transformation is needed.")
    
    # Final comparison with the true model
    print("\nComparison with the true model:")
    print("True model: Y = 5 + (X-1)² + ε, where ε ~ N(0, 0.625)")
    print(f"Our best parameterized model: Y = {param_model.params[0]:.4f} + {param_model.params[1]:.4f}(X-1)²")