visualization.py

'''
Functions to visualize plots of the data:
- Helper functions to create the results folder and to save and show data
- Entry point function for easier access from main.py
- Separate functions to plot the results and to visualize the raw data from training, test and ideal functions 
'''

import os
from datetime import datetime
import tempfile
from bokeh.plotting import figure, output_file, show, save
from bokeh.palettes import Category10
from bokeh.layouts import column
from config import PLOT_RESULTS, PLOT_TRAINING_FIT, PLOT_TRAINING, PLOT_IDEAL, PLOT_TESTDATA, SAVE_PLOTS

def _create_output_folder():
    '''
    Helper: Creates a timestamped results folder in the working directory that can be used to save plots in
    '''
    timestamp = datetime.now().strftime("%Y-%m-%d-%H-%M-%S")
    folder_name = f"{timestamp}_Plots"
    os.makedirs(folder_name, exist_ok=True)
    return folder_name


def _save_and_show(plot_or_layout, filename, folder=None):
    '''
    Helper: saves plot to file (or temp file) and opens it in the browser.
    '''
    if SAVE_PLOTS and folder:
        filepath = os.path.join(folder, filename)
    else:
        filepath = tempfile.mktemp(suffix=".html")
    
    output_file(filepath)
    save(plot_or_layout)
    show(plot_or_layout)


def create_plots(training, ideal_functions, test_data):
    '''
    Main visualization entry point. 
    Generates all activated plots and optionally saves them to a timestamped folder.
    '''

    if SAVE_PLOTS:
        folder = _create_output_folder()
    else:
        folder = None

    if PLOT_RESULTS:
        plot_results(test_data, ideal_functions, training.best_ideal_functions, folder)
    if PLOT_TRAINING_FIT:
        plot_training_fit(training, ideal_functions, folder)
    if PLOT_TRAINING:
        plot_training(training, folder)
    if PLOT_IDEAL:
        plot_ideal(ideal_functions, folder)
    if PLOT_TESTDATA:
        plot_test(test_data, folder)


def plot_results(test_data, ideal_functions, best_ideal_functions, folder=None):
    '''
    Plots the chosen ideal functions and the mapped test data points.
    Adds a column chart to show the deviations of each point
    Unmatched test points are shown in grey.
    '''
    if test_data.results is None or test_data.results.empty:
        print("Warning: no test results to plot — skipping.")
        return
    # Create the figure
    plot = figure(
        title="Test data mapping to ideal functions",
        x_axis_label="x",
        y_axis_label="y"
    )

    ideal_df = ideal_functions.data.set_index('x')
    x_vals = ideal_df.index.tolist()

    # Assign a colour to each of the ideal functions
    chosen = list(best_ideal_functions.values())
    n = len(chosen)
    colours = Category10[max(3, min(n, 10))]

    # Plot each chosen ideal function as a line
    for i, ideal_col in enumerate(chosen):
        plot.line(
            x=x_vals,
            y=ideal_df[ideal_col].tolist(),
            color=colours[i],
            legend_label=f"Ideal: {ideal_col}",
            line_width=2
        )

    # Plot matched test data points
    matched = test_data.results.dropna(subset=['ideal_funct'])
    for i, ideal_col in enumerate(chosen):
        subset = matched[matched['ideal_funct'] == ideal_col]
        plot.scatter(
            x=subset['x'].tolist(),
            y=subset['y'].tolist(),
            color=colours[i],
            size=8,
            legend_label=f"Test → {ideal_col}"
        )

    # Plot unmatched test data points in grey
    unmatched = test_data.results[test_data.results['ideal_funct'].isna()]
    plot.scatter(
        x=unmatched['x'].tolist(),
        y=unmatched['y'].tolist(),
        color='grey',
        size=6,
        legend_label="Unmatched"
    )

    plot.legend.click_policy = "hide"  # click legend to toggle visibility

    # Deviation bar chart (matched points only)
    dev_plot = figure(
        title="Deviation of matched test points from ideal functions",
        x_axis_label="x",
        y_axis_label="delta_y",
        x_range=plot.x_range,   # links x-axis to main plot for synced panning/zooming
        width=plot.width
    )

    # Map each bar to the colour of its ideal function
    color_map = {ideal_col: colours[i] for i, ideal_col in enumerate(chosen)}
    bar_colors = [color_map[funct] for funct in matched['ideal_funct'].tolist()]

    dev_plot.scatter(
        x=matched['x'].tolist(),
        y=matched['delta_y'].tolist(),
        color=bar_colors,
        size=6,
        alpha=0.7
    )

    # Stack both plots vertically
    layout = column(plot, dev_plot)

    _save_and_show(layout, "results_plot.html", folder)

def plot_training_fit(training_data, ideal_functions, folder=None):
    '''
    Plots each training dataset overlaid with its best-matching ideal function.
    Adds a separate deviation plot per training dataset.
    '''
    plot = figure(
        title="Training Data vs. Best-Fitting Ideal Functions",
        x_axis_label="x",
        y_axis_label="y"
    )

    train_df = training_data.data.set_index('x')
    ideal_df = ideal_functions.data.set_index('x')
    x_vals = train_df.index.tolist()
    n = len(training_data.best_ideal_functions)
    colours = Category10[max(3, min(n, 10))]

    for i, (train_col, ideal_col) in enumerate(training_data.best_ideal_functions.items()):
        plot.scatter(
            x=x_vals,
            y=train_df[train_col].tolist(),
            color=colours[i],
            size=8,
            legend_label=f"Training: {train_col}"
        )
        plot.line(
            x=x_vals,
            y=ideal_df[ideal_col].tolist(),
            color=colours[i],
            line_width=2,
            legend_label=f"Ideal: {ideal_col}"
        )

    plot.legend.click_policy = "hide"

    # Generate one deviation plot per training dataset
    dev_plots = []
    for i, (train_col, ideal_col) in enumerate(training_data.best_ideal_functions.items()):
        deviations = (ideal_df[ideal_col] - train_df[train_col]).tolist()
        dev_plot = figure(
            title=f"Deviation: {train_col} → {ideal_col}",
            x_axis_label="x",
            y_axis_label="|deviation|",
            x_range=plot.x_range,
            width=plot.width
        )
        dev_plot.scatter(
            x=x_vals,
            y=deviations,
            color=colours[i],
            size=4
        )
        dev_plots.append(dev_plot)

    layout = column(plot, *dev_plots)
    _save_and_show(layout, "training_fit_plot.html", folder)

def plot_training(training_data, folder=None):
    '''
    Plots the training datasets as scatter plots
    '''
    # Create the figure
    plot = figure(
        title="Training datasets",
        x_axis_label="x",
        y_axis_label="y"
    )

    training_df = training_data.data.set_index('x')
    x_vals = training_df.index.tolist()

    # Assign a colour to each of the training functions
    n = len(training_df.columns)
    colours = Category10[max(3, min(n, 10))]
    training_set = training_df.columns.tolist()

    # Plot each of the training functions as scatter
    for i, train_col in enumerate(training_set):
        plot.scatter(
            x=x_vals,
            y=training_df[train_col].tolist(),
            color=colours[i],
            legend_label=f"Set no.: {train_col}",
            size=2
        )
    
    plot.legend.click_policy = "hide" # click legend to toggle visibility

    _save_and_show(plot, "training_plot.html", folder)


def plot_ideal(ideal_data, folder=None):
    '''
    Plot the ideal functions as separate lines
    '''
    # Create the figure
    plot = figure(
        title="Ideal functions",
        x_axis_label="x",
        y_axis_label="y"
    )

    ideal_df = ideal_data.data.set_index('x')
    x_vals = ideal_df.index.tolist()

    # Assign a colour to each of the ideal functions
    ideal_set = ideal_df.columns.tolist()
    colours = [Category10[10][i % 10] for i in range(len(ideal_set))] # iterate over the Category10 set of colors
    
    # Plot each of the ideal functions as line
    for i, ideal_col in enumerate(ideal_set):
        plot.line(
            x=x_vals,
            y=ideal_df[ideal_col].tolist(),
            color=colours[i],
            legend_label=f"Ideal function no.: {ideal_col}",
            line_width=2
        )
    
    plot.legend.visible = False # hide legend due to high number of entries in this visualization 

    _save_and_show(plot, "ideal_plot.html", folder)


def plot_test(test_data, folder=None):
    '''
    Plots the test data as scatter
    '''
    # Create the figure
    plot = figure(
        title="Test data",
        x_axis_label="x",
        y_axis_label="y"
    )

    test_df = test_data.data

    # Plot each of the test functions as scatter
    plot.scatter(
        x=test_df['x'].tolist(),
        y=test_df['y'].tolist(),
        legend_label="Test Data",
        size=6,
        color='steelblue'
    )
    plot.legend.visible = False # hide legend since there is only one category of data points

    _save_and_show(plot, "test_plot.html", folder)