"""tkdataio.py

軽量な表データ入出力ユーティリティ。

既存の tkVariousData/tkFit を置き換える目的ではなく、独立した小さな解析スクリプトを
素早く作るための補助モジュール。
"""

from __future__ import annotations

from pathlib import Path
from typing import Dict, Mapping, Optional, Sequence, Tuple, Union
import json
import numpy as np

ColumnKey = Union[int, str]


def read_table(path: Union[str, Path], *, sheet_name=0):
    """CSV/TSV/TXT/XLSX を pandas.DataFrame として読む。"""
    try:
        import pandas as pd
    except Exception as exc:  # pragma: no cover
        raise ImportError("read_table requires pandas. Install with: pip install pandas openpyxl") from exc

    path = Path(path)
    ext = path.suffix.lower()
    if ext == ".csv":
        return pd.read_csv(path)
    if ext in [".tsv", ".txt", ".dat"]:
        # まずタブ、失敗しやすい空白区切りはユーザ側で pandas を直接使う想定
        return pd.read_csv(path, sep="\t")
    if ext in [".xlsx", ".xlsm", ".xls"]:
        return pd.read_excel(path, sheet_name=sheet_name)
    raise ValueError(f"Unsupported table format: {path.suffix}")


def _column_from_key(df, key: ColumnKey):
    """DataFrame から列を index または label で取り出す。"""
    if isinstance(key, int):
        return df.columns[key], df.iloc[:, key].to_numpy(dtype=float)

    # 数字文字列は index としても扱えるようにする
    if isinstance(key, str):
        ks = key.strip()
        if ks.lstrip("+-").isdigit():
            idx = int(ks)
            return df.columns[idx], df.iloc[:, idx].to_numpy(dtype=float)
        if key in df.columns:
            return key, df[key].to_numpy(dtype=float)

    raise KeyError(f"Column {key!r} not found")


def read_xy(
    path: Union[str, Path],
    x: ColumnKey = 0,
    y: ColumnKey = 1,
    *,
    sheet_name=0,
    xmin: float = -1e100,
    xmax: float = 1e100,
    dropna: bool = True,
) -> Tuple[np.ndarray, np.ndarray, str, str]:
    """表ファイルから x, y の2列を取り出す。"""
    df = read_table(path, sheet_name=sheet_name)
    xlabel, xv = _column_from_key(df, x)
    ylabel, yv = _column_from_key(df, y)
    xv = np.asarray(xv, dtype=float)
    yv = np.asarray(yv, dtype=float)

    mask = (xv >= xmin) & (xv <= xmax)
    if dropna:
        mask &= np.isfinite(xv) & np.isfinite(yv)
    return xv[mask], yv[mask], str(xlabel), str(ylabel)


def write_excel_tables(
    path: Union[str, Path],
    tables: Mapping[str, object],
    *,
    index: bool = False,
) -> None:
    """複数テーブルを Excel の複数シートに保存する。

    tables の値は pandas.DataFrame、dict、2D array のいずれかを想定。
    """
    try:
        import pandas as pd
    except Exception as exc:  # pragma: no cover
        raise ImportError("write_excel_tables requires pandas/openpyxl") from exc

    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    with pd.ExcelWriter(path, engine="openpyxl") as writer:
        for sheet, obj in tables.items():
            sheet_name = str(sheet)[:31] or "Sheet"
            if hasattr(obj, "to_excel"):
                df = obj
            elif isinstance(obj, Mapping):
                # dict of arrays/scalars
                maxlen = 1
                for v in obj.values():
                    if isinstance(v, (str, bytes)):
                        continue
                    try:
                        maxlen = max(maxlen, len(v))
                    except TypeError:
                        pass
                data = {}
                for k, v in obj.items():
                    if isinstance(v, (str, bytes)):
                        data[k] = [v] + [None] * (maxlen - 1)
                    else:
                        try:
                            vv = list(v)
                            data[k] = vv + [None] * (maxlen - len(vv))
                        except TypeError:
                            data[k] = [v] + [None] * (maxlen - 1)
                df = pd.DataFrame(data)
            else:
                arr = np.asarray(obj)
                df = pd.DataFrame(arr)
            df.to_excel(writer, sheet_name=sheet_name, index=index)


def save_json(path: Union[str, Path], data: object, *, indent: int = 2) -> None:
    """JSON保存。NumPy型はPython型へ変換する。"""
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)

    def convert(o):
        if isinstance(o, np.ndarray):
            return o.tolist()
        if isinstance(o, (np.floating, np.integer)):
            return o.item()
        if hasattr(o, "__dict__"):
            return o.__dict__
        raise TypeError(f"Object of type {type(o).__name__} is not JSON serializable")

    with open(path, "w", encoding="utf-8") as f:
        json.dump(data, f, ensure_ascii=False, indent=indent, default=convert)


def load_json(path: Union[str, Path]) -> object:
    """JSON読み込み。"""
    with open(path, "r", encoding="utf-8") as f:
        return json.load(f)