# -*- coding: utf-8 -*-
"""대안 1·2 계산 묶음 — 노트북은 이 함수만 부르고, 계산 논리는 ai_index·ai_varx·ai_forecast·varx_engine에 있다.

배분 형태(shape)와 파급 방식(mode)
  shape : 산업별 AI 효과의 상대 크기. 설문 강도 a_i(E4-1 2040년 − 1), 노출 지수를 설문 척도로 옮긴 a'_i, 대안 2의 사후평균 γ_i.
  mode  : '직접'  — 파급 없이 산업 자기 효과만(현행 4장 방식과 같은 구조)
          '장기'  — 기존 VARX 동학을 쓰되, 장기 효과가 shape가 되도록 동시항을 정한다(γ = (I − ΣAᵀ)·shape). 중심 사양.
          '동시항' — shape를 VARX 동시항 계수로 그대로 넣고 동학이 퍼뜨리게 둔다(진단용).
  모든 경우 2035년 부가가치 가중평균 프리미엄 = Λ가 되도록 배율을 맞추므로, 사양 간 차이는 산업 간 배분과 파급에서만 생긴다.
"""
from __future__ import annotations

import numpy as np
import pandas as pd

import ai_forecast as F
import ai_index as I
import ai_varx as X

MODES = ("직접", "장기", "동시항")


# ── 입력 정리 ───────────────────────────────────────────────────────────
def survey_groups(names: list[str], crosswalk_csv: str) -> list[str]:
    cw = pd.read_csv(crosswalk_csv).set_index("62번")
    return [cw.loc[int(c.split(".")[0]), "설문업종"] for c in names]


def survey_shape(sig: pd.DataFrame, grp: list[str]) -> np.ndarray:
    """설문 강도 a_i = E4-1 2040년 평균 점수 − 1(4장 배분 A1과 같다)."""
    return np.array([sig.loc[g, "AI성장_2040"] - 1.0 for g in grp])


def build_index(patent_csv: str, paper_csv: str, news_csv: str) -> pd.DataFrame:
    pat = pd.read_csv(patent_csv, index_col=0)["5개국_G06N"]
    pap = pd.read_csv(paper_csv, index_col=0)["AI_개념_전체"]
    new = pd.read_csv(news_csv, index_col=0, parse_dates=True).iloc[:, 0]
    return I.composite_index(pat, pap, new)


def build_exposure(news_csv, field_csv, field_map_csv, coipc_csv, coipc_sup_csv, ipc_map_csv,
                   va2025: pd.Series, grp: list[str]) -> tuple[pd.DataFrame, pd.Series]:
    """21개 업종 노출 지표·지수와, '전체'·'제조업'을 더한 업종별 z(62개 산업 대응용)."""
    er = pd.read_csv(news_csv)
    fs = pd.read_csv(field_csv).set_index("분야(OpenAlex field)")
    fmap = pd.read_csv(field_map_csv)
    imap = pd.read_csv(ipc_map_csv)
    co = I.merge_coipc(pd.read_csv(coipc_csv), pd.read_csv(coipc_sup_csv))
    va_g = va2025.groupby(grp).sum() / 1000.0                      # 조 원
    inds = sorted(er["설문업종"])
    ex = I.exposure_index([I.exposure_news(er), I.exposure_papers(fs, fmap), I.exposure_patents(co, imap, va_g)], inds)
    zg = I.exposure_by_group(ex["노출지수_z"], va_g)
    return ex, zg


def build_keis(scen_csv, map11_csv, short_csv, va2025: pd.Series, grp: list[str]):
    """한국고용정보원(2017) 혁신 시나리오의 업종별 ΔVA(%p)와 '전체'·'제조업'(부가가치 가중평균)을 더한 값, 전산업 값."""
    import keis_prior as KP
    tab, tot = KP.r14_dva(scen_csv, map11_csv, short_csv)
    va_g = va2025.groupby(grp).sum()
    g = I.exposure_by_group(tab["ΔVA_평균"], va_g)
    return tab, g, tot


# ── 시나리오 ───────────────────────────────────────────────────────────
def scenario(base: pd.DataFrame, A: list[np.ndarray], shape: np.ndarray, xq: pd.Series, lam: float, mode: str):
    n = len(shape)
    if mode == "직접":
        AA, g = [np.zeros((n, n))], shape
    elif mode == "장기":
        AA, g = A, X.impact_from_long_run(A, shape)
    elif mode == "동시항":
        AA, g = A, shape
    else:
        raise ValueError(mode)
    ai, prem, c = X.ai_path(base, AA, g, xq, lam)
    return ai, prem


def summarize(label: str, base: pd.DataFrame, ai: pd.DataFrame, prem: pd.Series, start=2025, end=2035) -> dict:
    b0, b1, a1 = base.loc[start].sum(), base.loc[end].sum(), ai.loc[end].sum()
    return {"사양": label, "AI_2035_조": a1 / 1e3, "기준_2035_조": b1 / 1e3,
            "AI_CAGR": F.cagr(b0, a1, end - start), "기준_CAGR": F.cagr(b0, b1, end - start),
            "차이_%p": F.cagr(b0, a1, end - start) - F.cagr(b0, b1, end - start),
            "수준차이_조": (a1 - b1) / 1e3, "프리미엄_최소": float(prem.min()), "프리미엄_최대": float(prem.max()),
            "음의_산업수": int((prem < 0).sum())}


def rank_corr(x: np.ndarray, y: np.ndarray) -> float:
    return float(pd.Series(x).rank().corr(pd.Series(y).rank()))


# ── 대안 2 ─────────────────────────────────────────────────────────────
def alt2_estimate(L, exo, Wio, gap, a: np.ndarray, se_i: np.ndarray) -> dict:
    """자료만의 AI 계수 → 경험적 베이즈 β·τ² → 두 사전분포(설문 표본오차만 / 계층)로 재추정."""
    g, prec = X.ai_coef_and_precision(L, exo, Wio, gap)
    beta, beta_se = X.empirical_bayes_beta(g, prec, a)
    tau2 = X.hierarchical_tau2(g, prec, a, beta)
    out = {"g_data": g, "prec": prec, "beta": beta, "beta_se": beta_se, "tau2": tau2, "specs": {}}
    for name, pv in [("설문 표본오차", (abs(beta) * se_i) ** 2), ("계층(경험적 베이즈)", (abs(beta) * se_i) ** 2 + tau2)]:
        pm = beta * a
        B2, groups, r0, sig = X.fit_with_ai(L, exo, Wio, gap, pm, pv)
        out["specs"][name] = {"prior_mean": pm, "prior_var": pv, "B": B2, "r0": r0, "post": B2[r0, :],
                              "data_weight": X.shrinkage_weight(prec, pv)}
    return out


def baseline_change(L, gdp, exo, B_old, B_new, gap, s1: pd.DataFrame, anchor: float, year: int = 2035) -> tuple[dict, pd.Series]:
    """미래 AI 갭 0(기준 시나리오)에서 재추정 모형이 기존 전망을 얼마나 바꾸는가.
    바텀업 합계의 변화와, 5장 비례 벤치마킹(해마다 같은 배율)을 거친 뒤의 산업별 변화(= 구성비 변화)를 돌려준다."""
    f_old = X.forecast_s1(L, gdp, exo, B_old, s1, anchor, False)
    exo_ai = exo.copy()
    exo_ai[X.AI_COL] = gap.reindex(exo.index).fillna(0.0)
    f_new = X.forecast_s1(L, gdp, exo_ai, B_new, s1, anchor, True)
    t_old, t_new = f_old.sum(axis=1), f_new.sum(axis=1)
    sh_old, sh_new = f_old.loc[year] / t_old[year], f_new.loc[year] / t_new[year]
    ch = 100.0 * (sh_new / sh_old - 1.0)
    return {"바텀업_기존_조": t_old[year] / 1e3, "바텀업_재추정_조": t_new[year] / 1e3,
            "바텀업_변화_%": 100.0 * (t_new[year] / t_old[year] - 1.0),
            "벤치마킹후_최대_%": float(ch.max()), "벤치마킹후_최소_%": float(ch.min()),
            "벤치마킹후_절대중앙_%": float(ch.abs().median())}, ch