#! /usr/bin/env python

if __name__ == "__main__":
    import datetime
    import pathlib

    import numpy as np
    import pandas as pd

    from skrub import TableReport, datasets

    reports_dir = pathlib.Path(__file__).resolve().parent / "_reports"
    reports_dir.mkdir(exist_ok=True)

    df = datasets.fetch_employee_salaries().X
    html = TableReport(df).html()
    (reports_dir / "employee_salaries.html").write_text(html, encoding="UTF-8")

    df = pd.DataFrame(
        {
            "c 0": list(range(7)),
            "c 1": [f"v 0{i}" for i in range(7)],
            "c 2": [i + 10 for i in range(7)],
            "c 3": [f"v {i + 10}" for i in range(7)],
        }
    )
    html = TableReport(df).html()
    (reports_dir / "mini.html").write_text(html, encoding="UTF-8")

    # Example from https://pandas.pydata.org/docs/user_guide/reshaping.html#pivot-table

    df = pd.DataFrame(
        {
            "A": ["one", "one", "two", "three"] * 6,
            "B": ["A", "B", "C"] * 8,
            "C": ["foo", "foo", "foo", "bar", "bar", "bar"] * 4,
            "D": np.random.randn(24),
            "E": np.random.randn(24),
            "F": [datetime.datetime(2013, i, 1) for i in range(1, 13)] + [
                datetime.datetime(2013, i, 15) for i in range(1, 13)
            ],
        }
    )

    df = pd.pivot_table(
        df,
        values="E",
        index=["B", "C"],
        columns=["A"],
        aggfunc=["sum", "mean"],
    )

    html = TableReport(df).html()
    (reports_dir / "multi_index.html").write_text(html, encoding="UTF-8")
