"""Bayesian linear instrumental-variable estimation via PyMC.

Jointly models the first stage (``D ~ Z + X``) and the structural
equation (``Y ~ D + X``) with a bivariate-Normal error structure so
the posterior over the LATE prices endogeneity + weak-instrument risk
automatically.
"""
from __future__ import annotations

from typing import List, Optional, Sequence, Tuple, Union

import numpy as np
import pandas as pd

from ._base import (
    BayesianCausalResult,
    BayesianIVResult,
    _require_pymc,
    _sample_model,
    _summarise_posterior,
)


def _prepare_iv_frame(
    data: pd.DataFrame,
    y: str,
    treat: str,
    instrument: Union[str, Sequence[str]],
    covariates: Optional[List[str]],
) -> dict:
    """Validate, drop NA, extract arrays for Bayesian IV."""
    if isinstance(instrument, str):
        iv_cols = [instrument]
    else:
        iv_cols = list(instrument)

    for c in [y, treat] + iv_cols:
        if c not in data.columns:
            raise ValueError(f"Column '{c}' not found in data")

    cov_cols = list(covariates) if covariates else []
    for c in cov_cols:
        if c not in data.columns:
            raise ValueError(f"Covariate '{c}' not found in data")

    all_cols = [y, treat] + iv_cols + cov_cols
    clean = data[all_cols].dropna().reset_index(drop=True)
    n = len(clean)
    if n < 30:
        raise ValueError(
            f"Bayesian IV needs at least 30 observations after dropping NA, "
            f"got {n}."
        )

    Y = clean[y].to_numpy(dtype=float)
    D = clean[treat].to_numpy(dtype=float)
    Z = clean[iv_cols].to_numpy(dtype=float)
    X = clean[cov_cols].to_numpy(dtype=float) if cov_cols else None

    return {
        'n': n,
        'Y': Y,
        'D': D,
        'Z': Z,
        'X': X,
        'iv_cols': iv_cols,
        'cov_cols': cov_cols,
    }


def bayes_iv(
    data: pd.DataFrame,
    y: str,
    treat: str,
    instrument: Union[str, Sequence[str]],
    covariates: Optional[List[str]] = None,
    *,
    per_instrument: bool = False,
    prior_late: Tuple[float, float] = (0.0, 10.0),
    prior_first_stage_sigma: float = 5.0,
    prior_coef_sigma: float = 10.0,
    prior_noise: float = 5.0,
    rope: Optional[Tuple[float, float]] = None,
    hdi_prob: float = 0.95,
    inference: str = 'nuts',
    advi_iterations: int = 20000,
    draws: int = 2000,
    tune: int = 1000,
    chains: int = 4,
    target_accept: float = 0.9,
    random_state: int = 42,
    progressbar: bool = False,
) -> BayesianIVResult:
    """Bayesian linear IV via jointly-modelled first stage + structural equation.

    The model:

    .. code-block:: text

        D_i = pi_0 + pi_Z' * Z_i + pi_X' * X_i + v_i
        Y_i = alpha + LATE * D_i + beta_X' * X_i + eps_i
        (v_i, eps_i) ~ BivariateNormal(0, Sigma)

    ``Sigma`` is parameterised via an LKJ prior on the correlation
    matrix and HalfNormal priors on the two scales, which lets the
    model identify the LATE from *exogenous* variation in Z even when
    D is endogenous. Under a weak instrument (``pi_Z ≈ 0``) the LATE
    posterior correctly widens — there's no "weak-instrument F < 10"
    footgun here; the posterior just gets more uncertain.

    Parameters
    ----------
    data : pd.DataFrame
    y : str
        Outcome column.
    treat : str
        Endogenous treatment / regressor column (continuous or binary).
    instrument : str or sequence of str
        One or more instruments. Must be excluded from the structural
        equation.
    covariates : list of str, optional
        Exogenous controls entering both stages.
    per_instrument : bool, default ``False``
        When ``True`` and multiple instruments are supplied, additionally
        fits one just-identified Bayesian IV sub-model per instrument
        and populates :attr:`BayesianIVResult.instrument_summaries`,
        letting ``tidy(terms='per_instrument')`` emit one LATE row per
        ``Z_j``. The top-level pooled LATE posterior remains the joint
        over-identified fit (v0.9.15 behaviour). Each per-instrument
        sub-fit reuses the same priors and sampler controls as the
        pooled fit (draws/tune/chains/target_accept/random_state), so
        runtime scales roughly as ``(K+1)×`` the pooled fit.
    prior_late : (float, float)
        Normal prior on the structural LATE coefficient.
    prior_first_stage_sigma, prior_coef_sigma, prior_noise : float
        Priors for first-stage coefficients, structural coefficients,
        and the two residual scales.
    rope : (float, float), optional
        Region of practical equivalence.
    hdi_prob, draws, tune, chains, target_accept, random_state, progressbar :
        Sampler controls — see :func:`bayes_did`.

    Returns
    -------
    BayesianCausalResult
        Posterior summary on the LATE coefficient.
    """
    pm, _ = _require_pymc()

    prep = _prepare_iv_frame(data, y, treat, instrument, covariates)
    n = prep['n']
    Y = prep['Y']
    D = prep['D']
    Z = prep['Z']
    X = prep['X']
    n_instr = Z.shape[1]

    mu_late, sigma_late = prior_late

    # Control-function pre-compute: first-stage residuals from a quick
    # OLS on (1, Z, X). Feeding the plug-in residuals into the
    # structural equation as an exogeneity control makes the LATE
    # posterior coincide (asymptotically) with 2SLS while giving us a
    # tractable single-likelihood model that PyMC can sample cleanly.
    W_fs = [np.ones((n, 1)), Z]
    if X is not None:
        W_fs.append(X)
    W_fs_mat = np.hstack(W_fs)
    pi_ols, *_ = np.linalg.lstsq(W_fs_mat, D, rcond=None)
    D_hat = W_fs_mat @ pi_ols
    v_hat = D - D_hat  # first-stage residuals

    with pm.Model() as model:
        # ------------------------------------------------------------------
        # First stage: D ~ Z + X  (modelled so we have a posterior on
        # pi_Z, but the residuals used in the control function are the
        # plug-in OLS residuals above, which is a valid 2SLS-equivalent
        # simplification under homoskedasticity and a linear first stage)
        # ------------------------------------------------------------------
        pi_intercept = pm.Normal('pi_intercept', mu=0.0, sigma=prior_coef_sigma)
        pi_Z = pm.Normal(
            'pi_Z', mu=0.0, sigma=prior_coef_sigma, shape=n_instr,
        )
        first_stage = pi_intercept + pm.math.dot(Z, pi_Z)
        if X is not None:
            pi_X = pm.Normal(
                'pi_X', mu=0.0, sigma=prior_coef_sigma, shape=X.shape[1],
            )
            first_stage = first_stage + pm.math.dot(X, pi_X)
        sigma_v = pm.HalfNormal('sigma_v', sigma=prior_noise)
        pm.Normal('d_obs', mu=first_stage, sigma=sigma_v, observed=D)

        # ------------------------------------------------------------------
        # Structural: Y ~ D + X + rho * v_hat  (control function)
        # ------------------------------------------------------------------
        alpha = pm.Normal('alpha', mu=0.0, sigma=prior_coef_sigma)
        late = pm.Normal('late', mu=mu_late, sigma=sigma_late)
        rho = pm.Normal('rho_cf', mu=0.0, sigma=prior_coef_sigma)
        structural = alpha + late * D + rho * v_hat
        if X is not None:
            beta_X = pm.Normal(
                'beta_X', mu=0.0, sigma=prior_coef_sigma, shape=X.shape[1],
            )
            structural = structural + pm.math.dot(X, beta_X)
        sigma_eps = pm.HalfNormal('sigma_eps', sigma=prior_noise)
        pm.Normal('y_obs', mu=structural, sigma=sigma_eps, observed=Y)

    trace = _sample_model(
        model,
        inference=inference,
        draws=draws,
        tune=tune,
        chains=chains,
        target_accept=target_accept,
        random_state=random_state,
        progressbar=progressbar,
        advi_iterations=advi_iterations,
    )

    summary = _summarise_posterior(
        trace, 'late', hdi_prob=hdi_prob, rope=rope,
    )

    model_info = {
        'inference': inference,
        'draws': draws,
        'tune': tune,
        'chains': chains,
        'target_accept': target_accept,
        'prior_late': prior_late,
        'prior_first_stage_sigma': prior_first_stage_sigma,
        'prior_coef_sigma': prior_coef_sigma,
        'prior_noise': prior_noise,
        'instruments': prep['iv_cols'],
        'covariates': prep['cov_cols'],
        'n_instruments': n_instr,
    }

    method_label = (
        f"Bayesian IV (joint 2SLS, {n_instr} instrument"
        f"{'s' if n_instr > 1 else ''})"
    )

    instrument_summaries: dict = {}
    instrument_labels: list = []
    # Per-instrument just-identified LATEs. Only meaningful with
    # multiple instruments — silently skip otherwise and return the
    # pooled result as v0.9.15 did, so passing per_instrument=True
    # with a single Z doesn't error (the single LATE IS the
    # per-instrument LATE). With K>=2 we loop the pooled fit with one
    # Z at a time; this is slow but transparent — each sub-fit is a
    # plain bayes_iv call under the hood.
    if per_instrument and n_instr >= 2:
        for z_name in prep['iv_cols']:
            sub = bayes_iv(
                data,
                y=y,
                treat=treat,
                instrument=z_name,        # scalar path = just-identified
                covariates=covariates,
                per_instrument=False,     # guard against recursion
                prior_late=prior_late,
                prior_first_stage_sigma=prior_first_stage_sigma,
                prior_coef_sigma=prior_coef_sigma,
                prior_noise=prior_noise,
                rope=rope,
                hdi_prob=hdi_prob,
                inference=inference,
                advi_iterations=advi_iterations,
                draws=draws,
                tune=tune,
                chains=chains,
                target_accept=target_accept,
                random_state=random_state,
                progressbar=progressbar,
            )
            instrument_summaries[z_name] = {
                'posterior_mean': sub.posterior_mean,
                'posterior_median': sub.posterior_median,
                'posterior_sd': sub.posterior_sd,
                'hdi_lower': sub.hdi_lower,
                'hdi_upper': sub.hdi_upper,
                'prob_positive': sub.prob_positive,
            }
        instrument_labels = list(prep['iv_cols'])
        method_label += f" + {len(instrument_labels)} per-Z sub-fits"

    return BayesianIVResult(
        method=method_label,
        estimand='LATE',
        posterior_mean=summary['posterior_mean'],
        posterior_median=summary['posterior_median'],
        posterior_sd=summary['posterior_sd'],
        hdi_lower=summary['hdi_lower'],
        hdi_upper=summary['hdi_upper'],
        prob_positive=summary['prob_positive'],
        prob_rope=summary.get('prob_rope'),
        rhat=summary['rhat'],
        ess=summary['ess'],
        n_obs=n,
        hdi_prob=hdi_prob,
        trace=trace,
        model_info=model_info,
        instrument_summaries=instrument_summaries,
        instrument_labels=instrument_labels,
    )
