#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""poly_lsq_app.py

order 次多項式回帰のサンプルアプリ。

前提:
  - このファイルと同じディレクトリ、または PYTHONPATH 上に tklsq.py を置く。
  - 入出力・描画・Excel保存はこのアプリ側で担当し、回帰計算は tklsq.py に任せる。

Modes:
  read          : 入力データを読み、Excel保存・散布図表示
  sim           : 多項式のテストデータを生成し、Excel保存・グラフ表示
  fit           : order 次多項式で最小二乗回帰し、誤差評価つきで保存・描画
  model_select  : 0..order 次、または min_order..max_order 次を AIC/BIC/AICc/evidence で比較

Examples:
  python poly_lsq_app.py --mode sim --order 3 --show 1
  python poly_lsq_app.py --mode fit --infile data.xlsx --xlabel 0 --ylabel 1 --order 3
  python poly_lsq_app.py --mode model_select --infile data.xlsx --order 6 --criterion BIC
  python poly_lsq_app.py --mode model_select --criterion evidence --beta 25
"""

from __future__ import annotations

import argparse
import csv
import math
import os
import sys
import traceback
from dataclasses import dataclass
from itertools import zip_longest
from pathlib import Path
from typing import Dict, Iterable, List, Optional, Sequence, Tuple

import numpy as np

try:
    import matplotlib.pyplot as plt
except ImportError:
    print("ERROR: matplotlib is required.")
    print("pip install matplotlib")
    traceback.print_exc()
    sys.exit(1)

try:
    import openpyxl
    from openpyxl import Workbook
except ImportError:
    print("ERROR: openpyxl is required.")
    print("pip install openpyxl")
    traceback.print_exc()
    sys.exit(1)

try:
    import tklsq
except ImportError:
    print("ERROR: tklsq.py was not found.")
    print("Place tklsq.py in the same directory as this script, or add it to PYTHONPATH.")
    traceback.print_exc()
    sys.exit(1)


# =====================================================================
# Data containers
# =====================================================================

@dataclass
class XYData:
    x: np.ndarray
    y: np.ndarray
    xlabel: str = "x"
    ylabel: str = "y"
    y_true: Optional[np.ndarray] = None
    source: str = ""


# =====================================================================
# argparse
# =====================================================================

def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Polynomial least-squares regression using tklsq.py"
    )

    parser.add_argument(
        "--mode",
        type=str,
        choices=["read", "sim", "fit", "model_select"],
        default="fit",
        help="Execution mode",
    )
    parser.add_argument("--infile", type=str, default="", help="Input CSV/TSV/TXT/XLSX file")
    parser.add_argument("--sheet", type=str, default="", help="Excel sheet name. Empty = active sheet")
    parser.add_argument("--xlabel", type=str, default="0", help="x column label or 0-based index")
    parser.add_argument("--ylabel", type=str, default="1", help="y column label or 0-based index")

    parser.add_argument("--order", type=int, default=3, help="Polynomial order for fit, or max order for model_select")
    parser.add_argument("--min_order", type=int, default=0, help="Minimum order for model_select")
    parser.add_argument("--max_order", type=int, default=-1, help="Maximum order for model_select. -1 = order")
    parser.add_argument(
        "--criterion",
        type=str,
        choices=["AIC", "AICc", "BIC", "evidence"],
        default="BIC",
        help="Model selection criterion",
    )
    parser.add_argument("--alpha", type=float, default=1.0, help="Bayesian prior precision for evidence")
    parser.add_argument(
        "--beta",
        type=float,
        default=0.0,
        help="Noise precision for evidence. <=0 means estimate from highest-order LSQ residual.",
    )

    parser.add_argument("--xmin", type=float, default=-1.0e300, help="Minimum x used for fitting/selection")
    parser.add_argument("--xmax", type=float, default=1.0e300, help="Maximum x used for fitting/selection")
    parser.add_argument("--xcalmin", type=str, default="*", help="Minimum x for calculation grid. '*' = min(x)")
    parser.add_argument("--xcalmax", type=str, default="*", help="Maximum x for calculation grid. '*' = max(x)")
    parser.add_argument("--ncal", type=int, default=301, help="Number of calculation-grid points")

    parser.add_argument("--n", type=int, default=100, help="Number of simulated data points")
    parser.add_argument("--sim_xmin", type=float, default=-3.0, help="Simulation x minimum")
    parser.add_argument("--sim_xmax", type=float, default=3.0, help="Simulation x maximum")
    parser.add_argument("--noise_sigma", type=float, default=1.0, help="Simulation noise standard deviation")
    parser.add_argument("--seed", type=int, default=0, help="Random seed for simulation")
    parser.add_argument(
        "--true_coeffs",
        type=str,
        default="1,-2,0.5,0.1",
        help="Comma-separated true coefficients for simulation, e.g. '1,-2,0.5,0.1'",
    )

    parser.add_argument("--outfile", type=str, default="*", help="Output Excel path. '*' = automatic")
    parser.add_argument("--figfile", type=str, default="*", help="Output figure path. '*' = automatic")
    parser.add_argument("--outdir", type=str, default="*", help="Output directory. '*' = input directory or current directory")

    # Launcher-friendly 0/1 options
    parser.add_argument("--show", type=int, choices=[0, 1], default=1, help="Show graph window")
    parser.add_argument("--savefig", type=int, choices=[0, 1], default=1, help="Save figure")
    parser.add_argument("--plot_sigma_param", type=int, choices=[0, 1], default=1, help="Plot ±sigma(param)")
    parser.add_argument("--plot_sigma_pred", type=int, choices=[0, 1], default=1, help="Plot ±sigma(param&resid)")
    parser.add_argument("--plot_param", type=int, choices=[0, 1], default=1, help="Plot parameter errorbar panel")
    parser.add_argument("--pause", type=int, choices=[0, 1], default=0, help="Pause before exit")

    parser.add_argument("--figsize", type=float, nargs=2, default=[10.0, 5.0], help="Figure size")
    parser.add_argument("--fontsize", type=int, default=14, help="Font size")
    parser.add_argument("--fontsize_legend", type=int, default=11, help="Legend font size")

    return parser


# =====================================================================
# Utilities
# =====================================================================

def pause_if_needed(flag: int) -> None:
    if flag:
        try:
            input("\nPress ENTER to terminate>>\n")
        except EOFError:
            pass


def parse_float_or_star(value: str, default: float) -> float:
    if value is None or str(value).strip() == "" or str(value).strip() == "*":
        return float(default)
    return float(value)


def parse_coeffs(text: str) -> np.ndarray:
    try:
        coeffs = [float(v.strip()) for v in text.split(",") if v.strip() != ""]
    except ValueError as exc:
        raise ValueError(f"Invalid --true_coeffs: {text}") from exc
    if len(coeffs) == 0:
        raise ValueError("--true_coeffs must contain at least one coefficient")
    return np.asarray(coeffs, dtype=float)


def poly_eval(x: np.ndarray, coeffs: Sequence[float]) -> np.ndarray:
    x = np.asarray(x, dtype=float)
    y = np.zeros_like(x, dtype=float)
    for i, c in enumerate(coeffs):
        y += float(c) * x ** i
    return y


def is_number(v) -> bool:
    if v is None:
        return False
    try:
        float(v)
        return True
    except (TypeError, ValueError):
        return False


def to_float_or_nan(v) -> float:
    try:
        if v is None or v == "":
            return float("nan")
        return float(v)
    except (TypeError, ValueError):
        return float("nan")


def sanitize_sheet_name(name: str) -> str:
    invalid = "[]:*?/\\"
    for ch in invalid:
        name = name.replace(ch, "_")
    return name[:31] if len(name) > 31 else name


def excel_value(v):
    if v is None:
        return None
    if isinstance(v, np.generic):
        return v.item()
    if isinstance(v, float):
        if math.isnan(v) or math.isinf(v):
            return None
    return v


def write_columns_sheet(wb: Workbook, sheet_name: str, columns: Dict[str, Sequence]) -> None:
    ws = wb.create_sheet(sanitize_sheet_name(sheet_name))
    headers = list(columns.keys())
    ws.append(headers)
    seqs = []
    for h in headers:
        v = columns[h]
        if v is None:
            v = []
        elif np.isscalar(v):
            v = [v]
        else:
            v = list(v)
        seqs.append(v)
    for row in zip_longest(*seqs, fillvalue=None):
        ws.append([excel_value(v) for v in row])


def write_rows_sheet(wb: Workbook, sheet_name: str, headers: Sequence[str], rows: Iterable[Sequence]) -> None:
    ws = wb.create_sheet(sanitize_sheet_name(sheet_name))
    ws.append(list(headers))
    for row in rows:
        ws.append([excel_value(v) for v in row])


def write_key_value_sheet(wb: Workbook, sheet_name: str, items: Dict[str, object]) -> None:
    ws = wb.create_sheet(sanitize_sheet_name(sheet_name))
    ws.append(["key", "value"])
    for k, v in items.items():
        ws.append([str(k), excel_value(v)])


def save_workbook(path: str, sheets_callback) -> None:
    wb = Workbook()
    # remove default sheet after adding our sheets
    default = wb.active
    wb.remove(default)
    sheets_callback(wb)
    Path(path).parent.mkdir(parents=True, exist_ok=True)
    wb.save(path)


def get_output_base(args, suffix: str) -> Tuple[str, str]:
    infile = Path(args.infile) if args.infile else None
    if args.outdir == "*":
        if infile and str(infile.parent) not in ("", "."):
            outdir = infile.parent
        else:
            outdir = Path(".")
    else:
        outdir = Path(args.outdir)

    if infile and infile.name:
        stem = infile.stem
    else:
        stem = "poly_lsq"
    base = outdir / f"{stem}_{suffix}"
    return str(outdir), str(base)


def get_output_paths(args, suffix: str) -> Tuple[str, str]:
    _outdir, base = get_output_base(args, suffix)
    outfile = f"{base}.xlsx" if args.outfile == "*" else args.outfile
    figfile = f"{base}.png" if args.figfile == "*" else args.figfile
    return outfile, figfile


# =====================================================================
# Input data
# =====================================================================

def rows_to_table(rows: List[List[object]]) -> Tuple[List[str], List[List[object]]]:
    rows = [list(r) for r in rows if r is not None and len(r) > 0]
    if not rows:
        raise ValueError("No data rows were found")

    first = rows[0]
    first_numeric = all(is_number(v) or v is None or v == "" for v in first)
    if first_numeric:
        labels = [str(i) for i in range(len(first))]
        data_rows = rows
    else:
        labels = [str(v) if v is not None else str(i) for i, v in enumerate(first)]
        data_rows = rows[1:]
    return labels, data_rows


def read_table(path: str, sheet: str = "") -> Tuple[List[str], List[List[object]]]:
    p = Path(path)
    if not p.exists():
        raise FileNotFoundError(path)
    ext = p.suffix.lower()

    if ext in [".csv"]:
        with open(p, "r", encoding="utf-8-sig", newline="") as f:
            rows = list(csv.reader(f))
        return rows_to_table(rows)

    if ext in [".txt", ".tsv"]:
        with open(p, "r", encoding="utf-8-sig", newline="") as f:
            rows = list(csv.reader(f, delimiter="\t"))
        return rows_to_table(rows)

    if ext in [".xlsx", ".xlsm"]:
        wb = openpyxl.load_workbook(p, data_only=True, read_only=True)
        if sheet:
            if sheet not in wb.sheetnames:
                raise ValueError(f"Sheet [{sheet}] was not found. Available: {wb.sheetnames}")
            ws = wb[sheet]
        else:
            ws = wb.active
        rows = [list(r) for r in ws.iter_rows(values_only=True)]
        return rows_to_table(rows)

    raise ValueError(f"Unsupported input format: {ext}")


def find_column(labels: Sequence[str], key: str) -> int:
    key_s = str(key).strip()
    try:
        idx = int(key_s)
        if 0 <= idx < len(labels):
            return idx
        raise IndexError
    except ValueError:
        pass
    except IndexError as exc:
        raise ValueError(f"Column index {key_s} is out of range. ncols={len(labels)}") from exc

    labels_str = [str(v).strip() for v in labels]
    if key_s in labels_str:
        return labels_str.index(key_s)
    lower = [v.lower() for v in labels_str]
    if key_s.lower() in lower:
        return lower.index(key_s.lower())
    raise ValueError(f"Column [{key}] was not found. labels={labels}")


def load_xy_from_file(args) -> XYData:
    labels, rows = read_table(args.infile, sheet=args.sheet)
    ix = find_column(labels, args.xlabel)
    iy = find_column(labels, args.ylabel)

    x_list = []
    y_list = []
    for r in rows:
        if ix >= len(r) or iy >= len(r):
            continue
        x = to_float_or_nan(r[ix])
        y = to_float_or_nan(r[iy])
        if np.isfinite(x) and np.isfinite(y) and args.xmin <= x <= args.xmax:
            x_list.append(x)
            y_list.append(y)

    if not x_list:
        raise ValueError("No valid numeric x-y data were found in the selected range")

    return XYData(
        x=np.asarray(x_list, dtype=float),
        y=np.asarray(y_list, dtype=float),
        xlabel=str(labels[ix]),
        ylabel=str(labels[iy]),
        source=args.infile,
    )


def simulate_xy(args) -> XYData:
    rng = np.random.default_rng(args.seed)
    coeffs = parse_coeffs(args.true_coeffs)
    x = np.linspace(args.sim_xmin, args.sim_xmax, args.n)
    y_true = poly_eval(x, coeffs)
    y = y_true + rng.normal(0.0, args.noise_sigma, size=args.n)
    return XYData(x=x, y=y, y_true=y_true, xlabel="x", ylabel="y", source="simulation")


def get_xy(args) -> XYData:
    if args.infile:
        return load_xy_from_file(args)
    if args.mode in ["sim", "fit", "model_select"]:
        print("No --infile specified. Simulated data will be used.")
        return simulate_xy(args)
    raise ValueError("--infile is required for mode=read")


# =====================================================================
# Analysis helpers
# =====================================================================

def make_xcal(x: np.ndarray, args) -> np.ndarray:
    xcalmin = parse_float_or_star(args.xcalmin, float(np.min(x)))
    xcalmax = parse_float_or_star(args.xcalmax, float(np.max(x)))
    if xcalmax == xcalmin:
        dx = 0.5 if xcalmax == 0 else abs(xcalmax) * 0.05
        xcalmin -= dx
        xcalmax += dx
    return np.linspace(xcalmin, xcalmax, args.ncal)


def fit_polynomial(xy: XYData, order: int):
    result = tklsq.polynomial_lsq(xy.x, xy.y, order=order)
    xcal = None
    pred = None
    return result


def estimate_beta_for_evidence(x: np.ndarray, y: np.ndarray, max_order: int) -> float:
    max_order = max(0, int(max_order))
    result = tklsq.polynomial_lsq(x, y, order=max_order)
    if result.sigma2_resid is not None and result.sigma2_resid > 0:
        return 1.0 / result.sigma2_resid
    var_y = float(np.var(y, ddof=1)) if len(y) > 1 else 1.0
    return 1.0 / max(var_y, np.finfo(float).eps)


def run_model_selection(xy: XYData, args):
    max_order = args.order if args.max_order < 0 else args.max_order
    max_order = max(max_order, args.min_order)
    candidates = tklsq.polynomial_candidates(xy.x, max_order=max_order, min_order=args.min_order)

    if args.criterion == "evidence":
        beta = args.beta if args.beta > 0 else estimate_beta_for_evidence(xy.x, xy.y, max_order)
        sel = tklsq.select_models_by_evidence(
            candidates,
            xy.y,
            alpha=args.alpha,
            beta=beta,
        )
        criterion_value_name = "log_evidence"
        used_beta = beta
    else:
        sel = tklsq.select_models_by_ic(candidates, xy.y, criterion=args.criterion)
        criterion_value_name = sel.criterion
        used_beta = None

    # Fit all models as LSQ too, for diagnostics and plotting.
    fit_results = []
    for label, X in candidates.items():
        fit_results.append(tklsq.linear_lsq(X, xy.y))

    return sel, fit_results, criterion_value_name, used_beta


# =====================================================================
# Excel output
# =====================================================================

def save_read_or_sim_excel(path: str, xy: XYData, args, mode: str) -> None:
    def sheets(wb: Workbook):
        cols = {xy.xlabel: xy.x, xy.ylabel: xy.y}
        if xy.y_true is not None:
            cols["y_true"] = xy.y_true
            cols["residual_noise"] = xy.y - xy.y_true
        write_columns_sheet(wb, "data", cols)
        write_key_value_sheet(
            wb,
            "info",
            {
                "mode": mode,
                "source": xy.source,
                "N": len(xy.x),
                "xlabel": xy.xlabel,
                "ylabel": xy.ylabel,
                "order": args.order,
                "noise_sigma": args.noise_sigma if mode == "sim" else None,
                "true_coeffs": args.true_coeffs if mode == "sim" else None,
            },
        )
    save_workbook(path, sheets)


def save_fit_excel(path: str, xy: XYData, result, xcal: np.ndarray, pred, args, label: str = "fit") -> None:
    X_data = tklsq.polynomial_design_matrix(xy.x, order=result.p - 1)
    y_fit_data = X_data @ result.beta
    scores = tklsq.regression_scores(xy.y, y_fit_data, p=result.p)

    def sheets(wb: Workbook):
        data_cols = {xy.xlabel: xy.x, xy.ylabel: xy.y, f"{xy.ylabel}_fit": y_fit_data, "residual": xy.y - y_fit_data}
        if xy.y_true is not None:
            data_cols["y_true"] = xy.y_true
        write_columns_sheet(wb, "fit_data", data_cols)

        pred_cols = {xy.xlabel: xcal, f"{xy.ylabel}_mean": pred.y_mean}
        if pred.sigma_param is not None:
            pred_cols["sigma_param"] = pred.sigma_param
            pred_cols["mean-sigma_param"] = pred.y_mean - pred.sigma_param
            pred_cols["mean+sigma_param"] = pred.y_mean + pred.sigma_param
        if pred.sigma_pred is not None:
            pred_cols["sigma_pred"] = pred.sigma_pred
            pred_cols["mean-sigma_pred"] = pred.y_mean - pred.sigma_pred
            pred_cols["mean+sigma_pred"] = pred.y_mean + pred.sigma_pred
        write_columns_sheet(wb, "prediction", pred_cols)

        rows = []
        for i, b in enumerate(result.beta):
            std = None if result.beta_std is None else result.beta_std[i]
            rows.append([i, b, std])
        write_rows_sheet(wb, "parameters", ["i", "coeff", "std"], rows)

        diag = {
            "label": label,
            "source": xy.source,
            "N": result.N,
            "p": result.p,
            "dof": result.dof,
            "rank": result.rank,
            "RSS": result.RSS,
            "WRSS": result.WRSS,
            "sigma2_resid": result.sigma2_resid,
            "sigma_resid": result.sigma_resid,
            "condition_number": result.condition_number,
            "error_estimation_enabled": int(result.error_estimation_enabled),
            "warning": result.warning,
        }
        diag.update(scores)
        write_key_value_sheet(wb, "diagnostics", diag)

    save_workbook(path, sheets)


def save_model_select_excel(
    path: str,
    xy: XYData,
    sel,
    fit_results,
    criterion_value_name: str,
    used_beta: Optional[float],
    best_result,
    xcal: np.ndarray,
    best_pred,
    args,
) -> None:
    def parse_order(label: str) -> Optional[int]:
        if label.startswith("poly"):
            try:
                return int(label.replace("poly", ""))
            except ValueError:
                return None
        return None

    def sheets(wb: Workbook):
        rows = []
        for i, label in enumerate(sel.labels):
            r = fit_results[i]
            score = sel.scores[i]
            rows.append([
                i,
                label,
                parse_order(label),
                criterion_value_name,
                score,
                sel.weights[i],
                r.N,
                r.p,
                r.dof,
                r.RSS,
                r.sigma2_resid,
                r.condition_number,
                int(r.error_estimation_enabled),
                r.warning,
            ])
        write_rows_sheet(
            wb,
            "model_selection",
            [
                "index",
                "label",
                "order",
                "criterion",
                "score",
                "weight",
                "N",
                "p",
                "dof",
                "RSS",
                "sigma2_resid",
                "condition_number",
                "error_estimation_enabled",
                "warning",
            ],
            rows,
        )

        X_data = tklsq.polynomial_design_matrix(xy.x, order=best_result.p - 1)
        best_y_fit_data = X_data @ best_result.beta
        data_cols = {xy.xlabel: xy.x, xy.ylabel: xy.y, f"{xy.ylabel}_best_fit": best_y_fit_data, "residual": xy.y - best_y_fit_data}
        if xy.y_true is not None:
            data_cols["y_true"] = xy.y_true
        write_columns_sheet(wb, "best_fit_data", data_cols)

        pred_cols = {xy.xlabel: xcal, f"{xy.ylabel}_best_mean": best_pred.y_mean}
        if best_pred.sigma_param is not None:
            pred_cols["sigma_param"] = best_pred.sigma_param
        if best_pred.sigma_pred is not None:
            pred_cols["sigma_pred"] = best_pred.sigma_pred
        write_columns_sheet(wb, "best_prediction", pred_cols)

        param_rows = []
        for i, b in enumerate(best_result.beta):
            std = None if best_result.beta_std is None else best_result.beta_std[i]
            param_rows.append([i, b, std])
        write_rows_sheet(wb, "best_parameters", ["i", "coeff", "std"], param_rows)

        write_key_value_sheet(
            wb,
            "info",
            {
                "source": xy.source,
                "criterion": sel.criterion,
                "best_index": sel.best_index,
                "best_label": sel.best_label,
                "used_alpha": args.alpha if args.criterion == "evidence" else None,
                "used_beta": used_beta,
                "min_order": args.min_order,
                "max_order": args.order if args.max_order < 0 else args.max_order,
            },
        )

    save_workbook(path, sheets)


# =====================================================================
# Plotting
# =====================================================================

def apply_axis_style(ax, args) -> None:
    ax.tick_params(labelsize=args.fontsize)
    ax.grid(True, alpha=0.3)


def finalize_figure(fig, figfile: str, args) -> None:
    fig.tight_layout()
    if args.savefig:
        Path(figfile).parent.mkdir(parents=True, exist_ok=True)
        fig.savefig(figfile, dpi=200)
        print(f"Save figure: {figfile}")
    if args.show:
        plt.show()
    else:
        plt.close(fig)


def plot_read_or_sim(xy: XYData, figfile: str, args, mode: str) -> None:
    fig, ax = plt.subplots(figsize=args.figsize)
    ax.plot(xy.x, xy.y, "o", markersize=4, label="data")
    if xy.y_true is not None:
        ax.plot(xy.x, xy.y_true, "-", linewidth=1.2, label="true")
    ax.set_xlabel(xy.xlabel, fontsize=args.fontsize)
    ax.set_ylabel(xy.ylabel, fontsize=args.fontsize)
    ax.set_title(mode, fontsize=args.fontsize)
    ax.legend(fontsize=args.fontsize_legend)
    apply_axis_style(ax, args)
    finalize_figure(fig, figfile, args)


def plot_fit(xy: XYData, result, xcal: np.ndarray, pred, figfile: str, args, title: str = "fit") -> None:
    if args.plot_param and result.beta_std is not None:
        fig, (ax, axp) = plt.subplots(1, 2, figsize=args.figsize)
    else:
        fig, ax = plt.subplots(figsize=args.figsize)
        axp = None

    ax.plot(xy.x, xy.y, "o", markersize=4, label="data")
    if xy.y_true is not None:
        ax.plot(xy.x, xy.y_true, "--", linewidth=1.0, label="true")
    ax.plot(xcal, pred.y_mean, "-", linewidth=1.4, label="fit")

    if args.plot_sigma_pred and pred.sigma_pred is not None:
        ax.fill_between(
            xcal,
            pred.y_mean - pred.sigma_pred,
            pred.y_mean + pred.sigma_pred,
            alpha=0.20,
            label="±σ(param&resid)",
        )
    if args.plot_sigma_param and pred.sigma_param is not None:
        ax.fill_between(
            xcal,
            pred.y_mean - pred.sigma_param,
            pred.y_mean + pred.sigma_param,
            alpha=0.30,
            label="±σ(param)",
        )

    if pred.sigma_param is None:
        ax.text(
            0.03,
            0.97,
            "parameter error disabled",
            transform=ax.transAxes,
            va="top",
            fontsize=args.fontsize_legend,
        )

    ax.set_xlabel(xy.xlabel, fontsize=args.fontsize)
    ax.set_ylabel(xy.ylabel, fontsize=args.fontsize)
    ax.set_title(title, fontsize=args.fontsize)
    ax.legend(fontsize=args.fontsize_legend)
    apply_axis_style(ax, args)

    if axp is not None:
        idx = np.arange(result.p)
        axp.errorbar(idx, result.beta, yerr=result.beta_std, fmt="o", capsize=3)
        axp.set_xlabel("coefficient index", fontsize=args.fontsize)
        axp.set_ylabel("coefficient", fontsize=args.fontsize)
        axp.set_title("parameters", fontsize=args.fontsize)
        apply_axis_style(axp, args)

    finalize_figure(fig, figfile, args)


def plot_model_selection(xy: XYData, sel, best_result, xcal: np.ndarray, best_pred, figfile: str, args) -> None:
    fig, (ax, axw) = plt.subplots(1, 2, figsize=args.figsize)

    ax.plot(xy.x, xy.y, "o", markersize=4, label="data")
    if xy.y_true is not None:
        ax.plot(xy.x, xy.y_true, "--", linewidth=1.0, label="true")
    ax.plot(xcal, best_pred.y_mean, "-", linewidth=1.4, label=f"best: {sel.best_label}")
    if best_pred.sigma_pred is not None and args.plot_sigma_pred:
        ax.fill_between(
            xcal,
            best_pred.y_mean - best_pred.sigma_pred,
            best_pred.y_mean + best_pred.sigma_pred,
            alpha=0.20,
            label="±σ(param&resid)",
        )
    ax.set_xlabel(xy.xlabel, fontsize=args.fontsize)
    ax.set_ylabel(xy.ylabel, fontsize=args.fontsize)
    ax.legend(fontsize=args.fontsize_legend)
    apply_axis_style(ax, args)

    xpos = np.arange(len(sel.labels))
    axw.bar(xpos, sel.weights)
    axw.set_xticks(xpos)
    axw.set_xticklabels(sel.labels, rotation=45, ha="right", fontsize=args.fontsize_legend)
    axw.set_ylabel("model weight", fontsize=args.fontsize)
    axw.set_title(f"criterion: {sel.criterion}", fontsize=args.fontsize)
    apply_axis_style(axw, args)

    finalize_figure(fig, figfile, args)


# =====================================================================
# Modes
# =====================================================================

def mode_read(args) -> None:
    xy = get_xy(args)
    outfile, figfile = get_output_paths(args, "read")
    save_read_or_sim_excel(outfile, xy, args, "read")
    print(f"Save Excel: {outfile}")
    print(f"N = {len(xy.x)}")
    plot_read_or_sim(xy, figfile, args, "read")


def mode_sim(args) -> None:
    xy = simulate_xy(args)
    outfile, figfile = get_output_paths(args, "sim")
    save_read_or_sim_excel(outfile, xy, args, "sim")
    print(f"Save Excel: {outfile}")
    print(f"N = {len(xy.x)}")
    print(f"true_coeffs = {parse_coeffs(args.true_coeffs)}")
    plot_read_or_sim(xy, figfile, args, "sim")


def mode_fit(args) -> None:
    xy = get_xy(args)
    result = tklsq.polynomial_lsq(xy.x, xy.y, order=args.order)
    xcal = make_xcal(xy.x, args)
    pred = tklsq.predict_polynomial(xcal, result, order=args.order)

    print("# Fit diagnostics")
    print(f"N={result.N}, p={result.p}, dof={result.dof}, rank={result.rank}, RSS={result.RSS:.6g}")
    print(f"sigma2_resid={result.sigma2_resid}")
    print(f"condition_number={result.condition_number:.6g}")
    if result.warning:
        print(result.warning)
    print("# Coefficients")
    for i, b in enumerate(result.beta):
        if result.beta_std is None:
            print(f"c{i} = {b:.12g}")
        else:
            print(f"c{i} = {b:.12g} +- {result.beta_std[i]:.6g}")

    outfile, figfile = get_output_paths(args, f"fit_order{args.order}")
    save_fit_excel(outfile, xy, result, xcal, pred, args, label=f"poly{args.order}")
    print(f"Save Excel: {outfile}")
    plot_fit(xy, result, xcal, pred, figfile, args, title=f"polynomial LSQ: order={args.order}")


def mode_model_select(args) -> None:
    xy = get_xy(args)
    sel, fit_results, criterion_value_name, used_beta = run_model_selection(xy, args)
    best_result = fit_results[sel.best_index]
    best_order = best_result.p - 1
    xcal = make_xcal(xy.x, args)
    best_pred = tklsq.predict_polynomial(xcal, best_result, order=best_order)

    print("# Model selection")
    print(f"criterion = {sel.criterion}")
    if used_beta is not None:
        print(f"alpha = {args.alpha:.6g}, beta = {used_beta:.6g}")
    for i, label in enumerate(sel.labels):
        print(f"{i:3d} {label:10s} score={sel.scores[i]: .6g} weight={sel.weights[i]:.6g}")
    print(f">> Best model: {sel.best_label}  weight={sel.weights[sel.best_index]:.6g}")

    outfile, figfile = get_output_paths(args, f"model_select_{args.criterion}")
    save_model_select_excel(
        outfile,
        xy,
        sel,
        fit_results,
        criterion_value_name,
        used_beta,
        best_result,
        xcal,
        best_pred,
        args,
    )
    print(f"Save Excel: {outfile}")
    plot_model_selection(xy, sel, best_result, xcal, best_pred, figfile, args)


# =====================================================================
# main
# =====================================================================

def main() -> None:
    parser = build_parser()
    args = parser.parse_args()

    try:
        if args.mode == "read":
            mode_read(args)
        elif args.mode == "sim":
            mode_sim(args)
        elif args.mode == "fit":
            mode_fit(args)
        elif args.mode == "model_select":
            mode_model_select(args)
        else:
            raise ValueError(f"Unknown mode: {args.mode}")
    except Exception as exc:
        print(f"ERROR: {exc}")
        traceback.print_exc()
        pause_if_needed(args.pause)
        sys.exit(1)

    pause_if_needed(args.pause)


if __name__ == "__main__":
    main()
