"""FastWoe.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1dGH4Z68kGfN3se1dMCcm2bmKViiT8-uW
"""

import warnings

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns
import statsmodels.api as sm
from scipy import stats
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import roc_auc_score
from sklearn.model_selection import train_test_split

warnings.filterwarnings("ignore")

# Set style
plt.style.use("default")
sns.set_palette("husl")


# Your FastWoe implementation (simplified for demonstration)
class FastWoeDemo:
    """Simplified FastWoe for demonstration"""

    def __init__(self):
        self.mappings_ = {}
        self.y_prior_ = None
        self.is_fitted_ = False

    def _calculate_woe_se(self, good_count, bad_count):
        """Calculate standard error of WOE"""
        if good_count <= 0 or bad_count <= 0:
            return np.inf
        return np.sqrt(1.0 / good_count + 1.0 / bad_count)

    def fit(self, X, y):
        """Fit the WOE encoder"""
        self.y_prior_ = y.mean()
        self.mappings_ = {}

        for col in X.columns:
            # Calculate WOE for each category
            mapping_data = []
            for category in X[col].unique():
                mask = X[col] == category
                good_count = (1 - y[mask]).sum()  # y=0 (good)
                bad_count = y[mask].sum()  # y=1 (bad)
                total_count = good_count + bad_count

                if total_count > 0:
                    event_rate = bad_count / total_count

                    # WOE calculation
                    if good_count > 0 and bad_count > 0:
                        good_rate = good_count / (len(y) - y.sum())  # P(category|y=0)
                        bad_rate = bad_count / y.sum()  # P(category|y=1)
                        woe = np.log(bad_rate / good_rate)
                    else:
                        woe = 0

                    woe_se = self._calculate_woe_se(good_count, bad_count)

                    mapping_data.append(
                        {
                            "category": category,
                            "count": total_count,
                            "good_count": good_count,
                            "bad_count": bad_count,
                            "event_rate": event_rate,
                            "woe": woe,
                            "woe_se": woe_se,
                        }
                    )

            self.mappings_[col] = pd.DataFrame(mapping_data).set_index("category")

        self.is_fitted_ = True
        return self

    def transform(self, X):
        """Transform to WOE values"""
        X_woe = X.copy()
        for col in X.columns:
            if col in self.mappings_:
                mapping = self.mappings_[col]
                X_woe[col] = X[col].map(mapping["woe"]).fillna(0)
        return X_woe

    def get_mapping(self, col):
        """Get mapping for a column"""
        return self.mappings_[col].reset_index()


def demonstrate_woe_logistic_equivalence():
    """Demonstrate the mathematical equivalence between WOE and logistic regression"""

    print("=" * 80)
    print("DEMONSTRATING WOE-LOGISTIC REGRESSION EQUIVALENCE")
    print("=" * 80)

    # Create example data from your document
    np.random.seed(42)

    # Example 1: Simple binary categorical
    print("\nExample 1: Binary Categorical Variable")
    print("-" * 50)

    # Create data based on contingency table: Color vs Outcome
    # Blue: 10 good, 30 bad
    # Red: 20 good, 10 bad
    data_rows = []
    data_rows.extend([{"Color": "Blue", "Target": 0}] * 10)  # good
    data_rows.extend([{"Color": "Blue", "Target": 1}] * 30)  # bad
    data_rows.extend([{"Color": "Red", "Target": 0}] * 20)  # good
    data_rows.extend([{"Color": "Red", "Target": 1}] * 10)  # bad

    df1 = pd.DataFrame(data_rows)
    X1 = df1[["Color"]]
    y1 = df1["Target"]

    # Fit WOE encoder
    woe_encoder = FastWoeDemo()
    woe_encoder.fit(X1, y1)

    # Display WOE mappings
    print("WOE Mappings:")
    color_mapping = woe_encoder.get_mapping("Color")
    print(color_mapping.round(4))

    # Fit logistic regression
    X1_dummy = pd.get_dummies(X1, drop_first=True)
    log_reg = LogisticRegression(fit_intercept=True, penalty=None, max_iter=1000)
    log_reg.fit(X1_dummy, y1)

    # Using statsmodels for detailed statistics
    X1_sm = sm.add_constant(X1_dummy.astype(float))
    logit_model = sm.Logit(y1, X1_sm)
    logit_result = logit_model.fit(disp=False)

    print("\nLogistic Regression Results:")
    print(
        f"Intercept (Red): {logit_result.params['const']:.4f} (SE: {logit_result.bse['const']:.4f})"
    )
    if "Color_Red" in logit_result.params:
        print(
            f"Coefficient (Red vs Blue): {logit_result.params['Color_Red']:.4f} (SE: {logit_result.bse['Color_Red']:.4f})"
        )

    # Calculate WOE difference
    woe_blue = color_mapping[color_mapping["category"] == "Blue"]["woe"].iloc[0]
    woe_red = color_mapping[color_mapping["category"] == "Red"]["woe"].iloc[0]
    woe_diff = woe_red - woe_blue

    # Calculate SE of WOE difference
    se_blue = color_mapping[color_mapping["category"] == "Blue"]["woe_se"].iloc[0]
    se_red = color_mapping[color_mapping["category"] == "Red"]["woe_se"].iloc[0]
    se_diff = np.sqrt(se_blue**2 + se_red**2)

    print(f"\nWOE Difference (Red - Blue): {woe_diff:.4f} (SE: {se_diff:.4f})")
    print(
        f"Coefficient matches WOE difference: {abs(woe_diff - logit_result.params.iloc[1]) < 0.001}"
    )

    return df1, woe_encoder, logit_result


def simulate_woe_variance_properties():
    """Simulate to show that WOE and log-odds have same variance"""

    print("\n" + "=" * 80)
    print("SIMULATING WOE VARIANCE PROPERTIES")
    print("=" * 80)

    np.random.seed(42)

    # True probabilities from your document
    p1 = 30 / (30 + 20)  # x=1 group
    p0 = 10 / (10 + 40)  # x=0 group
    prior_p = (30 + 10) / (30 + 10 + 20 + 40)

    # Sample sizes
    n1_total = 50
    n0_total = 50
    n_sim = 10000

    print(f"True probabilities: p1={p1:.3f}, p0={p0:.3f}, prior_p={prior_p:.3f}")
    print(f"Sample sizes: n1={n1_total}, n0={n0_total}")
    print(f"Simulations: {n_sim}")

    # Store results
    log_odds_1_samples = []
    woe_1_samples = []

    for _ in range(n_sim):
        # Simulate binary outcomes for x=1 group
        y1 = np.random.binomial(1, p1, size=n1_total)
        n11 = y1.sum()  # y=1, x=1
        n10 = n1_total - n11  # y=0, x=1

        # Compute log odds
        if n11 > 0 and n10 > 0:
            log_odds_1 = np.log(n11 / n10)
            log_odds_prior = np.log(prior_p / (1 - prior_p))
            woe_1 = log_odds_1 - log_odds_prior

            log_odds_1_samples.append(log_odds_1)
            woe_1_samples.append(woe_1)

    # Convert to arrays
    log_odds_1_samples = np.array(log_odds_1_samples)
    woe_1_samples = np.array(woe_1_samples)

    # Compare empirical variances
    variance_log_odds = np.var(log_odds_1_samples, ddof=1)
    variance_woe = np.var(woe_1_samples, ddof=1)

    print("\nEmpirical Results:")
    print(f"Var(Log Odds): {variance_log_odds:.6f}")
    print(f"Var(WOE): {variance_woe:.6f}")
    print(f"Difference: {abs(variance_log_odds - variance_woe):.6f}")
    print(f"Variances are equal: {abs(variance_log_odds - variance_woe) < 0.001}")

    # Theoretical calculation
    theoretical_var = 1 / 30 + 1 / 20  # 1/n_good + 1/n_bad for true values
    print(f"\nTheoretical SE: {np.sqrt(theoretical_var):.4f}")
    print(f"Empirical SE (Log Odds): {np.sqrt(variance_log_odds):.4f}")
    print(f"Empirical SE (WOE): {np.sqrt(variance_woe):.4f}")

    return log_odds_1_samples, woe_1_samples


def create_comprehensive_example():
    """Create a more comprehensive example with multiple categories"""

    print("\n" + "=" * 80)
    print("COMPREHENSIVE MULTI-CATEGORY EXAMPLE")
    print("=" * 80)

    np.random.seed(42)

    # Create synthetic dataset with multiple categorical features
    n_samples = 1000

    # Generate features
    regions = np.random.choice(
        ["North", "South", "East", "West"], n_samples, p=[0.3, 0.25, 0.25, 0.2]
    )
    job_types = np.random.choice(
        ["Tech", "Finance", "Healthcare", "Education"],
        n_samples,
        p=[0.4, 0.3, 0.2, 0.1],
    )

    # Create target with realistic relationships
    target_probs = np.zeros(n_samples)
    for i in range(n_samples):
        base_prob = 0.3

        # Regional effects
        if regions[i] == "North":
            base_prob += 0.1
        elif regions[i] == "South":
            base_prob -= 0.05

        # Job type effects
        if job_types[i] == "Tech":
            base_prob += 0.15
        elif job_types[i] == "Finance":
            base_prob += 0.05
        elif job_types[i] == "Healthcare":
            base_prob -= 0.1

        target_probs[i] = np.clip(base_prob, 0.05, 0.95)

    # Generate binary targets
    targets = np.random.binomial(1, target_probs)

    # Create DataFrame
    df = pd.DataFrame({"Region": regions, "JobType": job_types, "Target": targets})

    print(f"Dataset created: {len(df)} samples")
    print(f"Overall target rate: {df['Target'].mean():.3f}")
    print("\nTarget rates by category:")
    for col in ["Region", "JobType"]:
        rates = df.groupby(col)["Target"].agg(["count", "mean"]).round(3)
        print(f"\n{col}:")
        print(rates)

    # Split data
    X = df[["Region", "JobType"]]
    y = df["Target"]
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.3, random_state=42, stratify=y
    )

    # Fit WOE encoder
    woe_encoder = FastWoeDemo()
    woe_encoder.fit(X_train, y_train)

    # Transform data
    X_train_woe = woe_encoder.transform(X_train)
    X_test_woe = woe_encoder.transform(X_test)

    # Fit logistic regression on WOE-transformed data
    lr_woe = LogisticRegression(fit_intercept=True, penalty=None, max_iter=1000)
    lr_woe.fit(X_train_woe, y_train)

    # Fit logistic regression on dummy-encoded data
    X_train_dummy = pd.get_dummies(X_train, drop_first=True)
    X_test_dummy = pd.get_dummies(X_test, drop_first=True)
    lr_dummy = LogisticRegression(fit_intercept=True, penalty=None, max_iter=1000)
    lr_dummy.fit(X_train_dummy, y_train)

    # Predictions
    pred_woe = lr_woe.predict_proba(X_test_woe)[:, 1]
    pred_dummy = lr_dummy.predict_proba(X_test_dummy)[:, 1]

    # Calculate AUC
    auc_woe = roc_auc_score(y_test, pred_woe)
    auc_dummy = roc_auc_score(y_test, pred_dummy)

    print("\nModel Performance:")
    print(f"WOE-based model AUC: {auc_woe:.4f}")
    print(f"Dummy-encoded model AUC: {auc_dummy:.4f}")
    print(f"Performance difference: {abs(auc_woe - auc_dummy):.6f}")

    # Display WOE mappings with confidence intervals
    print("\nWOE Mappings with Statistical Significance:")
    for col in ["Region", "JobType"]:
        mapping = woe_encoder.get_mapping(col)
        print(f"\n{col}:")

        # Add statistical significance tests
        mapping["z_score"] = mapping["woe"] / mapping["woe_se"]
        mapping["p_value"] = 2 * (1 - stats.norm.cdf(abs(mapping["z_score"])))
        mapping["significant"] = mapping["p_value"] < 0.05

        display_cols = [
            "category",
            "count",
            "event_rate",
            "woe",
            "woe_se",
            "z_score",
            "p_value",
            "significant",
        ]
        print(mapping[display_cols].round(4))

    return df, woe_encoder, X_train, X_test, y_train, y_test


def visualize_woe_analysis(df, woe_encoder):
    """Create comprehensive visualizations"""

    print("\n" + "=" * 80)
    print("CREATING VISUALIZATIONS")
    print("=" * 80)

    fig, axes = plt.subplots(2, 3, figsize=(18, 12))

    # Plot 1: WOE values by Region
    region_mapping = woe_encoder.get_mapping("Region")
    ax1 = axes[0, 0]
    bars1 = ax1.bar(
        region_mapping["category"],
        region_mapping["woe"],
        color="skyblue",
        alpha=0.7,
        edgecolor="black",
    )
    ax1.axhline(y=0, color="red", linestyle="--", alpha=0.7)
    ax1.set_title("WOE Values by Region", fontweight="bold", fontsize=12)
    ax1.set_ylabel("Weight of Evidence")
    ax1.tick_params(axis="x", rotation=45)

    # Add error bars
    ax1.errorbar(
        region_mapping["category"],
        region_mapping["woe"],
        yerr=1.96 * region_mapping["woe_se"],
        fmt="none",
        color="red",
        capsize=5,
    )

    # Add value labels
    for bar, val in zip(bars1, region_mapping["woe"]):
        height = bar.get_height()
        ax1.text(
            bar.get_x() + bar.get_width() / 2.0,
            height + (0.05 if height >= 0 else -0.1),
            f"{val:.3f}",
            ha="center",
            va="bottom" if height >= 0 else "top",
            fontweight="bold",
        )

    # Plot 2: WOE values by JobType
    job_mapping = woe_encoder.get_mapping("JobType")
    ax2 = axes[0, 1]
    bars2 = ax2.bar(
        job_mapping["category"],
        job_mapping["woe"],
        color="lightgreen",
        alpha=0.7,
        edgecolor="black",
    )
    ax2.axhline(y=0, color="red", linestyle="--", alpha=0.7)
    ax2.set_title("WOE Values by Job Type", fontweight="bold", fontsize=12)
    ax2.set_ylabel("Weight of Evidence")
    ax2.tick_params(axis="x", rotation=45)

    # Add error bars
    ax2.errorbar(
        job_mapping["category"],
        job_mapping["woe"],
        yerr=1.96 * job_mapping["woe_se"],
        fmt="none",
        color="red",
        capsize=5,
    )

    # Add value labels
    for bar, val in zip(bars2, job_mapping["woe"]):
        height = bar.get_height()
        ax2.text(
            bar.get_x() + bar.get_width() / 2.0,
            height + (0.05 if height >= 0 else -0.1),
            f"{val:.3f}",
            ha="center",
            va="bottom" if height >= 0 else "top",
            fontweight="bold",
        )

    # Plot 3: Event rates by categories
    ax3 = axes[0, 2]
    rates_region = df.groupby("Region")["Target"].mean()
    bars3 = ax3.bar(
        rates_region.index,
        rates_region.values,
        color="orange",
        alpha=0.7,
        edgecolor="black",
    )
    overall_rate = df["Target"].mean()
    ax3.axhline(
        y=overall_rate,
        color="red",
        linestyle="--",
        alpha=0.7,
        label=f"Overall: {overall_rate:.3f}",
    )
    ax3.set_title("Event Rates by Region", fontweight="bold", fontsize=12)
    ax3.set_ylabel("Event Rate")
    ax3.legend()
    ax3.tick_params(axis="x", rotation=45)

    # Add value labels
    for bar, val in zip(bars3, rates_region.values):
        height = bar.get_height()
        ax3.text(
            bar.get_x() + bar.get_width() / 2.0,
            height + 0.01,
            f"{val:.3f}",
            ha="center",
            va="bottom",
            fontweight="bold",
        )

    # Plot 4: WOE vs Event Rate correlation
    ax4 = axes[1, 0]
    all_woe = np.concatenate([region_mapping["woe"], job_mapping["woe"]])
    all_rates = np.concatenate(
        [region_mapping["event_rate"], job_mapping["event_rate"]]
    )

    ax4.scatter(
        all_rates, all_woe, alpha=0.7, s=100, color="purple", edgecolors="black"
    )
    ax4.set_xlabel("Event Rate")
    ax4.set_ylabel("WOE Value")
    ax4.set_title("WOE vs Event Rate", fontweight="bold", fontsize=12)
    ax4.grid(True, alpha=0.3)

    # Add correlation coefficient
    corr = np.corrcoef(all_rates, all_woe)[0, 1]
    ax4.text(
        0.05,
        0.95,
        f"Correlation: {corr:.3f}",
        transform=ax4.transAxes,
        bbox=dict(boxstyle="round", facecolor="wheat", alpha=0.8),
    )

    # Plot 5: Sample sizes and WOE standard errors
    ax5 = axes[1, 1]
    all_counts = np.concatenate([region_mapping["count"], job_mapping["count"]])
    all_se = np.concatenate([region_mapping["woe_se"], job_mapping["woe_se"]])

    ax5.scatter(all_counts, all_se, alpha=0.7, s=100, color="red", edgecolors="black")
    ax5.set_xlabel("Sample Size")
    ax5.set_ylabel("WOE Standard Error")
    ax5.set_title("Sample Size vs WOE Standard Error", fontweight="bold", fontsize=12)
    ax5.grid(True, alpha=0.3)

    # Theoretical relationship
    theoretical_x = np.linspace(min(all_counts), max(all_counts), 100)
    # Approximate theoretical SE (assuming balanced good/bad split)
    theoretical_y = np.sqrt(2 / theoretical_x)
    ax5.plot(
        theoretical_x, theoretical_y, "r--", alpha=0.7, label="Theoretical (balanced)"
    )
    ax5.legend()

    # Plot 6: Z-scores (statistical significance)
    ax6 = axes[1, 2]
    region_mapping["z_score"] = region_mapping["woe"] / region_mapping["woe_se"]
    job_mapping["z_score"] = job_mapping["woe"] / job_mapping["woe_se"]

    all_categories = list(region_mapping["category"]) + list(job_mapping["category"])
    all_z_scores = np.concatenate([region_mapping["z_score"], job_mapping["z_score"]])
    colors = ["blue"] * len(region_mapping) + ["green"] * len(job_mapping)

    bars6 = ax6.bar(
        range(len(all_categories)),
        all_z_scores,
        color=colors,
        alpha=0.7,
        edgecolor="black",
    )
    ax6.axhline(
        y=1.96, color="red", linestyle="--", alpha=0.7, label="95% significance"
    )
    ax6.axhline(y=-1.96, color="red", linestyle="--", alpha=0.7)
    ax6.set_title("Statistical Significance (Z-scores)", fontweight="bold", fontsize=12)
    ax6.set_ylabel("Z-score")
    ax6.set_xticks(range(len(all_categories)))
    ax6.set_xticklabels(all_categories, rotation=45)
    ax6.legend()

    # Add significance markers
    for bar, z_score in zip(bars6, all_z_scores):
        if abs(z_score) > 1.96:
            ax6.text(
                bar.get_x() + bar.get_width() / 2.0,
                z_score + (0.1 if z_score > 0 else -0.2),
                "*",
                ha="center",
                va="center",
                fontweight="bold",
                fontsize=16,
                color="red",
            )

    plt.tight_layout()
    plt.show()


# Run the complete analysis
if __name__ == "__main__":
    # Demonstrate basic equivalence
    df1, woe_encoder1, logit_result1 = demonstrate_woe_logistic_equivalence()

    # Simulate variance properties
    log_odds_samples, woe_samples = simulate_woe_variance_properties()

    # Comprehensive example
    df2, woe_encoder2, X_train, X_test, y_train, y_test = create_comprehensive_example()

    # Create visualizations
    visualize_woe_analysis(df2, woe_encoder2)

    print("\n" + "=" * 80)
    print("KEY INSIGHTS SUMMARY")
    print("=" * 80)

    print("""
    1. MATHEMATICAL EQUIVALENCE:
       - WOE differences = Logistic regression coefficients
       - Both measure the same log-odds relationships
       - Standard errors are identical between approaches

    2. VARIANCE PROPERTIES:
       - WOE and log-odds have identical variance
       - Centering (subtracting prior log-odds) doesn't affect variance
       - SE(WOE) = sqrt(1/n_good + 1/n_bad)

    3. PRACTICAL IMPLICATIONS:
       - WOE provides intuitive interpretation of categorical variables
       - Can be used for feature engineering before logistic regression
       - Confidence intervals enable statistical significance testing
       - Identifies categories with insufficient data (high SE)

    4. STATISTICAL SIGNIFICANCE:
       - Z-scores (WOE/SE) indicate statistical significance
       - Categories with |z| > 1.96 are significant at 95% level
       - Helps identify truly predictive categories vs noise

    5. MODEL PERFORMANCE:
       - WOE-based and dummy-encoded models yield identical performance
       - WOE transformation preserves all predictive information
       - Provides more interpretable and stable features
    """)
