import json
from collections import defaultdict

from ScEpTIC.analysis.options import AnalysisResultFormat
from ScEpTIC.exceptions import ConfigurationException


class CustomMetricsManager:

    def __init__(self, vmstate):
        self.metrics = {}
        self.names = []
        self.vmstate = vmstate

        self.last_energy = defaultdict(float)
        self.last_time = defaultdict(float)
        self.last_cc = defaultdict(float)

        for metric_id, metric_settings in self.vmstate.config.custom_metrics.metrics.items():
            self._register(
                metric_id,
                metric_settings['metric_name'],
                metric_settings['collect_energy'],
                metric_settings['collect_time'],
                metric_settings['collect_cc'],
                metric_settings['data_diffs'],
                metric_settings['print_data'],
            )

    def _register(self, metric_id, metric_name, collect_energy=False, collect_time=False, collect_cc=False, data_diffs=False, print_data=False):
        """
        Register a custom metric
        :param metric_id: Metric ID
        :param metric_name: Metric name
        :param collect_energy: Records simulation energy when metric is incremented
        :param collect_time: Record simulation energy when metric is incremented
        """
        if metric_id in self.metrics:
            raise ValueError(f"Metric with id #{metric_id} already registered")

        if metric_name in self.names:
            raise ValueError(f"Metric with name {metric_name} already registered")

        self.names.append(metric_name)
        self.metrics[metric_id] = {
            'name': metric_name,
            'value': 0,
            'energy': [] if collect_energy else None,
            'time': [] if collect_time else None,
            'cc': [] if collect_cc else None,
            'data_diffs': data_diffs,
            'print_data': print_data,
        }

    def increment(self, metric_id, val=1):
        """
        Increment custom metric by val
        :param metric_id: metric_id
        :param val: value to increment
        """
        if metric_id not in self.metrics:
            raise ValueError(f"Metric with id #{metric_id} not registered")

        self.metrics[metric_id]['value'] += val

        #print(f"{metric_id}: {self.metrics[metric_id]['value']}")

        if self.metrics[metric_id]['energy'] is not None:
            e = self.vmstate.config.analysis.energy.system_model.get_used_energy()

            if self.metrics[metric_id]['data_diffs']:
                last_energy = self.last_energy[metric_id]
                self.last_energy[metric_id] = e
                e -= last_energy

            self.metrics[metric_id]['energy'].append(e)

        if self.metrics[metric_id]['time'] is not None:
            t = self.vmstate.config.analysis.energy.system_model.get_simulation_time()

            if self.metrics[metric_id]['data_diffs']:
                last_time = self.last_time[metric_id]
                self.last_time[metric_id] = t
                t -= last_time

            self.metrics[metric_id]['time'].append(t)

        if self.metrics[metric_id]['cc'] is not None:
            cc = self.vmstate.config.analysis.energy.system_model.get_elapsed_ticks()

            if self.metrics[metric_id]['data_diffs']:
                last_cc = self.last_cc[metric_id]
                self.last_cc[metric_id] = cc
                cc -= last_cc

            self.metrics[metric_id]['cc'].append(cc)

        if self.metrics[metric_id]['print_data']:
            m = self.metrics[metric_id]
            print(f"{m['name']} ({m['value']}) -> E={m['energy'][-1]}J; T={m['time'][-1]}s; CC={m['cc'][-1]}")

    def get_name(self, metric_id):
        """
        Returns metric name from metric_id
        :param metric_id: metric_id
        :return: metric name
        """
        if metric_id not in self.metrics:
            raise ValueError("Metric not registered")

        return self.metrics[metric_id]['name']

    def exists(self, metric_id):
        """
        Returns if metric exists
        :param metric_id: metric_id
        :return: bool
        """
        return metric_id in self.metrics

    def reset(self):
        """
        Resets all custom metrics
        """
        for values in self.metrics.values():
            values['value'] = 0
            collect_energy = values['energy'] is None
            collect_time = values['time'] is None

            # Free space
            del values['energy']
            del values['time']
            import gc
            gc.collect()

            values['energy'] = [] if collect_energy is not None else None
            values['time'] = [] if collect_time is not None else None


    def save_metrics(self, save_dir, result_format):
        """
        Saves custom metrics
        :param save_dir: the path where to save the metrics
        :param result_format: the metrics format
        """
        if not isinstance(result_format, AnalysisResultFormat):
            raise Exception(f"Wrong result_format for {self.__class__.__name__}: {result_format} is not of type AnalysisResultFormat")

        results = self.get_result(result_format)

        with open(f"{save_dir}_custom_metrics.{result_format}", "w") as fp:
            fp.write(results)

    def get_result(self, result_format):
        """
        Returns custom metrics values
        :param result_format: the result format
        """
        data = {}

        for values in self.metrics.values():
            data[values['name']] = {'value': values['value'], 'energy': values['energy'], 'time': values['time'], 'cc': values['cc']}

        if result_format == AnalysisResultFormat.JSON:
            retval = json.dumps(self.metrics)

        elif result_format == AnalysisResultFormat.TEXT:
            retval = ""

            for name, val in data.items():
                retval += f"{name}:\n"
                retval += f"  value: {val['value']}\n"

                if val['energy'] is not None:
                    retval += f"  energy:\n"
                    for e in val['energy']:
                        retval += f"    - {e}J\n"

                if val['time'] is not None:
                    retval += f"  time:\n"
                    for t in val['time']:
                        retval += f"    - {t}s\n"

        else:
            raise ConfigurationException(f"Invalid result format {result_format}")

        return retval