import json
import logging
from dataclasses import asdict, dataclass, field
from pathlib import Path
from typing import Optional

from idtrackerai.utils import create_dir, json_default


@dataclass
class NetworkParams:
    number_of_classes: int
    schedule: list[int]
    architecture: str
    model_name: str
    dataset: str
    image_size: list[int]
    optim_args: dict = field(default_factory=dict)
    epochs: int = 0
    optimizer: str = "SGD"
    loss: str = "CE"
    use_gpu: bool = True
    apply_mask: bool = False
    skip_eval: bool = False
    save_folder: Path = Path("")

    scopes_layers_to_optimize: Optional[list[str]] = field(default_factory=list)
    knowledge_transfer_folder: Path | None = None
    use_adam_optimiser: bool = False
    return_store_objects: bool = False
    restore_folder: Path = Path()

    @property
    def load_model_path(self) -> Path:
        # v5.0.0 compatibility
        v5_path = self.restore_folder / (self.model_file_name + "_.model.pth")
        if v5_path.is_file():
            return v5_path
        return self.restore_folder / (self.model_file_name + ".model.pth")

    @property
    def save_model_path(self) -> Path:
        return self.save_folder / self.model_file_name

    @property
    def model_file_name(self) -> str:
        return f"{self.dataset}_{self.model_name}"

    @property
    def knowledge_transfer_model_file(self) -> Path | None:
        if self.knowledge_transfer_folder is None:
            return None
        return (
            self.knowledge_transfer_folder
            / "supervised_identification_network.model.pth"
        )

    def save(self) -> None:
        path = self.save_folder / "model_params.json"
        logging.info(f"Saving NetworkParams at {path}")
        create_dir(self.save_folder)
        path.write_text(json.dumps(asdict(self), indent=4, default=json_default))
