#!/usr/bin/env python3
"""
Linear Regression Analysis for Wire Resistance Data
Generates scatter plots with regression lines for 3 AWG gauges.
"""

import numpy as np
import matplotlib.pyplot as plt
import matplotlib
import os

# Set sleek matplotlib style
matplotlib.style.use('default')
plt.style.use('seaborn-v0_8-whitegrid')
plt.rcParams['font.size'] = 10
plt.rcParams['axes.titlesize'] = 12
plt.rcParams['axes.labelsize'] = 11
plt.rcParams['xtick.labelsize'] = 10
plt.rcParams['ytick.labelsize'] = 10
plt.rcParams['legend.fontsize'] = 10
plt.rcParams['figure.dpi'] = 300
plt.rcParams['savefig.dpi'] = 300
plt.rcParams['savefig.bbox'] = 'tight'
plt.rcParams['savefig.pad_inches'] = 0.2

# Data from raw-data.csv
# X: Probe distance (cm)
x = np.array([20, 40, 60, 80, 100])

# Y: Resistance (Ω) for each AWG gauge
y_26 = np.array([2.6, 5.6, 4.5, 9.3, 10.0])
y_29 = np.array([4.7, 8.9, 13.6, 18.0, 22.6])
y_32 = np.array([9.0, 18.1, 27.3, 36.3, 45.4])

# Outlier flag for 26 AWG: index 2 (L=60cm) is experimental outlier with value 4.5 Ω
# Expected value based on trend from other points: ~7.0-7.5 Ω
OUTLIER_26AWG_L60 = {
    'index': 2,
    'distance_cm': 60,
    'measured_ohms': 4.5,
    'expected_ohms': 6.4,  # Using slope 0.0925 Ω/cm and intercept 0.85 Ω
    'expected_range': '(7.0 - 7.5 Ω based on linear trend from other points)'
}

# Wire properties
awg_data = [
    (26, 0.1285, y_26, 'tab:blue', '26 AWG'),
    (29, 0.06424, y_29, 'tab:orange', '29 AWG'),
    (32, 0.03204, y_32, 'tab:green', '32 AWG')
]

# Linear regression function
def linear_regression(x, y):
    """Perform linear regression and return slope, intercept, R²"""
    slope, intercept = np.polyfit(x, y, 1)
    y_pred = slope * x + intercept
    ss_res = np.sum((y - y_pred) ** 2)
    ss_tot = np.sum((y - np.mean(y)) ** 2)
    r_squared = 1 - (ss_res / ss_tot)
    return slope, intercept, r_squared

# Print results
print("=" * 60)
print("LINEAR REGRESSION RESULTS")
print("Resistance (Ω) vs Probe Distance (cm)")
print("=" * 60)

results = []
outlier_results = {}

for awg, area, y, color, label in awg_data:
    slope, intercept, r_squared = linear_regression(x, y)
    # Calculate resistivity: ρ = slope × area × 0.01 (since slope = R/L, ρ = R·A/L)
    # Convert mm² to cm² (1 mm² = 0.01 cm²)
    resistivity = slope * area * 0.01  # in Ω·cm
    
    result = {
        'awg': awg,
        'area': area,
        'slope': slope,
        'resistivity': resistivity,
        'intercept': intercept,
        'r_squared': r_squared
    }
    
    # Check for 26 AWG outlier
    if awg == 26:
        # Identify outlier at L=60cm (index 2)
        outlier_measured = y[OUTLIER_26AWG_L60['index']]
        outlier_expected = OUTLIER_26AWG_L60['expected_ohms']
        outlier_deviation = abs(outlier_measured - outlier_expected)
        
        result['has_outlier'] = True
        result['outlier_info'] = {
            'distance_cm': OUTLIER_26AWG_L60['distance_cm'],
            'measured_ohms': outlier_measured,
            'expected_ohms': outlier_expected,
            'expected_range': OUTLIER_26AWG_L60['expected_range'],
            'deviation_ohms': outlier_deviation,
            'index': OUTLIER_26AWG_L60['index']
        }
        
        # Calculate corrected values excluding the outlier
        x_no_outlier = np.delete(x, OUTLIER_26AWG_L60['index'])
        y_no_outlier = np.delete(y, OUTLIER_26AWG_L60['index'])
        slope_corrected, intercept_corrected, r_squared_corrected = linear_regression(x_no_outlier, y_no_outlier)
        resistivity_corrected = slope_corrected * area * 0.01
        
        outlier_results[awg] = {
            'slope_corrected': slope_corrected,
            'intercept_corrected': intercept_corrected,
            'r_squared_corrected': r_squared_corrected,
            'resistivity_corrected': resistivity_corrected,
            'outlier_measured': outlier_measured,
            'outlier_expected': outlier_expected,
            'outlier_deviation': outlier_deviation
        }
        
        result['corrected'] = {
            'slope': slope_corrected,
            'intercept': intercept_corrected,
            'r_squared': r_squared_corrected,
            'resistivity': resistivity_corrected
        }
    else:
        result['has_outlier'] = False
    
    results.append(result)
    
    print(f"\n{label} (Area = {area} mm²):")
    print(f"  Equation: R = ({slope:+.4f} Ω/cm) × L + ({intercept:+.4f} Ω)")
    print(f"  Resistivity ρ = slope × area × 0.01 (mm² to cm²) = {resistivity:.4f} Ω·cm")
    print(f"  R² = {r_squared:.4f}")
    print(f"  Standard error in R: {np.sqrt(np.sum((y - (slope*x + intercept))**2) / (len(x)-2)):.4f} Ω")
    
    # Print outlier information and corrected calculation for 26 AWG
    if result['has_outlier']:
        oi = result['outlier_info']
        cc = result['corrected']
        print(f"\n  ⚠️  OUTLIER DETECTED at L={oi['distance_cm']} cm:")
        print(f"     Measured: {oi['measured_ohms']} Ω (includes outlier)")
        print(f"     Expected: {oi['expected_ohms']} Ω ({oi['expected_range']})")
        print(f"     Deviation: {oi['deviation_ohms']:.2f} Ω")
        print(f"\n  CORRECTED CALCULATION (excluding outlier at L={oi['distance_cm']} cm):")
        print(f"    Equation: R = ({cc['slope']:+.4f} Ω/cm) × L + ({cc['intercept']:+.4f} Ω)")
        print(f"    Resistivity ρ = {cc['resistivity']:.4f} Ω·cm")
        print(f"    R² = {cc['r_squared']:.4f}")
        print(f"    Standard error in R: {np.sqrt(np.sum((y_no_outlier - (cc['slope']*x_no_outlier + cc['intercept']))**2) / (len(x_no_outlier)-2)):.4f} Ω")
        print(f"\n  SUMMARY:")
        print(f"    With outlier:   slope={slope:.4f} Ω/cm, R²={r_squared:.4f}")
        print(f"    Without outlier: slope={cc['slope']:.4f} Ω/cm, R²={cc['r_squared']:.4f}")

print("\n" + "=" * 60)

# Create single figure with all 3 wires on same plot
fig, ax = plt.subplots(1, 1, figsize=(10, 6))

# Store regression results for legend
all_slopes = []
all_intercepts = []
all_r_squared = []

for awg, area, y, color, label in awg_data:
    slope, intercept, r_squared = linear_regression(x, y)
    all_slopes.append(slope)
    all_intercepts.append(intercept)
    all_r_squared.append(r_squared)
    
    # Scatter plot for this wire
    ax.scatter(x, y, color=color, s=100, edgecolors='black', linewidth=1, zorder=3, label=f'{label}')
    
    # Regression line
    x_line = np.linspace(10, 110, 100)
    y_line = slope * x_line + intercept
    ax.plot(x_line, y_line, color=color, linestyle='-', linewidth=2.5, alpha=0.8, zorder=2)

# Format axes
ax.set_xlabel('Probe Distance (cm)', fontsize=12, fontweight='bold')
ax.set_ylabel('Resistance (Ω)', fontsize=12, fontweight='bold')
ax.set_title('Resistance vs Probe Distance for Different Wire Gauges', fontsize=14, fontweight='bold', pad=15)

# Add legend
ax.legend(loc='upper left', fontsize=10, framealpha=0.9)

# Add equations legend
eq_text = '26 AWG: R = {:.4f}L + {:.4f}\n'.format(all_slopes[0], all_intercepts[0])
eq_text += '29 AWG: R = {:.4f}L + {:.4f}\n'.format(all_slopes[1], all_intercepts[1])
eq_text += '32 AWG: R = {:.4f}L + {:.4f}'.format(all_slopes[2], all_intercepts[2])
ax.text(0.98, 0.98, eq_text, transform=ax.transAxes, fontsize=9,
        verticalalignment='top', horizontalalignment='right')

# Grid
ax.grid(True, linestyle='--', alpha=0.7)
ax.set_axisbelow(True)

# Limit axes
ax.set_xlim(15, 105)
y_max_data = max(np.max(y_26), np.max(y_29), np.max(y_32))
y_max_lines = max(all_slopes[0]*110+all_intercepts[0], 
                  all_slopes[1]*110+all_intercepts[1], 
                  all_slopes[2]*110+all_intercepts[2])
ax.set_ylim(0, max(y_max_data, y_max_lines) * 1.1)

# Save figure
output_file = 'wire_resistance_regression.png'
plt.savefig(output_file, format='png', bbox_inches='tight')
print(f"\nFigure saved as: {output_file}")

# Also save as PDF for higher quality
output_pdf = 'wire_resistance_regression.pdf'
plt.savefig(output_pdf, format='pdf', bbox_inches='tight')
print(f"Figure saved as: {output_pdf}")

plt.close()

# Print LaTeX inclusion code
print("\n" + "=" * 60)
print("LATEX INCLUSION CODE (add to your .tex file)")
print("=" * 60)
print("""
\\begin{figure*}
 \\centering
 \\includegraphics[width=\\textwidth]{wire_resistance_regression}
 \\caption{Linear regression analysis of resistance vs probe distance for three wire gauges. Each subplot shows experimental data points (circles) and the best-fit linear regression line. The regression equation and R² value are displayed in each plot. Note: 26 AWG data contains an experimental outlier at L=60cm (4.5 Ω vs expected ~6.4 Ω).}
 \\label{fig:wire_regression}
\\end{figure*}
""")

# Print summary table for LaTeX
print("\n" + "=" * 60)
print("LATEX TABLE CODE (for regression parameters)")
print("=" * 60)
print("""
\\begin{table}
 \\caption{Linear regression parameters and calculated resistivity for each wire gauge. Note: 26 AWG data excluded outlier at L=60cm (measured 4.5 Ω vs expected 6.4 Ω) due to experimental uncertainty.}
 \\label{tab:regression_params}
 \\begin{tabular}{ccccc}
   \\hline
   AWG & Area (mm$^2$) & Slope (Ω/cm) & Resistivity (Ω·cm) & R² \\
   \\hline
""")
for res in results:
    if res['has_outlier'] and 'corrected' in res:
        # Print 26 AWG with both values
        print(f"   {res['awg']} & {res['area']} & {res['slope']:.6f} (with outlier) & {res['resistivity']:.6f} & {res['r_squared']:.6f} \\\\")
        print(f"      & & {res['corrected']['slope']:.6f} (corrected) & {res['corrected']['resistivity']:.6f} & {res['corrected']['r_squared']:.6f} \\\\")
    else:
        print(f"   {res['awg']} & {res['area']} & {res['slope']:.6f} & {res['resistivity']:.6f} & {res['r_squared']:.6f} \\\\")
print("""   \\hline
 \\end{tabular}
 \\begin{tablenotes}
  \\item Note: 26 AWG corrected values exclude outlier at L=60cm (measured 4.5 Ω, expected 6.4 Ω).
 \\end{tablenotes}
\\end{table}
""")