import marimo

__generated_with = "0.16.5"
app = marimo.App(width="medium")


@app.cell
def _():
    import marimo as mo
    return


@app.cell
def _():
    from pif_compiler.services.srv_echa import orchestrator
    return (orchestrator,)


@app.cell
def _(orchestrator):
    result = orchestrator("57-55-6")
    return (result,)


@app.cell
def _(result):
    test = result['repeated_dose_toxicity']
    return (test,)


@app.cell
def _(result):
    acute = result['acute_toxicity']
    acute
    return (acute,)


@app.cell
def _():
    from pif_compiler.services.srv_echa import extract_levels, at_extractor, rdt_extractor
    return at_extractor, extract_levels, rdt_extractor


@app.cell
def _(acute, at_extractor, extract_levels):
    at = extract_levels(acute, at_extractor)
    return (at,)


@app.cell
def _(at):
    at
    return


@app.cell
def _(extract_levels, rdt_extractor, test):
    rdt = extract_levels(test, rdt_extractor)
    return (rdt,)


@app.cell
def _(rdt):
    rdt
    return


@app.cell
def _():
    from pydantic import BaseModel
    from typing import Optional
    return BaseModel, Optional


@app.cell
def _(BaseModel, Optional):
    class ToxIndicator(BaseModel):
        indicator : str
        value : int
        unit : str
        route : str
        toxicity_type : Optional[str] = None
        ref : Optional[str] = None

        @property
        def priority_rank(self):
            """Returns the numerical priority based on the toxicological indicator."""
            mapping = {
                'LD50': 1,
                'DL50': 1,
                'NOAEC': 2,
                'LOAEL': 3,
                'NOAEL': 4
            }
            return mapping.get(self.indicator, 99)

        @property
        def factor(self):
            """Returns the factor based on the toxicity type."""
            if self.priority_rank == 1:
                return 10
            elif self.priority_rank == 3:
                return 3
            return 1
    return (ToxIndicator,)


@app.cell
def _(ToxIndicator, rdt):
    lista = []
    for i in rdt:
        tesss = rdt.get(i)
        t = ToxIndicator(**tesss)
        lista.append(t)

    lista
    return


@app.cell
def _():
    from pydantic import model_validator
    return (model_validator,)


@app.cell
def _(
    BaseModel,
    Optional,
    ToxIndicator,
    at_extractor,
    extract_levels,
    model_validator,
    rdt_extractor,
):
    class Toxicity(BaseModel):
        cas: str
        indicators: list[ToxIndicator]
        best_case: Optional[ToxIndicator] = None
        factor: Optional[int] = None

        @model_validator(mode='after')
        def set_best_case(self) -> 'Toxicity':
            if self.indicators:
                self.best_case = max(self.indicators, key=lambda x: x.priority_rank)
                self.factor = self.best_case.factor
            return self

        @classmethod
        def from_result(cls, cas: str, result):
            toxicity_types = ['repeated_dose_toxicity', 'acute_toxicity']
            indicators_list = []

            for tt in toxicity_types:
                if tt not in result:
                    continue
                
                try:
                    extractor = at_extractor if tt == 'acute_toxicity' else rdt_extractor
                    fetch = extract_levels(result[tt], extractor=extractor)
                
                    link = result.get(f"{tt}_link", "")

                    for key, lvl in fetch.items():
                        lvl['ref'] = link 
                        elem = ToxIndicator(**lvl)
                        indicators_list.append(elem)
                    
                except Exception as e:
                    print(f"Errore durante l'estrazione di {tt}: {e}")
                    continue
                
            return cls(
                cas=cas,
                indicators=indicators_list
            )
    return (Toxicity,)


@app.cell
def _(Toxicity, result):
    tox = Toxicity.from_result("57-55-6", result)
    tox
    return


@app.cell
def _(result):
    result
    return


if __name__ == "__main__":
    app.run()