import importlib.resources
import json
from abc import ABC, abstractmethod
from typing import Iterable, List, Tuple, Union

import numpy as np

import physrisk.data.static.vulnerability

from ..data_objects.vulnerability_curve import VulnerabilityCurves
from .assets import Asset
from .curve import ExceedanceCurve
from .hazard_event_distrib import HazardEventDistrib
from .hazard_model import EventDataRequest, EventDataResponse
from .impact_curve import ImpactCurve
from .impact_distrib import ImpactDistrib
from .vulnerability_distrib import VulnerabilityDistrib

PLUGINS = dict()


def repeat(num_times):
    def decorator_repeat(func):
        ...  # Create and return a wrapper function

    return decorator_repeat


def applies_to_events(event_types):
    def decorator_events(func):
        PLUGINS[func.__name__] = func
        return func

    return decorator_events


def applies_to_assets(asset_types):
    def decorator_events(func):
        PLUGINS[func.__name__] = func
        return func

    return decorator_events


def get_vulnerability_curves_from_resource(id: str) -> VulnerabilityCurves:
    with importlib.resources.open_text(physrisk.data.static.vulnerability, id + ".json") as f:
        curve_set = VulnerabilityCurves(**json.load(f))
        return curve_set


class VulnerabilityModelBase(ABC):
    """Models generate the VulnerabilityDistrib and HazardEventDistrib of an
    Asset.
    """

    def __init__(self, model: str, event_type: type, year: int, scenario: str):
        self.model = model
        self.event_type = event_type
        self.year = year
        self.scenario = scenario
        self._event_types: List[type] = []
        self._asset_types: List[type] = []

    @abstractmethod
    def get_event_data_requests(
        self, asset: Asset, *, scenario: str, year: int
    ) -> Union[EventDataRequest, Iterable[EventDataRequest]]:
        """Provide the one or more hazard event data requests required in order to calculate
        the VulnerabilityDistrib and HazardEventDistrib for the asset."""
        ...

    @abstractmethod
    def get_distributions(
        self, asset: Asset, event_data_responses: Iterable[EventDataResponse]
    ) -> Tuple[VulnerabilityDistrib, HazardEventDistrib]:
        """Return distributions for asset: VulnerabilityDistrib and HazardEventDistrib.
        The hazard event data is used to do this.

        Args:
            asset: the asset.
            event_data_responses: the responses to the requests made by get_event_data_requests, in the same order.
        """
        ...

    def get_impact(
        self, asset: Asset, event_data_responses: Iterable[EventDataResponse]
    ) -> Tuple[ImpactDistrib, VulnerabilityDistrib, HazardEventDistrib]:
        """Return impact distribution along with vulnerability and hazard event distributions used to infer this.

        Args:
            asset: the asset.
            event_data_responses: the responses to the requests made by get_event_data_requests, in the same order.
        """
        vulnerability_dist, event_dist = self.get_distributions(asset, event_data_responses)
        impact_prob = vulnerability_dist.prob_matrix.T @ event_dist.prob
        return (
            ImpactDistrib(vulnerability_dist.event_type, vulnerability_dist.impact_bins, impact_prob),
            vulnerability_dist,
            event_dist,
        )

    def _check_event_type(self):
        if self.event_type not in self._event_types:
            raise NotImplementedError(f"model does not support events of type {self.event_type.__name__}")


class VulnerabilityModel(VulnerabilityModelBase):
    """A vulnerability model that requires only specification of distributions of impacts for given intensities,
    by implementing get_impact_curve."""

    def __init__(self, *, model: str = "", event_type: type, impact_bin_edges):
        self.event_type = event_type
        self.model = model
        self.impact_bin_edges = impact_bin_edges

    def get_event_data_requests(
        self, asset: Asset, *, scenario: str, year: int
    ) -> Union[EventDataRequest, Iterable[EventDataRequest]]:
        return EventDataRequest(
            self.event_type, asset.longitude, asset.latitude, scenario=scenario, year=year, model=self.model
        )

    def get_distributions(
        self, asset: Asset, event_data_responses: Iterable[EventDataResponse]
    ) -> Tuple[VulnerabilityDistrib, HazardEventDistrib]:

        (event_data,) = event_data_responses
        intensity_curve = ExceedanceCurve(1.0 / event_data.return_periods, event_data.intensities)
        intensity_bin_edges, probs = intensity_curve.get_probability_bins()
        intensity_bin_centres = (intensity_bin_edges[1:] + intensity_bin_edges[:-1]) / 2

        vul = VulnerabilityDistrib(
            self.event_type,
            intensity_bin_edges,
            self.impact_bin_edges,
            self.get_impact_curve(intensity_bin_centres, asset).to_prob_matrix(self.impact_bin_edges),
        )

        event = HazardEventDistrib(self.event_type, intensity_bin_edges, probs)
        return vul, event

    @abstractmethod
    def get_impact_curve(self, intensities, asset) -> ImpactCurve:
        """Get the impact curves for the intensities specified."""
        ...


class DeterministicVulnerabilityModel(VulnerabilityModelBase):
    """A vulnerability model that requires only specification of a damage/disruption curve.
    This simple model contains no uncertainty around damage/disruption."""

    def __init__(self, *, event_type: type, intensities, impacts, model=""):
        self.event_type = event_type
        self.intensities = intensities
        self.impacts = impacts
        self.model = model

    def get_event_data_requests(
        self, asset: Asset, *, scenario: str, year: int
    ) -> Union[EventDataRequest, Iterable[EventDataRequest]]:
        return EventDataRequest(
            self.event_type, asset.longitude, asset.latitude, scenario=scenario, year=year, model=self.model
        )

    def get_distributions(
        self, asset: Asset, event_data_responses: Iterable[EventDataResponse]
    ) -> Tuple[VulnerabilityDistrib, HazardEventDistrib]:

        (event_data,) = event_data_responses
        intensity_curve = ExceedanceCurve(1.0 / event_data.return_periods, event_data.intensities)
        intensity_bin_edges, probs = intensity_curve.get_probability_bins()

        impact_bins_edges = np.interp(intensity_bin_edges, self.intensities, self.impacts)

        vul = VulnerabilityDistrib(
            type(self.event_type), intensity_bin_edges, impact_bins_edges, np.eye(len(impact_bins_edges) - 1)
        )
        event = HazardEventDistrib(type(self.event_type), intensity_bin_edges, probs)
        return vul, event
