import io
import math

import pandas as pd
import numpy as np

import matplotlib.pyplot as plt
from matplotlib.figure import Figure
from matplotlib.patches import Circle
import matplotlib as mpl

from pypdf import PdfWriter

import panel as pn


# MARK: Functions
class PANES:
    TEMPLATE = "Template"
    RAW_DATA = "Experiment raw data"
    READY_DATA = "Experiment data"
    PLATE_VIEWER = "Plate viewer"
    PAGE_VIEWER = "Page viewer"


VERTICAL_PADDING = 4
CFO_SHOW_COORDINATES = "Show coordinates"
CFO_SHOW_BOLD_FONT = "Bold font"


def expand_position(row, col, discs_per_leaf):
    min_row = (row - 1) * discs_per_leaf + 1
    max_row = min_row + discs_per_leaf
    return (min_row, max_row), col


def is_valid_forbibden_position(fp):
    return fp is not None and fp[0] > 0 and fp[1] > 0


def is_forbibden_position(fp, r, c):
    return is_valid_forbibden_position(fp) and r == fp[0] and c == fp[1]


def get_next_position(
    tray, row, col, row_count: int, col_count: int, forbidden_position: tuple | None
):
    if col == 0 and row == 0:
        col = 1
        row = 1
        if is_forbibden_position(forbidden_position, row, col):
            return get_next_position(
                tray,
                col,
                row,
                row_count=row_count,
                col_count=col_count,
                forbidden_position=forbidden_position,
            )
        return tray, row, col
    if col < col_count:
        col += 1
    else:
        col = 1
        if row < row_count:
            row += 1
        else:
            row = 1
            tray += 1
    if is_forbibden_position(forbidden_position, row, col):
        return get_next_position(
            tray,
            col,
            row,
            row_count=row_count,
            col_count=col_count,
            forbidden_position=forbidden_position,
        )
    return tray, row, col


def draw_tray(
    forbidden_position: tuple | None,
    row_count: int,
    col_count: int,
    discs_per_leaf: int,
    tray_size: tuple,
    row: int = 0,
    col: int = 0,
    text_data: np.ndarray | None = None,
    color_data: np.ndarray | None = None,
    fig: Figure | None = None,
    fontsize=14,
    plaque_index: str | None = None,
    print_coordinates: bool = False,
    print_index: bool = False,
    coord_bf: bool = False,
    coord_font_size: int = 10,
    leaf_size: int | None = None,
):
    t = np.zeros((row_count * discs_per_leaf, col_count))
    if is_valid_forbibden_position(forbidden_position):
        t[
            forbidden_position[0] - 1 : forbidden_position[0] - 1 + discs_per_leaf,
            forbidden_position[1] - 1,
        ] = 1
    if col > 0 and row > 0:
        (min_row, max_row), col = expand_position(row, col, discs_per_leaf)
        t[min_row - 1 : max_row - 1, col - 1] = 2

    had_fig = fig is not None
    if fig is None:
        fig = Figure(figsize=tuple(i / 2.54 for i in tray_size))

    axii = fig.subplots(nrows=t.shape[0], ncols=t.shape[1])
    left, width = 0, 1
    bottom, height = 0, 1
    right = left + width
    top = bottom + height

    if color_data is not None and color_data.size > 0:
        cmap = mpl.colormaps["cool"]
        colors = cmap(np.linspace(0, 1, np.nanmax(color_data).astype(int) + 1))

    def is_valid(a, i, j):
        return a is not None and 0 <= i < a.shape[0] and 0 <= j < a.shape[1]

    leaf_index = 0

    for i, line in enumerate(axii):
        for j, ax in enumerate(line):
            if t[i, j] != 1:
                leaf_index += 1
            # Draw background
            if i == 0 and j == 0 and plaque_index is not None and t[i, j] == 1:
                ax.set_facecolor("xkcd:light grey")
            else:
                if t[i, j] == 1:
                    ax.set_facecolor("xkcd:black")
                elif color_data is None:
                    match t[i, j]:
                        case 0:
                            ax.set_facecolor("xkcd:white")
                        case 2:
                            ax.set_facecolor("xkcd:salmon")
                elif is_valid(color_data, i, j) and np.isnan(color_data[i, j]) != True:
                    ax.set_facecolor(colors[int(color_data[i, j].astype(int))])
            # Draw leafs
            if leaf_size is not None and leaf_size > 0 and t[i, j] != 1:
                ax.add_patch(
                    Circle(
                        xy=(0.5, 0.5),
                        radius=leaf_size / 2,
                        edgecolor="green",
                        facecolor="lime",
                        linewidth=1,
                    )
                )
            # Write text
            if print_coordinates is True:
                ax.text(
                    left + 0.02,
                    top - 0.02,
                    f"{chr(65 + i)}{j+1}",
                    dict(size=coord_font_size),
                    horizontalalignment="left",
                    verticalalignment="top",
                    fontweight="bold" if coord_bf is True else "normal",
                )
                ax.text(
                    right - 0.03,
                    bottom + 0.05,
                    f"{chr(65 + i)}{j+1}",
                    dict(size=coord_font_size),
                    horizontalalignment="right",
                    verticalalignment="bottom",
                    rotation=180,
                    transform_rotates_text=True,
                    fontweight="bold" if coord_bf is True else "normal",
                )
            if print_index is True:
                ax.text(
                    0.5,
                    0.5,
                    str(leaf_index),
                    dict(size=fontsize),
                    horizontalalignment="center",
                    verticalalignment="center",
                    transform=ax.transAxes,
                )
            if i == 0 and j == 0 and plaque_index is not None and t[i, j] == 1:
                ax.set_facecolor("xkcd:light grey")
                ax.text(
                    0.5,
                    0.5,
                    plaque_index,
                    dict(size=fontsize),
                    horizontalalignment="center",
                    verticalalignment="center",
                )
            if (
                text_data is not None
                and is_valid(text_data, i, j)
                and isinstance(text_data[i, j], str) is True
            ):
                ax.text(
                    0.5,
                    0.5,
                    text_data[i, j],
                    dict(size=fontsize),
                    horizontalalignment="center",
                    verticalalignment="center",
                    transform=ax.transAxes,
                )
            ax.set_xticklabels([])
            ax.set_yticklabels([])
            ax.set_xticks([])
            ax.set_yticks([])

    if had_fig is False:
        fig.tight_layout(pad=0.2)
    return fig


def draw_page(
    data: pd.DataFrame,
    figs_per_col: int,
    figs_per_row: int,
    forbibden_position: tuple,
    row_count: int,
    col_count: int,
    discs_per_leaf: int,
    orientation: str,
    font_size: int,
    page_index: int,
):
    main_fig = Figure(
        figsize=(21, 29.7) if orientation == "portrait" else (29.7, 21),
        layout="constrained",
    )
    subfigs = main_fig.subfigures(figs_per_row, figs_per_col, wspace=0.07)
    start_tray = (figs_per_col * figs_per_row * (page_index - 1)) + 1

    for i in range(figs_per_col * figs_per_row):
        if i >= data.plaque.max():
            break
        current_fig = subfigs.flatten()[i]
        draw_tray(
            forbidden_position=(
                None
                if forbibden_position is None
                else (forbibden_position[0], forbibden_position[1])
            ),
            row_count=row_count,
            col_count=col_count,
            discs_per_leaf=discs_per_leaf,
            tray_size=(row_count * discs_per_leaf, col_count),
            text_data=data[data.plaque == i + start_tray]
            .pivot(index="ligne", columns="colonne", values="Num_Plvt")
            .values,
            color_data=data[data.plaque == i + start_tray]
            .pivot(index="ligne", columns="colonne", values="group")
            .values,
            fontsize=font_size,
            fig=current_fig,
            plaque_index=f"P{i+start_tray}",
        )
        if is_valid_forbibden_position(forbibden_position) is False:
            current_fig.suptitle(f"P{i+start_tray}", fontsize=font_size)

    if "experience" in data:
        main_fig.suptitle(
            f"{data.iloc[0].experience}, page {page_index}", fontsize=font_size * 1.5
        )
    return main_fig


def prepare_experiment(
    raw_data: pd.DataFrame,
    row_count: int,
    col_count: int,
    discs_per_leaf: int,
    forbidden_position: tuple,
    control_blocks: int = 3,
):
    ctr_data = raw_data[raw_data.TYPE == "THD"].reset_index(drop=True)
    gen_data = raw_data[raw_data.TYPE == "G"].reset_index(drop=True)

    match control_blocks:
        case 0:
            all_data = gen_data
        case 1:
            all_data = pd.concat(
                [
                    ctr_data.sample(n=len(ctr_data), random_state=0).assign(
                        group=lambda x: 1
                    ),
                    gen_data,
                ]
            )
        case _:
            len_chunks = math.ceil(len(gen_data) / (control_blocks - 1))
            list_chunks = [
                gen_data[i : i + len_chunks]
                for i in range(0, len(gen_data), len_chunks)
            ]
            all_data = ctr_data.sample(n=len(ctr_data), random_state=0).assign(
                group=lambda x: 1
            )
            for i, chunk in enumerate(list_chunks):
                all_data = pd.concat(
                    [
                        all_data,
                        chunk,
                        ctr_data.sample(n=len(ctr_data), random_state=i).assign(
                            group=lambda x: i + 2
                        ),
                    ]
                )

    tray, row, col = 1, 0, 0
    out = []
    for _, r in all_data.reset_index(drop=True).iterrows():
        tray, row, col = get_next_position(
            tray=tray,
            row=row,
            col=col,
            row_count=row_count,
            col_count=col_count,
            forbidden_position=forbidden_position,
        )
        (min_row, max_row), col = expand_position(
            row, col, discs_per_leaf=discs_per_leaf
        )
        for i in range(min_row, max_row):
            out.append(
                dict(r)
                | {
                    "plaque": tray,
                    "ligne": i,
                    "colonne": col,
                    "alpha_line": chr(i + 64),
                }
            )

    return pd.DataFrame(out)


# MARK: Template
ii_row_count = pn.widgets.IntInput(
    name="Row count", start=1, end=100, value=3, sizing_mode="stretch_width"
)
ii_col_count = pn.widgets.IntInput(
    name="Column count", start=1, end=100, value=9, sizing_mode="stretch_width"
)
ii_discs_per_leaf = pn.widgets.IntInput(
    name="Discs per leaf", start=1, end=100, value=3, sizing_mode="stretch_width"
)

ii_forbibden_row = pn.widgets.IntInput(
    name="Row", start=0, end=100, value=1, sizing_mode="stretch_width"
)
ii_forbibden_col = pn.widgets.IntInput(
    name="Column", start=0, end=100, value=1, sizing_mode="stretch_width"
)

ii_template_width = pn.widgets.IntInput(
    name="Template width (in mm)",
    start=10,
    end=1000,
    value=215,
    sizing_mode="stretch_width",
)
ii_template_height = pn.widgets.IntInput(
    name="Template height (in mm)",
    start=10,
    end=1000,
    value=215,
    sizing_mode="stretch_width",
)

cbg_coordinates = pn.widgets.CheckBoxGroup(
    name="Coordiantes options",
    options=[CFO_SHOW_COORDINATES, CFO_SHOW_BOLD_FONT],
    value=[CFO_SHOW_COORDINATES],
    inline=False,
)
ii_coordinate_font_size = pn.widgets.IntInput(
    name="Coordinates font size",
    start=1,
    end=100,
    value=11,
    sizing_mode="stretch_width",
)
ii_leaf_size = pn.widgets.FloatInput(
    name="Leaf size",
    start=0,
    end=100,
    value=0,
    step=0.1,
    sizing_mode="stretch_width",
)

dwn_template = pn.widgets.FileDownload(
    file="",
    filename="",
    label="Download template",
    name="",
    # sizing_mode="stretch_width",
    icon="file-download",
    button_type="primary",
    align="end",
)

cd_template = pn.layout.Card(
    pn.Column(
        pn.Row(ii_row_count, ii_col_count),
        ii_discs_per_leaf,
        pn.layout.Divider(),
        pn.pane.Str("Forbidden position", margin=(0, 10, -10, 2), align="center"),
        pn.Row(ii_forbibden_row, ii_forbibden_col),
        pn.layout.Divider(),
        pn.Row(ii_template_width, ii_template_height),
        pn.layout.Divider(),
        pn.Row(cbg_coordinates, ii_coordinate_font_size),
        ii_leaf_size,
        pn.layout.Divider(),
        dwn_template,
    ),
    title="Template",
    collapsed=True,
    margin=(VERTICAL_PADDING, 0, VERTICAL_PADDING, 0),
)

# MARK: Load experiment
csv_separator = pn.widgets.Select(
    name="CSV separator", options=[";", ","], sizing_mode="stretch_width"
)
ii_control_blocks = pn.widgets.IntInput(
    name="Control blocks", start=0, end=100, value=3, sizing_mode="stretch_width"
)
fi_raw_exp = pn.widgets.FileInput(accept=".csv,.json", sizing_mode="stretch_width")


cd_load_exp = pn.layout.Card(
    pn.Column(pn.Row(csv_separator, ii_control_blocks), fi_raw_exp),
    title="Load experiment",
    collapsed=False,
    margin=(VERTICAL_PADDING, 0, VERTICAL_PADDING, 0),
)

# MARK: View plates
ii_plate_number = pn.widgets.IntInput(
    name="Plate number",
    start=1,
    end=1,
    value=1,
    sizing_mode="stretch_width",
)
ii_font_size = pn.widgets.IntInput(
    name="Font size",
    start=1,
    end=100,
    value=14,
    sizing_mode="stretch_width",
)

fd_genrate_exp_plates = pn.widgets.FileDownload(
    name="", button_type="primary", filename="exp_plates.pdf", align="end"
)

cd_view_plates = pn.layout.Card(
    pn.Column(pn.Row(ii_font_size, ii_plate_number), fd_genrate_exp_plates),
    title="View plates",
    collapsed=False,
    margin=(VERTICAL_PADDING, 0, VERTICAL_PADDING, 0),
)

# MARK: Page viewer
ii_figs_per_col = pn.widgets.IntInput(
    name="Plates per column", start=1, end=100, value=3, sizing_mode="stretch_width"
)
ii_figs_per_row = pn.widgets.IntInput(
    name="Plates per row", start=1, end=100, value=3, sizing_mode="stretch_width"
)
sel_orientation = pn.widgets.Select(
    name="Page orientation",
    options=["Landscape", "Portrait"],
    sizing_mode="stretch_width",
)
ii_page_font_size = pn.widgets.IntInput(
    name="Page font size",
    start=1,
    end=100,
    value=14,
    sizing_mode="stretch_width",
)
isl_page_index = pn.widgets.IntInput(
    name="Page index", start=1, end=100, value=1, sizing_mode="stretch_width"
)
bt_generate_page = pn.widgets.Button(
    name="View page", button_type="primary", align="center"
)

fd_genrate_exp_helper = pn.widgets.FileDownload(
    name="",
    button_type="primary",
    filename="exp_helper.pdf",
    align="center",
    sizing_mode="stretch_width",
)

cd_view_pages = pn.layout.Card(
    pn.Column(
        pn.Row(ii_figs_per_row, ii_figs_per_col),
        isl_page_index,
        pn.Row(sel_orientation, ii_page_font_size),
        pn.Row(bt_generate_page, fd_genrate_exp_helper),
    ),
    title="View pages",
    collapsed=False,
    margin=(VERTICAL_PADDING, 0, VERTICAL_PADDING, 0),
)


# MARK: Main form
rbg_active = pn.widgets.RadioButtonGroup(
    name="",
    options=[
        PANES.TEMPLATE,
        PANES.RAW_DATA,
        PANES.READY_DATA,
        PANES.PLATE_VIEWER,
        PANES.PAGE_VIEWER,
    ],
    value=PANES.TEMPLATE,
    button_type="primary",
    button_style="outline",
    sizing_mode="stretch_width",
)
plt_template = pn.pane.Matplotlib(align="center", sizing_mode="stretch_height")
tb_raw_experiment = pn.widgets.Tabulator(
    value=pd.DataFrame(), pagination="local", page_size=20, sizing_mode="stretch_width"
)
tb_experiment = pn.widgets.Tabulator(
    value=pd.DataFrame(), pagination="local", page_size=20, sizing_mode="stretch_width"
)
plt_plate = pn.pane.Matplotlib(align="center", sizing_mode="stretch_height")
plt_page = pn.pane.Matplotlib(align="center", sizing_mode="stretch_height")

ph_main = pn.pane.Placeholder(object=plt_template)


# MARK: Callbacks
def update_experiment():
    tb_experiment.value = prepare_experiment(
        raw_data=tb_raw_experiment.value,
        row_count=ii_row_count.value,
        col_count=ii_col_count.value,
        discs_per_leaf=ii_discs_per_leaf.value,
        forbidden_position=(ii_forbibden_row.value, ii_forbibden_col.value),
        control_blocks=ii_control_blocks.value,
    )
    ii_plate_number.end = int(tb_experiment.value.plaque.max())
    ii_plate_number.value = 1
    plt_plate.object = update_tray(data=tb_experiment.value)


def update_tray(data: pd.DataFrame | None = None, plate_number: int | None = None):
    if data is not None:
        p_number = (
            data.plaque == ii_plate_number.value
            if plate_number is None
            else plate_number
        )
        text_data = (
            data[data.plaque == p_number]
            .pivot(index="ligne", columns="colonne", values="Num_Plvt")
            .values
        )
        color_data = (
            data[data.plaque == p_number]
            .pivot(index="ligne", columns="colonne", values="group")
            .values
        )
    else:
        text_data = None
        color_data = None
    return draw_tray(
        forbidden_position=(ii_forbibden_row.value, ii_forbibden_col.value),
        row_count=ii_row_count.value,
        col_count=ii_col_count.value,
        discs_per_leaf=ii_discs_per_leaf.value,
        tray_size=(ii_template_width.value / 10, ii_template_height.value / 10),
        text_data=text_data,
        color_data=color_data,
        fontsize=ii_font_size.value,
        coord_font_size=ii_coordinate_font_size.value,
        print_coordinates=CFO_SHOW_COORDINATES in cbg_coordinates.value,
        coord_bf=CFO_SHOW_BOLD_FONT in cbg_coordinates.value,
        leaf_size=ii_leaf_size.value,
    )


def on_exp_plates_requested(data):
    merger = PdfWriter()
    for i in sorted(data.plaque.unique()):
        fig = update_tray(data, plate_number=int(i))
        s_buf = io.BytesIO()
        fig.savefig(s_buf, bbox_inches="tight", pad_inches=0, format="pdf")
        s_buf.seek(0)
        merger.append(s_buf)
    out_buf = io.BytesIO()
    merger.write(out_buf)
    out_buf.seek(0)
    return out_buf


fd_genrate_exp_plates.callback = pn.bind(on_exp_plates_requested, tb_experiment)


def on_exp_helper_requested(data):
    merger = PdfWriter()
    figs_per_col = ii_figs_per_col.value
    figs_per_row = ii_figs_per_row.value
    page_count = data.plaque.max() // (figs_per_row * figs_per_col) + 1
    for i in range(page_count):
        fig = draw_page(
            data=data,
            figs_per_col=figs_per_col,
            figs_per_row=figs_per_row,
            forbibden_position=(ii_forbibden_row.value, ii_forbibden_col.value),
            row_count=ii_row_count.value,
            col_count=ii_col_count.value,
            discs_per_leaf=ii_discs_per_leaf.value,
            orientation=sel_orientation.value.lower(),
            font_size=ii_page_font_size.value,
            page_index=i + 1,
        )
        s_buf = io.BytesIO()
        fig.savefig(s_buf, bbox_inches="tight", pad_inches=0, format="pdf")
        s_buf.seek(0)
        merger.append(s_buf)
    out_buf = io.BytesIO()
    merger.write(out_buf)
    out_buf.seek(0)
    return out_buf


fd_genrate_exp_helper.callback = pn.bind(on_exp_helper_requested, tb_experiment)


@pn.depends(ii_plate_number.param.value, watch=True)
def on_plate_number_changed(plate_number):
    plt_plate.object = update_tray(data=tb_experiment.value)


@pn.depends(ii_font_size.param.value, watch=True)
def on_plate_number_changed(font_size):
    plt_plate.object = update_tray(data=tb_experiment.value)


@pn.depends(
    *[
        p.param.value
        for p in [
            ii_row_count,
            ii_col_count,
            ii_discs_per_leaf,
            ii_forbibden_row,
            ii_forbibden_col,
            ii_template_width,
            ii_template_height,
            cbg_coordinates,
            ii_coordinate_font_size,
            ii_leaf_size,
            ii_control_blocks,
        ]
    ],
    watch=True,
)
def on_template_changed(
    row_count,
    col_count,
    discs_per_leaf,
    forbibden_row,
    forbibden_column,
    template_width,
    template_height,
    show_coordinates,
    coordinate_font_size,
    leaf_size,
    control_blocks,
):
    fig = update_tray()
    s_buf = io.BytesIO()
    fig.savefig(s_buf, bbox_inches="tight", pad_inches=0, format="pdf")
    s_buf.seek(0)
    dwn_template.filename = (
        f"[R{row_count}][C{col_count}][DPL{discs_per_leaf}]_template.pdf"
    )
    dwn_template.file = s_buf

    plt_template.object = fig

    if len(tb_raw_experiment.value) > 0:
        update_experiment()


@pn.depends(fi_raw_exp.param.value, watch=True)
def on_file_loaded(value):
    if not value or not isinstance(value, bytes):
        return pd.DataFrame()

    string_io = io.StringIO(value.decode("utf8"))

    ret = pd.read_csv(string_io, sep=csv_separator.value)
    # ret = pd.concat([ret.sample(n=len(ret)) for _ in range(60)]).reset_index(drop=True)
    tb_raw_experiment.value = ret
    update_experiment()
    fd_genrate_exp_plates.file = "exp_plates.pdf"
    fd_genrate_exp_helper.file = "exp_helper.pdf"
    rbg_active.value = PANES.READY_DATA


@pn.depends(rbg_active.param.value, watch=True)
def on_active_pane_changed(rbg_value):
    match rbg_value:
        case PANES.TEMPLATE:
            ph_main.object = plt_template
            cd_template.collapsed = False
        case PANES.RAW_DATA:
            ph_main.object = tb_raw_experiment
        case PANES.READY_DATA:
            ph_main.object = tb_experiment
        case PANES.PLATE_VIEWER:
            ph_main.object = plt_plate
            cd_view_plates.collapsed = False
        case PANES.PAGE_VIEWER:
            ph_main.object = plt_page
            cd_view_pages.collapsed = False


def on_page_requested(event):
    if not event:
        return
    rbg_active.value = PANES.PAGE_VIEWER
    plt_page.object = draw_page(
        data=tb_experiment.value,
        figs_per_col=ii_figs_per_col.value,
        figs_per_row=ii_figs_per_row.value,
        forbibden_position=(ii_forbibden_row.value, ii_forbibden_col.value),
        row_count=ii_row_count.value,
        col_count=ii_col_count.value,
        discs_per_leaf=ii_discs_per_leaf.value,
        orientation=sel_orientation.value.lower(),
        font_size=ii_page_font_size.value,
        page_index=isl_page_index.value,
    )


pn.bind(on_page_requested, bt_generate_page, watch=True)

on_template_changed(
    *[
        p.value
        for p in [
            ii_row_count,
            ii_col_count,
            ii_discs_per_leaf,
            ii_forbibden_row,
            ii_forbibden_col,
            ii_template_width,
            ii_template_height,
            cbg_coordinates,
            ii_coordinate_font_size,
            ii_leaf_size,
            ii_control_blocks,
        ]
    ]
)

# MARK: Layout
pn.extension("tabulator", "ipywidgets")
template = pn.template.BootstrapTemplate(title="Vitis Experiment Builder")
template.sidebar.append(
    pn.Column(
        cd_template,
        cd_load_exp,
        cd_view_plates,
        cd_view_pages,
    )
)
template.main.append(pn.Column(rbg_active, ph_main))
template.servable()