#!/usr/bin/env python
import numpy as np
import re
import pandas as pd


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

        研究室番号 | 値

    の DataFrame を生成する。

    Parameters
    ----------
    text : str
        Google Forms 結果ページからコピーしたテキスト

    Returns
    -------
    pandas.DataFrame
    """

    # 「第一希希望研究室...」の集計部分は除外したいので、
    # 最初の個別研究室ブロックから開始する
    m = re.search(
        r'^研究室\s*[0-9０-９]+.*?\n[0-9０-９]+\s*件の回答',
        text,
        flags=re.MULTILINE
    )

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

    text = text[m.start():]

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

    rows = []

    for match in pattern.finditer(text):
        lab_no_str = match.group(1)
        n_answers_str = match.group(2)
        body = match.group(3).strip()

        # 全角数字対応
        trans = str.maketrans("０１２３４５６７８９", "0123456789")

        lab_no = int(lab_no_str.translate(trans))
        n_answers = int(n_answers_str.translate(trans))

        # 回答数 0 はスキップ
        if n_answers == 0:
            continue

        # --------------------------------------------------
        # Case 1:
        # 「値  カウント」表がある場合
        #
        # 例:
        # 値    カウント
        # 300   2
        # 350   2
        # --------------------------------------------------
        if re.search(r'^値\s+カウント\s*$', body, flags=re.MULTILINE):
            table_part = re.split(
                r'^値\s+カウント\s*$',
                body,
                maxsplit=1,
                flags=re.MULTILINE
            )[1]

            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)

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

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

                if re.fullmatch(r'-?\d+(?:\.\d+)?', line):
                    values.append(float(line))

        # 回答数との整合性チェック
        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:
        if (df["値"] % 1 == 0).all():
            df["値"] = df["値"].astype(int)

    return df


if __name__ == "__main__":
    import argparse

    par = argparse.ArgumentParser(description="test")
    par.add_argument('googleform_str')
    args = par.parse_args()

    # コピペしたテキストをファイルに保存しておく
    with open(args.googleform_str, encoding="utf-8") as f:
        text = f.read()

    df = parse_google_form_results(text)
    df_r = pd.DataFrame()
    for n in np.unique(df["研究室番号"]):
        df_ = df.loc[df["研究室番号"] == n, :]
        df__ = df_.loc[df_["値"] > 10, :]
        ave, std = df__.mean(), df__.std()
        df_r.loc[n, "在校生"] = df_.shape[0]
        df_r.loc[n, "編入生"] = df_.shape[0] - df__.shape[0]
        df_r.loc[n, "GPA mean"] = df__["値"].mean()
        df_r.loc[n, "GPA std"] = df__["値"].std()
    df_r[["在校生", "編入生"]] = df_r[["在校生", "編入生"]].astype(int)
    # 表示用
    df_show = df_r.copy()
    df_show.index.name = "研究室"
    df_show["GPA mean"] = df_show["GPA mean"].round(1)
    df_show["GPA std"] = df_show["GPA std"].round(1)
    print(
        f"{'研究室':>6} "
        f"{'在校生':>6} "
        f"{'編入生':>6} "
        f"{'GPA mean':>10} "
        f"{'GPA std':>10}"
    )
    print("-" * 52)
    for n, row in df_r.iterrows():
        mean = f"{row['GPA mean']:.1f}"
        std = "-" if pd.isna(row["GPA std"]) else f"{row['GPA std']:.1f}"
        print(
        f"{n:>6} "
        f"{int(row['在校生']):>8} "
        f"{int(row['編入生']):>8} "
        f"{mean:>10} "
        f"{std:>10}")    

    # outfile = __file__.replace(".py", ".xlsx")
    # df.to_excel(outfile)
    # print(f"{outfile} was created.")