# -*- coding: utf-8 -*-
"""AI 확산을 산업모형(베이지언 VARX, 62분류·로그차분·산업연관표 사전분포)의 외생 항으로 넣는 대안 1·2.

공통
  산업 i의 분기 성장률(100·Δlog VA)에 AI 항 γ_i·x_t를 더한다. x_t는 AI 확산 시나리오가 기준 경로보다 앞선 정도이며,
  미래에는 설문 시점 가중치 r(t)(2025년 0 → 2040년 1)를 분기로 보간한 경로를 쓴다(현행 방식과 같은 시간 경로).
  AI 효과는 VAR 동학을 따라 다른 산업으로 퍼진다:  δ_t = Σ_l A_lᵀ δ_{t−l} + γ ⊙ x_t,  D_t = Σ δ (로그 수준 편차).
  연간 AI 확산 부가가치 = 기준 부가가치 × (그해 4개 분기 exp(D/100)의 평균).
  크기는 현행과 같이 2035년 부가가치 가중평균 연간 프리미엄 = Λ(0.5%p)가 되도록 γ 전체에 공통 배율을 곱해 맞춘다.
  따라서 대안과 현행의 차이는 산업 간 배분과 파급에서만 생긴다.

대안 1(고정 계수 위 추가)
  보고서1의 계수(A_l)를 그대로 쓰고 AI 계수는 사전 정보로 정한다(추정하지 않음). 기준 시나리오에서는 x = 0이므로 기존 전망이 바뀌지 않는다.
  설문은 2040년까지의 누적 효과를 묻으므로 사전 정보는 모형의 장기 반응에 둔다: 동시항 γ = (I − ΣA_lᵀ)·shape (impact_from_long_run).
  1a: shape = 설문 강도 a_i(E4-1 2040년 − 1),  1b: shape = 특허·논문·뉴스 노출 지수를 설문 척도로 옮긴 값.
  shape를 동시항에 그대로 넣는 사양(진단)은 축약형 동학의 교차계수를 통해 일부 산업에 음(−)의 효과를 만든다(spillover_decomp).

대안 2(일곱째 외생으로 재추정)
  과거 AI 확산 갭(총량 지수 − HP 추세, 2015년 이전 0)을 일곱째 외생으로 넣어 다시 추정한다.
  AI 동시항 계수의 사전분포: γ_i ~ N(β·a_i, τ_i²).  β는 설문 배분을 평균으로 두고 자료로 정한 공통 척도(경험적 베이즈),
  τ_i² = (β·SE_i)² [+ τ̂²], SE_i는 원자료에서 구한 업종 평균 E4-1의 표준오차(설문 표본 불확실성), τ̂²는 계층 모형의 적률 추정치이다.
  AI 시차항(1·2기)은 평균 0·분산 매우 작게 두어 사실상 뺀다. 미래 기준 시나리오에서는 x = 0이다.
"""
from __future__ import annotations

import numpy as np
import pandas as pd

import varx_engine as V

AI_COL = "AI확산갭"


# ── 공통: 시간 경로와 전파 ───────────────────────────────────────────────
def quarterly_ramp(ramp) -> pd.Series:
    """연간 시점 가중치 함수 r(t)를 분기 중앙 시점으로 보간(2026Q1–2035Q4)."""
    return pd.Series({q: float(ramp(int(q[:4]) + (int(q[-1]) - 0.5) / 4.0)) for q in V.FUT})


def lag_blocks(B: np.ndarray, n: int, p: int = V.P) -> list[np.ndarray]:
    """계수행렬에서 내생 시차 블록 A_l(n×n, [j, i] = 산업 j의 시차가 방정식 i에 주는 계수)."""
    return [B[(l - 1) * n:l * n, :] for l in range(1, p + 1)]


def propagate(A: list[np.ndarray], gamma: np.ndarray, x: pd.Series) -> pd.DataFrame:
    """δ_t = Σ_l A_lᵀ δ_{t−l} + γ x_t 를 2026Q1부터 재귀 계산해 로그 수준 편차 D_t(100·log)를 돌려준다."""
    n = len(gamma)
    hist = [np.zeros(n) for _ in A]
    D = np.zeros(n)
    out = {}
    for q in x.index:
        d = sum(A[l].T @ hist[l] for l in range(len(A))) + gamma * float(x[q])
        hist = [d] + hist[:-1]
        D = D + d
        out[q] = D.copy()
    return pd.DataFrame(out).T


def annual_factor(Dq: pd.DataFrame) -> pd.DataFrame:
    yr = pd.Series([int(q[:4]) for q in Dq.index], index=Dq.index)
    return np.exp(Dq / 100.0).groupby(yr).mean()


def ai_path(base: pd.DataFrame, A, gamma_shape: np.ndarray, x: pd.Series, lam: float, year: int = 2035):
    """γ 형태를 Λ로 보정해 AI 확산 경로(연간)와 산업별 2035년 프리미엄을 구한다. 반환: (AI VA, 프리미엄, 배율)."""
    cols = list(base.columns)
    f0 = annual_factor(propagate(A, gamma_shape, x))
    f0.columns = cols
    prem0 = 100.0 * np.log(f0.loc[year] / f0.loc[year - 1])
    w = base.loc[year] / base.loc[year].sum()
    c = lam / float((w * prem0).sum())
    f = annual_factor(propagate(A, c * gamma_shape, x))
    f.columns = cols
    ai = base.loc[2025:2035].copy()
    for y in range(2026, 2036):
        ai.loc[y] = base.loc[y] * f.loc[y]
    prem = 100.0 * np.log(f.loc[year] / f.loc[year - 1])
    return ai, prem, c


def direct_only_premium(base: pd.DataFrame, gamma_shape: np.ndarray, x: pd.Series, lam: float, year: int = 2035):
    """파급 없이(A = 0) 직접 효과만 둔 경우 — 현행 방식과 같은 구조의 점검용."""
    n = len(gamma_shape)
    return ai_path(base, [np.zeros((n, n))], gamma_shape, x, lam, year)


# ── 기존 모형 적합 ──────────────────────────────────────────────────────
def fit_existing(levels_csv, exog_hist_csv, io_csv):
    """보고서1 선택 모형(62분류·차분·산업연관표·VARX)을 전 표본으로 적합한다(5장 패키지 varx_engine와 같은 계산)."""
    return V.fit_full(levels_csv, exog_hist_csv, io_csv, spec="차분", prior="io", em="varx")


# ── 대안 2: AI 외생을 넣은 재추정 ────────────────────────────────────────
def design_with_ai(L: pd.DataFrame, exo: pd.DataFrame, ai_gap: pd.Series):
    Y = L.diff().dropna()
    ex = exo.copy()
    ex[AI_COL] = ai_gap.reindex(ex.index).fillna(0.0)
    exs = ex.loc[Y.index, V.EXOG + [AI_COL]]
    Ym, X, groups = V.design(Y, exs, True, None)
    return Ym, X, groups, Y


def fit_with_ai(L, exo, Wio, ai_gap, prior_mean_ai: np.ndarray | None, prior_var_ai: np.ndarray | None,
                lag_var: float = 1e-8):
    """AI 외생을 일곱째로 넣고 방정식별 사후평균을 구한다. AI 동시항 사전분포 N(prior_mean_ai, prior_var_ai).

    prior_mean_ai/var_ai가 None이면 다른 외생과 같은 기본 사전(평균 0, 분산 4λ²)을 쓴다(β 추정용).
    """
    Ym, X, groups, Y = design_with_ai(L, exo, ai_gap)
    n = L.shape[1]
    sig = V.ar_sigma(Y)
    Vp = V.prior_var(groups, sig, n, "io", Wio)
    ai_idx = len(V.EXOG)                                   # 외생 열 순서에서 AI의 위치
    k = X.shape[1]
    B = np.zeros((k, n))
    rows_ai = {s: c for c, (kind, s, j) in enumerate(groups) if kind == "exog" and j == ai_idx}
    XtX = X.T @ X
    for i in range(n):
        v = Vp[:, i].copy()
        b0 = np.zeros(k)
        for s, c in rows_ai.items():
            if s == 0:
                if prior_mean_ai is not None:
                    b0[c] = prior_mean_ai[i]
                    v[c] = prior_var_ai[i]
            else:
                v[c] = lag_var                             # AI 시차항은 사실상 0으로 고정
        s2 = sig[i] ** 2
        iV0 = np.diag(1.0 / v)
        B[:, i] = np.linalg.solve(iV0 + XtX / s2, iV0 @ b0 + (X.T @ Ym[:, i]) / s2)
    se_like = None
    return B, groups, rows_ai[0], sig


def ai_coef_and_precision(L, exo, Wio, ai_gap, diffuse_var: float = 1e6):
    """AI 동시항 계수에 사실상 정보가 없는 사전(평균 0, 분산 1e6)을 두고 자료만으로 정해지는 계수와 그 정밀도를 구한다.
    다른 계수는 기존 사전분포를 그대로 쓴다. 정밀도는 방정식별 사후공분산 (V0⁻¹ + X'X/σ²)⁻¹의 AI 대각원소의 역수로,
    AI 갭과 다른 설명변수(외생·시차) 사이의 상관을 반영한다. 경험적 베이즈(β, τ²)와 '자료의 비중' 계산에 쓴다."""
    n = L.shape[1]
    B, groups, r0, sig = fit_with_ai(L, exo, Wio, ai_gap, np.zeros(n), np.full(n, diffuse_var))
    Ym, X, groups, Y = design_with_ai(L, exo, ai_gap)
    Vp = V.prior_var(groups, sig, n, "io", Wio)
    XtX = X.T @ X
    prec = np.zeros(n)
    for i in range(n):
        v = Vp[:, i].copy()
        for c, (kind, s, j) in enumerate(groups):
            if kind == "exog" and j == len(V.EXOG):
                v[c] = diffuse_var if s == 0 else 1e-8
        cov = np.linalg.inv(np.diag(1.0 / v) + XtX / sig[i] ** 2)
        prec[i] = 1.0 / cov[r0, r0]
    return B[r0, :], prec


def empirical_bayes_beta(g_hat: np.ndarray, prec: np.ndarray, a: np.ndarray) -> tuple[float, float]:
    """γ̂_i ≈ β·a_i 의 가중최소제곱(가중치 = 자료 정밀도) — β 추정치와 표준오차."""
    w = prec
    beta = float((w * a * g_hat).sum() / (w * a * a).sum())
    se = float(np.sqrt(1.0 / (w * a * a).sum()))
    return beta, se


def hierarchical_tau2(g_hat: np.ndarray, prec: np.ndarray, a: np.ndarray, beta: float) -> float:
    """γ̂_i = β·a_i + u_i + e_i,  u_i ~ N(0, τ²)(설문으로 설명되지 않는 산업별 차이),  e_i ~ N(0, 1/정밀도_i)(추정 오차).
    적률법: τ² = max(0, 평균[(γ̂ − β·a)² − 1/정밀도]). 0이면 자료의 흩어짐이 추정 오차만으로 설명된다는 뜻이다."""
    r = g_hat - beta * a
    return max(0.0, float(np.mean(r ** 2 - 1.0 / prec)))


def shrinkage_weight(prec: np.ndarray, prior_var: np.ndarray) -> np.ndarray:
    """사후평균에서 자료가 차지하는 비중 = 정밀도 / (정밀도 + 1/사전분산). 0이면 사후 = 사전, 1이면 사후 = 자료."""
    return prec / (prec + 1.0 / prior_var)


def long_run_multiplier(A: list[np.ndarray]) -> np.ndarray:
    """영구적인 성장률 충격의 장기 파급 배수 M = (I − Σ_l A_lᵀ)⁻¹. M[i, j]는 산업 j의 직접 효과 1이 산업 i 성장률을 바꾸는 크기."""
    n = A[0].shape[0]
    return np.linalg.inv(np.eye(n) - sum(Al.T for Al in A))


def impact_from_long_run(A: list[np.ndarray], long_run: np.ndarray) -> np.ndarray:
    """장기 효과가 long_run이 되도록 하는 동시항 계수 γ = M⁻¹·long_run = (I − Σ_l A_lᵀ)·long_run.
    설문(E4-1)은 2040년까지의 누적 효과를 묻는다. 그래서 사전 정보는 동시항 γ가 아니라 모형의 장기 반응 M·γ에 둔다."""
    n = len(long_run)
    return (np.eye(n) - sum(Al.T for Al in A)) @ long_run


def spillover_decomp(A: list[np.ndarray], gamma: np.ndarray, names: list[str]) -> pd.DataFrame:
    """장기 배수로 본 산업별 AI 효과 = 자기 몫(M_ii γ_i) + 다른 산업에서 온 몫(Σ_{j≠i} M_ij γ_j). 가장 크게 끌어내린 산업도 적는다."""
    M = long_run_multiplier(A)
    C = M * gamma[None, :]
    own = np.diag(C).copy()
    other = C.sum(axis=1) - own
    Coff = C.copy(); np.fill_diagonal(Coff, 0.0)
    jmin = Coff.argmin(axis=1); jmax = Coff.argmax(axis=1)
    return pd.DataFrame({"직접": gamma, "자기몫_장기": own, "파급몫_장기": other, "합계_장기": own + other,
                         "가장_큰_음의_파급원": [names[j] for j in jmin], "그_크기": Coff[np.arange(len(names)), jmin],
                         "가장_큰_양의_파급원": [names[j] for j in jmax]}, index=names)


def forecast_s1(L, gdp, exo, B, exo_future, anchor_growth: float, with_ai: bool):
    """S1 외생 경로로 2026–2035 분기 전망(드리프트 재설정 포함) → 연간 합계(십억 원)."""
    ex_hist = exo.copy()
    ex_fut = exo_future.copy()
    if with_ai:
        ex_fut[AI_COL] = 0.0                               # 미래 기준 시나리오: AI 갭 0
    n = L.shape[1]
    B2, _ = V.anchor_drift(L, B, anchor_growth)
    cols_ex = V.EXOG + ([AI_COL] if with_ai else [])
    efull = pd.concat([ex_hist[cols_ex], ex_fut[cols_ex]])
    Gv = {q: L.diff().loc[q].values.astype(float) for q in L.index}
    Lv = {q: L.loc[q].values.astype(float) for q in L.index}
    allq = list(L.index) + V.FUT
    for t in V.FUT:
        ti = allq.index(t)
        cols = [Gv[allq[ti - l]] for l in range(1, V.P + 1)]
        cols += [efull.loc[allq[ti - s]].values.astype(float) for s in range(0, V.Q + 1)]
        cols.append(np.array([1.0]))
        yhat = np.concatenate(cols) @ B2
        Lv[t] = Lv[allq[ti - 1]] + yhat
        Gv[t] = yhat
    fq = pd.DataFrame({t: np.exp(Lv[t] / 100.0) for t in V.FUT}, index=list(L.columns)).T
    return V.q_to_annual(fq)


def survey_se_by_group(raw_xlsx: str) -> pd.DataFrame:
    """원자료에서 21개 업종별 E4-1 2040년 평균·표준편차·응답수·표준오차(사전분산의 근거)."""
    d = pd.read_excel(raw_xlsx, header=None, skiprows=4, usecols=[1, 9, 362])
    d.columns = ["업종코드", "KSIC", "E41_2040"]
    d = d[d["KSIC"].astype(str).str.match(r"^[A-Z]\d")]
    d["E41_2040"] = pd.to_numeric(d["E41_2040"], errors="coerce")
    g = d.groupby("업종코드")["E41_2040"].agg(["mean", "std", "count"])
    manuf = d[d["업종코드"].astype(str).str[:2].isin([f"{k:02d}" for k in range(1, 10)])]["E41_2040"]
    g.loc["전체"] = [d["E41_2040"].mean(), d["E41_2040"].std(), d["E41_2040"].count()]
    g.loc["제조업"] = [manuf.mean(), manuf.std(), manuf.count()]
    g["se"] = g["std"] / np.sqrt(g["count"])
    return g