"""Data models for energy sharing."""

from enum import Enum
from typing import Annotated, Any

import pandas as pd
from pydantic import BaseModel, Field, confloat

from openenergyid import TimeDataFrame

from .const import NET_INJECTION, NET_OFFTAKE, SHARED_ENERGY
from .data_formatting import create_multi_index_input_frame


class CalculationMethod(Enum):
    """Calculation method for energy sharing."""

    FIXED = "Fixed"
    RELATIVE = "Relative"
    OPTIMAL = "Optimal"


class KeyInput(TimeDataFrame):
    """Energy Sharing Keys."""

    data: Annotated[
        list[list[confloat(ge=0.0, le=1.0)]],  # type: ignore
        Field(
            description="Key data, column per participant. "
            "Must be between 0 and 1. "
            "Each row must sum to 1."
        ),
    ]

    def model_post_init(self, __context: Any) -> None:
        """Post-initialization validation."""
        for row in self.data:
            if round(sum(row), 3) != 1.0:
                raise ValueError("Each row must sum to 1.")
        return super().model_post_init(__context)


class EnergySharingInput(BaseModel):
    """Input data for energy sharing."""

    gross_injection: Annotated[
        TimeDataFrame,
        Field(alias="grossInjection", description="Gross injection data, column per participant"),
    ]
    gross_offtake: Annotated[
        TimeDataFrame,
        Field(alias="grossOfftake", description="Gross offtake data, column per participant"),
    ]
    key: KeyInput

    def to_pandas(self) -> pd.DataFrame:
        """Return the data as a combined DataFrame"""
        df = create_multi_index_input_frame(
            gross_injection=self.gross_injection.to_pandas(),
            gross_offtake=self.gross_offtake.to_pandas(),
            key=self.key.to_pandas(),
        )
        return df


class EnergySharingOutput(BaseModel):
    """Output data for energy sharing."""

    net_injection: TimeDataFrame = Field(alias="netInjection")
    net_offtake: TimeDataFrame = Field(alias="netOfftake")
    shared_energy: TimeDataFrame = Field(alias="sharedEnergy")

    @classmethod
    def from_calculation_result(cls, result: pd.DataFrame) -> "EnergySharingOutput":
        """Create an output model from a calculation result."""
        return cls.model_construct(
            net_injection=TimeDataFrame.from_pandas(result[NET_INJECTION]),
            net_offtake=TimeDataFrame.from_pandas(result[NET_OFFTAKE]),
            shared_energy=TimeDataFrame.from_pandas(result[SHARED_ENERGY]),
        )
