"""
시각화 및 리포트 모듈

Lifetime PD 분석 결과를 차트와 테이블로 시각화합니다.

차트 목록:
1. Zt 시계열 (과거 + 예측)
2. Zt vs 거시변수 산점도
3. 시나리오별 Marginal PD 곡선
4. 시나리오별 Cumulative PD 곡선
5. 시나리오 가중평균 PD Term Structure
6. 전이행렬 히트맵
7. 통계 검증 결과 요약
"""

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.ticker as mticker
import seaborn as sns
from pathlib import Path
from typing import Dict, List, Optional
import logging

logger = logging.getLogger(__name__)

# 한글 폰트 설정
plt.rcParams["font.family"] = "Malgun Gothic"
plt.rcParams["axes.unicode_minus"] = False
plt.rcParams["figure.dpi"] = 150

# 시나리오 색상
SCENARIO_COLORS = {
    "upside": "#2ecc71",    # 초록
    "base": "#3498db",      # 파랑
    "downside": "#e74c3c",  # 빨강
}

SCENARIO_LABELS = {
    "upside": "호황 (Upside)",
    "base": "중립 (Base)",
    "downside": "불황 (Downside)",
}


def plot_zt_timeseries(
    zt_history: Dict[int, float],
    z_paths: Dict[str, np.ndarray] = None,
    base_year: int = 2025,
    save_path: str = None
):
    """Zt 시계열 차트 (과거 + 시나리오별 예측)"""
    fig, ax = plt.subplots(figsize=(14, 6))
    
    # 과거 Zt
    years = sorted(zt_history.keys())
    values = [zt_history[y] for y in years]
    ax.plot(years, values, "ko-", markersize=4, linewidth=1.5, label="과거 Zt (추정)")
    ax.fill_between(years, values, 0, alpha=0.15, color="gray")
    
    # 시나리오별 예측
    if z_paths:
        for scenario, z_path in z_paths.items():
            future_years = list(range(base_year + 1, base_year + len(z_path) + 1))
            color = SCENARIO_COLORS.get(scenario, "gray")
            label = SCENARIO_LABELS.get(scenario, scenario)
            ax.plot(future_years, z_path, color=color, linewidth=1.5,
                   linestyle="--", alpha=0.8, label=label)
    
    ax.axhline(y=0, color="navy", linestyle=":", alpha=0.5, label="TTC (Z=0)")
    ax.set_xlabel("연도")
    ax.set_ylabel("Credit Cycle Index (Zt)")
    ax.set_title("Belkin & Suchower 신용사이클 인덱스 (Zt)")
    ax.legend(loc="best", fontsize=9)
    ax.grid(True, alpha=0.3)
    
    # 주요 이벤트 표시
    events = {2000: "IMF\n여파", 2003: "카드\n사태", 2008: "GFC", 2020: "COVID"}
    for yr, label in events.items():
        if yr in zt_history:
            ax.annotate(label, xy=(yr, zt_history[yr]),
                       xytext=(yr, zt_history[yr] - 0.5),
                       fontsize=7, ha="center", color="red",
                       arrowprops=dict(arrowstyle="->", color="red", lw=0.8))
    
    plt.tight_layout()
    if save_path:
        fig.savefig(save_path, bbox_inches="tight")
        logger.info(f"차트 저장: {save_path}")
    plt.close(fig)
    return fig


def plot_macro_vs_zt(
    zt_series: pd.Series,
    macro_data: pd.DataFrame,
    save_path: str = None
):
    """거시변수 vs Zt 산점도 (회귀선 포함)"""
    n_vars = min(len(macro_data.columns), 6)
    fig, axes = plt.subplots(2, 3, figsize=(16, 10))
    axes = axes.flatten()
    
    common = sorted(set(zt_series.index) & set(macro_data.index))
    
    for idx, col in enumerate(macro_data.columns[:n_vars]):
        ax = axes[idx]
        x = macro_data.loc[common, col].values
        y = zt_series.loc[common].values
        
        ax.scatter(x, y, color="#3498db", alpha=0.7, s=30)
        
        # 회귀선
        if len(x) > 2:
            z = np.polyfit(x, y, 1)
            p = np.poly1d(z)
            x_line = np.linspace(x.min(), x.max(), 50)
            ax.plot(x_line, p(x_line), "r--", alpha=0.7, linewidth=1.2)
            
            corr = np.corrcoef(x, y)[0, 1]
            ax.set_title(f"{col}\n(ρ = {corr:.3f})", fontsize=10)
        else:
            ax.set_title(col, fontsize=10)
        
        ax.set_xlabel(col, fontsize=9)
        ax.set_ylabel("Zt", fontsize=9)
        ax.grid(True, alpha=0.3)
    
    # 빈 서브플롯 숨기기
    for idx in range(n_vars, len(axes)):
        axes[idx].set_visible(False)
    
    fig.suptitle("거시경제변수 vs 신용사이클 인덱스 (Zt)", fontsize=13, y=1.02)
    plt.tight_layout()
    if save_path:
        fig.savefig(save_path, bbox_inches="tight")
    plt.close(fig)
    return fig


def plot_lifetime_pd(
    results: Dict,
    pd_type: str = "cumulative",
    grades_to_show: List[str] = None,
    base_year: int = 2025,
    save_path: str = None
):
    """
    시나리오별 Lifetime PD 곡선
    
    Parameters
    ----------
    pd_type : str - "cumulative" or "marginal"
    grades_to_show : List[str] - 표시할 등급 (기본: BBB, BB, B)
    """
    from data.transition_matrices import RATING_GRADES
    non_default = RATING_GRADES[:-1]
    
    if grades_to_show is None:
        grades_to_show = ["BBB", "BB", "B"]
    
    grade_indices = [non_default.index(g) for g in grades_to_show if g in non_default]
    n_grades = len(grade_indices)
    
    fig, axes = plt.subplots(1, n_grades, figsize=(6 * n_grades, 5))
    if n_grades == 1:
        axes = [axes]
    
    pd_key = f"{pd_type}_pd"
    
    for ax_idx, (ax, gi) in enumerate(zip(axes, grade_indices)):
        grade = non_default[gi]
        
        for scenario, data in results["by_scenario"].items():
            pds = data[pd_key][:, gi]
            horizon = len(pds)
            years = list(range(base_year + 1, base_year + horizon + 1))
            
            color = SCENARIO_COLORS.get(scenario, "gray")
            label = SCENARIO_LABELS.get(scenario, scenario)
            ax.plot(years, pds * 100, color=color, linewidth=1.5, label=label)
        
        # 가중평균
        weighted_key = f"weighted_{pd_type}_pd"
        if weighted_key in results:
            w_pds = results[weighted_key][:, gi]
            years = list(range(base_year + 1, base_year + len(w_pds) + 1))
            ax.plot(years, w_pds * 100, color="purple", linewidth=2.5,
                   linestyle="-.", label="가중평균", alpha=0.8)
        
        ax.set_xlabel("연도")
        ax.set_ylabel(f"{'누적' if pd_type == 'cumulative' else '한계'} PD (%)")
        ax.set_title(f"{grade} 등급")
        ax.legend(fontsize=8)
        ax.grid(True, alpha=0.3)
        ax.yaxis.set_major_formatter(mticker.FuncFormatter(lambda x, p: f"{x:.1f}%"))
    
    title = f"50년 {'누적' if pd_type == 'cumulative' else '한계'} PD — 시나리오별"
    fig.suptitle(title, fontsize=13, y=1.02)
    plt.tight_layout()
    if save_path:
        fig.savefig(save_path, bbox_inches="tight")
    plt.close(fig)
    return fig


def plot_pd_term_structure(
    results: Dict,
    base_year: int = 2025,
    save_path: str = None
):
    """전 등급 가중평균 누적 PD Term Structure"""
    from data.transition_matrices import RATING_GRADES
    non_default = RATING_GRADES[:-1]
    
    fig, ax = plt.subplots(figsize=(12, 7))
    
    weighted = results["weighted_cumulative_pd"]
    horizon = weighted.shape[0]
    years = list(range(1, horizon + 1))
    
    colors = plt.cm.RdYlGn_r(np.linspace(0.1, 0.9, len(non_default)))
    
    for gi, (grade, color) in enumerate(zip(non_default, colors)):
        ax.plot(years, weighted[:, gi] * 100, color=color, linewidth=1.5, label=grade)
    
    ax.set_xlabel("기간 (년)")
    ax.set_ylabel("누적 PD (%)")
    ax.set_title("가중평균 누적 PD Term Structure (전 등급)", fontsize=13)
    ax.legend(bbox_to_anchor=(1.05, 1), loc="upper left", fontsize=9)
    ax.grid(True, alpha=0.3)
    ax.set_xlim(1, horizon)
    
    plt.tight_layout()
    if save_path:
        fig.savefig(save_path, bbox_inches="tight")
    plt.close(fig)
    return fig


def plot_transition_heatmap(
    tm: np.ndarray,
    title: str = "전이행렬",
    grades: List[str] = None,
    save_path: str = None
):
    """전이행렬 히트맵"""
    from data.transition_matrices import RATING_GRADES
    if grades is None:
        grades = RATING_GRADES
    
    fig, ax = plt.subplots(figsize=(9, 7))
    
    sns.heatmap(
        tm * 100,
        annot=True, fmt=".2f",
        xticklabels=grades, yticklabels=grades,
        cmap="YlOrRd", ax=ax,
        cbar_kws={"label": "전이확률 (%)"},
        linewidths=0.5
    )
    
    ax.set_xlabel("전이 후 등급")
    ax.set_ylabel("전이 전 등급")
    ax.set_title(title, fontsize=13)
    
    plt.tight_layout()
    if save_path:
        fig.savefig(save_path, bbox_inches="tight")
    plt.close(fig)
    return fig


def plot_validation_summary(
    validation_df: pd.DataFrame,
    save_path: str = None
):
    """검증 결과 요약 테이블 이미지"""
    fig, ax = plt.subplots(figsize=(14, max(3, len(validation_df) * 0.6 + 1)))
    ax.axis("off")
    ax.set_title("통계적 검증 결과 요약", fontsize=14, pad=20)
    
    # 테이블 색상
    colors = []
    for _, row in validation_df.iterrows():
        if "Pass" in str(row.get("결과", "")):
            colors.append(["#d4edda"] * len(validation_df.columns))
        else:
            colors.append(["#f8d7da"] * len(validation_df.columns))
    
    table = ax.table(
        cellText=validation_df.values,
        colLabels=validation_df.columns,
        cellColours=colors if colors else None,
        loc="center",
        cellLoc="center"
    )
    table.auto_set_font_size(False)
    table.set_fontsize(8)
    table.scale(1.0, 1.5)
    
    plt.tight_layout()
    if save_path:
        fig.savefig(save_path, bbox_inches="tight")
    plt.close(fig)
    return fig


def generate_all_plots(
    zt_history: Dict[int, float],
    z_paths: Dict[str, np.ndarray],
    zt_series_pd: pd.Series,
    macro_data: pd.DataFrame,
    pd_results: Dict,
    ttc_matrix: np.ndarray,
    validation_df: pd.DataFrame,
    output_dir: str = "results",
    base_year: int = 2025
):
    """모든 차트 일괄 생성"""
    out = Path(output_dir)
    out.mkdir(parents=True, exist_ok=True)
    
    logger.info(f"차트 생성 중... (저장 경로: {out})")
    
    # 1. Zt 시계열
    plot_zt_timeseries(zt_history, z_paths, base_year,
                      str(out / "01_zt_timeseries.png"))
    
    # 2. 거시변수 vs Zt
    plot_macro_vs_zt(zt_series_pd, macro_data,
                    str(out / "02_macro_vs_zt.png"))
    
    # 3. 누적 PD (주요 등급)
    plot_lifetime_pd(pd_results, "cumulative", ["BBB", "BB", "B"],
                    base_year, str(out / "03_cumulative_pd.png"))
    
    # 4. 한계 PD (주요 등급)
    plot_lifetime_pd(pd_results, "marginal", ["BBB", "BB", "B"],
                    base_year, str(out / "04_marginal_pd.png"))
    
    # 5. 전 등급 Term Structure
    plot_pd_term_structure(pd_results, base_year,
                          str(out / "05_pd_term_structure.png"))
    
    # 6. TTC 전이행렬 히트맵
    plot_transition_heatmap(ttc_matrix, "TTC 전이행렬 (장기 평균)",
                          save_path=str(out / "06_ttc_heatmap.png"))
    
    # 7. 검증 결과
    plot_validation_summary(validation_df,
                           str(out / "07_validation_summary.png"))
    
    logger.info(f"총 7개 차트 생성 완료 → {out}/")