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 ..api.v1.common import VulnerabilityCurves
from .assets import Asset
from .curve import ExceedanceCurve
from .hazard_event_distrib import HazardEventDistrib
from .hazard_model import HazardDataRequest, HazardDataResponse, HazardEventDataResponse
from .impact_distrib import ImpactDistrib
from .vulnerability_distrib import VulnerabilityDistrib
from .vulnerability_matrix_provider import VulnMatrixProvider

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):
    def __init__(self, model: str, hazard_type: type):
        self.model = model
        self.hazard_type = hazard_type
        self._event_types: List[type] = []
        self._asset_types: List[type] = []

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

    @abstractmethod
    def get_impact(self, asset: Asset, event_data: Iterable[HazardDataResponse]) -> ImpactDistrib:
        ...


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

    def __init__(self, model: str, hazard_type: type):
        super().__init__(model, hazard_type)

    @abstractmethod
    def get_distributions(
        self, asset: Asset, event_data_responses: Iterable[HazardDataResponse]
    ) -> 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[HazardDataResponse]):
        impact, _, _ = self.get_impact_details(asset, event_data_responses)
        return impact

    def get_impact_details(
        self, asset: Asset, event_data_responses: Iterable[HazardDataResponse]
    ) -> 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.hazard_type not in self._event_types:
            raise NotImplementedError(f"model does not support events of type {self.hazard_type.__name__}")


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

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

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

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

        (event_data,) = event_data_responses
        assert isinstance(event_data, HazardEventDataResponse)

        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.hazard_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.hazard_type, intensity_bin_edges, probs)
        return vul, event

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


class DeterministicVulnerabilityModel(VulnerabilityModelAcuteBase):
    """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_data_requests(
        self, asset: Asset, *, scenario: str, year: int
    ) -> Union[HazardDataRequest, Iterable[HazardDataRequest]]:
        return HazardDataRequest(
            self.event_type, asset.longitude, asset.latitude, scenario=scenario, year=year, model=self.model
        )

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

        (event_data,) = event_data_responses
        assert isinstance(event_data, HazardEventDataResponse)

        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
