AlignRight = {"ALIGN": "RIGHT"}
import os
import tempfile
import itertools

import pandas as pd
import numpy as np
from reportlab.lib.styles import getSampleStyleSheet
from reportlab.lib import colors
from reportlab.platypus import *

from tia.rlab import *
import tia.util.fmt as fmt


def sample1():
    pdf_path = os.path.join(tempfile.gettempdir(), "pdf_test_4.pdf")
    cols = ["pct", "int", "k", "M", "$", "date"]
    df = pd.DataFrame(np.random.randn(40, len(cols)), columns=cols)
    df["int"] = 10000.0 * df["pct"]
    df["k"] = 50000.0 * df["pct"]
    df["M"] = 5000000.0 * df["pct"]
    df["$"] = 500000.0 * df["pct"]
    df["date"] = pd.date_range("1/1/2010", periods=len(df.index))
    df["id"] = "ID-1"
    # Make this a multi-index frame
    df2 = df.copy()
    df2["id"] = "ID-2"
    df = df.set_index("id", append=True).unstack().reorder_levels([1, 0], axis=1)
    df2 = df2.set_index("id", append=True).unstack().reorder_levels([1, 0], axis=1)
    aggdf = pd.concat([df, df2], axis=1)
    # Start building the pdf
    pdf = PdfBuilder(pdf_path)
    # build the templates to use
    gt = GridTemplate("T1", 100, 100)
    gt.define_frames(
        {
            "HEADER": gt[:10, :],
            "TBL": gt[10:],
        }
    )
    gt.register(pdf)

    # Build the pdf tables to marry with the template
    def make_builder(hdr=1, idx=1, cstyles=None):
        tf = TableFormatter(aggdf, inc_header=hdr, inc_index=idx)
        tf.apply_default_style(index_override={"BACKGROUND": colors.beige})
        tf.header.detect_colspans()
        tf.header.apply_style("ALIGN", "CENTER")
        tf.cells.match_any_labels("pct").apply_number_format(PercentFormatter)
        tf.cells.match_any_labels("k").apply_number_format(ThousandsFormatter)
        tf.cells.match_any_labels("int").apply_number_format(IntFormatter)
        tf.cells.match_any_labels("M").apply_number_format(MillionsFormatter)
        tf.cells.match_any_labels("$").apply_number_format(DollarCentsFormatter)

        def red_weekend(x):
            if x.dayofweek in (5, 6):
                return dict(
                    BACKGROUND=colors.HexColor("#800000"), TEXTCOLOR=colors.white
                )

        tf.cells.match_any_labels("date").apply(
            format=fmt.Y_m_dFormatter, cstyles=red_weekend
        )
        return tf

    # Build PDF
    for hon, ion in list(itertools.product([True, False], repeat=2)):
        offon = lambda v: v and "On" or "Off"
        for cstyle in [None, ConditionalRedBlack]:
            hdr = "Index=%s Header=%s Color=%s" % (
                offon(ion),
                offon(hon),
                offon(cstyle is not None),
            )
            data = {
                "HEADER": Paragraph(hdr, getSampleStyleSheet()["Normal"]),
                "TBL": make_builder(hon, ion, cstyle).build(),
            }
            pdf.build_page("T1", data)
    pdf.save()
    print(pdf_path)


def sample_long_table():
    pdf_path = os.path.join(tempfile.gettempdir(), "pdf_test_long_table.pdf")
    cols = ["pct", "int", "k", "M", "$", "date"]
    df = pd.DataFrame(np.random.randn(200, len(cols)), columns=cols)
    df["int"] = 10000.0 * df["pct"]
    df["k"] = 50000.0 * df["pct"]
    df["M"] = 5000000.0 * df["pct"]
    df["$"] = 500000.0 * df["pct"]
    df["date"] = pd.date_range("1/1/2010", periods=len(df.index))
    df["id"] = "ID-1"
    # Start building the pdf
    pdf = PdfBuilder(pdf_path)
    gt = GridTemplate("T1", 100, 100)
    gt.define_frames(
        {
            "HEADER": gt[:10, :],
            "TBL": gt[10:],
        }
    )
    gt.register(pdf)

    tf = TableFormatter(df)
    tf.apply_default_style()
    pdf.build_page("T1", {"HEADER": pdf.para("HEADER"), "TBL": tf.build()})
    pdf.save()
    print(pdf_path)


def sample_wide_table():
    # Should shrink to fit the page
    pdf_path = os.path.join(tempfile.gettempdir(), "pdf_test_wide_table.pdf")
    cols = ["C%s" % i for i in range(36)]
    df = pd.DataFrame(np.random.randn(20, len(cols)), columns=cols)
    pdf = PdfBuilder(pdf_path)
    gt = GridTemplate("T1", 100, 100)
    gt.define_frames(
        {
            "HEADER": gt[:10, :],
            "TBL": gt[10:],
        }
    )
    gt.register(pdf)
    tf = TableFormatter(df)
    tf.apply_default_style()
    pdf.build_page("T1", {"HEADER": pdf.para("HEADER"), "TBL": tf.build()})
    pdf.save()
    print(pdf_path)


def sample_dyn_col_row_table():
    # Should shrink to fit the page
    pdf_path = os.path.join(tempfile.gettempdir(), "pdf_test_dyn_col_row_table.pdf")
    cols = ["C%s" % i for i in range(4)]
    df = pd.DataFrame(np.random.randn(4, len(cols)), columns=cols)
    pdf = PdfBuilder(pdf_path)
    gt = GridTemplate("T1", 100, 100)
    gt.define_frames(
        {
            "HEADER": gt[:10, :],
            "TBL": gt[10:],
        }
    )
    gt.register(pdf)
    tf = TableFormatter(df)
    tf.apply_default_style()
    tf.set_col_widths(pcts=[0.2, 0.1, 0.2, 0.3, 0.4])
    tf.set_row_heights(pcts=[0.2, 0.1, 0.2, 0.3, 0.4])
    pdf.build_page("T1", {"HEADER": pdf.para("HEADER"), "TBL": tf.build()})
    pdf.save()
    print(pdf_path)


def sample_multi_page():
    # Expand the colwidths to fill page width, but allow rows to be split across pages
    #
    # TODO - make this right
    #
    pdf_path = os.path.join(tempfile.gettempdir(), "pdf_test_multi_page.pdf")
    cols = ["C%s" % i for i in range(4)]
    df = pd.DataFrame(np.random.randn(400, len(cols)), columns=cols)
    pdf = PdfBuilder(pdf_path)
    gt = GridTemplate("T1", 100, 100)
    gt.define_frames(
        {
            "HEADER": gt[:10, :],
            "TBL": gt[10:],
        }
    )
    gt.register(pdf)
    tf = TableFormatter(df)
    tf.apply_default_style()
    tf.set_col_widths(pcts=[0.2, 0.1, 0.2, 0.3, 0.4])
    pdf.build_page("T1", {"HEADER": pdf.para("HEADER"), "TBL": tf.build(shrink=None)})
    pdf.save()
    print(pdf_path)


def sample_heatmap():
    # Expand the colwidths to fill page width, but allow rows to be split across pages
    pdf_path = os.path.join(tempfile.gettempdir(), "pdf_test_heat_map.pdf")
    df = pd.DataFrame(
        np.arange(-100, 100, 10).reshape((5, 4)), columns=["A", "B", "C", "D"]
    )
    pdf = PdfBuilder(pdf_path)
    gt = GridTemplate("T1", 100, 100)
    gt.define_frames(
        {
            "HEADER": gt[:10, :],
            "TBL": gt[10:],
        }
    )
    gt.register(pdf)
    tf = TableFormatter(df)
    tf.apply_default_style()
    tf.cells.heat_map()
    pdf.build_page(
        "T1", {"HEADER": pdf.para("ALL CELLS"), "TBL": tf.build(shrink=None)}
    )
    tf = TableFormatter(df)
    tf.apply_default_style()
    tf.cells.iloc[:, 0].heat_map()
    pdf.build_page(
        "T1", {"HEADER": pdf.para("First Column"), "TBL": tf.build(shrink=None)}
    )
    tf = TableFormatter(df)
    tf.apply_default_style()
    # change font for fun
    tf.cells.iloc[0, :].heat_map(font_cmap="Greys")
    pdf.build_page(
        "T1", {"HEADER": pdf.para("First Row"), "TBL": tf.build(shrink=None)}
    )
    pdf.save()
    print(pdf_path)


def runall():
    sample1()
    sample_long_table()
    sample_wide_table()
    sample_dyn_col_row_table()
    # sample_multi_page()
    sample_heatmap()


if __name__ == "__main__":
    runall()