#!/usr/bin/env python
# coding: utf-8

"""
c2regbase.py
A minimal Python library that automates a structured linear regression workflow 
in six well-defined functions.

The library provides:
    1. Variable overview and data type inspection
    2. Data type conversion utilities
    3. Missing-value analysis with listwise deletion
    4. Ordinary Least Squares (OLS) model fitting with standardized output
    5. Export of model parameters in a compact, dataset-friendly structure
    6. Diagnostic statistics, including standard errors, t-values, p-values, 
       and confidence interval bounds

Dependencies (minimal):
    - pandas
    - numpy
    - statsmodels
    - scipy    # only required if manual ANOVA computation is used
"""


from typing import List, Dict, Optional, Tuple, Any
import pandas as pd
import numpy as np
import statsmodels.api as sm
from statsmodels.formula.api import ols
from scipy import stats 



# ----------------------------
# Function 1: Variable Check
# ----------------------------
def print_dataset_structure(df: pd.DataFrame, max_rows_preview: int = 5) -> None:
    """
    Print dataset columns and their data types (like R's str() or pandas.dtypes).
    Args:
        df: pandas DataFrame (raw dataset)
        max_rows_preview: number of rows to preview after dtypes
    Output:
        Prints column names, dtype, non-null count, and a small preview (head).
    """
    print("=== Dataset Structure Summary ===")
    info_df = pd.DataFrame({
        "column": df.columns,
        "dtype": [str(dt) for dt in df.dtypes],
        "non_null_count": [int(df[c].notna().sum()) for c in df.columns]
    })
    print(info_df.to_string(index=False))
    print("\n--- DataFrame preview (top {} rows) ---".format(max_rows_preview))
    if len(df) > 0:
        print(df.head(max_rows_preview).to_string(index=False))
    else:
        print("(empty dataframe)")


# ----------------------------
# Function 2: User-Directed Type Conversion
# ----------------------------
def convert_variable_types(df: pd.DataFrame, conversions: Dict[str, str]) -> pd.DataFrame:
    """
    Convert specified variables to requested types.
    Args:
        df: pandas DataFrame
        conversions: dict mapping column_name -> target_type_string
                     allowed target_type_string examples:
                         'float', 'int', 'str', 'category', 'bool'
    Returns:
        A new DataFrame with converted types (copy).
    Notes:
        - Numeric conversions use pandas.to_numeric(..., errors='coerce').
        - Conversions are applied only to columns present in df.
        - No automatic feature-engineering (no creation of x^2 etc) — strictly type casting.
    """
    df2 = df.copy()
    for col, t in conversions.items():
        if col not in df2.columns:
            print(f"Warning: column '{col}' not in DataFrame; skipping conversion.")
            continue
        t = t.lower()
        if t in ('float', 'double'):
            df2[col] = pd.to_numeric(df2[col], errors='coerce').astype('float64')
        elif t in ('int', 'integer'):
            # coerce then convert to Int64 (nullable integer) to preserve NA support
            df2[col] = pd.to_numeric(df2[col], errors='coerce').astype('Int64')
        elif t in ('str', 'object', 'string'):
            df2[col] = df2[col].astype('object')
        elif t in ('category', 'cat'):
            df2[col] = df2[col].astype('category')
        elif t in ('bool', 'boolean'):
            # try to coerce common boolean-like values
            df2[col] = df2[col].map({True: True, False: False, 'True': True, 'False': False,
                                      'true': True, 'false': False, '1': True, '0': False}).astype('boolean')
        else:
            print(f"Warning: requested target type '{t}' for column '{col}' not recognized. Skipping.")
    return df2


# ----------------------------
# Function 3: Missing Value Summary & Listwise Deletion
# ----------------------------
def mv_summary_and_listwise_deletion(
    df: pd.DataFrame,
    variables_to_consider: Optional[List[str]] = None
) -> Tuple[pd.DataFrame, Dict[str, Any]]:
    """
    Provide summary of missing values and perform listwise deletion across specified variables.
    Args:
        df: pandas DataFrame (after type conversions)
        variables_to_consider: list of variables that will be considered "required".
            - If None: consider ALL columns in df.
            
    Returns:
        (df_clean, summary_dict)
        where:
            df_clean: DataFrame after listwise deletion (rows with any missing in variables_to_consider removed)
            summary_dict: {
                'rows_before': int,
                'rows_after': int,
                'rows_removed': int,
                'vars_considered': list,
                'per_variable_missing_count': dict,
                'rows_with_any_missing': int,
                'example_deleted_rows_indices': list (up to 10 indices)
            }
    """
    if variables_to_consider is None:
        variables_to_consider = list(df.columns)

    vars_present = [v for v in variables_to_consider if v in df.columns]
    missing_counts = df[vars_present].isna().sum().to_dict()

    rows_before = df.shape[0]
    rows_with_any_missing = int(df[vars_present].isna().any(axis=1).sum())
    df_clean = df.dropna(subset=vars_present).copy()
    rows_after = df_clean.shape[0]
    rows_removed = rows_before - rows_after

    # identify which variables caused deletions (i.e., had at least one missing)
    vars_causing_deletion = [v for v, cnt in missing_counts.items() if cnt > 0]

    # sample up to 10 deleted row indices for user's inspection
    deleted_indices = df[vars_present].index[df[vars_present].isna().any(axis=1)].tolist()
    example_deleted_indices = deleted_indices[:10]

    summary = {
        'rows_before': rows_before,
        'rows_after': rows_after,
        'rows_removed': rows_removed,
        'vars_considered': vars_present,
        'per_variable_missing_count': missing_counts,
        'rows_with_any_missing': rows_with_any_missing,
        'vars_causing_deletion': vars_causing_deletion,
        'example_deleted_row_indices': example_deleted_indices
    }

    # Print human-readable summary
    print("=== Missing Value Summary & Listwise Deletion ===")
    print(f"Rows before: {rows_before}")
    print(f"Rows after (post listwise deletion): {rows_after}")
    print(f"Rows removed: {rows_removed}")
    if vars_causing_deletion:
        print("Variables causing deletion (have missing values):")
        for v in vars_causing_deletion:
            print(f"  - {v}: {missing_counts[v]} missing")
    else:
        print("No missing values in variables considered.")
    if example_deleted_indices:
        print(f"Example deleted row indices (up to 10): {example_deleted_indices}")

    return df_clean, summary


# ----------------------------
# Function 4: Fit OLS Model
# ----------------------------
def fit_ols(
    df: pd.DataFrame,
    dependent: str,
    independents: List[str],
    add_intercept: bool = True,
    NumDigits: int = 6  # Number of decimal digits for printing
):
    """
    Fit OLS regression with standardized output and optional rounding for printed results.
    Returns:
        results (RegressionResultsWrapper)
        summary_stats (dict)
    """
    # Extract Y
    y = pd.to_numeric(df[dependent], errors='coerce').astype(float)

    # Build X from independents
    X_blocks = []
    for col in independents:
        s = df[col]
        if s.dtype.kind in "biufc":
            X_blocks.append(pd.to_numeric(s, errors='coerce').astype(float))
        else:
            dummies = pd.get_dummies(s, prefix=col, drop_first=True).astype(float)
            X_blocks.append(dummies)
    X = pd.concat(X_blocks, axis=1)

    # Add intercept
    if add_intercept:
        X = sm.add_constant(X, has_constant='add')

    X = X.astype(float)
    y = y.astype(float)

    if X.isna().any().any() or y.isna().any():
        raise ValueError("Missing values remain after encoding / conversion.")

    # Fit model
    model = sm.OLS(y, X)
    results = model.fit()

    # ---------------------------
    # Manual ANOVA
    # ---------------------------
    y_mean = y.mean()
    ss_total = np.sum((y - y_mean) ** 2)
    ss_model = np.sum((results.fittedvalues - y_mean) ** 2)
    ss_error = np.sum(results.resid ** 2)

    df_model = results.df_model
    df_error = results.df_resid
    df_total = df_model + df_error

    ms_model = ss_model / df_model
    ms_error = ss_error / df_error

    F_val = ms_model / ms_error
    p_val = 1 - stats.f.cdf(F_val, df_model, df_error)

    anova = pd.DataFrame({
        "SS": [ss_model, ss_error, ss_total],
        "DF": [df_model, df_error, df_total],
        "MS": [ms_model, ms_error, ""],
        "F": [F_val, "", ""],
        "Prob>F": [p_val, "", ""]
    }, index=["Model", "Error", "Total"])

    # -------------------------------------
    # Fit statistics
    # -------------------------------------
    rmse = np.sqrt(ms_error)
    coeff_var = rmse / y_mean * 100

    coef_df = pd.DataFrame({
        "parameter": results.params.index,
        "estimate": results.params.values,
        "stderr": results.bse,
        "t": results.tvalues,
        "pvalue": results.pvalues
    })

    # =========================
    # PRINT OUTPUT
    # =========================
    print("\n=== OLS OUTPUT ===")
    print(f"Dependent Variable: {dependent}")
    print(f"Number of Observations Used: {int(results.nobs)}")
    print(f"Root MSE: {rmse:.{NumDigits}f}")
    print(f"R-Square: {results.rsquared:.{NumDigits}f}")
    print(f"Adj R-Sq: {results.rsquared_adj:.{NumDigits}f}")
    print(f"Dependent Mean: {y_mean:.{NumDigits}f}")
    print(f"Coeff Var (%): {coeff_var:.{NumDigits}f}")
    print("\n--- Parameter Estimates ---")
    print(coef_df.round(NumDigits).to_string(index=False))
    print("\n--- ANOVA ---")
    print(anova)

    summary_stats = {
        "anova_table": anova,
        "coefficients_table": coef_df,
        "rmse": float(rmse),
        "r_squared": float(results.rsquared),
        "adj_r_squared": float(results.rsquared_adj),
        "dependent_mean": float(y_mean),
        "coeff_var_percent": float(coeff_var),
        "nobs": int(results.nobs),
        "df_resid": int(df_error),
    }

    return results, summary_stats

# ----------------------------
# Function 5: Output SAS like
# ----------------------------
def specific_output(
    results: Any,
    dependent: str,
    model_label: str = "Model1",
    NumDigits: int = 6
) -> pd.DataFrame:
    """
    Create an OUTEST-like DataFrame with coefficients and metadata.
    """
    params = results.params.copy()
    rmse = np.sqrt(results.mse_resid)
    row = params.to_dict()

    # Add metadata
    row.update({
        '_MODEL_': model_label,
        '_TYPE_': 'PARMS',
        '_DEPVAR_': -1,  # SAS marker
        '_RMSE_': float(rmse)
    })

    meta_cols = ['_MODEL_', '_TYPE_', '_DEPVAR_', '_RMSE_']
    param_cols = [c for c in params.index]
    out_cols = meta_cols + param_cols

    outest_df = pd.DataFrame([row], columns=out_cols)

    # Round numeric columns
    for col in params.index:
        outest_df[col] = outest_df[col].round(NumDigits)
    outest_df['_RMSE_'] = round(outest_df['_RMSE_'], NumDigits)

    return outest_df

# ----------------------------
# Function 6: Diagnostics
# ----------------------------
def summary_table(
    results: Any,
    alpha: float = 0.05,
    NumDigits: int = 6
) -> pd.DataFrame:
    """
    Compute diagnostics (STDERR, T, PVALUE, CI) with optional rounding.
    """
    params = results.params
    bse = results.bse
    tvals = results.tvalues
    pvals = results.pvalues
    df_resid = getattr(results, 'df_resid', None)

    try:
        ci = results.conf_int(alpha=alpha)
        lnb_vals = ci.iloc[:, 0].values
        unb_vals = ci.iloc[:, 1].values
    except Exception:
        tcrit = stats.t.ppf(1 - alpha / 2, df_resid)
        lnb_vals = params - tcrit * bse
        unb_vals = params + tcrit * bse

    diagnostics_df = pd.DataFrame({
        'parameter': params.index,
        'estimate': params.values,
        'stderr': bse.values,
        't_stat': tvals.values,
        'pvalue': pvals.values,
        'lnb': lnb_vals,
        'unb': unb_vals
    }).reset_index(drop=True)

    # Round for printing
    numeric_cols = ['estimate','stderr','t_stat','pvalue','lnb','unb']
    diagnostics_df_rounded = diagnostics_df.copy()
    diagnostics_df_rounded[numeric_cols] = diagnostics_df_rounded[numeric_cols].round(NumDigits)

    print("=== Summaries ===")
    print(f"Alpha level for CI: {alpha}")
    print(diagnostics_df_rounded.to_string(index=False))

    return diagnostics_df
