plotly_utils.py

import sys
from pathlib import Path
local_python_path = str(Path(__file__).parents[1])
if local_python_path not in sys.path:
   sys.path.append(local_python_path)
from utils.utils import load_config, get_logger
logger = get_logger(__name__)
config = load_config(Path(local_python_path) / "config.json") 
import os
import subprocess
import tempfile

from utils.file_handler_utils import read_json
from matplotlib import cm
from matplotlib import pyplot as plt
import numpy as np
import plotly.graph_objects as go
from plotly.subplots import make_subplots

import plotly.express as px
from utils.file_handler_utils import (
    APPEND_SHEET,
    OVERWRITE_FILE,
    OVERWRITE_SHEET,
    read_excel,
    read_json,
    write_excel,
    write_json,
)

IMAGE = 'image'
HTML = 'html'
DEFAULT_FONT_SIZE = config.get('font_size', 36)
PLOTLY_IMAGE_TIMEOUT_SECONDS = 90
PLOTLY_WORKER_JSON_ENV = "PLOTLY_UTILS_WORKER_JSON"
PLOTLY_WORKER_FILENAME_ENV = "PLOTLY_UTILS_WORKER_FILENAME"
PLOTLY_WORKER_OUTPUT_DIR_ENV = "PLOTLY_UTILS_WORKER_OUTPUT_DIR"
PLOTLY_WORKER_WIDTH_FACTOR_ENV = "PLOTLY_UTILS_WORKER_WIDTH_FACTOR"
PLOTLY_WORKER_HEIGHT_FACTOR_ENV = "PLOTLY_UTILS_WORKER_HEIGHT_FACTOR"


def apply_layout(fig, layout_params, font_size):
    """Apply layout parameters including title, base font size, legend and colorbar."""
    if layout_params is None:
        layout_params = {}

    # Center title and apply font size if no explicit title supplied
    if 'title' not in layout_params:
        layout_params['title'] = {
            'x': 0.5,
            'font': {'size': font_size},
            'xanchor': 'center'
        }

    # Ensure the default layout font uses this size and set legend/colorbar defaults
    base_font = layout_params.get('font', {})
    layout_params['font'] = {
        'size': base_font.get('size', font_size),
        **{k: v for k, v in base_font.items() if k != 'size'}
    }

    # Merge or create legend settings
    legend_cfg = layout_params.get('legend', {})
    legend_font = legend_cfg.get('font', {})
    legend_cfg['font'] = {
        'size': legend_font.get('size', font_size),
        **{k: v for k, v in legend_font.items() if k != 'size'}
    }
    layout_params['legend'] = legend_cfg

    # Merge or create coloraxis_colorbar settings
    cab_cfg = layout_params.get('coloraxis_colorbar', {})
    cab_title = cab_cfg.get('title', {})
    cab_title_font = cab_title.get('font', {})
    cab_title['font'] = {
        'size': cab_title_font.get('size', font_size),
        **{k: v for k, v in cab_title_font.items() if k != 'size'}
    }
    cab_cfg['title'] = cab_title

    cab_tickfont = cab_cfg.get('tickfont', {})
    cab_cfg['tickfont'] = {
        'size': cab_tickfont.get('size', font_size),
        **{k: v for k, v in cab_tickfont.items() if k != 'size'}
    }
    layout_params['coloraxis_colorbar'] = cab_cfg

    fig.update_layout(**layout_params)
    return layout_params


def apply_axes(fig, xaxes, yaxes, font_size):
    """Apply x/y axis parameters including tick and title fonts."""
    if xaxes is None:
        xaxes = {}
    if yaxes is None:
        yaxes = {}

    t = dict(tickfont={'size': font_size}, title_font={'size': font_size})
    t.update(xaxes)
    fig.update_xaxes(**t)

    t = dict(tickfont={'size': font_size}, title_font={'size': font_size})
    t.update(yaxes)
    fig.update_yaxes(**t)


def apply_annotations(fig, anotations, font_size):
    """Apply annotations / subplot titles font settings."""
    if anotations is None:
        anotations = {}
    ann_kwargs = dict(font=dict(size=font_size))
    ann_kwargs.update(anotations)
    fig.update_annotations(**ann_kwargs)


def write_output(fig, filename, output_dir, output_type, width, height):
    """Write figure to disk as image or HTML."""
    if output_dir is None:
        output_dir = config['output_dir']

    if output_type == IMAGE:
        fn = output_dir / "{}.png".format(filename)
        func = fig.write_image
#        kw_args = dict(height=height, width=width, engine="orca")
        kw_args = dict(height=height, width=width, engine="kaleido")
        kw_args = dict(height=height, width=width)
    elif output_type == HTML:
        fn = output_dir / "{}.html".format(filename)
        func = fig.write_html
        kw_args = dict(include_plotlyjs=True)
    else:
        raise AssertionError("received illegal output_type {}".format(output_type))

    logger.info("Writing image to {}".format(fn))
    fn.unlink(missing_ok=True)
    fn.parents[0].mkdir(parents=True, exist_ok=True)
    func(fn, **kw_args)

def fix_and_write(fig,
                  filename,
                  traces=None,
                  layout_params=None,
                  output_dir=None,
                  width_factor=1,
                  height_factor=1,
                  xaxes=None,
                  yaxes=None,
                  anotations=None,
                  output_type=IMAGE,
                  font_size=None):
    """
    Fix common layout aspects of a Plotly figure and write it to disk.

    Parameters
    ----------
    fig : plotly.graph_objects.Figure
        The figure to modify and save.
    filename : str
        Base filename (without extension).
    traces : dict, optional
        Passed to fig.update_traces(**traces).
    layout_params : dict, optional
        Passed to fig.update_layout(**layout_params).
    output_dir : Path or str, optional
        Directory to write the file into. Defaults to config['output_dir'].
    width_factor, height_factor : float, optional
        Multipliers for base width/height from config.
    xaxes, yaxes : dict, optional
        Extra parameters for fig.update_xaxes / fig.update_yaxes.
    anotations : dict, optional
        Extra parameters for fig.update_annotations.
    output_type : {'image', 'html'}
        Output format.
    font_size : int, optional
        Font size to apply to all text in the figure (axes, titles,
        annotations, legend, colorbar, and layout default font). If None,
        falls back to DEFAULT_FONT_SIZE from config.
    """
    if font_size is None:
        font_size = DEFAULT_FONT_SIZE
    width = config.get('width', 1920) * width_factor
    height = config.get('height', 1280) * height_factor
    if traces is not None:
        fig.update_traces(**traces)

    # Layout and text handling
    layout_params = apply_layout(fig, layout_params, font_size)
    apply_axes(fig, xaxes, yaxes, font_size)
    apply_annotations(fig, anotations, font_size)

    # Finally, write to disk
    write_output(fig, filename, output_dir, output_type, width, height)

_SWATCH_W = 0.04
_SWATCH_H = 0.03
_LABEL_GAP = 0.008
_ENTRY_GAP = 0.04
_CHAR_W = 0.011

def make_custom_legend(fig, entries, y=-0.22, font_size=28,
                       swatch_w=_SWATCH_W, swatch_h=_SWATCH_H,
                       axis_start=10):
    """Add patterned swatches and labels to *fig* as a hand-drawn legend.

    Parameters
    ----------
    fig : plotly.graph_objects.Figure
    entries : list of (label, color, pattern_shape) tuples.
        pattern_shape can be "" or None for a solid swatch.
    y : float
        Vertical centre of the legend row in paper coordinates.
    font_size : int
    swatch_w, swatch_h : float
        Width and height of each swatch in paper coordinates.
    axis_start : int
        First axis index to use for hidden swatch axes (e.g. 10 → xaxis10).
        Must not collide with axes already on the figure.
    """
    entry_widths = [swatch_w + _LABEL_GAP + len(label) * _CHAR_W
                    for label, *_ in entries]
    total_w = sum(entry_widths) + (len(entries) - 1) * _ENTRY_GAP
    cx = 0.5 - total_w / 2

    annotations = []
    for i, (entry, ew) in enumerate(zip(entries, entry_widths)):
        label, color = entry[0], entry[1]
        pattern = entry[2] if len(entry) > 2 else ""

        ax_idx = axis_start + i
        x_key = f"xaxis{ax_idx}"
        y_key = f"yaxis{ax_idx}"
        xref = f"x{ax_idx}"
        yref = f"y{ax_idx}"

        fig.update_layout(**{
            x_key: dict(
                domain=[cx, cx + swatch_w],
                visible=False, fixedrange=True,
            ),
            y_key: dict(
                domain=[y, y + swatch_h],
                visible=False, fixedrange=True,
                anchor=xref,
            ),
        })

        fig.add_trace(go.Bar(
            x=["s"], y=[1],
            marker=dict(
                color=color,
                pattern_shape=pattern or "",
                pattern_solidity=0.5,
                line=dict(color="black", width=1),
            ),
            showlegend=False,
            xaxis=xref,
            yaxis=yref,
        ))

        annotations.append(dict(
            xref="paper", yref="paper",
            x=cx + swatch_w + _LABEL_GAP,
            y=y + swatch_h / 2,
            text=label,
            showarrow=False,
            xanchor="left",
            yanchor="middle",
            font=dict(size=font_size),
        ))
        cx += ew + _ENTRY_GAP

    existing = list(fig.layout.annotations or [])
    fig.update_layout(annotations=existing + annotations)


_LINE_SWATCH_W = 0.06

def make_custom_line_legend(fig, entries, y=0.02, font_size=24):
    """Add line swatches and labels to *fig* as a hand-drawn legend.

    Parameters
    ----------
    fig : plotly.graph_objects.Figure
    entries : list of (label, color, dash) tuples.
        dash is a plotly dash string: "solid", "dash", "dot", "dashdot", etc.
    y : float
        Vertical centre of the legend row in paper coordinates (must be in [0, 1]).
    font_size : int
    """
    entry_widths = [_LINE_SWATCH_W + _LABEL_GAP + len(label) * _CHAR_W
                    for label, *_ in entries]
    total_w = sum(entry_widths) + (len(entries) - 1) * _ENTRY_GAP
    cx = 0.5 - total_w / 2

    shapes = list(fig.layout.shapes or [])
    annotations = []
    for (entry, ew) in zip(entries, entry_widths):
        label, color = entry[0], entry[1]
        dash = entry[2] if len(entry) > 2 else "solid"

        shapes.append(dict(
            type="line",
            xref="paper", yref="paper",
            x0=cx, x1=cx + _LINE_SWATCH_W,
            y0=y, y1=y,
            line=dict(color=color, width=3, dash=dash),
        ))
        annotations.append(dict(
            xref="paper", yref="paper",
            x=cx + _LINE_SWATCH_W + _LABEL_GAP,
            y=y,
            text=label,
            showarrow=False,
            xanchor="left",
            yanchor="middle",
            font=dict(size=font_size),
        ))
        cx += ew + _ENTRY_GAP

    existing_annots = list(fig.layout.annotations or [])
    fig.update_layout(shapes=shapes, annotations=existing_annots + annotations)


def combine_figures(figs_list):


    # Extract titles from fig1 and fig2
    titles = [fig.layout.title.text if fig.layout.title.text else f"Figure {i+1}" for i, fig in enumerate(figs_list)]

    # Create a subplot figure with the extracted titles
    combined_fig = make_subplots(
        cols=1, rows=len(figs_list),  # 1 row and 2 columns
        subplot_titles=(titles)  # Use extracted titles
    )

    # Add traces from fig1 to the first subplot
    for i, fig in enumerate(figs_list):
        for trace in fig.data:
            combined_fig.add_trace(trace, col=1, row=i+1)
    
    # Update the layout
    combined_fig.update_layout(
        showlegend=False  # Set to True if you want a combined legend
    )

    # Show the combined figure
    return combined_fig


tab10 = [[31, 119, 180],
 [31, 119, 180],
 [255, 127, 14],
 [255, 127, 14],
 [44, 160, 44],
 [44, 160, 44],
 [214, 39, 40],
 [214, 39, 40],
 [148, 103, 189],
 [148, 103, 189],
 [140, 86, 75],
 [140, 86, 75],
 [227, 119, 194],
 [227, 119, 194],
 [127, 127, 127],
 [127, 127, 127],
 [188, 189, 34],
 [188, 189, 34],
 [23, 190, 207],
 [23, 190, 207],
 [23, 190, 207]]
    
tab20 = [[31, 119, 180],
 [174, 199, 232],
 [255, 127, 14],
 [255, 187, 120],
 [44, 160, 44],
 [152, 223, 138],
 [214, 39, 40],
 [255, 152, 150],
 [148, 103, 189],
 [197, 176, 213],
 [140, 86, 75],
 [196, 156, 148],
 [227, 119, 194],
 [247, 182, 210],
 [127, 127, 127],
 [199, 199, 199],
 [188, 189, 34],
 [219, 219, 141],
 [23, 190, 207],
 [158, 218, 229],
 [158, 218, 229]]

#def get_colors(N, cmap_name=None, with_faded=False):
    # if cmap_name is not None:
    #     cmap = cm.get_cmap(plt.get_cmap(cmap_name))
    #     colors = cmap(np.linspace(0, 1, N))
    #     colors = [[int(x) for x in y] for y in (colors[:, :3]*255).tolist()]
    # elif N <= 9:
    #     colors = [[int(y) for y in x[4:-1].split(",")] for x in plotly.colors.qualitative.Set1]
    #     # [[int(x[1:][i:i+2], 16) for i in [0,2,4]] for x in plotly.colors.qualitative.Set1]
    # elif N<= 24:
    #     colors = [[int(x[1:][i:i+2], 16) for i in [0,2,4]] for x in plotly.colors.qualitative.Light24]
    # elif N < 100:
    #     colors = distinctipy.get_colors(N)
    # else:
    #     raise AssertionError("get_colors got too large N")
    # if with_faded:
    #     return ['rgb({})'.format(",".join([str(x) for x in y])) for y in colors], \
    #         ['rgba({},0.2)'.format(",".join([str(x) for x in y])) for y in colors]
    # return ['rgb({})'.format(",".join([str(x) for x in y])) for y in colors]


def is_valid_plotly_png(output_dir, filename, previous_mtime_ns=None):
    output_path = Path(output_dir) / f"{filename}.png"
    if not output_path.exists() or output_path.stat().st_size <= 0:
        return False
    if previous_mtime_ns is None:
        return True
    return output_path.stat().st_mtime_ns != previous_mtime_ns


def get_plotly_worker_environment(
    figure_json_path,
    filename,
    output_dir,
    width_factor,
    height_factor,
):
    worker_environment = os.environ.copy()
    worker_environment.update({
        PLOTLY_WORKER_JSON_ENV: str(figure_json_path),
        PLOTLY_WORKER_FILENAME_ENV: str(filename),
        PLOTLY_WORKER_OUTPUT_DIR_ENV: str(output_dir),
        PLOTLY_WORKER_WIDTH_FACTOR_ENV: str(width_factor),
        PLOTLY_WORKER_HEIGHT_FACTOR_ENV: str(height_factor),
    })
    return worker_environment


def kill_process_tree(process):
    if os.name == "nt":
        subprocess.run(
            ["taskkill", "/PID", str(process.pid), "/T", "/F"],
            capture_output=True,
            text=True,
            check=False,
        )
    else:
        process.kill()
    process.wait()


def run_plotly_image_worker():
    figure_json_path = Path(os.environ[PLOTLY_WORKER_JSON_ENV])
    figure = go.Figure(read_json(figure_json_path))
    fix_and_write(
        figure,
        os.environ[PLOTLY_WORKER_FILENAME_ENV],
        output_dir=Path(os.environ[PLOTLY_WORKER_OUTPUT_DIR_ENV]),
        width_factor=float(os.environ[PLOTLY_WORKER_WIDTH_FACTOR_ENV]),
        height_factor=float(os.environ[PLOTLY_WORKER_HEIGHT_FACTOR_ENV]),
    )


def write_plotly_image_with_timeout(
    fig,
    filename,
    output_dir,
    width_factor,
    height_factor,
    timeout_seconds=PLOTLY_IMAGE_TIMEOUT_SECONDS,
):
    output_path = Path(output_dir) / f"{filename}.png"
    previous_mtime_ns = output_path.stat().st_mtime_ns if output_path.exists() else None
    with tempfile.TemporaryDirectory() as temp_dir:
        figure_json_path = Path(temp_dir) / "figure.json"
        figure_json_path.write_text(fig.to_json(), encoding="utf-8")
        worker_environment = get_plotly_worker_environment(
            figure_json_path,
            filename,
            output_dir,
            width_factor,
            height_factor,
        )
        process = subprocess.Popen(
            [sys.executable, "-m", "utils.plotly_utils"],
            cwd=local_python_path,
            env=worker_environment,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
        )
        try:
            stdout, stderr = process.communicate(timeout=timeout_seconds)
        except subprocess.TimeoutExpired:
            logger.error(
                f"Kaleido timed out after {timeout_seconds} seconds while writing "
                f"{filename}.png; terminating its process tree"
            )
            kill_process_tree(process)
            if is_valid_plotly_png(output_dir, filename, previous_mtime_ns):
                logger.warning(
                    f"Keeping {filename}.png because Kaleido completed the image before hanging"
                )
                return True
            return False

        if process.returncode == 0:
            return True
        logger.error(
            f"Plotly image worker failed for {filename}.png with exit code "
            f"{process.returncode}: {stderr.strip() or stdout.strip()}"
        )
        return is_valid_plotly_png(output_dir, filename, previous_mtime_ns)


if __name__ == "__main__" and PLOTLY_WORKER_JSON_ENV in os.environ:
    run_plotly_image_worker()