# Databricks notebook source
# MAGIC %md
# MAGIC # Smoke Test: insurance-causal, insurance-conformal, insurance-monitoring
# MAGIC End-to-end quality audit: PyPI install, README quickstart, import check, pytest.

# COMMAND ----------

# Install test runner deps and the three libraries from PyPI in the current cluster env
# (Databricks doesn't support venv activation in notebook cells, so we install to the cluster)

# COMMAND ----------

# MAGIC %pip install pytest insurance-causal insurance-conformal insurance-monitoring --quiet

# COMMAND ----------

dbutils.library.restartPython()

# COMMAND ----------

import subprocess, sys, textwrap, json

results = {}

# ============================================================
# LIBRARY 1: insurance-monitoring
# ============================================================
print("=" * 70)
print("SMOKE TEST: insurance-monitoring")
print("=" * 70)

# --- 1a. Import check ---
monitoring_imports_ok = True
monitoring_import_errors = []
try:
    from insurance_monitoring import (
        MonitoringReport,
        psi, csi, ks_test, wasserstein_distance,
        DriftAttributor, DriftAttributionResult,
        InterpretableDriftDetector, InterpretableDriftResult,
        ae_ratio, ae_ratio_ci, calibration_curve, hosmer_lemeshow,
        check_balance, check_auto_calibration, murphy_decomposition,
        rectify_balance, isotonic_recalibrate, deviance,
        BalanceResult, AutoCalibResult, MurphyResult, CalibrationReport, CalibrationChecker,
        PITMonitor, PITAlarm, PITSummary,
        gini_coefficient, gini_drift_test, gini_drift_test_onesample,
        GiniDriftResult, GiniDriftOneSampleResult,
        GiniBootstrapResult, GiniDriftBootstrapTest,
        GiniDriftTest, GiniDriftTestResult,
        lorenz_curve,
        SequentialTest, SequentialTestResult, sequential_test_from_df,
        AERatioThresholds, GiniDriftThresholds, MonitoringThresholds, PSIThresholds,
    )
    print("OK: All public imports succeeded")
except Exception as e:
    monitoring_imports_ok = False
    monitoring_import_errors.append(str(e))
    print(f"FAIL imports: {e}")

# --- 1b. README quickstart (full version with feature DataFrames) ---
monitoring_quickstart_ok = True
monitoring_quickstart_error = None
try:
    import polars as pl
    import numpy as np
    from insurance_monitoring import MonitoringReport

    rng = np.random.default_rng(42)

    n_ref = 10_000
    pred_ref = rng.uniform(0.05, 0.20, n_ref)
    act_ref = rng.poisson(pred_ref).astype(float)

    n_cur = 4_000
    pred_cur = rng.uniform(0.05, 0.20, n_cur)
    act_cur = rng.poisson(pred_cur * 1.08).astype(float)

    feat_ref = pl.DataFrame({
        "driver_age":  rng.integers(18, 80, n_ref).tolist(),
        "vehicle_age": rng.integers(0, 15, n_ref).tolist(),
        "ncd_years":   rng.integers(0, 9, n_ref).tolist(),
    })
    feat_cur = pl.DataFrame({
        "driver_age":  rng.integers(25, 85, n_cur).tolist(),
        "vehicle_age": rng.integers(0, 15, n_cur).tolist(),
        "ncd_years":   rng.integers(0, 9, n_cur).tolist(),
    })

    report = MonitoringReport(
        reference_actual=act_ref,
        reference_predicted=pred_ref,
        current_actual=act_cur,
        current_predicted=pred_cur,
        feature_df_reference=feat_ref,
        feature_df_current=feat_cur,
        features=["driver_age", "vehicle_age", "ncd_years"],
        murphy_distribution="poisson",
    )

    rec = report.recommendation
    df_out = report.to_polars()
    print(f"OK: README quickstart passed. Recommendation: {rec}")
    print(df_out)
except Exception as e:
    monitoring_quickstart_ok = False
    monitoring_quickstart_error = str(e)
    import traceback
    print(f"FAIL quickstart: {e}")
    traceback.print_exc()

# --- 1c. Simple quickstart (no features) ---
monitoring_simple_ok = True
monitoring_simple_error = None
try:
    import numpy as np
    from insurance_monitoring import MonitoringReport

    rng = np.random.default_rng(42)
    pred_ref = rng.uniform(0.05, 0.20, 10_000)
    act_ref = rng.poisson(pred_ref).astype(float)
    pred_cur = rng.uniform(0.05, 0.20, 4_000)
    act_cur = rng.poisson(pred_cur * 1.08).astype(float)

    report = MonitoringReport(
        reference_actual=act_ref,
        reference_predicted=pred_ref,
        current_actual=act_cur,
        current_predicted=pred_cur,
        murphy_distribution="poisson",
    )
    print(f"OK: Simple quickstart recommendation: {report.recommendation}")
except Exception as e:
    monitoring_simple_ok = False
    monitoring_simple_error = str(e)
    import traceback
    print(f"FAIL simple quickstart: {e}")
    traceback.print_exc()

results["monitoring"] = {
    "imports_ok": monitoring_imports_ok,
    "import_errors": monitoring_import_errors,
    "quickstart_ok": monitoring_quickstart_ok,
    "quickstart_error": monitoring_quickstart_error,
    "simple_ok": monitoring_simple_ok,
    "simple_error": monitoring_simple_error,
}

# COMMAND ----------

# ============================================================
# LIBRARY 2: insurance-conformal
# ============================================================
print("=" * 70)
print("SMOKE TEST: insurance-conformal")
print("=" * 70)

# --- 2a. Import check ---
conformal_imports_ok = True
conformal_import_errors = []
try:
    from insurance_conformal import (
        InsuranceConformalPredictor,
        NonconformityScore,
        raw_score, pearson_score, pearson_weighted_score,
        deviance_score, anscombe_score,
        CoverageDiagnostics,
        LocallyWeightedConformal,
        HongModelFree, HongTransformConformal,
        SCRReport,
        ert_coverage_gap, subgroup_coverage, width_efficiency_comparison,
        RetroAdj,
        ConformalisedQuantileRegression,
        solvency_capital_range, SolvencyCapitalRange,
    )
    # Sub-package imports
    from insurance_conformal.risk import (
        PremiumSufficiencyController,
        IntervalWidthController,
        SelectiveRiskController,
        conformal_risk_calibration,
        shortfall_loss,
        premium_sufficiency_report,
    )
    from insurance_conformal.claims import FrequencySeverityConformal
    from insurance_conformal.multivariate import JointConformalPredictor
    print("OK: All public imports succeeded")
except Exception as e:
    conformal_imports_ok = False
    conformal_import_errors.append(str(e))
    print(f"FAIL imports: {e}")
    import traceback; traceback.print_exc()

# --- 2b. README quickstart (uses catboost) ---
conformal_quickstart_ok = True
conformal_quickstart_error = None
try:
    import numpy as np
    from insurance_conformal import InsuranceConformalPredictor

    rng = np.random.default_rng(42)
    n_train, n_cal, n_test = 50_000, 10_000, 10_000
    n_features = 6
    X_train = rng.standard_normal((n_train, n_features))
    X_cal   = rng.standard_normal((n_cal,   n_features))
    X_test  = rng.standard_normal((n_test,  n_features))
    y_train = rng.gamma(shape=1.5, scale=500, size=n_train)
    y_cal   = rng.gamma(shape=1.5, scale=500, size=n_cal)
    y_test  = rng.gamma(shape=1.5, scale=500, size=n_test)

    import catboost
    model = catboost.CatBoostRegressor(
        loss_function="Tweedie:variance_power=1.5",
        iterations=300,
        learning_rate=0.05,
        depth=6,
        verbose=0,
    )
    model.fit(X_train, y_train)

    cp = InsuranceConformalPredictor(
        model=model,
        nonconformity="pearson_weighted",
        distribution="tweedie",
        tweedie_power=1.5,
    )

    cp.calibrate(X_cal, y_cal)
    intervals = cp.predict_interval(X_test, alpha=0.10)
    print("OK: README quickstart passed.")
    print(intervals.head())
except Exception as e:
    conformal_quickstart_ok = False
    conformal_quickstart_error = str(e)
    import traceback
    print(f"FAIL quickstart: {e}")
    traceback.print_exc()

# --- 2c. CoverageDiagnostics test ---
conformal_diag_ok = True
conformal_diag_error = None
try:
    from insurance_conformal import CoverageDiagnostics
    import numpy as np

    rng = np.random.default_rng(42)
    n = 1000
    y_true = rng.gamma(1.5, 500, n)
    y_pred = rng.gamma(1.5, 480, n)
    y_lower = np.maximum(y_pred - 200, 0)
    y_upper = y_pred + 800

    diag = CoverageDiagnostics(
        y_true=y_true,
        y_lower=y_lower,
        y_upper=y_upper,
        y_pred=y_pred,
        alpha=0.10,
    )
    # coverage_by_decile is on the predictor, but CoverageDiagnostics has coverage_plot
    import matplotlib
    matplotlib.use("Agg")
    fig = diag.coverage_plot()
    print("OK: CoverageDiagnostics.coverage_plot() succeeded")
    fig2 = diag.interval_width_distribution()
    print("OK: CoverageDiagnostics.interval_width_distribution() succeeded")
except Exception as e:
    conformal_diag_ok = False
    conformal_diag_error = str(e)
    import traceback
    print(f"FAIL CoverageDiagnostics: {e}")
    traceback.print_exc()

results["conformal"] = {
    "imports_ok": conformal_imports_ok,
    "import_errors": conformal_import_errors,
    "quickstart_ok": conformal_quickstart_ok,
    "quickstart_error": conformal_quickstart_error,
    "diag_ok": conformal_diag_ok,
    "diag_error": conformal_diag_error,
}

# COMMAND ----------

# ============================================================
# LIBRARY 3: insurance-causal
# ============================================================
print("=" * 70)
print("SMOKE TEST: insurance-causal")
print("=" * 70)

# --- 3a. Import check ---
causal_imports_ok = True
causal_import_errors = []
try:
    from insurance_causal import (
        CausalPricingModel,
        AverageTreatmentEffect,
        RateChangeEvaluator,
        make_rate_change_data,
    )
    import insurance_causal.treatments as treatments_mod
    import insurance_causal.diagnostics as diagnostics_mod
    import insurance_causal.autodml as autodml_mod
    import insurance_causal.elasticity as elasticity_mod
    import insurance_causal.causal_forest as causal_forest_mod
    import insurance_causal.rate_change as rate_change_mod
    print("OK: All top-level public imports succeeded")
except Exception as e:
    causal_imports_ok = False
    causal_import_errors.append(str(e))
    print(f"FAIL imports: {e}")
    import traceback; traceback.print_exc()

# --- 3b. README quickstart: autodml subpackage ---
causal_autodml_ok = True
causal_autodml_error = None
try:
    import numpy as np
    from insurance_causal.autodml import PremiumElasticity, DoseResponseCurve

    rng = np.random.default_rng(42)
    n = 5_000

    driver_age  = rng.integers(25, 70, n).astype(float)
    vehicle_age = rng.integers(1, 12, n).astype(float)
    ncb_years   = rng.integers(0, 9, n).astype(float)

    X = np.column_stack([driver_age, vehicle_age, ncb_years,
                         rng.standard_normal(n)])

    risk_score = 0.02 * np.maximum(30 - driver_age, 0) + 0.05 * vehicle_age - 0.08 * ncb_years
    D = 400 + 200 * risk_score + rng.normal(0, 40, n)

    lam = np.exp(-2.0 + 0.3 * risk_score - 0.0001 * D)
    Y = rng.poisson(lam).astype(float)
    exposure = rng.uniform(0.7, 1.0, n)

    model = PremiumElasticity(outcome_family="poisson", n_folds=5)
    model.fit(X, D, Y, exposure=exposure)
    result = model.estimate()
    print("OK: PremiumElasticity README quickstart passed.")
    print(result.summary())

    dr = DoseResponseCurve(outcome_family="poisson")
    dr.fit(X, D, Y)
    curve = dr.predict(d_grid=np.linspace(200, 800, 20))
    print(f"OK: DoseResponseCurve passed, curve shape: {curve.shape}")
except Exception as e:
    causal_autodml_ok = False
    causal_autodml_error = str(e)
    import traceback
    print(f"FAIL autodml quickstart: {e}")
    traceback.print_exc()

# --- 3c. README quickstart: elasticity subpackage ---
causal_elasticity_ok = True
causal_elasticity_error = None
try:
    from insurance_causal.elasticity import RenewalElasticityEstimator, RenewalPricingOptimiser
    from insurance_causal.elasticity.data import make_renewal_data

    df = make_renewal_data(n=10_000)
    confounders = ["age", "ncd_years", "vehicle_group", "region", "channel"]

    est = RenewalElasticityEstimator()
    est.fit(df, confounders=confounders)
    ate, lb, ub = est.ate()
    print(f"OK: RenewalElasticityEstimator ATE: {ate:.3f}  95% CI: [{lb:.3f}, {ub:.3f}]")

    opt = RenewalPricingOptimiser(est)
    result = opt.optimise(df, budget_constraint_pct=0.0)
    print(f"OK: RenewalPricingOptimiser passed")
except Exception as e:
    causal_elasticity_ok = False
    causal_elasticity_error = str(e)
    import traceback
    print(f"FAIL elasticity quickstart: {e}")
    traceback.print_exc()

# --- 3d. README quickstart: causal_forest subpackage ---
causal_forest_ok = True
causal_forest_error = None
try:
    import numpy as np
    from insurance_causal.causal_forest import (
        HeterogeneousElasticityEstimator,
        HeterogeneousInference,
        TargetingEvaluator,
        make_hte_renewal_data,
    )

    df = make_hte_renewal_data(n=10_000, seed=42)
    confounders = ["age", "ncd_years", "vehicle_group", "channel"]

    est = HeterogeneousElasticityEstimator(n_estimators=200, catboost_iterations=200)
    est.fit(df, outcome="renewed", treatment="log_price_change", confounders=confounders)

    ate, lb, ub = est.ate()
    print(f"OK: HeterogeneousElasticityEstimator ATE: {ate:.3f}  95% CI: [{lb:.3f}, {ub:.3f}]")
    cates = est.cate(df)
    gates = est.gate(df, by="ncd_years")
    print(f"OK: CATE shape: {len(cates)}, GATE shape: {gates.shape}")
except Exception as e:
    causal_forest_ok = False
    causal_forest_error = str(e)
    import traceback
    print(f"FAIL causal_forest quickstart: {e}")
    traceback.print_exc()

# --- 3e. README quickstart: rate_change subpackage ---
causal_rate_change_ok = True
causal_rate_change_error = None
try:
    from insurance_causal.rate_change import RateChangeEvaluator, make_rate_change_data

    df = make_rate_change_data(n_segments=40, true_att=-0.05, seed=0)
    evaluator = RateChangeEvaluator(
        outcome_col='outcome',
        treatment_period=9,
        unit_col='segment',
        weight_col='earned_exposure',
    )
    result = evaluator.fit(df)
    print("OK: RateChangeEvaluator quickstart passed.")
    print(result.summary())
except Exception as e:
    causal_rate_change_ok = False
    causal_rate_change_error = str(e)
    import traceback
    print(f"FAIL rate_change quickstart: {e}")
    traceback.print_exc()

results["causal"] = {
    "imports_ok": causal_imports_ok,
    "import_errors": causal_import_errors,
    "autodml_ok": causal_autodml_ok,
    "autodml_error": causal_autodml_error,
    "elasticity_ok": causal_elasticity_ok,
    "elasticity_error": causal_elasticity_error,
    "causal_forest_ok": causal_forest_ok,
    "causal_forest_error": causal_forest_error,
    "rate_change_ok": causal_rate_change_ok,
    "rate_change_error": causal_rate_change_error,
}

# COMMAND ----------

# ============================================================
# RUN PYTEST FOR EACH LIBRARY
# ============================================================

print("=" * 70)
print("RUNNING PYTEST: insurance-monitoring")
print("=" * 70)

# Install the repos in editable mode for pytest
subprocess.run(
    [sys.executable, "-m", "pip", "install", "-e",
     "/Workspace/Repos/pricing.frontier@gmail.com/insurance-monitoring[dev]", "--quiet"],
    capture_output=True
)

r_monitoring_pytest = subprocess.run(
    [sys.executable, "-m", "pytest",
     "/Workspace/Repos/pricing.frontier@gmail.com/insurance-monitoring/tests/",
     "-v", "--tb=short", "-q", "--no-header", "-x"],
    capture_output=True, text=True
)
print(r_monitoring_pytest.stdout[-8000:])
if r_monitoring_pytest.stderr:
    print("STDERR:", r_monitoring_pytest.stderr[-2000:])
results["monitoring"]["pytest_returncode"] = r_monitoring_pytest.returncode
results["monitoring"]["pytest_output"] = r_monitoring_pytest.stdout[-3000:]

# COMMAND ----------

print("=" * 70)
print("RUNNING PYTEST: insurance-conformal")
print("=" * 70)

subprocess.run(
    [sys.executable, "-m", "pip", "install", "-e",
     "/Workspace/Repos/pricing.frontier@gmail.com/insurance-conformal[dev]", "--quiet"],
    capture_output=True
)

r_conformal_pytest = subprocess.run(
    [sys.executable, "-m", "pytest",
     "/Workspace/Repos/pricing.frontier@gmail.com/insurance-conformal/tests/",
     "-v", "--tb=short", "-q", "--no-header", "-x"],
    capture_output=True, text=True
)
print(r_conformal_pytest.stdout[-8000:])
if r_conformal_pytest.stderr:
    print("STDERR:", r_conformal_pytest.stderr[-2000:])
results["conformal"]["pytest_returncode"] = r_conformal_pytest.returncode
results["conformal"]["pytest_output"] = r_conformal_pytest.stdout[-3000:]

# COMMAND ----------

print("=" * 70)
print("RUNNING PYTEST: insurance-causal")
print("=" * 70)

subprocess.run(
    [sys.executable, "-m", "pip", "install", "-e",
     "/Workspace/Repos/pricing.frontier@gmail.com/insurance-causal[dev]", "--quiet"],
    capture_output=True
)

r_causal_pytest = subprocess.run(
    [sys.executable, "-m", "pytest",
     "/Workspace/Repos/pricing.frontier@gmail.com/insurance-causal/tests/",
     "-v", "--tb=short", "-q", "--no-header", "-x"],
    capture_output=True, text=True
)
print(r_causal_pytest.stdout[-8000:])
if r_causal_pytest.stderr:
    print("STDERR:", r_causal_pytest.stderr[-2000:])
results["causal"]["pytest_returncode"] = r_causal_pytest.returncode
results["causal"]["pytest_output"] = r_causal_pytest.stdout[-3000:]

# COMMAND ----------

# ============================================================
# SUMMARY
# ============================================================
print("=" * 70)
print("SMOKE TEST SUMMARY")
print("=" * 70)

for lib, res in results.items():
    ok = all([
        res.get("imports_ok", True),
        res.get("quickstart_ok", True),
        res.get("simple_ok", True),
        res.get("diag_ok", True),
        res.get("autodml_ok", True),
        res.get("elasticity_ok", True),
        res.get("causal_forest_ok", True),
        res.get("rate_change_ok", True),
        res.get("pytest_returncode", 0) == 0,
    ])
    status = "PASS" if ok else "FAIL"
    pytest_rc = res.get("pytest_returncode", "not run")
    print(f"\n{lib.upper()}: {status}")
    print(f"  pytest returncode: {pytest_rc}")
    for k, v in res.items():
        if "error" in k and v:
            print(f"  ERROR [{k}]: {v[:300]}")
        if "errors" in k and v:
            print(f"  ERRORS [{k}]: {v}")

print("\nFull results JSON:")
print(json.dumps({k: {k2: v2 for k2, v2 in v.items() if "output" not in k2} for k, v in results.items()}, indent=2))
