process_frame_expansion.py

import json
import logging
from rich.logging import RichHandler

import pandas as pd
from numpy import polyfit, polyval
from matplotlib import pyplot as plt
from pathlib import Path

import ruptures as rpt

logging.basicConfig(level=logging.INFO, 
                    format="%(message)s",
                    datefmt="[%X]",
                    handlers=[RichHandler()])
logger = logging.getLogger(__name__)


class ThermalProfile():
    NON_TEMP_COLS = ['datetime', 'mcu_z', 'sample_index', 'elapsed_min']

    def __init__(self, json_results_path: str) -> None:
        self._json_results = self.import_json_results(json_results_path)
        self._expansion_thermals = self.temperature_to_pd(self._json_results)
        self._json_path = json_results_path

    def import_json_results(self, path: str) -> dict:
        logger.info('Reading results JSON file: "%s"', path)
        with open(path, 'r') as result_file:
            results = json.load(result_file)
        if results['metadata']['script']['data_structure'] < 3:
            err_msg = 'Dataset generated by incompatible measurement script.\n'\
            'Please rerun the thermal profiling measuring script.'
            raise ValueError(err_msg)
        return results

    def temperature_to_pd(self, results: dict) -> pd.DataFrame:
        'Organize JSON temperature results into a pandas dataframe'
        df_temperatures_raw = (pd.DataFrame(results['temp_data'])
            .reset_index()
            .rename(columns={'index': 'measurement'})
            .melt(
                id_vars='measurement',
                var_name='datetime')
            .pivot(
                index='datetime',
                columns='measurement',
                values='value')
            .reset_index()
        )
        df_temperatures_raw.datetime = pd.to_datetime(df_temperatures_raw.datetime, 
                                                    format="%Y/%m/%d-%H:%M:%S")

        df_temperatures_raw['elapsed_min'] = (df_temperatures_raw.datetime
                                            .diff()
                                            .dt
                                            .total_seconds()\
                                            .cumsum().fillna(0)
                                            .div(60)
        )

        return df_temperatures_raw

    def export_to_csv(self, path: Path, filename: str = None) -> None:
        if not filename:
            timestamp = self.user_data['timestamp']
            username = self.user_data['id']
            filename = f'thermal_profile_raw_{username}_{timestamp}.csv'

        out_path = path/filename
        logging.info('Exported thermal profile CSV to:\n"%s"', out_path.absolute())
        self._expansion_thermals.to_csv(out_path, index = False)

    def _get_temp_columns(self):
        raw = self._expansion_thermals
        raw_temps = raw.loc[ : ,~raw.columns.isin(self.NON_TEMP_COLS)]
        raw_temps = raw_temps.loc[: , [col for col in raw_temps.columns if '_target' not in col]]
        raw_temps = raw_temps.loc[:,
                                  ([(raw_temps[col] != -180.)
                                    .all() for col in raw_temps.columns])
                                  ]
        return raw_temps.columns

    def export_metadata_to_txt(self, path: Path, filename: str = None,) -> None:
        if not filename:
            timestamp = self.user_data['timestamp']
            username = self.user_data['id']
            filename = f'metadata_{username}_{timestamp}'
        out_path = path/f"{filename}.txt"
        self.metadata
        
        with open(out_path, 'w') as file:
            file.write("----- Measurement Metadata -----\n\n")

            # Write User Metadata
            file.write("User:\n")
            file.write(f"\tBackers: {self.metadata['user']['backers']}\n")
            file.write(f"\tHome Type: {self.metadata['user']['home_type']}\n")
            file.write(f"\tID: {self.metadata['user']['id']}\n")
            file.write(f"\tPrinter: {self.metadata['user']['printer']}\n")
            file.write(f"\tProbe Type: {self.metadata['user']['probe_type']}\n")
            file.write(f"\tTimestamp: {self.metadata['user']['timestamp']}\n")
            file.write(f"\tX Rails: {self.metadata['user']['x_rails']}\n\n")

            # Write Z Axis Metadata
            file.write("Z Axis:\n")
            file.write(f"\tHoming Speed: {self.metadata['z_axis']['homing_speed']}\n")
            file.write(f"\tMax Z: {self.metadata['z_axis']['max_z']}\n")
            file.write(f"\tStep Distance: {self.metadata['z_axis']['step_dist']:.5e}\n\n")

            # Write Script Metadata
            file.write("Measurement Script:\n")
            file.write(f"\tData Structure: {self.metadata['script']['data_structure']}\n")
            file.write(f"\tHot Duration: {self.metadata['script']['hot_duration']}\n")

        logging.info('Exported thermal profile metadata to:\n"%s"',
                     out_path.absolute())

    @property
    def z_axis(self) -> dict:
        'Z-axis configuration metadata.'
        return self._json_results['metadata']['z_axis']
    
    @property
    def raw_expansion_thermals(self) -> pd.DataFrame:
        'Unmodified Pandas dataframe object'
        return self._expansion_thermals
    
    @property
    def json_results(self) -> dict:
        return self._json_results

    @property
    def metadata(self) -> dict:
        return self._json_results['metadata']

    @property
    def user_data(self) -> dict:
        return self._json_results['metadata']['user']

    @property
    def timestamp(self) -> str:
        return self._json_results['metadata']['user']['timestamp']

    @property
    def user_id(self) -> str:
        id = str(self._json_results['metadata']['user']['id']).replace('#', '')
        return id if id else 'unknown_user'

    @property
    def results_filename(self) -> str:
        return self._json_path.strip('\\').strip('\\.').replace('.json', '')

    @property
    def temp_sensors(self) -> list:
        return list(self._get_temp_columns())


class ThermalAnalysis():
    def __init__(self, profile: ThermalProfile) -> None:
        self.raw_thermals = profile.raw_expansion_thermals
        self.profile = profile
        self._json_results = profile.json_results
        self.breakpoints = None
        self._manual_range = False
        self._delta_z_sd_cutoff = None
        
        # Calculate the mean and standard deviation of samples
        self.mean_thermals = (self.raw_thermals
            .groupby(['sample_index'])
            .mean(numeric_only=False)
            .join((self.raw_thermals.groupby(['sample_index'])
                .std()),
                    on='sample_index', lsuffix='_mean', rsuffix='_sd')
        )

        self.mean_thermals['delta_z_mean'] = (
            self.mean_thermals['mcu_z_mean']
                .transform(lambda x: (x-x.iloc[0]) * 
                           self._json_results['metadata']['z_axis']['step_dist'])
        )
        self.mean_thermals['delta_z_sd'] = (
            self.mean_thermals['mcu_z_sd'] * 
                self._json_results['metadata']['z_axis']['step_dist']
        )

        self.mean_thermals = (self.mean_thermals
            .drop(columns='elapsed_min_sd')
            .rename(columns={'elapsed_min_mean': 'elapsed_min'})
        )

        self._time_range = [0, max(self.mean_thermals['elapsed_min'])]
        self.filtered_thermals = self.mean_thermals
        logger.debug('Max delta Z SD = %f', self.mean_thermals.delta_z_sd.max())

    def fit_temp_coeff(self, expansion_df: pd.DataFrame, 
                          elapsed_time_range: tuple = None) -> tuple:
        df = expansion_df

        logger.debug('Fitting temp_coeff..')

        # Subset the dataset by elapsed time range
        if elapsed_time_range:
            df = (df.loc[df.elapsed_min >= min(elapsed_time_range)]
                .loc[df.elapsed_min <= max(elapsed_time_range)]
            )

        if any(df['frame_temp_mean'] == -180.):
            logger.warning(('No frame temperature data collected. Cannot'
                           ' calculate [bold italic]temp_coeff[/]'),
                           extra={"markup": True})
            return df, None

        # Linear fitting, store slope (m) and Y-intercept(c)
        m, c = polyfit(
            df['frame_temp_mean'],
            df['delta_z_mean'],
            1
        )

        # Create column of fitted Z values (for plotting)
        df['fit_delta_z'] = (
            polyval([m, c], 
                    df['frame_temp_mean'])
        )
        logger.info(('Calculated '
                     '[bold italic]temp_coeff[/] = '
                     f'{(m * -1): .6f} mm/degC'), extra={"markup": True})
        return df, m, c
    
    def calculate(self):
        self.filtered_thermals = self.filter_std(self.filtered_thermals,
                                                self._delta_z_sd_cutoff)
        if not self._manual_range:
            self.calc_breakpoints(self.filtered_thermals)
            # logger.debug('Breakpoint calculated: %s', self.breakpoints)
            if len(self.breakpoints) >= 2:
                self._time_range[0] = self.breakpoints[0]
                self._time_range[1] = self.breakpoints[1]
                logger.info('Fit range auto set to: %i-%i min.', *self._time_range)
            elif len(self.breakpoints) == 1:
                logger.info('Could not auto set time range. Using full dataset')
        self.filtered_thermals, self._m, self._c = self.fit_temp_coeff(
            self.filtered_thermals,
            self._time_range)

    def calc_breakpoints(self, df: pd.DataFrame):
        window_size = 30
        values = df.delta_z_mean.values

        algo = rpt.Window(width=window_size)
        results = algo.fit_predict(values, n_bkps=2, pen=5)
        results = [idx - 1 for idx in results] # 0-index results
        self.breakpoints = df['elapsed_min'].iloc[results].tolist()
        logger.debug('Breakpoints identified at: %s min,', self.breakpoints)

    def filter_std(self, df: pd.DataFrame, cutoff: float = 0.05):
        # If user used single-point sampling, all sd values are NaN
        if all(df['delta_z_sd'].isna()):
            subset = df
            subset['delta_z_sd'] = 0
            logger.info('Single Z measurements detected.' \
                        '[bold] Ignoring SD cutoff.[/]',
                        extra={"markup": True})
        else:
            subset = df.loc[df['delta_z_sd'] <= cutoff]
            logger.info('Filtered out %i data points by delta_z SD <= %f mm',
                        (len(df)-len(subset)), cutoff)
        return subset

    @property
    def mean_expansion_thermals(self) -> pd.DataFrame:
        return self.mean_thermals

    @property
    def temp_coeff(self) -> float:
        if self._m:
            return -1*self._m
        return None

    @property
    def fit_c(self) -> float:
        return self._c

    @property
    def raw_profile(self) -> ThermalProfile:
        return self.profile
    
    @property
    def fit_range(self) -> list:
        return self._time_range
    
    @fit_range.setter
    def fit_range(self, range):
        self._time_range[0] = range[0] if (range[0] is not None) else self._time_range[0]
        self._time_range[1] = range[1] if range[1] else self._time_range[1]
        if any(range):
            self._manual_range = True
            logger.info('User fit range set to: %i-%i min.', *self._time_range)
        self.calculate()

    @property
    def delta_z_sd_cutoff(self) -> float:
        return self._delta_z_sd_cutoff
    
    @delta_z_sd_cutoff.setter
    def delta_z_sd_cutoff(self, cutoff: float):
        if cutoff <= 0.:
            logger.error('Delta Z SD cutoff must be > 0.0 mm. Using default (0.05 mm).')
            self._delta_z_sd_cutoff = 0.05
        else:
            logger.debug('Delta Z cutoff set to: %f mm', cutoff)
            self._delta_z_sd_cutoff = cutoff


class Plotter():
    def __init__(self, profile: ThermalAnalysis) -> None:
        self._thermal_analysis = profile
        if not self._thermal_analysis.temp_coeff:
            self._thermal_analysis.calculate()

    def plot_timeseries(self, exclude_sensors=['he_temp', 'bed_temp']):
        df = self._thermal_analysis.mean_thermals
        user_data = self._thermal_analysis.raw_profile.user_data
        fit_range = self._thermal_analysis.fit_range
        fig, host = plt.subplots(figsize=(8, 5), dpi=300)
        par1 = host.twinx()

        host.set_xlabel("Elapsed Time [min]")
        host.set_ylabel("Delta Z (+/- S.D.) [mm]")
        par1.set_ylabel("Temperature [degC]")

        colour1 = "black"
        
        # Overlay fitting region
        host.axvspan(fit_range[0],
                     fit_range[1],
                     color='gray', alpha=0.2,
                     label='Fitting Region')
        
        # if not self._thermal_analysis._manual_range:
            # for b in self._thermal_analysis.breakpoints:
                # plt.axvline(x=b, color='gray', linestyle='--')

        # Plot delta Z
        host.plot(df['elapsed_min'], df['delta_z_mean'],
                        label='Delta Z', color=colour1,
                        zorder=4)
        host.errorbar(df['elapsed_min'],
                        df['delta_z_mean'],
                        yerr=df['delta_z_sd'],
                        color='grey',
                        marker=" ",
                        #   fmt=".k",
                        #  ecolor='frame_temp',
                        )

        # Add the text label "Fitting Region"
        label_x_position = (fit_range[0] + fit_range[1]) / 2
        label_y_position = host.get_ylim()[1] - 0.03 * (host.get_ylim()[1] - host.get_ylim()[0])
        host.text(label_x_position, label_y_position, "temp_coeff Fitting Region", 
              ha='center', va='top', fontsize=8, fontstyle='italic',
              color='black',
              bbox=dict(boxstyle='round,pad=1', ec='none', fill=False))

        # Plot temperature sensor data
        for sensor in self._thermal_analysis.raw_profile._get_temp_columns():
            if sensor not in exclude_sensors:
                par1.plot(df['elapsed_min'], df[sensor+'_mean'],
                            label=sensor)

        par1.legend(title='Temperature Sensors', bbox_to_anchor=(1.1,0.5), loc=2)
        title = f"{user_data['id']} ({user_data['timestamp']})\nFrame Expansion Time Series"
        host.set_title(title)
        fig.set_facecolor('white') # Plot background fill
        fig.set_tight_layout(tight=True) # Keep legend within bounds of output
        return plt

    def plot_coeff_fitting_context(self):
        df = self._thermal_analysis.filtered_thermals
        df_all = self._thermal_analysis.mean_thermals
        fit = self._thermal_analysis.temp_coeff * -1
        yint = self._thermal_analysis._c
        user_data = self._thermal_analysis.raw_profile.user_data

        plt.figure(1, (6, 6), dpi=300)
        plt.scatter('frame_temp_mean', 'delta_z_mean', c='grey',
                    data=self._thermal_analysis.mean_thermals,
                    zorder=4,
                    # edgecolors="black",
                    label='Unused/Filtered Points',
                    alpha=0.5)
        plt.scatter('frame_temp_mean', 'delta_z_mean', c='elapsed_min',
                    cmap='viridis', data=df,
                    zorder=4,
                    edgecolors="black",
                    label='Fitting Point',
                    alpha=0.8)
        plt.errorbar('frame_temp_mean',
                     'delta_z_mean',
                     yerr=df['delta_z_sd'],
                     fmt=".k",
                     #  ecolor='frame_temp',
                     marker=None,
                     data=df,
                     alpha=0.8)
        plt.axline(xy1=(0.0, yint),
                   slope=fit,
                   linestyle="--",
                    c='black')
        title = ("%s (%s)\nTemperature Coefficient Fitting\n"
                 "All Data Points (Unused/Removed Points in Grey)" %
                    (user_data['id'],
                     user_data['timestamp']
                     )) 
        plt.title(title)
       
        # Adjust limits to avoid plotting the fit overlay Y-intercept point
        xrange = df_all.frame_temp_mean.max() - df_all.frame_temp_mean.min()
        yrange =  df_all.delta_z_mean.max() - df_all.delta_z_mean.min()
        plt.xlim(df_all.frame_temp_mean.min() - xrange * 0.1,
                 df_all.frame_temp_mean.max() + xrange * 0.1)
        plt.ylim(df_all.delta_z_mean.min() - yrange * 0.1,
                 (df_all.delta_z_mean + df_all.delta_z_sd).max() + yrange * 0.1)
        
        plt.xlabel('Frame Temperature [degC]')
        plt.ylabel('Delta Z (Mean ± S.D.) [mm] ')

        # Try to avoid plotting the temp_coeff overtop of data
        coeff_annot_x = 0.2 if (fit > 0) else 0.75
        plt.annotate(text="temp_coeff:\n%.4f mm/K" % (-1*fit),
                     xy=(coeff_annot_x, 0.85), xycoords='figure fraction',
                     ha='left' if (fit > 0) else 'right',
                     va='top',
                     zorder=10,
                     bbox=dict(boxstyle='square, pad=0.5', facecolor='white',
                               edgecolor='black', alpha=0.7))
        cbar = plt.colorbar()
        cbar.set_label('Elapsed Time [min]')
        cbar.set_ticks(range(0, int(df_all.elapsed_min.max()), 10))
        plt.tight_layout()
        return plt

    def plot_coeff_fitting(self):
        df = self._thermal_analysis.filtered_thermals
        fit = self._thermal_analysis.temp_coeff * -1
        yint = self._thermal_analysis._c
        user_data = self._thermal_analysis.raw_profile.user_data

        plt.figure(1, (6, 6), dpi=300)
        plt.scatter('frame_temp_mean', 'delta_z_mean', c='elapsed_min',
                    cmap='viridis', data=df,
                    zorder=4,
                    edgecolors="black",
                    alpha=0.8)
        plt.errorbar('frame_temp_mean',
                     'delta_z_mean',
                     yerr=df['delta_z_sd'],
                     fmt=".k",
                     #  ecolor='frame_temp',
                     marker=None,
                     data=df,
                     alpha=0.8)
        plt.axline(xy1=(0.0, yint),
                   slope=fit,
                   linestyle="--",
                    c='black')
        title = ("%s (%s)\nTemperature Coefficient Fitting\n"
                 "Time Range Subset: %i-%i min." %
                    (user_data['id'],
                     user_data['timestamp'],
                     *self._thermal_analysis.fit_range)) 
        cbar = plt.colorbar()
        cbar.set_label('Elapsed Time [min]')
        cbar.set_ticks(range(int(df.elapsed_min.min()),
                             int(df.elapsed_min.max()),
                             10))
        plt.title(title)
       
        # Adjust limits to avoid plotting the fit overlay Y-intercept point
        xrange = df.frame_temp_mean.max() - df.frame_temp_mean.min()
        yrange =  df.delta_z_mean.max() - df.delta_z_mean.min()
        plt.xlim(df.frame_temp_mean.min() - xrange * 0.1,
                 df.frame_temp_mean.max() + xrange * 0.1)
        plt.ylim(df.delta_z_mean.min() - yrange * 0.1,
                 (df.delta_z_mean + df.delta_z_sd).max() + yrange * 0.1)
        
        plt.xlabel('Frame Temperature [degC]')
        plt.ylabel('Delta Z (Mean ± S.D.) [mm] ')

        # Try to avoid plotting the temp_coeff overtop of data
        coeff_annot_x = 0.2 if (fit > 0) else 0.75
        plt.annotate(text="temp_coeff:\n%.4f mm/K" % (-1*fit),
                     xy=(coeff_annot_x, 0.85), xycoords='figure fraction',
                     ha='left' if (fit > 0) else 'right',
                     va='top',
                     zorder=10,
                     bbox=dict(boxstyle='square, pad=0.5', facecolor='white',
                               edgecolor='black', alpha=0.7))
        plt.tight_layout()
        return plt

    def save_all_plots(self, dir: Path):
        plt.close()
        self.plot_coeff_fitting().savefig(dir/'coeff_fit_plot.png')
        plt.close()
        self.plot_timeseries().savefig(dir/'timeseries.png')
        plt.close()
        self.plot_coeff_fitting_context().savefig(dir/'coff_fit_context_plot.png')
        plt.close()
        logger.info('All plots saved to:\n"%s"', dir.absolute())

desc_result_json = 'Path to thermal behaviour JSON results input file.'
desc_fit_start = 'Exclude measurement points collected [bold]BEFORE[/] this time when fitting [italic]temp_coeff[/] [bold]\[minutes][/].'
desc_fit_end = 'Exclude measurement points collected [bold]AFTER[/] this time when fitting [italic]temp_coeff[/] [bold]\[minutes][/].'
desc_sd_cutoff = 'Standard deviation cutoff for delta_z data point filter. Default = 0.05. [bold]\[mm][/]'
if __name__ == "__main__":
    import argparse
    from rich_argparse import RichHelpFormatter
    # logger.setLevel(logging.DEBUG)

    parser = argparse.ArgumentParser(formatter_class=RichHelpFormatter)
    parser.add_argument('results_json', type=str, help=desc_result_json)
    parser.add_argument('--start', type=int, required=False, default=None, 
                        dest='fit_start_minutes', help=desc_fit_start)
    parser.add_argument('--end', type=int, required=False, default=None,
                        dest='fit_end_minutes', help=desc_fit_end)
    parser.add_argument('--sd_cutoff', type=float, required=False, default=0.05,
                        dest='sd_cutoff', help=desc_sd_cutoff)
    parser.add_argument('-v', '--verbose', action='count', required=False, default=0,
                        dest='verbose')
    args = parser.parse_args()

    if(args.verbose > 0):
        logger.setLevel(logging.DEBUG)

    dataset_name = Path(args.results_json.strip('.\\'))
    profile = ThermalProfile(dataset_name)
    analysis = ThermalAnalysis(profile)
    analysis.delta_z_sd_cutoff = args.sd_cutoff
    analysis.fit_range = [args.fit_start_minutes, args.fit_end_minutes]
    plotter = Plotter(analysis)
    
    output_path = Path(profile.user_id, profile.timestamp)
    output_path.mkdir(parents=True, exist_ok=True)

    profile.export_to_csv(output_path)
    profile.export_metadata_to_txt(output_path, 'metadata')
    plotter.save_all_plots(output_path)