import pandas as pd
from diff_diff import DifferenceInDifferences

# Create sample data
data = pd.DataFrame({
    'outcome': [10, 11, 15, 18, 9, 10, 12, 13],
    'treated': [1, 1, 1, 1, 0, 0, 0, 0],
    'post': [0, 0, 1, 1, 0, 0, 1, 1]
})

# Fit the model
did = DifferenceInDifferences()
results = did.fit(data, outcome='outcome', treatment='treated', time='post')

# View results
print(results)  # DiDResults(ATT=3.0000, SE=1.7321, p=0.1583)
results.print_summary()

from diff_diff import generate_did_data, DifferenceInDifferences

# Generate panel data with 100 units, 4 periods, and a treatment effect of 5
data = generate_did_data(
    n_units=100,
    n_periods=4,
    treatment_effect=5.0,
    treatment_fraction=0.5,  # 50% of units are treated
    treatment_period=2,       # Treatment starts at period 2
    seed=42
)

# Verify the estimator recovers the treatment effect
did = DifferenceInDifferences()
results = did.fit(data, outcome='outcome', treatment='treated', time='post')
print(f"Estimated ATT: {results.att:.2f} (true: 5.0)")

from diff_diff import make_treatment_indicator

# From categorical variable
df = make_treatment_indicator(
    data,
    column='state',
    treated_values=['CA', 'NY', 'TX']  # These states are treated
)

# From numeric threshold (e.g., firms above median size)
df = make_treatment_indicator(
    data,
    column='firm_size',
    threshold=data['firm_size'].median()
)

# Treat units below threshold
df = make_treatment_indicator(
    data,
    column='income',
    threshold=50000,
    above_threshold=False  # Units with income <= 50000 are treated
)

from diff_diff import make_post_indicator

# From specific post-treatment periods
df = make_post_indicator(
    data,
    time_column='year',
    post_periods=[2020, 2021, 2022]
)

# From treatment start date
df = make_post_indicator(
    data,
    time_column='year',
    treatment_start=2020  # All years >= 2020 are post-treatment
)

# Works with datetime columns
df = make_post_indicator(
    data,
    time_column='date',
    treatment_start='2020-01-01'
)

from diff_diff import wide_to_long

# Wide format: columns like sales_2019, sales_2020, sales_2021
wide_df = pd.DataFrame({
    'firm_id': [1, 2, 3],
    'industry': ['tech', 'retail', 'tech'],
    'sales_2019': [100, 150, 200],
    'sales_2020': [110, 160, 210],
    'sales_2021': [120, 170, 220]
})

# Convert to long format for DiD
long_df = wide_to_long(
    wide_df,
    value_columns=['sales_2019', 'sales_2020', 'sales_2021'],
    id_column='firm_id',
    time_name='year',
    value_name='sales',
    time_values=[2019, 2020, 2021]
)
# Result: 9 rows (3 firms × 3 years), columns: firm_id, year, sales, industry

from diff_diff import balance_panel

# Keep only units with complete data (drop incomplete units)
balanced = balance_panel(
    data,
    unit_column='firm_id',
    time_column='year',
    method='inner'
)

# Include all unit-period combinations (creates NaN for missing)
balanced = balance_panel(
    data,
    unit_column='firm_id',
    time_column='year',
    method='outer'
)

# Fill missing values
balanced = balance_panel(
    data,
    unit_column='firm_id',
    time_column='year',
    method='fill',
    fill_value=0  # Or None for forward/backward fill
)

from diff_diff import validate_did_data

# Validate and get informative error messages
result = validate_did_data(
    data,
    outcome='sales',
    treatment='treated',
    time='post',
    unit='firm_id',      # Optional: for panel-specific validation
    raise_on_error=False  # Return dict instead of raising
)

if result['valid']:
    print("Data is ready for DiD analysis!")
    print(f"Summary: {result['summary']}")
else:
    print("Issues found:")
    for error in result['errors']:
        print(f"  - {error}")

for warning in result['warnings']:
    print(f"Warning: {warning}")

from diff_diff import summarize_did_data

summary = summarize_did_data(
    data,
    outcome='sales',
    treatment='treated',
    time='post'
)
print(summary)

from diff_diff import create_event_time

# Add event-time column relative to treatment timing
df = create_event_time(
    data,
    time_column='year',
    treatment_time_column='treatment_year'
)
# Result: event_time = -2, -1, 0, 1, 2 relative to treatment

from diff_diff import aggregate_to_cohorts

cohort_data = aggregate_to_cohorts(
    data,
    unit_column='firm_id',
    time_column='year',
    treatment_column='treated',
    outcome='sales'
)
# Result: mean outcome by treatment group and period

from diff_diff import rank_control_units, generate_did_data

# Generate sample data
data = generate_did_data(n_units=50, n_periods=6, seed=42)

# Rank control units by their similarity to treated units
ranking = rank_control_units(
    data,
    unit_column='unit',
    time_column='period',
    outcome_column='outcome',
    treatment_column='treated',
    n_top=10  # Return top 10 controls
)

print(ranking[['unit', 'quality_score', 'pre_trend_rmse']])

# Add covariate-based matching
ranking = rank_control_units(
    data,
    unit_column='unit',
    time_column='period',
    outcome_column='outcome',
    treatment_column='treated',
    covariates=['size', 'age'],  # Match on these too
    outcome_weight=0.7,          # 70% weight on outcome trends
    covariate_weight=0.3         # 30% weight on covariate similarity
)

from diff_diff import SyntheticDiD

# Get top control units
top_controls = ranking['unit'].tolist()

# Filter data to treated + top controls
filtered_data = data[
    (data['treated'] == 1) | (data['unit'].isin(top_controls))
]

# Fit SyntheticDiD with selected controls
sdid = SyntheticDiD()
results = sdid.fit(
    filtered_data,
    outcome='outcome',
    treatment='treated',
    unit='unit',
    time='period',
    post_periods=[3, 4, 5]
)

from diff_diff import DifferenceInDifferences

did = DifferenceInDifferences(robust=True, alpha=0.05)
results = did.fit(
    data,
    outcome='sales',
    treatment='treated',
    time='post_policy'
)

# Access results
print(f"ATT: {results.att:.4f}")
print(f"Standard Error: {results.se:.4f}")
print(f"P-value: {results.p_value:.4f}")
print(f"95% CI: {results.conf_int}")
print(f"Significant: {results.is_significant}")

# R-style formula syntax
results = did.fit(data, formula='outcome ~ treated * post')

# Explicit interaction syntax
results = did.fit(data, formula='outcome ~ treated + post + treated:post')

# With covariates
results = did.fit(data, formula='outcome ~ treated * post + age + income')

results = did.fit(
    data,
    outcome='outcome',
    treatment='treated',
    time='post',
    covariates=['age', 'income', 'education']
)

# State and industry fixed effects
results = did.fit(
    data,
    outcome='sales',
    treatment='treated',
    time='post',
    fixed_effects=['state', 'industry']
)

# Access fixed effect coefficients
state_coefs = {k: v for k, v in results.coefficients.items() if k.startswith('state_')}

# Absorb firm-level fixed effects (efficient for many firms)
results = did.fit(
    data,
    outcome='sales',
    treatment='treated',
    time='post',
    absorb=['firm_id']
)

results = did.fit(
    data,
    outcome='sales',
    treatment='treated',
    time='post',
    covariates=['size', 'age'],           # Linear controls
    fixed_effects=['industry'],            # Low-dimensional FE (dummies)
    absorb=['firm_id']                     # High-dimensional FE (absorbed)
)

did = DifferenceInDifferences(cluster='state')
results = did.fit(
    data,
    outcome='outcome',
    treatment='treated',
    time='post'
)

# Use wild bootstrap for inference
did = DifferenceInDifferences(
    cluster='state',
    inference='wild_bootstrap',
    n_bootstrap=999,
    bootstrap_weights='rademacher',  # or 'webb' for <10 clusters, 'mammen'
    seed=42
)
results = did.fit(data, outcome='y', treatment='treated', time='post')

# Results include bootstrap-based SE and p-value
print(f"ATT: {results.att:.3f} (SE: {results.se:.3f})")
print(f"P-value: {results.p_value:.4f}")
print(f"95% CI: {results.conf_int}")
print(f"Inference method: {results.inference_method}")
print(f"Number of clusters: {results.n_clusters}")

from diff_diff import TwoWayFixedEffects

twfe = TwoWayFixedEffects()
results = twfe.fit(
    panel_data,
    outcome='outcome',
    treatment='treated',
    time='year',
    unit='firm_id'
)

from diff_diff import MultiPeriodDiD

# Fit with multiple time periods
did = MultiPeriodDiD()
results = did.fit(
    panel_data,
    outcome='sales',
    treatment='treated',
    time='period',
    post_periods=[3, 4, 5],      # Periods 3-5 are post-treatment
    reference_period=0           # Reference period for comparison
)

# View period-specific treatment effects
for period, effect in results.period_effects.items():
    print(f"Period {period}: {effect.effect:.3f} (SE: {effect.se:.3f})")

# View average treatment effect across post-periods
print(f"Average ATT: {results.avg_att:.3f}")
print(f"Average SE: {results.avg_se:.3f}")

# Full summary with all period effects
results.print_summary()

from diff_diff import CallawaySantAnna

# Panel data with staggered treatment
# 'first_treat' = period when unit was first treated (0 if never treated)
cs = CallawaySantAnna()
results = cs.fit(
    panel_data,
    outcome='sales',
    unit='firm_id',
    time='year',
    first_treat='first_treat',  # 0 for never-treated, else first treatment year
    aggregate='event_study'      # Compute event study effects
)

# View results
results.print_summary()

# Access group-time effects ATT(g,t)
for (group, time), effect in results.group_time_effects.items():
    print(f"Cohort {group}, Period {time}: {effect['effect']:.3f}")

# Event study effects (averaged by relative time)
for rel_time, effect in results.event_study_effects.items():
    print(f"e={rel_time}: {effect['effect']:.3f} (SE: {effect['se']:.3f})")

# Convert to DataFrame
df = results.to_dataframe(level='event_study')

CallawaySantAnna(
    control_group='never_treated',  # or 'not_yet_treated'
    anticipation=0,                  # Periods before treatment with effects
    estimation_method='dr',          # 'dr', 'ipw', or 'reg'
    alpha=0.05,                      # Significance level
    cluster=None,                    # Column for cluster SEs
    n_bootstrap=0,                   # Bootstrap iterations (0 = analytical SEs)
    bootstrap_weight_type='rademacher',  # 'rademacher', 'mammen', or 'webb'
    seed=None                        # Random seed
)

# Bootstrap inference with 999 iterations
cs = CallawaySantAnna(
    n_bootstrap=999,
    bootstrap_weight_type='rademacher',  # or 'mammen', 'webb'
    seed=42
)
results = cs.fit(
    data,
    outcome='sales',
    unit='firm_id',
    time='year',
    first_treat='first_treat',
    aggregate='event_study'
)

# Access bootstrap results
print(f"Overall ATT: {results.overall_att:.3f}")
print(f"Bootstrap SE: {results.bootstrap_results.overall_att_se:.3f}")
print(f"Bootstrap 95% CI: {results.bootstrap_results.overall_att_ci}")
print(f"Bootstrap p-value: {results.bootstrap_results.overall_att_p_value:.4f}")

# Event study bootstrap inference
for rel_time, se in results.bootstrap_results.event_study_ses.items():
    ci = results.bootstrap_results.event_study_cis[rel_time]
    print(f"e={rel_time}: SE={se:.3f}, 95% CI=[{ci[0]:.3f}, {ci[1]:.3f}]")

# Doubly robust estimation with covariates
cs = CallawaySantAnna(estimation_method='dr')  # 'dr', 'ipw', or 'reg'
results = cs.fit(
    data,
    outcome='sales',
    unit='firm_id',
    time='year',
    first_treat='first_treat',
    covariates=['size', 'age', 'industry'],  # Covariates for conditional PT
    aggregate='event_study'
)

from diff_diff import SunAbraham

# Basic usage
sa = SunAbraham()
results = sa.fit(
    panel_data,
    outcome='sales',
    unit='firm_id',
    time='year',
    first_treat='first_treat'  # 0 for never-treated, else first treatment year
)

# View results
results.print_summary()

# Event study effects (by relative time to treatment)
for rel_time, effect in results.event_study_effects.items():
    print(f"e={rel_time}: {effect['effect']:.3f} (SE: {effect['se']:.3f})")

# Overall ATT
print(f"Overall ATT: {results.overall_att:.3f} (SE: {results.overall_se:.3f})")

# Cohort weights (how each cohort contributes to each event-time estimate)
for rel_time, weights in results.cohort_weights.items():
    print(f"e={rel_time}: {weights}")

SunAbraham(
    control_group='never_treated',  # or 'not_yet_treated'
    anticipation=0,                  # Periods before treatment with effects
    alpha=0.05,                      # Significance level
    cluster=None,                    # Column for cluster SEs
    n_bootstrap=0,                   # Bootstrap iterations (0 = analytical SEs)
    bootstrap_weights='rademacher',  # 'rademacher', 'mammen', or 'webb'
    seed=None                        # Random seed
)

# Bootstrap inference with 999 iterations
sa = SunAbraham(
    n_bootstrap=999,
    bootstrap_weights='rademacher',
    seed=42
)
results = sa.fit(
    data,
    outcome='sales',
    unit='firm_id',
    time='year',
    first_treat='first_treat'
)

# Access bootstrap results
print(f"Overall ATT: {results.overall_att:.3f}")
print(f"Bootstrap SE: {results.bootstrap_results.overall_att_se:.3f}")
print(f"Bootstrap 95% CI: {results.bootstrap_results.overall_att_ci}")
print(f"Bootstrap p-value: {results.bootstrap_results.overall_att_p_value:.4f}")

from diff_diff import CallawaySantAnna, SunAbraham

# Callaway-Sant'Anna
cs = CallawaySantAnna()
cs_results = cs.fit(data, outcome='y', unit='unit', time='time', first_treat='first_treat')

# Sun-Abraham
sa = SunAbraham()
sa_results = sa.fit(data, outcome='y', unit='unit', time='time', first_treat='first_treat')

# Compare
print(f"Callaway-Sant'Anna ATT: {cs_results.overall_att:.3f}")
print(f"Sun-Abraham ATT: {sa_results.overall_att:.3f}")

# If results differ substantially, investigate heterogeneity

from diff_diff import TripleDifference, triple_difference

# Basic usage
ddd = TripleDifference(estimation_method='dr')  # doubly robust (recommended)
results = ddd.fit(
    data,
    outcome='wages',
    group='policy_state',       # 1=state enacted policy, 0=control state
    partition='female',         # 1=women (affected by policy), 0=men
    time='post'                 # 1=post-policy, 0=pre-policy
)

# View results
results.print_summary()
print(f"ATT: {results.att:.3f} (SE: {results.se:.3f})")

# With covariates (properly incorporated, unlike naive DDD)
results = ddd.fit(
    data,
    outcome='wages',
    group='policy_state',
    partition='female',
    time='post',
    covariates=['age', 'education', 'experience']
)

# Compare estimation methods
for method in ['reg', 'ipw', 'dr']:
    est = TripleDifference(estimation_method=method)
    res = est.fit(data, outcome='y', group='g', partition='p', time='t')
    print(f"{method}: ATT={res.att:.3f} (SE={res.se:.3f})")

# One-liner estimation
results = triple_difference(
    data,
    outcome='wages',
    group='policy_state',
    partition='female',
    time='post',
    covariates=['age', 'education'],
    estimation_method='dr'
)

from diff_diff import plot_event_study, MultiPeriodDiD, CallawaySantAnna, SunAbraham

# From MultiPeriodDiD
did = MultiPeriodDiD()
results = did.fit(data, outcome='y', treatment='treated',
                  time='period', post_periods=[3, 4, 5])
plot_event_study(results, title="Treatment Effects Over Time")

# From CallawaySantAnna (with event study aggregation)
cs = CallawaySantAnna()
results = cs.fit(data, outcome='y', unit='unit', time='period',
                 first_treat='first_treat', aggregate='event_study')
plot_event_study(results, title="Staggered DiD Event Study (CS)")

# From SunAbraham
sa = SunAbraham()
results = sa.fit(data, outcome='y', unit='unit', time='period',
                 first_treat='first_treat')
plot_event_study(results, title="Staggered DiD Event Study (SA)")

# From a DataFrame
df = pd.DataFrame({
    'period': [-2, -1, 0, 1, 2],
    'effect': [0.1, 0.05, 0.0, 2.5, 2.8],
    'se': [0.3, 0.25, 0.0, 0.4, 0.45]
})
plot_event_study(df, reference_period=0)

# With customization
ax = plot_event_study(
    results,
    title="Dynamic Treatment Effects",
    xlabel="Years Relative to Treatment",
    ylabel="Effect on Sales ($1000s)",
    color="#2563eb",
    marker="o",
    shade_pre=True,           # Shade pre-treatment region
    show_zero_line=True,      # Horizontal line at y=0
    show_reference_line=True, # Vertical line at reference period
    figsize=(10, 6),
    show=False                # Don't call plt.show(), return axes
)

from diff_diff import SyntheticDiD

# Fit Synthetic DiD model
sdid = SyntheticDiD()
results = sdid.fit(
    panel_data,
    outcome='gdp_growth',
    treatment='treated',
    unit='state',
    time='year',
    post_periods=[2015, 2016, 2017, 2018]
)

# View results
results.print_summary()
print(f"ATT: {results.att:.3f} (SE: {results.se:.3f})")

# Examine unit weights (which control units matter most)
weights_df = results.get_unit_weights_df()
print(weights_df.head(10))

# Examine time weights
time_weights_df = results.get_time_weights_df()
print(time_weights_df)

SyntheticDiD(
    lambda_reg=0.0,     # Regularization toward uniform weights (0 = no reg)
    zeta=1.0,           # Time weight regularization (higher = more uniform)
    alpha=0.05,         # Significance level
    n_bootstrap=200,    # Bootstrap iterations for SE (0 = placebo-based)
    seed=None           # Random seed for reproducibility
)

# As dictionary
results.to_dict()
# {'att': 3.5, 'se': 1.26, 'p_value': 0.037, ...}

# As DataFrame
df = results.to_dataframe()

if results.is_significant:
    print(f"Effect is significant at {did.alpha} level")

# Get significance stars
print(f"ATT: {results.att}{results.significance_stars}")
# ATT: 3.5000*

# All coefficients
results.coefficients
# {'const': 9.5, 'treated': 1.0, 'post': 2.5, 'treated:post': 3.5}

# Variance-covariance matrix
results.vcov

# Residuals and fitted values
results.residuals
results.fitted_values

# R-squared
results.r_squared

from diff_diff.utils import check_parallel_trends

trends = check_parallel_trends(
    data,
    outcome='outcome',
    time='period',
    treatment_group='treated'
)

print(f"Treated trend: {trends['treated_trend']:.4f}")
print(f"Control trend: {trends['control_trend']:.4f}")
print(f"Difference p-value: {trends['p_value']:.4f}")

from diff_diff.utils import check_parallel_trends_robust

results = check_parallel_trends_robust(
    data,
    outcome='outcome',
    time='period',
    treatment_group='treated',
    unit='firm_id',              # Unit identifier for panel data
    pre_periods=[2018, 2019],    # Pre-treatment periods
    n_permutations=1000          # Permutations for p-value
)

print(f"Wasserstein distance: {results['wasserstein_distance']:.4f}")
print(f"Wasserstein p-value: {results['wasserstein_p_value']:.4f}")
print(f"KS test p-value: {results['ks_p_value']:.4f}")
print(f"Parallel trends plausible: {results['parallel_trends_plausible']}")

from diff_diff.utils import equivalence_test_trends

results = equivalence_test_trends(
    data,
    outcome='outcome',
    time='period',
    treatment_group='treated',
    unit='firm_id',
    equivalence_margin=0.5       # Define "practically equivalent"
)

print(f"Mean difference: {results['mean_difference']:.4f}")
print(f"TOST p-value: {results['tost_p_value']:.4f}")
print(f"Trends equivalent: {results['equivalent']}")

from diff_diff import HonestDiD, MultiPeriodDiD

# First, fit a standard event study
did = MultiPeriodDiD()
event_results = did.fit(
    data,
    outcome='outcome',
    treatment='treated',
    time='period',
    post_periods=[5, 6, 7, 8, 9]
)

# Compute honest bounds with relative magnitudes restriction
# M=1 means post-treatment violations can be up to 1x the worst pre-treatment violation
honest = HonestDiD(method='relative_magnitude', M=1.0)
honest_results = honest.fit(event_results)

print(honest_results.summary())
print(f"Original estimate: {honest_results.original_estimate:.4f}")
print(f"Robust 95% CI: [{honest_results.ci_lb:.4f}, {honest_results.ci_ub:.4f}]")
print(f"Effect robust to violations: {honest_results.is_significant}")

# How do results change as we allow larger violations?
sensitivity = honest.sensitivity_analysis(
    event_results,
    M_grid=[0, 0.5, 1.0, 1.5, 2.0]
)

print(sensitivity.summary())
print(f"Breakdown value: M = {sensitivity.breakdown_M}")
# Breakdown = smallest M where the robust CI includes zero

breakdown = honest.breakdown_value(event_results)
if breakdown >= 1.0:
    print("Result holds even if post-treatment violations are as bad as pre-treatment")
else:
    print(f"Result requires violations smaller than {breakdown:.1f}x pre-treatment")

# Bounds second differences of trend violations
# M=0 means linear extrapolation of pre-trends
honest_smooth = HonestDiD(method='smoothness', M=0.5)
smooth_results = honest_smooth.fit(event_results)

from diff_diff import plot_sensitivity, plot_honest_event_study

# Plot sensitivity analysis
plot_sensitivity(sensitivity, title="Sensitivity to Parallel Trends Violations")

# Event study with honest confidence intervals
plot_honest_event_study(event_results, honest_results)

from diff_diff import PreTrendsPower, MultiPeriodDiD

# First, fit an event study
did = MultiPeriodDiD()
event_results = did.fit(
    data,
    outcome='outcome',
    treatment='treated',
    time='period',
    post_periods=[5, 6, 7, 8, 9]
)

# Analyze pre-trends test power
pt = PreTrendsPower(alpha=0.05, power=0.80)
power_results = pt.fit(event_results)

print(power_results.summary())
print(f"Minimum Detectable Violation (MDV): {power_results.mdv:.4f}")
print(f"Power to detect violations of size MDV: {power_results.power:.1%}")

from diff_diff import plot_pretrends_power

# Generate power curve across violation magnitudes
curve = pt.power_curve(event_results)

# Plot the power curve
plot_pretrends_power(curve, title="Pre-Trends Test Power Curve")

# Or from the curve object directly
curve.plot()

# Linear trend violations (default) - most common assumption
pt_linear = PreTrendsPower(violation_type='linear')

# Constant violation in all pre-periods
pt_constant = PreTrendsPower(violation_type='constant')

# Violation only in the last pre-period (sharp break)
pt_last = PreTrendsPower(violation_type='last_period')

# Custom violation pattern
custom_weights = np.array([0.1, 0.3, 0.6])  # Increasing violations
pt_custom = PreTrendsPower(violation_type='custom', violation_weights=custom_weights)

from diff_diff import HonestDiD, PreTrendsPower

# If MDV is large relative to your estimated effect, be cautious
pt = PreTrendsPower()
power_results = pt.fit(event_results)
sensitivity = pt.sensitivity_to_honest_did(event_results)
print(sensitivity['interpretation'])

# Use HonestDiD for robust inference
honest = HonestDiD(method='relative_magnitude', M=1.0)
honest_results = honest.fit(event_results)

from diff_diff import run_placebo_test

# Test: Is there an effect before treatment actually occurred?
# Actual treatment is at period 3 (post_periods=[3, 4, 5])
# We test if a "fake" treatment at period 1 shows an effect
results = run_placebo_test(
    data,
    outcome='outcome',
    treatment='treated',
    time='period',
    test_type='fake_timing',
    fake_treatment_period=1,  # Pretend treatment was in period 1
    post_periods=[3, 4, 5]    # Actual post-treatment periods
)

print(results.summary())
# If parallel trends hold, placebo_effect should be ~0 and not significant
print(f"Placebo effect: {results.placebo_effect:.3f} (p={results.p_value:.3f})")
print(f"Is significant (bad): {results.is_significant}")

# Test: Is there an effect among never-treated units?
# Get some control unit IDs to use as "fake treated"
control_units = data[data['treated'] == 0]['firm_id'].unique()[:5]

results = run_placebo_test(
    data,
    outcome='outcome',
    treatment='treated',
    time='period',
    unit='firm_id',
    test_type='fake_group',
    fake_treatment_group=list(control_units),  # List of control unit IDs
    post_periods=[3, 4, 5]
)

# Randomly reassign treatment and compute distribution of effects
# Note: requires binary post indicator (use 'post' column, not 'period')
results = run_placebo_test(
    data,
    outcome='outcome',
    treatment='treated',
    time='post',           # Binary post-treatment indicator
    unit='firm_id',
    test_type='permutation',
    n_permutations=1000,
    seed=42
)

print(f"Original effect: {results.original_effect:.3f}")
print(f"Permutation p-value: {results.p_value:.4f}")
# Low p-value indicates the effect is unlikely to be due to chance

# Test sensitivity to individual treated units
# Note: requires binary post indicator (use 'post' column, not 'period')
results = run_placebo_test(
    data,
    outcome='outcome',
    treatment='treated',
    time='post',           # Binary post-treatment indicator
    unit='firm_id',
    test_type='leave_one_out'
)

# Check if any single unit drives the result
print(results.leave_one_out_effects)  # Effect when each unit is dropped

from diff_diff import run_all_placebo_tests

# Comprehensive diagnostic suite
# Note: This function runs fake_timing tests on pre-treatment periods.
# The permutation and leave_one_out tests require a binary post indicator,
# so they may return errors if the data uses multi-period time column.
all_results = run_all_placebo_tests(
    data,
    outcome='outcome',
    treatment='treated',
    time='period',
    unit='firm_id',
    pre_periods=[0, 1, 2],
    post_periods=[3, 4, 5],
    n_permutations=500,
    seed=42
)

for test_name, result in all_results.items():
    if hasattr(result, 'p_value'):
        print(f"{test_name}: p={result.p_value:.3f}, significant={result.is_significant}")
    elif isinstance(result, dict) and 'error' in result:
        print(f"{test_name}: Error - {result['error']}")

DifferenceInDifferences(
    robust=True,      # Use HC1 robust standard errors
    cluster=None,     # Column for cluster-robust SEs
    alpha=0.05        # Significance level for CIs
)

MultiPeriodDiD(
    robust=True,      # Use HC1 robust standard errors
    cluster=None,     # Column for cluster-robust SEs
    alpha=0.05        # Significance level for CIs
)

SyntheticDiD(
    lambda_reg=0.0,     # L2 regularization for unit weights
    zeta=1.0,           # Regularization for time weights
    alpha=0.05,         # Significance level for CIs
    n_bootstrap=200,    # Bootstrap iterations for SE
    seed=None           # Random seed for reproducibility
)

SunAbraham(
    control_group='never_treated',  # or 'not_yet_treated'
    anticipation=0,           # Periods of anticipation effects
    alpha=0.05,               # Significance level for CIs
    cluster=None,             # Column for cluster-robust SEs
    n_bootstrap=0,            # Bootstrap iterations (0 = analytical SEs)
    bootstrap_weights='rademacher',  # 'rademacher', 'mammen', or 'webb'
    seed=None                 # Random seed
)

TripleDifference(
    estimation_method='dr',   # 'dr' (doubly robust), 'reg', or 'ipw'
    robust=True,              # Use HC1 robust standard errors
    cluster=None,             # Column for cluster-robust SEs
    alpha=0.05,               # Significance level for CIs
    pscore_trim=0.01          # Propensity score trimming threshold
)

HonestDiD(
    method='relative_magnitude',  # 'relative_magnitude' or 'smoothness'
    M=None,               # Restriction parameter (default: 1.0 for RM, 0.0 for SD)
    alpha=0.05,           # Significance level for CIs
    l_vec=None            # Linear combination vector for target parameter
)

PreTrendsPower(
    alpha=0.05,           # Significance level for pre-trends test
    power=0.80,           # Target power for MDV calculation
    violation_type='linear',  # 'linear', 'constant', 'last_period', 'custom'
    violation_weights=None    # Custom weights (required if violation_type='custom')
)

generate_did_data(
    n_units=100,          # Number of units
    n_periods=4,          # Number of time periods
    treatment_effect=5.0, # True ATT
    treatment_fraction=0.5,  # Fraction treated
    treatment_period=2,   # First post-treatment period
    unit_fe_sd=2.0,       # Unit fixed effect std dev
    time_trend=0.5,       # Linear time trend
    noise_sd=1.0,         # Idiosyncratic noise std dev
    seed=None             # Random seed
)

make_treatment_indicator(
    data,                 # Input DataFrame
    column,               # Column to create treatment from
    treated_values=None,  # Value(s) indicating treatment
    threshold=None,       # Numeric threshold for treatment
    above_threshold=True, # If True, >= threshold is treated
    new_column='treated'  # Output column name
)

make_post_indicator(
    data,                  # Input DataFrame
    time_column,           # Time/period column
    post_periods=None,     # Specific post-treatment period(s)
    treatment_start=None,  # First post-treatment period
    new_column='post'      # Output column name
)

wide_to_long(
    data,                  # Wide-format DataFrame
    value_columns,         # List of time-varying columns
    id_column,             # Unit identifier column
    time_name='period',    # Name for time column
    value_name='value',    # Name for value column
    time_values=None       # Values for time periods
)

balance_panel(
    data,                  # Panel DataFrame
    unit_column,           # Unit identifier column
    time_column,           # Time period column
    method='inner',        # 'inner', 'outer', or 'fill'
    fill_value=None        # Value for filling (if method='fill')
)

validate_did_data(
    data,                  # DataFrame to validate
    outcome,               # Outcome column name
    treatment,             # Treatment column name
    time,                  # Time/post column name
    unit=None,             # Unit column (for panel validation)
    raise_on_error=True    # Raise ValueError or return dict
)

summarize_did_data(
    data,                  # Input DataFrame
    outcome,               # Outcome column name
    treatment,             # Treatment column name
    time,                  # Time/post column name
    unit=None              # Unit column (optional)
)

create_event_time(
    data,                  # Panel DataFrame
    time_column,           # Calendar time column
    treatment_time_column, # Column with treatment timing
    new_column='event_time' # Output column name
)

aggregate_to_cohorts(
    data,                  # Unit-level panel data
    unit_column,           # Unit identifier column
    time_column,           # Time period column
    treatment_column,      # Treatment indicator column
    outcome,               # Outcome variable column
    covariates=None        # Additional columns to aggregate
)

rank_control_units(
    data,                          # Panel data in long format
    unit_column,                   # Unit identifier column
    time_column,                   # Time period column
    outcome_column,                # Outcome variable column
    treatment_column=None,         # Treatment indicator column (0/1)
    treated_units=None,            # Explicit list of treated unit IDs
    pre_periods=None,              # Pre-treatment periods (default: first half)
    covariates=None,               # Covariate columns for matching
    outcome_weight=0.7,            # Weight for outcome trend similarity (0-1)
    covariate_weight=0.3,          # Weight for covariate distance (0-1)
    exclude_units=None,            # Units to exclude from control pool
    require_units=None,            # Units that must appear in output
    n_top=None,                    # Return only top N controls
    suggest_treatment_candidates=False,  # Identify treatment candidates
    n_treatment_candidates=5,      # Number of treatment candidates
    lambda_reg=0.0                 # Regularization for synthetic weights
)