from typing import Dict, List

import plotly.express as px
import plotly
import pandas as pd
from IPython.display import HTML
from matplotlib import pyplot as plt

from efootprint.abstract_modeling_classes.modeling_object import ModelingObject
from efootprint.constants.units import u
from efootprint.core.hardware.network import Network
from efootprint.core.hardware.servers.server_base_class import Server
from efootprint.core.hardware.storage import Storage
from efootprint.core.service import Service
from efootprint.core.usage.usage_pattern import UsagePattern
from efootprint.core.usage.user_journey import UserJourney
from efootprint.abstract_modeling_classes.explainable_objects import ExplainableQuantity, EmptyExplainableObject, \
    ExplainableHourlyQuantities
from efootprint.logger import logger
from efootprint.utils.plot_emission_diffs import EmissionPlotter
from efootprint.utils.tools import format_co2_amount, display_co2_amount


class System(ModelingObject):
    def __init__(self, name: str, usage_patterns: List[UsagePattern]):
        super().__init__(name)
        self.usage_patterns = usage_patterns
        self.previous_change = None
        self.previous_total_energy_footprints = None
        self.previous_total_fabrication_footprints = None
        self.all_changes = []
        self.initial_total_energy_footprints = None
        self.initial_total_fabrication_footprints = None

    @property
    def user_journeys(self) -> List[UserJourney]:
        output_set = set()
        for usage_pattern in self.usage_patterns:
            output_set.update({usage_pattern.user_journey})

        return list(output_set)

    @property
    def modeling_objects_whose_attributes_depend_directly_on_me(self):
        return self.user_journeys

    @property
    def systems(self) -> List:
        return []

    def after_init(self):
        self.init_has_passed = True
        self.launch_attributes_computation_chain(self.attributes_computation_chain)
        logger.info(f"Finished computing {self.name} modeling")
        self.initial_total_energy_footprints = self.total_energy_footprint_sum_over_period
        self.initial_total_fabrication_footprints = self.total_fabrication_footprint_sum_over_period

    @property
    def servers(self) -> List[Server]:
        output_set = set()
        for usage_pattern in self.usage_patterns:
            output_set.update(usage_pattern.user_journey.servers)

        return list(output_set)

    @property
    def storages(self) -> List[Storage]:
        output_set = set()
        for usage_pattern in self.usage_patterns:
            output_set.update(usage_pattern.user_journey.storages)

        return list(output_set)

    @property
    def services(self) -> List[Service]:
        output_set = set()
        for usage_pattern in self.usage_patterns:
            output_set.update(usage_pattern.user_journey.services)

        return list(output_set)

    @property
    def networks(self) -> List[Network]:
        output_set = set()
        for usage_pattern in self.usage_patterns:
            output_set.update({usage_pattern.network})

        return list(output_set)

    def get_storage_by_name(self, storage_name) -> Storage:
        for storage in self.storages:
            if storage.name == storage_name:
                return storage

    def get_server_by_name(self, server_name) -> Server:
        for server in self.servers:
            if server.name == server_name:
                return server

    def get_usage_pattern_by_name(self, usage_pattern_name) -> UsagePattern:
        for usage_pattern in self.usage_patterns:
            if usage_pattern.name == usage_pattern_name:
                return usage_pattern

    @property
    def fabrication_footprints(self) -> Dict[str, Dict[str, ExplainableHourlyQuantities]]:
        fab_footprints = {
            "Servers": {server.name: server.instances_fabrication_footprint for server in self.servers},
            "Storage": {storage.name: storage.instances_fabrication_footprint for storage in self.storages},
            "Network": {"networks": EmptyExplainableObject()},
            "Devices": {usage_pattern.name: usage_pattern.devices_fabrication_footprint
                        for usage_pattern in self.usage_patterns},
        }

        return fab_footprints

    @property
    def energy_footprints(self) -> Dict[str, Dict[str, ExplainableHourlyQuantities]]:
        energy_footprints = {
            "Servers": {server.name: server.energy_footprint for server in self.servers},
            "Storage": {storage.name: storage.energy_footprint for storage in self.storages},
            "Network": {network.name: network.energy_footprint for network in self.networks},
            "Devices": {usage_pattern.name: usage_pattern.devices_energy_footprint
                        for usage_pattern in self.usage_patterns},
        }

        return energy_footprints

    @property
    def total_fabrication_footprints(self) -> Dict[str, ExplainableHourlyQuantities]:
        fab_footprints = {
            "Servers": sum(server.instances_fabrication_footprint for server in self.servers),
            "Storage": sum(storage.instances_fabrication_footprint for storage in self.storages),
            "Network": EmptyExplainableObject(),
            "Devices": sum(usage_pattern.devices_fabrication_footprint
                           for usage_pattern in self.usage_patterns)
        }

        return fab_footprints

    @property
    def total_energy_footprints(self) -> Dict[str, ExplainableHourlyQuantities]:
        energy_footprints = {
            "Servers": sum(server.energy_footprint for server in self.servers),
            "Storage": sum(storage.energy_footprint for storage in self.storages),
            "Network": sum(network.energy_footprint for network in self.networks),
            "Devices": sum(usage_pattern.devices_energy_footprint for usage_pattern in self.usage_patterns)
        }

        return energy_footprints

    @staticmethod
    def sum_and_remove_empty_explainable_object(expl_obj):
        tmp_sum = expl_obj.sum()
        if isinstance(tmp_sum, EmptyExplainableObject):
            tmp_sum = ExplainableQuantity(0 * u.kg, "null value")

        return tmp_sum

    @property
    def fabrication_footprint_sum_over_period(self) -> Dict[str, Dict[ModelingObject, ExplainableQuantity]]:
        fab_footprints_sum = {}
        for key, dict_value in self.fabrication_footprints.items():
            fab_footprints_sum[key] = {
                obj_key: self.sum_and_remove_empty_explainable_object(obj_value)
                for obj_key, obj_value in dict_value.items()
            }

        return fab_footprints_sum

    @property
    def energy_footprint_sum_over_period(self) -> Dict[str, Dict[ModelingObject, ExplainableQuantity]]:
        energy_footprints_sum = {}
        for key, dict_value in self.energy_footprints.items():
            energy_footprints_sum[key] = {
                obj_key: self.sum_and_remove_empty_explainable_object(obj_value)
                for obj_key, obj_value in dict_value.items()
            }

        return energy_footprints_sum

    @property
    def total_fabrication_footprint_sum_over_period(self) -> Dict[str, ExplainableQuantity]:
        fab_footprints = {
            obj_key: self.sum_and_remove_empty_explainable_object(obj_value)
            for obj_key, obj_value in self.total_fabrication_footprints.items()
        }
        return fab_footprints

    @property
    def total_energy_footprint_sum_over_period(self) -> Dict[str, ExplainableQuantity]:
        energy_footprints = {
            obj_key: self.sum_and_remove_empty_explainable_object(obj_value)
            for obj_key, obj_value in self.total_energy_footprints.items()
        }
        return energy_footprints

    @property
    def total_footprint(self):
        return (
            sum(
                sum(
                    self.fabrication_footprints[key].values()) + sum(self.energy_footprints[key].values())
                for key in self.fabrication_footprints.keys()
            )
        ).set_label(f"{self.name} total carbon footprint")

    def plot_footprints_by_category_and_object(self, filename=None, height=400, width=800, return_only_html=False):
        fab_footprints = self.fabrication_footprint_sum_over_period
        energy_footprints = self.energy_footprint_sum_over_period
        categories = list(fab_footprints.keys())

        rows_as_dicts = []

        value_colname = "tonnes CO2 emissions"
        for category in categories:
            fab_objects = sorted(fab_footprints[category].items(), key=lambda x: x[0])
            energy_objects = sorted(energy_footprints[category].items(), key=lambda x: x[0])

            for objs, color in zip([energy_objects, fab_objects], ["Electricity", "Fabrication"]):
                data_dicts = [
                    {"Type": color, "Category": category, "Object": obj[0],
                     value_colname: obj[1].value.magnitude / 1000,
                     "Amount": f"{display_co2_amount(format_co2_amount(obj[1].value.magnitude))}"}
                    for obj in objs]
                rows_as_dicts += data_dicts

        df = pd.DataFrame.from_records(rows_as_dicts)

        total_co2 = df[value_colname].sum()

        total_footprint = self.total_footprint

        fig = px.bar(
            df, x="Category", y=value_colname, color='Type', barmode='group', height=height, width=width,
            hover_data={"Type": False, "Category": False, "Object": True, value_colname: False, "Amount": True},
            template="plotly_white",
            title=f"Total CO2 emissions from "
                  f"{total_footprint.value.index.min().to_timestamp().date()} "
                  f"to {total_footprint.value.index.max().to_timestamp().date()}"
                  f": {display_co2_amount(format_co2_amount(total_co2 * 1000, rounding_value=0))}"
                  )

        if (max(sum(energy_footprints["Servers"].values()), sum(fab_footprints["Servers"].values())) >
                max(sum(energy_footprints["Devices"].values()), sum(fab_footprints["Devices"].values()))):
            legend_alignment = "right"
            legend_x = 0.98
        else:
            legend_alignment = "left"
            legend_x = 0.02

        fig.update_layout(
            legend={"orientation": "v", "yanchor": "top", "y": 1.02, "xanchor": legend_alignment, "x": legend_x,
                    "title": ""},
            title={"x": 0.5, "y": 0.9, "xanchor": 'center', "yanchor": 'top'})

        total_co2_per_category_and_type = df.groupby(["Category", "Type"])[value_colname].sum()

        for category, source_type in total_co2_per_category_and_type.keys():
            height = total_co2_per_category_and_type.loc[category, source_type]
            x_shift_direction = 1 if source_type == 'Fabrication' else -1

            fig.add_annotation(
                x=category,
                y=height,
                text=f"{int((height / total_co2) * 100)}%",  # Format the label as a percentage
                showarrow=False,
                yshift=10,  # Shift the label slightly above the stack
                xshift=30 * x_shift_direction
            )

        if return_only_html:
            return fig.to_html(full_html=False, include_plotlyjs=False)

        else:
            if filename is None:
                filename = f"{self.name} footprints.html"

            plotly.offline.plot(fig, filename=filename, auto_open=False)

            return HTML(filename)

    def plot_emission_diffs(self, filepath=None, figsize=(10, 5), from_start=False, plt_show=False):
        if self.previous_change is None:
            raise ValueError(
                f"There has been no change to the system yet so no diff to plot.\n"
                f"Use System.plot_footprints_by_category_and_object() to visualize footprints")

        if from_start and len(self.all_changes) > 1:
            changes_list = "\n- ".join([change.replace('changed', 'changing') for change in self.all_changes])
            print(f"Plotting the impact of:\n\n- {changes_list}")
            emissions_dict__old = [self.initial_total_energy_footprints, self.initial_total_fabrication_footprints]
        else:
            print(f"Plotting the impact of {self.previous_change.replace('changed', 'changing')}")
            emissions_dict__old = [self.previous_total_energy_footprints, self.previous_total_fabrication_footprints]

        emissions_dict__new = [self.total_energy_footprint_sum_over_period, self.total_fabrication_footprint_sum_over_period]

        fig, ax = plt.subplots(nrows=1, ncols=1, figsize=figsize)

        EmissionPlotter(
            ax, emissions_dict__old, emissions_dict__new, rounding_value=0).plot_emission_diffs()

        if filepath is not None:
            plt.savefig(filepath, bbox_inches='tight')

        if plt_show:
            plt.show()
