#!/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 としては読まない
    - 「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 を生成する。
    """

    text = z2a(text)

    # 「研究室第1希望」などの集計部分を除外し、
    # 最初の個別研究室ブロックから開始する。
    #
    # 以下の両方に対応:
    #
    # 研究室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()

        if n_answers == 0:
            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)

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

    return df


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

    rows = []

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

        # 編入生は 0 として入れてある。
        df_regular = df_lab.loc[df_lab["値"] > 10]

        rows.append({
            "研究室": lab_no,
            "在校生": df_regular.shape[0],
            "編入生": df_lab.shape[0] - df_regular.shape[0],
            "GP mean": df_regular["値"].mean(),
            "GP std": df_regular["値"].std(),
        })

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


def print_summary(df, df_r):
    print("-" * 52)
    print("回答者数:", df.shape[0])
    print(
        f"{'研究室':>6} "
        f"{'在校生':>6} "
        f"{'編入生':>6} "
        f"{'GP mean':>10} "
        f"{'GP std':>10}"
    )
    print("-" * 52)

    for n, 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"{n:>6} "
            f"{int(row['在校生']):>8} "
            f"{int(row['編入生']):>8} "
            f"{mean:>10} "
            f"{std:>10}"
        )


def main():
    parser = argparse.ArgumentParser(
        description="Google Forms のコピペ結果から研究室別GPを集計する"
    )
    parser.add_argument("googleform_str", help="Google Forms からコピーしたテキストファイル")
    parser.add_argument(
        "-o", "--output",
        help="Excel 出力ファイル名。指定しない場合は出力しない"
    )

    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.output:
        with pd.ExcelWriter(args.output) as writer:
            df.to_excel(writer, sheet_name="raw")
            df_r.to_excel(writer, sheet_name="summary")

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


if __name__ == "__main__":
    main()