import pandas as pd
import numpy as np # Need numpy for calculations

# --- Configuration ---
csv_filename = 'precip_data.csv'  # <--- CHANGE THIS TO YOUR CSV FILENAME
date_column = 'Date'
site1_precip_col = 'Parleys_IN' # Column name for Parley's precipitation
site2_precip_col = 'SLC_IN'     # Column name for SLC precipitation
precip_threshold = 0.0          # Value to determine 'wet' (e.g., > 0.0)
# --- End Configuration ---

print(f"Reading data from: {csv_filename}")

try:
    # Read the CSV file, parse dates, and set Date as index
    df = pd.read_csv(csv_filename, parse_dates=[date_column], index_col=date_column)
    print(f"Successfully read {len(df)} rows.")

    # --- Data Preparation ---
    # Ensure precipitation columns are numeric, coercing errors to NaN
    df[site1_precip_col] = pd.to_numeric(df[site1_precip_col], errors='coerce')
    df[site2_precip_col] = pd.to_numeric(df[site2_precip_col], errors='coerce')

    # Drop rows where either precipitation value is missing after coercion
    initial_rows = len(df)
    df.dropna(subset=[site1_precip_col, site2_precip_col], inplace=True)
    dropped_rows = initial_rows - len(df)
    if dropped_rows > 0:
        print(f"Dropped {dropped_rows} rows due to missing/non-numeric precipitation data.")

    # Create binary Wet/Dry columns (1 if Precip > threshold, 0 otherwise)
    df['Site1_Wet'] = (df[site1_precip_col] > precip_threshold).astype(int)
    df['Site2_Wet'] = (df[site2_precip_col] > precip_threshold).astype(int)
    print("Created binary Wet/Dry columns.")

    # --- Contingency Table Calculation ---
    # Calculate counts for the 2x2 table
    A = len(df[(df['Site1_Wet'] == 1) & (df['Site2_Wet'] == 1)])  # Both Wet
    B = len(df[(df['Site1_Wet'] == 1) & (df['Site2_Wet'] == 0)])  # Site 1 Wet, Site 2 Dry
    C = len(df[(df['Site1_Wet'] == 0) & (df['Site2_Wet'] == 1)])  # Site 1 Dry, Site 2 Wet
    D = len(df[(df['Site1_Wet'] == 0) & (df['Site2_Wet'] == 0)])  # Both Dry
    N = A + B + C + D

    print("\nContingency Table:")
    print(f"                     Site 2 Wet | Site 2 Dry")
    print(f"---------------------------------------------")
    print(f" Site 1 Wet        | {A:10d} | {B:10d}")
    print(f" Site 1 Dry        | {C:10d} | {D:10d}")
    print(f"---------------------------------------------")
    print(f" Total Days (N): {N}")

    # Basic check
    if N != len(df):
        print("\nWarning: Sum of contingency table cells does not match dataframe length. Check data.")

    # Calculate Phi Coefficient using pandas correlation on binary columns (for comparison)
    phi_coefficient = df['Site1_Wet'].corr(df['Site2_Wet'], method='pearson')
    print(f"\nPhi Coefficient (Pearson on 0/1 data): {phi_coefficient:.4f}")


    # --- Approximate Tetrachoric Correlation Calculation ---
    print("\nCalculating Approximate Tetrachoric Correlation using Odds Ratio...")

    # Calculate Odds Ratio (handle potential division by zero)
    if B == 0 or C == 0:
        print("\nWarning: Cannot calculate Odds Ratio because B or C count is zero.")
        approx_tetra_corr_value = np.nan # Assign NaN if calculation not possible
    else:
        odds_ratio = (A * D) / (B * C)
        print(f"Odds Ratio (AD/BC): {odds_ratio:.4f}")
        # Calculate approximation based on Digby (1983) / Edwards (1957) form
        # r_tetra ≈ sin( (pi/2) * (sqrt(OR) - 1) / (sqrt(OR) + 1) )
        approx_tetra_corr_value = np.sin(
            (np.pi / 2) * (np.sqrt(odds_ratio) - 1) / (np.sqrt(odds_ratio) + 1)
            )

    print(f"\n======================================================")
    print(f" Approximate Tetrachoric Correlation: {approx_tetra_corr_value:.4f}")
    print(f"======================================================")
    print("\nNote: This is an approximation based on the Odds Ratio.")

except FileNotFoundError:
    print(f"\nError: Could not find the file '{csv_filename}'.")
    print("Please make sure the file exists and is in the same directory as the script,")
    print("or provide the full path.")
except KeyError as e:
    print(f"\nError: Column not found in CSV: {e}")
    print(f"Please ensure your CSV has columns named '{date_column}', '{site1_precip_col}', and '{site2_precip_col}'.")
except Exception as e:
    print(f"\nAn unexpected error occurred: {e}")