from enum import Enum
from typing import Optional, List


class IMType(Enum):
    """Available IMs to use"""

    PGA = "PGA"
    PGV = "PGV"
    pSA = "pSA"
    CAV = "CAV"
    AI = "AI"
    ASI = "ASI"
    DSI = "DSI"
    SI = "SI"
    Ds575 = "Ds575"
    Ds595 = "Ds595"

    def __str__(self):
        """Ensures that we return PGA not IMType.PGA"""
        return self.name

    @classmethod
    def has_value(cls, value: str):
        """Checks if the value is in the IMType set"""
        available_names = set(im.name for im in IMType)
        if value.startswith("pSA"):
            value = "pSA"
        return value in available_names


class IMComponent(Enum):
    """Available components to use"""

    RotD50 = "RotD50"
    RotD100 = "RotD100"
    Larger = "Larger"

    def __str__(self):
        """Ensures that we return RotD50 not IMComponent.RotD50"""
        return self.name


class IM:
    """
    Represents an IM to use for calculations
    """

    def __init__(
        self,
        im_type: IMType,
        period: Optional[float] = None,
        component: Optional[IMComponent] = IMComponent.RotD50,
    ):
        self.im_type = im_type
        self.period = period
        self.component = component

        if im_type == IMType.pSA and period is None:
            raise ValueError("Creation of pSA IM does not have a specified period")

    @classmethod
    def from_str(cls, im_string: str, im_component: Optional[str] = None):
        """Converts a given string to an IM object"""
        period = None
        if im_string.startswith("pSA") and "_" in im_string:
            im_string, period = im_string.split("_")
            period = float(period.replace("p", "."))

        return (
            cls(IMType[im_string], period)
            if im_component is None
            else cls(IMType[im_string], period, IMComponent(im_component))
        )

    def __str__(self):
        """Overrides the string method by just
        returning the name instead of the object"""
        if self.period:
            return f"{self.im_type}_{self.period}"
        else:
            return f"{self.im_type}"

    def __hash__(self):
        return hash((self.im_type, self.period, self.component))

    def __eq__(self, other):
        return (self.im_type, self.period, self.component) == (
            other.im_type,
            other.period,
            other.component,
        )

    def __ne__(self, other):
        return not self.__eq__(other)

    def file_format(self):
        """
        Outputs the normal str version of the IM
        Except does a replace on the period for p instead of . for saving in file formats
        """
        return str(self).replace(".", "p")

    def is_pSA(self):
        """Returns True if IM is of type pSA otherwise False"""
        return self.im_type == IMType.pSA


def to_string_list(IMs: List[IM]):
    """Converts a list of IM Objects to their string form"""
    return [str(im) for im in IMs]


def to_im_list(IMs: List[str]):
    """Converts a list of string to IM Objects"""
    return [IM.from_str(im) for im in IMs]


# Available IM Components per IMType
IM_COMPONENT_MAPPING = {
    IMType.PGA: [IMComponent.RotD50, IMComponent.RotD100, IMComponent.Larger],
    IMType.pSA: [IMComponent.RotD50, IMComponent.RotD100, IMComponent.Larger],
    IMType.PGV: [IMComponent.RotD50],
    IMType.AI: [IMComponent.RotD50],
    IMType.ASI: [IMComponent.RotD50],
    IMType.CAV: [IMComponent.RotD50],
    IMType.Ds575: [IMComponent.RotD50],
    IMType.Ds595: [IMComponent.RotD50],
    IMType.DSI: [IMComponent.RotD50],
    IMType.SI: [IMComponent.RotD50],
}

# Default pSA periods
DEFAULT_PSA_PERIODS = [
    0.01,
    0.02,
    0.03,
    0.04,
    0.05,
    0.075,
    0.1,
    0.12,
    0.15,
    0.17,
    0.2,
    0.25,
    0.3,
    0.4,
    0.5,
    0.6,
    0.7,
    0.75,
    0.8,
    0.9,
    1.0,
    1.25,
    1.5,
    2.0,
    2.5,
    3.0,
    4.0,
    5.0,
    6.0,
    7.5,
    10.0,
]

from enum import Enum
from typing import Optional, List


class IMType(Enum):
    """Available IMs to use"""

    PGA = "PGA"
    PGV = "PGV"
    pSA = "pSA"
    CAV = "CAV"
    AI = "AI"
    ASI = "ASI"
    DSI = "DSI"
    SI = "SI"
    Ds575 = "Ds575"
    Ds595 = "Ds595"

    def __str__(self):
        """Ensures that we return PGA not IMType.PGA"""
        return self.name

    @classmethod
    def has_value(cls, value: str):
        """Checks if the value is in the IMType set"""
        available_names = set(im.name for im in IMType)
        if value.startswith("pSA"):
            value = "pSA"
        return value in available_names


class IMComponent(Enum):
    """Available components to use"""

    RotD50 = "RotD50"
    RotD100 = "RotD100"
    Larger = "Larger"

    def __str__(self):
        """Ensures that we return RotD50 not IMComponent.RotD50"""
        return self.name


class IM:
    """
    Represents an IM to use for calculations
    """

    def __init__(
        self,
        im_type: IMType,
        period: Optional[float] = None,
        component: Optional[IMComponent] = IMComponent.RotD50,
    ):
        self.im_type = im_type
        self.period = period
        self.component = component

        if im_type == IMType.pSA and period is None:
            raise ValueError("Creation of pSA IM does not have a specified period")

    @classmethod
    def from_str(cls, im_string: str, im_component: Optional[str] = None):
        """Converts a given string to an IM object"""
        period = None
        if im_string.startswith("pSA") and "_" in im_string:
            im_string, period = im_string.split("_")
            period = float(period.replace("p", "."))

        return (
            cls(IMType[im_string], period)
            if im_component is None
            else cls(IMType[im_string], period, IMComponent(im_component))
        )

    def __str__(self):
        """Overrides the string method by just
        returning the name instead of the object"""
        if self.period:
            return f"{self.im_type}_{self.period}"
        else:
            return f"{self.im_type}"

    def __hash__(self):
        return hash((self.im_type, self.period, self.component))

    def __eq__(self, other):
        return (self.im_type, self.period, self.component) == (
            other.im_type,
            other.period,
            other.component,
        )

    def __ne__(self, other):
        return not self.__eq__(other)

    def file_format(self):
        """
        Outputs the normal str version of the IM
        Except does a replace on the period for p instead of . for saving in file formats
        """
        return str(self).replace(".", "p")

    def is_pSA(self):
        """Returns True if IM is of type pSA otherwise False"""
        return self.im_type == IMType.pSA


def to_string_list(IMs: List[IM]):
    """Converts a list of IM Objects to their string form"""
    return [str(im) for im in IMs]


def to_im_list(IMs: List[str]):
    """Converts a list of string to IM Objects"""
    return [IM.from_str(im) for im in IMs]


# Available IM Components per IMType
IM_COMPONENT_MAPPING = {
    IMType.PGA: [IMComponent.RotD50, IMComponent.RotD100, IMComponent.Larger],
    IMType.pSA: [IMComponent.RotD50, IMComponent.RotD100, IMComponent.Larger],
    IMType.PGV: [IMComponent.RotD50],
    IMType.AI: [IMComponent.RotD50],
    IMType.ASI: [IMComponent.RotD50],
    IMType.CAV: [IMComponent.RotD50],
    IMType.Ds575: [IMComponent.RotD50],
    IMType.Ds595: [IMComponent.RotD50],
    IMType.DSI: [IMComponent.RotD50],
    IMType.SI: [IMComponent.RotD50],
}

# Default pSA periods
DEFAULT_PSA_PERIODS = [
    0.01,
    0.02,
    0.03,
    0.04,
    0.05,
    0.075,
    0.1,
    0.12,
    0.15,
    0.17,
    0.2,
    0.25,
    0.3,
    0.4,
    0.5,
    0.6,
    0.7,
    0.75,
    0.8,
    0.9,
    1.0,
    1.25,
    1.5,
    2.0,
    2.5,
    3.0,
    4.0,
    5.0,
    6.0,
    7.5,
    10.0,
]

Available IMs to use Ensures that we return PGA not IMType.PGA Checks if the value is in the IMType set Available components to use Ensures that we return RotD50 not IMComponent.RotD50 Represents an IM to use for calculations Converts a given string to an IM object Overrides the string method by just         returning the name instead of the object Outputs the normal str version of the IM         Except does a replace on the period for p instead of . for saving in file formats Returns True if IM is of type pSA otherwise False Converts a list of IM Objects to their string form Converts a list of string to IM Objects # Available IM Components per IMType # Default pSA periods