import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import statsmodels.api as sm
from statsmodels.formula.api import ols
from scipy import stats
from scipy.optimize import minimize
import seaborn as sns

# Create DataFrame from the tree data in Table 2.18
data = {
    'diameter': [8.3, 8.6, 8.8, 10.5, 10.7, 10.8, 11.0, 11.0, 11.1, 11.2, 11.3, 11.4, 11.4, 11.7, 
                 12.0, 12.9, 12.9, 13.3, 13.7, 13.8, 14.0, 14.2, 14.5, 16.0, 16.3, 17.3, 17.5, 17.9, 
                 18.0, 18.0, 20.6],
    'height': [70, 65, 63, 72, 81, 83, 66, 75, 80, 75, 79, 76, 76, 69, 75, 74, 85, 86, 71, 64, 78, 
               80, 74, 72, 77, 81, 82, 80, 80, 80, 87],
    'volume': [10.3, 10.3, 10.2, 16.4, 18.8, 19.7, 15.6, 18.2, 22.6, 19.9, 24.2, 21.0, 21.4, 21.3, 
               19.1, 22.2, 33.8, 27.4, 25.7, 24.9, 34.5, 31.7, 36.3, 38.3, 42.6, 55.4, 55.7, 58.3, 
               51.5, 51.0, 77.0]
}

df = pd.DataFrame(data)

# Display basic statistics
print("Dataset summary:")
print(df.describe())

# Correlation analysis
print("\nCorrelation matrix:")
print(df.corr())

# Scatter plot matrix
print("\nCreating scatter plot matrix...")
sns.pairplot(df)
plt.suptitle('Scatter Plot Matrix of Tree Measurements')
plt.tight_layout()
plt.savefig('scatter_matrix.png')

# Question 1: Fit linear regression model and analyze residuals
print("\n\nQuestion 1: Linear Regression Model Analysis")
print("="*60)

# Fit the linear model
model = ols('volume ~ diameter + height', data=df).fit()
print(model.summary())

# Residual analysis
residuals = model.resid
fitted_values = model.fittedvalues

# Plot residuals vs fitted values
plt.figure(figsize=(10, 6))
plt.scatter(fitted_values, residuals)
plt.axhline(y=0, color='r', linestyle='-')
plt.xlabel('Fitted Values')
plt.ylabel('Residuals')
plt.title('Residuals vs Fitted Values')
plt.savefig('residuals_vs_fitted.png')

# Plot residuals vs predictors
fig, axes = plt.subplots(1, 2, figsize=(15, 6))
axes[0].scatter(df['diameter'], residuals)
axes[0].axhline(y=0, color='r', linestyle='-')
axes[0].set_xlabel('Diameter')
axes[0].set_ylabel('Residuals')
axes[0].set_title('Residuals vs Diameter')

axes[1].scatter(df['height'], residuals)
axes[1].axhline(y=0, color='r', linestyle='-')
axes[1].set_xlabel('Height')
axes[1].set_ylabel('Residuals')
axes[1].set_title('Residuals vs Height')
plt.tight_layout()
plt.savefig('residuals_vs_predictors.png')

# Q-Q plot for residuals
fig = sm.qqplot(residuals, line='45', fit=True)
plt.title('Q-Q Plot of Residuals')
plt.savefig('qq_plot.png')

# Shapiro-Wilk test for normality
shapiro_test = stats.shapiro(residuals)
print(f"\nShapiro-Wilk Test for Normality of Residuals:")
print(f"W-statistic: {shapiro_test[0]:.4f}")
print(f"p-value: {shapiro_test[1]:.4f}")
if shapiro_test[1] < 0.05:
    print("Residuals are not normally distributed (p < 0.05)")
    print("Data transformation might be needed.")
else:
    print("Residuals appear to be normally distributed (p >= 0.05)")
    print("No transformation is needed based on normality.")

# Check for heteroscedasticity
plt.figure(figsize=(10, 6))
plt.scatter(fitted_values, np.abs(residuals))
plt.xlabel('Fitted Values')
plt.ylabel('Absolute Residuals')
plt.title('Scale-Location Plot (Check for Heteroscedasticity)')
plt.savefig('scale_location.png')

# Breusch-Pagan test for heteroscedasticity
from statsmodels.stats.diagnostic import het_breuschpagan
bp_test = het_breuschpagan(residuals, model.model.exog)
print(f"\nBreusch-Pagan Test for Heteroscedasticity:")
print(f"LM-statistic: {bp_test[0]:.4f}")
print(f"p-value: {bp_test[1]:.4f}")
if bp_test[1] < 0.05:
    print("Heteroscedasticity detected (p < 0.05)")
    print("Data transformation might help stabilize variance.")
else:
    print("Homoscedasticity assumption is met (p >= 0.05)")

# Question 2: Box-Cox Transformation
print("\n\nQuestion 2: Box-Cox Transformation")
print("="*60)

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

# Define function to optimize - the negative log-likelihood of the Box-Cox transformation
def neg_log_likelihood(lambda_val):
    y = df['volume']
    if any(y <= 0):
        return np.inf  # Box-Cox requires positive values
    
    # Transform data
    if lambda_val == 0:
        y_transformed = np.log(y)
    else:
        y_transformed = (y**lambda_val - 1) / lambda_val
    
    # Calculate residuals from linear model with transformed y
    X = sm.add_constant(df[['diameter', 'height']])
    model = sm.OLS(y_transformed, X).fit()
    residuals = model.resid
    
    # Calculate log-likelihood
    n = len(residuals)
    sigma_squared = np.sum(residuals**2) / n
    log_likelihood = -n/2 * np.log(sigma_squared) + (lambda_val - 1) * np.sum(np.log(y))
    
    return -log_likelihood

# Find optimal lambda value using optimization
lambda_range = np.linspace(-2, 2, 100)
log_likelihoods = [neg_log_likelihood(l) for l in lambda_range]
optimal_lambda_approx = lambda_range[np.argmin(log_likelihoods)]

# Refine search with minimize
result = minimize(neg_log_likelihood, optimal_lambda_approx, method='BFGS')
optimal_lambda = result.x[0]

print(f"Optimal Box-Cox transformation parameter λ: {optimal_lambda:.4f}")

# Transform the dependent variable using the optimal lambda
if abs(optimal_lambda) < 1e-10:  # Close to zero
    y_transformed = np.log(df['volume'])
    transformation_description = "log transformation (λ ≈ 0)"
else:
    y_transformed = boxcox_transform(df['volume'], optimal_lambda)
    transformation_description = f"Box-Cox transformation with λ = {optimal_lambda:.4f}"

# Create a new DataFrame with the transformed variable
df_transformed = df.copy()
df_transformed['volume_transformed'] = y_transformed

# Fit the linear model with transformed data
model_transformed = ols('volume_transformed ~ diameter + height', data=df_transformed).fit()
print(f"\nLinear Regression Model with {transformation_description}:")
print(model_transformed.summary())

# Residual analysis for transformed model
residuals_transformed = model_transformed.resid
fitted_values_transformed = model_transformed.fittedvalues

# Plot residuals vs fitted values for transformed model
plt.figure(figsize=(10, 6))
plt.scatter(fitted_values_transformed, residuals_transformed)
plt.axhline(y=0, color='r', linestyle='-')
plt.xlabel('Fitted Values (Transformed Model)')
plt.ylabel('Residuals (Transformed Model)')
plt.title(f'Residuals vs Fitted Values after {transformation_description}')
plt.savefig('residuals_vs_fitted_transformed.png')

# Q-Q plot for residuals of transformed model
fig = sm.qqplot(residuals_transformed, line='45', fit=True)
plt.title(f'Q-Q Plot of Residuals after {transformation_description}')
plt.savefig('qq_plot_transformed.png')

# Shapiro-Wilk test for normality of transformed residuals
shapiro_test_transformed = stats.shapiro(residuals_transformed)
print(f"\nShapiro-Wilk Test for Normality of Transformed Residuals:")
print(f"W-statistic: {shapiro_test_transformed[0]:.4f}")
print(f"p-value: {shapiro_test_transformed[1]:.4f}")
if shapiro_test_transformed[1] < 0.05:
    print("Transformed residuals are not normally distributed (p < 0.05)")
else:
    print("Transformed residuals appear to be normally distributed (p >= 0.05)")

# Check for heteroscedasticity in transformed model
plt.figure(figsize=(10, 6))
plt.scatter(fitted_values_transformed, np.abs(residuals_transformed))
plt.xlabel('Fitted Values (Transformed Model)')
plt.ylabel('Absolute Residuals (Transformed Model)')
plt.title(f'Scale-Location Plot after {transformation_description}')
plt.savefig('scale_location_transformed.png')

# Breusch-Pagan test for heteroscedasticity in transformed model
bp_test_transformed = het_breuschpagan(residuals_transformed, model_transformed.model.exog)
print(f"\nBreusch-Pagan Test for Heteroscedasticity (Transformed Model):")
print(f"LM-statistic: {bp_test_transformed[0]:.4f}")
print(f"p-value: {bp_test_transformed[1]:.4f}")
if bp_test_transformed[1] < 0.05:
    print("Heteroscedasticity still detected in transformed model (p < 0.05)")
else:
    print("Homoscedasticity assumption is met in transformed model (p >= 0.05)")

# Compare the original and transformed models
print("\nComparison of Original and Transformed Models:")
print(f"Original Model R-squared: {model.rsquared:.4f}")
print(f"Transformed Model R-squared: {model_transformed.rsquared:.4f}")

# Conclusion about the effect of Box-Cox transformation
print("\nConclusion about Box-Cox Transformation Effect:")
if shapiro_test_transformed[1] > shapiro_test[1] or bp_test_transformed[1] > bp_test[1]:
    print(f"The Box-Cox transformation with λ = {optimal_lambda:.4f} improved the model by:")
    if shapiro_test_transformed[1] > shapiro_test[1]:
        print("- Improving the normality of residuals")
    if bp_test_transformed[1] > bp_test[1]:
        print("- Reducing heteroscedasticity")
    if model_transformed.rsquared > model.rsquared:
        print(f"- Increasing R-squared from {model.rsquared:.4f} to {model_transformed.rsquared:.4f}")
else:
    print(f"The Box-Cox transformation with λ = {optimal_lambda:.4f} did not significantly improve the model.")