"""Models for generator devices."""

from typing import Annotated, Any

from pint import Quantity
from pydantic import Field, NonNegativeFloat, field_serializer

from r2x.enums import PrimeMoversType, StorageTechs, ThermalFuels
from r2x.models.core import Device, InputOutput, MinMax, UpDown
from r2x.models.costs import HydroGenerationCost, RenewableGenerationCost, StorageCost, ThermalGenerationCost
from r2x.models.load import PowerLoad
from r2x.models.topology import ACBus
from r2x.units import (
    ActivePower,
    ApparentPower,
    Energy,
    Percentage,
    PowerRate,
    Time,
    VOMPrice,
    ureg,
)


class Generator(Device):
    """Abstract generator class."""

    bus: Annotated[ACBus, Field(description="Bus where the generator is connected.")] | None = None
    rating: Annotated[
        ApparentPower | None,
        Field(ge=0, description="Maximum output power rating of the unit (MVA)."),
    ] = ApparentPower(1, "MVA")
    active_power: Annotated[
        ActivePower,
        Field(
            description=(
                "Initial active power set point of the unit in MW. For power flow, this is the steady "
                "state operating point of the system."
            ),
        ),
    ] = ActivePower(0.0, "MW")
    reactive_power: Annotated[
        ApparentPower | None,
        Field(
            description=(
                "Reactive power set point of the unit in MW. For power flow, this is the steady "
                "state operating point of the system."
            ),
        ),
    ] = ApparentPower(0.0, "MVA")
    base_mva: float = 1
    base_power: Annotated[
        ApparentPower | None,
        Field(
            gt=0,
            description="Base power of the unit (MVA) for per unitization.",
        ),
    ] = None
    must_run: Annotated[int | None, Field(description="If we need to force the dispatch of the device.")] = (
        None
    )
    vom_price: Annotated[VOMPrice, Field(description="Variable operational price $/MWh.")] | None = None
    prime_mover_type: (
        Annotated[PrimeMoversType, Field(description="Prime mover technology according to EIA 923.")] | None
    ) = None
    unit_type: Annotated[
        PrimeMoversType | None, Field(description="Prime mover technology according to EIA 923.")
    ] = None
    min_rated_capacity: Annotated[ActivePower, Field(description="Minimum rated power generation.")] = (
        0 * ureg.MW
    )
    ramp_up: (
        Annotated[
            PowerRate,
            Field(description="Ramping rate on the positve direction."),
        ]
        | None
    ) = None
    ramp_down: (
        Annotated[
            PowerRate,
            Field(description="Ramping rate on the negative direction."),
        ]
        | None
    ) = None
    min_up_time: (
        Annotated[
            Time,
            Field(ge=0, description="Minimum up time in hours for UC decision."),
        ]
        | None
    ) = None
    min_down_time: (
        Annotated[
            Time,
            Field(ge=0, description="Minimum down time in hours for UC decision."),
        ]
        | None
    ) = None
    mean_time_to_repair: (
        Annotated[
            Time,
            Field(gt=0, description="Total hours to repair after outage occur."),
        ]
        | None
    ) = None
    forced_outage_rate: (
        Annotated[
            Percentage,
            Field(description="Expected level of unplanned outages in percent."),
        ]
        | None
    ) = None
    planned_outage_rate: (
        Annotated[
            Percentage,
            Field(description="Expected level of planned outages in percent."),
        ]
        | None
    ) = None
    startup_cost: Annotated[NonNegativeFloat, Field(description="Cost in $ of starting a unit.")] | None = (
        None
    )
    shutdown_cost: (
        Annotated[NonNegativeFloat, Field(description="Cost in $ of shuting down a unit.")] | None
    ) = None
    active_power_limits: Annotated[
        MinMax | None, Field(description="Maximum output power rating of the unit (MVA).")
    ] = None
    reactive_power_limits: Annotated[
        MinMax | None, Field(description="Maximum output power rating of the unit (MVA).")
    ] = None

    @field_serializer("active_power_limits")
    def serialize_active_power_limits(self, min_max: MinMax | dict | None) -> dict[str, Any] | None:
        if min_max is None:
            return None
        if not isinstance(min_max, MinMax):
            min_max = MinMax(**min_max)
        if min_max is not None:
            return {
                "min": min_max.min.magnitude if isinstance(min_max.min, Quantity) else min_max.min,
                "max": min_max.max.magnitude if isinstance(min_max.max, Quantity) else min_max.max,
            }


class RenewableGen(Generator):
    """Abstract class for renewable generators."""


class RenewableDispatch(RenewableGen):
    """Curtailable renewable generator.

    This type of generator have a hourly capacity factor profile.
    """

    power_factor: float = 1.0
    operation_cost: RenewableGenerationCost | None = None


class RenewableNonDispatch(RenewableGen):
    """Non-curtailable renewable generator.

    Renewable technologies w/o operational cost.
    """

    power_factor: float = 1.0


class HydroGen(Generator):
    """Hydroelectric generator."""


class HydroDispatch(HydroGen):
    """Class representing flexible hydro generators."""

    operation_cost: HydroGenerationCost | None = None
    ramp_up: (
        Annotated[
            PowerRate,
            Field(ge=0, description="Ramping rate on the positve direction."),
        ]
        | None
    ) = None
    ramp_down: (
        Annotated[
            PowerRate,
            Field(ge=0, description="Ramping rate on the negative direction."),
        ]
        | None
    ) = None


class HydroFix(HydroGen):
    """Class representing unflexible hydro."""


class HydroEnergyReservoir(HydroGen):
    """Class representing hydro system with reservoirs."""

    ramp_limits: UpDown | None = None
    time_limits: UpDown | None = None
    operation_cost: HydroGenerationCost | None = None
    inflow: float | None = None
    initial_energy: (
        Annotated[NonNegativeFloat, Field(description="Initial water volume or percentage.")] | None
    ) = 0
    initial_storage: float | None = None
    storage_capacity: (
        Annotated[
            Energy,
            Field(description="Total water volume or percentage."),
        ]
        | None
    ) = None
    min_storage_capacity: (
        Annotated[
            Energy,
            Field(description="Minimum water volume or percentage."),
        ]
        | None
    ) = None
    storage_target: (
        Annotated[
            Energy,
            Field(description="Maximum energy limit."),
        ]
        | None
    ) = None


class HydroPumpedStorage(HydroGen):
    """Class representing pumped hydro generators."""

    operation_cost: HydroGenerationCost | StorageCost | None = None
    storage_duration: (
        Annotated[
            Time,
            Field(description="Storage duration in hours."),
        ]
        | None
    ) = None
    initial_volume: (
        Annotated[Energy, Field(gt=0, description="Initial water volume or percentage.")] | None
    ) = None
    storage_capacity: Annotated[
        Energy,
        Field(gt=0, description="Total water volume or percentage."),
    ]
    min_storage_capacity: (
        Annotated[
            Energy,
            Field(description="Minimum water volume or percentage."),
        ]
        | None
    ) = None
    pump_efficiency: Annotated[Percentage, Field(ge=0, le=1, description="Pumping efficiency.")] | None = None
    pump_load: (
        Annotated[
            ActivePower,
            Field(description="Load related to the usage of the pump."),
        ]
        | None
    ) = None

    @classmethod
    def example(cls) -> "HydroPumpedStorage":
        return HydroPumpedStorage(
            name="HydroPumpedStorage",
            active_power=ActivePower(100, "MW"),
            pump_load=ActivePower(100, "MW"),
            bus=ACBus.example(),
            prime_mover_type=PrimeMoversType.PS,
            storage_duration=Time(10, "h"),
            storage_capacity=Energy(1000, "MWh"),
            min_storage_capacity=Energy(10, "MWh"),
            pump_efficiency=Percentage(85, "%"),
            initial_volume=Energy(500, "MWh"),
            ext={"description": "Pumped hydro unit with 10 hour of duration"},
        )


class ThermalGen(Generator):
    """Class representing fuel based thermal generator."""

    fuel: Annotated[str, Field(description="Fuel category")] | None = None
    operation_cost: ThermalGenerationCost | None = None


class ThermalStandard(ThermalGen):
    """Class representing a standard thermal generator."""

    status: bool = True
    ramp_limits: UpDown | None = None
    time_limits: UpDown | None = None
    fuel: ThermalFuels = ThermalFuels.OTHER

    @classmethod
    def example(cls) -> "ThermalStandard":
        return ThermalStandard(
            name="ThermalStandard",
            bus=ACBus.example(),
            fuel=ThermalFuels.NATURAL_GAS,
            active_power=100.0 * ureg.MW,
            ext={"Additional data": "Additional value"},
        )


class ThermalMultiStart(ThermalGen):
    """We will fill this class once we have the need for it."""


class Storage(Generator):
    """Default Storage class."""

    storage_duration: (
        Annotated[
            Time,
            Field(description="Storage duration in hours."),
        ]
        | None
    ) = None
    storage_capacity: Annotated[
        Energy,
        Field(description="Maximum allowed volume or state of charge."),
    ]
    initial_energy: Annotated[Percentage, Field(description="Initial state of charge.")] | None = None
    min_storage_capacity: Annotated[Percentage, Field(description="Minimum state of charge")] = Percentage(
        0, "%"
    )
    max_storage_capacity: Annotated[Percentage, Field(description="Minimum state of charge")] = Percentage(
        100, "%"
    )


class GenericBattery(Storage):
    """Battery energy storage model."""

    charge_efficiency: Annotated[Percentage, Field(ge=0, description="Charge efficiency.")] | None = None
    discharge_efficiency: Annotated[Percentage, Field(ge=0, description="Discharge efficiency.")] | None = (
        None
    )


class EnergyReservoirStorage(Storage):
    """Battery energy storage model."""

    storage_technology_type: StorageTechs
    operation_cost: StorageCost | None = None
    charge_efficiency: Annotated[Percentage, Field(ge=0, description="Charge efficiency.")] | None = None
    discharge_efficiency: Annotated[Percentage, Field(ge=0, description="Discharge efficiency.")] | None = (
        None
    )
    storage_level_limits: MinMax | None = None
    initial_storage_capacity_level: float
    input_active_power_limits: MinMax
    output_active_power_limits: MinMax
    efficiency: InputOutput


class HybridSystem(Device):
    """Representation of hybrid system with renewable generation, load, thermal generation and storage.

    This class is just a link between two components.
    For the implementation see:
    - https://github.com/NREL-Sienna/PowerSystems.jl/blob/main/src/models/HybridSystem.jl
    """

    storage_unit: Storage | None = None
    renewable_unit: RenewableGen | None = None
    thermal_unit: ThermalGen | None = None
    electric_load: PowerLoad | None = None
