import os
import numpy as np
import pandas as pd
from scipy.stats import chi2_contingency

csv_filename = "Data For Research !0.0.csv"

if not os.path.exists(csv_filename):
    raise FileNotFoundError(f"Could not find '{csv_filename}' in the local directory.")

df = pd.read_csv(csv_filename)

df['log_sSFR'] = df['sSFR']

if df['Mbar'].max() > 100:
    df['log_Mbar'] = np.log10(df['Mbar'])
else:
    df['log_Mbar'] = df['Mbar']

df['is_outlier_mcg'] = df['is_outlier'].astype(bool)

log_Vexp_mcg = (df['log_Mbar'] - 1.70) / 4.0

inc_radians = np.radians(df['Inc'])
log_Vobs_reconstructed = np.log10(
    df['w50'] / (2.0 * np.sin(inc_radians) + np.finfo(float).eps) + np.finfo(float).eps
)
sign_mask = np.sign(log_Vobs_reconstructed - log_Vexp_mcg)

df['delta_V_mcg'] = df['DeltaV_abs'] * sign_mask

df['delta_V_lel'] = df['delta_V_mcg'] - 0.0075
df['is_outlier_lel'] = df['delta_V_lel'].abs() > 0.3

df = df.sort_values(by='log_sSFR').reset_index(drop=True)

target_bin_sizes = [111, 136, 318, 558, 364, 103, 50]
bin_labels = []
for idx, size in enumerate(target_bin_sizes, start=1):
    bin_labels.extend([f"Bin {idx}"] * size)

df['sSFR_bin'] = bin_labels

summary = df.groupby('sSFR_bin', observed=False).agg(
    Total_Galaxies=('log_sSFR', 'count'),
    Outliers_McGaugh=('is_outlier_mcg', 'sum'),
    Outliers_Lelli=('is_outlier_lel', 'sum')
)

summary['Pct_McGaugh'] = (summary['Outliers_McGaugh'] / summary['Total_Galaxies'] * 100).round(1).astype(str) + '%'
summary['Pct_Lelli'] = (summary['Outliers_Lelli'] / summary['Total_Galaxies'] * 100).round(1).astype(str) + '%'

def execute_chi2_test(outlier_column):
    observed_outliers = summary[outlier_column].values
    observed_normals = summary['Total_Galaxies'].values - observed_outliers
    contingency_matrix = np.array([observed_outliers, observed_normals])
    chi2_stat, p_value, _, _ = chi2_contingency(contingency_matrix)
    return chi2_stat, p_value

chi2_mcg, p_mcg = execute_chi2_test('Outliers_McGaugh')
chi2_lel, p_lel = execute_chi2_test('Outliers_Lelli')

print("\n" + "="*85)
print("             BARYONIC TULLY-FISHER RELATION ROBUSTNESS ANALYSIS               ")
print("="*85)
print(summary[['Total_Galaxies', 'Outliers_McGaugh', 'Pct_McGaugh', 'Outliers_Lelli', 'Pct_Lelli']].to_string())
print("-"*85)
print(f"-> McGaugh (2012) Calibration  [log10(A) = 1.70]: Chi2 = {chi2_mcg:.2f}, p-value = {p_mcg:.4e}")
print(f"-> Lelli et al. (2016) Calibration [log10(A) = 1.67]: Chi2 = {chi2_lel:.2f}, p-value = {p_lel:.4e}")
print("=====================================================================================\n")