from enum import Enum

from worker.data.serialization import serialization
from worker.data.unit_conversions import KG_M3_to_PPG


@serialization
class HoleType(Enum):
    CASED_HOLE = "Cased Hole"
    OPEN_HOLE = "Open Hole"


@serialization
class PipeType(Enum):
    Unknown = "Unknown"
    DP = "Drillpipe"
    HWDP = "Heavy Weight Drillpipe"
    DC = "Drill Collar"
    STABILIZER = "Stabilizer"
    AGITATOR = "Agitator"
    PDM = "Positive Displacement Motor"
    MWD = "Measurement While Drilling"
    RSS = "Rotary Steerable System"
    BIT = "Bit"

    @staticmethod
    def determine_type(family: str):
        family = family.lower()
        mapping = {
            'dp': PipeType.DP,
            'hwdp': PipeType.HWDP,
            'dc': PipeType.DC,
            'stabilizer': PipeType.STABILIZER,
            'agitator': PipeType.AGITATOR,
            'pdm': PipeType.PDM,
            'mwd': PipeType.MWD,
            'rss': PipeType.RSS,
            'bit': PipeType.BIT
        }
        _type = mapping.get(family)
        if not _type:
            _type = PipeType.Unknown

        return _type


@serialization
class PipeMaterial(Enum):
    Steel = {
        'density': 7850 * KG_M3_to_PPG,
        'modulus_of_elasticity': 29.9e6,
        'poisson_ratio': 0.30
    }
    Aluminium = {
        'density': 2700 * KG_M3_to_PPG,
        'modulus_of_elasticity': 10.5e6,
        'poisson_ratio': 0.32
    }
    Titanium = {
        'density': 4506 * KG_M3_to_PPG,
        'modulus_of_elasticity': 16.5e6,
        'poisson_ratio': 0.31
    }
    NonMagnetic = {
        'density': 7850 * KG_M3_to_PPG,
        'modulus_of_elasticity': 29.9e6,
        'poisson_ratio': 0.30
    }

    def get_density(self) -> float:
        """
        Get the density of the pipe material
        :return: pipe density PPG
        """
        return self['density']

    def get_modulus_of_elasticity(self) -> float:
        """
        Get the pipe modulus of elasticity
        :return: modulus of elasticity in PSI
        """
        return self['modulus_of_elasticity']

    def get_poisson_ratio(self) -> float:
        """
        Get the Poisson ratio
        :return: Poisson ratio
        """
        return self['poisson_ratio']
