import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np
from scipy import stats
import warnings
warnings.filterwarnings('ignore')

# Set style for better looking plots
plt.style.use('seaborn-v0_8')
sns.set_palette("husl")

def load_and_analyze_data():
    """Load the CSV data and perform analysis"""
    try:
        # Load the data
        df = pd.read_csv('data/RCarson_20250812_GFET001_A001001_noniso.csv')
        print(f"Data loaded successfully! Shape: {df.shape}")
        print(f"Columns: {list(df.columns)}")
        print(f"First few rows:")
        print(df.head())
        print(f"\nData types:")
        print(df.dtypes)
        print(f"\nBasic statistics:")
        print(df.describe())
        
        return df
    except Exception as e:
        print(f"Error loading data: {e}")
        return None

def create_vrse_vs_vk_plots(df):
    """Create various plots of vrse vs vk and calculate R²"""
    if df is None:
        return
    
    # Clean the data - remove railed values and warm-up period
    print("\n=== DATA CLEANING ===")
    print(f"Original data shape: {df.shape}")
    
    # Remove rows with NaN values
    df_clean = df.dropna(subset=['Vrse', 'Vk'])
    print(f"After removing NaN: {df_clean.shape}")
    
    # Remove railed values (-2.048)
    railed_mask = (df_clean['Vrse'] == -2.048) | (df_clean['Vk'] == -2.048)
    df_clean = df_clean[~railed_mask]
    print(f"After removing railed values (-2.048): {df_clean.shape}")
    
    # Remove warm-up period (first 1000 measurements or first 10% of data)
    warmup_cutoff = min(1000, int(len(df_clean) * 0.1))
    df_clean = df_clean.iloc[warmup_cutoff:]
    print(f"After removing warm-up period (first {warmup_cutoff} measurements): {df_clean.shape}")
    
    # Additional filtering: remove extreme outliers (beyond 3 standard deviations)
    vrse_mean, vrse_std = df_clean['Vrse'].mean(), df_clean['Vrse'].std()
    vk_mean, vk_std = df_clean['Vk'].mean(), df_clean['Vk'].std()
    
    outlier_mask = (
        (abs(df_clean['Vrse'] - vrse_mean) > 3 * vrse_std) |
        (abs(df_clean['Vk'] - vk_mean) > 3 * vk_std)
    )
    df_clean = df_clean[~outlier_mask]
    print(f"After removing extreme outliers (3σ): {df_clean.shape}")
    
    print(f"Final cleaned data shape: {df_clean.shape}")
    print(f"Data range - Vrse: [{df_clean['Vrse'].min():.6f}, {df_clean['Vrse'].max():.6f}]")
    print(f"Data range - Vk: [{df_clean['Vk'].min():.6f}, {df_clean['Vk'].max():.6f}]")
    
    # Calculate correlation and R²
    correlation = df_clean['Vrse'].corr(df_clean['Vk'])
    r_squared = correlation ** 2
    
    # Linear regression
    slope, intercept, r_value, p_value, std_err = stats.linregress(df_clean['Vrse'], df_clean['Vk'])
    r_squared_regression = r_value ** 2
    
    print(f"\n=== CORRELATION ANALYSIS (CLEANED DATA) ===")
    print(f"Pearson correlation coefficient: {correlation:.6f}")
    print(f"R² (correlation): {r_squared:.6f}")
    print(f"R² (regression): {r_squared_regression:.6f}")
    print(f"P-value: {p_value:.2e}")
    print(f"Slope: {slope:.6f}")
    print(f"Intercept: {intercept:.6f}")
    
    # Create figure with multiple subplots
    fig, axes = plt.subplots(2, 2, figsize=(15, 12))
    fig.suptitle('Vrse vs Vk Analysis (Cleaned Data)', fontsize=16, fontweight='bold')
    
    # Plot 1: Basic scatter plot
    axes[0, 0].scatter(df_clean['Vrse'], df_clean['Vk'], alpha=0.6, s=20)
    axes[0, 0].set_xlabel('Vrse (V)')
    axes[0, 0].set_ylabel('Vk (V)')
    axes[0, 0].set_title(f'Vrse vs Vk Scatter Plot (Cleaned)\nR² = {r_squared:.6f}')
    axes[0, 0].grid(True, alpha=0.3)
    
    # Add regression line
    x_range = np.linspace(df_clean['Vrse'].min(), df_clean['Vrse'].max(), 100)
    y_pred = slope * x_range + intercept
    axes[0, 0].plot(x_range, y_pred, 'r-', linewidth=2, label=f'y = {slope:.4f}x + {intercept:.4f}')
    axes[0, 0].legend()
    
    # Plot 2: Scatter plot with color by VbiasPos
    scatter = axes[0, 1].scatter(df_clean['Vrse'], df_clean['Vk'], 
                                 c=df_clean['VbiasPos'], cmap='viridis', 
                                 alpha=0.7, s=20)
    axes[0, 1].set_xlabel('Vrse (V)')
    axes[0, 1].set_ylabel('Vk (V)')
    axes[0, 1].set_title('Vrse vs Vk Colored by VbiasPos (Cleaned)')
    axes[0, 1].grid(True, alpha=0.3)
    plt.colorbar(scatter, ax=axes[0, 1], label='VbiasPos (V)')
    
    # Plot 3: Scatter plot with error bars (if std values are meaningful)
    if 'Vrse_std' in df_clean.columns and 'Vk_std' in df_clean.columns:
        # Only plot error bars for non-zero standard deviations
        mask = (df_clean['Vrse_std'] > 0) & (df_clean['Vk_std'] > 0)
        if mask.sum() > 0:
            df_with_errors = df_clean[mask]
            axes[1, 0].errorbar(df_with_errors['Vrse'], df_with_errors['Vk'],
                               xerr=df_with_errors['Vrse_std'], 
                               yerr=df_with_errors['Vk_std'],
                               fmt='o', alpha=0.6, capsize=3, capthick=1)
        else:
            axes[1, 0].scatter(df_clean['Vrse'], df_clean['Vk'], alpha=0.6, s=20)
    else:
        axes[1, 0].scatter(df_clean['Vrse'], df_clean['Vk'], alpha=0.6, s=20)
    
    axes[1, 0].set_xlabel('Vrse (V)')
    axes[1, 0].set_ylabel('Vk (V)')
    axes[1, 0].set_title('Vrse vs Vk with Error Bars (Cleaned)')
    axes[1, 0].grid(True, alpha=0.3)
    
    # Plot 4: Residuals plot
    y_pred_actual = slope * df_clean['Vrse'] + intercept
    residuals = df_clean['Vk'] - y_pred_actual
    
    axes[1, 1].scatter(df_clean['Vrse'], residuals, alpha=0.6, s=20)
    axes[1, 1].axhline(y=0, color='r', linestyle='--', alpha=0.7)
    axes[1, 1].set_xlabel('Vrse (V)')
    axes[1, 1].set_ylabel('Residuals (V)')
    axes[1, 1].set_title('Residuals Plot (Cleaned)')
    axes[1, 1].grid(True, alpha=0.3)
    
    plt.tight_layout()
    plt.savefig('vrse_vs_vk_analysis_cleaned.png', dpi=300, bbox_inches='tight')
    plt.show()
    
    return r_squared, slope, intercept

def create_additional_plots(df):
    """Create additional analysis plots"""
    if df is None:
        return
    
    # Clean the data using the same criteria as the main analysis
    df_clean = df.dropna(subset=['Vrse', 'Vk'])
    
    # Remove railed values (-2.048)
    railed_mask = (df_clean['Vrse'] == -2.048) | (df_clean['Vk'] == -2.048)
    df_clean = df_clean[~railed_mask]
    
    # Remove warm-up period
    warmup_cutoff = min(1000, int(len(df_clean) * 0.1))
    df_clean = df_clean.iloc[warmup_cutoff:]
    
    # Remove extreme outliers
    vrse_mean, vrse_std = df_clean['Vrse'].mean(), df_clean['Vrse'].std()
    vk_mean, vk_std = df_clean['Vk'].mean(), df_clean['Vk'].std()
    
    outlier_mask = (
        (abs(df_clean['Vrse'] - vrse_mean) > 3 * vrse_std) |
        (abs(df_clean['Vk'] - vk_mean) > 3 * vk_std)
    )
    df_clean = df_clean[~outlier_mask]
    
    # Create a figure with additional plots
    fig, axes = plt.subplots(2, 2, figsize=(15, 12))
    fig.suptitle('Additional Vrse vs Vk Analysis (Cleaned Data)', fontsize=16, fontweight='bold')
    
    # Plot 1: Time series of Vrse and Vk
    if 'Datetime' in df_clean.columns:
        try:
            df_clean['Datetime'] = pd.to_datetime(df_clean['Datetime'])
            df_clean = df_clean.sort_values('Datetime')
            
            axes[0, 0].plot(df_clean['Datetime'], df_clean['Vrse'], label='Vrse', alpha=0.7)
            axes[0, 0].plot(df_clean['Datetime'], df_clean['Vk'], label='Vk', alpha=0.7)
            axes[0, 0].set_xlabel('Time')
            axes[0, 0].set_ylabel('Voltage (V)')
            axes[0, 0].set_title('Vrse and Vk vs Time (Cleaned)')
            axes[0, 0].legend()
            axes[0, 0].tick_params(axis='x', rotation=45)
            axes[0, 0].grid(True, alpha=0.3)
        except:
            axes[0, 0].text(0.5, 0.5, 'Time series plot not available', 
                           ha='center', va='center', transform=axes[0, 0].transAxes)
    
    # Plot 2: Histogram of Vrse
    axes[0, 1].hist(df_clean['Vrse'], bins=50, alpha=0.7, edgecolor='black')
    axes[0, 1].set_xlabel('Vrse (V)')
    axes[0, 1].set_ylabel('Frequency')
    axes[0, 1].set_title('Histogram of Vrse Values (Cleaned)')
    axes[0, 1].grid(True, alpha=0.3)
    
    # Plot 3: Histogram of Vk
    axes[1, 0].hist(df_clean['Vk'], bins=50, alpha=0.7, edgecolor='black')
    axes[1, 0].set_xlabel('Vk (V)')
    axes[1, 0].set_ylabel('Frequency')
    axes[1, 0].set_title('Histogram of Vk Values (Cleaned)')
    axes[1, 0].grid(True, alpha=0.3)
    
    # Plot 4: Box plot comparison
    data_to_plot = [df_clean['Vrse'], df_clean['Vk']]
    axes[1, 1].boxplot(data_to_plot, labels=['Vrse', 'Vk'])
    axes[1, 1].set_ylabel('Voltage (V)')
    axes[1, 1].set_title('Box Plot Comparison: Vrse vs Vk (Cleaned)')
    axes[1, 1].grid(True, alpha=0.3)
    
    plt.tight_layout()
    plt.savefig('vrse_vs_vk_additional_analysis_cleaned.png', dpi=300, bbox_inches='tight')
    plt.show()

def main():
    """Main function to run the analysis"""
    print("=== Vrse vs Vk Analysis ===")
    print("Loading data...")
    
    # Load data
    df = load_and_analyze_data()
    
    if df is not None:
        print("\nCreating plots and calculating R²...")
        
        # Create main analysis plots
        r_squared, slope, intercept = create_vrse_vs_vk_plots(df)
        
        # Create additional analysis plots
        create_additional_plots(df)
        
        print(f"\n=== SUMMARY ===")
        print(f"R² value: {r_squared:.6f}")
        print(f"Linear relationship: Vk = {slope:.6f} × Vrse + {intercept:.6f}")
        
        if r_squared > 0.8:
            print("Strong correlation between Vrse and Vk")
        elif r_squared > 0.6:
            print("Moderate correlation between Vrse and Vk")
        elif r_squared > 0.4:
            print("Weak correlation between Vrse and Vk")
        else:
            print("Very weak or no correlation between Vrse and Vk")
            
        print(f"\nAnalysis complete! Check the generated PNG files for plots.")
    else:
        print("Failed to load data. Please check the file path and format.")

if __name__ == "__main__":
    main()