Aug 25 Bug Fixes

models/its.py

@@ -1,228 +1,534 @@
-import numpy as np
-import pandas as pd
-from io import BytesIO
-from PIL import Image
-import causalpy as cp
-import matplotlib.pyplot as plt
-from sklearn.linear_model import LinearRegression
-import statsmodels.api as sm # For sm.add_constant
-from scipy import stats
-
-
-def enhanced_its_analysis(file, target_col, date_col, pre_dates, post_dates, freq_input="D", control_vars=""):  
-    """
-    Performs interrupted time series analysis using causalpy on user-provided frequency.
-    Uses DatetimeIndex for CausalPy compatibility with Timestamp treatment_time.
-    Ensures predicted counterfactual is 1D for statistics.
-    
-    Parameters:
-    -----------
-    file : file object
-        The uploaded CSV file
-    target_col : str
-        The column name for the outcome variable
-    date_col : str
-        The column name for the date variable
-    pre_dates : str
-        Comma-separated start and end dates for pre-intervention period
-    post_dates : str
-        Comma-separated start and end dates for post-intervention period
-    freq_input : str, default="D"
-        Pandas frequency alias for the time series
-    control_vars : str, default=""
-        Comma-separated list of control variable column names
-    """
-    if file is None:
-        return "Error: No file uploaded.", None
-    try:
-        df = pd.read_csv(file.name)
-        if date_col not in df.columns:
-            return f"Error: Date column '{date_col}' not found.", None
-        
-        df[date_col] = pd.to_datetime(df[date_col], errors='coerce')
-        df = df.dropna(subset=[date_col]) 
-        if df.empty:
-            return f"Error: No valid dates found in column '{date_col}' after parsing.", None
-
-        if target_col not in df.columns:
-            return f"Error: Column '{target_col}' not found.", None
-        
-        if not pd.api.types.is_numeric_dtype(df[target_col]):
-            df[target_col] = pd.to_numeric(df[target_col], errors='coerce')
-        df = df.dropna(subset=[target_col]) 
-        if df.empty:
-            return f"Error: No valid data in target column '{target_col}' after NA removal.", None
-
-        # Process control variables
-        control_columns = []
-        if control_vars and control_vars.strip():
-            control_columns = [col.strip() for col in control_vars.split(',')]
-            # Validate control variables exist in the dataframe
-            missing_cols = [col for col in control_columns if col not in df.columns]
-            if missing_cols:
-                return f"Error: Control variable column(s) not found: {', '.join(missing_cols)}", None
-            
-            # Convert control variables to numeric if needed
-            for col in control_columns:
-                if not pd.api.types.is_numeric_dtype(df[col]):
-                    df[col] = pd.to_numeric(df[col], errors='coerce')
-                # Check if conversion resulted in all NaN values
-                if df[col].isna().all():
-                    return f"Error: Control variable '{col}' could not be converted to numeric values.", None
-
-        pre_list  = [d.strip() for d in pre_dates.split(',')]
-        post_list = [d.strip() for d in post_dates.split(',')]
-        
-        start_pre, end_pre   = pd.to_datetime(pre_list[0], errors='coerce'), pd.to_datetime(pre_list[1], errors='coerce')
-        start_post, end_post = pd.to_datetime(post_list[0], errors='coerce'), pd.to_datetime(post_list[1], errors='coerce')
-
-        if pd.NaT in [start_pre, end_pre, start_post, end_post]:
-            return "Error: One or more pre/post period boundary dates are invalid. Use YYYY-MM-DD.", None
-        
-        if start_post <= end_pre: # Basic sanity check for period ordering
-            return f"Error: Post-intervention start date ({start_post.date()}) must be after pre-intervention end date ({end_pre.date()}).", None
-
-        mask = (df[date_col] >= start_pre) & (df[date_col] <= end_post)
-        analysis_df_filtered = df.loc[mask].copy()
-
-        if analysis_df_filtered.empty:
-            return "Error: No data in the specified overall date range (pre-start to post-end).", None
-
-        analysis_df_for_cp = analysis_df_filtered.sort_values(date_col)
-        analysis_df_for_cp = analysis_df_for_cp.set_index(date_col, drop=True)
-        
-        analysis_df_for_cp['time_index'] = range(len(analysis_df_for_cp))
-        analysis_df_for_cp = analysis_df_for_cp.rename(columns={target_col: 'y'})
-
-        pre_df  = analysis_df_for_cp.loc[start_pre:end_pre]
-        post_df = analysis_df_for_cp.loc[start_post:end_post]
-
-        if pre_df.empty:
-            return f"Error: Pre-intervention period ({start_pre.date()} to {end_pre.date()}) contains no data after filtering.", None
-        if post_df.empty:
-             return f"Error: Post-intervention period ({start_post.date()} to {end_post.date()}) contains no data after filtering.", None
-
-        # Build formula with control variables if provided
-        formula = 'y ~ 1 + time_index'
-        if control_columns:
-            formula += ' + ' + ' + '.join(control_columns)
-
-        # Check for missing values in control variables for the analysis period
-        if control_columns:
-            for col in control_columns:
-                if analysis_df_for_cp[col].isna().any():
-                    return f"Error: Control variable '{col}' contains missing values in the analysis period.", None
-
-        its_model = cp.InterruptedTimeSeries(
-            data=analysis_df_for_cp,
-            formula=formula,  # Now includes control variables if specified
-            treatment_time=start_post,
-            model=LinearRegression(),
-            freq=freq_input
-        )
-
-        # For prediction, we need to prepare X_post with control variables
-        if not control_columns:
-            X_post = sm.add_constant(post_df[['time_index']])
-        else:
-            # Ensure control columns are present in post_df (they should be as post_df is a slice of analysis_df_for_cp)
-            missing_in_post_df = [col for col in control_columns if col not in post_df.columns]
-            if missing_in_post_df: # Should ideally not happen if logic is correct
-                return f"Error: Control variable(s) {', '.join(missing_in_post_df)} missing from post-intervention data slice.", None
-            X_post = sm.add_constant(post_df[['time_index'] + control_columns])
-            
-        pred_cf_array = its_model.model.predict(X_post)
-        
-        # --- FIX: Ensure pred_cf_array is 1D ---
-        if pred_cf_array.ndim > 1:
-            pred_cf_array = pred_cf_array.squeeze() 
-        # --- End of FIX ---
-
-        pred_cf = pd.Series(pred_cf_array, index=post_df.index, name='y_fc')
-        observed_post = post_df['y']
-
-        post_mean = observed_post.mean()
-        cf_mean   = pred_cf.mean()
-        effect    = post_mean - cf_mean
-
-        if len(observed_post) < 2 or len(pred_cf) < 2:
-            post_se, cf_se, eff_se, t_stat, p_value, ci_low, ci_high = [np.nan] * 7
-            df_t = 0
-        else:
-            post_se   = observed_post.std(ddof=1) / np.sqrt(len(observed_post))
-            cf_se     = pred_cf.std(ddof=1) / np.sqrt(len(pred_cf))
-            eff_se    = np.sqrt(post_se**2 + cf_se**2) if (not np.isnan(post_se) and not np.isnan(cf_se)) else np.nan
-            df_t      = min(len(post_df) - 1, len(pred_cf) -1) 
-            if df_t < 1 : df_t=1
-
-            if np.isnan(eff_se) or eff_se == 0:
-                t_stat, p_value, ci_low, ci_high = [np.nan] * 4
-            else:
-                t_stat    = effect / eff_se
-                p_value   = 2 * (1 - stats.t.cdf(abs(t_stat), df=df_t))
-                ci_low, ci_high = stats.t.interval(0.95, df_t, loc=effect, scale=eff_se)
-
-        # Enhanced report to include control variables
-        report_lines = [
-            f"ITS Analysis for: 'y' (originally '{target_col}')",
-            f"Intervention at: {start_post.strftime('%Y-%m-%d')}",
-            f"Pre-period: {start_pre.strftime('%Y-%m-%d')} to {end_pre.strftime('%Y-%m-%d')}",
-            f"Post-period: {start_post.strftime('%Y-%m-%d')} to {end_post.strftime('%Y-%m-%d')}"
-        ]
-        
-        if control_columns:
-            report_lines.append(f"Control Variables: {', '.join(control_columns)}")
-            
-        report_lines.extend([
-            "--- Average Effect Estimation ---",
-            f"Observed post-intervention mean: {post_mean:.3f}",
-            f"Estimated counterfactual mean: {cf_mean:.3f}",
-            f"Estimated average effect: {effect:.3f}"
-        ])
-        
-        if not np.isnan(p_value):
-            report_lines.append(f"  95% CI: [{ci_low:.3f}, {ci_high:.3f}]")
-            report_lines.append(f"  t-statistic: {t_stat:.3f}, p-value: {p_value:.4f} (df={df_t})")
-        else:
-            report_lines.append("  (CI/p-value not computed due to insufficient data or variability)")
-        report = "\n".join(report_lines)
-
-        fig, ax = its_model.plot(plot_predict_all=False, plot_show_params=True)
-        buf = BytesIO()
-        fig.savefig(buf, format='png', bbox_inches='tight')
-        plt.close(fig)
-        buf.seek(0)
-        img = Image.open(buf)
-
-        return report, img
-
-    except Exception as e:
-        # import traceback # For debugging
-        # print("--- TRACEBACK ---")
-        # traceback.print_exc()
-        # print("--- END TRACEBACK ---")
-        return f"An unexpected error occurred: {str(e)}", None
-
-
-def run_its_analysis(file, target_col, date_col, pre_dates, post_dates, freq_input, control_vars=""):
-    """
-    Wrapper function for the enhanced_its_analysis function.
-    
-    Parameters:
-    -----------
-    file : file object
-        The uploaded CSV file
-    target_col : str
-        The column name for the outcome variable
-    date_col : str
-        The column name for the date variable
-    pre_dates : str
-        Comma-separated start and end dates for pre-intervention period
-    post_dates : str
-        Comma-separated start and end dates for post-intervention period
-    freq_input : str
-        Pandas frequency alias for the time series
-    control_vars : str, default=""
-        Comma-separated list of control variable column names
-    """
-    return enhanced_its_analysis(file, target_col, date_col, pre_dates, post_dates, freq_input, control_vars)
+from io import BytesIO
+from typing import List, Optional, Tuple, Union
+
+import numpy as np
+import pandas as pd
+from PIL import Image
+
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+import matplotlib.gridspec as gridspec
+
+import causalpy as cp
+import patsy
+import statsmodels.api as sm
+from scipy import stats
+from sklearn.linear_model import LinearRegression
+from statsmodels.stats.diagnostic import acorr_ljungbox
+from statsmodels.tsa.stattools import acf
+import statsmodels.stats.stattools as smt
+import traceback
+
+# ==================== Global knobs ====================
+SEASONALITY_ENABLED = True      
+SEASONALITY_METHOD  = "fourier"  # "fourier" or "dummies"
+FOURIER_WEEKLY_K    = 3          # number of sin/cos pairs for short cycle
+FOURIER_YEARLY_K    = 5          # number of sin/cos pairs for long cycle
+
+
+HAC_ENABLED: bool = True        
+HAC_MAXLAGS: Union[str, int] = "auto"  # "auto" = plug-in; or set an int (e.g., 8 for ~two months on weekly data)
+HAC_SMALL_SAMPLE_CORR: bool = True     # use finite-sample correction in statsmodels
+
+
+# -------------------- rendering helpers --------------------
+
+def _fig_to_pil(fig: plt.Figure, dpi: int = 110) -> Image.Image:
+    """Save a figure to a PIL image with outer padding and opaque background."""
+    buf = BytesIO()
+    fig.savefig(buf, format="png", dpi=dpi, bbox_inches="tight", pad_inches=0.40, facecolor="white")
+    plt.close(fig)
+    buf.seek(0)
+    return Image.open(buf).convert("RGB")
+
+
+def _stack_images_vertical(images: List[Image.Image], pad: int = 22, bg=(255, 255, 255)) -> Optional[Image.Image]:
+    if not images:
+        return None
+    max_w = max(im.width for im in images)
+    total_h = sum(im.height for im in images) + pad * (len(images) - 1)
+    out = Image.new("RGB", (max_w, total_h), bg)
+    y = 0
+    for im in images:
+        x = (max_w - im.width) // 2
+        out.paste(im, (x, y))
+        y += im.height + pad
+    return out
+
+
+def _rotate_all_xticklabels(fig: plt.Figure) -> None:
+    """Rotate x-tick labels on every axes in a figure."""
+    for ax in fig.axes:
+        for lbl in ax.get_xticklabels():
+            lbl.set_rotation(45)
+            lbl.set_ha("right")
+    try:
+        fig.autofmt_xdate()
+    except Exception:
+        pass
+
+
+# -------------------- frequency-aware seasonality (opt-in originally; default ON now) --------------------
+
+def _add_frequency_aware_seasonality(df: pd.DataFrame, freq_input: str) -> Tuple[pd.DataFrame, List[str]]:
+    """
+    Return (df_with_terms, season_terms_list) based on freq_input.
+    df index must be DatetimeIndex and df must contain 'time_index'.
+    No side effects outside df copy. Purely pre-period features, no leakage.
+    """
+    df = df.copy()
+    added: List[str] = []
+
+    # Helpers to add Fourier pairs
+    def add_fourier(prefix: str, period: Optional[float], K: int) -> None:
+        if K <= 0 or period is None:
+            return
+        for k in range(1, K + 1):
+            s = f"{prefix}_sin_{k}"
+            c = f"{prefix}_cos_{k}"
+            df[s] = np.sin(2 * np.pi * k * df["time_index"] / period)
+            df[c] = np.cos(2 * np.pi * k * df["time_index"] / period)
+            added.extend([s, c])
+
+    f = (freq_input or "").upper()
+
+    if SEASONALITY_METHOD == "dummies":
+        if f == "M":
+            df["month"] = df.index.month
+            added.append("C(month)")
+        elif f == "Q":
+            df["quarter"] = df.index.quarter
+            added.append("C(quarter)")
+        elif f == "D":
+            df["dow"] = df.index.dayofweek
+            df["month"] = df.index.month
+            added.extend(["C(dow)", "C(month)"])
+        elif f == "W":
+            df["weekofyear"] = df.index.isocalendar().week.astype(int)
+            added.append("C(weekofyear)")
+        else:
+            # Unknown/other: no-op
+            pass
+
+    else:  # "fourier" (smooth & compact)
+        if f == "D":
+            add_fourier("wk", period=7.0, K=FOURIER_WEEKLY_K)           # weekly cycle
+            add_fourier("yr", period=365.25, K=FOURIER_YEARLY_K)        # yearly cycle
+        elif f == "W":
+            add_fourier("yr", period=52.1775, K=FOURIER_YEARLY_K)       # annual on weekly cadence
+        elif f == "M":
+            add_fourier("yr", period=12.0, K=FOURIER_YEARLY_K)          # annual on monthly cadence
+        elif f == "Q":
+            add_fourier("yr", period=4.0, K=max(1, min(FOURIER_YEARLY_K, 2)))  # annual on quarterly cadence
+        else:
+            # Fallback: do nothing if alias not recognized
+            pass
+
+    return df, added
+
+
+# -------------------- HAC utilities --------------------
+
+def _nw_auto_maxlags(n: int) -> int:
+    """
+    Newey–West plug-in bandwidth: floor(4 * (n/100)^(2/9)), at least 1.
+    """
+    if n <= 1:
+        return 1
+    return max(1, int(np.floor(4.0 * (n / 100.0) ** (2.0 / 9.0))))
+
+
+# Bartlett-weighted (Newey–West) SE for the mean of a time series
+def _nw_se_of_mean(series: pd.Series, maxlags: int) -> float:
+    """
+    Compute Newey–West standard error of the sample mean with Bartlett weights.
+    Var(mean) ≈ (1/n) * [γ0 + 2 * sum_{k=1..L} w_k * γ_k], w_k = 1 - k/(L+1).
+    γ_k are sample autocovariances at lag k with divisor n (not n-1).
+    """
+    x = np.asarray(series, dtype=float)
+    n = x.shape[0]
+    if n <= 1:
+        return np.nan
+    x = x - x.mean()
+    # autocovariances γ_k
+    gamma0 = np.dot(x, x) / n
+    lrvar = gamma0
+    L = min(maxlags, n - 1) if n > 1 else 0
+    for k in range(1, L + 1):
+        w = 1.0 - k / (L + 1.0)
+        cov = np.dot(x[k:], x[:-k]) / n
+        lrvar += 2.0 * w * cov
+    var_mean = lrvar / n
+    return float(np.sqrt(var_mean))
+
+
+# -------------------- diagnostics & comparisons --------------------
+
+def add_diagnostic_tests(sm_model, pre_data, formula, report_lines):
+    """
+    Build diagnostic figures and append textual tests to report_lines.
+    Returns a dict of {name: Figure}.
+    """
+    diagnostic_plots = {}
+
+    report_lines.append("\n--- Diagnostic Tests ---")
+    try:
+        residuals = sm_model.resid
+
+        # Durbin–Watson
+        dw = smt.durbin_watson(residuals)
+        report_lines.append(f"Durbin-Watson statistic: {dw:.3f}")
+        if dw < 1.5:
+            report_lines.append("  ⚠️ Positive autocorrelation detected (DW < 1.5)")
+        elif dw > 2.5:
+            report_lines.append("  ⚠️ Negative autocorrelation detected (DW > 2.5)")
+        else:
+            report_lines.append("  ✓ No significant autocorrelation (1.5 < DW < 2.5)")
+
+        # Ljung–Box
+        if len(residuals) > 10:
+            lb = acorr_ljungbox(residuals, lags=min(10, len(residuals)//4), return_df=True)
+            sig = lb[lb["lb_pvalue"] < 0.05]
+            if len(sig) > 0:
+                report_lines.append(f"  Ljung-Box: Autocorrelation at lags {list(sig.index)}")
+            else:
+                report_lines.append("  Ljung-Box: No significant autocorrelation up to lag 10")
+
+        # ACF
+        if len(residuals) > 20:
+            fig_acf, ax = plt.subplots(figsize=(11, 6))
+            acf_vals = acf(residuals, nlags=min(20, len(residuals)//4))
+            ax.bar(range(len(acf_vals)), acf_vals, alpha=0.85)
+            ax.axhline(0, linewidth=0.5)
+            ci = 1.96/np.sqrt(len(residuals))
+            ax.axhline(ci, linestyle="--", alpha=0.7)
+            ax.axhline(-ci, linestyle="--", alpha=0.7)
+            ax.set_title("Autocorrelation Function (ACF) of Residuals")
+            ax.set_xlabel("Lag"); ax.set_ylabel("Autocorrelation")
+            ax.grid(True, alpha=0.3)
+            fig_acf.tight_layout(pad=1.2)
+            diagnostic_plots["acf"] = fig_acf
+
+    except Exception as e:
+        report_lines.append(f"  Could not perform autocorrelation test: {e}")
+
+    # Model fit stats
+    report_lines.append("\n--- Model Fit Statistics ---")
+    report_lines.append(f"R-squared: {sm_model.rsquared:.3f}")
+    report_lines.append(f"Adjusted R-squared: {sm_model.rsquared_adj:.3f}")
+    report_lines.append(f"AIC: {sm_model.aic:.2f}")
+    report_lines.append(f"BIC: {sm_model.bic:.2f}")
+
+    # Residuals figure (6 panels) with generous spacing
+    try:
+        fig_resid = plt.figure(figsize=(13.5, 10.5), constrained_layout=False)  # MODIFIED: bigger
+        gs = gridspec.GridSpec(3, 2, figure=fig_resid, hspace=0.85, wspace=0.55)  # MODIFIED: more space
+
+        # Residuals vs Fitted
+        ax1 = fig_resid.add_subplot(gs[0, 0])
+        ax1.scatter(sm_model.fittedvalues, sm_model.resid, alpha=0.65)
+        ax1.axhline(0, linestyle="--", alpha=0.7)
+        ax1.set_title("Residuals vs Fitted Values"); ax1.set_xlabel("Fitted Values"); ax1.set_ylabel("Residuals")
+        ax1.grid(True, alpha=0.3)
+
+        # Normal Q–Q
+        ax2 = fig_resid.add_subplot(gs[0, 1])
+        stats.probplot(sm_model.resid, dist="norm", plot=ax2)
+        ax2.set_title("Normal Q-Q Plot"); ax2.grid(True, alpha=0.3)
+
+        # Histogram
+        ax3 = fig_resid.add_subplot(gs[1, 0])
+        ax3.hist(sm_model.resid, bins=20, edgecolor="black", alpha=0.75)
+        ax3.set_title("Histogram of Residuals"); ax3.set_xlabel("Residuals"); ax3.set_ylabel("Density")
+        ax3.grid(True, alpha=0.3)
+
+        # Residuals over time
+        ax4 = fig_resid.add_subplot(gs[1, 1])
+        ax4.plot(pre_data.index, sm_model.resid, marker="o", alpha=0.7)
+        ax4.axhline(0, linestyle="--", alpha=0.7)
+        ax4.set_title("Residuals Over Time"); ax4.set_xlabel("Date"); ax4.set_ylabel("Residuals")
+        for lbl in ax4.get_xticklabels():
+            lbl.set_rotation(45); lbl.set_ha("right")
+        ax4.grid(True, alpha=0.3)
+
+        # Scale–Location
+        ax5 = fig_resid.add_subplot(gs[2, 0])
+        std_resid = sm_model.resid / sm_model.resid.std()
+        ax5.scatter(sm_model.fittedvalues, np.sqrt(np.abs(std_resid)), alpha=0.65)
+        ax5.set_title("Scale-Location Plot"); ax5.set_xlabel("Fitted Values"); ax5.set_ylabel("√|Standardized Residuals|")
+        ax5.grid(True, alpha=0.3)
+
+        # Influence (Cook’s Distance)
+        ax6 = fig_resid.add_subplot(gs[2, 1])
+        try:
+            from statsmodels.stats.outliers_influence import OLSInfluence
+            infl = OLSInfluence(sm_model)
+            ax6.scatter(range(len(infl.cooks_distance[0])), infl.cooks_distance[0], alpha=0.65)
+            ax6.axhline(4/len(sm_model.resid), linestyle="--", alpha=0.7, label="4/n threshold")
+            ax6.legend()
+        except Exception:
+            ax6.text(0.5, 0.5, "Influence plot unavailable", ha="center", va="center")
+        ax6.set_title("Cook's Distance (Influence Plot)"); ax6.set_xlabel("Observation Index"); ax6.set_ylabel("Cook's Distance")
+        ax6.grid(True, alpha=0.3)
+
+        fig_resid.subplots_adjust(top=0.92, bottom=0.20, left=0.10, right=0.98, hspace=0.85, wspace=0.55)  # MODIFIED
+        diagnostic_plots["residuals"] = fig_resid
+
+    except Exception as e:
+        report_lines.append(f"  Could not create residual diagnostic plots: {e}")
+
+    return diagnostic_plots
+
+
+def compare_model_specifications(pre_data: pd.DataFrame, formula_base: str,
+                                 control_columns: List[str], report_lines: List[str]):
+    """
+    Fit linear vs polynomial pre-period trends; return dict possibly
+    containing 'comparison_plot' Figure. Append summary lines to report.
+    """
+    report_lines.append("\n--- Model Specification Comparison ---")
+    try:
+        formula_linear = formula_base
+        formula_quad = 'y ~ 1 + time_index + I(time_index**2)' + ('' if not control_columns else ' + ' + ' + '.join(control_columns))
+        formula_cubic = 'y ~ 1 + time_index + I(time_index**2) + I(time_index**3)' + ('' if not control_columns else ' + ' + ' + '.join(control_columns))
+
+        models, formulas = {}, {'Linear': formula_linear, 'Quadratic': formula_quad, 'Cubic': formula_cubic}
+        for name, fml in formulas.items():
+            try:
+                y, X = patsy.dmatrices(fml, data=pre_data, return_type='dataframe')
+                models[name] = sm.OLS(y, X).fit()
+                report_lines.append(f"{name}: R² {models[name].rsquared:.3f}, AIC {models[name].aic:.1f}, BIC {models[name].bic:.1f}")
+            except Exception:
+                report_lines.append(f"{name}: could not fit")
+
+        out = {}
+        if 'Linear' in models:
+            fig, ax = plt.subplots(figsize=(11, 6))
+            ax.scatter(pre_data.index, pre_data['y'], alpha=0.6, label='Actual', color='black')
+            colors = {'Linear': 'tab:blue', 'Quadratic': 'tab:red', 'Cubic': 'tab:green'}
+            for name, mdl in models.items():
+                if hasattr(mdl, 'fittedvalues'):
+                    ax.plot(pre_data.index, mdl.fittedvalues, label=f'{name} fit', color=colors.get(name, None), linewidth=2, alpha=0.85)
+            ax.set_title('Model Specification Comparison (Pre Period)')
+            ax.set_xlabel('Date'); ax.set_ylabel('Outcome'); ax.grid(True, alpha=0.3); ax.legend()
+            fig.tight_layout(pad=1.2)
+            out['comparison_plot'] = fig
+
+        if len(models) >= 2:
+            best = min(models.items(), key=lambda kv: kv[1].bic if hasattr(kv[1], 'bic') else np.inf)
+            report_lines.append(f"Recommended (by BIC): {best[0]}")
+
+        return out
+
+    except Exception as e:
+        report_lines.append(f"  Could not compare model specifications: {e}")
+        return {}
+
+
+# -------------------- analysis --------------------
+
+def enhanced_its_analysis(file, target_col, date_col, pre_dates, post_dates, freq_input="D",
+                          control_vars="", run_diagnostics=True, show_formulas=False):
+    """
+    ITS analysis using CausalPy with diagnostics and optional model-spec comparison.
+    Returns (report_text, stacked_image).
+    """
+    if file is None:
+        return "Error: No file uploaded.", None
+
+    try:
+        # Load & validate
+        df = pd.read_csv(file.name)
+        if date_col not in df.columns:
+            return f"Error: Date column '{date_col}' not found.", None
+        if target_col not in df.columns:
+            return f"Error: Column '{target_col}' not found.", None
+
+        df[date_col] = pd.to_datetime(df[date_col], errors="coerce")
+        df = df.dropna(subset=[date_col]).sort_values(date_col).set_index(date_col)
+        df[target_col] = pd.to_numeric(df[target_col], errors="coerce")
+        df = df.dropna(subset=[target_col]).rename(columns={target_col: "y"})
+        df["time_index"] = np.arange(len(df), dtype=int)
+
+        # Periods
+        try:
+            pre_s, pre_e = [pd.to_datetime(s.strip(), errors="raise") for s in pre_dates.split(",")]
+            post_s, post_e = [pd.to_datetime(s.strip(), errors="raise") for s in post_dates.split(",")]
+        except Exception:
+            return "Error: Use 'YYYY-MM-DD,YYYY-MM-DD' for pre/post.", None
+        if not (pre_s <= pre_e < post_s <= post_e):
+            return "Error: Must satisfy pre_start <= pre_end < post_start <= post_end.", None
+
+        df = df.loc[(df.index >= pre_s) & (df.index <= post_e)].copy()
+        if df.empty:
+            return "Error: No data in the specified overall date range.", None
+
+        # Controls
+        control_columns: List[str] = []
+        if control_vars and control_vars.strip():
+            control_columns = [c.strip() for c in control_vars.split(",") if c.strip()]
+            missing = [c for c in control_columns if c not in df.columns]
+            if missing:
+                return f"Error: Control variable(s) not found: {', '.join(missing)}", None
+            for c in control_columns:
+                df[c] = pd.to_numeric(df[c], errors="coerce")
+            df = df.dropna(subset=control_columns)
+            if df.empty:
+                return "Error: Data empty after removing NA rows for controls.", None
+
+        season_terms: List[str] = []
+        if SEASONALITY_ENABLED:
+            df, season_terms = _add_frequency_aware_seasonality(df, freq_input)
+
+        # Formula (base + seasonal + controls)
+        formula = "y ~ 1 + time_index"
+        if season_terms:
+            season_rhs = " + ".join(season_terms)  # 'C(...)' tokens or column names
+            formula += " + " + season_rhs
+        if control_columns:
+            formula += " + " + " + ".join(control_columns)
+
+        # Fit pre-period OLS for inference components
+        pre_df = df.loc[df.index < post_s]
+        if pre_df.empty:
+            return "Error: Pre-intervention period is empty after filtering.", None
+        y_pre, X_pre = patsy.dmatrices(formula, data=pre_df, return_type="dataframe")
+        if X_pre.shape[0] <= X_pre.shape[1]:
+            return f"Error: Not enough pre-period observations ({X_pre.shape[0]}) to estimate {X_pre.shape[1]} parameters.", None
+
+        # HAC bandwidth (auto or user-provided)
+        if HAC_ENABLED:
+            if isinstance(HAC_MAXLAGS, str) and HAC_MAXLAGS.lower() == "auto":
+                hac_lags = _nw_auto_maxlags(len(pre_df))
+            else:
+                hac_lags = int(HAC_MAXLAGS)
+            sm_ols = sm.OLS(y_pre, X_pre).fit(
+                cov_type="HAC",
+                cov_kwds={"maxlags": hac_lags, "use_correction": HAC_SMALL_SAMPLE_CORR}
+            )
+            inference_note = f"HAC (Newey–West), maxlags={hac_lags}"
+        else:
+            sm_ols = sm.OLS(y_pre, X_pre).fit()
+            inference_note = "OLS (iid errors assumption)"
+
+        # Post design for counterfactual mean + inference
+        post_df = df.loc[(df.index >= post_s) & (df.index <= post_e)]
+        X_post = patsy.dmatrix(formula.split("~", 1)[1], data=post_df, return_type="dataframe")
+        pred_cf = sm_ols.predict(X_post)
+        observed_post = post_df['y']
+
+        post_mean = float(observed_post.mean())
+        cf_mean = float(np.asarray(pred_cf).mean())
+        effect = post_mean - cf_mean
+
+        if HAC_ENABLED:
+            post_se = _nw_se_of_mean(observed_post, maxlags=_nw_auto_maxlags(len(observed_post)) if (isinstance(HAC_MAXLAGS, str) and HAC_MAXLAGS.lower()=="auto") else int(HAC_MAXLAGS))
+        else:
+            post_se = float(observed_post.std(ddof=1) / np.sqrt(len(observed_post))) if len(observed_post) >= 2 else np.nan
+
+        # SE(counterfactual mean) via delta method using (robust) cov(beta)
+        cov_beta = sm_ols.cov_params()  # MODIFIED: robust if HAC_ENABLED=True
+        X_bar = X_post.mean(axis=0).reindex(cov_beta.columns)
+        var_cf_mean = float(X_bar.T @ cov_beta @ X_bar)
+        cf_se = float(np.sqrt(var_cf_mean)) if var_cf_mean >= 0 else np.nan
+
+        # Combine SEs (independence approximation between observed and cf)
+        eff_se = float(np.sqrt(post_se**2 + cf_se**2)) if (np.isfinite(post_se) and np.isfinite(cf_se)) else np.nan
+
+        # test statistic & CI — normal (z) under HAC, t otherwise
+        if np.isfinite(eff_se) and eff_se > 0:
+            if HAC_ENABLED:
+                z_stat = effect / eff_se
+                p_value = 2 * (1 - stats.norm.cdf(abs(z_stat)))
+                ci_margin = 1.96 * eff_se
+                ci_low, ci_high = effect - ci_margin, effect + ci_margin
+                test_line = f"  z-statistic:               {z_stat:.3f}"
+            else:
+                df_t = max(int(sm_ols.df_resid), 1)
+                t_stat = effect / eff_se
+                p_value = 2 * (1 - stats.t.cdf(abs(t_stat), df=df_t))
+                ci_low, ci_high = stats.t.interval(0.95, df_t, loc=effect, scale=eff_se)
+                test_line = f"  t-statistic:               {t_stat:.3f}"
+        else:
+            p_value = ci_low = ci_high = np.nan
+            test_line = "  (Test statistic unavailable)"
+
+        # Build report
+        report_lines = [
+            "=" * 60,
+            "INTERRUPTED TIME SERIES ANALYSIS REPORT",
+            "=" * 60,
+            f"\nOutcome: {target_col}",
+            f"Pre: {pre_s.date()} to {pre_e.date()}  |  Post: {post_s.date()} to {post_e.date()}",
+            f"Frequency: {freq_input}",
+            f"Formula: {formula}",
+            f"Inference method: {inference_note}",  # MODIFIED: document inference method
+            "\n" + "=" * 60,
+            "MAIN RESULTS",
+            "=" * 60,
+            f"Observed post-intervention mean: {post_mean:.3f}",
+            f"Estimated counterfactual mean:  {cf_mean:.3f}",
+            f"**Estimated average effect:     {effect:.3f}**",
+        ]
+        if np.isfinite(eff_se):
+            report_lines += [
+                "\nStatistical inference:",
+                f"  Standard error:            {eff_se:.3f}",
+                f"  95% Confidence interval:  [{ci_low:.3f}, {ci_high:.3f}]",
+                test_line,
+                f"  p-value:                   {p_value:.4f}",
+            ]
+        else:
+            report_lines.append("\n(Statistical inference unavailable due to insufficient data)")
+
+        # Main ITS (CausalPy) figure
+        its = cp.InterruptedTimeSeries(
+            data=df, formula=formula, treatment_time=post_s, model=LinearRegression(), freq=freq_input
+        )
+        result = its.plot(plot_predict_all=False, plot_show_params=True)
+        fig_main = result[0] if isinstance(result, tuple) else result
+        # MODIFIED: make sure the ITS composite is large, rotate ticks & add spacing
+        try:
+            fig_main.set_size_inches(14, 9, forward=True)
+        except Exception:
+            pass
+        _rotate_all_xticklabels(fig_main)
+        try:
+            fig_main.tight_layout(pad=1.3)
+            fig_main.subplots_adjust(top=0.92, bottom=0.20, left=0.08, right=0.98, hspace=0.42)
+        except Exception:
+            pass
+
+        images: List[Image.Image] = [_fig_to_pil(fig_main)]
+
+        # Diagnostics + comparison
+        if run_diagnostics:
+            diag_figs = add_diagnostic_tests(sm_ols, pre_df, formula, report_lines)
+            if "acf" in diag_figs:
+                images.append(_fig_to_pil(diag_figs["acf"]))
+            if "residuals" in diag_figs:
+                images.append(_fig_to_pil(diag_figs["residuals"]))
+        else:
+            report_lines.append("\n(Diagnostic plots disabled)")
+
+        if show_formulas:
+            cmp_figs = compare_model_specifications(pre_df, formula, control_columns, report_lines)
+            if "comparison_plot" in cmp_figs:
+                images.append(_fig_to_pil(cmp_figs["comparison_plot"]))
+        else:
+            report_lines.append("\n(Model specification comparison disabled)")
+
+        final_img = images[0] if len(images) == 1 else _stack_images_vertical(images, pad=22)
+
+        return "\n".join(report_lines), final_img
+
+    except Exception as e:
+        return f"An unexpected error occurred: {e}\n{traceback.format_exc()}", None
+
+
+def run_its_analysis(file, target_col, date_col, pre_dates, post_dates, freq_input,
+                     control_vars="", run_diagnostics=True, show_formulas=False):
+    """Public entrypoint used by the UI."""
+    return enhanced_its_analysis(file, target_col, date_col, pre_dates, post_dates,
+                                 freq_input, control_vars, run_diagnostics, show_formulas)