#!/usr/bin/env python

import argparse
import re

import numpy as np
import pandas as pd


ZEN_TO_ASCII = str.maketrans(
    "０１２３４５６７８９",
    "0123456789",
)


def z2a(s):
    """全角数字を半角数字に変換する。"""
    return s.translate(ZEN_TO_ASCII)


def extract_values_from_body(body):
    """
    研究室ごとの回答本文からGP値を取り出す。

    - 「値 カウント」表がある場合は表を優先
    - 表がない場合は1行1回答として読む
    - 「編入生」はGPとしては読まず、内部値0として扱う
    - 「365（修正済）」のような値にも対応
    """

    body = z2a(body)

    # --------------------------------------------------
    # Case 1: 「値 カウント」表がある場合
    # --------------------------------------------------
    m = re.search(
        r"^値\s+カウント\s*$",
        body,
        flags=re.MULTILINE,
    )

    if m:
        table_part = body[m.end():]

        count_rows = re.findall(
            r"^\s*(-?\d+(?:\.\d+)?)\s+(\d+)\s*$",
            table_part,
            flags=re.MULTILINE,
        )

        values = []

        for value_str, count_str in count_rows:
            value = float(value_str)
            count = int(count_str)
            values.extend([value] * count)

        return values

    # --------------------------------------------------
    # Case 2: 回答値が1行ずつ並んでいる場合
    # --------------------------------------------------
    values = []

    for line in body.splitlines():
        line = line.strip()

        if not line:
            continue

        if "編入生" in line:
            values.append(0.0)
            continue

        if "この質問にはまだ回答がありません" in line:
            continue

        # 例: 365（修正済）
        m = re.match(
            r"^(-?\d+(?:\.\d+)?)",
            line,
        )

        if m:
            values.append(float(m.group(1)))

    return values


def parse_google_form_results(text):
    """
    Google Formsのアンケート結果をコピペしたテキストから、

        研究室番号 | 値

    のDataFrameを生成する。

    回答者が0人の研究室についても、

        研究室番号 | NaN

    という行を作成してDataFrameに残す。
    """

    text = z2a(text)

    # 「第一希望研究室」などの集計部分を除外し、
    # 最初の個別研究室ブロックから開始する。
    #
    # 以下の両方に対応する。
    #
    # 研究室1
    # GP7件の回答
    #
    # 研究室1
    # 7件の回答
    m = re.search(
        r"^研究室\s*[0-9]+[^\n]*\n"
        r"(?:[^\d\n]*?)"
        r"[0-9]+\s*件の回答",
        text,
        flags=re.MULTILINE,
    )

    if not m:
        raise ValueError("個別研究室ブロックが見つかりません")

    text = text[m.start():]

    # 各研究室ブロックを取得する。
    pattern = re.compile(
        r"^研究室\s*([0-9]+)[^\n]*\n"
        r"(?:[^\d\n]*?)"
        r"([0-9]+)\s*件の回答\s*\n"
        r"(.*?)"
        r"(?="
        r"^研究室\s*[0-9]+[^\n]*\n"
        r"(?:[^\d\n]*?)"
        r"[0-9]+\s*件の回答"
        r"|\Z"
        r")",
        flags=re.MULTILINE | re.DOTALL,
    )

    rows = []

    for match in pattern.finditer(text):
        lab_no = int(match.group(1))
        n_answers = int(match.group(2))
        body = match.group(3).strip()

        # 回答者が0人の研究室もDataFrameに残す。
        #
        # 編入生を表す内部値0と区別するため、
        # 回答者なしはNaNで表現する。
        if n_answers == 0:
            rows.append(
                {
                    "研究室番号": lab_no,
                    "値": np.nan,
                }
            )
            continue

        values = extract_values_from_body(body)

        if len(values) != n_answers:
            print(
                f"警告: 研究室{lab_no}: "
                f"回答数={n_answers}, 抽出数={len(values)}"
            )

        for value in values:
            rows.append(
                {
                    "研究室番号": lab_no,
                    "値": value,
                }
            )

    df = pd.DataFrame(
        rows,
        columns=["研究室番号", "値"],
    )

    if df.empty:
        return df

    # すべての有効値が整数の場合は、NaNを保持できる
    # pandasのnullable整数型Int64に変換する。
    non_na_values = df["値"].dropna()

    if (
        not non_na_values.empty
        and (non_na_values % 1 == 0).all()
    ):
        df["値"] = df["値"].astype("Int64")

    return df


def summarize_by_lab(df):
    """
    研究室ごとの在校生数・編入生数・GP平均・標準偏差を計算する。

    - NaNは回答者0人を表すため、人数には含めない
    - 値が10より大きい回答を在校生として扱う
    - 値が10以下の回答を編入生として扱う
    """

    required_columns = {
        "研究室番号",
        "値",
    }

    missing_columns = required_columns - set(df.columns)

    if missing_columns:
        raise KeyError(
            "必須列がありません: "
            + ", ".join(sorted(missing_columns))
        )

    rows = []

    for lab_no in sorted(df["研究室番号"].unique()):
        df_lab = df.loc[
            df["研究室番号"] == lab_no
        ].copy()

        # 回答者0人を表すNaN行を除外する。
        df_answers = df_lab.loc[
            df_lab["値"].notna()
        ].copy()

        # 在校生と編入生を分類する。
        df_regular = df_answers.loc[
            df_answers["値"] > 10
        ]

        df_transfer = df_answers.loc[
            df_answers["値"] <= 10
        ]

        rows.append(
            {
                "研究室": lab_no,
                "回答者数": df_answers.shape[0],
                "在校生": df_regular.shape[0],
                "編入生": df_transfer.shape[0],
                "GP mean": df_regular["値"].mean(),
                "GP std": df_regular["値"].std(),
            }
        )

    return pd.DataFrame(rows).set_index("研究室")


def print_summary(df, df_r):
    """研究室ごとの集計結果を端末に表示する。"""

    # NaNは回答者0人の研究室を表すため数えない。
    n_respondents = int(df["値"].notna().sum())

    print("-" * 66)
    print("回答者数:", n_respondents)

    print(
        f"{'研究室':>6} "
        f"{'回答者':>8} "
        f"{'在校生':>8} "
        f"{'編入生':>8} "
        f"{'GP mean':>10} "
        f"{'GP std':>10}"
    )

    print("-" * 66)

    for lab_no, row in df_r.iterrows():
        mean = (
            "-"
            if pd.isna(row["GP mean"])
            else f"{row['GP mean']:.1f}"
        )

        std = (
            "-"
            if pd.isna(row["GP std"])
            else f"{row['GP std']:.1f}"
        )

        print(
            f"{lab_no:>6} "
            f"{int(row['回答者数']):>8} "
            f"{int(row['在校生']):>8} "
            f"{int(row['編入生']):>8} "
            f"{mean:>10} "
            f"{std:>10}"
        )

    print("-" * 66)


def apply_plot_style():
    """
    論文図として使いやすいmatplotlibの描画設定を適用する。

    仕様:
    - seabornやscienceplotsには依存しない
    - タイトルは付けず、軸ラベルと目盛だけで情報を伝える
    - フォントサイズ、線幅、余白、保存時DPIを論文図向けに調整する
    - カラーマップは既定でtab10を使う
    """

    import matplotlib.pyplot as plt

    plt.rcParams.update(
        {
            "figure.dpi": 120,
            "savefig.dpi": 300,
            "savefig.bbox": "tight",
            "font.family": "sans-serif",
            "font.sans-serif": [
                "Arial",
                "Helvetica",
                "DejaVu Sans",
            ],
            "font.size": 11,
            "axes.labelsize": 12,
            "axes.linewidth": 1.0,
            "axes.spines.top": False,
            "axes.spines.right": False,
            "xtick.labelsize": 10,
            "ytick.labelsize": 10,
            "xtick.direction": "out",
            "ytick.direction": "out",
            "xtick.major.size": 4,
            "ytick.major.size": 4,
            "xtick.major.width": 1.0,
            "ytick.major.width": 1.0,
            "legend.frameon": False,
            "pdf.fonttype": 42,
            "ps.fonttype": 42,
        }
    )


def filter_plot_data(data, lower=10, upper=None):
    """
    GPの描画対象データを抽出する。

    Parameters
    ----------
    data : pandas.DataFrame
        「研究室番号」と「値」列を持つデータフレーム。
    lower : float, default 10
        この値より大きい回答を描画対象にする。
        既定値では編入生と回答者なしのNaNを除く。
    upper : float or None, default None
        指定した場合、この値未満の回答だけを描画対象にする。

    Returns
    -------
    pandas.DataFrame
        フィルタ後のデータ。元データは変更しない。
    """

    required_columns = {
        "研究室番号",
        "値",
    }

    missing_columns = required_columns - set(data.columns)

    if missing_columns:
        raise KeyError(
            "必須列がありません: "
            + ", ".join(sorted(missing_columns))
        )

    # NaNは比較結果がFalseになるため、ここで自動的に除外される。
    df = data.loc[
        data["値"] > lower,
        ["研究室番号", "値"],
    ].copy()

    if upper is not None:
        df = df.loc[
            df["値"] < upper
        ].copy()

    if df.empty:
        raise ValueError(
            "描画対象のデータがありません。"
            "lower/upperを確認してください。"
        )

    # matplotlib/numpyで扱いやすい通常のfloat型に変換する。
    df["値"] = df["値"].astype(float)

    return df


def make_histogram_bins(values, bin_width):
    """
    GPヒストグラム用のビン境界を作成する。

    仕様:
    - 最小値側はbin_widthの倍数に切り下げる
    - 最大値側はbin_widthの倍数に切り上げる
    - 最大値が最後のビンに入るようにする
    - bin_widthは正の数でなければならない
    """

    if bin_width <= 0:
        raise ValueError(
            "bin_widthは正の数を指定してください。"
        )

    min_edge = (
        np.floor(values.min() / bin_width)
        * bin_width
    )

    max_edge = (
        np.ceil(values.max() / bin_width)
        * bin_width
        + bin_width
    )

    return np.arange(
        min_edge,
        max_edge,
        bin_width,
    )


def style_axis(ax, grid_axis="y"):
    """
    軸まわりの体裁を統一する。

    仕様:
    - 上枠と右枠を消す
    - 目盛を外向きにする
    - 指定軸方向に薄いグリッドを入れる
    """

    ax.spines["top"].set_visible(False)
    ax.spines["right"].set_visible(False)

    ax.tick_params(
        direction="out",
        length=4,
        width=1,
    )

    ax.grid(
        axis=grid_axis,
        color="0.88",
        linewidth=0.8,
    )

    ax.set_axisbelow(True)


def draw_lab_boxplot(ax, df, labs, color, rng):
    """
    研究室別GP分布を箱ひげ図と実測点で描画する。

    仕様:
    - 箱ひげ図は外れ値を非表示にする
    - 実測点を重ねて全回答の分布を見せる
    - 実測点には固定乱数のjitterを与える
    - x軸のカテゴリ順はlabsの順序に従う
    """

    positions = np.arange(
        1,
        len(labs) + 1,
    )

    grouped_values = [
        df.loc[
            df["研究室番号"] == lab,
            "値",
        ].to_numpy(dtype=float)
        for lab in labs
    ]

    box = ax.boxplot(
        grouped_values,
        positions=positions,
        widths=0.55,
        patch_artist=True,
        showfliers=False,
        medianprops={
            "color": "black",
            "linewidth": 1.2,
        },
        boxprops={
            "facecolor": color,
            "edgecolor": "black",
            "linewidth": 1.0,
        },
        whiskerprops={
            "color": "black",
            "linewidth": 1.0,
        },
        capprops={
            "color": "black",
            "linewidth": 1.0,
        },
    )

    for patch in box["boxes"]:
        patch.set_alpha(0.55)

    for x_position, values in zip(
        positions,
        grouped_values,
    ):
        jitter = rng.uniform(
            -0.16,
            0.16,
            size=len(values),
        )

        ax.scatter(
            np.full(len(values), x_position) + jitter,
            values,
            s=18,
            color="black",
            alpha=0.45,
            linewidths=0,
            zorder=3,
        )

    ax.set_xlim(
        0.4,
        len(labs) + 0.6,
    )

    ax.set_xticks(positions)

    ax.set_xticklabels(
        [str(lab) for lab in labs]
    )

    ax.set_xlabel("Laboratory")
    ax.set_ylabel("GP")


def annotate_target_values(
    ax,
    target_values,
    color,
):
    """
    ヒストグラム上に対象研究室のGPを縦線と数値で示す。

    仕様:
    - 同じGP値は1本の線にまとめる
    - ラベルは軸上端の内側に置く
    - 近い値が複数ある場合はラベル高さをずらす
    """

    unique_values = (
        pd.Series(target_values)
        .value_counts()
        .sort_index()
    )

    for i, (value, count) in enumerate(
        unique_values.items()
    ):
        ax.axvline(
            value,
            color=color,
            linewidth=1.3,
            alpha=0.85,
            linestyle="--",
        )

        if count == 1:
            label = f"{value:g}"
        else:
            label = f"{value:g} x{count}"

        ax.text(
            value,
            0.96 - 0.11 * (i % 3),
            label,
            transform=ax.get_xaxis_transform(),
            horizontalalignment="center",
            verticalalignment="top",
            rotation=90,
            color=color,
            fontsize=9,
            bbox={
                "facecolor": "white",
                "edgecolor": "none",
                "alpha": 0.75,
                "pad": 1.2,
            },
        )


def draw_gp_histogram(
    ax,
    df,
    target_values,
    bin_width,
    color,
    target_color,
):
    """
    全研究室のGPヒストグラムを描画する。

    仕様:
    - bin_widthごとの頻度を表示する
    - 対象研究室のGPは縦線で示す
    - 研究室別箱ひげ図と同じ主色を使う
    """

    bins = make_histogram_bins(
        df["値"],
        bin_width,
    )

    ax.hist(
        df["値"],
        bins=bins,
        color=color,
        alpha=0.62,
        edgecolor="black",
        linewidth=0.8,
    )

    if len(target_values) > 0:
        annotate_target_values(
            ax,
            target_values,
            target_color,
        )

    ax.set_xlabel("GP")
    ax.set_ylabel("Frequency")


def plot_data(
    data,
    lower=10,
    upper=None,
    target_lab=None,
    bin_width=10,
):
    """
    研究室配属希望調査のGP分布を描画する。

    Parameters
    ----------
    data : pandas.DataFrame
        「研究室番号」と「値」列を持つデータフレーム。
    lower : float, default 10
        値がlowerより大きい回答だけを描く。
    upper : float or None, default None
        GPの上限。Noneの場合は上限を設けない。
    target_lab : int, default None
        GP値を縦線表示する研究室番号。
    bin_width : float, default 10
        ヒストグラムのビン幅。

    Returns
    -------
    tuple
        生成したFigureと2個のAxes。

    仕様:
    - matplotlibとpandas/numpyのみで描画する
    - 図タイトルは付けない
    - 上段に研究室別箱ひげ図と個別点を描く
    - 下段に全体GPヒストグラムを描く
    - 編入生および回答者なしのNaNは描画しない
    """

    import matplotlib.pyplot as plt

    df = filter_plot_data(
        data.copy(),
        lower=lower,
        upper=upper,
    )

    labs = sorted(
        df["研究室番号"].unique()
    )

    target_values = df.loc[
        df["研究室番号"] == target_lab,
        "値",
    ].to_numpy(dtype=float)

    apply_plot_style()

    colors = plt.get_cmap("tab10").colors

    rng = np.random.default_rng(
        20260709
    )

    fig, axs = plt.subplots(
        2,
        1,
        figsize=(8.0, 6.2),
        gridspec_kw={
            "height_ratios": [1.1, 1],
        },
        constrained_layout=True,
    )

    draw_lab_boxplot(
        axs[0],
        df,
        labs,
        colors[0],
        rng,
    )

    draw_gp_histogram(
        axs[1],
        df,
        target_values,
        bin_width,
        colors[0],
        colors[3],
    )

    describe_text = (
        df["値"]
        .describe()
        .to_string(
            float_format=lambda x: f"{x:.1f}"
        )
    )

    # 下段ヒストグラムの左上にdescribe()を表示する。
    axs[1].text(
        0.01,
        0.99,
        describe_text,
        transform=axs[1].transAxes,
        horizontalalignment="left",
        verticalalignment="top",
        fontsize=8,
        fontfamily="monospace",
        bbox={
            "facecolor": "white",
            "edgecolor": "gray",
            "alpha": 0.85,
            "boxstyle": "round,pad=0.2",
        },
        zorder=1,
    )

    for ax in axs:
        style_axis(ax)

    plt.show()

    return fig, axs


def make_output_filename(
    input_filename,
    extension,
):
    """
    入力ファイル名の拡張子を変更した出力ファイル名を作る。
    """

    if "." in input_filename:
        base = input_filename.rsplit(".", 1)[0]
    else:
        base = input_filename

    return base + extension


def main():
    parser = argparse.ArgumentParser(
        description=(
            "Google Formsのコピペ結果から"
            "研究室別GPを集計する"
        )
    )

    parser.add_argument(
        "googleform_str",
        help=(
            "Google Formsからコピーした"
            "テキストファイル"
        ),
    )

    parser.add_argument(
        "-o",
        "--output",
        help=(
            "Excel出力ファイル名。"
            "指定しない場合は出力しない"
        ),
    )

    parser.add_argument(
        "-p",
        "--plot",
        action="store_true",
        help=(
            "各種plotを作成する。"
            "指定しない場合は作成しない"
        ),
    )

    args = parser.parse_args()

    with open(
        args.googleform_str,
        encoding="utf-8",
    ) as f:
        text = f.read()

    df = parse_google_form_results(text)
    df_r = summarize_by_lab(df)

    print_summary(
        df,
        df_r,
    )

    if args.plot:
        fig, axs = plot_data(df)

        plot_filename = make_output_filename(
            args.googleform_str,
            ".png",
        )

        fig.savefig(
            plot_filename,
            dpi=300,
        )

        print(
            f"{plot_filename} was created."
        )

    if args.output:
        with pd.ExcelWriter(args.output) as writer:
            df.to_excel(
                writer,
                sheet_name="raw",
                index=False,
            )

            df_r.to_excel(
                writer,
                sheet_name="summary",
            )

        print(
            f"{args.output} was created."
        )


if __name__ == "__main__":
    main()