from __future__ import annotations

from survey_kit import logger
from survey_kit.statistics.adapters import (
    statsmodels_adapter,
    linearmodels_adapter,
    pyfixest_adapter,
    polars_ds_adapter,
    mi_ses_from_statsmodels,
    mi_ses_from_linearmodels,
    mi_ses_from_pyfixest,
    mi_ses_from_polars_ds,
)
from survey_kit.utilities.dataframe import summary
from survey_kit.statistics.replicates import Replicates
from sample_data import make_implicates, with_bootstrap_weights

logger.info("survey_kit.statistics.adapters has four pure-Python regression")
logger.info("adapters, one per package - no R/Stata/rpy2/pystata needed. Every")
logger.info("adapter (these four, plus the R/Stata ones) returns the same")
logger.info("normalized shape:")
logger.info("    (df_estimates, df_ses, df_vcov, df_tidy)")
logger.info("and each has a matching mi_ses_from_<package>(...) - same shape as")
logger.info("mi_ses_from_r_fixest/mi_ses_from_stata - that runs it across")
logger.info("implicates, combined via Rubin's rules, taking the adapter's own")
logger.info("arguments directly instead of an arguments={} dict.")

survey_kit.statistics.adapters has four pure-Python regression

adapters, one per package - no R/Stata/rpy2/pystata needed. Every

adapter (these four, plus the R/Stata ones) returns the same

normalized shape:

    (df_estimates, df_ses, df_vcov, df_tidy)

and each has a matching mi_ses_from_<package>(...) - same shape as

mi_ses_from_r_fixest/mi_ses_from_stata - that runs it across

implicates, combined via Rubin's rules, taking the adapter's own

arguments directly instead of an arguments={} dict.

df_implicates = make_implicates()
df = df_implicates[0]
logger.info("\n\nSample data (y = 1 + 2*x1 - 1.5*x2 + noise) - see sample_data.py:")
summary(df)

Sample data (y = 1 + 2*x1 - 1.5*x2 + noise) - see sample_data.py:

┌──────────┬─────┬─────────────┬───────────┬──────────┬───────────┬──────────┐
│ Variable ┆   n ┆ n (missing) ┆      mean ┆      std ┆       min ┆      max │
╞══════════╪═════╪═════════════╪═══════════╪══════════╪═══════════╪══════════╡
│       x1 ┆ 300 ┆           0 ┆ -0.035817 ┆  1.01952 ┆ -3.106337 ┆ 3.066037 │
│       x2 ┆ 300 ┆           0 ┆ -0.009609 ┆ 0.976813 ┆ -3.899422 ┆ 2.118803 │
│        y ┆ 300 ┆           0 ┆  0.931985 ┆ 2.489877 ┆  -5.57149 ┆ 7.779213 │
└──────────┴─────┴─────────────┴───────────┴──────────┴───────────┴──────────┘

naive plan: (run LazyFrame.explain(optimized=True) to see the optimized plan)

SELECT [col("Variable"), col("n"), col("n (missing)"), col("mean"), col("std"), col("min"), col("max")]

UNION

PLAN 0:

SELECT [col("Variable"), col("n"), col("n (missing)"), col("mean"), col("std"), col("min"), col("max")]

WITH_COLUMNS:

["x1".alias("Variable")]

SELECT [col("n (missing)"), col("max"), col("std"), col("n"), col("mean"), col("min")]

SELECT [col("___index___"), col("x1_rawn_missing").alias("n (missing)"), col("x1_max").alias("max"), col("x1_std").alias("std"), col("x1_rawn").alias("n"), col("x1_mean").alias("mean"), col("x1_min").alias("min")]

SELECT [col("___index___"), col("x1_rawn_missing"), col("x1_max"), col("x1_std"), col("x1_rawn"), col("x1_mean"), col("x1_min")]

DF ["___index___", "x1_rawn", "x1_mean", "x1_std", ...]; PROJECT */19 COLUMNS

PLAN 1:

SELECT [col("Variable"), col("n"), col("n (missing)"), col("mean"), col("std"), col("min"), col("max")]

WITH_COLUMNS:

["x2".alias("Variable")]

SELECT [col("n (missing)"), col("max"), col("std"), col("n"), col("mean"), col("min")]

SELECT [col("___index___"), col("x2_rawn_missing").alias("n (missing)"), col("x2_max").alias("max"), col("x2_std").alias("std"), col("x2_rawn").alias("n"), col("x2_mean").alias("mean"), col("x2_min").alias("min")]

SELECT [col("___index___"), col("x2_rawn_missing"), col("x2_max"), col("x2_std"), col("x2_rawn"), col("x2_mean"), col("x2_min")]

DF ["___index___", "x1_rawn", "x1_mean", "x1_std", ...]; PROJECT */19 COLUMNS

PLAN 2:

SELECT [col("Variable"), col("n"), col("n (missing)"), col("mean"), col("std"), col("min"), col("max")]

WITH_COLUMNS:

["y".alias("Variable")]

SELECT [col("n (missing)"), col("max"), col("std"), col("n"), col("mean"), col("min")]

SELECT [col("___index___"), col("y_rawn_missing").alias("n (missing)"), col("y_max").alias("max"), col("y_std").alias("std"), col("y_rawn").alias("n"), col("y_mean").alias("mean"), col("y_min").alias("min")]

SELECT [col("___index___"), col("y_rawn_missing"), col("y_max"), col("y_std"), col("y_rawn"), col("y_mean"), col("y_min")]

DF ["___index___", "x1_rawn", "x1_mean", "x1_std", ...]; PROJECT */19 COLUMNS

END UNION

logger.info("\n\nstatsmodels - y/x column lists rather than a formula, HC3")
logger.info("(heteroskedasticity-robust) SEs by default:")
(df_estimates, df_ses, df_vcov, df_tidy) = statsmodels_adapter(df, y="y", x=["x1", "x2"])
logger.info("\n   Single df")
logger.info(df_estimates)

mi_sm = mi_ses_from_statsmodels(df_implicates=df_implicates, y="y", x=["x1", "x2"])
logger.info("\n   Multiple imputation")
mi_sm.print()

statsmodels - y/x column lists rather than a formula, HC3

(heteroskedasticity-robust) SEs by default:

   Single df

shape: (3, 2)
┌──────────┬───────────┐
│ Variable ┆ estimate  │
│ ---      ┆ ---       │
│ str      ┆ f64       │
╞══════════╪═══════════╡
│ const    ┆ 0.990009  │
│ x1       ┆ 2.016691  │
│ x2       ┆ -1.478573 │
└──────────┴───────────┘

Implicate #1

Implicate #2

Implicate #3

Implicate #4

Implicate #5

   Multiple imputation

┌──────────┬───────────┐
│ Variable ┆  estimate │
╞══════════╪═══════════╡
│    const ┆  0.983256 │
│          ┆  0.028769 │
│       x1 ┆  1.996435 │
│          ┆  0.047977 │
│       x2 ┆ -1.488951 │
│          ┆  0.030939 │
└──────────┴───────────┘

logger.info("\n\nlinearmodels - formula syntax, IV/panel-capable (plain OLS via")
logger.info("IV2SLS with no instruments, as here):")
(df_estimates, df_ses, df_vcov, df_tidy) = linearmodels_adapter(
    df, formula="y ~ 1 + x1 + x2"
)
logger.info("\n   Single df")
logger.info(df_estimates)

mi_lm = mi_ses_from_linearmodels(df_implicates=df_implicates, formula="y ~ 1 + x1 + x2")
logger.info("\n   Multiple imputation")
mi_lm.print()

linearmodels - formula syntax, IV/panel-capable (plain OLS via

IV2SLS with no instruments, as here):

   Single df

shape: (3, 2)
┌───────────┬───────────┐
│ Variable  ┆ estimate  │
│ ---       ┆ ---       │
│ str       ┆ f64       │
╞═══════════╪═══════════╡
│ Intercept ┆ 0.990009  │
│ x1        ┆ 2.016691  │
│ x2        ┆ -1.478573 │
└───────────┴───────────┘

Implicate #1

Implicate #2

Implicate #3

Implicate #4

Implicate #5

   Multiple imputation

┌───────────┬───────────┐
│  Variable ┆  estimate │
╞═══════════╪═══════════╡
│ Intercept ┆  0.983256 │
│           ┆  0.028573 │
│        x1 ┆  1.996435 │
│           ┆  0.047803 │
│        x2 ┆ -1.488951 │
│           ┆   0.03055 │
└───────────┴───────────┘

logger.info("\n\npyfixest - fixest-syntax formula, fixed effects supported")
logger.info("directly (e.g. 'y ~ x1 + x2 | firm') - generally the best default")
logger.info("of the four unless you specifically need something it doesn't")
logger.info("cover (see its docstring):")
(df_estimates, df_ses, df_vcov, df_tidy) = pyfixest_adapter(df, formula="y ~ x1 + x2")
logger.info("\n   Single df")
logger.info(df_estimates)

mi_pf = mi_ses_from_pyfixest.feols(df_implicates=df_implicates, fml="y ~ x1 + x2")
logger.info("\n   Multiple imputation")
mi_pf.print()

pyfixest - fixest-syntax formula, fixed effects supported

directly (e.g. 'y ~ x1 + x2 | firm') - generally the best default

of the four unless you specifically need something it doesn't

cover (see its docstring):

   Single df

shape: (3, 2)
┌───────────┬───────────┐
│ Variable  ┆ estimate  │
│ ---       ┆ ---       │
│ str       ┆ f64       │
╞═══════════╪═══════════╡
│ Intercept ┆ 0.990009  │
│ x1        ┆ 2.016691  │
│ x2        ┆ -1.478573 │
└───────────┴───────────┘

Implicate #1

Implicate #2

Implicate #3

Implicate #4

Implicate #5

   Multiple imputation

┌───────────┬───────────┐
│  Variable ┆  estimate │
╞═══════════╪═══════════╡
│ Intercept ┆  0.983256 │
│           ┆  0.028668 │
│        x1 ┆  1.996435 │
│           ┆  0.047854 │
│        x2 ┆ -1.488951 │
│           ┆  0.030649 │
└───────────┴───────────┘

logger.info("\n\nReplicate-weight bootstrapping instead of pyfixest's own vcov: pass")
logger.info("replicates=, and pyfixest_adapter runs once per replicate weight column")
logger.info("(point estimates only, vcov forced to \"iid\" since it's discarded")
logger.info("anyway) - the spread of estimates across replicates IS the SE, computed")
logger.info("by survey_kit's own Replicates/StatCalculator machinery, the same")
logger.info("approach mi_ses_from_stata/mi_ses_from_r_fixest's replicates= use. Every")
logger.info("mi_ses_from_<package> here (except .feglm/.femlm, whose underlying")
logger.info("estimators have no weights= argument to substitute a replicate column")
logger.info("into) supports this the same way.")

N_REPLICATES = 20
mi_pf_boot = mi_ses_from_pyfixest.feols(
    df_implicates=with_bootstrap_weights(df_implicates, n_replicates=N_REPLICATES),
    fml="y ~ x1 + x2",
    replicates=Replicates(weight_stub="replicate_", n_replicates=N_REPLICATES, bootstrap=True),
)
mi_pf_boot.print()

Replicate-weight bootstrapping instead of pyfixest's own vcov: pass

replicates=, and pyfixest_adapter runs once per replicate weight column

(point estimates only, vcov forced to "iid" since it's discarded

anyway) - the spread of estimates across replicates IS the SE, computed

by survey_kit's own Replicates/StatCalculator machinery, the same

approach mi_ses_from_stata/mi_ses_from_r_fixest's replicates= use. Every

mi_ses_from_<package> here (except .feglm/.femlm, whose underlying

estimators have no weights= argument to substitute a replicate column

into) supports this the same way.

Implicate #1

Running _point_estimate

0....5....10....15....20

Implicate #2

Running _point_estimate

0....5....10....15....20

Implicate #3

Running _point_estimate

0

....5....10....15....20

Implicate #4

Running _point_estimate

0....5....10....15....20

Implicate #5

Running _point_estimate

0

....5....10....15....20

┌───────────┬───────────┐
│  Variable ┆  estimate │
╞═══════════╪═══════════╡
│ Intercept ┆  0.983256 │
│           ┆  0.041516 │
│        x1 ┆  1.996435 │
│           ┆  0.059038 │
│        x2 ┆ -1.488951 │
│           ┆  0.049244 │
└───────────┴───────────┘

logger.info("\n\npolars_ds - stays entirely in polars/narwhals, no pandas")
logger.info("conversion at all; doesn't compute a covariance matrix, so df_vcov")
logger.info("is always None here:")
(df_estimates, df_ses, df_vcov, df_tidy) = polars_ds_adapter(df, y="y", x=["x1", "x2"])
logger.info("\n   Single df")
logger.info(df_estimates)

mi_pds = mi_ses_from_polars_ds(df_implicates=df_implicates, y="y", x=["x1", "x2"])
logger.info("\n   Multiple imputation")
mi_pds.print()

polars_ds - stays entirely in polars/narwhals, no pandas

conversion at all; doesn't compute a covariance matrix, so df_vcov

is always None here:

   Single df

shape: (3, 2)
┌──────────┬───────────┐
│ Variable ┆ estimate  │
│ ---      ┆ ---       │
│ str      ┆ f64       │
╞══════════╪═══════════╡
│ x1       ┆ 2.016691  │
│ x2       ┆ -1.478573 │
│ __bias__ ┆ 0.990009  │
└──────────┴───────────┘

Implicate #1

Implicate #2

Implicate #3

Implicate #4

Implicate #5

   Multiple imputation

┌──────────┬───────────┐
│ Variable ┆  estimate │
╞══════════╪═══════════╡
│       x1 ┆  1.996435 │
│          ┆  0.047977 │
│       x2 ┆ -1.488951 │
│          ┆  0.030939 │
│ __bias__ ┆  0.983256 │
│          ┆  0.028769 │
└──────────┴───────────┘

logger.info("\n\nFor R or Stata instead, see basic_r.py/basic_stata.py (or")
logger.info("r_arbitrary_estimators.py/stata_arbitrary_estimators.py for the generic")
logger.info("escape hatch into any function either package provides).")

For R or Stata instead, see basic_r.py/basic_stata.py (or

r_arbitrary_estimators.py/stata_arbitrary_estimators.py for the generic

escape hatch into any function either package provides).