from pydantic import BaseModel, ValidationError
from typing import List
import random


class ModelOutput(BaseModel):

    @classmethod
    def classname(cls) -> str:
        """Return classname."""
        return cls.__name__

    @classmethod
    def list_fields(cls) -> List[str]:
        """List options that are stored as attributes."""
        return list(cls.model_fields.keys())

    @classmethod
    def from_random(cls):
        """Instantiate with random numbers."""
        kwargs = {field: random.random() for field in cls.list_fields()}
        return cls(**kwargs)

    @classmethod
    def from_choice(cls, option: str):
        """Instantiate from choice."""
        if option is None:
            raise ValidationError()
        assert isinstance(option, str)
        allowed_options_list = cls.list_fields()
        assert option in allowed_options_list, \
            f"{option} is not among allowed fields {allowed_options_list}"
        kwargs = {field: 0 for field in cls.list_fields()}
        kwargs[option] = 1
        return cls(**kwargs)

    def __repr__(self) -> str:
        model_dict = self.model_dump()
        model_repr_str = f"{self.classname()}("
        model_repr_str += ", ".join([
            f"{field}={value:.3f}"
            for field, value
            in model_dict.items()
        ])
        model_repr_str += ")"
        return model_repr_str

    def highest_score_field(self) -> str:
        """Return name of field with highest score."""
        model_dict = self.model_dump()
        return max(model_dict, key=lambda k: model_dict[k])

    def highest_score_value(self) -> float:
        """Return value of field with highest score."""
        model_dict = self.model_dump()
        return max(model_dict.values())


class InformationValueModelOutput(ModelOutput):
    given_new: float
    ideal_real: float
    central_marginal: float


class FramingModelOutput(ModelOutput):
    frame_lines: float
    empty_space: float
    colour_contrast: float
    form_contrast: float


class SalienceModelOutput(ModelOutput):
    size: float
    colour: float
    tone: float
    form: float
    positioning: float


if __name__ == '__main__':
    m = InformationValueModelOutput.from_random()
    print(repr(m))
    print(m.highest_score_field())
    print(m.highest_score_value())
    m = FramingModelOutput.from_random()
    print(repr(m))
    print(m.highest_score_field())
    print(m.highest_score_value())
    m = SalienceModelOutput.from_random()
    print(repr(m))
    print(m.highest_score_field())
    print(m.highest_score_value())