from __future__ import annotations

import sys
import time
from typing import Tuple, List

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import statsmodels.api as sm
from statsmodels.sandbox.regression.predstd import wls_prediction_std
from . .queries import *
from . .regression import *
from . .figures import *
import time

# -----------------------------------------------------------------------------
# Utilities expected to exist in the caller's namespace
# -----------------------------------------------------------------------------
# * append_to_log(msg)   → writes to the session log
# * neg_response(txt)    → True if txt is a negative user response
# * pos_response(txt)    → True if txt is a positive user response
# * wls_regression(x, y, log_file_path, weights=None)  → robust WLS helper that
#     returns  (slope, intercept, adj_r2, p, std_err, model)
# -----------------------------------------------------------------------------

REQ_COLS_STD: List[str] = ["time_rel", "area", "chain"]

# -----------------------------------------------------------------------------
# 1 – user prompt
# -----------------------------------------------------------------------------

def q_drift(log_file_path: str) -> str:
    """Prompt the user whether to **apply** a drift correction.

    Returns the raw lower‑case response (``'y'`` | ``'n'``).
    Continues prompting until a valid flag is entered.
    """
    valid = {"yes", "y", "true", "t", "no", "n", "false", "f"}

    while True:
        choice = input("\nApply drift correction? (Y/N)\n").strip().lower()
        if choice in valid:
            append_to_log(log_file_path, f"- Drift correction applied: {choice}")
            return choice[0]  # just 'y' or 'n'
        print("\n[ERROR] Invalid response – please type Y or N.\n")

# -----------------------------------------------------------------------------
# 2 – public wrapper (kept mostly for backward compatibility)
# -----------------------------------------------------------------------------

# def process_drift_correction(
#     cfg: CorrectionConfig,
#     samp: pd.DataFrame,
#     lin_std: pd.DataFrame,
#     drift_std: pd.DataFrame,
#     correction_log: pd.DataFrame,
#     fig_path: str,
#     log_file_path: str,
#     isotope: str = "dD") -> Tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame, str, pd.DataFrame]:
#     """Thin wrapper that asks the user, runs :pyfunc:`drift_correction`,
#     and updates *correction_log*.
#     """
#     append_to_log(log_file_path, f"Drift Correction")
#     if q_drift(log_file_path) == "n":
#         # No correction
#         for df in (samp, lin_std, drift_std):
#             df["drift_corrected_dD"] = df.get("dD", np.nan)
#         return samp, lin_std, drift_std, "dD", correction_log

#     # Yes correction
#     if result is None:  # user aborted inside drift_correction
#         print("[INFO] Drift correction NOT applied – data reverted.")
#         for df in (samp, lin_std, drift_std):
#         return samp, lin_std, drift_std, "dD", correction_log

#     samp, lin_std, drift_std = result
#     correction_log.loc["Drift", "samples"] = 1
#     return samp, lin_std, drift_std, "drift_corrected_dD", correction_log
def process_drift_correction(
    cfg: CorrectionConfig,
    samp: pd.DataFrame,
    lin_std: pd.DataFrame,
    drift_std: pd.DataFrame,
    correction_log: pd.DataFrame,
    fig_path: str,
    log_file_path: str,
    isotope: str = "dD"
) -> Tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame, str, pd.DataFrame]:
    """Thin wrapper that asks the user, runs :py:func:`drift_correction`,
    updates *cfg* and also writes to *correction_log*."""
    append_to_log(log_file_path, "Drift Correction")

    # ask Y/N
    choice = q_drift(log_file_path).strip().lower()
    if choice in ("n", "no", "false", "f"):
        # mark config
        cfg.drift_applied = False

        # fill a drift_corrected_dD column with raw dD
        for df in (samp, lin_std, drift_std):
            df["drift_corrected_dD"] = df.get("dD", np.nan)
        return samp, lin_std, drift_std, "dD", correction_log

    # user said “yes” → do the heavy lifting
    result = drift_correction(
        samp, lin_std, drift_std,
        fig_path, isotope, log_file_path, fig=True
    )
    if result is None:
        # user aborted inside drift_correction()
        cfg.drift_applied = False
        print("[INFO] Drift correction NOT applied – data reverted.")
        for df in (samp, lin_std, drift_std):
            df["drift_corrected_dD"] = df.get("dD", np.nan)
        return samp, lin_std, drift_std, "dD", correction_log

    # drift actually applied!
    samp, lin_std, drift_std = result
    cfg.drift_applied = True                # <— flip the flag here
    correction_log.loc["Drift", "samples"] = 1

    return samp, lin_std, drift_std, "drift_corrected_dD", correction_log
# -----------------------------------------------------------------------------
# 3 – core routine
# -----------------------------------------------------------------------------

def _require_cols(df: pd.DataFrame, name: str, cols: List[str]):
    missing = [c for c in cols if c not in df.columns]
    if missing:
        raise ValueError(f"{name} is missing required column(s): {', '.join(missing)}")


def drift_correction(
    samp: pd.DataFrame,
    lin_std: pd.DataFrame,
    drift_std: pd.DataFrame,
    fig_path: str,
    isotope: str,
    log_file_path: str,
    fig: bool = True):
    """Apply a *weighted* time‑drift correction.

    The regression weight is proportional to peak *area* so that small peaks
    contribute less to the slope estimation.
    """
    # 0. Validation & set‑up
    for df, name in ((samp, "samples"), (lin_std, "lin_std"), (drift_std, "drift_std")):
        _require_cols(df, name, REQ_COLS_STD + [isotope])
    samp, lin_std, drift_std = samp.copy(), lin_std.copy(), drift_std.copy()

    # 1. Chain‑wise centring of isotope value (removes constant offset)
    norm_parts = []
    for chain in drift_std["chain"].unique():
        block = drift_std.loc[drift_std["chain"] == chain, [isotope, "time_rel", "area", "chain"]].copy()
        block[isotope] -= block[isotope].mean()
        norm_parts.append(block)
    norm = pd.concat(norm_parts, ignore_index=True)

    # 2. Global centring of time axis
    t_mean = norm["time_rel"].mean()
    for df in (norm, drift_std, lin_std, samp):
        df["t_ctr"] = df["time_rel"] - t_mean

    # 3. Weighted least‑squares fit  δ = b₀ + b₁·t_ctr
    #    Weight wᵢ = areaᵢ / max(area)  (scaled 0–1)
    w = norm["area"]**2 / norm["area"].max()
    slope, intercept, r2, p, se, model = wls_regression(
        norm["t_ctr"], norm[isotope], log_file_path, weights=w)
    append_to_log(log_file_path, f"- WLS fit completed: b0={intercept:.4f}, b1={slope:.4f}, R²={r2:.3f}")

    # 4. Predict + 1σ band
    X_drift = sm.add_constant(drift_std["t_ctr"])
    pred_drift = model.get_prediction(X_drift)
    drift_pred = pred_drift.predicted_mean
    drift_sd   = pred_drift.se_obs * 2 # 2‑sigma (already weighted)

    # 5. Anchor – average of last three standards
    last3_idx = drift_std["time_rel"].nlargest(3).reset_index().index
    anchor = drift_pred[last3_idx].mean()
    drift_offset = drift_pred - anchor
    norm['corrected_norm'] = norm[isotope]-drift_pred
    drift_std["drift_corrected_dD"] = drift_std[isotope] - drift_offset
    drift_std["drift_error"]        = drift_sd

    # 6. Propagate to lin_std and samples (vectorised)
    def _apply(df: pd.DataFrame, name: str):
        X = sm.add_constant(df["t_ctr"])
        pred = model.predict(X)
        sd   = wls_prediction_std(model, exog=X)[0]  # returns prstd
        df["drift_corrected_dD"] = df[isotope] - (pred - anchor)
        df["drift_error"]        = sd
        append_to_log(log_file_path, f"- Drift correction applied to {name}: n={len(df)}")

    _apply(lin_std, "lin_std")
    _apply(samp, "samples")

    # ---------------------------------------------------------------------
    # 7. Plot diagnostic (optional)
    # ---------------------------------------------------------------------
    if fig:
        drift_std_corr(norm, isotope, drift_std, t_mean, intercept, slope, fig_path)

    # ---------------------------------------------------------------------
    # 8. Ask user to confirm
    # ---------------------------------------------------------------------
    print(f"\nWeighted‑LS model: {isotope} = {slope:.3f}·t_ctr + {intercept:.3f}")
    print(f"Adj R² = {r2:.3f} | p = {p:.3g} | SE = {se:.3f}")

    # if not sys.stdin.isatty():  # non‑interactive run – auto‑apply
    #     append_to_log(log_file_path, "Non‑interactive session – drift correction auto‑applied.")
    #     return samp, lin_std, drift_std

    if pos_response(input("Apply the correction? (Y/N): ")):
        append_to_log(log_file_path, "- User accepted drift correction.")
        append_to_log(log_file_path, f"- Weighted‑LS model: {isotope} = {slope:.3f} t_ctr + {intercept:.3f}")
        append_to_log(log_file_path, f"- Weighted‑LS model stats: Adj R² = {r2:.3f} | p = {p:.3g} | SE = {se:.3f}")
        return samp, lin_std, drift_std

    append_to_log(log_file_path, "- User *rejected* drift correction. No changes saved.")
    return None

