#!/usr/bin/env python3
import os
import numpy as np
import pandas as pd
import seaborn as sns

import matplotlib.pyplot as plt


def td2hr(t): return t.total_seconds() / 3600  # Timedelta to hour


class VisTool(object):

    def __init__(self):
        super().__init__()
        self.fig = None
        self.ax = None

    def clear_fig(self):
        self.fig = None
        self.ax = None

    def export_fig(self, vh, export_path='./data/fig'):
        """Export deployment MTBCF plot.
        
        :param vh: vehicle name (string)
        :param export_path: path to export folder 
        :return: N/A
         
        """
        # create folder if needed
        if not os.path.exists(export_path):
            os.makedirs(export_path)

        export_path = os.path.join(export_path, vh + '.png')
        self.fig.savefig(export_path, dpi=600, bbox_inches='tight')

    def generate_plot(self, logs, cr, vh):
        """Generate deployment MTBCF plot.
        
        :param logs: logs DataFrame
        :param cr: critical DataFrame
        :param vh: vehicle name (string)         
        :return: N/A
        
        """

        self.clear_fig()

        # Prep data for visualization
        # ------------------------------------------------------------------------------------------
        cr['tbcf'] = cr['tbcf'].apply(td2hr)
        logs['runtime'] = logs['runtime'].apply(td2hr)

        # deployment MTBCF
        dep = pd.DataFrame(columns=['yr', 'name', 'MTBCF', 'annual'])
        dep['yr'] = logs['year'].apply(int).groupby(logs['deployment']).first()
        dep['name'] = logs['deployment_name'].groupby(logs['deployment']).first()
        dep['MTBCF'] = np.nan

        # compute deployment MTBCF
        for d in dep.index:
            if cr['deployment'].str.contains(d).any():
                dep.loc[d, 'MTBCF'] = cr['tbcf'].loc[cr['deployment'] == d].mean()

        # parse deployment names
        dep['dep_names'] = [' '.join(ln.split('_')[-2:]) for ln in list(dep['name'])]

        # annual MTBCF
        for yr in logs['year'].unique():
            if cr['year'].str.contains(yr).any():
                dep.loc[dep['yr'] == int(yr), 'annual'] = cr['tbcf'].loc[cr['year'] == yr].mean()

        # slice deployments past 2012
        d = dep[dep.yr > 2011]
        d.index = np.arange(len(d))

        # Visualize
        # ------------------------------------------------------------------------------------------
        sns.set(style="darkgrid")
        sns.set_color_codes("dark")

        # Initialize the matplotlib figure
        f, ax = plt.subplots(figsize=(20, 9))
        ax.set_position([0.05, 0.2, 0.9, 0.7])

        # plot deployment MTBCF
        sns.barplot(d.index, y=d.MTBCF, color="b")
        plt.xticks(d.index, d.dep_names, rotation='vertical')

        # overlay annual MTBCF
        for y in d['yr'].unique():

            # compute annual statistics
            mtbcf = d[d['yr'] == y].annual.iloc[0]
            runtime = logs[logs['year'] == str(y)].runtime.sum()
            num_cr = logs[logs['year'] == str(y)].fail_count.sum()
            max_tbcf = cr[cr['year'] == str(y)].tbcf.max()

            # assemble label text
            s = " - MTBCF: {:4.1f} hr " \
                "| RUNTIME: {:6.1f} hr " \
                "| Fail count: {:3.0f} " \
                "| Best run: {:5.1f} hr".format(mtbcf, runtime, num_cr, max_tbcf)

            # plot annual MTBCF line
            plt.plot(d[d['yr'] == y].index, d[d['yr'] == y].annual, linewidth=3, label=str(y)+s)

        # plot annual MTBCF lines
        empty = d[d['MTBCF'].isnull()].index
        plt.scatter(empty, np.zeros(len(empty)), marker='*', clip_on=False, label='No failures')

        # axis labels
        ax.grid(True)
        plt.ylim(ymin=0)
        sns.plt.title(vh.capitalize() + ' - Deployment Mean Time Between Critical Failure (MTBCF)', size=22)
        sns.plt.xlabel('Deployment', size=16)
        sns.plt.ylabel('MTBCF (hour)', size=16)

        plt.legend(loc='upper left', frameon=True, facecolor='white', edgecolor='black', prop={'family': 'monospace'})

        # update self
        self.fig = f
        self.ax = ax


if __name__ == '__main__':

    def readin(vehicle: str, dat: str, json_path='../data/json/') -> pd.DataFrame:
        dtype = {
            'vehicle': 'object',
            'year': 'object',
            'deployment_name': 'object',
            'deployment': 'object',
            'log': 'object',
            'timestamp': 'datetime64[ns]',
            'tbcf': 'timedelta64[ns]',
            'mtbcf': 'timedelta64[ns]',
            'runtime': 'timedelta64[ns]',
            'fail_count': 'int',
            'comp': 'object',
            'msg': 'object',
        }

        df = pd.read_json(json_path + vehicle.lower() + '_' + dat.lower() + '.json', dtype=dtype, orient='records')
        return df

    vis = VisTool()

    vh_ = 'Tethys'
    logs_ = readin(vh_, 'logs')
    cr_ = readin(vh_, 'critical')

    vis.generate_plot(logs_, cr_, vh_)
    vis.export_fig(vh_, export_path='../data/fig')