# sampling_sim_app.py
import os
import io

import matplotlib
matplotlib.use("Agg")

import numpy as np
import pandas as pd
import geopandas as gpd
import matplotlib.pyplot as plt
from shapely.geometry import Point
from shapely import wkt as shp_wkt
from PIL import Image
import contextily as ctx  # ← 추가

import gradio as gr

# ----------------- 경로 설정 -----------------
BASE_DIR = os.path.dirname(__file__)
DATA_DIR = os.path.join(BASE_DIR, "Data")

SHP_PATH = os.path.join(DATA_DIR, "inha_boundary.shp")
ZONES_PATH = os.path.join(DATA_DIR, "zones_최종.geojson")

# 측정결과 CSV (lat, lon, dB 또는 위치, dB)
CSV_PATH = os.path.join(DATA_DIR, "측정결과_최종.csv")


# ---------- 유틸: Figure -> numpy ----------
def fig_to_array(fig, dpi=100):
    buf = io.BytesIO()
    fig.savefig(buf, format="png", dpi=dpi, bbox_inches="tight")
    buf.seek(0)
    img = np.array(Image.open(buf))
    plt.close(fig)
    return img


# ---------- 경계 SHP 로딩 ----------
def load_boundary():
    gdf4326 = gpd.read_file(SHP_PATH).to_crs(epsg=4326)
    gdf3857 = gdf4326.to_crs(epsg=3857)

    poly3857 = gdf3857.geometry.unary_union
    if poly3857.geom_type == "MultiPolygon":
        poly3857 = max(poly3857.geoms, key=lambda g: g.area)

    return gdf4326, gdf3857, poly3857


# ---------- 측정결과 CSV 로딩 ----------
def load_points_from_csv():
    try:
        df_u = pd.read_csv(CSV_PATH, encoding="cp949")
    except UnicodeDecodeError:
        df_u = pd.read_csv(CSV_PATH, encoding="utf-8")

    df_u.columns = [c.strip() for c in df_u.columns]
    df_u = df_u.loc[:, ~df_u.columns.str.contains("^Unnamed")]

    # 케이스 B: 위치, dB
    if "위치" in df_u.columns and "dB" in df_u.columns:
        pos_series = df_u["위치"].astype(str).str.strip()

        def safe_load_wkt(s):
            s_up = s.upper()
            if "POINT" not in s_up:
                return None
            try:
                return shp_wkt.loads(s)
            except Exception:
                return None

        geom = pos_series.apply(safe_load_wkt)
        mask = geom.notnull()
        if not mask.any():
            raise ValueError("'위치' 컬럼에 유효한 POINT WKT가 없다.")

        geom_valid = geom[mask]
        lons = geom_valid.apply(lambda g: round(g.x, 7))
        lats = geom_valid.apply(lambda g: round(g.y, 7))
        dBs = pd.to_numeric(df_u.loc[mask, "dB"], errors="coerce")

        df = pd.DataFrame({"lat": lats, "lon": lons, "dB": dBs}).dropna(
            subset=["lat", "lon"]
        )
        return df

    # 케이스 A: lat, lon, dB
    if {"lat", "lon", "dB"}.issubset(df_u.columns):
        df = df_u[["lat", "lon", "dB"]].copy()
        df["lat"] = pd.to_numeric(df["lat"], errors="coerce").round(7)
        df["lon"] = pd.to_numeric(df["lon"], errors="coerce").round(7)
        df["dB"] = pd.to_numeric(df["dB"], errors="coerce")
        df = df.dropna(subset=["lat", "lon"])
        return df

    raise ValueError("CSV는 (lat, lon, dB) 또는 (위치, dB) 형식이어야 한다.")


# ---------- 구역 GeoJSON 로딩 ----------
def load_zones(poly3857):
    zones4326 = gpd.read_file(ZONES_PATH).to_crs(epsg=4326)
    zones3857 = zones4326.to_crs(epsg=3857)

    if "zone_id" not in zones3857.columns:
        zones3857["zone_id"] = np.arange(1, len(zones3857) + 1)

    zones3857["geom_clip"] = zones3857.geometry.intersection(poly3857)
    zones3857["area"] = zones3857["geom_clip"].area
    zones3857 = zones3857[zones3857["area"] > 0].copy()

    A_total = zones3857["area"].sum()
    zones3857["W_area"] = zones3857["area"] / A_total

    return zones4326, zones3857


# ---------- 포인트에 zone_id 붙이기 ----------
def attach_zone_id(points_df, zones3857):
    pts4326 = gpd.GeoDataFrame(
        points_df.copy(),
        geometry=gpd.points_from_xy(points_df["lon"], points_df["lat"]),
        crs=4326,
    )
    pts3857 = pts4326.to_crs(epsg=3857)

    zones_clip = zones3857[["zone_id", "geom_clip"]].copy()
    zones_clip = zones_clip.set_geometry("geom_clip")
    zones_clip = zones_clip.rename_geometry("geometry")

    joined = gpd.sjoin(
        pts3857,
        zones_clip[["zone_id", "geometry"]],
        how="inner",
        predicate="within",
    )

    out = joined[["lat", "lon", "dB", "zone_id"]].copy().reset_index(drop=True)
    return out


# ---------- 공통 준비: 경계/구역/포인트/층정보 ----------
def prepare_sampling_data():
    boundary4326, boundary3857, poly3857 = load_boundary()
    zones4326, zones3857 = load_zones(poly3857)
    df_points = load_points_from_csv()
    points_zone = attach_zone_id(df_points, zones3857)

    # 층별 모집단 크기 N_h
    zone_sizes = points_zone.groupby("zone_id").size().rename("N_h")
    zones_info = zones3857[["zone_id", "area", "W_area"]].merge(
        zone_sizes, on="zone_id", how="left"
    )
    zones_info["N_h"] = zones_info["N_h"].fillna(0).astype(int)

    # 층별 표준편차 S_h (네이만 배분용)
    strata_std = points_zone.groupby("zone_id")["dB"].std(ddof=1).rename("S_h")
    zones_info = zones_info.merge(strata_std, on="zone_id", how="left")
    zones_info["S_h"] = zones_info["S_h"].fillna(0.0)

    return boundary4326, zones4326, df_points, points_zone, zones_info


# ---------- 면적비례 n_h 계산 ----------
def allocate_n_by_area(zones_info, n_total):
    zones = zones_info.copy()
    zones["n_raw"] = n_total * zones["W_area"]
    zones["n"] = zones["n_raw"].round().astype(int)

    diff = n_total - zones["n"].sum()
    if diff != 0:
        order = zones["n_raw"].sort_values(ascending=(diff < 0)).index
        for idx in order[:abs(diff)]:
            zones.loc[idx, "n"] += 1 if diff > 0 else -1

    zones["n"] = zones[["n", "N_h"]].min(axis=1)
    return zones


# ---------- 네이만 배분 n_h 계산 ----------
def allocate_n_neyman(zones_info, n_total):
    zones = zones_info.copy()
    if "S_h" not in zones.columns:
        raise ValueError("zones_info에 S_h가 없다. 네이만 배분 전에 S_h를 계산해야 한다.")

    zones["S_h"] = zones["S_h"].fillna(0.0)

    # 가중치 ∝ N_h * S_h
    weight = zones["N_h"] * zones["S_h"]
    if weight.sum() <= 0:
        # 분산이 전부 0이면 N_h 비례로 배분
        weight = zones["N_h"].copy()

    zones["n_raw"] = n_total * weight / weight.sum()
    zones["n"] = zones["n_raw"].round().astype(int)

    diff = n_total - zones["n"].sum()
    if diff != 0:
        order = zones["n_raw"].sort_values(ascending=(diff < 0)).index
        for idx in order[:abs(diff)]:
            zones.loc[idx, "n"] += 1 if diff > 0 else -1

    zones["n"] = zones[["n", "N_h"]].min(axis=1)
    return zones


# ---------- 시뮬레이션 ----------
def simulate_sampling(points_zone_df, zones_info, n, R, strat_label):
    y_all = points_zone_df["dB"].to_numpy()
    N = len(y_all)
    mu_pop = float(np.mean(y_all))

    zones = zones_info.set_index("zone_id")
    W_h = zones["N_h"] / zones["N_h"].sum()
    n_h = zones["n"]

    idx_by_zone = {}
    for z in zones.index:
        idx_by_zone[z] = points_zone_df.index[
            points_zone_df["zone_id"] == z
        ].to_numpy()

    srs_means = []
    strat_means = []

    for _ in range(R):
        # SRS
        n_eff = min(n, N)
        idx_srs = np.random.choice(N, size=n_eff, replace=False)
        mu_srs = float(np.mean(y_all[idx_srs]))
        srs_means.append(mu_srs)

        # Stratified
        mu_str = 0.0
        for z in zones.index:
            N_h = int(zones.loc[z, "N_h"])
            nh = int(n_h.loc[z])
            if N_h == 0 or nh == 0:
                continue
            idx_pool = idx_by_zone[z]
            if len(idx_pool) < nh:
                nh = len(idx_pool)
            choose = np.random.choice(idx_pool, size=nh, replace=False)
            y_h = points_zone_df.loc[choose, "dB"].to_numpy()
            ybar_h = float(np.mean(y_h))
            mu_str += float(W_h.loc[z]) * ybar_h
        strat_means.append(mu_str)

    srs_means = np.array(srs_means)
    strat_means = np.array(strat_means)

    def metrics(arr, name):
        mean_hat = float(arr.mean())
        var_hat = float(arr.var(ddof=1))
        bias = mean_hat - mu_pop
        mse = float(((arr - mu_pop) ** 2).mean())
        return {
            "design": name,
            "mean_hat": mean_hat,
            "var_hat": var_hat,
            "bias": bias,
            "MSE": mse,
        }

    m_srs = metrics(srs_means, "SRS")
    m_str = metrics(strat_means, strat_label)

    df_metrics = pd.DataFrame([m_srs, m_str])
    df_metrics["deff_vs_SRS"] = df_metrics["var_hat"] / df_metrics["var_hat"].iloc[0]

    # 숫자 4째 자리까지 반올림
    num_cols = ["mean_hat", "var_hat", "bias", "MSE", "deff_vs_SRS"]
    df_metrics[num_cols] = df_metrics[num_cols].round(4)

    return mu_pop, srs_means, strat_means, df_metrics


# ---------- 구역 요약 테이블 ----------
def build_zones_table(zones_alloc):
    zones_table = zones_alloc[["zone_id", "area", "W_area", "N_h", "n"]].copy()
    zones_table = zones_table.rename(
        columns={
            "area": "area_m2",
            "W_area": "area_weight",
            "N_h": "N_pop",
            "n": "n_alloc",
        }
    )
    zones_table["sampling_frac"] = zones_table["n_alloc"] / zones_table["N_pop"]

    total_area_m2 = zones_table["area_m2"].sum()
    total_area_weight = zones_table["area_weight"].sum()
    total_N = zones_table["N_pop"].sum()
    total_n = zones_table["n_alloc"].sum()
    total_sampling_frac = total_n / total_N if total_N > 0 else np.nan

    total_row = {
        "zone_id": "합계",
        "area_m2": total_area_m2,
        "area_weight": total_area_weight,
        "N_pop": total_N,
        "n_alloc": total_n,
        "sampling_frac": total_sampling_frac,
    }

    zones_table = zones_table[
        ["zone_id", "area_m2", "area_weight", "N_pop", "n_alloc", "sampling_frac"]
    ]
    zones_table = pd.concat([zones_table, pd.DataFrame([total_row])], ignore_index=True)

    # 숫자 컬럼 4째 자리까지 반올림
    num_cols_zone = ["area_m2", "area_weight", "N_pop", "n_alloc", "sampling_frac"]
    for col in num_cols_zone:
        zones_table[col] = pd.to_numeric(zones_table[col], errors="coerce").round(4)

    return zones_table, float(total_N), float(total_n)


# ---------- 시각화 ----------
def plot_population_map(boundary4326, zones4326, points_df):
    g_boundary = boundary4326.to_crs(epsg=3857)
    g_zones = zones4326.to_crs(epsg=3857)
    g_pts = gpd.GeoDataFrame(
        points_df.copy(),
        geometry=gpd.points_from_xy(points_df["lon"], points_df["lat"]),
        crs=4326,
    ).to_crs(epsg=3857)

    fig, ax = plt.subplots(figsize=(8, 8))

    # 베이스맵 먼저 깔기
    minx, miny, maxx, maxy = g_boundary.total_bounds
    ax.set_xlim(minx, maxx)
    ax.set_ylim(miny, maxy)
    ctx.add_basemap(ax, source=ctx.providers.OpenStreetMap.Mapnik, alpha=0.7)

    # 그 위에 경계/구역/포인트
    g_boundary.plot(ax=ax, facecolor="none", edgecolor="black", linewidth=2)
    if len(g_zones) > 0:
        g_zones.plot(
            ax=ax,
            facecolor="none",
            edgecolor="red",
            linewidth=1.5,
            alpha=0.7,
        )

    g_pts.plot(
        ax=ax,
        column="dB",
        cmap="viridis",
        markersize=20,
        edgecolor="black",
        linewidth=0.2,
        legend=False,
        alpha=0.9,
    )

    ax.set_title("Population points with zones", fontsize=12)
    ax.axis("off")
    plt.tight_layout()
    return fig_to_array(fig)


def plot_histograms(srs_means, strat_means, mu_pop, strat_label):
    fig, ax = plt.subplots(figsize=(8, 5))

    ax.hist(srs_means, bins=20, alpha=0.5, label="SRS")
    ax.hist(strat_means, bins=20, alpha=0.5, label=strat_label)
    ax.axvline(mu_pop, color="k", linestyle="--", label="Population mean")

    ax.set_xlabel("Sample mean (dB)")
    ax.set_ylabel("Frequency")
    ax.legend()
    ax.set_title("Sampling distribution of mean")
    plt.tight_layout()
    return fig_to_array(fig)


# ---------- Gradio 콜백: 면적비례 ----------
def run_sim_area(n, R):
    n = int(n)
    R = int(R)

    boundary4326, zones4326, df_points, points_zone, zones_info = prepare_sampling_data()

    # 면적비례 배분
    zones_alloc = allocate_n_by_area(zones_info, n)
    zones_table, total_N, total_n = build_zones_table(zones_alloc)

    mu_pop, srs_means, strat_means, df_metrics = simulate_sampling(
        points_zone, zones_alloc, n, R, "Stratified (area-proportional)"
    )

    map_img = plot_population_map(boundary4326, zones4326, df_points)
    hist_img = plot_histograms(
        srs_means, strat_means, mu_pop, "Stratified (area-proportional)"
    )

    summary = (
        f"[면적비례 층화]\n"
        f"모집단 포인트 수: {int(total_N)}개\n"
        f"구역 수: {len(zones_alloc)}개\n"
        f"모집단 평균 dB: {mu_pop:.4f}\n"
        f"총 표본크기 n = {int(total_n)} (입력값 {n}), 반복 R = {R}"
    )

    return map_img, hist_img, df_metrics, zones_table, summary


# ---------- Gradio 콜백: 네이만 ----------
def run_sim_neyman(n, R):
    n = int(n)
    R = int(R)

    boundary4326, zones4326, df_points, points_zone, zones_info = prepare_sampling_data()

    # 네이만 배분
    zones_alloc = allocate_n_neyman(zones_info, n)
    zones_table, total_N, total_n = build_zones_table(zones_alloc)

    mu_pop, srs_means, strat_means, df_metrics = simulate_sampling(
        points_zone, zones_alloc, n, R, "Stratified (Neyman)"
    )

    map_img = plot_population_map(boundary4326, zones4326, df_points)
    hist_img = plot_histograms(
        srs_means, strat_means, mu_pop, "Stratified (Neyman)"
    )

    summary = (
        f"[네이만 층화]\n"
        f"모집단 포인트 수: {int(total_N)}개\n"
        f"구역 수: {len(zones_alloc)}개\n"
        f"모집단 평균 dB: {mu_pop:.4f}\n"
        f"총 표본크기 n = {int(total_n)} (입력값 {n}), 반복 R = {R}"
    )

    return map_img, hist_img, df_metrics, zones_table, summary


# ---------- Gradio 앱 ----------
def build_app():
    with gr.Blocks(title="표본추출 시뮬레이션 (SRS vs 층화)") as demo:

        with gr.Tab("SRS vs 면적비례 층화"):
            with gr.Row():
                n_area = gr.Number(label="표본크기 n", value=60, precision=0)
                R_area = gr.Number(label="시뮬레이션 반복 횟수 R", value=500, precision=0)

            run_btn_area = gr.Button("면적비례 층화 시뮬레이션 실행", variant="primary")

            with gr.Row():
                map_img_area = gr.Image(label="Population + zones")
                hist_img_area = gr.Image(label="Sampling distribution")

            metrics_table_area = gr.Dataframe(
                headers=[
                    "design",
                    "mean_hat",
                    "var_hat",
                    "bias",
                    "MSE",
                    "deff_vs_SRS",
                ],
                label="성능 비교 지표",
                interactive=False,
            )

            zones_table_area = gr.Dataframe(
                label="구역별 면적 및 모집단/배정 표본수",
                interactive=False,
            )

            summary_md_area = gr.Markdown()

            run_btn_area.click(
                fn=run_sim_area,
                inputs=[n_area, R_area],
                outputs=[
                    map_img_area,
                    hist_img_area,
                    metrics_table_area,
                    zones_table_area,
                    summary_md_area,
                ],
            )

        with gr.Tab("SRS vs 네이만 층화"):
            with gr.Row():
                n_neyman = gr.Number(label="표본크기 n", value=60, precision=0)
                R_neyman = gr.Number(label="시뮬레이션 반복 횟수 R", value=500, precision=0)

            run_btn_neyman = gr.Button("네이만 층화 시뮬레이션 실행", variant="primary")

            with gr.Row():
                map_img_neyman = gr.Image(label="Population + zones")
                hist_img_neyman = gr.Image(label="Sampling distribution")

            metrics_table_neyman = gr.Dataframe(
                headers=[
                    "design",
                    "mean_hat",
                    "var_hat",
                    "bias",
                    "MSE",
                    "deff_vs_SRS",
                ],
                label="성능 비교 지표",
                interactive=False,
            )

            zones_table_neyman = gr.Dataframe(
                label="구역별 면적 및 모집단/배정 표본수",
                interactive=False,
            )

            summary_md_neyman = gr.Markdown()

            run_btn_neyman.click(
                fn=run_sim_neyman,
                inputs=[n_neyman, R_neyman],
                outputs=[
                    map_img_neyman,
                    hist_img_neyman,
                    metrics_table_neyman,
                    zones_table_neyman,
                    summary_md_neyman,
                ],
            )

    return demo


if __name__ == "__main__":
    app = build_app()
    app.launch(share=True)