import json
import logging
from pathlib import Path

from pydantic import BaseModel, Field
import yaml

from pilz.model.pilz import Pilz, Pilze

logger = logging.getLogger(__name__)


class EvalSettings(BaseModel):
    in_folders: list[Path] = Field(description="Name where the pilze was stored")
    out_folder: Path = Field(description="Name where the evaluation is stored")
    max_parallel_where: int = Field(
        description="Number of points for roc curve", default=100
    )


class TrainSettings(BaseModel):
    n: int = Field(description="Number of Pilze per target value", default=1)
    out_folder: Path = Field(description="Name where the pilz will be stored")
    n_cat: int = Field(description="Number of categories in one featerue", default=3)
    max_depth: int = Field(description="Max depth of the pilz", default=10)

    max_eval_fit: int = Field(
        description="Wieviel sollen max für training gleichzeitig benutzt werden",
        default=1000,
    )
    min_eval_fit: int = Field(
        description="Wieviel müssen fürs training mindestens benutzt werden", default=10
    )
    n_dims: int = Field(
        description="Wieviele dimensionen sollen auf einmal benutzt werden", default=3
    )
    calcs_per_dim: int | None = Field(
        description="Wieviele berechnungen sollen pro dimension gemacht werden",
        default=5000,
    )
    frac_eval_cat: float = Field(
        description="Wie groß ist der anteil der für eval benutzt werden soll",
        default=0.5,
    )

    neutral_faktor: float = Field(
        description="Faktor wie groß der Unterschied für Neutral sein soll. Wert kleiner als 0 beduet kein Neutral state",
        default=-1.0,
    )

    def pilz_exists(self, target: str | int, n: int) -> bool:
        if self.out_folder.exists():
            if (self.out_folder / f"{target}" / f"{n}.json").exists():
                return True
        return False

    def save_pilz(self, pilz: Pilz, n: int):
        target = pilz.target
        out_folder = self.out_folder / f"{target}"
        logger.info(f"save pilz {n} for target {target} to {out_folder}")
        out_folder.mkdir(parents=True, exist_ok=True)
        out_file = out_folder / f"{n}.json"
        with out_file.open("w") as f:
            f.write(json.dumps(pilz.model_dump(), ensure_ascii=False, indent=2))
