import json
from datetime import datetime

import numpy
import plotly
import plotly.graph_objs as go
from plotly.subplots import make_subplots

from fedn.network.storage.statestore.mongostatestore import MongoStateStore


class Plot:
    """

    """

    def __init__(self, statestore):
        try:
            statestore_config = statestore.get_config()
            statestore = MongoStateStore(
                statestore_config['network_id'], statestore_config['mongo_config'])
            self.mdb = statestore.connect()
            self.status = self.mdb['control.status']
            self.round_time = self.mdb["control.round_time"]
            self.combiner_round_time = self.mdb["control.combiner_round_time"]
            self.psutil_usage = self.mdb["control.psutil_monitoring"]
            self.network_clients = self.mdb["network.clients"]

        except Exception as e:
            print("FAILED TO CONNECT TO MONGO, {}".format(e), flush=True)
            self.collection = None
            raise

    # plot metrics from DB
    def _scalar_metrics(self, metrics):
        """ Extract all scalar valued metrics from a MODEL_VALIDATON. """

        data = json.loads(metrics['data'])
        data = json.loads(data['data'])

        valid_metrics = []
        for metric, val in data.items():
            # If it can be converted to a float it is a valid, scalar metric
            try:
                val = float(val)
                valid_metrics.append(metric)
            except Exception:
                pass

        return valid_metrics

    def create_table_plot(self):
        """

        :return:
        """
        metrics = self.status.find_one({'type': 'MODEL_VALIDATION'})
        if metrics is None:
            fig = go.Figure(data=[])
            fig.update_layout(
                title_text='No data currently available for table mean metrics')
            table = json.dumps(fig, cls=plotly.utils.PlotlyJSONEncoder)
            return False

        valid_metrics = self._scalar_metrics(metrics)
        if valid_metrics == []:
            fig = go.Figure(data=[])
            fig.update_layout(title_text='No scalar metrics found')
            table = json.dumps(fig, cls=plotly.utils.PlotlyJSONEncoder)
            return False

        all_vals = []
        models = []
        for metric in valid_metrics:
            validations = {}
            for post in self.status.find({'type': 'MODEL_VALIDATION'}):
                e = json.loads(post['data'])
                try:
                    validations[e['modelId']].append(
                        float(json.loads(e['data'])[metric]))
                except KeyError:
                    validations[e['modelId']] = [
                        float(json.loads(e['data'])[metric])]

            vals = []
            models = []
            for model, data in validations.items():
                vals.append(numpy.mean(data))
                models.append(model)
            all_vals.append(vals)

        header_vals = valid_metrics
        models.reverse()
        values = [models]

        for vals in all_vals:
            vals.reverse()
            values.append(vals)

        fig = go.Figure(data=[go.Table(
            header=dict(values=['Model ID'] + header_vals,
                        line_color='darkslategray',
                        fill_color='lightskyblue',
                        align='left'),

            cells=dict(values=values,  # 2nd column
                       line_color='darkslategray',
                       fill_color='lightcyan',
                       align='left'))
        ])

        fig.update_layout(title_text='Summary: mean metrics')
        table = json.dumps(fig, cls=plotly.utils.PlotlyJSONEncoder)
        return table

    def create_timeline_plot(self):
        """

        :return:
        """
        trace_data = []
        x = []
        y = []
        base = []
        for p in self.status.find({'type': 'MODEL_UPDATE_REQUEST'}):
            e = json.loads(p['data'])
            cid = e['correlationId']
            for cc in self.status.find({'sender': p['sender'], 'type': 'MODEL_UPDATE'}):
                da = json.loads(cc['data'])
                if da['correlationId'] == cid:
                    cp = cc

            cd = json.loads(cp['data'])
            tr = datetime.strptime(e['timestamp'], '%Y-%m-%d %H:%M:%S.%f')
            tu = datetime.strptime(cd['timestamp'], '%Y-%m-%d %H:%M:%S.%f')
            ts = tu - tr
            base.append(tr.timestamp())
            x.append(ts.total_seconds() / 60.0)
            y.append(p['sender']['name'])

        trace_data.append(go.Bar(
            x=y,
            y=x,
            marker=dict(color='royalblue'),
            name="Training",
        ))

        x = []
        y = []
        base = []
        for p in self.status.find({'type': 'MODEL_VALIDATION_REQUEST'}):
            e = json.loads(p['data'])
            cid = e['correlationId']
            for cc in self.status.find({'sender': p['sender'], 'type': 'MODEL_VALIDATION'}):
                da = json.loads(cc['data'])
                if da['correlationId'] == cid:
                    cp = cc
            cd = json.loads(cp['data'])
            tr = datetime.strptime(e['timestamp'], '%Y-%m-%d %H:%M:%S.%f')
            tu = datetime.strptime(cd['timestamp'], '%Y-%m-%d %H:%M:%S.%f')
            ts = tu - tr
            base.append(tr.timestamp())
            x.append(ts.total_seconds() / 60.0)
            y.append(p['sender']['name'])

        trace_data.append(go.Bar(
            x=y,
            y=x,
            marker=dict(color='lightskyblue'),
            name="Validation",
        ))

        layout = go.Layout(
            barmode='stack',
            showlegend=True,
        )

        fig = go.Figure(data=trace_data, layout=layout)
        fig.update_xaxes(title_text='Alliance/client')
        fig.update_yaxes(title_text='Time (Min)')
        fig.update_layout(title_text='Alliance timeline')
        timeline = json.dumps(fig, cls=plotly.utils.PlotlyJSONEncoder)
        return timeline

    def create_client_training_distribution(self):
        """

        :return:
        """
        training = []
        for p in self.status.find({'type': 'MODEL_UPDATE'}):
            e = json.loads(p['data'])
            meta = json.loads(e['meta'])
            training.append(meta['exec_training'])

        if not training:
            return False
        fig = go.Figure(data=go.Histogram(x=training))
        fig.update_layout(
            title_text='Client model training time, mean: {}'.format(numpy.mean(training)))
        histogram = json.dumps(fig, cls=plotly.utils.PlotlyJSONEncoder)
        return histogram

    def create_client_histogram_plot(self):
        """

        :return:
        """
        training = []
        for p in self.status.find({'type': 'MODEL_UPDATE'}):
            e = json.loads(p['data'])
            meta = json.loads(e['meta'])
            training.append(meta['exec_training'])

        fig = go.Figure()

        fig.update_layout(
            template="simple_white",
            xaxis=dict(title_text="Time (s)"),
            yaxis=dict(title_text='Number of updates'),
            title="Mean client training time: {}".format(
                numpy.mean(training)),
            # showlegend=True
        )
        if not training:
            return False

        fig.add_trace(go.Histogram(x=training))

        histogram_plot = json.dumps(fig, cls=plotly.utils.PlotlyJSONEncoder)
        return histogram_plot

    def create_client_plot(self):
        """

        :return:
        """
        processing = []
        upload = []
        download = []
        training = []
        for p in self.status.find({'type': 'MODEL_UPDATE'}):
            e = json.loads(p['data'])
            meta = json.loads(e['meta'])
            upload.append(meta['upload_model'])
            download.append(meta['fetch_model'])
            training.append(meta['exec_training'])
            processing.append(meta['processing_time'])

        fig = go.Figure()
        fig.update_layout(
            template="simple_white",
            title="Mean client processing time: {}".format(
                numpy.mean(processing)),
            showlegend=True
        )
        if not processing:
            return False
        data = [numpy.mean(training), numpy.mean(upload), numpy.mean(download)]
        labels = ["Training execution", "Model upload (to combiner)", "Model download (from combiner)"]
        fig.add_trace(go.Pie(labels=labels, values=data))

        client_plot = json.dumps(fig, cls=plotly.utils.PlotlyJSONEncoder)
        return client_plot

    def create_combiner_plot(self):
        """

        :return:
        """
        waiting = []
        aggregation = []
        model_load = []
        combination = []
        for round in self.mdb['control.round'].find():
            try:
                for combiner in round['combiners']:
                    data = combiner
                    stats = data['local_round']['1']
                    ml = stats['aggregation_time']['time_model_load']
                    ag = stats['aggregation_time']['time_model_aggregation']
                    combination.append(stats['time_combination'])
                    waiting.append(stats['time_combination'] - ml - ag)
                    model_load.append(ml)
                    aggregation.append(ag)
            except Exception:
                pass

        labels = ['Waiting for client updates',
                  'Aggregation', 'Loading model updates from disk']
        val = [numpy.mean(waiting), numpy.mean(
            aggregation), numpy.mean(model_load)]
        fig = go.Figure()

        fig.update_layout(
            template="simple_white",
            title="Mean combiner round time: {}".format(
                numpy.mean(combination)),
            showlegend=True
        )
        if not combination:
            return False
        fig.add_trace(go.Pie(labels=labels, values=val))
        combiner_plot = json.dumps(fig, cls=plotly.utils.PlotlyJSONEncoder)
        return combiner_plot

    def fetch_valid_metrics(self):
        """

        :return:
        """
        metrics = self.status.find_one({'type': 'MODEL_VALIDATION'})
        valid_metrics = self._scalar_metrics(metrics)
        return valid_metrics

    def create_box_plot(self, metric):
        """

        :param metric:
        :return:
        """
        metrics = self.status.find_one({'type': 'MODEL_VALIDATION'})
        if metrics is None:
            fig = go.Figure(data=[])
            fig.update_layout(title_text='No data currently available for metric distribution over  '
                                         'participants')
            box = json.dumps(fig, cls=plotly.utils.PlotlyJSONEncoder)
            return box

        valid_metrics = self._scalar_metrics(metrics)
        if valid_metrics == []:
            fig = go.Figure(data=[])
            fig.update_layout(title_text='No scalar metrics found')
            box = json.dumps(fig, cls=plotly.utils.PlotlyJSONEncoder)
            return box

        validations = {}
        for post in self.status.find({'type': 'MODEL_VALIDATION'}):
            e = json.loads(post['data'])
            try:
                validations[e['modelId']].append(
                    float(json.loads(e['data'])[metric]))
            except KeyError:
                validations[e['modelId']] = [
                    float(json.loads(e['data'])[metric])]

        # Make sure validations are plotted in chronological order
        model_trail = self.mdb.control.model.find_one({'key': 'model_trail'})
        model_trail_ids = model_trail['model']
        validations_sorted = []
        for model_id in model_trail_ids:
            try:
                validations_sorted.append(validations[model_id])
            except Exception:
                pass

        validations = validations_sorted

        box = go.Figure()

        y = []
        for j, acc in enumerate(validations):
            # x.append(j)
            y.append(numpy.mean([float(i) for i in acc]))
            if len(acc) >= 2:
                box.add_trace(go.Box(y=acc, name=str(j), marker_color="royalblue", showlegend=False,
                                     boxpoints=False))
            else:
                box.add_trace(go.Scatter(
                    x=[str(j)], y=[y[j]], showlegend=False))

        rounds = list(range(len(y)))
        box.add_trace(go.Scatter(
            x=rounds,
            y=y,
            name='Mean'
        ))

        box.update_xaxes(title_text='Rounds')
        box.update_yaxes(
            tickvals=[0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0])
        box.update_layout(title_text='Metric distribution over clients: {}'.format(metric),
                          margin=dict(l=20, r=20, t=45, b=20))
        box = json.dumps(box, cls=plotly.utils.PlotlyJSONEncoder)
        return box

    def create_round_plot(self):
        """

        :return:
        """
        trace_data = []
        metrics = self.round_time.find_one({'key': 'round_time'})
        if metrics is None:
            fig = go.Figure(data=[])
            fig.update_layout(
                title_text='No data currently available for round time')
            return False

        for post in self.round_time.find({'key': 'round_time'}):
            rounds = post['round']
            traces_data = post['round_time']

        trace_data.append(go.Scatter(
            x=rounds,
            y=traces_data,
            mode='lines+markers',
            name='Reducer'
        ))

        for rec in self.combiner_round_time.find({'key': 'combiner_round_time'}):
            c_traces_data = rec['round_time']

        trace_data.append(go.Scatter(
            x=rounds,
            y=c_traces_data,
            mode='lines+markers',
            name='Combiner'
        ))

        fig = go.Figure(data=trace_data)
        fig.update_xaxes(title_text='Round')
        fig.update_yaxes(title_text='Time (s)')
        fig.update_layout(title_text='Round time')
        round_t = json.dumps(fig, cls=plotly.utils.PlotlyJSONEncoder)
        return round_t

    def create_cpu_plot(self):
        """

        :return:
        """
        metrics = self.psutil_usage.find_one({'key': 'cpu_mem_usage'})
        if metrics is None:
            fig = go.Figure(data=[])
            fig.update_layout(
                title_text='No data currently available for MEM and CPU usage')
            cpu = json.dumps(fig, cls=plotly.utils.PlotlyJSONEncoder)
            return False

        for post in self.psutil_usage.find({'key': 'cpu_mem_usage'}):
            cpu = post['cpu']
            mem = post['mem']
            ps_time = post['time']
            round = post['round']

        # Create figure with secondary y-axis
        fig = make_subplots(specs=[[{"secondary_y": True}]])
        fig.add_trace(go.Scatter(
            x=ps_time,
            y=cpu,
            mode='lines+markers',
            name='CPU (%)'
        ))

        fig.add_trace(go.Scatter(
            x=ps_time,
            y=mem,
            mode='lines+markers',
            name='MEM (%)'
        ))

        fig.add_trace(go.Scatter(
            x=ps_time,
            y=round,
            mode='lines+markers',
            name='Round',
        ), secondary_y=True)

        fig.update_xaxes(title_text='Date Time')
        fig.update_yaxes(title_text='Percentage (%)')
        fig.update_yaxes(title_text="Round", secondary_y=True)
        fig.update_layout(title_text='CPU loads and memory usage')
        cpu = json.dumps(fig, cls=plotly.utils.PlotlyJSONEncoder)
        return cpu
